xbackend/db/ckdb.py
2021-07-27 16:08:49 +08:00

62 lines
2.1 KiB
Python

import datetime
from aioch import Client
import pandas as pd
class CKDrive:
ClientPool = set()
@classmethod
async def _execute(cls, *args, **kwargs):
client = cls.ClientPool.pop()
res = await client.execute(*args, **kwargs)
CKDrive.ClientPool.add(client)
return res
async def execute(self, sql) -> dict:
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._execute(sql, with_column_types=True, columnar=True)
df = pd.DataFrame({col[0]: d for d, col in zip(data, columns)})
return df
async def count(self, db: str, tb: str):
sql = f'select count() as `count` from {db}.{tb}'
res = await self.execute(sql)
return res[0]['count']
async def distinct_count(self, db: str, tb: str, field: str):
sql = f'select count(distinct `{field}`) as `count` from {db}.{tb}'
res = await self.execute(sql)
return res[0]['count']
async def field_count(self, db: str, tb: str):
sql = f"select count(name) as `count` from system.columns where database='{db}' and table='{tb}'"
res = await self.execute(sql)
return res[0]['count']
async def distinct(self, db: str, tb: str, field: str):
sql = f'select distinct `{field}` as v from {db}.{tb}'
res = await self.query_dataframe(sql)
return res['v'].to_list()
async def yesterday_event_count(self, db: str):
today = datetime.date.today()
yesterday = today - datetime.timedelta(days=1)
today_str = today.strftime('%Y-%m-%d %H:%M:%S')
yesterday_str = yesterday.strftime('%Y-%m-%d %H:%M:%S')
sql = f"select `#event_name` as event_name, count() as v from {db}.event where `#event_time`>='{yesterday_str}' and `#event_time`<'{today_str}' group by `#event_name`"
df = await self.query_dataframe(sql)
return df.set_index('event_name').T.to_dict()
ckdb = CKDrive()
def get_ck_db() -> CKDrive:
return ckdb