poyee-data-warehouse/bin/spark-sql-starter.py

#!/usr/bin/env /usr/bin/python3
# -*- coding:utf-8 -*-
"""
  -- 设置 SparkSession 名称（血缘分析）
  SET spark.app.name=crl_dwd.dwd_crl_xxx;
  -- 设置 Spark 配置
  SET spark.xxx.yyy.zzz=xyz;
  -- 引用 UDF
  ADD FILE dw_base/udf/business/spark_xxx_udf.py;
  -- 声明变量
  SET DT_START=20210101;
  SET TOPIC=xxx;
  SET KEY=CourtAnn;
  SET ODS_TABLE=crl_ods.ods_crl_xxx;
  -- 查看数据行数
  SET LIMIT=1000;
"""
import os
import sys
from typing import Dict, List

base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_dir)
from dw_base import *
from dw_base.common.config_constants import K_DT, KL_HELP
from dw_base.common.container import ValueContainer
from dw_base.spark.spark_sql import SparkSQL
from dw_base.spark.spark_utils import analyse_session_name
from dw_base.utils.common_utils import exist
from dw_base.utils.config_utils import parse_args
from dw_base.utils.datetime_utils import get_yesterday, get_date_range
from dw_base.utils.file_utils import list_files
from dw_base.utils.log_utils import pretty_print, CURRENT_LOG_FILE, get_log_file_path


def usage(code: int):
    print(
        f'{NORM_MGT}Usage: {sys.argv[0]}\n'
        f'{NORM_CYN}\t[-h/-H/--h/--H/--help]         打印脚本使用方法{DO_RESET}'
    )
    print(
        f'{NORM_MGT}Usage: {sys.argv[0]}\n'
        f'{NORM_GRN}\t<[-]-f< /=>sql file>           要执行的SQL文件，可传入多个，按传入顺序执行，与参数-d互斥\n'
        f'{NORM_GRN}\t<[-]-d< /=>sql file directory> 要执行的SQL文件夹，可传入多个，按传入顺序执行\n'
        f'{NORM_CYN}\t[[-]-u< /=>udf file]           被引用到的udf文件（.py），注意多个udf文件内定义的方法有重名时，只会取靠前的\n'
        f'{NORM_CYN}\t[[-]-sc spark-config]          Spark自定义配置，格式为config:config-value或config=config-value\n'
        f'{NORM_CYN}\t[[-]-dt< /=>date with format]  %Y%m%d 或 yyyyMMdd 格式的日期（命令行 > 默认）\n'
        f'{NORM_CYN}\t                               可以以四种形式传入日期：\n'
        f'{NORM_CYN}\t                                   1. 20211101，表示具体日期\n'
        f'{NORM_CYN}\t                                   2. 20211101-，表示20211101至昨天\n'
        f'{NORM_CYN}\t                                   3. 20211101-20211107，表示20211101至20211107\n'
        f'{NORM_CYN}\t                                   4. 20211101,20211103，表示离散的日期20211101、20211103\n'
        f'{NORM_CYN}\t[[-]-p sql-parameter]          SQL脚本的参数，格式为parameter:parameter-value或parameter=parameter-value\n'
        f'{NORM_CYN}\t[[-]-n]                        在串行执行多个SQL文件时，一个失败并不影响后续文件的执行，但最后程序会非正常结束，即返回非0值（失败的文件个数）\n'
        f'{NORM_CYN}\t[[-]-0]                        在串行执行多个SQL文件时，一个失败并不影响后续文件的执行，且最后程序正常结束（请谨慎使用）\n'
        f'{NORM_CYN}\t                               上述两个参数都不传的情况下，某个SQL文件执行失败时，后续文件都不再执行'
        f'{DO_RESET}'
    )
    exit(code)


def parse_cli_spark_config() -> Dict[str, str]:
    cli_spark_config = {}
    spark_configs = CONFIG.get('sc', [])
    if isinstance(spark_configs, str):
        spark_configs = [spark_configs]
    for elem in spark_configs:
        if elem.__contains__('='):
            spark_config, config_value = elem.split('=')
        elif elem.__contains__(':'):
            spark_config, config_value = elem.split(':')
        else:
            pretty_print(f'{NORM_YEL}无效的 Spark 配置 {NORM_GRN}{elem}')
            continue
        if cli_spark_config.__contains__(spark_config):
            pretty_print(f'{NORM_YEL}命令行多次传入了 Spark 配置 {NORM_GRN}{spark_config}'
                         f'{NORM_YEL}, 原值 {NORM_GRN}{cli_spark_config[spark_config]} '
                         f'{NORM_YEL}将被覆盖为新值 {NORM_GRN}{config_value}')
        cli_spark_config[spark_config] = config_value
    return cli_spark_config


def parse_sql_parameters() -> Dict[str, str]:
    sql_parameters = {}
    parameter_configs = CONFIG.get('p', [])
    if isinstance(parameter_configs, str):
        parameter_configs = [parameter_configs]
    for elem in parameter_configs:
        if elem.__contains__('='):
            parameter_key, parameter_value = elem.split('=')
        elif elem.__contains__(':'):
            parameter_key, parameter_value = elem.split(':')
        else:
            pretty_print(f'{NORM_YEL}无效的 SQL 参数 {NORM_GRN}{elem}')
            continue
        if parameter_key == 'dt':
            continue
        if sql_parameters.__contains__(parameter_key):
            pretty_print(f'{NORM_YEL}命令行多次传入了 SQL 参数 {NORM_GRN}{parameter_key}'
                         f'{NORM_YEL}, 原值 {NORM_GRN}{sql_parameters[parameter_key]} '
                         f'{NORM_YEL}将被覆盖为新值 {NORM_GRN}{parameter_value}')
        sql_parameters[parameter_key] = parameter_value
    return sql_parameters


def run_sql_file(sql_file: str, dt: str, is_last: bool):
    try:
        spark_session_name = CONFIG.get('ssn', analyse_session_name(sql_file))
        sql_file_name = os.path.splitext(os.path.basename(sql_file))[0]
        CURRENT_LOG_FILE.set(get_log_file_path('spark-sql', dt, sql_file_name))
        spark_sql = SparkSQL(spark_session_name, udf_files=udf_files, extra_spark_config=extra_spark_config)
        is_export = CONFIG.get('export')
        if is_last and is_export is True:
            sql_file_name = os.path.splitext(os.path.basename(sql_file))[0]
            delimiter = CONFIG.get('delimiter', ',')
            spark_sql.export_data(sql_file_name, sql_file, truncate=False,
                                  delimiter=delimiter, dt=dt, partition=1, **sql_parameters)
        else:
            spark_sql.execute(sql_file, check_parameter=True, dt=dt, **sql_parameters)
    except Exception as e:
        message = f'执行SQL文件 {sql_file} 失败（dt={dt}）'
        if IS_RUN_IN_RELEASE_DIR and IS_RUN_BY_RELEASE_USER:
            message += f'（可访问网址 http://{HOST}/log/spark-sql/{dt}/{os.path.basename(CURRENT_LOG_FILE.get())} 查看日志）'
            print(message)
        ERROR_COUNT.set(ERROR_COUNT.get() + 1)
        if CONTINUE_ALL_ON_ERROR and CONTINUE_ALL_ON_ERROR is True:
            return
        elif CONTINUE_NEXT_ON_ERROR and CONTINUE_NEXT_ON_ERROR is True:
            return
        else:
            raise e


def run_sql_files(sql_files: List[str], dt: str):
    for index in range(len(sql_files)):
        sql_file = sql_files[index]
        if sql_file == 'create_table.sql' or sql_file == 'create-table.sql':
            continue
        sql_file_dir = os.path.basename(os.path.basename(sql_file))
        if sql_file_dir == 'create_table' or sql_file_dir == 'create-table':
            continue
        run_sql_file(sql_file, dt, len(sql_files) == index + 1)


def run_sql_directories(dt: str):
    directories = CONFIG.get('d', [])
    if isinstance(directories, str):
        directories = [directories]
    for each_dir in directories:
        sql_files_in_dir = list_files(each_dir, extension='.sql')
        sql_files_in_dir = sorted(sql_files_in_dir)
        if len(sql_files_in_dir) == 0:
            pretty_print(f'{NORM_RED}文件夹 {NORM_GRN}{each_dir}{NORM_RED} 中未找到任何可执行的SQL文件')
            usage(1)
        run_sql_files(sql_files_in_dir, dt)


if __name__ == '__main__':
    pretty_print(f'{NORM_MGT}{sys.argv[0]} 收到参数：{NORM_GRN}{" ".join(sys.argv[1:])}')

    CONFIG, _ = parse_args(sys.argv[1:])
    # 未提供任何参数或查看帮助
    if len(sys.argv) == 1 or exist(CONFIG, KL_HELP):
        usage(0)
    CONTINUE_NEXT_ON_ERROR = CONFIG.get('n')
    CONTINUE_ALL_ON_ERROR = CONFIG.get('0')
    ERROR_COUNT = ValueContainer(0)
    date_range = get_date_range(CONFIG.get(K_DT, get_yesterday()))
    udf_files = CONFIG.get('u', [])
    if isinstance(udf_files, str):
        udf_files = [udf_files]
    if not udf_files.__contains__(COMMON_SPARK_UDF_FILE):
        udf_files.insert(0, COMMON_SPARK_UDF_FILE)
    extra_spark_config = parse_cli_spark_config()
    sql_parameters = parse_sql_parameters()
    if CONFIG.__contains__('f'):
        arg_sql_files = CONFIG.get('f', [])
        if isinstance(arg_sql_files, str):
            arg_sql_files = [arg_sql_files]
        if len(arg_sql_files) == 0:
            pretty_print(f'{NORM_RED}未提供任何可执行的SQL文件')
            usage(1)
        for each_dt in date_range:
            run_sql_files(arg_sql_files, each_dt)
    elif CONFIG.__contains__('d'):
        for each_dt in date_range:
            run_sql_directories(each_dt)
    else:
        usage(0)
    if not CONTINUE_ALL_ON_ERROR or CONTINUE_ALL_ON_ERROR is False:
        exit(ERROR_COUNT.get())