BIDI_EXTRACTION_AI/BiddingKG/dl/infra/db.py

# -*- coding: utf-8 -*-
"""统一 PostgreSQL 连接池。

合并原 common/Connection.py 与 interface/Connection.py（两者完全一致）。
所有 PG 连接经此模块，配置走 infra/config.py。

向后兼容：
    - common/Connection.getConnection() 继续可用
    - interface/Connection.getConnection() 继续可用
    二者均 re-export 自本模块。
"""
from __future__ import absolute_import

import logging

import psycopg2
from DBUtils.PooledDB import PooledDB

from . import config

__all__ = ["get_connection", "get_db_connection", "get_pool"]

logger = logging.getLogger(__name__)

_pools = {}


def get_pool(dbname="BiddingKG"):
    """获取指定库的连接池，按 dbname 缓存。"""
    if dbname not in _pools:
        cfg = config.pg_db_config(dbname)
        _pools[dbname] = PooledDB(
            psycopg2,
            int(cfg.get("pool_size", 10)),
            host=cfg.get("host", "127.0.0.1"),
            port=str(cfg.get("port", 5432)),
            user=cfg.get("user", "postgres"),
            password=cfg.get("password", ""),
            dbname=cfg.get("dbname", dbname),
        )
    return _pools[dbname]


def get_connection(dbname="BiddingKG"):
    """从连接池取一个连接。默认取 BiddingKG 库。

    向后兼容：等价于原 common/Connection.getConnection()。
    """
    return get_pool(dbname).connection()


def get_db_connection(dbname):
    """显式按库名取连接。用于训练/数据脚本访问非默认库。"""
    return get_connection(dbname)


def close_all():
    """关闭所有连接池，用于测试与进程退出。"""
    for dbname, pool in _pools.items():
        try:
            pool.close()
        except Exception as e:  # pragma: no cover
            logger.warning("close pool %s failed: %s", dbname, e)
    _pools.clear()