Update get_asset_key implementation

dagster-io · Nov 20, 2024 · c274578 · c274578
1 parent 4498d11
commit c274578
Showing 1 changed file with 4 additions and 27 deletions.
diff --git a/python_modules/libraries/dagster-tableau/dagster_tableau/translator.py b/python_modules/libraries/dagster-tableau/dagster_tableau/translator.py
@@ -111,14 +111,7 @@ def workspace_data(self) -> TableauWorkspaceData:
         return self._context
 
     def get_asset_key(self, data: TableauContentData) -> AssetKey:
-        if data.content_type == TableauContentType.SHEET:
-            return self.get_sheet_asset_key(data)
-        elif data.content_type == TableauContentType.DASHBOARD:
-            return self.get_dashboard_asset_key(data)
-        elif data.content_type == TableauContentType.DATA_SOURCE:
-            return self.get_data_source_asset_key(data)
-        else:
-            check.assert_never(data.content_type)
+        return self.get_asset_spec(data).key
 
     def get_asset_spec(self, data: TableauContentData) -> AssetSpec:
         if data.content_type == TableauContentType.SHEET:
@@ -131,15 +124,7 @@ def get_asset_spec(self, data: TableauContentData) -> AssetSpec:
             check.assert_never(data.content_type)
 
     def get_sheet_asset_key(self, data: TableauContentData) -> AssetKey:
-        workbook_id = data.properties["workbook"]["luid"]
-        workbook_data = self.workspace_data.workbooks_by_id[workbook_id]
-        return AssetKey(
-            [
-                _coerce_input_to_valid_name(workbook_data.properties["name"]),
-                "sheet",
-                _coerce_input_to_valid_name(data.properties["name"]),
-            ]
-        )
+        return self.get_sheet_spec(data).key
 
     def get_sheet_spec(self, data: TableauContentData) -> AssetSpec:
         sheet_embedded_data_sources = data.properties.get("parentEmbeddedDatasources", [])
@@ -176,15 +161,7 @@ def get_sheet_spec(self, data: TableauContentData) -> AssetSpec:
         )
 
     def get_dashboard_asset_key(self, data: TableauContentData) -> AssetKey:
-        workbook_id = data.properties["workbook"]["luid"]
-        workbook_data = self.workspace_data.workbooks_by_id[workbook_id]
-        return AssetKey(
-            [
-                _coerce_input_to_valid_name(workbook_data.properties["name"]),
-                "dashboard",
-                _coerce_input_to_valid_name(data.properties["name"]),
-            ]
-        )
+        return self.get_dashboard_spec(data).key
 
     def get_dashboard_spec(self, data: TableauContentData) -> AssetSpec:
         dashboard_upstream_sheets = data.properties.get("sheets", [])
@@ -217,7 +194,7 @@ def get_dashboard_spec(self, data: TableauContentData) -> AssetSpec:
         )
 
     def get_data_source_asset_key(self, data: TableauContentData) -> AssetKey:
-        return AssetKey([_coerce_input_to_valid_name(data.properties["name"])])
+        return self.get_data_source_spec(data).key
 
     def get_data_source_spec(self, data: TableauContentData) -> AssetSpec:
         return AssetSpec(