From 2638e5126b144c2bee0a83fcb4175d702cb3ec0a Mon Sep 17 00:00:00 2001
From: cspark <git@cspark.dev>
Date: Thu, 18 Apr 2024 18:08:49 +0100
Subject: [PATCH] Move database handling functionality to dbHandler, helper
 functions

---
 dbHandler.py | 50 ++++++++++++++++++++++++++++++++
 main.py      | 80 +++++++++++++++++++++++++++++-----------------------
 2 files changed, 95 insertions(+), 35 deletions(-)
 create mode 100644 dbHandler.py

diff --git a/dbHandler.py b/dbHandler.py
new file mode 100644
index 0000000..d59013e
--- /dev/null
+++ b/dbHandler.py
@@ -0,0 +1,50 @@
+import psycopg2
+
+def printNotice(dbConnection, i):
+    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 initTable(dbConnection, tableName, tableFormat):
+    dbCursor = dbConnection.cursor()
+
+    dbCursor.execute("""
+    DO $$
+    BEGIN
+         IF (EXISTS (SELECT *
+                            FROM INFORMATION_SCHEMA.TABLES
+                            WHERE TABLE_NAME = '""" + tableName.lower() + """'))
+         THEN
+              RAISE NOTICE '""" + tableName + """ Table does exist! Skipping creating table.';
+         ELSE
+              RAISE NOTICE '""" + tableName + """ Table does not exist! Creating table.';
+              CREATE TABLE """ + tableName.lower() + """ (
+              """ + tableFormat + """
+              );
+         END IF;
+    END;
+    $$""")
+    printNotice(dbConnection, -1)
+
+    dbConnection.commit()
+    dbCursor.close()
+
+def commitQuery(dbConnection, query):
+    dbCursor = dbConnection.cursor()
+    dbCursor.execute(query)
+    dbConnection.commit()
+    dbResults = dbCursor.fetchall()
+    dbCursor.close()
+    return dbResults
+
+def execQuery(dbConnection, query):
+    dbCursor = dbConnection.cursor()
+    dbCursor.execute(query)
+    dbResults = dbCursor.fetchall()
+    dbCursor.close()
+    return dbResults
diff --git a/main.py b/main.py
index 88694fd..fce2add 100644
--- a/main.py
+++ b/main.py
@@ -5,45 +5,55 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from pydantic import BaseModel
 
-import psycopg2
+import dbHandler
 
-dbconn = psycopg2.connect(database="blorgdb",
-                          host="172.20.0.10",
-                          user="dev",
-                          password="dev",
-                          port="5432")
+dbConnection = dbHandler.connect(database="blorgdb",
+                                 host="172.20.0.10",
+                                 user="dev",
+                                 password="dev",
+                                 port="5432")
 
-dbcursor = dbconn.cursor()
+dbHandler.initTable(dbConnection, "Users", """
+ID SERIAL PRIMARY KEY,
+Username VARCHAR(255),
+FirstName VARCHAR(255),
+LastName VARCHAR(255),
+Description VARCHAR(255),
+Country VARCHAR(255),
+Theme VARCHAR(255),
+AccentColor VARCHAR(255),
+PasswordHash VARCHAR(255)
+""")
+dbHandler.initTable(dbConnection, "SignOns", """
+ID SERIAL PRIMARY KEY,
+UserID VARCHAR(255),
+LoginSuccess BOOLEAN,
+DateAttempted VARCHAR(255),
+IPLocationAttempted VARCHAR(255)
+""")
+dbHandler.initTable(dbConnection, "AuthTokens", """
+ID SERIAL PRIMARY KEY,
+Token VARCHAR(255),
+OwnerID INTEGER,
+DateCreated TIMESTAMP,
+DateExpiry TIMESTAMP,
+IPLocationCreated VARCHAR(255)
+""")
+dbHandler.initTable(dbConnection, "Blogs", """
+ID SERIAL PRIMARY KEY,
+AuthorID INTEGER,
+CategoryID INTEGER,
+DatePosted TIMESTAMP,
+Description VARCHAR(255)
+""")
+dbHandler.initTable(dbConnection, "Categories", """
+ID SERIAL PRIMARY KEY,
+Name VARCHAR(255)
+""")
 
-dbcursor.execute("""
-DO $$
-BEGIN
-  IF (EXISTS (SELECT *
-                   FROM INFORMATION_SCHEMA.TABLES
-                   WHERE TABLE_NAME = 'users'))
-  THEN
-    RAISE NOTICE 'Table does exist!';
-  ELSE
-    RAISE NOTICE 'User Table does not exist! Creating table.';
-    CREATE TABLE users (
-      UserID int,
-      FirstName varchar(255),
-      LastName varchar(255),
-      PasswordHash varchar(255)
-    );
-  END IF;
-END;
-$$""")
+print(dbHandler.execQuery(dbConnection, "SELECT * FROM categories"))
 
-
-#print(dbcursor.fetchall())
-
-for notice in dbconn.notices:
-    print("(SQL SERVER) " + notice)
-
-dbconn.commit()
-dbcursor.close()
-dbconn.close()
+dbConnection.close()
 
 app = FastAPI()