diff --git a/db-setup.py b/db-setup.py index ff03421..ed4ed85 100644 --- a/db-setup.py +++ b/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 -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) +if __name__ == '__main__': -############################################################ -# create database -############################################################ + config_filepath, feeds_db, feeds_db_user, feeds_url = get_dbconfig() -conn = None + create_error = create_db() -try: + if create_error == '': - conn = psycopg2.connect(dbname='postgres', - user=feeds_db_user, host='', - password='') + check_db_conn() - 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( - sql.Identifier(feeds_db)) - ) - print("Database " + feeds_db + " created!") + else: -except (Exception, psycopg2.DatabaseError) as error: + os.remove(config_filepath) - print(error) + sys.exit() -finally: + ############################################################ + # Create needed tables + ############################################################ - if conn is not None: + db = feeds_db + db_user = feeds_db_user + table = "id" + sql = f'create table {table} (toot_id bigint PRIMARY KEY, tweet_id bigint)' - conn.close() + create_table(db, db_user, table, sql) -############################################################################################# - -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)" -create_table(db, db_user, table, sql) - -##################################### - -print("Done!") -print("Now you can run setup.py!") -print("\n") + print(f'Done!\nNow you can run setup.py\n')