fix(ingestion/transformer): new transformer to clean user URN for DatasetUsageStatistics aspect

metadata-ingestion/docs/transformer/dataset_transformer.md

@@ -925,6 +925,24 @@ transformers:
       replacement: "sub"
 ```
 
+## Clean User URN in DatasetUsageStatistics Aspect
+### Config Details
+| Field                       | Required | Type    | Default       | Description                                 |
+|-----------------------------|----------|---------|---------------|---------------------------------------------|
+| `pattern_for_cleanup`                 | ✅         | list[string]    |    | List of suffix/prefix to remove from the Owner URN(s) |
+
+
+Matches against a User URN in DatasetUsageStatistics aspect and remove the matching part from it
+```yaml
+transformers:
+  - type: "pattern_cleanup_dataset_usage_user"
+    config:
+      pattern_for_cleanup:
+        - "ABCDEF"
+        - (?<=_)(\w+)
+```
+
+
 ## Simple Add Dataset domains 
 ### Config Details
 | Field              | Required | Type                   | Default       | Description                                                      |

metadata-ingestion/setup.py

@@ -706,6 +706,7 @@
         "simple_add_dataset_dataproduct = datahub.ingestion.transformer.add_dataset_dataproduct:SimpleAddDatasetDataProduct",
         "pattern_add_dataset_dataproduct = datahub.ingestion.transformer.add_dataset_dataproduct:PatternAddDatasetDataProduct",
         "replace_external_url = datahub.ingestion.transformer.replace_external_url:ReplaceExternalUrl",
+        "pattern_cleanup_dataset_usage_user = datahub.ingestion.transformer.pattern_cleanup_dataset_usage_user:PatternCleanupDatasetUsageUser",
     ],
     "datahub.ingestion.sink.plugins": [
         "file = datahub.ingestion.sink.file:FileSink",

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

@@ -123,3 +123,8 @@ def aspect_name(self) -> str:
 class DatasetDataproductTransformer(DatasetTransformer, metaclass=ABCMeta):
     def aspect_name(self) -> str:
         return "dataProductProperties"
+
+
+class DatasetUsageStatisticsTransformer(DatasetTransformer, metaclass=ABCMeta):
+    def aspect_name(self) -> str:
+        return "datasetUsageStatistics"

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

@@ -0,0 +1,67 @@
+import copy
+import re
+from typing import Any, Dict, List, Optional, cast
+
+from datahub.configuration.common import ConfigModel
+from datahub.emitter.mce_builder import Aspect
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.transformer.dataset_transformer import (
+    DatasetUsageStatisticsTransformer,
+)
+from datahub.metadata.schema_classes import DatasetUsageStatisticsClass
+
+_USER_URN_PREFIX: str = "urn:li:corpuser:"
+
+
+class PatternCleanupDatasetUsageUserConfig(ConfigModel):
+    pattern_for_cleanup: List[str]
+
+
+class PatternCleanupDatasetUsageUser(DatasetUsageStatisticsTransformer):
+    """Transformer that clean the user URN for DatasetUsageStatistics aspect."""
+
+    ctx: PipelineContext
+    config: PatternCleanupDatasetUsageUserConfig
+
+    def __init__(
+        self,
+        config: PatternCleanupDatasetUsageUserConfig,
+        ctx: PipelineContext,
+        **resolver_args: Dict[str, Any],
+    ):
+        super().__init__()
+        self.config = config
+        self.ctx = ctx
+        self.resolver_args = resolver_args
+
+    @classmethod
+    def create(
+        cls, config_dict: dict, ctx: PipelineContext
+    ) -> "PatternCleanupDatasetUsageUser":
+        config = PatternCleanupDatasetUsageUserConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    def transform_aspect(
+        self, entity_urn: str, aspect_name: str, aspect: Optional[Aspect]
+    ) -> Optional[Aspect]:
+        in_dataset_properties_aspect: DatasetUsageStatisticsClass = cast(
+            DatasetUsageStatisticsClass, aspect
+        )
+
+        if in_dataset_properties_aspect.userCounts is not None:
+            out_dataset_properties_aspect: DatasetUsageStatisticsClass = copy.deepcopy(
+                in_dataset_properties_aspect
+            )
+
+            if out_dataset_properties_aspect.userCounts is not None:
+                for user in out_dataset_properties_aspect.userCounts:
+                    user_id: str = user.user.split(_USER_URN_PREFIX)[1]
+                    for value in self.config.pattern_for_cleanup:
+                        cleaned_user_id = re.sub(value, "", user_id)
+                        user.user = _USER_URN_PREFIX + cleaned_user_id
+
+                return cast(Aspect, out_dataset_properties_aspect)
+            else:
+                return cast(Aspect, out_dataset_properties_aspect)
+        else:
+            return cast(Aspect, in_dataset_properties_aspect)

metadata-ingestion/tests/unit/test_transform_dataset.py

@@ -1,5 +1,6 @@
 import json
 import re
+from datetime import datetime, timezone
 from typing import (
     Any,
     Callable,
@@ -72,6 +73,9 @@
     ExtractOwnersFromTagsTransformer,
 )
 from datahub.ingestion.transformer.mark_dataset_status import MarkDatasetStatus
+from datahub.ingestion.transformer.pattern_cleanup_dataset_usage_user import (
+    PatternCleanupDatasetUsageUser,
+)
 from datahub.ingestion.transformer.pattern_cleanup_ownership import (
     PatternCleanUpOwnership,
 )
@@ -82,6 +86,7 @@
 from datahub.metadata.schema_classes import (
     BrowsePathsClass,
     DatasetPropertiesClass,
+    DatasetUserUsageCountsClass,
     GlobalTagsClass,
     MetadataChangeEventClass,
     OwnershipClass,
@@ -3291,3 +3296,165 @@ def test_replace_external_regex_replace_2(
         output[0].record.aspect.externalUrl
         == "https://test.com/test/looker-demo/blob/master/foo.view.lkml"
     )
+
+
+def test_pattern_cleanup_usage_statistics_user_1(
+    mock_datahub_graph,
+):
+    pipeline_context: PipelineContext = PipelineContext(
+        run_id="test_pattern_cleanup_usage_statistics_user"
+    )
+    pipeline_context.graph = mock_datahub_graph(DatahubClientConfig)
+
+    TS_1 = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc)
+
+    output = run_dataset_transformer_pipeline(
+        transformer_type=PatternCleanupDatasetUsageUser,
+        aspect=models.DatasetUsageStatisticsClass(
+            timestampMillis=int(TS_1.timestamp() * 1000),
+            userCounts=[
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("IAM:user1"),
+                    count=1,
+                    userEmail="user1@exaple.com",
+                ),
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("user2"),
+                    count=2,
+                    userEmail="user2@exaple.com",
+                ),
+            ],
+        ),
+        config={"pattern_for_cleanup": ["IAM:"]},
+        pipeline_context=pipeline_context,
+    )
+
+    expectedUsageStatistics = models.DatasetUsageStatisticsClass(
+        timestampMillis=int(TS_1.timestamp() * 1000),
+        userCounts=[
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("user1"),
+                count=1,
+                userEmail="user1@exaple.com",
+            ),
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("user2"),
+                count=2,
+                userEmail="user2@exaple.com",
+            ),
+        ],
+    )
+
+    assert len(output) == 2
+    assert output[0].record
+    assert output[0].record.aspect
+    assert len(output[0].record.aspect.userCounts) == 2
+    assert output[0].record.aspect.userCounts == expectedUsageStatistics.userCounts
+
+
+def test_pattern_cleanup_usage_statistics_user_2(
+    mock_datahub_graph,
+):
+    pipeline_context: PipelineContext = PipelineContext(
+        run_id="test_pattern_cleanup_usage_statistics_user"
+    )
+    pipeline_context.graph = mock_datahub_graph(DatahubClientConfig)
+
+    TS_1 = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc)
+
+    output = run_dataset_transformer_pipeline(
+        transformer_type=PatternCleanupDatasetUsageUser,
+        aspect=models.DatasetUsageStatisticsClass(
+            timestampMillis=int(TS_1.timestamp() * 1000),
+            userCounts=[
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("test_user_1"),
+                    count=1,
+                    userEmail="user1@exaple.com",
+                ),
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("test_user_2"),
+                    count=2,
+                    userEmail="user2@exaple.com",
+                ),
+            ],
+        ),
+        config={"pattern_for_cleanup": ["_user"]},
+        pipeline_context=pipeline_context,
+    )
+
+    expectedUsageStatistics = models.DatasetUsageStatisticsClass(
+        timestampMillis=int(TS_1.timestamp() * 1000),
+        userCounts=[
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("test_1"),
+                count=1,
+                userEmail="user1@exaple.com",
+            ),
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("test_2"),
+                count=2,
+                userEmail="user2@exaple.com",
+            ),
+        ],
+    )
+
+    assert len(output) == 2
+    assert output[0].record
+    assert output[0].record.aspect
+    assert len(output[0].record.aspect.userCounts) == 2
+    assert output[0].record.aspect.userCounts == expectedUsageStatistics.userCounts
+
+
+def test_pattern_cleanup_usage_statistics_user_3(
+    mock_datahub_graph,
+):
+    pipeline_context: PipelineContext = PipelineContext(
+        run_id="test_pattern_cleanup_usage_statistics_user"
+    )
+    pipeline_context.graph = mock_datahub_graph(DatahubClientConfig)
+
+    TS_1 = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc)
+
+    output = run_dataset_transformer_pipeline(
+        transformer_type=PatternCleanupDatasetUsageUser,
+        aspect=models.DatasetUsageStatisticsClass(
+            timestampMillis=int(TS_1.timestamp() * 1000),
+            userCounts=[
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("abc_user_1"),
+                    count=1,
+                    userEmail="user1@exaple.com",
+                ),
+                DatasetUserUsageCountsClass(
+                    user=builder.make_user_urn("xyz_user_2"),
+                    count=2,
+                    userEmail="user2@exaple.com",
+                ),
+            ],
+        ),
+        config={"pattern_for_cleanup": [r"_user_\d+"]},
+        pipeline_context=pipeline_context,
+    )
+
+    expectedUsageStatistics = models.DatasetUsageStatisticsClass(
+        timestampMillis=int(TS_1.timestamp() * 1000),
+        userCounts=[
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("abc"),
+                count=1,
+                userEmail="user1@exaple.com",
+            ),
+            DatasetUserUsageCountsClass(
+                user=builder.make_user_urn("xyz"),
+                count=2,
+                userEmail="user2@exaple.com",
+            ),
+        ],
+    )
+
+    assert len(output) == 2
+    assert output[0].record
+    assert output[0].record.aspect
+    assert len(output[0].record.aspect.userCounts) == 2
+    assert output[0].record.aspect.userCounts == expectedUsageStatistics.userCounts