make the DBConnector class usable as a standalone object or in a context manager

This commit is contained in:
Camryn Thomas 2024-01-14 14:16:03 -06:00
parent 1ef515e265
commit 83f1dc6a9f
Signed by: cptlobster
GPG Key ID: 6D341D688163A176

View File

@ -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)