poyee-data-warehouse/bin/datax-sync-template-gen.py

#!/usr/bin/env /usr/bin/python3
# -*- coding:utf-8 -*-
"""
PG → HDFS DataX sync ini 模板生成器 + raw 建模 metadata 表 + 表探查。

一次跑同时产出三件：
  1. PG 表探查段（行数估值 + 锚点字段维护质量 + 软删字段命中），落 md 头部
  2. PG 全字段 metadata markdown 表（序号/字段名/中文名/数据类型/主键标识/
     脱敏类型）—— 用于 kb/24 raw 建模文档
  3. 全字段 sync ini 模板 —— 开发者按 md 讨论结果手动裁剪字段 / 改 where /
     加 [mask] / 调 splitPk / 改 writer.path 表名后缀等，再提交到 jobs/raw/{域}/

CLI:
  python3 bin/datax-sync-template-gen.py \\
    -ds postgresql/prod-hobby \\
    -t public.card_group_order_info \\
    [-mask-conf <PATH>] [-o [DIR]]

参数:
  -ds         数据源 ref，形如 {db_type}/{env}-{实例简称}（同 sync ini
              里 dataSource 字段格式）。暂只支持 postgresql。
  -t          schema 限定的表名（如 public.card_group_order_info）。
  -mask-conf  mask 配置 ini 路径（{table}.mask.ini，可选）。传入时按配置
              剔除 trim 字段 + 渲染 [mask] 段，md 脱敏类型列填好；不传时
              全字段输出，md 脱敏类型列空白。**文件不存在直接报错**。
  -o          输出目录（可选；任意三态下 stdout 都同时打印 md + ini）：
              - 不传：仅 stdout
              - 传 -o 不带值：stdout + 落盘 workspace/{yyyymmdd}/{table}.{md,ini}
              - 传 -o <DIR>：stdout + 落盘 <DIR>/{table}.{md,ini}
"""
import argparse
import os
import re
import sys
from configparser import ConfigParser
from datetime import datetime

project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

from dw_base.datax.datasources.data_source_factory import DataSourceFactory
from dw_base.datax.datax_constants import DS_POSTGRE_SQL_JDBC_URL


WORKSPACE_DEFAULT = os.path.join(
    project_root, 'workspace', datetime.now().strftime('%Y%m%d'),
)

# 探查硬编码：增量同步标准锚点字段（推行后端命名标准）
ANCHOR_FIELDS = ('create_time', 'update_time')
# 抽样上限：TABLESAMPLE SYSTEM(1) 按存储页跳跃后再 LIMIT 截断
PROBE_SAMPLE_LIMIT = 1000


def resolve_datasource(ds_ref):
    """复用 plugin.py:34-42 的 ref → DataSource 解析逻辑。

    ds_ref 形如 'postgresql/prod-hobby'，首段为 db_type（同父目录名）。
    datasource ini 落点：项目同级 ../datasource/{ds_ref}.ini。
    """
    ds_type = ds_ref.split('/')[0]
    if ds_type != 'postgresql':
        raise NotImplementedError('暂只支持 postgresql 数据源，收到: ' + ds_type)
    ds_file_path = os.path.normpath(
        os.path.join(project_root, '..', 'datasource', ds_ref + '.ini'))
    if not os.path.isfile(ds_file_path):
        raise FileNotFoundError('数据源 ini 不存在: ' + ds_file_path)
    return DataSourceFactory.get_data_source(ds_type, ds_file_path)


def parse_jdbc_url(jdbc_url):
    """从 jdbc:postgresql://host:port/database 抽 (host, port, database)。"""
    m = re.match(r'jdbc:postgresql://([^:/]+)(?::(\d+))?/(.+)', jdbc_url)
    if not m:
        raise ValueError('无法解析 PG jdbcUrl: ' + jdbc_url)
    return m.group(1), int(m.group(2) or 5432), m.group(3)


def query_columns_full(conn, schema, table):
    """带序号 / 类型 / 主键标识的全字段 metadata 查询，按 attnum 排序。

    返回 [(attnum, attname, comment, pg_type, pk_flag), ...]
    """
    cur = conn.cursor()
    cur.execute("""
        SELECT
            a.attnum,
            a.attname,
            pg_catalog.col_description(a.attrelid, a.attnum),
            pg_catalog.format_type(a.atttypid, a.atttypmod),
            CASE WHEN EXISTS (
                SELECT 1 FROM pg_index i
                WHERE i.indrelid = a.attrelid AND i.indisprimary
                  AND a.attnum = ANY(i.indkey)
            ) THEN 'PK' ELSE '' END
        FROM pg_catalog.pg_attribute a
        JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
        JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
        WHERE n.nspname = %s AND c.relname = %s
          AND a.attnum > 0 AND NOT a.attisdropped
        ORDER BY a.attnum
    """, (schema, table))
    return cur.fetchall()


def probe_table(conn, schema, table, full_rows):
    """对表做行数估值 + PK + 锚点抽样 + 近期 update_time + 主键序时间范围 + 软删命中。

    - 行数：pg_class.reltuples 估值
    - PK：单/复合/无 + 是否自增（pg_get_serial_sequence）
    - 锚点：create_time / update_time 存在性 + 抽样非空率（TABLESAMPLE SYSTEM(1) LIMIT 1000）
    - 近期 update_time 非空率：仅当单 PK + update_time 存在；ORDER BY pk DESC LIMIT 1000
    - create_time 主键序范围：仅当单列自增 PK + create_time 存在；ORDER BY pk ASC/DESC LIMIT 1
    - 软删：full_rows 筛 'del' 子串（不区分大小写）

    返回 dict 见 render_probe_md 引用字段。
    """
    cur = conn.cursor()

    cur.execute("""
        SELECT c.reltuples::bigint
        FROM pg_catalog.pg_class c
        JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
        WHERE n.nspname = %s AND c.relname = %s
    """, (schema, table))
    row = cur.fetchone()
    reltuples = int(row[0]) if row and row[0] is not None else 0

    pk_cols = [r[1] for r in full_rows if r[4] == 'PK']
    pk_auto_increment = False
    if len(pk_cols) == 1:
        # pg_get_serial_sequence 只识别 OWNED BY 关联的 sequence——
        # 业务库手工建的 sequence 没 OWNED BY 标记会漏判，所以同时查 attidentity
        # （PG 10+ IDENTITY 列）和 default 表达式（含 nextval 即视为自增）。
        cur.execute("""
            SELECT a.attidentity, pg_get_expr(ad.adbin, ad.adrelid)
            FROM pg_attribute a
            LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum
            JOIN pg_class c ON c.oid = a.attrelid
            JOIN pg_namespace n ON n.oid = c.relnamespace
            WHERE n.nspname = %s AND c.relname = %s AND a.attname = %s
        """, (schema, table, pk_cols[0]))
        r = cur.fetchone()
        if r:
            attidentity, default_expr = r[0], r[1]
            pk_auto_increment = (
                attidentity in ('a', 'd')
                or (default_expr is not None and 'nextval' in default_expr.lower())
            )

    field_names = {r[1] for r in full_rows}
    anchor = {col: {'exists': col in field_names, 'notnull': None}
              for col in ANCHOR_FIELDS}

    sample_total = 0
    present = [c for c in ANCHOR_FIELDS if anchor[c]['exists']]
    if present:
        notnull_select = ', '.join('count("{}")'.format(c) for c in present)
        sql = (
            'SELECT count(*), {nn} FROM '
            '(SELECT * FROM "{schema}"."{table}" '
            ' TABLESAMPLE SYSTEM(1) LIMIT {lim}) AS sub'
        ).format(nn=notnull_select, schema=schema, table=table,
                 lim=PROBE_SAMPLE_LIMIT)
        cur.execute(sql)
        result = cur.fetchone()
        sample_total = int(result[0])
        for i, c in enumerate(present):
            anchor[c]['notnull'] = int(result[i + 1])

    recent_total = None
    recent_update_notnull = None
    if len(pk_cols) == 1 and anchor['update_time']['exists']:
        sql = (
            'SELECT count(*), count("update_time") FROM '
            '(SELECT update_time FROM "{schema}"."{table}" '
            ' ORDER BY "{pk}" DESC LIMIT {lim}) AS sub'
        ).format(schema=schema, table=table, pk=pk_cols[0],
                 lim=PROBE_SAMPLE_LIMIT)
        cur.execute(sql)
        result = cur.fetchone()
        recent_total = int(result[0])
        recent_update_notnull = int(result[1])

    create_time_earliest = None
    create_time_latest = None
    if pk_auto_increment and anchor['create_time']['exists']:
        sql = (
            'SELECT '
            '(SELECT create_time FROM "{schema}"."{table}" '
            ' ORDER BY "{pk}" ASC LIMIT 1), '
            '(SELECT create_time FROM "{schema}"."{table}" '
            ' ORDER BY "{pk}" DESC LIMIT 1)'
        ).format(schema=schema, table=table, pk=pk_cols[0])
        cur.execute(sql)
        result = cur.fetchone()
        create_time_earliest = result[0]
        create_time_latest = result[1]

    del_candidates = sorted(r[1] for r in full_rows if 'del' in r[1].lower())

    return {
        'reltuples': reltuples,
        'pk_cols': pk_cols,
        'pk_auto_increment': pk_auto_increment,
        'sample_total': sample_total,
        'anchor': anchor,
        'recent_total': recent_total,
        'recent_update_notnull': recent_update_notnull,
        'create_time_earliest': create_time_earliest,
        'create_time_latest': create_time_latest,
        'del_candidates': del_candidates,
    }


def render_probe_md(stats):
    """渲染探查段 markdown。"""
    lines = ['### 探查', '']
    lines.append('- 行数估值（pg_class.reltuples）：{:,}'.format(stats['reltuples']))

    pk_cols = stats['pk_cols']
    if not pk_cols:
        pk_desc = '无（DataX channel 无法并行）'
    elif len(pk_cols) > 1:
        pk_desc = '复合 ({}) （DataX splitPk 不支持复合，退串行）'.format(
            ', '.join('`{}`'.format(c) for c in pk_cols))
    elif stats['pk_auto_increment']:
        pk_desc = '`{}`（自增）'.format(pk_cols[0])
    else:
        pk_desc = '`{}`（非自增，DataX channel 切分分布可能不均）'.format(pk_cols[0])
    lines.append('- 主键：' + pk_desc)

    lines.append('- 锚点字段：')
    total = stats['sample_total']
    for col in ANCHOR_FIELDS:
        s = stats['anchor'][col]
        if not s['exists']:
            lines.append('  - `{}`：缺失'.format(col))
            continue
        if total > 0 and s['notnull'] is not None:
            pct = 100.0 * s['notnull'] / total
            base = '`{}`：存在；整体非空率 {:.1f}% ({}/{} 抽样)'.format(
                col, pct, s['notnull'], total)
        else:
            base = '`{}`：存在；抽样无数据'.format(col)
        if col == 'create_time' and stats['create_time_earliest']:
            base += '；按主键序范围 {} ~ {}'.format(
                stats['create_time_earliest'], stats['create_time_latest'])
        lines.append('  - ' + base)
        if col == 'update_time' and stats['recent_total'] is not None:
            rt = stats['recent_total']
            rnn = stats['recent_update_notnull']
            rpct = 100.0 * rnn / rt if rt else 0.0
            lines.append('    - 近期非空率 {:.1f}% ({}/{} 最近 1000 行)'.format(
                rpct, rnn, rt))

    if stats['del_candidates']:
        lines.append('- 软删字段（含 `del` 子串）：' + ', '.join(
            '`{}`'.format(c) for c in stats['del_candidates']))
    else:
        lines.append('- 软删字段（含 `del` 子串）：未命中')
    return '\n'.join(lines) + '\n'


def _resolve_to_project_root(path):
    """相对路径按项目根解析，绝对路径原样返回。

    复用 dw_base.datax.entry._resolve_relative_to_base 的逻辑——
    任何 cwd 跑此脚本都能找到 mask conf 等相对路径资源，
    与项目其他 bin 入口（datax-hive-import-starter 等）行为一致。
    """
    if os.path.isabs(path):
        return path
    return os.path.join(project_root, path)


def load_mask_conf(path):
    """读 mask 配置 ini，返回 {field: method} dict。

    格式（与 jobs/raw/{域}/{table}.mask.ini 同款）：
        [mask]
        field1 = method1
        field2 = method2

    method ∈ trim / md5 / month_trunc / mask_middle / keep_first_n / keep_last_n
    - trim：整字段不入 raw（reader column 不查询）
    - 其他：字段入 raw，由 dw_base.datax.mask 在 reader 端脱敏

    文件不存在直接 raise FileNotFoundError（不静默失败）。
    """
    if not os.path.isfile(path):
        raise FileNotFoundError('mask 配置不存在: ' + path)
    cp = ConfigParser()
    cp.read(path, encoding='utf-8')
    if not cp.has_section('mask'):
        return {}
    return dict(cp.items('mask'))


def render_schema_md(rows, mask_dict=None):
    """输出 markdown 表格：序号 / 字段名 / 中文名 / 数据类型 / 主键标识 / 脱敏类型。

    mask_dict 不传时脱敏类型列为空白；传入时填字段对应的 method（含 trim）。
    """
    lines = [
        '| 序号 | 字段名 | 中文名 | 数据类型 | 主键标识 | 脱敏类型 |',
        '| --- | --- | --- | --- | --- | --- |',
    ]
    methods = mask_dict or {}
    for num, name, comment, typ, pk in rows:
        method = methods.get(name, '')
        lines.append('| {} | `{}` | {} | {} | {} | {} |'.format(
            num, name, comment or '', typ, pk, method))
    return '\n'.join(lines) + '\n'


def render_template(ds_ref, database, schema, table, columns, pk, mask_methods=None):
    """渲染 sync ini 模板。

    columns: [(name, comment), ...] 已剔除 trim 字段，保持 PG 原顺序
    mask_methods: {field: method} 仅含非 trim 方法（mask_middle / month_trunc 等），
                  渲染 [mask] 段；空 dict 或 None 时不渲染 [mask] 段
    """
    column_str = ','.join(c for c, _ in columns)
    today = datetime.now().strftime('%Y-%m-%d')

    if mask_methods:
        mask_lines = '\n'.join('{} = {}'.format(f, m) for f, m in mask_methods.items())
        mask_section = '[mask]\n' + mask_lines + '\n\n'
    else:
        mask_section = ''

    return (
        '; 作者：<TODO>\n'
        '; 日期：{today}\n'
        '; 工单：<TODO>\n'
        '; 目的：PG {database}.{schema}.{table} → Hive raw.<TODO> 同步模板\n'
        '; 状态：[待执行]\n'
        '; 备注：自动生成的全字段参考模板。开发者按需裁剪字段 / 改 where / 加 mask 段 /\n'
        ';       调 splitPk / 改 writer.path 表名后缀（_inc_d / _his_o 等）\n'
        ';\n'
        '; 配套 DDL：manual/ddl/raw/<TODO_domain>/raw_<TODO>_create.sql\n'
        '\n'
        '[reader]\n'
        'dataSource = {ds_ref}\n'
        'database = {database}\n'
        'table = {schema}.{table}\n'
        'column = {column_str}\n'
        'columnType =\n'
        "where = update_time >= '${{start_date}}' AND update_time < '${{stop_date}}'\n"
        'querySql =\n'
        'splitPk = {pk}\n'
        'fetchSize = 1000\n'
        '\n'
        '{mask_section}'
        '[writer]\n'
        'dataSource = hdfs/<TODO>\n'
        'path = /user/hive/warehouse/raw.db/{table}_TODO_d/dt=${{dt}}/\n'
        'column = {column_str}\n'
        'columnType =\n'
        'fileType = orc\n'
        'fileName = {table}_TODO_d\n'
        'encoding = UTF-8\n'
        'writeMode = truncate\n'
        'fieldDelimiter = \\t\n'
    ).format(
        today=today, ds_ref=ds_ref, database=database, schema=schema,
        table=table, column_str=column_str, pk=pk, mask_section=mask_section,
    )


def main():
    parser = argparse.ArgumentParser(
        prog='datax-sync-template-gen',
        description='PG → HDFS DataX sync ini 模板生成器（全字段参考模板）',
    )
    parser.add_argument('-ds', required=True, metavar='DS_REF',
                        help='数据源 ref，形如 postgresql/prod-hobby（同 sync ini dataSource 字段）')
    parser.add_argument('-t', required=True, metavar='SCHEMA.TABLE',
                        help='schema 限定的表名（如 public.card_group_order_info）')
    parser.add_argument('-o', nargs='?', const=WORKSPACE_DEFAULT, default=None, metavar='DIR',
                        help='输出目录（任意三态 stdout 始终打印 md + ini；不传仅 stdout；不带值额外落盘 workspace/{yyyymmdd}/；带值额外落盘 <DIR>/）')
    parser.add_argument('-mask-conf', default=None, metavar='PATH', dest='mask_conf',
                        help='mask 配置 ini 路径（{table}.mask.ini）。传入时按配置剔除 trim 字段 + 渲染 [mask] 段，md 脱敏类型列填好；不传时全字段输出，md 脱敏类型列空白')
    args = parser.parse_args()

    if '.' not in args.t:
        print('-t 必须 schema.table 格式，收到: ' + args.t, file=sys.stderr)
        sys.exit(2)
    schema, table = args.t.split('.', 1)

    ds = resolve_datasource(args.ds)
    ds_dict = ds.parse()
    jdbc_url = ds_dict[DS_POSTGRE_SQL_JDBC_URL]
    user = ds_dict['username']
    password = ds_dict['password']
    host, port, database = parse_jdbc_url(jdbc_url)

    import pg8000.dbapi
    conn = pg8000.dbapi.connect(
        host=host, port=port, database=database,
        user=user, password=password,
    )
    try:
        full_rows = query_columns_full(conn, schema, table)
        if not full_rows:
            raise ValueError('表不存在或无字段: {}.{}'.format(schema, table))
        probe_stats = probe_table(conn, schema, table, full_rows)
    finally:
        conn.close()

    # full_rows: [(attnum, attname, comment, pg_type, pk_flag), ...]
    if args.mask_conf:
        mask_path = _resolve_to_project_root(args.mask_conf)
        mask_dict = load_mask_conf(mask_path)
    else:
        mask_dict = {}

    # mask 配置含表中不存在字段时 stderr 警告（不阻断）
    pg_field_set = {r[1] for r in full_rows}
    unknown_fields = [f for f in mask_dict if f not in pg_field_set]
    if unknown_fields:
        print('警告：mask 配置含表中不存在字段（已忽略）: ' + ', '.join(unknown_fields),
              file=sys.stderr)

    trim_set = {f for f, m in mask_dict.items() if m == 'trim'}
    non_trim_mask = {f: m for f, m in mask_dict.items() if m != 'trim'}

    # 已剔除 trim 字段的 column 列表，保持 PG 原顺序（attnum 升序）
    columns = [(r[1], r[2] or '') for r in full_rows if r[1] not in trim_set]

    pk_names = [r[1] for r in full_rows if r[4] == 'PK']
    pk = pk_names[0] if len(pk_names) == 1 and pk_names[0] not in trim_set else ''

    probe_md = render_probe_md(probe_stats)
    schema_md = render_schema_md(full_rows, mask_dict)
    md_content = probe_md + '\n### 字段\n\n' + schema_md
    ini_content = render_template(args.ds, database, schema, table, columns, pk, non_trim_mask)

    # stdout 始终打印（先 md 表后 ini 模板），传 -o 时再额外落盘
    sys.stdout.write(md_content)
    sys.stdout.write('\n')
    sys.stdout.write(ini_content)

    if args.o is not None:
        os.makedirs(args.o, exist_ok=True)
        md_path = os.path.join(args.o, table + '.md')
        ini_path = os.path.join(args.o, table + '.ini')
        with open(md_path, 'w', encoding='utf-8') as f:
            f.write(md_content)
        with open(ini_path, 'w', encoding='utf-8') as f:
            f.write(ini_content)
        print('已写入: ' + md_path, file=sys.stderr)
        print('已写入: ' + ini_path, file=sys.stderr)


if __name__ == '__main__':
    main()