"""公共数据库连接池

依赖：psycopg[binary]>=3, psycopg_pool>=3
说明：
    1. 用 psycopg v3（异步/同步双栈 + 更好的连接池）
    2. DSN 通过环境变量 PG_DSN 覆盖，默认连 docker-compose 里的 pg
    3. 提供 with_tenant() 上下文管理器：在连接上设置 app.tenant_id 供 RLS 使用
"""
from __future__ import annotations

import os
from contextlib import contextmanager
from typing import Iterator, Optional

import psycopg
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

DSN = os.environ.get(
    "PG_DSN",
    "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres password=postgres",
)

# 全局连接池：最多 10 个连接，idle 超过 30s 回收
_pool: Optional[ConnectionPool] = None


def get_pool() -> ConnectionPool:
    """懒加载全局连接池。"""
    global _pool
    if _pool is None:
        _pool = ConnectionPool(
            conninfo=DSN,
            min_size=1,
            max_size=10,
            kwargs={"row_factory": dict_row, "options": "-c search_path=blog,public"},
            open=True,
        )
    return _pool


@contextmanager
def get_conn(tenant_id: Optional[int] = None) -> Iterator[psycopg.Connection]:
    """借一个连接。若传 tenant_id 则自动注入 app.tenant_id，供 RLS 策略读取。"""
    pool = get_pool()
    with pool.connection() as conn:
        if tenant_id is not None:
            with conn.cursor() as cur:
                cur.execute("SELECT set_config('app.tenant_id', %s, true)",
                            (str(tenant_id),))
        yield conn


def fetch_all(sql: str, params: tuple | list | dict | None = None,
              tenant_id: Optional[int] = None) -> list[dict]:
    with get_conn(tenant_id) as conn, conn.cursor() as cur:
        cur.execute(sql, params or ())
        return cur.fetchall()


def fetch_one(sql: str, params: tuple | list | dict | None = None,
              tenant_id: Optional[int] = None) -> dict | None:
    with get_conn(tenant_id) as conn, conn.cursor() as cur:
        cur.execute(sql, params or ())
        return cur.fetchone()


def execute(sql: str, params: tuple | list | dict | None = None,
            tenant_id: Optional[int] = None) -> int:
    """返回受影响行数。"""
    with get_conn(tenant_id) as conn, conn.cursor() as cur:
        cur.execute(sql, params or ())
        return cur.rowcount


def close_pool() -> None:
    global _pool
    if _pool is not None:
        _pool.close()
        _pool = None


if __name__ == "__main__":
    # 自检：能连上就打印版本
    row = fetch_one("SELECT version() AS v, current_database() AS db")
    print("连接成功:", row)
    close_pool()