import getpass import os 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" replicator_db = get_parameter("replicator_db", config_filepath) replicator_db_user = get_parameter("replicator_db_user", config_filepath) return (config_filepath, replicator_db, replicator_db_user) 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) write_parameter( parameter, file_path ) #sys.exit(0) # Find parameter in file with open( file_path ) as f: for line in f: if line.startswith( parameter ): return line.replace(parameter + ":", "").strip() # Cannot find parameter, exit print(file_path + " Missing parameter %s "%parameter) sys.exit(0) def create_db(): create_error = '' conn = None try: conn = psycopg2.connect(dbname='postgres', user=replicator_db_user, host='', password='') conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) cur = conn.cursor() print("Creating database " + replicator_db + ". Please wait...") cur.execute(sql.SQL("CREATE DATABASE {}").format( sql.Identifier(replicator_db)) ) cur.execute(sql.SQL('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')) print(f'Database {replicator_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 = replicator_db, user = replicator_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("replicator 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 replicator parameters...") print("\n") replicator_db = input("replicator db name: ") replicator_db_user = input("replicator db user: ") with open(file_path, "w") as text_file: print("replicator_db: {}".format(replicator_db), file=text_file) print("replicator_db_user: {}".format(replicator_db_user), file=text_file) def create_table(db, db_user, table, sql): conn = None try: conn = psycopg2.connect(database = db, user = db_user, password = "", host = "/var/run/postgresql", port = "5432") cur = conn.cursor() print(f'Creating table.. {table}') # Create the table in PostgreSQL database cur.execute(sql) conn.commit() print(f'Table {table} created!\n') except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() def write_dbconfig(db, db_user, table, sql): conn = None try: conn = psycopg2.connect(database = db, user = db_user, password = "", host = "/var/run/postgresql", port = "5432") cur = conn.cursor() print(f'Saving dbconfig to {table}') cur.execute(sql, (db, db_user)) conn.commit() print(f'dbconfig params saved to {table}!\n') except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() ############################################################################################# # main if __name__ == '__main__': config_filepath, replicator_db, replicator_db_user = get_dbconfig() create_error = create_db() if create_error == '': check_db_conn() else: if create_error == '42P04': sys.exit() else: os.remove(config_filepath) sys.exit() ############################################################ # Create needed tables ############################################################ db = replicator_db db_user = replicator_db_user table = "mastodon_id" sql = f'create table {table} (toot_id bigint PRIMARY KEY, tweet_id bigint)' create_table(db, db_user, table, sql) table = "pleroma_id" sql = f'create table {table} (toot_id varchar(18) PRIMARY KEY, tweet_id bigint)' create_table(db, db_user, table, sql) table = "maccounts" sql = f'create table {table} (bot_id uuid, hostname varchar(20), username varchar(45), client_id varchar(45), client_secret varchar(45), client_token varchar(45), hostname_soft varchar(10), active boolean default True, PRIMARY KEY (bot_id))' create_table(db, db_user, table, sql) table = "taccounts" sql = f'create table {table} (bot_id uuid, username varchar(15), PRIMARY KEY (bot_id))' create_table(db, db_user, table, sql) table = "tsecrets" sql = f'create table {table} (api_key varchar(25) PRIMARY KEY, api_key_secret varchar(50))' create_table(db, db_user, table, sql) print(f'Done!\nNow you can run setup.py\n')