from pyspark.sql import functions as F, DataFrame


class DorisHelper(object):
    """
    导入导出工具https://doris.apache.org/docs/2.0/ecosystem/spark-doris-connector/
    """
    __connection_info__ = {
        "adv": {
            "ip": "192.168.10.156",
            "http_port": 8030,
            "jdbc_port": 19030,
            "user": "selection",
            "pwd": "1w&@Js12a3",
            "def_db": "advertising_manager"
        }
    }

    @staticmethod
    def get_connection_info(use_type='adv'):
        """
        获取链接信息
        :param use_type:  adv 广告系统
        :return:
        """
        return DorisHelper.__connection_info__[use_type]

    @classmethod
    def spark_export(cls, df_save: DataFrame, table_name, use_type='adv'):
        """
        :param df_save:
        :param db_name:
        :param table_name:
        :param use_type:
        :return:
        """
        print(f"执行导出数据到Doris任务,导出数据： {df_save}, Doris表名： {table_name}")
        try:
            connection_info = DorisHelper.get_connection_info(use_type)
            options = {
                "doris.fenodes": f"{connection_info['ip']}:{connection_info['jdbc_port']}",
                "user": {connection_info['user']},
                "password": {connection_info['pwd']},
                "doris.table.identifier": table_name,
            }
            df_save.write.format("doris") \
                .options(**options) \
                .mode("append") \
                .save()
            print("DataFrame数据导出任务完毕")
        except Exception as e:
            raise e
