diff --git a/db/ckdb.py b/db/ckdb.py index 6f3cb76..397f111 100644 --- a/db/ckdb.py +++ b/db/ckdb.py @@ -5,13 +5,13 @@ import pandas as pd class CKDrive: - ClientPool = dict() + ClientPool = set() @classmethod async def _execute(cls, *args, **kwargs): - key, client = cls.ClientPool.popitem() + client = cls.ClientPool.pop() res = await client.execute(*args, **kwargs) - CKDrive.ClientPool[key] = client + CKDrive.ClientPool.add(client) return res async def execute(self, sql) -> dict: diff --git a/db/ckdb_utils.py b/db/ckdb_utils.py index 0f04dcc..51ca270 100644 --- a/db/ckdb_utils.py +++ b/db/ckdb_utils.py @@ -7,8 +7,9 @@ from .ckdb import CKDrive async def connect_to_ck(pool_size=100): for i in range(pool_size): client = Client(**settings.CK_CONFIG) - CKDrive.ClientPool[i] = client + CKDrive.ClientPool.add(client) async def close_ck_connection(): - await CKDrive.client.disconnect() + for c in CKDrive.ClientPool: + await c.disconnect()