#!/usr/bin/env python3 """Own the D1 scratch-database lifecycle and its durable recovery proof.""" from __future__ import annotations import argparse import json import os import re import secrets import shlex import shutil import subprocess from pathlib import Path DATABASE_RE = re.compile(r"^nexuscrm_d1_[0-9]+_[0-9]+$") TOKEN_RE = re.compile(r"^[0-9a-f]{64}$") REQUIRED_CONNECTION = {"TIDB_HOST", "TIDB_PORT", "TIDB_USER", "TIDB_PASSWORD"} CONNECTION_INPUTS = REQUIRED_CONNECTION | {"TIDB_TLS"} LOCAL_TIDB_HOSTS = frozenset({"127.0.0.1", "localhost", "host.docker.internal"}) def fail(message: str) -> None: raise SystemExit(message) def decode_shell_env_value(value: str) -> str: """Decode one shell-quoted scalar without evaluating shell syntax.""" try: decoded = shlex.split(value, posix=True) except ValueError as exc: raise ValueError(f"invalid shell-quoted environment value: {exc}") from exc if len(decoded) != 1: raise ValueError("environment value must encode exactly one shell word") return decoded[0] def load_env(path: Path, *, allow_insecure_local: bool = False) -> dict[str, str]: values: dict[str, str] = {} for raw_line in path.read_text(encoding="utf-8").splitlines(): line = raw_line.strip() if not line or line.startswith("#") or "=" not in line: continue if line.startswith("export "): line = line[len("export "):].strip() key, value = line.split("=", 1) if key not in CONNECTION_INPUTS: continue if key in values: fail(f"duplicate scratch connection input: {key}") if key == "TIDB_TLS" and not value.strip(): values[key] = "" else: try: values[key] = decode_shell_env_value(value) except ValueError as exc: fail(str(exc)) missing = sorted(REQUIRED_CONNECTION - values.keys()) if missing: fail(f"missing scratch connection inputs: {', '.join(missing)}") non_secret_inputs = REQUIRED_CONNECTION - {"TIDB_PASSWORD"} if any(not values[key] for key in non_secret_inputs): fail("scratch connection endpoint and user inputs must be non-empty") if not values.get("TIDB_TLS", "").strip(): values["TIDB_TLS"] = "true" tls_policy = values["TIDB_TLS"].strip().lower() if tls_policy not in {"true", "false"}: fail("TIDB_TLS must be exactly true or false") values["TIDB_TLS"] = tls_policy tls_enabled = tls_policy == "true" if not tls_enabled and ( not allow_insecure_local or values["TIDB_HOST"] not in LOCAL_TIDB_HOSTS ): fail( "D1 test2 requires the TLS-backed TiDB contract; " "insecure TiDB is allowed only for explicit local validation" ) if not values["TIDB_PASSWORD"] and not ( allow_insecure_local and values["TIDB_HOST"] in LOCAL_TIDB_HOSTS ): fail("an empty TiDB password is allowed only for explicit local validation") if not values["TIDB_PORT"].isdigit(): fail("invalid scratch connection TiDB port") return values class Scratch: def __init__( self, run_root: Path, marker: str, database: str, recovery_path: Path, *, allow_insecure_local: bool = False, ) -> None: self.run_root = run_root.resolve() self.marker = marker self.database = database self.recovery_path = recovery_path.resolve() self.allow_insecure_local = allow_insecure_local if not DATABASE_RE.fullmatch(database): fail("invalid D1 scratch database name") if marker != database.replace("nexuscrm_d1_", "nexuscrm-d1-").replace("_", "-"): fail("D1 scratch database does not belong to the run marker") if self.recovery_path.parent != self.run_root.parent: fail("D1 scratch recovery record must be outside and adjacent to the run root") if self.recovery_path.name != f"{marker}-scratch-recovery.json": fail("invalid D1 scratch recovery record name") def prepare(self, source_env: Path) -> None: if self.recovery_path.exists(): fail("D1 scratch recovery record already exists") values = load_env(source_env, allow_insecure_local=self.allow_insecure_local) tls_enabled = values["TIDB_TLS"] == "true" payload = { "schema": "nexuscrm.d1-scratch-database-recovery/v1", "run": {"container_prefix": self.marker, "root": str(self.run_root)}, "tidb": { "host": values["TIDB_HOST"], "port": values["TIDB_PORT"], "tls": tls_enabled, "database": self.database, }, "guard_token": secrets.token_hex(32), } temporary = self.recovery_path.with_name( f".{self.recovery_path.name}.{os.getpid()}.{secrets.token_hex(8)}.tmp" ) try: descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, "w", encoding="utf-8") as handle: json.dump(payload, handle, sort_keys=True, separators=(",", ":")) handle.write("\n") handle.flush() os.fsync(handle.fileno()) temporary.chmod(0o400) # Publish without a check-then-replace race: an existing recovery # owner must never be overwritten by another attempt. os.link(temporary, self.recovery_path) finally: temporary.unlink(missing_ok=True) def load(self) -> tuple[dict[str, object], str]: try: payload = json.loads(self.recovery_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: fail(f"invalid D1 scratch recovery record: {exc}") if payload.get("schema") != "nexuscrm.d1-scratch-database-recovery/v1": fail("invalid D1 scratch recovery schema") if payload.get("run") != {"container_prefix": self.marker, "root": str(self.run_root)}: fail("D1 scratch recovery record belongs to another run") tidb = payload.get("tidb") if ( not isinstance(tidb, dict) or tidb.get("database") != self.database or not isinstance(tidb.get("tls"), bool) ): fail("D1 scratch recovery record has the wrong database or TLS policy") if tidb["tls"] is not True and ( not self.allow_insecure_local or tidb.get("host") not in LOCAL_TIDB_HOSTS ): fail("insecure D1 scratch recovery is valid only for explicit local validation") for key in ("host", "port"): if not isinstance(tidb.get(key), str) or not tidb[key]: fail(f"invalid D1 scratch recovery field: {key}") token = payload.get("guard_token") if not str(tidb["port"]).isdigit() or not isinstance(token, str) or not TOKEN_RE.fullmatch(token): fail("invalid D1 scratch recovery port or guard") return tidb, token def connection(self, source_env: Path) -> dict[str, object]: tidb, _ = self.load() credentials = load_env( source_env, allow_insecure_local=self.allow_insecure_local ) credential_tls = credentials["TIDB_TLS"] == "true" if ( credentials["TIDB_HOST"] != tidb["host"] or credentials["TIDB_PORT"] != tidb["port"] or credential_tls is not tidb["tls"] ): fail("current TiDB credential source does not match the D1 recovery endpoint") return { "host": tidb["host"], "port": tidb["port"], "user": credentials["TIDB_USER"], "password": credentials["TIDB_PASSWORD"], "tls": credential_tls, "database": self.database, } def mysql(self, sql: str, source_env: Path) -> str: tidb = self.connection(source_env) client = shutil.which("mariadb") or shutil.which("mysql") if not client: fail("mysql or mariadb client is required") help_result = subprocess.run([client, "--help"], text=True, capture_output=True, check=False) if tidb["tls"]: tls_arg = ( "--ssl-mode=REQUIRED" if "--ssl-mode" in help_result.stdout + help_result.stderr else "--ssl" ) else: tls_arg = ( "--ssl-mode=DISABLED" if "--ssl-mode" in help_result.stdout + help_result.stderr else "--skip-ssl" ) result = subprocess.run( [client, tls_arg, f"-h{tidb['host']}", f"-P{tidb['port']}", f"-u{tidb['user']}", "-N", "-B"], input=sql + "\n", text=True, capture_output=True, env=os.environ | {"MYSQL_PWD": str(tidb["password"])}, check=False, ) if result.returncode != 0: fail(f"D1 scratch database command failed: {result.stderr.strip()}") return result.stdout.strip() def exists(self, source_env: Path) -> bool: result = self.mysql( f"SELECT COUNT(*) FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME='{self.database}';", source_env, ) if result not in {"0", "1"}: fail("unexpected D1 scratch database existence result") return result == "1" def verify(self, source_env: Path) -> None: _, token = self.load() stored = self.mysql( f"SELECT `GuardToken` FROM `{self.database}`.`_D1_ScratchGuard` LIMIT 1;", source_env, ) if stored != token: fail("refuse D1 scratch database without matching resource-local guard") def create(self, source_env: Path) -> None: _, token = self.load() if self.exists(source_env): fail("refuse to reuse an existing D1 scratch database") self.mysql( f"CREATE DATABASE `{self.database}`; " f"CREATE TABLE `{self.database}`.`_D1_ScratchGuard` (`GuardToken` VARCHAR(64) NOT NULL); " f"INSERT INTO `{self.database}`.`_D1_ScratchGuard` (`GuardToken`) VALUES ('{token}');", source_env, ) self.verify(source_env) def drop(self, source_env: Path) -> None: if not self.exists(source_env): return self.verify(source_env) self.mysql(f"DROP DATABASE `{self.database}`;", source_env) if self.exists(source_env): fail("D1 scratch database still exists after DROP") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("command", choices=("prepare", "validate", "project", "create", "verify", "drop")) parser.add_argument("--run-root", required=True) parser.add_argument("--marker", required=True) parser.add_argument("--database", required=True) parser.add_argument("--recovery-file", required=True) parser.add_argument("--source-env", required=True) parser.add_argument( "--allow-insecure-local", action="store_true", help="allow a non-TLS loopback/host.docker.internal TiDB only for local validation", ) return parser.parse_args() def main() -> int: args = parse_args() scratch = Scratch( Path(args.run_root), args.marker, args.database, Path(args.recovery_file), allow_insecure_local=args.allow_insecure_local, ) if args.command == "prepare": scratch.prepare(Path(args.source_env)) elif args.command == "validate": scratch.load() scratch.connection(Path(args.source_env)) elif args.command == "project": print(json.dumps(scratch.connection(Path(args.source_env)), sort_keys=True, separators=(",", ":"))) elif args.command == "create": scratch.create(Path(args.source_env)) elif args.command == "verify": scratch.verify(Path(args.source_env)) else: scratch.drop(Path(args.source_env)) return 0 if __name__ == "__main__": raise SystemExit(main())