from __future__ import annotations import json import os import re import shutil import subprocess import tempfile import time from collections import defaultdict from datetime import UTC, datetime from pathlib import Path from typing import Any from urllib.parse import urlparse import boto3 import pyarrow as pa import pyarrow.parquet as pq from databend_driver import BlockingDatabendClient from google.cloud import bigquery, storage from .config import PipelineConfig from .gcs import download_text, list_objects, parse_gs_uri, upload_text from .models import table_to_manifest from .naming import safe_table_id from .source_inspect import ( bigquery_table_types, include_table, inspect_exact_tidb_tables, inspect_tables, table_selection_contract, tidb_connect, tidb_export_patterns, ) TERMINAL_EXPORT_STATES = {"SUCCEEDED", "FAILED", "CANCELED"} FAULT_DOMAIN_NAME = re.compile(r"^[a-z][a-z0-9_-]{0,62}$") PROVIDER_DIAGNOSTIC_FIELDS = ( "clusterId", "completeTime", "createTime", "displayName", "exportId", "name", "snapshotTime", "state", "updateTime", ) PROVIDER_REASON_FIELDS = ("errorMessage", "reason", "message", "detail", "error") PROVIDER_ERROR_CODE_FIELDS = ("errorCode", "code") class ExtractRunError(RuntimeError): def __init__(self, message: str, *, summary: dict[str, Any]) -> None: super().__init__(message) self.summary = summary def run_extract( config: PipelineConfig, *, run_date: str | None = None, run_id: str | None = None, tables: list[str] | None = None, ) -> dict[str, Any]: run_date = run_date or datetime.now(UTC).strftime("%Y-%m-%d") run_id = run_id or datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") selected = tables if selected is not None and not selected: raise ValueError("explicit table selection must not be empty") try: if config.source_type == "bigquery": extract_result = extract_bigquery(config, run_date=run_date, run_id=run_id, tables=selected) elif config.source_type == "tidbcloud": extract_result = extract_tidbcloud(config, run_date=run_date, run_id=run_id, tables=selected) elif config.source_type == "s3": extract_result = extract_s3_snapshot(config, run_date=run_date, run_id=run_id) elif config.source_type == "databend": extract_result = extract_databend_snapshot(config, run_date=run_date, run_id=run_id) elif config.source_type == "mysql_query": extract_result = extract_mysql_query_snapshot( config, run_date=run_date, run_id=run_id, tables=selected, ) else: raise ValueError(f"unsupported source type: {config.source_type}") manifest = build_manifest( config, run_date=run_date, run_id=run_id, extract_result=extract_result, selected=selected, ) except ExtractRunError as error: record_extract_failure(config, run_date=run_date, run_id=run_id, error=error) raise manifest_uri = config.manifest_gcs_uri(run_date, run_id) text = json.dumps(manifest, indent=2, sort_keys=True) + "\n" write_local_run_file(config, run_date, run_id, "manifest.json", text) upload_text(manifest_uri, text) return {"run_date": run_date, "run_id": run_id, "manifest_uri": manifest_uri, "manifest": manifest} def record_extract_failure( config: PipelineConfig, *, run_date: str, run_id: str, error: ExtractRunError, ) -> None: event = { "event_schema_version": 1, "event_type": "extract_failed", "status": "failed", "source_name": config.name, "run_date": run_date, "run_id": run_id, "created_at": datetime.now(UTC).isoformat(), **error.summary, } runtime_log_uri = config.runtime_log_gcs_uri(run_date, run_id) text = json.dumps(event, sort_keys=True) + "\n" try: write_local_run_file(config, run_date, run_id, "events.jsonl", text) upload_text(runtime_log_uri, text, content_type="application/jsonl") except Exception as recording_error: event["diagnostic_recording"] = { "status": "failed", "error_type": type(recording_error).__name__, } else: event["diagnostic_recording"] = { "status": "recorded", "runtime_log_uri": runtime_log_uri, } error.summary = event def extract_bigquery( config: PipelineConfig, *, run_date: str, run_id: str, tables: list[str] | None, ) -> dict[str, Any]: client = bigquery.Client(project=config.target_project, location=config.location) table_names = tables or [table.name for table in inspect_tables(config)] table_types = bigquery_table_types(client, config, table_names) if table_names else {} jobs = [] for table_name in table_names: destination_uri = f"{config.table_gcs_uri(table_name, run_date, run_id)}{table_name}-*.parquet" if table_types.get(table_name) == "VIEW": job_id = extract_bigquery_view(client, config, table_name, destination_uri, run_id) jobs.append( { "table": table_name, "table_type": "VIEW", "job_id": job_id, "destination_uri": destination_uri, } ) else: source_table = f"{config.source_project}.{config.source_dataset}.{table_name}" job_config = bigquery.ExtractJobConfig( destination_format=bigquery.DestinationFormat.PARQUET, compression=config.parquet_compression, ) job = client.extract_table(source_table, destination_uri, job_config=job_config, location=config.location) job.result() jobs.append( { "table": table_name, "table_type": table_types.get(table_name, "BASE TABLE"), "job_id": job.job_id, "destination_uri": destination_uri, } ) return {"engine": "bigquery", "jobs": jobs, "gcs_uri": config.gcs_run_uri(run_date, run_id)} def extract_bigquery_view( client: bigquery.Client, config: PipelineConfig, table_name: str, destination_uri: str, run_id: str, ) -> str: stage_table = ( f"{config.target_project}.{config.control_dataset}." f"__extract_{safe_table_id(config.name)}_{safe_table_id(table_name)}_{safe_table_id(run_id)}" ) query_job = client.query( f"SELECT * FROM `{config.source_project}.{config.source_dataset}.{table_name}`", job_config=bigquery.QueryJobConfig( destination=stage_table, write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, ), location=config.location, ) query_job.result() try: job_config = bigquery.ExtractJobConfig( destination_format=bigquery.DestinationFormat.PARQUET, compression=config.parquet_compression, ) extract_job = client.extract_table( stage_table, destination_uri, job_config=job_config, location=config.location, ) extract_job.result() return extract_job.job_id finally: client.delete_table(stage_table, not_found_ok=True) def tidb_export_filter_bytes(patterns: list[str]) -> int: payload = {"table": {"patterns": patterns, "where": "", "partitions": []}} return len(json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode()) def chunk_tidb_export_patterns(patterns: list[str], *, max_filter_bytes: int) -> list[list[str]]: chunks: list[list[str]] = [] current: list[str] = [] for pattern in patterns: if tidb_export_filter_bytes([pattern]) > max_filter_bytes: raise ValueError(f"TiDB Cloud export filter pattern exceeds {max_filter_bytes} bytes: {pattern!r}") candidate = [*current, pattern] if tidb_export_filter_bytes(candidate) <= max_filter_bytes: current = candidate continue chunks.append(current) current = [pattern] if current: chunks.append(current) return chunks def plan_tidb_export_batches( patterns: list[str], *, fault_domains: list[dict[str, Any]], max_filter_bytes: int, require_all_fault_domains: bool = True, ) -> list[dict[str, Any]]: seen: set[str] = set() duplicates: set[str] = set() for pattern in patterns: if pattern in seen: duplicates.add(pattern) seen.add(pattern) if duplicates: raise ValueError(f"duplicate TiDB Cloud export patterns are not allowed: {sorted(duplicates)}") configured_domains = normalize_tidb_fault_domains(fault_domains) if configured_domains: grouped: dict[str, list[str]] = {str(domain["name"]): [] for domain in configured_domains} for pattern in patterns: table_name = pattern.split(".", 1)[1] if "." in pattern else pattern matches = [ str(domain["name"]) for domain in configured_domains if include_table( table_name, selected=list(domain["include"]), excluded=list(domain["exclude"]), ) ] if len(matches) != 1: match_description = ", ".join(matches) if matches else "none" raise ValueError( f"TiDB Cloud export pattern {pattern!r} must match exactly one fault domain; " f"matched {match_description}" ) grouped[matches[0]].append(pattern) empty_domains = [name for name, domain_patterns in grouped.items() if not domain_patterns] if empty_domains and require_all_fault_domains: raise ValueError(f"TiDB Cloud export fault domains matched no patterns: {empty_domains}") domain_groups = [ (str(domain["name"]), grouped[str(domain["name"])]) for domain in configured_domains if grouped[str(domain["name"])] ] else: domain_groups = [("default", list(patterns))] batches: list[dict[str, Any]] = [] domain_count = len(domain_groups) for domain_index, (domain_name, domain_patterns) in enumerate(domain_groups, start=1): chunks = chunk_tidb_export_patterns(domain_patterns, max_filter_bytes=max_filter_bytes) for chunk_index, chunk in enumerate(chunks, start=1): batches.append( { "fault_domain": domain_name, "fault_domain_index": domain_index, "fault_domain_count": domain_count, "chunk_index": chunk_index, "chunk_count": len(chunks), "filter_patterns": chunk, } ) for batch_index, batch in enumerate(batches, start=1): batch["batch_index"] = batch_index batch["batch_count"] = len(batches) return batches def normalize_tidb_fault_domains(fault_domains: list[dict[str, Any]]) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] names: set[str] = set() for index, raw_domain in enumerate(fault_domains, start=1): if not isinstance(raw_domain, dict): raise ValueError(f"TiDB Cloud export fault domain {index} must be a mapping") name = str(raw_domain.get("name", "")) if not FAULT_DOMAIN_NAME.fullmatch(name): raise ValueError( f"TiDB Cloud export fault domain name {name!r} must match {FAULT_DOMAIN_NAME.pattern}" ) if name in names: raise ValueError(f"duplicate TiDB Cloud export fault domain name: {name}") names.add(name) include = raw_domain.get("include") exclude = raw_domain.get("exclude", []) if not isinstance(include, list) or not include or not all(isinstance(item, str) and item for item in include): raise ValueError(f"TiDB Cloud export fault domain {name!r} must have a non-empty string include list") if not isinstance(exclude, list) or not all(isinstance(item, str) and item for item in exclude): raise ValueError(f"TiDB Cloud export fault domain {name!r} exclude must be a string list") normalized.append({"name": name, "include": list(include), "exclude": list(exclude)}) return normalized def extract_tidbcloud( config: PipelineConfig, *, run_date: str, run_id: str, tables: list[str] | None = None, ) -> dict[str, Any]: ensure_ticloud_profile(config) gcs_uri = config.gcs_run_uri(run_date, run_id) service_account_key = os.environ.get(config.tidb_gcs_service_account_key_env) if not service_account_key: raise RuntimeError(f"required environment variable is not set: {config.tidb_gcs_service_account_key_env}") patterns = tidb_export_patterns(config, selected=tables) batches = plan_tidb_export_batches( patterns, fault_domains=config.tidb_export_fault_domains, max_filter_bytes=config.tidb_filter_max_bytes, require_all_fault_domains=tables is None, ) if not batches: raise RuntimeError(f"no TiDB Cloud tables selected for source {config.name}") export_chunks = [] configured_fault_domains = bool(config.tidb_export_fault_domains) for batch in batches: chunk = list(batch["filter_patterns"]) chunk_index = int(batch["chunk_index"]) chunk_count = int(batch["chunk_count"]) batch_index = int(batch["batch_index"]) batch_count = int(batch["batch_count"]) fault_domain = str(batch["fault_domain"]) if configured_fault_domains: chunk_gcs_uri = f"{gcs_uri}fault_domain={fault_domain}/chunk={chunk_index:04d}/" description = ( f"TiDB Cloud CreateExport fault domain {fault_domain} chunk {chunk_index}/{chunk_count} " f"(batch {batch_index}/{batch_count})" ) else: chunk_gcs_uri = gcs_uri if batch_count == 1 else f"{gcs_uri}chunk={batch_index:04d}/" description = f"TiDB Cloud CreateExport chunk {batch_index}/{batch_count}" command = [ ticloud_bin(), "serverless", "export", "create", "--profile", config.tidb_cli_profile, "--cluster-id", config.tidb_cluster_id, "--target-type", "GCS", "--gcs.uri", chunk_gcs_uri, "--gcs.service-account-key", service_account_key, "--filter", ",".join(chunk), "--file-type", "PARQUET", "--parquet.compression", config.parquet_compression, "--force", ] max_attempts = int(config.data.get("export", {}).get("tidb_max_attempts", 1)) if not 1 <= max_attempts <= 3: raise ValueError("tidb_max_attempts must be between 1 and 3") base_gcs_uri = chunk_gcs_uri failed_attempts = [] for attempt in range(1, max_attempts + 1): chunk_gcs_uri = base_gcs_uri if max_attempts == 1 else f"{base_gcs_uri}attempt={attempt:04d}/" command[command.index("--gcs.uri") + 1] = chunk_gcs_uri completed = run_sensitive_command( command, description=description if max_attempts == 1 else f"{description} attempt {attempt}/{max_attempts}", sensitive_values=(service_account_key,), ) export_id = parse_export_id(completed.stdout) if not export_id: raise RuntimeError("TiDB Cloud export started but no export id was found") export = wait_for_ticloud_export(config, export_id) if export.get("state") != "FAILED" or attempt == max_attempts: break failed_attempts.append(sanitize_ticloud_export_failure(export, sensitive_values=(service_account_key,))) time.sleep(int(config.data.get("export", {}).get("poll_interval_seconds", 30))) if export.get("state") != "SUCCEEDED": diagnostics = sanitize_ticloud_export_failure( export, sensitive_values=(service_account_key,), ) provider_reason = diagnostics.get("provider_reason") if configured_fault_domains: chunk_label = ( f"fault domain {fault_domain} chunk {chunk_index}/{chunk_count} " f"(batch {batch_index}/{batch_count})" ) else: chunk_label = f"chunk {batch_index}/{batch_count}" reason_suffix = f"; provider reason: {provider_reason}" if provider_reason else "" raise ExtractRunError( f"TiDB Cloud export {chunk_label} ({export_id}) ended in state " f"{export.get('state')}{reason_suffix}", summary={ "engine": "tidbcloud", "failure_stage": "provider_terminal_state", "fault_domain": fault_domain, "fault_domain_index": int(batch["fault_domain_index"]), "fault_domain_count": int(batch["fault_domain_count"]), "chunk_index": chunk_index, "chunk_count": chunk_count, "batch_index": batch_index, "batch_count": batch_count, "filter_pattern_count": len(chunk), "filter_payload_bytes": tidb_export_filter_bytes(chunk), "gcs_uri": chunk_gcs_uri, **diagnostics, }, ) empty_tables = [] if config.data.get("export", {}).get("tidb_materialize_empty_tables", False): empty_tables = materialize_empty_tidb_tables(config, chunk_gcs_uri, chunk, export_id) export_chunks.append( { "chunk_index": chunk_index, "chunk_count": chunk_count, "batch_index": batch_index, "batch_count": batch_count, "fault_domain": fault_domain, "fault_domain_index": int(batch["fault_domain_index"]), "fault_domain_count": int(batch["fault_domain_count"]), "filter_patterns": chunk, "gcs_uri": chunk_gcs_uri, "export_id": export_id, "export": sanitize_ticloud_export(export), "empty_tables": empty_tables, "failed_attempts": failed_attempts, } ) result = { "engine": "tidbcloud", "selection": table_selection_contract( config, [pattern.split(".", 1)[1] for pattern in patterns], requested=tables, ), "export_count": len(export_chunks), "export_ids": [chunk["export_id"] for chunk in export_chunks], "export_chunks": export_chunks, "fault_domains": summarize_export_fault_domains(export_chunks), "tidb_filter_max_bytes": config.tidb_filter_max_bytes, "gcs_uri": gcs_uri, } if len(export_chunks) == 1: result["export_id"] = export_chunks[0]["export_id"] result["export"] = export_chunks[0]["export"] return result def tidb_empty_arrow_type(column: dict[str, Any]) -> pa.DataType: """Match Dumpling's Parquet logical types (including DATE/TIME as strings).""" kind = str(column["data_type"]).lower() column_type = str(column["column_type"]).lower() unsigned = "unsigned" in column_type if kind == "bigint": return pa.decimal128(20, 0) if unsigned else pa.int64() if kind in {"int", "integer"}: return pa.int64() if unsigned else pa.int32() if kind in {"tinyint", "smallint", "mediumint", "year"}: return pa.int32() if kind in {"float", "double"}: return pa.float64() if kind in {"datetime", "timestamp"}: return pa.timestamp("us") if kind == "decimal": match = re.fullmatch(r"decimal\((\d+),(\d+)\)(?: unsigned)?", column_type) if not match: raise ValueError("invalid decimal metadata for empty TiDB table") precision, scale = map(int, match.groups()) return pa.decimal128(precision, scale) if precision <= 38 else pa.string() if kind in {"binary", "varbinary", "blob", "tinyblob", "mediumblob", "longblob", "bit"}: return pa.binary() if kind in {"char", "varchar", "text", "tinytext", "mediumtext", "longtext", "date", "time", "set", "json", "enum", "geometry"}: return pa.string() raise ValueError(f"unsupported empty TiDB column type: {kind}") def materialize_empty_tidb_tables( config: PipelineConfig, gcs_uri: str, patterns: list[str], export_id: str, ) -> list[dict[str, Any]]: """Fill schema-only exports only after checking emptiness at the export TSO. Called only for SUCCEEDED exports. Missing schema, snapshot expiry, a nonempty table, or unknown types fail closed; estimates are never evidence. """ objects = list_objects(gcs_uri) names = {item["name"] for item in objects} data_tables = {object_table_name(config, item["name"]) for item in objects} missing = [pattern.split(".", 1)[1] for pattern in patterns if pattern.split(".", 1)[1] not in data_tables] if not missing: return [] parsed = parse_gs_uri(gcs_uri) schema_uris = {} for name in missing: schema_name = f"{parsed.name}{config.source_database}.{name}-schema.sql" if schema_name not in names: raise RuntimeError(f"missing schema artifact for TiDB table {name}") schema_uris[name] = f"gs://{parsed.bucket}/{schema_name}" metadata = download_text(f"{gcs_uri}metadata") match = re.search(r"(?m)^\s*Pos:\s*(\d+)\s*$", metadata) if not match or "Finished dump at:" not in metadata: raise RuntimeError("empty-table verification requires a completed export snapshot TSO") snapshot = match.group(1) tables = load_table_metadata(config, selected=missing) schemas = { name: pa.schema([ pa.field(str(col["name"]), tidb_empty_arrow_type(col), nullable=bool(col["nullable"]) or col["data_type"] in {"datetime", "timestamp"}) for col in tables[name]["columns"] ]) for name in missing } with tidb_connect(config) as connection, connection.cursor() as cursor: cursor.execute("SET @@tidb_snapshot = %s", (snapshot,)) for name in missing: database = config.source_database.replace("`", "``") table_name = name.replace("`", "``") cursor.execute(f"SELECT 1 FROM `{database}`.`{table_name}` LIMIT 1") if cursor.fetchone() is not None: raise RuntimeError(f"TiDB table {name} has rows at export snapshot but no Parquet") client = storage.Client(project=config.target_project) receipts = [] for name in missing: destination = f"{parsed.name}{config.source_database}.{name}.empty.parquet" with tempfile.NamedTemporaryFile(suffix=".parquet") as tmp: pq.write_table(pa.Table.from_pylist([], schema=schemas[name]), tmp.name, compression=config.parquet_compression.lower()) client.bucket(parsed.bucket).blob(destination).upload_from_filename( tmp.name, content_type="application/octet-stream", if_generation_match=0, ) receipts.append({"table_name": name, "row_count": 0, "snapshot_tso": snapshot, "export_id": export_id, "schema_uri": schema_uris[name], "uri": f"gs://{parsed.bucket}/{destination}"}) return receipts def summarize_export_fault_domains(export_chunks: list[dict[str, Any]]) -> list[dict[str, Any]]: summaries: dict[str, dict[str, Any]] = {} for chunk in export_chunks: name = str(chunk["fault_domain"]) summary = summaries.setdefault( name, { "name": name, "fault_domain_index": int(chunk["fault_domain_index"]), "fault_domain_count": int(chunk["fault_domain_count"]), "chunk_count": 0, "filter_pattern_count": 0, "export_ids": [], }, ) summary["chunk_count"] += 1 summary["filter_pattern_count"] += len(chunk["filter_patterns"]) summary["export_ids"].append(chunk["export_id"]) return sorted(summaries.values(), key=lambda item: int(item["fault_domain_index"])) def extract_s3_snapshot(config: PipelineConfig, *, run_date: str, run_id: str) -> dict[str, Any]: source = parse_s3_uri(config.external_uri) table_name = config.external_table s3_client = build_s3_client(config) gcs_client = storage.Client(project=config.target_project) destination_bucket = gcs_client.bucket(config.bucket) copied = [] for item in list_s3_objects(s3_client, source["bucket"], source["prefix"]): key = item["key"] if not key.endswith(".parquet"): continue relative = key.removeprefix(source["prefix"]).lstrip("/") destination_name = ( f"{config.snapshot_prefix}/source={config.name}/run_date={run_date}/run_id={run_id}/" f"table={table_name}/{relative or Path(key).name}" ) copy_s3_object_to_gcs( s3_client=s3_client, source_bucket=source["bucket"], source_key=key, destination_bucket=destination_bucket, destination_name=destination_name, ) copied.append( { "source_uri": f"s3://{source['bucket']}/{key}", "destination_uri": f"gs://{config.bucket}/{destination_name}", "size": item["size"], } ) if not copied: raise RuntimeError(f"no parquet objects found under {config.external_uri}") return { "engine": "s3", "source_uri": config.external_uri, "gcs_uri": config.gcs_run_uri(run_date, run_id), "copied_objects": copied, } def extract_databend_snapshot(config: PipelineConfig, *, run_date: str, run_id: str) -> dict[str, Any]: dsn = os.environ.get(config.databend_dsn_env) if not dsn: raise RuntimeError(f"required environment variable is not set: {config.databend_dsn_env}") table_name = config.external_table columns = [str(column["name"]) for column in config.external_columns] if not columns: raise RuntimeError(f"Databend source {config.name} has no external columns configured") client = BlockingDatabendClient(databend_dsn(dsn)) cursor = client.cursor() source_table = f"{quote_databend(config.databend_database)}.{quote_databend(config.databend_table)}" cursor.execute(f"SELECT {', '.join(quote_databend(column) for column in columns)} FROM {source_table}") rows = [dict(zip(columns, tuple(row), strict=True)) for row in cursor.fetchall()] destination_name = ( f"{config.snapshot_prefix}/source={config.name}/run_date={run_date}/run_id={run_id}/" f"table={table_name}/{run_id}-part-0.parquet" ) destination_uri = f"gs://{config.bucket}/{destination_name}" write_rows_to_parquet_gcs(config, rows, destination_name) return { "engine": "databend", "source_table": f"{config.databend_database}.{config.databend_table}", "source_uri": config.external.get("uri", ""), "gcs_uri": config.gcs_run_uri(run_date, run_id), "copied_objects": [{"destination_uri": destination_uri, "row_count": len(rows)}], } def extract_mysql_query_snapshot( config: PipelineConfig, *, run_date: str, run_id: str, tables: list[str] | None = None, ) -> dict[str, Any]: selected = (config.selected_tables or ["*"]) if tables is None else tables query_tables = [ table for table in config.query_tables if include_table(str(table["table"]), selected=selected, excluded=config.excluded_tables) ] if not query_tables: raise RuntimeError(f"MySQL query source {config.name} has no selected query tables") copied_objects: list[dict[str, Any]] = [] query_files: list[str] = [] connection = tidb_connect(config) try: with connection.cursor() as cursor: for query_table in query_tables: table_name = str(query_table["table"]) columns = list(query_table.get("columns", [])) expected_columns = [str(column["name"]) for column in columns] if not expected_columns: raise RuntimeError(f"MySQL query table {config.name}.{table_name} has no columns configured") query_path = config.query_sql_path_for(query_table) cursor.execute(read_only_query_sql(query_path)) actual_columns = [str(description[0]) for description in (cursor.description or [])] if actual_columns != expected_columns: raise RuntimeError( f"MySQL query table {config.name}.{table_name} returned columns " f"{actual_columns!r}; expected {expected_columns!r}" ) rows = normalize_external_rows( config, [dict(row) for row in cursor.fetchall()], columns=columns, ) destination_name = ( f"{config.snapshot_prefix}/source={config.name}/run_date={run_date}/run_id={run_id}/" f"table={table_name}/{run_id}-part-0.parquet" ) destination_uri = f"gs://{config.bucket}/{destination_name}" write_rows_to_parquet_gcs(config, rows, destination_name, columns=columns) copied_objects.append( { "destination_uri": destination_uri, "row_count": len(rows), "table_name": table_name, } ) query_files.append(query_path.name) finally: connection.close() return { "engine": "mysql_query", "source_tables": [str(table["table"]) for table in query_tables], "query_files": query_files, "gcs_uri": config.gcs_run_uri(run_date, run_id), "copied_objects": copied_objects, } def read_only_query_sql(path: Path) -> str: sql = path.read_text().strip() statement = sql[:-1].rstrip() if sql.endswith(";") else sql if not re.match(r"^(SELECT|WITH)\b", statement, flags=re.IGNORECASE): raise RuntimeError(f"query file must contain a read-only SELECT statement: {path}") if ";" in statement: raise RuntimeError(f"query file must contain exactly one statement: {path}") if re.search( r"\b(ALTER|CALL|CREATE|DELETE|DROP|GRANT|INSERT|LOAD|RENAME|REPLACE|REVOKE|TRUNCATE|UPDATE)\b", statement, flags=re.IGNORECASE, ): raise RuntimeError(f"query file contains a non-read-only SQL keyword: {path}") return statement def normalize_external_rows( config: PipelineConfig, rows: list[dict[str, Any]], *, columns: list[dict[str, Any]] | None = None, ) -> list[dict[str, Any]]: configured_columns = config.external_columns if columns is None else columns boolean_columns = { str(column["name"]) for column in configured_columns if str(column.get("data_type", column.get("column_type", ""))).lower() in {"boolean", "bool"} } if not boolean_columns: return rows return [ {key: bool(value) if key in boolean_columns and value is not None else value for key, value in row.items()} for row in rows ] def databend_dsn(dsn: str) -> str: if dsn.startswith("lake://"): return "databend://" + dsn.removeprefix("lake://") return dsn def quote_databend(identifier: str) -> str: return f"`{identifier.replace('`', '``')}`" def write_rows_to_parquet_gcs( config: PipelineConfig, rows: list[dict[str, Any]], destination_name: str, *, columns: list[dict[str, Any]] | None = None, ) -> None: table = pa.Table.from_pylist(rows, schema=external_arrow_schema(config, columns=columns)) with tempfile.NamedTemporaryFile(suffix=".parquet") as tmp: pq.write_table(table, tmp.name, compression=config.parquet_compression.lower()) storage.Client(project=config.target_project).bucket(config.bucket).blob(destination_name).upload_from_filename( tmp.name, content_type="application/octet-stream", ) def external_arrow_schema( config: PipelineConfig, *, columns: list[dict[str, Any]] | None = None, ) -> pa.Schema: configured_columns = config.external_columns if columns is None else columns return pa.schema( [ pa.field( str(column["name"]), arrow_type_for(str(column.get("data_type", column.get("column_type", "string")))), nullable=bool(column.get("nullable", True)), ) for column in configured_columns ] ) def arrow_type_for(data_type: str) -> pa.DataType: normalized = data_type.lower().split("(", 1)[0] if normalized in {"varchar", "char", "text", "string"}: return pa.string() if normalized in {"double", "float", "float64"}: return pa.float64() if normalized in {"boolean", "bool"}: return pa.bool_() if normalized == "date": return pa.date32() if normalized in {"timestamp", "datetime"}: return pa.timestamp("us") if normalized in {"int", "integer", "int64", "bigint"}: return pa.int64() return pa.string() def build_manifest( config: PipelineConfig, *, run_date: str, run_id: str, extract_result: dict[str, Any], selected: list[str] | None, ) -> dict[str, Any]: gcs_uri = config.gcs_run_uri(run_date, run_id) chunks = extract_result.get("export_chunks", []) # Failed attempts can leave complete-looking Parquet files. Admit only the # exact successful attempt prefixes, never the enclosing run prefix. if config.source_type == "tidbcloud" and chunks: if any(chunk.get("export", {}).get("state") != "SUCCEEDED" for chunk in chunks): raise RuntimeError("manifest requires every TiDB export chunk to succeed") objects = [item for chunk in chunks for item in list_objects(chunk["gcs_uri"])] else: objects = list_objects(gcs_uri) table_metadata = load_table_metadata(config, selected=selected) tables = summarize_tables(config, run_date=run_date, run_id=run_id, objects=objects, table_metadata=table_metadata) manifest_uri = config.manifest_gcs_uri(run_date, run_id) drift = schema_drift( exported=(str(item["table_name"]) for item in tables), expected=table_metadata.keys(), ) selection = table_selection_contract( config, list(table_metadata), requested=selected, ) if drift["files_without_metadata"] or drift["metadata_without_files"]: raise ExtractRunError( "exported objects do not match the selected manifest table contract", summary={ "engine": str(extract_result.get("engine", config.source_type)), "failure_stage": "manifest_table_coverage", "selection": selection, "schema_drift": drift, "gcs_uri": gcs_uri, "export_ids": list(extract_result.get("export_ids", [])), "fault_domains": list(extract_result.get("fault_domains", [])), }, ) return { "manifest_schema_version": 3, "source_name": config.name, "source_type": config.source_type, "target_project": config.target_project, "source_project": config.source_project if config.source_type == "bigquery" else "", "source_database": config.source_database, "run_date": run_date, "run_id": run_id, "product_status": "exported", "gcs_uri": gcs_uri, "manifest_uri": manifest_uri, "runtime_log_uri": config.runtime_log_gcs_uri(run_date, run_id), "created_at": datetime.now(UTC).isoformat(), "selection": selection, "extract": extract_result, "objects": { "file_count": len(objects), "total_bytes": sum(item["size"] for item in objects), "files": [{"uri": item["uri"], "name": item["name"], "size": item["size"]} for item in objects], }, "tables": tables, "schema_drift": drift, } def summarize_tables( config: PipelineConfig, *, run_date: str, run_id: str, objects: list[dict[str, Any]], table_metadata: dict[str, dict[str, Any]], ) -> list[dict[str, Any]]: grouped: dict[str, dict[str, Any]] = defaultdict(lambda: {"file_count": 0, "total_bytes": 0, "uris": []}) for item in objects: if not item["name"].endswith(".parquet"): continue table_name = object_table_name(config, item["name"]) if table_name is None: continue grouped[table_name]["table_name"] = table_name grouped[table_name]["file_count"] += 1 grouped[table_name]["total_bytes"] += int(item["size"]) grouped[table_name]["uris"].append(item["uri"]) summaries = [] for table_name, summary in grouped.items(): summaries.append( { **table_metadata.get(table_name, {}), **summary, "uris": sorted(summary["uris"]), "gcs_uri": table_gcs_uri_for_manifest(config, table_name, run_date, run_id), } ) return sorted(summaries, key=lambda item: str(item["table_name"])) def object_table_name(config: PipelineConfig, object_name: str) -> str | None: if config.source_type in {"bigquery", "s3", "databend", "mysql_query"}: for segment in object_name.split("/"): if segment.startswith("table="): return segment.removeprefix("table=") return None name = Path(object_name).name if name.endswith("-schema-create.sql") or not name.endswith(".parquet"): return None parts = name.split(".") if len(parts) >= 3 and parts[0] == config.source_database: return parts[1] if len(parts) >= 2: return parts[1] return None def table_gcs_uri_for_manifest(config: PipelineConfig, table_name: str, run_date: str, run_id: str) -> str: if config.source_type in {"bigquery", "s3", "databend", "mysql_query"}: return config.table_gcs_uri(table_name, run_date, run_id) return config.gcs_run_uri(run_date, run_id) def load_table_metadata(config: PipelineConfig, selected: list[str] | None) -> dict[str, dict[str, Any]]: if config.source_type in {"s3", "databend"}: table_name = config.external_table if selected and table_name not in selected: return {} return { table_name: { "source_database": config.source_database, "table_name": table_name, "row_estimate": config.external_row_estimate, "primary_key": config.primary_key_for(table_name), "watermark_columns": config.watermark_candidates, "classifications": list(config.data.get("governance", {}).get("classifications", [])), "columns": config.external_columns, } } if config.source_type == "mysql_query": metadata: dict[str, dict[str, Any]] = {} for table in inspect_tables(config, selected=selected): item = table_to_manifest(table) configured = config.query_table_for(table.name) classifications = ( configured.get("classifications") if configured and configured.get("classifications") is not None else config.data.get("governance", {}).get("classifications", []) ) item["classifications"] = list(classifications) metadata[table.name] = item return metadata if config.source_type == "tidbcloud" and selected is not None: tables = inspect_exact_tidb_tables(config, selected) else: tables = inspect_tables(config, selected=selected) return {table.name: table_to_manifest(table) for table in tables} def schema_drift(exported, expected) -> dict[str, list[str]]: exported_set = set(exported) expected_set = set(expected) return { "files_without_metadata": sorted(exported_set - expected_set), "metadata_without_files": sorted(expected_set - exported_set), } def ensure_ticloud_profile(config: PipelineConfig) -> None: profile = config.tidb_cli_profile if profile in existing_ticloud_profiles(): return public_key = os.environ.get(config.tidb_public_key_env) private_key = os.environ.get(config.tidb_private_key_env) if not public_key or not private_key: raise RuntimeError( f"ticloud profile {profile!r} is missing and {config.tidb_public_key_env}/" f"{config.tidb_private_key_env} are not both set" ) subprocess.run( [ ticloud_bin(), "config", "create", "--profile-name", profile, "--public-key", public_key, "--private-key", private_key, ], check=True, capture_output=True, text=True, ) def existing_ticloud_profiles() -> set[str]: try: completed = subprocess.run([ticloud_bin(), "config", "list"], capture_output=True, text=True, check=False) except OSError: return set() if completed.returncode != 0: return set() return {line.split()[0] for line in completed.stdout.splitlines()[1:] if line.strip()} def wait_for_ticloud_export(config: PipelineConfig, export_id: str) -> dict[str, Any]: interval_seconds = int(config.data.get("export", {}).get("poll_interval_seconds", 30)) timeout_seconds = int(config.data.get("export", {}).get("timeout_seconds", 7200)) deadline = time.monotonic() + timeout_seconds while True: export = find_ticloud_export(config, export_id) if export.get("state") in TERMINAL_EXPORT_STATES: return export if time.monotonic() >= deadline: raise TimeoutError(f"timed out waiting for TiDB Cloud export {export_id}") time.sleep(interval_seconds) def find_ticloud_export(config: PipelineConfig, export_id: str) -> dict[str, Any]: payload = run_json( [ ticloud_bin(), "serverless", "export", "list", "--profile", config.tidb_cli_profile, "--cluster-id", config.tidb_cluster_id, "--output", "json", ] ) for export in payload.get("exports", []): if export.get("exportId") == export_id: return export raise RuntimeError(f"export {export_id} was not found for cluster {config.tidb_cluster_id}") def run_json(args: list[str]) -> dict[str, Any]: completed = subprocess.run(args, check=True, capture_output=True, text=True) return json.loads(completed.stdout or "{}") def run_sensitive_command( args: list[str], *, description: str, sensitive_values: tuple[str, ...], ) -> subprocess.CompletedProcess[str]: completed = subprocess.run(args, check=False, capture_output=True, text=True) if completed.returncode == 0: return completed stdout = redact_sensitive_text(completed.stdout, sensitive_values) stderr = redact_sensitive_text(completed.stderr, sensitive_values) raise RuntimeError( f"{description} failed with exit code {completed.returncode}; stdout={stdout!r}; stderr={stderr!r}" ) def redact_sensitive_text(text: str | None, sensitive_values: tuple[str, ...]) -> str: redacted = text or "" for value in sensitive_values: if value: redacted = redacted.replace(value, "[REDACTED]") return redacted def parse_export_id(stdout: str) -> str | None: try: payload = json.loads(stdout) except json.JSONDecodeError: payload = None if isinstance(payload, dict): nested = payload.get("export") nested_id = nested.get("exportId") if isinstance(nested, dict) else None return payload.get("exportId") or payload.get("id") or nested_id match = re.search(r"\bexp-[a-z0-9]+\b", stdout or "") return match.group(0) if match else None def sanitize_ticloud_export(export: dict[str, Any]) -> dict[str, Any]: allowed = { "clusterId", "completeTime", "createTime", "displayName", "exportId", "exportOptions", "name", "snapshotTime", "state", "updateTime", } return {key: export[key] for key in allowed if key in export} def sanitize_ticloud_export_failure( export: dict[str, Any], *, sensitive_values: tuple[str, ...] = (), ) -> dict[str, Any]: provider_export = {} for key in PROVIDER_DIAGNOSTIC_FIELDS: value = sanitize_provider_scalar(export.get(key), sensitive_values=sensitive_values) if value is not None: provider_export[key] = value provider_reason = first_provider_field( export, PROVIDER_REASON_FIELDS, sensitive_values=sensitive_values, ) provider_error_code = first_provider_field( export, PROVIDER_ERROR_CODE_FIELDS, sensitive_values=sensitive_values, ) nested_error = export.get("error") if isinstance(nested_error, dict): provider_reason = provider_reason or first_provider_field( nested_error, PROVIDER_REASON_FIELDS, sensitive_values=sensitive_values, ) provider_error_code = provider_error_code or first_provider_field( nested_error, PROVIDER_ERROR_CODE_FIELDS, sensitive_values=sensitive_values, ) return { "provider_export": provider_export, "provider_reason": provider_reason, "provider_error_code": provider_error_code, } def first_provider_field( payload: dict[str, Any], fields: tuple[str, ...], *, sensitive_values: tuple[str, ...], ) -> str | int | float | bool | None: for field in fields: value = sanitize_provider_scalar(payload.get(field), sensitive_values=sensitive_values) if value is not None: return value return None def sanitize_provider_scalar( value: Any, *, sensitive_values: tuple[str, ...], ) -> str | int | float | bool | None: if isinstance(value, bool | int | float): return value if not isinstance(value, str) or not value.strip(): return None sanitized = redact_sensitive_text(value, sensitive_values) sanitized = re.sub( r"(?is)-----BEGIN [^-\n]+-----.*?-----END [^-\n]+-----", "[REDACTED]", sanitized, ) sanitized = re.sub( r"(?i)\b(password|passwd|token|secret|private[_ -]?key|service[_ -]?account[_ -]?key)" r"(\s*[:=]\s*)(\"[^\"\n]*\"|'[^'\n]*'|[^\s,;]+)", r"\1\2[REDACTED]", sanitized, ) sanitized = re.sub( r"(?i)\b([a-z][a-z0-9+.-]*://[^:/@\s]+:)[^@\s]+@", r"\1[REDACTED]@", sanitized, ) return sanitized[:2000] def ticloud_bin() -> str: configured = os.environ.get("TICLOUD_BIN") if configured: return configured found = shutil.which("ticloud") if found: return found raise RuntimeError("ticloud CLI is required for TiDB Cloud sources") def parse_s3_uri(uri: str) -> dict[str, str]: parsed = urlparse(uri) if parsed.scheme != "s3" or not parsed.netloc: raise ValueError(f"not an S3 URI: {uri}") return {"bucket": parsed.netloc, "prefix": parsed.path.lstrip("/")} def build_s3_client(config: PipelineConfig): access_key = os.environ.get(config.s3_access_key_id_env) secret_key = os.environ.get(config.s3_secret_access_key_env) if not access_key or not secret_key: if config.s3_profile: return boto3.Session(profile_name=config.s3_profile, region_name=config.s3_region).client("s3") raise RuntimeError( f"required S3 credentials are not set: {config.s3_access_key_id_env}, " f"{config.s3_secret_access_key_env}" ) kwargs = { "region_name": config.s3_region, "aws_access_key_id": access_key, "aws_secret_access_key": secret_key, } if config.s3_session_token_env and os.environ.get(config.s3_session_token_env): kwargs["aws_session_token"] = os.environ[config.s3_session_token_env] return boto3.client("s3", **kwargs) def list_s3_objects(s3_client, bucket: str, prefix: str) -> list[dict[str, Any]]: paginator = s3_client.get_paginator("list_objects_v2") objects = [] for page in paginator.paginate(Bucket=bucket, Prefix=prefix): for item in page.get("Contents", []): objects.append( { "key": item["Key"], "size": int(item.get("Size", 0)), "last_modified": item.get("LastModified").isoformat() if item.get("LastModified") is not None else None, } ) return objects def copy_s3_object_to_gcs( *, s3_client, source_bucket: str, source_key: str, destination_bucket, destination_name: str, ) -> None: response = s3_client.get_object(Bucket=source_bucket, Key=source_key) blob = destination_bucket.blob(destination_name) with response["Body"] as body: with blob.open("wb", content_type="application/octet-stream") as output: for chunk in body.iter_chunks(chunk_size=8 * 1024 * 1024): if chunk: output.write(chunk) def write_local_run_file(config: PipelineConfig, run_date: str, run_id: str, filename: str, text: str) -> None: output_path = config.root / "runs" / run_date / run_id / filename output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(text)