feat(ingestion/transformer): Add dataset dataproduct transformer (dat…

…ahub-project#9491)

Co-authored-by: Harshal Sheth <[email protected]>

metadata-ingestion/docs/transformer/dataset_transformer.md

@@ -14,6 +14,7 @@ The below table shows transformer which can transform aspects of entity [Dataset
 | `schemaMetadata`    | - [Pattern Add Dataset Schema Field glossaryTerms](#pattern-add-dataset-schema-field-glossaryterms)<br/> - [Pattern Add Dataset Schema Field globalTags](#pattern-add-dataset-schema-field-globaltags)            |
 | `datasetProperties` | - [Simple Add Dataset datasetProperties](#simple-add-dataset-datasetproperties)<br/> - [Add Dataset datasetProperties](#add-dataset-datasetproperties)                                                            |
 | `domains`           | - [Simple Add Dataset domains](#simple-add-dataset-domains)<br/> - [Pattern Add Dataset domains](#pattern-add-dataset-domains)                                                                                      | 
+| `dataProduct`       | - [Simple Add Dataset dataProduct ](#simple-add-dataset-dataproduct)<br/> - [Pattern Add Dataset dataProduct](#pattern-add-dataset-dataproduct)<br/> - [Add Dataset dataProduct](#add-dataset-dataproduct)  
 
 ## Extract Ownership from Tags
 ### Config Details
@@ -961,6 +962,75 @@ in both of the cases domain should be provisioned on DataHub GMS
                 'urn:li:dataset:\(urn:li:dataPlatform:postgres,postgres\.public\.n.*': ["hr"]
                 'urn:li:dataset:\(urn:li:dataPlatform:postgres,postgres\.public\.t.*': ["urn:li:domain:finance"] 
   ```
+## Simple Add Dataset dataProduct
+### Config Details
+| Field                         | Required | Type            | Default       | Description                                                                            |
+|-------------------------------|----------|-----------------|---------------|----------------------------------------------------------------------------------------|
+| `dataset_to_data_product_urns`| ✅       | Dict[str, str]  |               | Dataset Entity urn as key and dataproduct urn as value to create with dataset as asset.|
+
+Let’s suppose we’d like to add a set of dataproduct with specific datasets as its assets. To do so, we can use the `simple_add_dataset_dataproduct` transformer that’s included in the ingestion framework.
+
+The config, which we’d append to our ingestion recipe YAML, would look like this:
+
+  ```yaml
+  transformers:
+    - type: "simple_add_dataset_dataproduct"
+      config:
+        dataset_to_data_product_urns:
+          "urn:li:dataset:(urn:li:dataPlatform:bigquery,example1,PROD)": "urn:li:dataProduct:first"
+          "urn:li:dataset:(urn:li:dataPlatform:bigquery,example2,PROD)": "urn:li:dataProduct:second"
+  ```
+
+## Pattern Add Dataset dataProduct
+### Config Details
+| Field                                 | Required | Type                 | Default     | Description                                                                                 |
+|---------------------------------------|----------|----------------------|-------------|---------------------------------------------------------------------------------------------|
+| `dataset_to_data_product_urns_pattern`| ✅       | map[regx, urn]       |             | Dataset Entity urn with regular expression and dataproduct urn apply to matching entity urn.|
+
+Let’s suppose we’d like to append a series of dataproducts with specific datasets as its assets. To do so, we can use the `pattern_add_dataset_dataproduct` module that’s included in the ingestion framework.  This will match the regex pattern to `urn` of the dataset and create the data product entity with given urn and matched datasets as its assets. 
+
+The config, which we’d append to our ingestion recipe YAML, would look like this:
+
+  ```yaml
+  transformers:
+    - type: "pattern_add_dataset_dataproduct"
+      config:
+        dataset_to_data_product_urns_pattern:
+          rules:
+            ".*example1.*": "urn:li:dataProduct:first"
+            ".*example2.*": "urn:li:dataProduct:second"
+  ```
+
+## Add Dataset dataProduct
+### Config Details
+| Field                       | Required | Type                              | Default       | Description                                                                              |
+|-----------------------------|----------|-----------------------------------|---------------|------------------------------------------------------------------------------------------|
+| `get_data_product_to_add`   | ✅        | callable[[str], Optional[str]]   |               | A function which takes dataset entity urn as input and return dataproduct urn to create. |
+
+If you'd like to add more complex logic for creating dataproducts, you can use the more generic add_dataset_dataproduct transformer, which calls a user-provided function to determine the dataproduct to create with specified datasets as its asset.
+
+```yaml
+transformers:
+  - type: "add_dataset_dataproduct"
+    config:
+      get_data_product_to_add: "<your_module>.<your_function>"
+```
+
+Then define your function to return a dataproduct entity urn, for example:
+
+```python
+import datahub.emitter.mce_builder as builder
+
+def custom_dataproducts(entity_urn: str) -> Optional[str]:
+    """Compute the dataproduct urn to a given dataset urn."""
+
+    dataset_to_data_product_map = {
+        builder.make_dataset_urn("bigquery", "example1"): "urn:li:dataProduct:first"
+    }
+    return dataset_to_data_product_map.get(dataset_urn)
+```
+Finally, you can install and use your custom transformer as [shown here](#installing-the-package).
+
 ## Relationship Between replace_existing and semantics
 The transformer behaviour mentioned here is in context of `simple_add_dataset_ownership`, however it is applicable for all dataset transformers which are supporting `replace_existing`
 and `semantics` configuration attributes, for example `simple_add_dataset_tags` will add or remove tags as per behaviour mentioned in this section.

metadata-ingestion/setup.py

@@ -649,6 +649,9 @@
         "pattern_add_dataset_schema_terms = datahub.ingestion.transformer.add_dataset_schema_terms:PatternAddDatasetSchemaTerms",
         "pattern_add_dataset_schema_tags = datahub.ingestion.transformer.add_dataset_schema_tags:PatternAddDatasetSchemaTags",
         "extract_ownership_from_tags = datahub.ingestion.transformer.extract_ownership_from_tags:ExtractOwnersFromTagsTransformer",
+        "add_dataset_dataproduct = datahub.ingestion.transformer.add_dataset_dataproduct:AddDatasetDataProduct",
+        "simple_add_dataset_dataproduct = datahub.ingestion.transformer.add_dataset_dataproduct:SimpleAddDatasetDataProduct",
+        "pattern_add_dataset_dataproduct = datahub.ingestion.transformer.add_dataset_dataproduct:PatternAddDatasetDataProduct",
     ],
     "datahub.ingestion.sink.plugins": [
         "file = datahub.ingestion.sink.file:FileSink",

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

@@ -0,0 +1,133 @@
+import logging
+from typing import Callable, Dict, List, Optional, Union
+
+import pydantic
+
+from datahub.configuration.common import ConfigModel, KeyValuePattern
+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 (
+    DatasetDataproductTransformer,
+)
+from datahub.metadata.schema_classes import MetadataChangeProposalClass
+from datahub.specific.dataproduct import DataProductPatchBuilder
+
+logger = logging.getLogger(__name__)
+
+
+class AddDatasetDataProductConfig(ConfigModel):
+    # dataset_urn -> data product urn
+    get_data_product_to_add: Callable[[str], Optional[str]]
+
+    _resolve_data_product_fn = pydantic_resolve_key("get_data_product_to_add")
+
+
+class AddDatasetDataProduct(DatasetDataproductTransformer):
+    """Transformer that adds dataproduct entity for provided dataset as its asset according to a callback function."""
+
+    ctx: PipelineContext
+    config: AddDatasetDataProductConfig
+
+    def __init__(self, config: AddDatasetDataProductConfig, ctx: PipelineContext):
+        super().__init__()
+        self.ctx = ctx
+        self.config = config
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> "AddDatasetDataProduct":
+        config = AddDatasetDataProductConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    def transform_aspect(
+        self, entity_urn: str, aspect_name: str, aspect: Optional[Aspect]
+    ) -> Optional[Aspect]:
+        return None
+
+    def handle_end_of_stream(
+        self,
+    ) -> List[Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]]:
+        data_products: Dict[str, DataProductPatchBuilder] = {}
+
+        logger.debug("Generating dataproducts")
+        for entity_urn in self.entity_map.keys():
+            data_product_urn = self.config.get_data_product_to_add(entity_urn)
+            if data_product_urn:
+                if data_product_urn not in data_products:
+                    data_products[data_product_urn] = DataProductPatchBuilder(
+                        data_product_urn
+                    ).add_asset(entity_urn)
+                else:
+                    data_products[data_product_urn] = data_products[
+                        data_product_urn
+                    ].add_asset(entity_urn)
+
+        mcps: List[
+            Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]
+        ] = []
+        for data_product in data_products.values():
+            mcps.extend(list(data_product.build()))
+        return mcps
+
+
+class SimpleDatasetDataProductConfig(ConfigModel):
+    dataset_to_data_product_urns: Dict[str, str]
+
+
+class SimpleAddDatasetDataProduct(AddDatasetDataProduct):
+    """Transformer that adds a specified dataproduct entity for provided dataset as its asset."""
+
+    def __init__(self, config: SimpleDatasetDataProductConfig, ctx: PipelineContext):
+
+        generic_config = AddDatasetDataProductConfig(
+            get_data_product_to_add=lambda dataset_urn: config.dataset_to_data_product_urns.get(
+                dataset_urn
+            ),
+        )
+        super().__init__(generic_config, ctx)
+
+    @classmethod
+    def create(
+        cls, config_dict: dict, ctx: PipelineContext
+    ) -> "SimpleAddDatasetDataProduct":
+        config = SimpleDatasetDataProductConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+
+class PatternDatasetDataProductConfig(ConfigModel):
+    dataset_to_data_product_urns_pattern: KeyValuePattern = KeyValuePattern.all()
+
+    @pydantic.root_validator(pre=True)
+    def validate_pattern_value(cls, values: Dict) -> Dict:
+        rules = values["dataset_to_data_product_urns_pattern"]["rules"]
+        for key, value in rules.items():
+            if isinstance(value, list) and len(value) > 1:
+                raise ValueError(
+                    "Same dataset cannot be an asset of two different data product."
+                )
+            elif isinstance(value, str):
+                rules[key] = [rules[key]]
+        return values
+
+
+class PatternAddDatasetDataProduct(AddDatasetDataProduct):
+    """Transformer that adds a specified dataproduct entity for provided dataset as its asset."""
+
+    def __init__(self, config: PatternDatasetDataProductConfig, ctx: PipelineContext):
+        dataset_to_data_product = config.dataset_to_data_product_urns_pattern
+        generic_config = AddDatasetDataProductConfig(
+            get_data_product_to_add=lambda dataset_urn: dataset_to_data_product.value(
+                dataset_urn
+            )[0]
+            if dataset_to_data_product.value(dataset_urn)
+            else None,
+        )
+        super().__init__(generic_config, ctx)
+
+    @classmethod
+    def create(
+        cls, config_dict: dict, ctx: PipelineContext
+    ) -> "PatternAddDatasetDataProduct":
+        config = PatternDatasetDataProductConfig.parse_obj(config_dict)
+        return cls(config, ctx)

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

@@ -1,5 +1,5 @@
 import logging
-from typing import Callable, List, Optional, cast
+from typing import Callable, List, Optional, Union, cast
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import (
@@ -13,6 +13,7 @@
 from datahub.ingestion.transformer.dataset_transformer import DatasetTagsTransformer
 from datahub.metadata.schema_classes import (
     GlobalTagsClass,
+    MetadataChangeProposalClass,
     TagAssociationClass,
     TagKeyClass,
 )
@@ -65,9 +66,13 @@ def transform_aspect(
             self.config, self.ctx.graph, entity_urn, out_global_tags_aspect
         )
 
-    def handle_end_of_stream(self) -> List[MetadataChangeProposalWrapper]:
+    def handle_end_of_stream(
+        self,
+    ) -> List[Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]]:
 
-        mcps: List[MetadataChangeProposalWrapper] = []
+        mcps: List[
+            Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]
+        ] = []
 
         logger.debug("Generating tags")
 
@@ -121,7 +126,6 @@ class PatternAddDatasetTags(AddDatasetTags):
     """Transformer that adds a specified set of tags to each dataset."""
 
     def __init__(self, config: PatternDatasetTagsConfig, ctx: PipelineContext):
-        config.tag_pattern.all
         tag_pattern = config.tag_pattern
         generic_config = AddDatasetTagsConfig(
             get_tags_to_add=lambda _: [

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

@@ -1,6 +1,6 @@
 import logging
 from abc import ABCMeta, abstractmethod
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
 
 import datahub.emitter.mce_builder as builder
 from datahub.emitter.aspect import ASPECT_MAP
@@ -28,7 +28,9 @@ def _update_work_unit_id(
 
 
 class HandleEndOfStreamTransformer:
-    def handle_end_of_stream(self) -> List[MetadataChangeProposalWrapper]:
+    def handle_end_of_stream(
+        self,
+    ) -> Sequence[Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]]:
         return []
 
 
@@ -206,15 +208,19 @@ def _handle_end_of_stream(
         ):
             return
 
-        mcps: List[MetadataChangeProposalWrapper] = self.handle_end_of_stream()
+        mcps: Sequence[
+            Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]
+        ] = self.handle_end_of_stream()
 
         for mcp in mcps:
-            if mcp.aspect is None or mcp.entityUrn is None:  # to silent the lint error
+            if (
+                mcp.aspect is None or mcp.aspectName is None or mcp.entityUrn is None
+            ):  # to silent the lint error
                 continue
 
             record_metadata = _update_work_unit_id(
                 envelope=envelope,
-                aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
+                aspect_name=mcp.aspectName,
                 urn=mcp.entityUrn,
             )
 

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

@@ -118,3 +118,8 @@ def aspect_name(self) -> str:
 class DatasetSchemaMetadataTransformer(DatasetTransformer, metaclass=ABCMeta):
     def aspect_name(self) -> str:
         return "schemaMetadata"
+
+
+class DatasetDataproductTransformer(DatasetTransformer, metaclass=ABCMeta):
+    def aspect_name(self) -> str:
+        return "dataProductProperties"