#!/usr/bin/env python3 import argparse import base64 import json import re import shutil import subprocess import sys from datetime import datetime from pathlib import Path SCRIPT_DIR = Path(__file__).resolve().parent INIT_SCRIPT = SCRIPT_DIR / "dumpling_packed_backup_init.py" NAMESPACE = "tidb-serverless" SOURCE_STS = "tikv-api" SOURCE_CONTAINER = "tikv-api" SERVICE_ACCOUNT = "backup-restore" STARTER = "/usr/local/bin/dumpling_packed_backup_init.py" DEFAULT_RUNNER_IMAGE = "public.ecr.aws/lambda/python:3.13" RUNNER_PYTHON = "/var/lang/bin/python3.13" DEFAULT_DUMPLING_IMAGE = ( "us-docker.pkg.dev/pingcap-testing-account/dev/pingcap/tidb/images/" "dumpling:exp-export-packed-aws-profile-proof-d79db4a" ) DEFAULT_CSE_CTL_IMAGE = ( "us-docker.pkg.dev/pingcap-testing-account/dev/tikv/tikv/" "image:juncen-exp-export-packed-csv-aws-profile-proof-aad1b20-nextgen" ) DEFAULT_MAX_RETRIES = 16 DEFAULT_TTL_SECONDS_AFTER_FINISHED = 30 * 24 * 60 * 60 LOCAL_GPG_CONFIG_KEY = "gpg-public-key" LOCAL_GPG_MOUNT_DIR = "/var/run/dumpling-packed-backup-gpg" LOCAL_GPG_PATH = f"{LOCAL_GPG_MOUNT_DIR}/public-key" JOB_INSTANCE_NAME_ENV = "DUMPLING_PACKED_BACKUP_JOB_INSTANCE_NAME" JOB_INSTANCE_LABEL = "app.kubernetes.io/instance" WORKFLOW_NAME_LABEL = "cse.tidbcloud.com/workflow-name" STAGE_LABEL = "cse.tidbcloud.com/stage" JOB_NAME_LABEL = "cse.tidbcloud.com/job-name" RESOURCE_ROLE_LABEL = "cse.tidbcloud.com/resource-role" REQUIRED_ENV_NAMES = ( "DFS_PREFIX", "DFS_S3_BUCKET", "DFS_S3_REGION", "CSE_MASTER_KEY_ID", "CSE_MASTER_KEY_CIPHER_TEXT", "CSE_MASTER_KEY_VENDOR", ) class GeneratorError(Exception): pass class KubectlFailed(Exception): def __init__(self, returncode): self.returncode = returncode class Parser(argparse.ArgumentParser): def error(self, message): raise GeneratorError(message) def non_empty(value): if not value: raise argparse.ArgumentTypeError("must not be empty") return value def non_negative_integer(value): try: parsed = int(value) except ValueError as exc: raise argparse.ArgumentTypeError("must be a non-negative integer") from exc if parsed < 0: raise argparse.ArgumentTypeError("must be a non-negative integer") return parsed def build_parser(): parser = Parser( usage="%(prog)s [OPTIONS] -- STARTER-ARGS...", description=( "Generate Dumpling packed-backup export Job YAML. Output is written " "to stdout and is never applied to the cluster." ), epilog=( "Arguments after -- are passed unchanged to " "dumpling_packed_backup_init.py. A second -- passes the remaining " "arguments to dumpling. Batch entries may be S3 prefix strings or " "objects containing pack_backup_prefix and optional task_name. " "--local-gpg embeds and mounts a local OpenPGP public-key file.\n\n" "Examples:\n" " generate_dumpling_packed_backup_export.py --node-group=adhoc -- " "--s3-prefix s3://bucket/prefix/packed_backup/native/backup.meta/" "batch/meta --out-prefix s3://bucket/export --parallel 4 -- " "--threads 2\n\n" " list-completed-pack-backup-offline-tasks.sh | " "generate_dumpling_packed_backup_export.py --batch-task-json=- -- " "--out-prefix s3://bucket/export --parallel 4 -- --threads 2" ), formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False, ) def add(*flags, **kwargs): parser.add_argument(*flags, type=non_empty, **kwargs) add("--name", default=None, help="Job group name (default: generated from time)") add("--workflow-name", default=None, help="stable workflow RUN_ID label") add("--batch-task-json", default=None, metavar="PATH", help="task JSON; use - for stdin") parser.add_argument( "--output", choices=("yaml", "json"), default="yaml", help="output format (default: yaml)", ) add( "--local-gpg", default=None, metavar="PATH", help="local OpenPGP public-key file to embed and mount", ) add( "--runner-image", default=DEFAULT_RUNNER_IMAGE, metavar="IMAGE", help="Lambda-compatible Python 3.13 image containing Boto3", ) add( "--dumpling-image", default=DEFAULT_DUMPLING_IMAGE, metavar="IMAGE", help="image containing /dumpling", ) add( "--cse-ctl-image", default=DEFAULT_CSE_CTL_IMAGE, metavar="IMAGE", help="image containing /cse-ctl", ) add( "--node-group", default="general", metavar="GROUP", help="node group label and toleration (default: general)", ) parser.add_argument( "--max-retries", type=non_negative_integer, default=DEFAULT_MAX_RETRIES, metavar="COUNT", help=f"maximum Job retries (default: {DEFAULT_MAX_RETRIES})", ) add( "--cpu.request", dest="cpu_request", default=None, metavar="QUANTITY", help="CPU request", ) add( "--cpu.limit", dest="cpu_limit", default=None, metavar="QUANTITY", help="CPU limit (defaults to request)", ) add( "--mem.request", "--memory.request", dest="mem_request", default=None, metavar="QUANTITY", help="memory request", ) add( "--mem.limit", "--memory.limit", dest="mem_limit", default=None, metavar="QUANTITY", help="memory limit (defaults to request)", ) return parser def split_arguments(argv): argv = list(argv) if "--" not in argv: return argv, [] separator = argv.index("--") return argv[:separator], argv[separator + 1 :] def reject_duplicate_generator_options(arguments): groups = ( ("--name",), ("--workflow-name",), ("--batch-task-json",), ("--output",), ("--local-gpg",), ("--runner-image",), ("--dumpling-image",), ("--cse-ctl-image",), ("--node-group",), ("--max-retries",), ("--cpu.request",), ("--cpu.limit",), ("--mem.request", "--memory.request"), ("--mem.limit", "--memory.limit"), ) option_groups = { option: group[0] for group in groups for option in group } seen = set() for argument in arguments: option = argument.split("=", 1)[0] canonical = option_groups.get(option) if canonical is None: continue if canonical in seen: raise GeneratorError( f"generator option {canonical!r} must not be repeated" ) seen.add(canonical) def validate_options(args, starter_args): for name in ( "name", "workflow_name", "runner_image", "dumpling_image", "cse_ctl_image", "node_group", "cpu_request", "cpu_limit", "mem_request", "mem_limit", ): value = getattr(args, name) if value and value.startswith("-"): raise GeneratorError(f"{name.replace('_', '.')} must not start with '-'") node_group_pattern = r"[A-Za-z0-9](?:[-A-Za-z0-9_.]*[A-Za-z0-9])?" if len(args.node_group) > 63 or re.fullmatch(node_group_pattern, args.node_group) is None: raise GeneratorError( f"invalid Kubernetes node-group label value '{args.node_group}'" ) if not starter_args: raise GeneratorError("missing starter arguments after --") before_dumpling = starter_args[:] if "--" in before_dumpling: before_dumpling = before_dumpling[: before_dumpling.index("--")] has_s3_prefix = any( value == "--s3-prefix" or value.startswith("--s3-prefix=") for value in before_dumpling ) has_out_prefix = any( value == "--out-prefix" or value.startswith("--out-prefix=") for value in before_dumpling ) has_gpg = any( value == "--gpg" or value.startswith("--gpg=") for value in before_dumpling ) if not has_out_prefix: raise GeneratorError("starter arguments must include --out-prefix") if args.batch_task_json and has_s3_prefix: raise GeneratorError( "do not pass --s3-prefix with --batch-task-json; it is injected per task" ) if not args.batch_task_json and not has_s3_prefix: raise GeneratorError("starter arguments must include --s3-prefix") if args.local_gpg and has_gpg: raise GeneratorError("--local-gpg must not be combined with starter --gpg") if args.cpu_limit and not args.cpu_request: raise GeneratorError("--cpu.limit requires --cpu.request") if args.mem_limit and not args.mem_request: raise GeneratorError("--mem.limit requires --mem.request") if not args.name: args.name = f"dumpling-packed-export-{datetime.now():%Y%m%d-%H%M%S}" job_name_pattern = r"[a-z0-9](?:[-a-z0-9]*[a-z0-9])?" if len(args.name) > 63 or re.fullmatch(job_name_pattern, args.name) is None: raise GeneratorError(f"invalid Kubernetes Job group name '{args.name}'") if args.workflow_name and ( len(args.workflow_name) > 63 or re.fullmatch(job_name_pattern, args.workflow_name) is None ): raise GeneratorError( f"invalid workflow name '{args.workflow_name}'" ) def parse_json(text, description): try: return json.loads(text) except json.JSONDecodeError as exc: raise GeneratorError( f"invalid {description} JSON at line {exc.lineno} column {exc.colno}: {exc.msg}" ) from exc def read_local_gpg(path): local_path = Path(path) if not local_path.is_file(): raise GeneratorError(f"--local-gpg is not a regular file: '{path}'") try: return local_path.read_bytes() except OSError as exc: raise GeneratorError(f"cannot read --local-gpg file '{path}'") from exc def load_batch_tasks(path): if path == "-": payload = parse_json(sys.stdin.read(), "batch task") else: try: payload = parse_json(Path(path).read_text(encoding="utf-8"), "batch task") except OSError as exc: raise GeneratorError(f"cannot read --batch-task-json file '{path}'") from exc if not isinstance(payload, list): raise GeneratorError("batch task JSON must be an array") if not payload: raise GeneratorError("batch task JSON must not be empty") tasks = [] for item in payload: if isinstance(item, str): task_name, prefix = None, item elif isinstance(item, dict): task_name = item.get("task_name") if not isinstance(task_name, str) or not task_name: task_name = None prefix = item.get("pack_backup_prefix") else: raise GeneratorError("each batch task must be an S3 prefix string or an object") if not isinstance(prefix, str): raise GeneratorError("each pack_backup_prefix must be a string") if re.fullmatch(r"s3://[^/]+/.+/", prefix) is None: raise GeneratorError( "each pack_backup_prefix must be an s3:// batch prefix ending in /" ) tasks.append({"task_name": task_name, "pack_backup_prefix": prefix}) return tasks def kubectl_json(kubectl, *arguments): result = subprocess.run( [kubectl, "--request-timeout=30s", *arguments], stdout=subprocess.PIPE, text=True, encoding="utf-8", check=False, ) if result.returncode: raise KubectlFailed(result.returncode) return parse_json(result.stdout, "kubectl response") def source_environment(kubectl): stateful_set = kubectl_json( kubectl, "-n", NAMESPACE, "get", "statefulset", SOURCE_STS, "-o", "json" ) if not isinstance(stateful_set, dict): raise GeneratorError("kubectl returned a non-object StatefulSet") try: containers = stateful_set["spec"]["template"]["spec"].get("containers", []) except (KeyError, TypeError, AttributeError) as exc: raise GeneratorError("invalid source StatefulSet pod template") from exc containers = [] if containers is None else containers if not isinstance(containers, list): raise GeneratorError("source StatefulSet containers must be an array") container = next( (item for item in containers if isinstance(item, dict) and item.get("name") == SOURCE_CONTAINER), None, ) if container is None: raise GeneratorError(f"container {SOURCE_CONTAINER} not found in source StatefulSet") env = container.get("env", []) env = [] if env is None else env if not isinstance(env, list): raise GeneratorError("source container env must be an array") selected = [] for name in REQUIRED_ENV_NAMES: matches = [item for item in env if isinstance(item, dict) and item.get("name") == name] if len(matches) != 1: raise GeneratorError(f"expected exactly one {name} entry in the source container") selected.append(matches[0]) return selected def build_resources(args): requests, limits = {}, {} if args.cpu_request: requests["cpu"] = args.cpu_request limits["cpu"] = args.cpu_limit or args.cpu_request if args.mem_request: requests["memory"] = args.mem_request limits["memory"] = args.mem_limit or args.mem_request return {"requests": requests, "limits": limits} if requests else {} def build_entries(name, starter_args, tasks): if tasks is None: return [{ "job_name": name, "task_name": None, "pack_backup_prefix": None, "args": list(starter_args), }] count = len(tasks) return [{ "job_name": f"{name}-task-{index}-of-{count}", "task_name": task["task_name"], "pack_backup_prefix": task["pack_backup_prefix"], "args": ["--s3-prefix", task["pack_backup_prefix"] + "meta", *starter_args], } for index, task in enumerate(tasks)] def build_job(entry, labels, args, env, resources, config_map_name, has_local_gpg): annotations = {} if entry["pack_backup_prefix"] is not None: annotations["cse.tidbcloud.com/source-pack-backup-prefix"] = entry["pack_backup_prefix"] if entry["task_name"] is not None: annotations["cse.tidbcloud.com/source-task"] = entry["task_name"] job_labels = {**labels, JOB_NAME_LABEL: entry["job_name"]} metadata = { "name": entry["job_name"], "namespace": NAMESPACE, "labels": job_labels, } if annotations: metadata["annotations"] = annotations container = { "name": "offline-dumpling", "image": args.runner_image, "imagePullPolicy": "Always", "command": [RUNNER_PYTHON, STARTER], "args": entry["args"], "env": [ *env, { "name": JOB_INSTANCE_NAME_ENV, "value": labels[JOB_INSTANCE_LABEL], }, ], "volumeMounts": [ {"name": "dumpling", "mountPath": "/dumpling", "subPath": "dumpling", "readOnly": True}, {"name": "cse-ctl", "mountPath": "/cse-ctl", "subPath": "cse-ctl", "readOnly": True}, { "name": "init-script", "mountPath": STARTER, "subPath": "dumpling_packed_backup_init.py", "readOnly": True, }, ], } if has_local_gpg: container["volumeMounts"].append({ "name": "local-gpg", "mountPath": LOCAL_GPG_MOUNT_DIR, "readOnly": True, }) if resources: container["resources"] = resources pod_spec = { "restartPolicy": "Never", "serviceAccountName": SERVICE_ACCOUNT, "nodeSelector": { "kubernetes.io/arch": "arm64", "serverless.tidbcloud.com/node": args.node_group, }, "tolerations": [{ "key": "serverless.tidbcloud.com/node", "operator": "Equal", "value": args.node_group, "effect": "NoSchedule", }], "initContainers": [ { "name": "stage-dumpling", "image": args.dumpling_image, "imagePullPolicy": "Always", "command": ["/bin/sh", "-ec"], "args": ["cp /dumpling /stage/dumpling && chmod 0755 /stage/dumpling"], "volumeMounts": [{"name": "dumpling", "mountPath": "/stage"}], }, { "name": "stage-cse-ctl", "image": args.cse_ctl_image, "imagePullPolicy": "Always", "command": ["/bin/sh", "-ec"], "args": ["cp /cse-ctl /stage/cse-ctl && chmod 0755 /stage/cse-ctl"], "volumeMounts": [{"name": "cse-ctl", "mountPath": "/stage"}], }, ], "containers": [container], "volumes": [ {"name": "dumpling", "emptyDir": {}}, {"name": "cse-ctl", "emptyDir": {}}, { "name": "init-script", "configMap": {"name": config_map_name, "defaultMode": 493}, }, ], } if has_local_gpg: pod_spec["volumes"].append({ "name": "local-gpg", "configMap": { "name": config_map_name, "defaultMode": 292, "items": [{"key": LOCAL_GPG_CONFIG_KEY, "path": "public-key"}], }, }) job_spec = { "backoffLimit": args.max_retries, "ttlSecondsAfterFinished": DEFAULT_TTL_SECONDS_AFTER_FINISHED, "template": {"metadata": {"labels": job_labels}, "spec": pod_spec}, } return { "apiVersion": "batch/v1", "kind": "Job", "metadata": metadata, "spec": job_spec, } def build_manifest(args, starter_args, tasks, env, init_script, local_gpg): labels = { "app.kubernetes.io/name": "dumpling-packed-backup-export", JOB_INSTANCE_LABEL: args.name, STAGE_LABEL: "export", } if args.workflow_name: labels[WORKFLOW_NAME_LABEL] = args.workflow_name config_map_name = f"{args.name}-init" config_map = { "apiVersion": "v1", "kind": "ConfigMap", "metadata": { "name": config_map_name, "namespace": NAMESPACE, "labels": { **labels, RESOURCE_ROLE_LABEL: "export-init", }, }, "data": {"dumpling_packed_backup_init.py": init_script}, } if local_gpg is not None: config_map["binaryData"] = { LOCAL_GPG_CONFIG_KEY: base64.b64encode(local_gpg).decode("ascii") } resources = build_resources(args) jobs = [ build_job( entry, labels, args, env, resources, config_map_name, local_gpg is not None, ) for entry in build_entries(args.name, starter_args, tasks) ] return {"apiVersion": "v1", "kind": "List", "items": [config_map, *jobs]} def main(argv=None): generator_args, starter_args = split_arguments(sys.argv[1:] if argv is None else argv) try: reject_duplicate_generator_options(generator_args) args = build_parser().parse_args(generator_args) validate_options(args, starter_args) kubectl = shutil.which("kubectl") if kubectl is None: raise GeneratorError("kubectl is required") try: init_script = INIT_SCRIPT.read_text(encoding="utf-8") except OSError as exc: raise GeneratorError(f"cannot read init script '{INIT_SCRIPT}'") from exc local_gpg = read_local_gpg(args.local_gpg) if args.local_gpg else None if local_gpg is not None: starter_args = ["--gpg", LOCAL_GPG_PATH, *starter_args] tasks = load_batch_tasks(args.batch_task_json) if args.batch_task_json else None if tasks is not None: longest_name = f"{args.name}-task-{len(tasks) - 1}-of-{len(tasks)}" if len(longest_name) > 63: raise GeneratorError(f"batch Job name '{longest_name}' exceeds 63 characters") manifest = build_manifest( args, starter_args, tasks, source_environment(kubectl), init_script, local_gpg, ) manifest_json = json.dumps( manifest, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ) if args.output == "json": print(manifest_json) return 0 result = subprocess.run( [kubectl, "create", "--dry-run=client", "--validate=false", "-f", "-", "-o", "yaml"], input=manifest_json, text=True, encoding="utf-8", check=False, ) return result.returncode except GeneratorError as exc: print(f"error: {exc}", file=sys.stderr) return 1 except KubectlFailed as exc: return exc.returncode except KeyboardInterrupt: return 130 if __name__ == "__main__": sys.exit(main())