poyee-data-warehouse/dw_base/ds/cli.py

# -*- coding:utf-8 -*-
"""
DolphinScheduler API CLI 入口。

用法：
  python -m dw_base.ds.cli get  <path> [-p k=v ...]
  python -m dw_base.ds.cli post <path> [-j '<json-body>' | -d k=v ...]
  python -m dw_base.ds.cli put  <path> [-j '<json-body>' | -d k=v ...]

示例：
  python -m dw_base.ds.cli get /projects -p pageSize=10 -p pageNo=1
  python -m dw_base.ds.cli post /projects/123/workflow-definition -d name=wf -d desc=hi
  python -m dw_base.ds.cli put /projects/123/workflow-definition/100 -j '{"name":"wf"}'
"""
import argparse
import json
import sys

from dw_base.ds.api import DSClient


def _parse_kv(items):
    d = {}
    for it in items or []:
        if '=' not in it:
            sys.stderr.write('参数 -p 需要 k=v 格式，跳过：' + it + '\n')
            continue
        k, v = it.split('=', 1)
        d[k] = v
    return d


def main(argv=None):
    parser = argparse.ArgumentParser(prog='dw_base.ds.cli', description='DolphinScheduler API CLI')
    sub = parser.add_subparsers(dest='cmd')
    sub.required = True

    g = sub.add_parser('get', help='HTTP GET')
    g.add_argument('path', help='API path，如 /projects')
    g.add_argument('-p', action='append', default=[], metavar='k=v', help='query 参数（可多次）')
    g.add_argument('-o', dest='output', default=None, metavar='FILE', help='写 JSON 到文件（不打 stdout）')

    p = sub.add_parser('post', help='HTTP POST')
    p.add_argument('path')
    p.add_argument('-j', dest='body', default=None, metavar='JSON', help='JSON body（与 -d 互斥）')
    p.add_argument('-d', dest='form', action='append', default=[], metavar='k=v',
                   help='form-data 参数（可多次，DS 3.x 多用 form-encoded；与 -j 互斥）')
    p.add_argument('-o', dest='output', default=None, metavar='FILE', help='写 JSON 到文件（不打 stdout）')

    u = sub.add_parser('put', help='HTTP PUT（如 workflow-definition 改造）')
    u.add_argument('path')
    u.add_argument('-j', dest='body', default=None, metavar='JSON', help='JSON body（与 -d 互斥）')
    u.add_argument('-d', dest='form', action='append', default=[], metavar='k=v',
                   help='form-data 参数（可多次；与 -j 互斥）')
    u.add_argument('-o', dest='output', default=None, metavar='FILE', help='写 JSON 到文件（不打 stdout）')

    args = parser.parse_args(argv)
    client = DSClient()

    if args.cmd == 'get':
        result = client.get(args.path, params=_parse_kv(args.p))
    else:
        # post / put 共用 body / form 解析与互斥校验
        body = json.loads(args.body) if args.body else None
        form = _parse_kv(args.form) if args.form else None
        if body is not None and form is not None:
            sys.stderr.write('-j 与 -d 互斥\n')
            sys.exit(1)
        method = client.put if args.cmd == 'put' else client.post
        result = method(args.path, json_body=body, form_data=form)

    text = json.dumps(result, ensure_ascii=False, indent=2)
    if args.output:
        with open(args.output, 'w', encoding='utf-8') as f:
            f.write(text)
            f.write('\n')
        sys.stderr.write('[ds.cli] 已写到 {}\n'.format(args.output))
    else:
        print(text)


if __name__ == '__main__':
    main()