poyee-data-warehouse/dw_base/datax/job_config_generator.py

# -*- coding:utf-8 -*-

import json
import os
from configparser import ConfigParser
from typing import Dict, Optional

from dw_base.datax.datax_constants import *
from dw_base.datax.plugins.plugin_factory import PluginFactory
from dw_base.datax.tuning import load_tuning_conf, merge_speed
from dw_base.utils.file_utils import delete_file, get_abs_path


class JobConfigGenerator(object):
    """
    生成 DataX 作业配置文件（json）。

    speed 三参走三级合并：L1 conf/datax-tuning.conf < L2 ini [speed] 段 < L3 cli_speed_overrides。
    合并后：channel 写 job.setting.speed（DataX 按 channel 数并发）；byte/record 只写
    core.transport.channel.speed 当每通道限速——不可再塞进 job.setting.speed，否则并发被压成 1。
    """

    def __init__(self, base_dir: str, generator_config: str, start_date: str, stop_date: str, output: str,
                 cli_speed_overrides: Optional[Dict[str, int]] = None):
        """
        Args:
            base_dir: 项目目录
            generator_config: DataX 作业配置生成器配置文件路径（.ini）
            start_date / stop_date: 内部日期
            output: 生成的 DataX json 输出路径
            cli_speed_overrides: L3 CLI 覆盖，形如 {'channel': 20, 'byte': None, 'record': None}
        """
        self.generator_config = get_abs_path(generator_config)
        self.base_dir = base_dir
        self.start_date = start_date
        self.stop_date = stop_date
        self.output = output
        self.cli_speed_overrides = cli_speed_overrides or {}
        self.config_parser = ConfigParser()
        self.config_parser.read(self.generator_config)

    def get_reader(self):
        reader = PluginFactory.get_plugin('reader', self.base_dir, self.config_parser, self.start_date, self.stop_date)
        return reader.configure()

    def get_writer(self):
        writer = PluginFactory.get_plugin('writer', self.base_dir, self.config_parser, self.start_date, self.stop_date)
        return writer.configure()

    def assemble(self):
        # speed 三级合并（L1 conf < L2 ini [speed] < L3 CLI）
        tuning_conf_path = os.path.join(self.base_dir, 'conf', 'datax-tuning.conf')
        l1 = load_tuning_conf(tuning_conf_path)
        merged = merge_speed(l1, self.config_parser, self.cli_speed_overrides)
        # job.setting.speed 只放 channel：DataX 并发数按「全局 byte/record ÷ 每通道 byte/record」推导，
        # 且 byte/record 限制优先级高于 channel。若这里也放 byte/record（与 core 每通道同值），
        # 会算出 全局÷每通道=1 → 把 channel 压成 1 通道。byte/record 只作每通道限速放进 core.transport。
        speed = {
            JOB_SETTING_SPEED_CHANNEL: merged['channel'],
        }
        core_speed = {
            'transport': {
                'channel': {
                    'speed': {
                        'byte': merged['byte'],
                        'record': merged['record'],
                    }
                }
            }
        }
        job_config_json = {
            'job': {
                'content': [
                    {
                        'reader': self.get_reader(),
                        'writer': self.get_writer()
                    }
                ],
                'setting': {
                    'speed': speed
                }
            },
            'core': core_speed
        }
        return job_config_json

    def run(self):
        job_config_json = self.assemble()
        # HDFS Mount的覆盖写入貌似有问题
        delete_file(self.output)
        with open(self.output, 'w') as w:
            json.dump(job_config_json, w, ensure_ascii=False)