diff --git a/database.py b/database.py index e69de29..267c87a 100644 --- a/database.py +++ b/database.py @@ -0,0 +1,27 @@ +import os +import psycopg2 + +DB_ADDRESS = os.environ['DB_ADDRESS'] +DB_USER = os.environ['DB_USER'] +DB_PASSWORD = os.environ['DB_PASSWORD'] +DB_NAME = os.environ['DB_NAME'] + + +class DBConnector: + """Context managed database class. Use with statements to automatically open and close the database connection, like + so: + + .. code-block:: python + with DBConnector() as db: + db.read() + """ + def __enter__(self): + self.conn = psycopg2.connect() + self.cur = self.conn.cursor() + + def __exit__(self): + self.cur.close() + self.conn.close() + + def _query(self, sql): + self.cur.execute(sql) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 59c1021..10dce06 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ opencv-python pypdf2==2.12.1 alive-progress requests +psycopg2 # Development matplotlib \ No newline at end of file