poyee-data-warehouse/dw_base/__init__.py

#!/usr/bin/env /usr/bin/python3
# -*- coding:utf-8 -*-

import os
import socket
import sys
import time

from dw_base.utils.env_loader import bootstrap_env

# bootstrap_env() 在 driver/边缘节点加载 conf/env.sh（返回 True）；Spark executor 上 dw_base 以 zip 分发、
# 不含 conf/，返回 False。下方 driver 专属初始化（读 env 变量 / 路径推断 / banner）整块跳过——
# executor 只需能 import 子模块（如 dw_base.tracking.mask，纯 stdlib）供 Python UDF 运行。
if bootstrap_env():
    # HADOOP_CONF_DIR / SPARK_CONF_DIR 由 conf/env.sh 提供默认值经 bootstrap_env 注入（shell 侧 export 优先）：
    #   - HADOOP_CONF_DIR：spark-submit 启动 YARN 校验需要；DataX JVM 不读 classpath conf，HA 由 ini [hadoop_config] 节显式注入
    #   - SPARK_CONF_DIR：pip pyspark 默认指向自身空 conf/，显式指到集群配置才能加载 hive-site.xml，否则 enableHiveSupport 回落 in-memory metastore
    # PYSPARK_*：复用 PYTHON3_PATH，避免与 conf/env.sh 双份硬编码
    os.environ.setdefault('PYSPARK_DRIVER_PYTHON', os.environ['PYTHON3_PATH'])
    os.environ.setdefault('PYSPARK_PYTHON', os.environ['PYTHON3_PATH'])
    os.environ['PYTHONUNBUFFERED'] = 'x'
    PROJECT_ROOT_PATH = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
    PROJECT_NAME = os.path.basename(PROJECT_ROOT_PATH)
    sys.path.append(PROJECT_ROOT_PATH)
    # 公用的Spark UDF文件
    COMMON_SPARK_UDF_FILE = 'dw_base/udf/common/spark_common_udf.py'
    BANNED_USER = 'root'
    RELEASE_USER = os.environ['RELEASE_USER']
    USER = os.environ['USER']
    HOME = os.environ['HOME']
    if USER == BANNED_USER and HOME.startswith('/home'):
        USER = os.path.basename(HOME)
    HOST = socket.gethostname()
    RELEASE_ROOT_DIR = os.environ['RELEASE_ROOT_DIR']

    if not PROJECT_ROOT_PATH.startswith(RELEASE_ROOT_DIR) or USER != RELEASE_USER:
        DO_RESET: str = '\033[0m'
        NORM_RED: str = '\033[0;31m'
        NORM_GRN: str = '\033[0;32m'
        NORM_YEL: str = '\033[0;33m'
        NORM_MGT: str = '\033[0;35m'
        NORM_CYN: str = '\033[0;36m'
    else:
        DO_RESET: str = ''
        NORM_RED: str = ''
        NORM_GRN: str = ''
        NORM_YEL: str = ''
        NORM_MGT: str = ''
        NORM_CYN: str = ''
    IS_RUN_BY_RELEASE_USER = False
    LOG_ROOT_DIR = os.environ['LOG_ROOT_DIR']
    if USER == RELEASE_USER:
        IS_RUN_BY_RELEASE_USER = True
    elif USER == BANNED_USER:
        ERROR_CODE = 18
        print(f'{NORM_MGT}{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} '
              f'{NORM_RED}Project {NORM_GRN}{PROJECT_NAME} '
              f'{NORM_RED}is running by banned user {NORM_GRN}{BANNED_USER}'
              f'{NORM_RED}, exit with error code {NORM_GRN}{ERROR_CODE}'
              f'{DO_RESET}')
        exit(ERROR_CODE)
    else:
        print(f'{NORM_CYN}{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} '
              f'{NORM_MGT}Project {NORM_GRN}{PROJECT_NAME} '
              f'{NORM_MGT}is running in normal user {NORM_GRN}{USER}')
    if PROJECT_ROOT_PATH.startswith(f'{RELEASE_ROOT_DIR}/{PROJECT_NAME}'):
        IS_RUN_IN_RELEASE_DIR = True
        print(f'{NORM_CYN}{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} '
              f'{NORM_MGT}Project {NORM_GRN}{PROJECT_NAME} '
              f'{NORM_MGT}is running in release dir {NORM_GRN}{RELEASE_ROOT_DIR}/{PROJECT_NAME}')
    else:
        IS_RUN_IN_RELEASE_DIR = False
        print(f'{NORM_CYN}{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} '
              f'{NORM_MGT}Project {NORM_GRN}{PROJECT_NAME} '
              f'{NORM_MGT}is running in normal user dir {NORM_GRN}{PROJECT_ROOT_PATH}')

    if not IS_RUN_IN_RELEASE_DIR or USER != RELEASE_USER:
        os.system(f'echo -en "{NORM_GRN}"')
    os.system(f'echo -en "{DO_RESET}"')