59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
from clickhouse_connect import get_async_client
|
||
from clickhouse_connect.driver import AsyncClient as AsyncClickHouseClient
|
||
from app.settings.config import settings
|
||
import urllib3
|
||
from typing import Any, Dict
|
||
|
||
|
||
async def get_clickhouse_client() -> AsyncClickHouseClient:
|
||
"""获取ClickHouse异步客户端"""
|
||
# 创建自定义连接池管理器，适配多worker模式
|
||
# 考虑到多worker环境，每个worker的连接池应该适中
|
||
# maxsize参数控制每个主机的连接池大小，block=True防止连接池溢出
|
||
pool_mgr = urllib3.PoolManager(
|
||
num_pools=2, # 连接池数量，减少以适配多worker
|
||
maxsize=5, # 每个连接池的最大连接数，每个worker最多10个连接
|
||
block=True # 当连接池满时阻塞而不是创建新连接
|
||
)
|
||
|
||
return await get_async_client(
|
||
host=settings.CLICKHOUSE_HOST,
|
||
username=settings.CLICKHOUSE_USER,
|
||
password=settings.CLICKHOUSE_PASS,
|
||
database=settings.CLICKHOUSE_DB,
|
||
port=settings.CLICKHOUSE_PORT,
|
||
pool_mgr=pool_mgr,
|
||
connect_timeout=30,
|
||
send_receive_timeout=120
|
||
)
|
||
|
||
|
||
class ClickHouseManager:
|
||
"""ClickHouse连接管理器"""
|
||
|
||
def __init__(self):
|
||
self._client: AsyncClickHouseClient = None
|
||
|
||
async def get_client(self) -> AsyncClickHouseClient:
|
||
"""获取或创建ClickHouse客户端"""
|
||
if self._client is None:
|
||
self._client = await get_clickhouse_client()
|
||
return self._client
|
||
|
||
async def execute(self, query: str, parameters: Dict[str, Any] = None):
|
||
"""执行SQL查询"""
|
||
client = await self.get_client()
|
||
return await client.query(query, parameters=parameters)
|
||
|
||
async def close(self):
|
||
"""关闭ClickHouse连接"""
|
||
if self._client:
|
||
close_result = self._client.close()
|
||
if close_result is not None:
|
||
await close_result
|
||
self._client = None
|
||
|
||
|
||
# 全局ClickHouse管理器实例
|
||
clickhouse_manager = ClickHouseManager()
|