# -*- coding: utf-8 -*-
"""数据库连接管理"""
import psycopg2
import psycopg2.extras

class DatabaseConnection:
    """数据库连接管理器"""
    
    def __init__(self, dsn: str, session: dict = None, connect_timeout: int = None):
        self.conn = psycopg2.connect(dsn, connect_timeout=(connect_timeout or 5))
        self.conn.autocommit = False
        
        # 设置会话参数
        if session:
            with self.conn.cursor() as c:
                if session.get("timezone"):
                    c.execute("SET TIME ZONE %s", (session["timezone"],))
                if session.get("statement_timeout_ms") is not None:
                    c.execute("SET statement_timeout = %s", (int(session["statement_timeout_ms"]),))
                if session.get("lock_timeout_ms") is not None:
                    c.execute("SET lock_timeout = %s", (int(session["lock_timeout_ms"]),))
                if session.get("idle_in_tx_timeout_ms") is not None:
                    c.execute("SET idle_in_transaction_session_timeout = %s", 
                             (int(session["idle_in_tx_timeout_ms"]),))
    
    def query(self, sql: str, args=None):
        """执行查询并返回结果"""
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as c:
            c.execute(sql, args)
            return c.fetchall()
    
    def execute(self, sql: str, args=None):
        """执行SQL语句"""
        with self.conn.cursor() as c:
            c.execute(sql, args)
    
    def commit(self):
        """提交事务"""
        self.conn.commit()
    
    def rollback(self):
        """回滚事务"""
        self.conn.rollback()
    
    def close(self):
        """关闭连接"""
        try:
            self.conn.close()
        except Exception:
            pass