From adf440f977151e6b70dbb04766db75dbfa44b174 Mon Sep 17 00:00:00 2001
From: wuaho <wuaho@legu.cc>
Date: Tue, 27 Jul 2021 16:08:49 +0800
Subject: [PATCH] =?UTF-8?q?ck=20=E8=BF=9E=E6=8E=A5=E6=B1=A0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 db/ckdb.py       | 6 +++---
 db/ckdb_utils.py | 5 +++--
 2 files changed, 6 insertions(+), 5 deletions(-)

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()