doris-sync/db/pool.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库连接池
支持 MySQL、Doris、PostgreSQL
"""

import pymysql
import pymysql.cursors
from typing import Any, Optional, Dict
from dbutils.pooled_db import PooledDB

from config import get_config


class DBPool:
    """数据库连接池管理器"""
    
    _pools: Dict[str, PooledDB] = {}
    
    @classmethod
    def get_pool(cls, db_type: str) -> PooledDB:
        """
        获取指定类型的连接池
        
        Args:
            db_type: 数据库类型 (mysql/doris/pgsql)
            
        Returns:
            连接池实例
            
        Raises:
            ValueError: 不支持的数据库类型
        """
        if db_type in cls._pools:
            return cls._pools[db_type]
        
        cfg: Dict[str, Any] = get_config(f"database.{db_type}", {})
        if not cfg:
            raise ValueError(f"未配置数据库: {db_type}")
        
        pool_size: int = cfg.get("pool_size", 10)
        
        if db_type == "mysql":
            pool = cls._create_mysql_pool(cfg, pool_size)
        elif db_type == "doris":
            pool = cls._create_doris_pool(cfg, pool_size)
        elif db_type == "pgsql":
            pool = cls._create_pgsql_pool(cfg, pool_size)
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")
        
        cls._pools[db_type] = pool
        return pool
    
    @classmethod
    def _create_mysql_pool(cls, cfg: Dict[str, Any], pool_size: int) -> PooledDB:
        """创建 MySQL 连接池"""
        return PooledDB(
            creator=pymysql,
            maxconnections=pool_size,
            mincached=2,
            maxcached=pool_size,
            blocking=True,
            ping=1,
            host=str(cfg.get("host", "")),
            port=int(cfg.get("port", 3306)),
            user=str(cfg.get("user", "")),
            password=str(cfg.get("password", "")),
            database=str(cfg.get("database", "")),
            charset=str(cfg.get("charset", "utf8mb4")),
            cursorclass=pymysql.cursors.DictCursor,
        )
    
    @classmethod
    def _create_doris_pool(cls, cfg: Dict[str, Any], pool_size: int) -> PooledDB:
        """创建 Doris 连接池 (使用 MySQL 协议)"""
        return PooledDB(
            creator=pymysql,
            maxconnections=pool_size,
            mincached=2,
            maxcached=pool_size,
            blocking=True,
            ping=1,
            host=str(cfg.get("host", "")),
            port=int(cfg.get("port", 9030)),
            user=str(cfg.get("user", "")),
            password=str(cfg.get("password", "")),
            database=str(cfg.get("database", "")),
            charset=str(cfg.get("charset", "utf8mb4")),
            cursorclass=pymysql.cursors.DictCursor,
        )
    
    @classmethod
    def _create_pgsql_pool(cls, cfg: Dict[str, Any], pool_size: int) -> PooledDB:
        """创建 PostgreSQL 连接池"""
        try:
            import psycopg2
            import psycopg2.pool
        except ImportError:
            raise ImportError("需要安装 psycopg2: pip install psycopg2-binary")
        
        return PooledDB(
            creator=psycopg2,
            maxconnections=pool_size,
            mincached=2,
            maxcached=pool_size,
            blocking=True,
            host=str(cfg.get("host", "")),
            port=int(cfg.get("port", 5432)),
            user=str(cfg.get("user", "")),
            password=str(cfg.get("password", "")),
            database=str(cfg.get("database", "")),
        )
    
    @classmethod
    def close_all(cls) -> None:
        """关闭所有连接池"""
        for pool in cls._pools.values():
            try:
                pool.close()
            except Exception:
                pass
        cls._pools.clear()


def get_connection(db_type: str):
    """
    获取数据库连接
    
    Args:
        db_type: 数据库类型 (mysql/doris/pgsql)
        
    Returns:
        数据库连接对象
    """
    return DBPool.get_pool(db_type).connection()


def get_mysql_conn():
    """获取 MySQL 连接（兼容旧API）"""
    return get_connection("mysql")


def get_doris_conn():
    """获取 Doris 连接（兼容旧API）"""
    return get_connection("doris")


def get_pgsql_conn():
    """获取 PostgreSQL 连接"""
    return get_connection("pgsql")