feat(ingest): migrate Cassandra source to new SDK (#12695)

datahub-project · Feb 24, 2025 · 8dfd8fb · 8dfd8fb
1 parent 005a9b0
commit 8dfd8fb
Show file tree

Hide file tree

Showing 11 changed files with 1,435 additions and 1,261 deletions.
diff --git a/metadata-ingestion/scripts/avro_codegen.py b/metadata-ingestion/scripts/avro_codegen.py
@@ -714,7 +714,7 @@ def from_key_aspect(cls, key_aspect: "{key_aspect_class}") -> "{class_name}":
         code += f"""
     @property
     def {field_name(field)}(self) -> {field_type(field)}:
-        return self.entity_ids[{i}]
+        return self._entity_ids[{i}]
 """
 
     return code

diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -43,6 +43,7 @@
     SystemMetadataTransformer,
 )
 from datahub.ingestion.transformer.transform_registry import transform_registry
+from datahub.sdk._attribution import KnownAttribution, change_default_attribution
 from datahub.telemetry import stats
 from datahub.telemetry.telemetry import telemetry_instance
 from datahub.utilities._custom_package_loader import model_version_name
@@ -410,6 +411,10 @@ def run(self) -> None:
                     )
                 )
 
+            self.exit_stack.enter_context(
+                change_default_attribution(KnownAttribution.INGESTION)
+            )
+
             self.final_status = PipelineStatus.UNKNOWN
             self._notify_reporters_on_ingestion_start()
             callback = None