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" cleanserver_db = get_parameter("cleanserver_db", config_filepath) cleanserver_db_user = get_parameter("cleanserver_db_user", config_filepath) return (config_filepath, cleanserver_db, cleanserver_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=cleanserver_db_user, host='', password='') conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) cur = conn.cursor() print("Creating database " + cleanserver_db + ". Please wait...") cur.execute(sql.SQL("CREATE DATABASE {}").format( sql.Identifier(cleanserver_db)) ) cur.execute(sql.SQL('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')) print(f'Database {cleanserver_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 = cleanserver_db, user = cleanserver_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("cleanserver 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 cleanserver parameters...") print("\n") cleanserver_db = input("cleanserver db name: ") cleanserver_db_user = input("cleanserver db user: ") with open(file_path, "w") as text_file: print("cleanserver_db: {}".format(cleanserver_db), file=text_file) print("cleanserver_db_user: {}".format(cleanserver_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, cleanserver_db, cleanserver_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 = cleanserver_db db_user = cleanserver_db_user table = "peers" sql = "create table "+table+" (server varchar(200) PRIMARY KEY, updated_at timestamptz, saved_at timestamptz, alive boolean, checked boolean, downs int)" create_table(db, db_user, table, sql) table = "totals" sql = "create table "+table+" (datetime timestamptz PRIMARY KEY, alive int, dead int, total int)" create_table(db, db_user, table, sql) print(f'Done!\nNow you can run setup.py\n')