Initial commit

db.py

@@ -0,0 +1,108 @@
+from typing import Any, NamedTuple
+from psycopg import Connection
+from psycopg.rows import namedtuple_row, dict_row
+from psycopg.sql import SQL, Identifier, Literal
+from psycopg_pool import ConnectionPool
+from columns import ColumnDefinition
+
+
+class _ExtendedConnection(Connection):
+    def __init__(self, *args, **kwargs):
+        super(_ExtendedConnection, self).__init__(*args, **kwargs)
+        self.row_factory = dict_row
+
+
+class DBConnector:
+    def __init__(self, conninfo: str):
+        self._pool = ConnectionPool(conninfo, connection_class=_ExtendedConnection)
+
+    def connection(self):
+        return self._pool.connection()
+
+    def addTableInfoToHstore(self, table_name: str, columns: list[ColumnDefinition]):
+        # TODO: implement this
+        ...
+
+    def tableExists(self, table_name: str):
+        stmt = SQL(
+            """SELECT EXISTS (
+                SELECT FROM information_schema.tables 
+                WHERE table_schema LIKE 'public' AND table_type LIKE 'BASE TABLE' AND table_name = {}
+            );"""
+        ).format(Literal(table_name))
+        with self.connection() as conn:
+            result = conn.execute(stmt).fetchone()
+            return False if result is None else result["exists"]
+
+    def createTable(self, table_name: str, columns: list[ColumnDefinition]):
+        stmt = SQL("CREATE TABLE IF NOT EXISTS {} ({})").format(
+            Identifier(table_name), SQL(", ").join(map(lambda c: c.sql(), columns))
+        )
+        with self.connection() as conn:
+            conn.execute(stmt)
+
+    def dropTable(self, table_name: str):
+        stmt = SQL("DROP TABLE IF EXISTS {}").format(Identifier(table_name))
+        with self.connection() as conn:
+            conn.execute(stmt)
+
+    def insertIntoTable(self, table_name: str, columns: dict[str, Any]):
+        if len(columns) == 0:
+            stmt = SQL("INSERT INTO {} DEFAULT VALUES").format(Identifier(table_name))
+        else:
+            stmt = SQL("INSERT INTO {} ({}) VALUES ({})").format(
+                Identifier(table_name),
+                SQL(", ").join(map(lambda c: Identifier(c), columns.keys())),
+                SQL(", ").join(map(lambda c: Literal(c), columns.values())),
+            )
+        with self.connection() as conn:
+            conn.execute(stmt)
+
+    def selectFromTable(
+        self, table_name: str, columns: list[str]
+    ) -> list[dict[str, Any]]:
+        if len(columns) == 1 and columns[0] == "*":
+            stmt = SQL("SELECT * FROM {}").format(Identifier(table_name))
+        else:
+            stmt = SQL("SELECT {} FROM {}").format(
+                SQL(", ").join(map(lambda c: Identifier(c), columns)),
+                Identifier(table_name),
+            )
+        with self.connection() as conn:
+            return conn.execute(stmt).fetchall()
+
+    def updateDataInTable(
+        self, table_name: str, columns: dict[str, Any], where: dict[str, Any]
+    ):
+        stmt = SQL("UPDATE {} SET {} WHERE {}").format(
+            Identifier(table_name),
+            SQL(", ").join(
+                map(
+                    lambda c: Identifier(c) + SQL(" = ") + Literal(columns[c]),
+                    columns.keys(),
+                )
+            ),
+            SQL(", ").join(
+                map(
+                    lambda c: Identifier(c) + SQL(" LIKE ") + Literal(where[c]),
+                    where.keys(),
+                )
+            ),
+        )
+        with self.connection() as conn:
+            conn.execute(stmt)
+
+    def selectJoinedTables(
+        self, base_table: str, join_table: str, columns: list[str], join_column: str
+    ):
+        stmt = SQL("SELECT {} FROM {} JOIN {} ON {}.{} = {}.{}").format(
+            SQL(", ").join(map(lambda c: Identifier(c), columns)),
+            Identifier(base_table),
+            Identifier(join_table),
+            Identifier(base_table),
+            Identifier(join_column),
+            Identifier(join_table),
+            Identifier(join_column),
+        )
+        with self.connection() as conn:
+            return conn.execute(stmt).fetchall()