diff --git a/database.py b/database.py index 47a351b..48c90ac 100644 --- a/database.py +++ b/database.py @@ -1,3 +1,6 @@ +"""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.""" import os import psycopg2 @@ -17,14 +20,27 @@ class DBConnector: with DBConnector() as db: db.read() """ - def __enter__(self): - self.conn = psycopg2.connect() + def _db_start(self): + self.conn = psycopg2.connect( + f"host={DB_ADDRESS} port={DB_PORT} dbname={DB_NAME} user={DB_USER} password={DB_PASSWORD}") self.cur = self.conn.cursor() - def __exit__(self): + def _db_stop(self): self.cur.close() self.conn.close() + def __init__(self): + self._db_start() + + def __del__(self): + self._db_stop() + + def __enter__(self): + self._db_start() + + def __exit__(self): + self._db_stop() + def _query(self, sql): try: self.cur.execute(sql)