diff --git a/config.py b/config.py index 9ca78d5..52ff09c 100644 --- a/config.py +++ b/config.py @@ -8,7 +8,7 @@ class Config: self = format else: self.data_path: str = "" - self.pg_config = None + self.pg_dsn = None self.sort_by_match_strength: bool = False def __repr__(self): @@ -63,11 +63,9 @@ class Config: self.data_path = data_path def set_pg_config(self): - """Determine if data should be associated with a PostgreSQL instance, and, if so, record the required connection info""" + """Determine if data should be associated with a PostgreSQL instance, and, if so, record the required connection info as dsn string""" elect_for_pg = input("Connect this program to a PostgreSQL instance? y/n ").lower() if elect_for_pg == "y": - self.pg_config = {} - print("Please provide the relevant connection info. This sensitive information will only be recorded locally on your own machine.") dbname = input("Provide the DB name: ") user = input("Provide the username: ") @@ -91,7 +89,7 @@ class Config: conn.close() print("Connection successful!") - self.pg_config['dsn'] = dsn + self.pg_dsn = dsn except psycopg2.Error as e: print("There was an error connecting: ") print(e) diff --git a/pgutil.py b/pgutil.py new file mode 100644 index 0000000..e3e87b6 --- /dev/null +++ b/pgutil.py @@ -0,0 +1,63 @@ +import psycopg2, json +from config import Config + +class PGUTIL: + def __init__(self, config_file: Config = None): + if config_file is None: + load_config = Config() + config_file = load_config.run() + + pg_dsn = config_file['pg_dsn'] + + if pg_dsn is None: + raise Exception("Insufficient data to establish PostgreSQL connection.") + + self.config_file = config_file + + def create_tables(self): + pg_dsn = self.config_file['pg_dsn'] + + print("Connecting to PostgreSQL...") + + # establish connection and create cursor + conn = psycopg2.connect(pg_dsn) + cur = conn.cursor() + + # create base tables + create_label_table = """ + CREATE TABLE IF NOT EXISTS label ( + id INT PRIMARY KEY, + name varchar + ); + """ + + create_photo_table = """ + CREATE TABLE IF NOT EXISTS photo ( + id INT PRIMARY KEY, + path varchar, + label varchar, + matchstrength decimal, + labelid INT REFERENCES label(id) + ); + """ + + # attempt to insert new tables + print("Creating tables...") + + try: + cur.execute(create_label_table) + cur.execute(create_photo_table) + except psycopg2.Error: + raise Exception("Problem executing database table creation") + + print("Success!") + + # commit changes and close connection + conn.commit() + cur.close() + conn.close() + + print("Connection closed.") + + def insert_data(self): + pass \ No newline at end of file diff --git a/postgresutil.py b/postgresutil.py deleted file mode 100644 index 48bc63a..0000000 --- a/postgresutil.py +++ /dev/null @@ -1,29 +0,0 @@ -import psycopg2 -from config import Config - -def postgresutil(config: Config): - pg_config = config.pg_config - - if pg_config is None: - raise Exception("Insufficient data to establish PostgreSQL connection.") - - conn = psycopg2.connect(pg_config.dsn) - - # TO DO: script to create these tables and interact with them - - """ - CREATE TABLE IF NOT EXISTS label ( - id INT PRIMARY KEY, - name varchar - ); - - CREATE TABLE IF NOT EXISTS photo ( - id INT PRIMARY KEY, - path varchar, - label varchar, - matchstrength decimal, - labelid INT REFERENCES label(id) - ); - """ - - return conn \ No newline at end of file