zlwl-algos-new/COMMON/service/datawarehouse/mysql_to_doris.py

from sshtunnel import SSHTunnelForwarder
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
import pandas as pd
from urllib import parse
from pydoris.doris_client import DorisClient
import json
import datetime
import numpy as np

# 数仓相关库函数
class DwService():
    
    def __init__(self, tables, doris_host, doris_port, doris_user, doris_password, doris_db, 
                    doris_table_format, mysql_host, mysql_port, mysql_user, mysql_password, mysql_db, start_time="", end_time="", ssh_host="", ssh_name="", ssh_password="",
                    settings={}):
        self.tables = tables
        self.doris_host = doris_host
        self.doris_port = doris_port
        self.doris_user = doris_user
        self.doris_password = doris_password
        self.doris_db = doris_db
        self.doris_table_format = doris_table_format
        self.mysql_host = mysql_host
        self.mysql_port = int(mysql_port)
        self.mysql_user = mysql_user
        self.mysql_password = mysql_password
        self.mysql_db = mysql_db
        self.start_time = start_time
        self.end_time = end_time
        self.ssh_host = ssh_host
        self.ssh_name = ssh_name
        self.ssh_password = ssh_password
        settings_default = {"skip_endtime_0000":"true","skip_endtime_1970":"true"}
        for k,v in settings.items():
            settings_default.update({k:v})
        self.settings = settings_default
        
    def _get_db(self):
        self.doris_client = DorisClient(self.doris_host, self.doris_port, self.doris_user, self.doris_password)
        self.doris_client.options\
            .set_json_format()\
            .set_auto_uuid_label()\
            .set_option('strip_outer_array', 'true')
        self.db_engine = create_engine(
            "mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8".format(self.mysql_user,parse.quote_plus(self.mysql_password),
                                                                 self.mysql_host, int(self.mysql_port), self.mysql_db),pool_recycle=3600,pool_size=5)

        self.Session = sessionmaker(self.db_engine)
    
    def _get_and_write(self, session, sql, table, ):
        # 从mysql取数，处理然后写入doris
        print(f'{datetime.datetime.now()}-{self.doris_table_format.format(table)} 数据开始同步...............')
        self.doris_client.options.set_auto_uuid_label()
        df_desc = pd.read_sql(f'desc {table}', self.db_engine)
        res = session.execute(sql).fetchall()
        columns = df_desc['Field'].tolist()
        types = df_desc['Type'].tolist()
        res_dict = []
        # 字典列表组装
        for r in res:
            mp = {}
            i = 0
            skip_flag = False
            for item in r:
                # 数据类型特殊处理
                if (str(item) == '0000-00-00 00:00:00'): # done表中存在endtime=0000的数据，影响索引字段
                    if self.settings.get("skip_endtime_0000") == "true": # 跳过endtime0000的数据
                        skip_flag = True
                if ('1970-01-01' in str(item) ): # done表中存在endtime=0000的数据，影响索引字段
                    if self.settings.get("skip_endtime_1970") == "true": # 跳过endtime0000的数据
                        skip_flag = True
                if types[i] == 'datetime' or types[i] == 'date' or types[i] == 'timestamp':
                    mp.update({columns[i]: str(item)})
                elif types[i] == 'tinyint':
                    mp.update({columns[i]:  0 if item == 'false' else 1 if item == 'true' else item})
                elif 'decimal' in types[i] :
                    mp.update({columns[i]: str(item)})
                else:
                    mp.update({columns[i]:item})
                i = i + 1
            if not skip_flag:
                res_dict.append(mp)
        if len(res_dict) > 0 :
            print(f'{datetime.datetime.now()}-{self.doris_table_format.format(table)} 数据开始导入doris...............')
            self.doris_client._session.should_strip_auth = lambda old_url, new_url: False
            resp = self.doris_client._session.request(
                'PUT', url=self.doris_client._build_url(self.doris_db, self.doris_table_format.format(table)),
                data=json.dumps(res_dict).encode('utf-8'), 
                headers=self.doris_client.options.get_options(),
                auth=self.doris_client._auth
            )
            load_status = json.loads(resp.text)['Status'] == 'Success'
            if resp.status_code == 200 and resp.reason == 'OK' and load_status:
                print(f'{datetime.datetime.now()}-{self.doris_table_format.format(table)} 数据导入成功')
                print(resp.text)
            else:
                print(f'{datetime.datetime.now()}-{self.doris_table_format.format(table)} 数据导入失败')
                raise Exception(resp.text)
        else:
            print(f"{datetime.datetime.now()}-本次需同步的数据条数0")

    def full_table_sync(self):
        # 全量更新
        sql_format = 'select * from {}'
        self._get_db()
        with self.Session() as session:
            for table in self.tables:
                sql = sql_format.format(table)
                self._get_and_write(session, sql, table)
                
    def inc_table_sync(self):
        # 增量更新
        if self.start_time == "" or  self.end_time == "":
            raise Exception("输入参数错误")
        sql_format = "select * from {} where {} >= '{}' and {} < '{}'"
        self._get_db()
        with self.Session() as session:
            for table_dict in self.tables:
                for table, value in table_dict.items():
                    time_field = value.get('time_field')
                    sql = sql_format.format(table, time_field, self.start_time, time_field, self.end_time)
                    self._get_and_write(session, sql, table)
                    
    def full_table_sync_with_ssh(self):
        with SSHTunnelForwarder(self.ssh_host, 22, ssh_username=self.ssh_name, ssh_password=self.ssh_password, remote_bind_address=(self.mysql_host, self.mysql_port)) as tunnel:
            self.mysql_host = '127.0.0.1'
            self.mysql_port = tunnel.local_bind_port
            self.full_table_sync()
            
    def inc_table_sync_with_ssh(self):
        if self.ssh_host == ""  or  self.ssh_name == "" or self.ssh_password == "":
            raise Exception("输入参数错误")
        with SSHTunnelForwarder(self.ssh_host, 22, ssh_username=self.ssh_name, ssh_password=self.ssh_password, remote_bind_address=(self.mysql_host, self.mysql_port)) as tunnel:
            self.mysql_host = '127.0.0.1'
            self.mysql_port = tunnel.local_bind_port
            self.inc_table_sync()