CoolBotDataSys/database.py

"""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()