diff --git a/database.py b/database.py index c5b0448..2e5ba6f 100644 --- a/database.py +++ b/database.py @@ -1,6 +1,8 @@ """This module contains functionality for interacting with a PostgreSQL database. It will automatically handle error conditions (i.e. missing columns) without terminating the entire program. Use the :py:class:`DBConnector` class to handle database interactions, either as a standalone object or in a context manager.""" +from __future__ import annotations + import os import psycopg2 from psycopg2 import DatabaseError, OperationalError @@ -22,7 +24,9 @@ class DBConnector: with DBConnector() as db: db.read() """ + def _db_start(self): + """Setup the database connection and cursor.""" try: self.conn = psycopg2.connect( f"host={DB_ADDRESS} port={DB_PORT} dbname={DB_NAME} user={DB_USER} password={DB_PASSWORD}") @@ -31,6 +35,7 @@ class DBConnector: raise e def _db_stop(self): + """Close the cursor and connection.""" self.cur.close() self.conn.close() @@ -47,13 +52,19 @@ class DBConnector: self._db_stop() def _query(self, sql) -> list[dict]: + """Basic function for running queries. + + :param sql: SQL query as plaintext. + :return: Results of the query, or an empty list if none.""" + result = [] try: self.cur.execute(sql) - result = self.cur.fetchall() + result = self._read_dict() except DatabaseError as e: - print(f"DB ERROR [{e.pgcode}]: {e.pgerror}") - result = [] - return result + print(f"ERROR {e.pgcode}: {e.pgerror}\n" + f"Caused by query: {sql}") + finally: + return result def read(self, **kwargs) -> list[dict]: """Read rows from a database that match the specified filters. @@ -90,4 +101,4 @@ class DBConnector: res0 = self.write(**i) if res0: successes.append(res0) - return successes \ No newline at end of file + return successes