#!/usr/bin/env python3 """Generate the ConfigMap and Job for one packed-backup watermark stage.""" import argparse import json import re import sys from datetime import datetime from pathlib import Path SCRIPT_DIR = Path(__file__).resolve().parent WORKER_SCRIPT = SCRIPT_DIR / "packed_backup_watermark.py" NAMESPACE = "tidb-serverless" SERVICE_ACCOUNT = "backup-restore" CONTAINER_NAME = "packed-backup-watermark" INPUT_MOUNT_DIR = "/var/run/packed-backup-watermark/input" MYSQL_MOUNT_DIR = "/var/run/packed-backup-watermark/mysql" TASKS_PATH = f"{INPUT_MOUNT_DIR}/pack-tasks.json" MYSQL_DEFAULTS_PATH = f"{MYSQL_MOUNT_DIR}/mysql.cnf" WORKER_PATH = f"{INPUT_MOUNT_DIR}/packed_backup_watermark.py" BACKUP_TIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ" DEFAULT_TTL_SECONDS_AFTER_FINISHED = 30 * 24 * 60 * 60 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" WORKER_OWNED_OPTIONS = { "--watermark-instance", "--export-instance", "--tasks-json", "--out-prefix", "--backup-time", "--mysql-defaults-file", "--aws-region", } class GeneratorError(Exception): pass 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 kubernetes_name(value): if ( len(value) > 63 or re.fullmatch(r"[a-z0-9]([-a-z0-9]*[a-z0-9])?", value) is None ): raise argparse.ArgumentTypeError( "must be a Kubernetes DNS label of at most 63 characters" ) return value def volume_key(value): if len(value) > 253 or re.fullmatch(r"[A-Za-z0-9._-]+", value) is None: raise argparse.ArgumentTypeError( "must be a valid Secret or ConfigMap data key" ) return value def s3_prefix(value): if re.fullmatch(r"s3://[^/\s]+/.+", value) is None: raise argparse.ArgumentTypeError( "must be an s3:/// URI" ) return value def backup_time(value): try: datetime.strptime(value, BACKUP_TIME_FORMAT) except ValueError as exc: raise argparse.ArgumentTypeError( f"must use {BACKUP_TIME_FORMAT}" ) from exc return value def build_parser(): parser = Parser( description=( "Generate a Kubernetes JSON List containing the input ConfigMap " "and Job for one packed-backup watermark stage." ), formatter_class=argparse.ArgumentDefaultsHelpFormatter, allow_abbrev=False, ) parser.add_argument("--name", type=kubernetes_name, required=True) parser.add_argument( "--workflow-name", type=kubernetes_name, default=None, metavar="NAME", help="stable workflow RUN_ID label", ) parser.add_argument( "--tasks-json", type=non_empty, required=True, metavar="PATH", help="Pack task JSON; use - for stdin", ) parser.add_argument( "--mysql-secret", type=kubernetes_name, required=True, metavar="NAME", ) parser.add_argument( "--mysql-secret-key", type=volume_key, default="mysql.cnf", metavar="KEY", ) parser.add_argument( "--out-prefix", type=s3_prefix, required=True, metavar="S3_URI", ) parser.add_argument( "--export-instance", type=kubernetes_name, required=True, metavar="NAME", ) parser.add_argument( "--backup-time", type=backup_time, required=True, metavar="UTC_TIME", ) parser.add_argument( "--aws-region", type=non_empty, required=True, metavar="REGION", ) parser.add_argument( "--runner-image", type=non_empty, required=True, metavar="IMAGE", help="image containing Python, Boto3, and PyMySQL", ) parser.add_argument( "--node-group", type=non_empty, default="general", metavar="GROUP", ) parser.add_argument( "--cpu.request", dest="cpu_request", type=non_empty, default=None, metavar="QUANTITY", ) parser.add_argument( "--memory.request", dest="memory_request", type=non_empty, default=None, metavar="QUANTITY", ) return parser def split_arguments(argv): arguments = list(argv) if "--" not in arguments: return arguments, [] separator = arguments.index("--") return arguments[:separator], arguments[separator + 1 :] def option_name(argument): if not argument.startswith("-") or argument == "-": return None return argument.split("=", 1)[0] def reject_duplicate_generator_options(arguments): options = { "--name", "--workflow-name", "--tasks-json", "--mysql-secret", "--mysql-secret-key", "--out-prefix", "--export-instance", "--backup-time", "--aws-region", "--runner-image", "--node-group", "--cpu.request", "--memory.request", } seen = set() for argument in arguments: option = option_name(argument) if option not in options: continue if option in seen: raise GeneratorError( f"generator option {option!r} must not be repeated" ) seen.add(option) 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} " f"column {exc.colno}: {exc.msg}" ) from exc def load_tasks(path): if path == "-": payload = parse_json(sys.stdin.read(), "Pack task") else: try: text = Path(path).read_text(encoding="utf-8") except OSError as exc: raise GeneratorError(f"cannot read Pack task file {path!r}") from exc payload = parse_json(text, "Pack task") if not isinstance(payload, list) or not payload: raise GeneratorError("Pack task JSON must be a non-empty array") if not all(isinstance(task, dict) for task in payload): raise GeneratorError("Pack task JSON contains a non-object task") return payload def read_worker(): try: return WORKER_SCRIPT.read_text(encoding="utf-8") except OSError as exc: raise GeneratorError( f"cannot read watermark worker {str(WORKER_SCRIPT)!r}" ) from exc def validate_options(args, worker_args): input_config_map = f"{args.name}-input" if len(input_config_map) > 63: raise GeneratorError( f"watermark input ConfigMap name {input_config_map!r} " "exceeds 63 characters" ) 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!r}" ) for index, argument in enumerate(worker_args): if not argument: raise GeneratorError( f"worker argument {index} must not be empty" ) if argument == "--": raise GeneratorError( "worker arguments must not introduce another argv boundary" ) option = option_name(argument) if option in WORKER_OWNED_OPTIONS: raise GeneratorError( f"worker argument {index} uses generator-owned option " f"{option!r}" ) def build_resources(args): requests = {} if args.cpu_request: requests["cpu"] = args.cpu_request if args.memory_request: requests["memory"] = args.memory_request if not requests: return {} return {"requests": requests, "limits": dict(requests)} def build_manifest(args, tasks, worker, worker_args): input_config_map = f"{args.name}-input" labels = { "app.kubernetes.io/name": "packed-backup-watermark", "app.kubernetes.io/instance": args.name, STAGE_LABEL: "watermark", } if args.workflow_name: labels[WORKFLOW_NAME_LABEL] = args.workflow_name job_labels = {**labels, JOB_NAME_LABEL: args.name} container = { "name": CONTAINER_NAME, "image": args.runner_image, "imagePullPolicy": "Always", "terminationMessagePath": "/dev/termination-log", "terminationMessagePolicy": "File", "command": ["python3", WORKER_PATH], "args": [ "--watermark-instance", args.name, "--export-instance", args.export_instance, "--tasks-json", TASKS_PATH, "--out-prefix", args.out_prefix, "--backup-time", args.backup_time, "--mysql-defaults-file", MYSQL_DEFAULTS_PATH, "--aws-region", args.aws_region, *worker_args, ], "volumeMounts": [ { "name": "watermark-input", "mountPath": INPUT_MOUNT_DIR, "readOnly": True, }, { "name": "mysql", "mountPath": MYSQL_MOUNT_DIR, "readOnly": True, }, ], } resources = build_resources(args) if resources: container["resources"] = resources config_map = { "apiVersion": "v1", "kind": "ConfigMap", "metadata": { "name": input_config_map, "namespace": NAMESPACE, "labels": { **labels, RESOURCE_ROLE_LABEL: "watermark-input", }, }, "data": { "packed_backup_watermark.py": worker, "pack-tasks.json": json.dumps( tasks, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ), }, } job = { "apiVersion": "batch/v1", "kind": "Job", "metadata": { "name": args.name, "namespace": NAMESPACE, "labels": job_labels, }, "spec": { "backoffLimit": 0, "ttlSecondsAfterFinished": DEFAULT_TTL_SECONDS_AFTER_FINISHED, "template": { "metadata": {"labels": job_labels}, "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", } ], "containers": [container], "volumes": [ { "name": "watermark-input", "configMap": {"name": input_config_map}, }, { "name": "mysql", "secret": { "secretName": args.mysql_secret, "defaultMode": 256, "items": [ { "key": args.mysql_secret_key, "path": "mysql.cnf", } ], }, }, ], }, }, }, } return { "apiVersion": "v1", "kind": "List", "items": [config_map, job], } def main(argv=None): try: generator_args, worker_args = split_arguments( sys.argv[1:] if argv is None else argv ) reject_duplicate_generator_options(generator_args) args = build_parser().parse_args(generator_args) validate_options(args, worker_args) manifest = build_manifest( args, load_tasks(args.tasks_json), read_worker(), worker_args, ) print( json.dumps( manifest, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ) ) return 0 except GeneratorError as exc: print(f"error: {exc}", file=sys.stderr) return 1 except KeyboardInterrupt: return 130 if __name__ == "__main__": sys.exit(main())