feat(ingest/looker): cleanup usage generation code (datahub-project#1…

…0153)
shubhamjagtap639 · Apr 3, 2024 · 786c776 · 786c776
1 parent 3c7c3ec
commit 786c776
Showing 1 changed file with 16 additions and 29 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_usage.py
@@ -37,7 +37,6 @@
 )
 from datahub.metadata.schema_classes import (
     CalendarIntervalClass,
-    ChangeTypeClass,
     ChartUsageStatisticsClass,
     ChartUserUsageCountsClass,
     DashboardUsageStatisticsClass,
@@ -207,7 +206,7 @@ def get_entity_stat_key(self, row: Dict) -> Tuple[str, str]:
         pass
 
     @abstractmethod
-    def _get_mcp_attributes(self, model: ModelForUsage) -> Dict:
+    def _get_urn(self, model: ModelForUsage) -> str:
         pass
 
     @abstractmethod
@@ -228,8 +227,8 @@ def create_mcp(
         self, model: ModelForUsage, aspect: Aspect
     ) -> MetadataChangeProposalWrapper:
         return MetadataChangeProposalWrapper(
+            entityUrn=self._get_urn(model=model),
             aspect=aspect,
-            **self._get_mcp_attributes(model=model),
         )
 
     def _round_time(self, date_time: str) -> int:
@@ -450,20 +449,14 @@ def get_entity_stat_key(self, row: Dict) -> Tuple[str, str]:
             row[HistoryViewField.HISTORY_CREATED_DATE],
         )
 
-    def _get_mcp_attributes(self, model: ModelForUsage) -> Dict:
-        dashboard: Dashboard = cast(Dashboard, model)
-        if dashboard is None or dashboard.id is None:  # to pass mypy lint
-            return {}
+    def _get_urn(self, model: ModelForUsage) -> str:
+        assert isinstance(model, LookerDashboardForUsage)
+        assert model.id is not None
 
-        return {
-            "entityUrn": builder.make_dashboard_urn(
-                self.config.platform_name,
-                looker_common.get_urn_looker_dashboard_id(dashboard.id),
-            ),
-            "entityType": "dashboard",
-            "changeType": ChangeTypeClass.UPSERT,
-            "aspectName": "dashboardUsageStatistics",
-        }
+        return builder.make_dashboard_urn(
+            self.config.platform_name,
+            looker_common.get_urn_looker_dashboard_id(model.id),
+        )
 
     def to_entity_absolute_stat_aspect(
         self, looker_object: ModelForUsage
@@ -569,20 +562,14 @@ def get_entity_stat_key(self, row: Dict) -> Tuple[str, str]:
             row[HistoryViewField.HISTORY_CREATED_DATE],
         )
 
-    def _get_mcp_attributes(self, model: ModelForUsage) -> Dict:
-        look: LookerChartForUsage = cast(LookerChartForUsage, model)
-        if look is None or look.id is None:
-            return {}
+    def _get_urn(self, model: ModelForUsage) -> str:
+        assert isinstance(model, LookerChartForUsage)
+        assert model.id is not None
 
-        return {
-            "entityUrn": builder.make_chart_urn(
-                self.config.platform_name,
-                looker_common.get_urn_looker_element_id(str(look.id)),
-            ),
-            "entityType": "chart",
-            "changeType": ChangeTypeClass.UPSERT,
-            "aspectName": "chartUsageStatistics",
-        }
+        return builder.make_chart_urn(
+            self.config.platform_name,
+            looker_common.get_urn_looker_element_id(str(model.id)),
+        )
 
     def to_entity_absolute_stat_aspect(
         self, looker_object: ModelForUsage