"""
数据库连接

使用 psycopg2 直连 PostgreSQL，不引入 ORM。
连接参数从环境变量读取（经 config 模块加载）。

提供两类连接：
- get_connection()：zqyy_app 读写连接（用户/队列/调度等业务数据）
- get_etl_readonly_connection(site_id)：etl_feiqiu 只读连接（数据库查看器），
  自动设置 RLS site_id 隔离
"""

import psycopg2
from psycopg2.extensions import connection as PgConnection

from app.config import (
    APP_DB_NAME,
    DB_HOST,
    DB_PASSWORD,
    DB_PORT,
    DB_USER,
    ETL_DB_HOST,
    ETL_DB_NAME,
    ETL_DB_PASSWORD,
    ETL_DB_PORT,
    ETL_DB_USER,
)


def get_connection() -> PgConnection:
    """
    获取 zqyy_app 数据库连接。

    调用方负责关闭连接（推荐配合 contextmanager 或 try/finally 使用）。
    """
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        dbname=APP_DB_NAME,
    )


def get_etl_readonly_connection(site_id: int | str) -> PgConnection:
    """
    获取 ETL 数据库（etl_feiqiu）的只读连接。

    连接建立后自动执行：
    1. SET default_transaction_read_only = on  — 禁止写操作
    2. SET LOCAL app.current_site_id = '{site_id}'  — 启用 RLS 门店隔离

    调用方负责关闭连接。典型用法::

        conn = get_etl_readonly_connection(site_id)
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT ...")
        finally:
            conn.close()
    """
    conn = psycopg2.connect(
        host=ETL_DB_HOST,
        port=ETL_DB_PORT,
        user=ETL_DB_USER,
        password=ETL_DB_PASSWORD,
        dbname=ETL_DB_NAME,
    )
    try:
        conn.autocommit = False
        with conn.cursor() as cur:
            # 会话级只读：防止任何写操作
            cur.execute("SET default_transaction_read_only = on")
            # 事务级 RLS 隔离：设置当前门店 ID
            cur.execute(
                "SET LOCAL app.current_site_id = %s", (str(site_id),)
            )
        conn.commit()
    except Exception:
        conn.close()
        raise
    return conn