feat: add local and gcp path IO adapter for id_mapping; add unit test for id_registrar

Author: BLMgithub

data_pipeline/contract/contract_executor.py

@@ -130,7 +130,14 @@ def apply_contract(
         valid_ids = set(df.get_column("order_id"))
 
     temp_path = run_context.contracted_path / "id_mapping"
-    storage_mapping = Path(run_context.storage_mapping_path)
+    storage_mapping_str = run_context.storage_mapping_path
+
+    # Switch between local and gcp path
+    storage_mapping = (
+        storage_mapping_str
+        if storage_mapping_str.startswith("gs://")
+        else Path(storage_mapping_str)
+    )
 
     df = id_mapping(
         df=df,
@@ -146,6 +153,8 @@ def apply_contract(
         report["status"] = "failed"
         report["errors"].append("Export failed")
 
+        return report, invalid_ids, valid_ids
+
     report["status"] = "success"
 
     return report, invalid_ids, valid_ids

data_pipeline/contract/id_registrar.py

@@ -4,7 +4,7 @@
 
 import polars as pl
 from pathlib import Path
-import shutil
+from data_pipeline.shared.storage_adapter import promote_new_mapping_files
 
 ID_COLUMNS_TO_MAP = {
     "df_orders": ["order_id", "customer_id"],
@@ -87,7 +87,7 @@ def id_mapping(
     table_name: str,
     mapping_dict: dict,
     runtime_dir: Path,
-    destination: Path,
+    destination: Path | str,
 ) -> pl.LazyFrame:
     """
     Orchestrates the two-phase transformation of UUID strings to persistent Integer keys.
@@ -113,30 +113,39 @@ def id_mapping(
 
     for id_column in cols_to_map:
         mapping_filename = f"{id_column}_mapping.parquet"
-        storage_path = destination / mapping_filename
         temp_path = runtime_dir / mapping_filename
 
-        # Check if mapping exists in storage else create
-        target_path = storage_path if storage_path.exists() else temp_path
+        # When destination is a GCS URI
+        if str(destination).startswith("gs://"):
+            target_path = temp_path
+        else:
+            storage_path = destination / mapping_filename  # type: ignore
+            target_path = storage_path if storage_path.exists() else temp_path
 
         if not target_path.parent.exists():
             target_path.parent.mkdir(parents=True, exist_ok=True)
 
         map_uuid_to_int(df, target_path, id_column)
 
-    # Promote new mapping files from runtime directory to central storage
-    if runtime_dir.exists():
-        destination.mkdir(parents=True, exist_ok=True)
-        for file in runtime_dir.glob("*_mapping.parquet"):
-            shutil.copy2(file, destination)
+    # Promote new mapping to storage
+    promote_new_mapping_files(runtime_dir=runtime_dir, destination=destination)
 
     lf_mapped = df.lazy()
 
     for id_column in cols_to_map:
         mapping_filename = f"{id_column}_mapping.parquet"
-        storage_path = destination / mapping_filename
 
-        # Enrich DataFrame with integer surrogates from the registry
+        # Switch between local and gcp IO
+        if str(destination).startswith("gs://"):
+            destination_str = str(destination)
+            if not destination_str.endswith("/"):
+                destination_str += "/"
+
+            storage_path = f"{destination_str}{mapping_filename}"
+        else:
+            storage_path = str(destination / mapping_filename)  # type: ignore
+
+        # Enriched existing id registry with new mapping
         registry_lf = pl.scan_parquet(storage_path)
         lf_mapped = lf_mapped.join(registry_lf, on=id_column, how="left")
 

data_pipeline/shared/storage_adapter.py

@@ -141,6 +141,7 @@ def upload_contracted_directory(run_context: RunContext) -> None:
 
     Contract:
     - Synchronizes the local 'contracted/' directory to 'storage_contracted_path'.
+    - Excludes the 'id_mapping' directory to prevent cross-contamination.
     - Purpose: Archives newly cleaned data for delta accumulation and historical lineage.
     """
 
@@ -167,7 +168,7 @@ def upload_contracted_directory(run_context: RunContext) -> None:
     bucket = client.bucket(bucket_name)
 
     for file in source.rglob("*"):
-        if file.is_file():
+        if file.is_file() and "id_mapping" not in file.parts:
 
             blob = bucket.blob(f"{prefix}/{file.relative_to(source)}")
             blob.upload_from_filename(file)
@@ -204,3 +205,36 @@ def download_contracted_datasets(run_context: RunContext) -> None:
         target.parent.mkdir(parents=True, exist_ok=True)
 
         blob.download_to_filename(target)
+
+
+def promote_new_mapping_files(runtime_dir: Path, destination: Path | str) -> None:
+    """
+    Synchronizes new UUID mapping files from the local temporary directory to central storage.
+
+    Contract:
+    - Identifies all '*_mapping.parquet' files in the local 'runtime_dir'.
+    - Promotes them to the persistent 'destination' (local directory or GCS bucket).
+    """
+
+    if not runtime_dir.exists():
+        return
+
+    destination_str = str(destination).replace("\\", "/")
+
+    # Local filesystem case
+    if not destination_str.startswith("gs://"):
+        Path(destination).mkdir(parents=True, exist_ok=True)
+        for file in runtime_dir.glob("*_mapping.parquet"):
+            shutil.copy2(file, Path(destination))
+        return
+
+    # GCS case
+    client = storage.Client()
+    bucket_name, prefix = _split_gcs_path(destination_str)
+    bucket = client.bucket(bucket_name)
+
+    for file in runtime_dir.glob("*_mapping.parquet"):
+        if file.is_file():
+            # Create a blob with the target filename in the bucket
+            blob = bucket.blob(f"{prefix}/{file.name}")
+            blob.upload_from_filename(str(file))

tests/test_contract_stage.py

@@ -15,6 +15,7 @@
     enforce_parent_reference,
     enforce_schema,
 )
+from data_pipeline.contract.id_registrar import map_uuid_to_int, id_mapping
 
 # ------------------------------------------------------------
 # FIXTURES
@@ -156,6 +157,64 @@ def test_enforce_schema():
     assert filtered["state"].dtype == pl.Categorical
 
 
+# ------------------------------------------------------------
+# ID REGISTRAR UNIT TESTS
+# ------------------------------------------------------------
+
+
+def test_map_uuid_to_int_new_file(tmp_path):
+    df = pl.DataFrame({"user_id": ["u1", "u2", "u3"]})
+    mapping_file = tmp_path / "user_id_mapping.parquet"
+
+    map_uuid_to_int(df, mapping_file, "user_id")
+
+    assert mapping_file.exists()
+    mapping_df = pl.read_parquet(mapping_file)
+    assert mapping_df.height == 3
+    assert "user_id" in mapping_df.columns
+    assert "user_id_int" in mapping_df.columns
+    assert mapping_df["user_id_int"].to_list() == [1, 2, 3]
+    assert mapping_df["user_id_int"].dtype == pl.UInt32
+
+
+def test_map_uuid_to_int_update_existing(tmp_path):
+    mapping_file = tmp_path / "user_id_mapping.parquet"
+    initial_df = pl.DataFrame({"user_id": ["u1", "u2"], "user_id_int": [1, 2]}).cast(
+        {"user_id_int": pl.UInt32}
+    )
+    initial_df.write_parquet(mapping_file)
+
+    new_df = pl.DataFrame({"user_id": ["u2", "u3", "u4"]})
+    map_uuid_to_int(new_df, mapping_file, "user_id")
+
+    mapping_df = pl.read_parquet(mapping_file).sort("user_id_int")
+    assert mapping_df.height == 4
+    assert set(mapping_df["user_id"].to_list()) == {"u1", "u2", "u3", "u4"}
+    assert mapping_df["user_id_int"].to_list() == [1, 2, 3, 4]
+
+
+def test_id_mapping_orchestration(tmp_path):
+    runtime_dir = tmp_path / "runtime"
+    destination = tmp_path / "destination"
+    runtime_dir.mkdir()
+    destination.mkdir()
+
+    df = pl.DataFrame({"order_id": ["o1", "o2"], "customer_id": ["c1", "c2"]})
+
+    mapping_dict = {"df_orders": ["order_id", "customer_id"]}
+
+    lf_mapped = id_mapping(df, "df_orders", mapping_dict, runtime_dir, destination)
+    result_df = lf_mapped.collect()
+
+    assert "order_id_int" in result_df.columns
+    assert "customer_id_int" in result_df.columns
+    assert result_df.height == 2
+
+    # Check if mapping files were promoted to destination
+    assert (destination / "order_id_mapping.parquet").exists()
+    assert (destination / "customer_id_mapping.parquet").exists()
+
+
 # ------------------------------------------------------------
 # EXECUTOR INTEGRATION TESTS
 # ------------------------------------------------------------
@@ -170,7 +229,6 @@ def test_apply_contract_orders_success(tmp_path, sample_orders_df):
         run_context.raw_snapshot_path / f"df_orders_{suffix}.csv"
     )
 
-    # New 3-tuple return signature
     report, inv_ids, val_ids = apply_contract(run_context, "df_orders")
 
     assert report["status"] == "success"
@@ -186,7 +244,6 @@ def test_apply_contract_cascade_and_valid_propagation(
     run_context = RunContext.create(base=tmp_path, storage=tmp_path / "storage")
     run_context.initialize_directories()
 
-    # o1: valid, o2: unparsable, o3: impossible
     sample_orders_df = sample_orders_df.with_columns(
         pl.when(pl.col("order_id") == "o2")
         .then(pl.lit("garbage"))
@@ -206,13 +263,11 @@ def test_apply_contract_cascade_and_valid_propagation(
         run_context.raw_snapshot_path / f"df_payments_{suffix}.csv"
     )
 
-    # 1. Process Orders
     rep_o, inv_o, val_o = apply_contract(run_context, "df_orders")
     assert "o2" in inv_o  # unparsable
     assert "o3" in inv_o  # impossible
     assert "o1" in val_o  # only one valid
 
-    # 2. Process Payments (should cascade drop o2, o3 and only keep o1)
     rep_p, inv_p, val_p = apply_contract(
         run_context, "df_payments", invalid_order_ids=inv_o, valid_order_ids=val_o
     )