43 lines
1.1 KiB
Python
43 lines
1.1 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():
|
|
"""Database connection context manager"""
|
|
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()
|
|
|
|
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()
|