feat(ingestion/transformer): create tag if not exist

metadata-ingestion/src/datahub/ingestion/graph/client.py

@@ -787,9 +787,11 @@ def get_aspect_counts(self, aspect: str, urn_like: Optional[str] = None) -> int:
 
     def execute_graphql(self, query: str, variables: Optional[Dict] = None) -> Dict:
         url = f"{self.config.server}/api/graphql"
+
         body: Dict = {
             "query": query,
         }
+
         if variables:
             body["variables"] = variables
 
@@ -1065,6 +1067,22 @@ def parse_sql_lineage(
             default_schema=default_schema,
         )
 
+    def create_tag(self, tag_name: str) -> Dict[Any, Any]:
+        graph_query: str = """
+            mutation($tag_detail: CreateTagInput!) {
+                createTag(input: $tag_detail)
+            }
+        """
+
+        variables = {
+            "tag_detail": {"name": tag_name},
+        }
+
+        return self.execute_graphql(
+            query=graph_query,
+            variables=variables,
+        )
+
     def close(self) -> None:
         self._make_schema_resolver.cache_clear()
         super().close()

metadata-ingestion/src/datahub/ingestion/transformer/add_dataset_tags.py

@@ -1,14 +1,24 @@
+import logging
 from typing import Callable, List, Optional, cast
 
+import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import (
     KeyValuePattern,
     TransformerSemanticsConfigModel,
 )
 from datahub.configuration.import_resolver import pydantic_resolve_key
 from datahub.emitter.mce_builder import Aspect
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.transformer.dataset_transformer import DatasetTagsTransformer
-from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
+from datahub.metadata.schema_classes import (
+    GlobalTagsClass,
+    TagAssociationClass,
+    TagKeyClass,
+)
+from datahub.utilities.urns.tag_urn import TagUrn
+
+logger = logging.getLogger(__name__)
 
 
 class AddDatasetTagsConfig(TransformerSemanticsConfigModel):
@@ -50,6 +60,36 @@ def transform_aspect(
             self.config, self.ctx.graph, entity_urn, out_global_tags_aspect
         )
 
+    def handle_end_of_stream(
+        self, entity_urn: str
+    ) -> List[MetadataChangeProposalWrapper]:
+        tags_to_add: List[TagAssociationClass] = self.config.get_tags_to_add(entity_urn)
+
+        mcps: List[MetadataChangeProposalWrapper] = []
+
+        if tags_to_add is None:
+            return mcps
+
+        logger.debug("Generating tags")
+
+        for tag_association in tags_to_add:
+            ids: List[str] = TagUrn.create_from_string(
+                tag_association.tag
+            ).get_entity_id()
+
+            assert len(ids) == 1, "Invalid Tag Urn"
+
+            tag_name: str = ids[0]
+
+            mcps.append(
+                MetadataChangeProposalWrapper(
+                    entityUrn=builder.make_tag_urn(tag=tag_name),
+                    aspect=TagKeyClass(name=tag_name),
+                )
+            )
+
+        return mcps
+
 
 class SimpleDatasetTagConfig(TransformerSemanticsConfigModel):
     tag_urns: List[str]

metadata-ingestion/src/datahub/ingestion/transformer/base_transformer.py

@@ -17,13 +17,22 @@
 log = logging.getLogger(__name__)
 
 
-class LegacyMCETransformer(Transformer, metaclass=ABCMeta):
+class HandleEndOfStreamTransformer:
+    def handle_end_of_stream(
+        self, entity_urn: str
+    ) -> List[MetadataChangeProposalWrapper]:
+        return []
+
+
+class LegacyMCETransformer(
+    Transformer, HandleEndOfStreamTransformer, metaclass=ABCMeta
+):
     @abstractmethod
     def transform_one(self, mce: MetadataChangeEventClass) -> MetadataChangeEventClass:
         pass
 
 
-class SingleAspectTransformer(metaclass=ABCMeta):
+class SingleAspectTransformer(HandleEndOfStreamTransformer, metaclass=ABCMeta):
     @abstractmethod
     def aspect_name(self) -> str:
         """Implement this method to specify a single aspect that the transformer is interested in subscribing to. No default provided."""
@@ -180,6 +189,28 @@ def _transform_or_record_mcpw(
             self._record_mcp(envelope.record)
         return envelope if envelope.record.aspect is not None else None
 
+    def _handle_end_of_stream(
+        self, envelope: RecordEnvelope, entity_urn: str
+    ) -> Iterable[RecordEnvelope]:
+
+        if not isinstance(self, SingleAspectTransformer) and not isinstance(
+            self, LegacyMCETransformer
+        ):
+            return
+
+        mcps: List[MetadataChangeProposalWrapper] = self.handle_end_of_stream(
+            entity_urn
+        )
+
+        for mcp in mcps:
+            if mcp.aspect is None or mcp.entityUrn is None:  # to silent the lint error
+                continue
+
+            yield RecordEnvelope(
+                record=mcp,
+                metadata=envelope.metadata,
+            )
+
     def transform(
         self, record_envelopes: Iterable[RecordEnvelope]
     ) -> Iterable[RecordEnvelope]:
@@ -216,15 +247,7 @@ def transform(
                             else None,
                         )
                         if transformed_aspect:
-                            # for end of stream records, we modify the workunit-id
                             structured_urn = Urn.create_from_string(urn)
-                            simple_name = "-".join(structured_urn.get_entity_id())
-                            record_metadata = envelope.metadata.copy()
-                            record_metadata.update(
-                                {
-                                    "workunit_id": f"txform-{simple_name}-{self.aspect_name()}"
-                                }
-                            )
                             yield RecordEnvelope(
                                 record=MetadataChangeProposalWrapper(
                                     entityUrn=urn,
@@ -235,7 +258,12 @@ def transform(
                                     aspectName=self.aspect_name(),
                                     aspect=transformed_aspect,
                                 ),
-                                metadata=record_metadata,
+                                metadata=envelope.metadata,
                             )
+                        yield from self._handle_end_of_stream(
+                            envelope=envelope, entity_urn=urn
+                        )
+
                     self._mark_processed(urn)
+
             yield envelope

metadata-ingestion/tests/unit/test_transform_dataset.py

@@ -775,14 +775,19 @@ def test_simple_dataset_tags_transformation(mock_time):
             ]
         )
     )
-    assert len(outputs) == 3
+
+    assert len(outputs) == 5
 
     # Check that tags were added.
     tags_aspect = outputs[1].record.aspect
     assert tags_aspect
     assert len(tags_aspect.tags) == 2
     assert tags_aspect.tags[0].tag == builder.make_tag_urn("NeedsDocumentation")
 
+    # Check new tag entity should be there
+    assert outputs[2].record.entityUrn == builder.make_tag_urn("NeedsDocumentation")
+    assert outputs[3].record.entityUrn == builder.make_tag_urn("Legacy")
+
 
 def dummy_tag_resolver_method(dataset_snapshot):
     return []
@@ -815,7 +820,7 @@ def test_pattern_dataset_tags_transformation(mock_time):
         )
     )
 
-    assert len(outputs) == 3
+    assert len(outputs) == 5
     tags_aspect = outputs[1].record.aspect
     assert tags_aspect
     assert len(tags_aspect.tags) == 2
@@ -1324,7 +1329,7 @@ def test_mcp_add_tags_missing(mock_time):
     ]
     input_stream.append(RecordEnvelope(record=EndOfStream(), metadata={}))
     outputs = list(transformer.transform(input_stream))
-    assert len(outputs) == 3
+    assert len(outputs) == 5
     assert outputs[0].record == dataset_mcp
     # Check that tags were added, this will be the second result
     tags_aspect = outputs[1].record.aspect