fixing up kafka and hana sources

datahub-project · Jul 10, 2022 · d00ccaa · d00ccaa
1 parent e0a1513
commit d00ccaa
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka.py b/metadata-ingestion/src/datahub/ingestion/source/kafka.py
@@ -47,6 +47,7 @@
     JobStatusClass,
     SubTypesClass,
 )
+from datahub.utilities.registries.domain_registry import DomainRegistry
 
 logger = logging.getLogger(__name__)
 
@@ -150,6 +151,11 @@ def __init__(self, config: KafkaSourceConfig, ctx: PipelineContext):
         self.schema_registry_client: KafkaSchemaRegistryBase = (
             KafkaSource.create_schema_registry(config, self.report)
         )
+        if self.source_config.domain:
+            self.domain_registry = DomainRegistry(
+                cached_domains=[k for k in self.source_config.domain],
+                graph=self.ctx.graph,
+            )
 
     def is_checkpointing_enabled(self, job_id: JobId) -> bool:
         if (
@@ -333,7 +339,9 @@ def _extract_record(self, topic: str) -> Iterable[MetadataWorkUnit]:
         # 6. Emit domains aspect MCPW
         for domain, pattern in self.source_config.domain.items():
             if pattern.allowed(dataset_name):
-                domain_urn = make_domain_urn(domain)
+                domain_urn = make_domain_urn(
+                    self.domain_registry.get_domain_urn(domain)
+                )
 
         if domain_urn:
             wus = add_domain_to_entity_wu(

diff --git a/metadata-ingestion/tests/integration/hana/hana_to_file.yml b/metadata-ingestion/tests/integration/hana/hana_to_file.yml
@@ -32,7 +32,7 @@ source:
       include_field_histogram: true
       include_field_sample_values: true
     domain:
-      sales:
+      "urn:li:domain:sales":
         allow:
           - "HOTEL"
 sink:

diff --git a/metadata-ingestion/tests/integration/kafka/kafka_to_file.yml b/metadata-ingestion/tests/integration/kafka/kafka_to_file.yml
@@ -7,7 +7,7 @@ source:
         bootstrap: "localhost:59092"
         schema_registry_url: "http://localhost:58081"
     domain:
-      sales:
+      "urn:li:domain:sales":
         allow:
           - "key_value_topic"
 sink:
-Original file line number
+Diff line change
@@ Expand Up / @@ -32,7 +32,7 @@ source: @@
           include_field_histogram: true
           include_field_sample_values: true
         domain:
-          sales:
+          "urn:li:domain:sales":
             allow:
               - "HOTEL"
     sink:
@@ Expand Down @@