poyee-data-warehouse/dw_base/datax/datax_utils.py

# -*- coding:utf-8 -*-

from typing import Dict, List

from dw_base import NORM_GRN, NORM_YEL
from dw_base.database.mysql_utils import MySQLColumn
from dw_base.utils.log_utils import pretty_print


def convert_mysql_column_types(columns: List[MySQLColumn]) -> Dict[str, str]:
    column_types = {}
    bool_types = ['bool', 'bit']
    double_types = ['float', 'double']
    string_types = ['text', 'longtext', 'mediumtext', 'time']
    timestamp_types = ['datetime', 'timestamp']
    for mysql_column in columns:
        column_name = mysql_column.COLUMN_NAME
        origin_type = mysql_column.COLUMN_TYPE
        datax_type = None
        if origin_type.startswith('bigint'):
            datax_type = 'bigint'
        elif origin_type.__contains__('int'):
            datax_type = 'int'
        elif timestamp_types.__contains__(origin_type):
            datax_type = 'timestamp'
        elif origin_type == 'date':
            datax_type = 'date'
        elif bool_types.__contains__(origin_type):
            datax_type = 'boolean'
        elif double_types.__contains__(origin_type) or origin_type.startswith('decimal'):
            datax_type = 'double'
        elif string_types.__contains__(origin_type) or origin_type.startswith('varchar') or origin_type.startswith(
                'char') or origin_type.startswith('enum'):
            # string不用管
            pass
        else:
            pretty_print(f'{NORM_YEL}遇到了未处理MySQL——DataX类型映射的MySQL类型：{NORM_GRN}{origin_type}')
        if datax_type is not None:
            column_types[column_name] = datax_type
    return column_types