diff --git a/dbHandler.py b/dbHandler.py index d59013e..0d17ba6 100644 --- a/dbHandler.py +++ b/dbHandler.py @@ -1,16 +1,21 @@ +import typing import psycopg2 +from psycopg2 import sql -def printNotice(dbConnection, i): +def debugPrint(msg: str) -> None: + print(msg) + +def debugPrintNotice(dbConnection: psycopg2.extensions.connection, i: int) -> None: print("(DB HANDLER) " + dbConnection.notices[i]) -def connect(**options): - return psycopg2.connect(database=options["database"], - host=options["host"], - user=options["user"], - password=options["password"], - port=options["port"]) +def connect(databaseOption: str, hostOption: str, userOption: str, passwordOption: str, portOption: str) -> psycopg2.extensions.connection: + return psycopg2.connect(database=databaseOption, + host=hostOption, + user=userOption, + password=passwordOption, + port=portOption) -def initTable(dbConnection, tableName, tableFormat): +def initTable(dbConnection: psycopg2.extensions.connection, tableName: str, tableFormat: str): dbCursor = dbConnection.cursor() dbCursor.execute(""" @@ -29,12 +34,31 @@ def initTable(dbConnection, tableName, tableFormat): END IF; END; $$""") - printNotice(dbConnection, -1) + debugPrintNotice(dbConnection, -1) dbConnection.commit() dbCursor.close() -def commitQuery(dbConnection, query): +def insertRow(dbConnection: psycopg2.extensions.connection, tableName: str, tableFormat: list[str], tableValues: list): + debugPrint("Attempting to insert new row...") + sanitisedQuery = sql.SQL(""" + INSERT INTO {table} ({format}) + VALUES ({values}) + RETURNING *; + """).format( + table=sql.Identifier(tableName), + format=sql.SQL(", ").join( + sql.Identifier(value) for value in tableFormat + ), + values=sql.SQL(", ").join( + sql.Literal(value) for value in tableValues + ) + ) + debugPrint(sanitisedQuery.as_string(dbConnection)) + commitQuery(dbConnection, sanitisedQuery) + +def commitQuery(dbConnection: psycopg2.extensions.connection, query: sql.Composable) -> list: + debugPrint("Commit query executing...") dbCursor = dbConnection.cursor() dbCursor.execute(query) dbConnection.commit() @@ -42,7 +66,8 @@ def commitQuery(dbConnection, query): dbCursor.close() return dbResults -def execQuery(dbConnection, query): +def execQuery(dbConnection: psycopg2.extensions.connection, query: sql.Composable) -> list: + debugPrint("Exec query executing...") dbCursor = dbConnection.cursor() dbCursor.execute(query) dbResults = dbCursor.fetchall() diff --git a/main.py b/main.py index fce2add..85fe09f 100644 --- a/main.py +++ b/main.py @@ -7,12 +7,11 @@ from pydantic import BaseModel import dbHandler -dbConnection = dbHandler.connect(database="blorgdb", - host="172.20.0.10", - user="dev", - password="dev", - port="5432") - +dbConnection = dbHandler.connect("blorgdb", + "172.20.0.10", + "dev", + "dev", + "5432") dbHandler.initTable(dbConnection, "Users", """ ID SERIAL PRIMARY KEY, Username VARCHAR(255), @@ -51,7 +50,12 @@ ID SERIAL PRIMARY KEY, Name VARCHAR(255) """) -print(dbHandler.execQuery(dbConnection, "SELECT * FROM categories")) +dbHandler.insertRow(dbConnection, + 'users', + ['username', 'firstname', 'lastname', 'description', 'country', 'theme', 'accentcolor', 'passwordhash'], + ['cspark', 'Curt', 'Spark', 'A short description', 'United Kingdom', 'light', 'purple', 'hash256']) + +print(dbHandler.execQuery(dbConnection, "SELECT * FROM users")) dbConnection.close()