106 lines
2.9 KiB
Python
106 lines
2.9 KiB
Python
"""PostgreSQL database connection module"""
|
|
import psycopg2
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
DB_CONFIG = {
|
|
"host": os.environ.get("DB_HOST", "pgm-bp1t1008h019ez6c.pg.rds.aliyuncs.com"),
|
|
"port": int(os.environ.get("DB_PORT", 5432)),
|
|
"user": os.environ.get("DB_USER", "coolbot"),
|
|
"password": os.environ.get("DB_PASSWORD", "Coolbot123"),
|
|
"database": os.environ.get("DB_NAME", "coolbot_data"),
|
|
}
|
|
|
|
@contextmanager
|
|
def get_db():
|
|
conn = None
|
|
try:
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
yield conn
|
|
conn.commit()
|
|
except Exception as e:
|
|
if conn:
|
|
conn.rollback()
|
|
raise e
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
|
|
class DictCursor:
|
|
def __init__(self, cursor):
|
|
self._cursor = cursor
|
|
|
|
def __iter__(self):
|
|
columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
|
|
for row in self._cursor:
|
|
yield dict(zip(columns, row))
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
self.close()
|
|
|
|
def execute(self, *args, **kwargs):
|
|
return self._cursor.execute(*args, **kwargs)
|
|
|
|
def fetchone(self):
|
|
row = self._cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
|
|
return dict(zip(columns, row))
|
|
|
|
def fetchall(self):
|
|
columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
|
|
return [dict(zip(columns, row)) for row in self._cursor.fetchall()]
|
|
|
|
def fetchmany(self, size=None):
|
|
if size is None:
|
|
size = self._cursor.arraysize
|
|
columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
|
|
rows = self._cursor.fetchmany(size)
|
|
return [dict(zip(columns, row)) for row in rows]
|
|
|
|
def close(self):
|
|
return self._cursor.close()
|
|
|
|
class Database:
|
|
@contextmanager
|
|
def get_cursor(self, dictionary=True):
|
|
conn = None
|
|
try:
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
raw_cursor = conn.cursor()
|
|
try:
|
|
if dictionary:
|
|
cursor = DictCursor(raw_cursor)
|
|
else:
|
|
cursor = raw_cursor
|
|
yield cursor
|
|
conn.commit()
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
finally:
|
|
raw_cursor.close()
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
|
|
db = Database()
|
|
|
|
def test_connection():
|
|
try:
|
|
with get_db() as conn:
|
|
cur = conn.cursor()
|
|
cur.execute("SELECT version()")
|
|
print(f"DB OK: {str(cur.fetchone()[0])[:50]}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"DB ERROR: {e}")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
test_connection()
|