from __future__ import annotations from pathlib import Path import pytest from pingcap_datapipelines.bq_load import load_manifest_uri, validate_manifest_selection_contract from pingcap_datapipelines.config import PipelineConfig from pingcap_datapipelines.extract import ExtractRunError, build_manifest, plan_tidb_export_batches, run_extract from pingcap_datapipelines.models import Table from pingcap_datapipelines.source_inspect import ( normalize_exact_tidb_selection, table_selection_contract, tidb_export_patterns, validate_exact_tidb_selection, ) def nexuscrm_config() -> PipelineConfig: return PipelineConfig( path=Path("source.yaml"), data={ "source": { "name": "nexuscrm", "type": "tidbcloud", "database": "nexuscrm", }, "target": { "project": "gcp-digital-bi", "location": "us-west2", "datasets": { "raw": "nexuscrm_raw", "final": "nexuscrm_final", "control": "nexuscrm_control", }, }, "gcs": { "bucket": "bucket", "snapshot_prefix": "snapshots", "manifest_prefix": "manifests", "runtime_log_prefix": "logs", }, "load": {"materialize_final": True, "full_snapshot_delete": True}, }, ) def table(name: str) -> Table: return Table("nexuscrm", name, 1, [], ["Id"]) @pytest.mark.parametrize( ("selected", "match"), [ (["Account", "Account"], "duplicate"), (["Account", "nexuscrm.Account"], "duplicate"), (["Acc*"], "exact table names"), (["other.Account"], "expected 'nexuscrm'"), ], ) def test_exact_tidb_selection_rejects_duplicates_patterns_and_wrong_database( selected: list[str], match: str, ) -> None: with pytest.raises(ValueError, match=match): normalize_exact_tidb_selection(nexuscrm_config(), selected) def test_exact_tidb_selection_fails_closed_on_omission_or_unexpected_table() -> None: with pytest.raises(ValueError, match=r"missing=\['Order'\].*unexpected=\['User'\]"): validate_exact_tidb_selection( ["Account", "Order"], [table("Account"), table("User")], ) def test_explicit_subset_becomes_the_exact_provider_filter(monkeypatch) -> None: config = nexuscrm_config() calls: list[list[str]] = [] monkeypatch.setattr( "pingcap_datapipelines.source_inspect.inspect_exact_tidb_tables", lambda _config, selected: calls.append(selected) or [table("Account"), table("Order")], ) patterns = tidb_export_patterns(config, selected=["Account", "Order"]) assert calls == [["Account", "Order"]] assert patterns == ["nexuscrm.Account", "nexuscrm.Order"] def test_run_extract_passes_explicit_subset_to_tidb_provider_path(monkeypatch) -> None: calls: list[list[str] | None] = [] monkeypatch.setattr( "pingcap_datapipelines.extract.extract_tidbcloud", lambda _config, *, run_date, run_id, tables: calls.append(tables) or {"engine": "tidbcloud", "run_date": run_date, "run_id": run_id}, ) monkeypatch.setattr( "pingcap_datapipelines.extract.build_manifest", lambda *_args, **_kwargs: {"run_id": "subset-1"}, ) monkeypatch.setattr("pingcap_datapipelines.extract.write_local_run_file", lambda *_: None) monkeypatch.setattr("pingcap_datapipelines.extract.upload_text", lambda *_: None) run_extract( nexuscrm_config(), run_date="2026-09-02", run_id="subset-1", tables=["Account", "Order"], ) assert calls == [["Account", "Order"]] def test_run_extract_rejects_empty_explicit_subset_before_provider(monkeypatch) -> None: called: list[str] = [] monkeypatch.setattr( "pingcap_datapipelines.extract.extract_tidbcloud", lambda *_args, **_kwargs: called.append("provider"), ) with pytest.raises(ValueError, match="must not be empty"): run_extract(nexuscrm_config(), tables=[]) assert called == [] def test_explicit_subset_uses_only_nonempty_fault_domains_and_preserves_every_table_once() -> None: patterns = ["nexuscrm.Account", "nexuscrm.Order"] batches = plan_tidb_export_batches( patterns, fault_domains=[ {"name": "business", "include": ["*"], "exclude": ["_System_*"]}, {"name": "system", "include": ["_System_*"]}, ], max_filter_bytes=10_000, require_all_fault_domains=False, ) assert [batch["fault_domain"] for batch in batches] == ["business"] assert [pattern for batch in batches for pattern in batch["filter_patterns"]] == patterns def test_subset_manifest_contains_only_exact_selected_tables(monkeypatch) -> None: config = nexuscrm_config() metadata = { "Account": {"table_name": "Account", "primary_key": ["Id"], "columns": []}, "Order": {"table_name": "Order", "primary_key": ["Id"], "columns": []}, } monkeypatch.setattr( "pingcap_datapipelines.extract.list_objects", lambda _: [ {"name": "nexuscrm.Account.000001.parquet", "uri": "gs://bucket/account.parquet", "size": 10}, {"name": "nexuscrm.Order.000001.parquet", "uri": "gs://bucket/order.parquet", "size": 20}, ], ) monkeypatch.setattr("pingcap_datapipelines.extract.load_table_metadata", lambda *_args, **_kwargs: metadata) manifest = build_manifest( config, run_date="2026-09-02", run_id="subset-1", extract_result={"engine": "tidbcloud"}, selected=["Order", "Account"], ) assert manifest["manifest_schema_version"] == 3 assert manifest["selection"]["mode"] == "explicit_subset" assert manifest["selection"]["resolved_tables"] == ["Account", "Order"] assert [item["table_name"] for item in manifest["tables"]] == ["Account", "Order"] assert manifest["schema_drift"] == {"files_without_metadata": [], "metadata_without_files": []} def test_subset_manifest_is_not_published_when_provider_objects_include_unselected_table(monkeypatch) -> None: config = nexuscrm_config() monkeypatch.setattr( "pingcap_datapipelines.extract.list_objects", lambda _: [ {"name": "nexuscrm.Account.000001.parquet", "uri": "gs://bucket/account.parquet", "size": 10}, {"name": "nexuscrm.User.000001.parquet", "uri": "gs://bucket/user.parquet", "size": 20}, ], ) monkeypatch.setattr( "pingcap_datapipelines.extract.load_table_metadata", lambda *_args, **_kwargs: { "Account": {"table_name": "Account", "primary_key": ["Id"], "columns": []} }, ) with pytest.raises(ExtractRunError) as error: build_manifest( config, run_date="2026-09-02", run_id="subset-1", extract_result={"engine": "tidbcloud"}, selected=["Account"], ) assert error.value.summary["failure_stage"] == "manifest_table_coverage" assert error.value.summary["schema_drift"] == { "files_without_metadata": ["User"], "metadata_without_files": [], } def test_load_calls_only_tables_bound_by_subset_manifest(monkeypatch) -> None: config = nexuscrm_config() table_names = ["Account", "Order"] manifest = { "manifest_schema_version": 3, "source_name": "nexuscrm", "source_database": "nexuscrm", "run_date": "2026-09-02", "run_id": "subset-1", "manifest_uri": "gs://bucket/manifest.json", "selection": table_selection_contract(config, table_names, requested=table_names), "schema_drift": {"files_without_metadata": [], "metadata_without_files": []}, "tables": [ {"table_name": "Account", "uris": ["gs://bucket/account.parquet"]}, {"table_name": "Order", "uris": ["gs://bucket/order.parquet"]}, ], } loaded: list[str] = [] statuses: list[str] = [] monkeypatch.setattr("pingcap_datapipelines.bq_load.read_manifest", lambda _: manifest) monkeypatch.setattr("pingcap_datapipelines.bq_load.ensure_control_tables", lambda *_: None) monkeypatch.setattr( "pingcap_datapipelines.bq_load.load_table", lambda _client, _config, _manifest, item: loaded.append(item["table_name"]) or {"source_table": item["table_name"], "status": "loaded"}, ) monkeypatch.setattr( "pingcap_datapipelines.bq_load.mark_run", lambda _client, _config, _manifest, *, status, **_kwargs: statuses.append(status), ) monkeypatch.setattr("pingcap_datapipelines.bq_load.write_runtime_log", lambda *_: None) result = load_manifest_uri( object(), config, "gs://bucket/manifest.json", expected_tables=["Order", "Account"], ) assert loaded == table_names assert statuses == ["running", "loaded"] assert [item["source_table"] for item in result["tables"]] == table_names def test_load_rejects_unselected_manifest_table_before_marking_run(monkeypatch) -> None: config = nexuscrm_config() manifest = { "manifest_schema_version": 3, "source_name": "nexuscrm", "source_database": "nexuscrm", "selection": table_selection_contract(config, ["Account"], requested=["Account"]), "schema_drift": {"files_without_metadata": [], "metadata_without_files": []}, "tables": [{"table_name": "Account"}, {"table_name": "User"}], } marked: list[str] = [] monkeypatch.setattr( "pingcap_datapipelines.bq_load.mark_run", lambda *_args, **_kwargs: marked.append("called"), ) with pytest.raises(RuntimeError, match="do not match the selection contract"): validate_manifest_selection_contract(config, manifest) assert marked == [] def test_load_rejects_requested_subset_that_differs_from_resolved_tables() -> None: config = nexuscrm_config() selection = table_selection_contract(config, ["Account"], requested=["Account"]) selection["requested_tables"] = ["User"] manifest = { "manifest_schema_version": 3, "source_name": "nexuscrm", "source_database": "nexuscrm", "selection": selection, "schema_drift": {"files_without_metadata": [], "metadata_without_files": []}, "tables": [{"table_name": "Account"}], } with pytest.raises(RuntimeError, match="requested tables do not match"): validate_manifest_selection_contract(config, manifest) def test_load_rejects_manifest_that_differs_from_operator_bound_subset() -> None: config = nexuscrm_config() table_names = ["Account", "Order"] manifest = { "manifest_schema_version": 3, "source_name": "nexuscrm", "source_database": "nexuscrm", "selection": table_selection_contract(config, table_names, requested=table_names), "schema_drift": {"files_without_metadata": [], "metadata_without_files": []}, "tables": [{"table_name": name} for name in table_names], } with pytest.raises(RuntimeError, match="operator-bound expected subset"): validate_manifest_selection_contract( config, manifest, expected_tables=["Account", "Order", "User"], ) def test_operator_bound_subset_rejects_legacy_manifest_without_selection_contract() -> None: config = nexuscrm_config() manifest = { "manifest_schema_version": 2, "source_name": "nexuscrm", "source_database": "nexuscrm", "tables": [{"table_name": "Account"}], } with pytest.raises(RuntimeError, match="schema-v3 selection contract"): validate_manifest_selection_contract(config, manifest, expected_tables=["Account"]) def test_manifest_never_lists_failed_attempt_prefix(monkeypatch): config = nexuscrm_config() prefixes = [] def objects(prefix): prefixes.append(prefix) return [{'name': 'nexuscrm.Account.000001.parquet', 'uri': prefix+'account.parquet', 'size': 1}] monkeypatch.setattr('pingcap_datapipelines.extract.list_objects', objects) monkeypatch.setattr('pingcap_datapipelines.extract.load_table_metadata', lambda *a, **kw: { 'Account': {'table_name': 'Account', 'primary_key': ['Id'], 'columns': []}, }) prefix = 'gs://bucket/run/attempt=0002/' manifest = build_manifest(config, run_date='2026-09-08', run_id='retry', selected=['Account'], extract_result={'engine': 'tidbcloud', 'export_chunks': [ {'gcs_uri': prefix, 'export': {'state': 'SUCCEEDED'}}, ]}) assert prefixes == [prefix] assert manifest['tables'][0]['uris'] == [prefix+'account.parquet']