energy-data-trans/utils/log/trans_log.py

# -*- coding: utf-8 -*-
# @Time    : 2024/5/16
# @Author  : 魏志亮

import datetime
import logging
import os
import sys

from utils.conf.read_conf import read_conf, yaml_conf


def set_trance_id(trace_id):
    """设置当前线程的链路ID"""
    os.environ['trace_id'] = trace_id


class ContextFilter(logging.Filter):
    """一个自定义的日志过滤器，用于在日志记录中添加链路ID"""

    def filter(self, record):
        record.trace_id = ''
        if 'trace_id' in os.environ.keys():
            record.trace_id = os.environ['trace_id']

        return True


logger = logging.getLogger("etl_tools")
logger.setLevel(logging.INFO)
stout_handle = logging.StreamHandler(sys.stdout)
stout_handle.setFormatter(
    logging.Formatter("%(asctime)s-%(trace_id)s: %(message)s"))
stout_handle.setLevel(logging.INFO)
stout_handle.addFilter(ContextFilter())
logger.addHandler(stout_handle)

config = yaml_conf(os.environ['ETL_CONF'])
log_path_dir = read_conf(config, 'log_path_dir', "/data/logs")

log_path = log_path_dir + os.sep + r'etl_tools_' + (os.environ['env'] if 'env' in os.environ else 'dev')
file_path = os.path.join(log_path)

if not os.path.exists(file_path):
    os.makedirs(file_path, exist_ok=True)
file_name = file_path + os.sep + str(datetime.date.today()) + '.log'

file_handler = logging.FileHandler(file_name, encoding='utf-8')
file_handler.setFormatter(
    logging.Formatter("%(asctime)s-%(trace_id)s: %(message)s"))
file_handler.setLevel(logging.INFO)
file_handler.addFilter(ContextFilter())
logger.addHandler(file_handler)


def trans_print(*args):
    logger.info("  ".join([str(a) for a in args]))