QingLong/ClickHouse/mappers/data_mapper.py

from datetime import datetime

class DataMapper:
    def __init__(self, ch_conn, table_name):
        self.ch_conn = ch_conn
        self.table_name = table_name
        self.date_columns = []
        self.uint64_columns = []
        self.min_date = datetime(1970, 1, 1)
        self.max_date = datetime(2105, 12, 31, 23, 59, 59)

    def _load_columns(self):
        """从ClickHouse表中加载列信息并设置日期和无符号整数列"""
        result = self.ch_conn.connect().execute(f"DESCRIBE TABLE {self.table_name}")
        for row in result:
            column_name = row[0]
            column_type = row[1]
            if 'DateTime' in column_type or 'Date' in column_type:
                self.date_columns.append(column_name)
            elif 'UInt' in column_type:
                self.uint64_columns.append(column_name)

    def map_row(self, columns, row):
        row_dict = dict(zip(columns, row))
        return {col: self._map_value(col, val) for col, val in row_dict.items()}

    def _map_value(self, col, value):
        if col in self.uint64_columns:
            return self._handle_uint64(value)
        elif col in self.date_columns:
            return self._handle_datetime(value)
        elif isinstance(value, str):
            return value.strip()
        return value

    def _handle_uint64(self, value):
        try:
            return int(float(value)) if value not in (None, '', 'NULL') else 0
        except:
            return 0

    def _handle_datetime(self, value):
        dt = self._parse_datetime(value)
        return dt if dt else self.min_date

    def _parse_datetime(self, value):
        if value in (None, 0, '0', '0.0', '0.00', '', 'null', 'NULL'):
            return self.min_date

        try:
            str_value = str(value).strip()
            for fmt in ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d', '%Y%m%d%H%M%S', '%Y/%m/%d %H:%M:%S'):
                try:
                    parsed = datetime.strptime(str_value, fmt)
                    return self._clamp_datetime(parsed)
                except ValueError:
                    continue

            if str_value.isdigit():
                ts = int(str_value)
                if 1e12 < ts < 1e13:  # 毫秒级时间戳
                    parsed = datetime.fromtimestamp(ts / 1000)
                elif 1e9 < ts < 1e10:  # 秒级时间戳
                    parsed = datetime.fromtimestamp(ts)
                return self._clamp_datetime(parsed)

            return self.min_date
        except:
            return self.min_date

    def _clamp_datetime(self, dt):
        if dt < self.min_date:
            return self.min_date
        elif dt > self.max_date:
            return self.max_date
        return dt