from __future__ import annotations import json import subprocess from pathlib import Path from unittest.mock import patch import pytest from pingcap_datapipelines.config import PipelineConfig from pingcap_datapipelines.extract import ( extract_mysql_query_snapshot, normalize_external_rows, read_only_query_sql, run_sensitive_command, sanitize_ticloud_export_failure, ) def test_run_sensitive_command_preserves_error_evidence_without_secret() -> None: secret = "base64-service-account-key" completed = subprocess.CompletedProcess( args=["ticloud", "--gcs.service-account-key", secret], returncode=1, stdout=f'{{"requestId":"req-123","credential":"{secret}"}}', stderr=f"HTTP 503 request-id=req-123 credential={secret}", ) with patch("pingcap_datapipelines.extract.subprocess.run", return_value=completed) as run: with pytest.raises(RuntimeError) as raised: run_sensitive_command( completed.args, description="TiDB Cloud CreateExport", sensitive_values=(secret,), ) message = str(raised.value) assert "TiDB Cloud CreateExport failed with exit code 1" in message assert "HTTP 503" in message assert "req-123" in message assert secret not in message assert message.count("[REDACTED]") == 2 run.assert_called_once_with(completed.args, check=False, capture_output=True, text=True) def test_run_sensitive_command_returns_successful_process() -> None: completed = subprocess.CompletedProcess( args=["ticloud", "serverless", "export", "create"], returncode=0, stdout='{"exportId":"exp-123"}', stderr="", ) with patch("pingcap_datapipelines.extract.subprocess.run", return_value=completed): actual = run_sensitive_command( completed.args, description="TiDB Cloud CreateExport", sensitive_values=("unused-secret",), ) assert actual is completed def test_terminal_export_diagnostics_allowlist_and_redact_provider_reason() -> None: secret = "base64-service-account-key" diagnostics = sanitize_ticloud_export_failure( { "clusterId": "cluster-1", "exportId": "exp-123", "state": "FAILED", "completeTime": "2026-09-02T01:03:55Z", "errorMessage": f"Export failed with unknown error; token={secret}", "exportOptions": {"gcsServiceAccountKey": secret}, "internalDebugPayload": secret, }, sensitive_values=(secret,), ) assert diagnostics == { "provider_export": { "clusterId": "cluster-1", "completeTime": "2026-09-02T01:03:55Z", "exportId": "exp-123", "state": "FAILED", }, "provider_reason": "Export failed with unknown error; token=[REDACTED]", "provider_error_code": None, } assert secret not in json.dumps(diagnostics) def test_terminal_export_diagnostics_accept_nested_safe_error_fields() -> None: diagnostics = sanitize_ticloud_export_failure( { "exportId": "exp-123", "state": "FAILED", "error": {"code": "UNKNOWN", "message": "Export failed with unknown error"}, } ) assert diagnostics["provider_reason"] == "Export failed with unknown error" assert diagnostics["provider_error_code"] == "UNKNOWN" def test_read_only_query_sql_accepts_one_select(tmp_path: Path) -> None: path = tmp_path / "query.sql" path.write_text("WITH rows AS (SELECT 1 AS id) SELECT id FROM rows;\n") assert read_only_query_sql(path) == "WITH rows AS (SELECT 1 AS id) SELECT id FROM rows" @pytest.mark.parametrize( "sql", [ "DELETE FROM accounts", "SELECT 1; SELECT 2", "WITH doomed AS (SELECT id FROM accounts) DELETE FROM accounts WHERE id IN (SELECT id FROM doomed)", ], ) def test_read_only_query_sql_rejects_mutation_or_multiple_statements(tmp_path: Path, sql: str) -> None: path = tmp_path / "query.sql" path.write_text(sql) with pytest.raises(RuntimeError): read_only_query_sql(path) def test_normalize_external_rows_converts_mysql_boolean_integers() -> None: config = PipelineConfig( path=Path("source.yaml"), data={ "source": { "name": "helix", "type": "mysql_query", "external": { "table": "cloud_accounts", "columns": [ {"name": "scope_id", "data_type": "string"}, {"name": "is_deleted", "data_type": "boolean"}, ], }, } }, ) assert normalize_external_rows( config, [{"scope_id": "active", "is_deleted": 0}, {"scope_id": "deleted", "is_deleted": 1}], ) == [ {"scope_id": "active", "is_deleted": False}, {"scope_id": "deleted", "is_deleted": True}, ] def test_extract_mysql_query_snapshot_exports_all_selected_query_tables(tmp_path: Path) -> None: (tmp_path / "cloud_accounts.sql").write_text("SELECT 'scope' AS scope_id, 0 AS is_deleted") (tmp_path / "users.sql").write_text("SELECT 'user' AS user_id, 1 AS is_deleted") config = PipelineConfig( path=tmp_path / "source.yaml", data={ "source": { "name": "helix", "type": "mysql_query", "database": "helix", "query": { "tables": [ { "table": "cloud_accounts", "sql_file": "cloud_accounts.sql", "columns": [ {"name": "scope_id", "data_type": "string"}, {"name": "is_deleted", "data_type": "boolean"}, ], }, { "table": "users", "sql_file": "users.sql", "columns": [ {"name": "user_id", "data_type": "string"}, {"name": "is_deleted", "data_type": "boolean"}, ], }, ] }, }, "target": { "project": "gcp-digital-bi", "location": "us-west2", "datasets": {"raw": "raw", "final": "helix", "control": "control"}, }, "tables": {"include": ["cloud_accounts", "users"], "exclude": []}, "gcs": { "bucket": "bucket", "snapshot_prefix": "snapshots", "manifest_prefix": "manifests", "runtime_log_prefix": "logs", }, }, ) class Cursor: description = [] def __init__(self) -> None: self.rows: list[dict[str, object]] = [] def __enter__(self): return self def __exit__(self, *_args): return False def execute(self, sql: str) -> None: if "scope_id" in sql: self.description = [("scope_id",), ("is_deleted",)] self.rows = [{"scope_id": "scope", "is_deleted": 0}] else: self.description = [("user_id",), ("is_deleted",)] self.rows = [{"user_id": "user", "is_deleted": 1}] def fetchall(self): return self.rows class Connection: closed = False def cursor(self): return Cursor() def close(self) -> None: self.closed = True connection = Connection() with ( patch("pingcap_datapipelines.extract.tidb_connect", return_value=connection), patch("pingcap_datapipelines.extract.write_rows_to_parquet_gcs") as write_rows, ): result = extract_mysql_query_snapshot( config, run_date="2026-08-05", run_id="20260805T000000Z", ) assert connection.closed is True assert result["source_tables"] == ["cloud_accounts", "users"] assert result["query_files"] == ["cloud_accounts.sql", "users.sql"] assert [item["row_count"] for item in result["copied_objects"]] == [1, 1] assert write_rows.call_count == 2 assert write_rows.call_args_list[0].args[1] == [{"scope_id": "scope", "is_deleted": False}] assert write_rows.call_args_list[1].args[1] == [{"user_id": "user", "is_deleted": True}] def test_extract_mysql_query_snapshot_rejects_query_schema_drift(tmp_path: Path) -> None: (tmp_path / "users.sql").write_text("SELECT 'user' AS wrong_id") config = PipelineConfig( path=tmp_path / "source.yaml", data={ "source": { "name": "helix", "type": "mysql_query", "query": { "tables": [ { "table": "users", "sql_file": "users.sql", "columns": [{"name": "user_id", "data_type": "string"}], } ] }, }, "target": { "project": "gcp-digital-bi", "location": "us-west2", "datasets": {"raw": "raw", "final": "helix", "control": "control"}, }, "tables": {"include": ["users"], "exclude": []}, "gcs": { "bucket": "bucket", "snapshot_prefix": "snapshots", "manifest_prefix": "manifests", "runtime_log_prefix": "logs", }, }, ) cursor = type( "Cursor", (), { "description": [("wrong_id",)], "__enter__": lambda self: self, "__exit__": lambda self, *_args: False, "execute": lambda self, _sql: None, "fetchall": lambda self: [{"wrong_id": "user"}], }, )() connection = type( "Connection", (), {"cursor": lambda self: cursor, "close": lambda self: None}, )() with patch("pingcap_datapipelines.extract.tidb_connect", return_value=connection): with pytest.raises(RuntimeError, match="returned columns"): extract_mysql_query_snapshot( config, run_date="2026-08-05", run_id="20260805T000000Z", )