diff options
Diffstat (limited to 'src/home/database/clickhouse.py')
-rw-r--r-- | src/home/database/clickhouse.py | 29 |
1 files changed, 29 insertions, 0 deletions
diff --git a/src/home/database/clickhouse.py b/src/home/database/clickhouse.py index 4a2a247..ca81628 100644 --- a/src/home/database/clickhouse.py +++ b/src/home/database/clickhouse.py @@ -1,4 +1,9 @@ +import logging + +from zoneinfo import ZoneInfo +from datetime import datetime, timedelta from clickhouse_driver import Client as ClickhouseClient +from ..config import is_development_mode _links = {} @@ -8,3 +13,27 @@ def get_clickhouse(db: str) -> ClickhouseClient: _links[db] = ClickhouseClient.from_url(f'clickhouse://localhost/{db}') return _links[db] + + +class ClickhouseDatabase: + def __init__(self, db: str): + self.db = get_clickhouse(db) + + self.server_timezone = self.db.execute('SELECT timezone()')[0][0] + self.logger = logging.getLogger(self.__class__.__name__) + + def query(self, *args, **kwargs): + settings = {'use_client_time_zone': True} + kwargs['settings'] = settings + + if 'no_tz_fix' not in kwargs and len(args) > 1 and isinstance(args[1], dict): + for k, v in args[1].items(): + if isinstance(v, datetime): + args[1][k] = v.astimezone(tz=ZoneInfo(self.server_timezone)) + + result = self.db.execute(*args, **kwargs) + + if is_development_mode(): + self.logger.debug(args[0] if len(args) == 1 else args[0] % args[1]) + + return result |