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()


# HADOOP_CONF_DIR：spark-submit 启动 YARN 校验需要；DataX JVM 不读 classpath conf，HA 由 ini [hadoop_config] 节显式注入
os.environ['HADOOP_CONF_DIR'] = '/etc/hadoop/conf'
# SPARK_CONF_DIR：pip pyspark 默认指向自身空 conf/，显式指到集群配置才能加载 hive-site.xml，否则 enableHiveSupport 回落 in-memory metastore
os.environ.setdefault('SPARK_CONF_DIR', '/etc/spark/conf')
os.environ["PYSPARK_DRIVER_PYTHON"] = "/usr/bin/python3"
os.environ["PYSPARK_PYTHON"] = "/usr/bin/python3"
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}"')