es_update_with_usr_mask.py

# author : wangrui
# data : 2024/7/17 14:28
import os
import sys

sys.path.append(os.path.dirname(sys.path[0]))  # 上级目录
from utils.db_util import DBUtil, DbTypes
from pyspark.sql import SparkSession
import pandas as pd
from yswg_utils.common_df import get_user_mask_type_asin_sql
from utils.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils, HdfsError
from pyspark.sql import Row
from utils.spark_util import SparkUtil
from utils.DorisHelper import DorisHelper

__es_ip__ = "192.168.10.217"
__es_port__ = "9200"
__es_user__ = "elastic"
__es_passwd__ = "selection2021.+"
__warehouse_dir__ = "hdfs://nameservice1:8020/home/big_data_selection"
__metastore_uris__ = "thrift://hadoop16:9083"



def get_es_index_name(site_name):
    engine_mysql = DBUtil.get_db_engine(db_type=DbTypes.mysql.name, site_name=site_name)
    sql = f"""SELECT REPLACE(report_date, '-', '_') as date_info FROM workflow_everyday WHERE site_name='{site_name}' AND date_type='month' AND page='流量选品' AND status_val=14 ORDER BY report_date DESC;"""
    print("查询需要更新的elaticsearch索引： ")
    print("sql=", sql)
    df = pd.read_sql(sql, con=engine_mysql)
    if df.shape[0]:
        date_info_list = list(df.date_info)
        filter_date_info_list = [date_info for date_info in date_info_list if date_info >= '2024-01']
        if filter_date_info_list:
            return filter_date_info_list
        else:
            print("没有需要更新的索引信息")
            sys.exit(0)


def update_es_fileds(spark, df_main, date_info_list, site_name, run_type):
    if run_type == "real":
        real_date_info_list = [max(date_info_list)]
        print("更新模式针对elasticsearch上最新索引，待更新索引的日期包括：", real_date_info_list)
    else:
        real_date_info_list = list(set(date_info_list) - set([max(date_info_list)]))
        print("更新模式针对elasticsearch上历史索引，待更新索引的日期包括：", real_date_info_list)
    for date_info in real_date_info_list:
        index_name = f"{site_name}_st_detail_month_{date_info}"
        es_asin_sql = f"""
              SELECT asin from es_selection.default_db.{index_name}
          """
        df_es = DorisHelper.spark_import_with_sql(spark, es_asin_sql)
        df_es = df_es.repartition(40)
        df_need_update = df_es.join(
            df_main, on=['asin'], how='inner'
        )
        print(f"Elasticsearch上{site_name} {date_info}需要更新的数据为：")
        df_need_update.cache()
        df_need_update.show(10, truncate=False)
        es_options = {
            "es.nodes": __es_ip__,
            "es.port": __es_port__,
            "es.net.http.auth.user": __es_user__,
            "es.net.http.auth.pass": __es_passwd__,
            "es.mapping.id": "asin",
            "es.resource": f"{index_name}/_doc",
            "es.batch.write.refresh": "false",
            "es.batch.write.retry.wait": "60s",
            "es.batch.size.entries": "5000",
            "es.nodes.wan.only": "false",
            "es.batch.write.concurrency": "60",
            "es.write.operation": "upsert"
        }
        try:
            df_need_update.write.format("org.elasticsearch.spark.sql") \
                .options(**es_options) \
                .mode("append") \
                .save()
            print(f"elasticsearch {index_name} 更新完毕!")
        except Exception as e:
            print("An error occurred while writing to Elasticsearch:", str(e))
            CommonUtil.send_wx_msg(['wujicang', 'wangrui4'], '\u26A0 es用户标记信息更新失败', f'es更新用户标记信息失败：{site_name}, {date_info}')
            pass
    print("elasticsearch 所有数据全部更新完毕")


def get_date_info(spark, file_path, date_info, run_type):
    if run_type == 'real':
        print("最新索引更新模式对应日期为：", date_info)
        return date_info
    else:
        cilent = HdfsUtils.get_hdfs_cilent()
        try:
            status = cilent.status(file_path)
            if status is not None:
                df_read = spark.read.parquet(file_path)
                hdfs_date_info = str(df_read.select("date_info").first()[0])
                print("hdfs上记录的日期为：", hdfs_date_info)
                return hdfs_date_info
        except HdfsError as e:
            print("hdfs上没有记录日期, 传入的日期参数为： ", date_info)
            return date_info


def save_date_info_to_hdfs(spark, date_info, file_path):
    if date_info <= '2024-01-01':
        pass
    else:
        data = [Row(date_info=date_info)]
        df_save = spark.createDataFrame(data)
        df_save.write.mode('overwrite').parquet(file_path)


def get_main_asin(spark, site_name, date_info):
    pg_con_info = DBUtil.get_connection_info("postgresql", "us")
    main_sql = get_user_mask_type_asin_sql(site_name, date_info)
    print("用户标记信息查询sql为： ", main_sql)
    if pg_con_info is not None:
        df_main = SparkUtil.read_jdbc_query(session=spark, url=pg_con_info['url'],
                                                        pwd=pg_con_info['pwd'], username=pg_con_info['username'],
                                                        query=main_sql)
        df_main = df_main.repartition(40).cache()
        count = df_main.count()
        if count == 0:
            print("没有需要更新的数据")
            sys.exit(0)
        else:
            print("待更新的数据量为：", count)
            df_main.show(10, truncate=False)
            return df_main
    else:
       pass


def main(site_name, date_info, run_type):
    file_path = f"/tmp/wangrui/usr_mask/{site_name}/date_info.parquet"
    spark = create_spark(site_name)
    cur_date_info = get_date_info(spark, file_path, date_info, run_type)
    df_main = get_main_asin(spark, site_name, cur_date_info)
    date_info_list = get_es_index_name(site_name)
    update_es_fileds(spark, df_main, date_info_list, site_name, run_type)
    save_date_info_to_hdfs(spark, date_info, file_path)
    spark.stop()


def create_spark(site_name):
    return SparkSession.builder \
                .master("yarn") \
                .appName(f"es_update_with_usr_mask: {site_name}") \
                .config("spark.sql.warehouse.dir", __warehouse_dir__) \
                .config("spark.metastore.uris", __metastore_uris__) \
                .config("spark.network.timeout", 1000000) \
                .config("spark.sql.orc.mergeSchema", True) \
                .config("spark.sql.parquet.compression.codec", "lzo") \
                .config("spark.driver.maxResultSize", "10g") \
                .config("spark.sql.autoBroadcastJoinThreshold", -1) \
                .config("spark.sql.shuffle.partitions", 100) \
                .config("spark.executor.memory", "15g") \
                .config("spark.executor.cores", "4") \
                .config("spark.executor.instances", "10") \
                .config("spark.driver.memory", "15g") \
                .config("spark.yarn.queue", "spark") \
                .enableHiveSupport() \
                .getOrCreate()


if __name__ == '__main__':
    arguments = sys.argv[1:]
    site_name = sys.argv[1]  # 参数1：站点
    run_type = sys.argv[2]   # 参数2：运行类型(real:最新索引;history:历史索引)
    if run_type == 'history':
        if len(arguments) == 3:  # 参数3：数据日期范围：如果针对最新索引,不用传入
            date_info = sys.argv[3]
        else:
            date_info = '2024-01-01'
    else:
        date_info = '2024-01-01'
    main(site_name, date_info, run_type)