diff --git a/db/ckdb.py b/db/ckdb.py index d627b24..6f3cb76 100644 --- a/db/ckdb.py +++ b/db/ckdb.py @@ -5,15 +5,22 @@ import pandas as pd class CKDrive: - client: Client = None + ClientPool = dict() + + @classmethod + async def _execute(cls, *args, **kwargs): + key, client = cls.ClientPool.popitem() + res = await client.execute(*args, **kwargs) + CKDrive.ClientPool[key] = client + return res async def execute(self, sql) -> dict: - data, columns = await self.client.execute(sql, with_column_types=True, columnar=True) + data, columns = await self._execute(sql, with_column_types=True, columnar=True) df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)}) return df.T.to_dict() async def query_dataframe(self, sql): - data, columns = await self.client.execute(sql, with_column_types=True, columnar=True) + data, columns = await self._execute(sql, with_column_types=True, columnar=True) df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)}) return df diff --git a/db/ckdb_utils.py b/db/ckdb_utils.py index 9d1518f..0f04dcc 100644 --- a/db/ckdb_utils.py +++ b/db/ckdb_utils.py @@ -4,8 +4,10 @@ from core.config import settings from .ckdb import CKDrive -async def connect_to_ck(): - CKDrive.client = Client(**settings.CK_CONFIG) +async def connect_to_ck(pool_size=100): + for i in range(pool_size): + client = Client(**settings.CK_CONFIG) + CKDrive.ClientPool[i] = client async def close_ck_connection():