refactored and added some error handles

db-setup.py

@@ -4,9 +4,20 @@ import sys
 import psycopg2
 from psycopg2 import sql
 from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
+import pdb
+
+def get_dbconfig():
+
+    # Load configuration from config file
+    config_filepath = "config/db_config.txt"
+    feeds_db = get_parameter("feeds_db", config_filepath)
+    feeds_db_user = get_parameter("feeds_db_user", config_filepath)
+    feeds_url = get_parameter("feeds_url", config_filepath)
+
+    return (config_filepath, feeds_db, feeds_db_user, feeds_url)
 
-# Returns the parameter from the specified file
 def get_parameter( parameter, file_path ):
+
     # Check if secrets file exists
     if not os.path.isfile(file_path):
         print("File %s not found, asking."%file_path)
@@ -23,7 +34,67 @@ def get_parameter( parameter, file_path ):
     print(file_path + "  Missing parameter %s "%parameter)
     sys.exit(0)
 
+def create_db():
+
+    create_error = ''
+
+    conn = None
+
+    try:
+
+      conn = psycopg2.connect(dbname='postgres',
+          user=feeds_db_user, host='',
+          password='')
+
+      conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+
+      cur = conn.cursor()
+
+      print("Creating database " + feeds_db + ". Please wait...")
+
+      cur.execute(sql.SQL("CREATE DATABASE {}").format(
+              sql.Identifier(feeds_db))
+          )
+      print("Database " + feeds_db + " created!")
+
+    except (Exception, psycopg2.DatabaseError) as error:
+
+        create_error = error.pgcode
+
+        sys.stdout.write(f'\n{str(error)}\n')
+
+    finally:
+
+        if conn is not None:
+
+            conn.close()
+
+    return create_error
+
+def check_db_conn():
+
+    try:
+
+        conn = None
+
+        conn = psycopg2.connect(database = feeds_db, user = feeds_db_user, password = "", host = "/var/run/postgresql", port = "5432")
+
+    except (Exception, psycopg2.DatabaseError) as error:
+
+        print(error)
+
+        os.remove(config_filepath)
+
+        sys.exit('Exiting. Run db-setup again with right parameters')
+
+    if conn is not None:
+
+        print("\n")
+        print("mastotuit parameters saved to db-config.txt!")
+        print("\n")
+
 def write_parameter( parameter, file_path ):
+
     if not os.path.exists('config'):
         os.makedirs('config')
     print("Setting up mastotuit parameters...")
@@ -64,84 +135,39 @@ def create_table(db, db_user, table, sql):
       conn.close()
 
 #############################################################################################
+# main
 
-# Load configuration from config file
+if __name__ == '__main__':
-config_filepath = "config/db_config.txt"
-feeds_db = get_parameter("feeds_db", config_filepath)
-feeds_db_user = get_parameter("feeds_db_user", config_filepath)
-feeds_url = get_parameter("feeds_url", config_filepath)
 
-############################################################
+    config_filepath, feeds_db, feeds_db_user, feeds_url = get_dbconfig()
-# create database
-############################################################
 
-conn = None
+    create_error = create_db()
 
-try:
+    if create_error == '':
 
-  conn = psycopg2.connect(dbname='postgres',
+        check_db_conn()
-      user=feeds_db_user, host='',
-      password='')
 
-  conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+    else:
 
-  cur = conn.cursor()
+        if create_error == '42P04':
 
-  print("Creating database " + feeds_db + ". Please wait...")
+            sys.exit()
 
-  cur.execute(sql.SQL("CREATE DATABASE {}").format(
+        else:
-          sql.Identifier(feeds_db))
-      )
-  print("Database " + feeds_db + " created!")
 
-except (Exception, psycopg2.DatabaseError) as error:
+            os.remove(config_filepath)
 
-  print(error)
+            sys.exit()
-
-finally:
-
-  if conn is not None:
-
-    conn.close()
-
-#############################################################################################
-
-try:
-
-  conn = None
-  conn = psycopg2.connect(database = feeds_db, user = feeds_db_user, password = "", host = "/var/run/postgresql", port = "5432")
-
-except (Exception, psycopg2.DatabaseError) as error:
-
-  print(error)
-  # Load configuration from config file
-  os.remove("db_config.txt")
-  print("Exiting. Run db-setup again with right parameters")
-  sys.exit(0)
-
-if conn is not None:
-
-  print("\n")
-  print("mastotuit parameters saved to db-config.txt!")
-  print("\n")
 
     ############################################################
     # Create needed tables
     ############################################################
 
-print("Creating table...")
-
-########################################
-
     db = feeds_db
     db_user = feeds_db_user
-
     table = "id"
-sql = "create table "+table+" (toot_id bigint PRIMARY KEY, tweet_id bigint)"
+    sql = f'create table {table} (toot_id bigint PRIMARY KEY, tweet_id bigint)'
+
     create_table(db, db_user, table, sql)
 
-#####################################
+    print(f'Done!\nNow you can run setup.py\n')
-
-print("Done!")
-print("Now you can run setup.py!")
-print("\n")