From e58d236ae645db54b5778ef69dddfee6a1ac516e Mon Sep 17 00:00:00 2001
From: Maxime Armstrong <maxime@retain.ai>
Date: Fri, 27 Dec 2024 10:57:39 -0500
Subject: [PATCH] [dagster-powerbi] Use Power BI translator instance in spec
 loader and state-backed defs

---
 .../power-bi/customize-power-bi-asset-defs.py     |  2 +-
 .../dagster-powerbi/dagster_powerbi/resource.py   | 15 ++++++++-------
 .../dagster_powerbi_tests/test_asset_specs.py     |  4 +++-
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/examples/docs_snippets/docs_snippets/integrations/power-bi/customize-power-bi-asset-defs.py b/examples/docs_snippets/docs_snippets/integrations/power-bi/customize-power-bi-asset-defs.py
index 091121656e147..f9ee870bfc334 100644
--- a/examples/docs_snippets/docs_snippets/integrations/power-bi/customize-power-bi-asset-defs.py
+++ b/examples/docs_snippets/docs_snippets/integrations/power-bi/customize-power-bi-asset-defs.py
@@ -37,7 +37,7 @@ def get_asset_spec(self, data: PowerBIContentData) -> dg.AssetSpec:
 
 
 power_bi_specs = load_powerbi_asset_specs(
-    power_bi_workspace, dagster_powerbi_translator=MyCustomPowerBITranslator
+    power_bi_workspace, dagster_powerbi_translator=MyCustomPowerBITranslator  # type: ignore
 )
 defs = dg.Definitions(
     assets=[*power_bi_specs], resources={"power_bi": power_bi_workspace}
diff --git a/python_modules/libraries/dagster-powerbi/dagster_powerbi/resource.py b/python_modules/libraries/dagster-powerbi/dagster_powerbi/resource.py
index 16d5f7bcaeb58..473b6fb0de6ab 100644
--- a/python_modules/libraries/dagster-powerbi/dagster_powerbi/resource.py
+++ b/python_modules/libraries/dagster-powerbi/dagster_powerbi/resource.py
@@ -385,7 +385,7 @@ def build_defs(
                 if PowerBITagSet.extract(spec.tags).asset_type == "semantic_model"
                 else spec
                 for spec in load_powerbi_asset_specs(
-                    self, dagster_powerbi_translator, use_workspace_scan=False
+                    self, dagster_powerbi_translator(), use_workspace_scan=False
                 )
             ],
             resources={resource_key: self},
@@ -395,13 +395,16 @@ def build_defs(
 @experimental
 def load_powerbi_asset_specs(
     workspace: PowerBIWorkspace,
-    dagster_powerbi_translator: Type[DagsterPowerBITranslator] = DagsterPowerBITranslator,
+    dagster_powerbi_translator: Optional[DagsterPowerBITranslator] = None,
     use_workspace_scan: bool = True,
 ) -> Sequence[AssetSpec]:
     """Returns a list of AssetSpecs representing the Power BI content in the workspace.
 
     Args:
         workspace (PowerBIWorkspace): The Power BI workspace to load assets from.
+        dagster_powerbi_translator (Optional[DagsterPowerBITranslator]): The translator to use
+            to convert Power BI content into :py:class:`dagster.AssetSpec`.
+            Defaults to :py:class:`DagsterPowerBITranslator`.
         use_workspace_scan (bool): Whether to scan the entire workspace using admin APIs
             at once to get all content. Defaults to True.
 
@@ -412,7 +415,7 @@ def load_powerbi_asset_specs(
         return check.is_list(
             PowerBIWorkspaceDefsLoader(
                 workspace=initialized_workspace,
-                translator_cls=dagster_powerbi_translator,
+                translator=dagster_powerbi_translator or DagsterPowerBITranslator(),
                 use_workspace_scan=use_workspace_scan,
             )
             .build_defs()
@@ -424,7 +427,7 @@ def load_powerbi_asset_specs(
 @dataclass
 class PowerBIWorkspaceDefsLoader(StateBackedDefinitionsLoader[PowerBIWorkspaceData]):
     workspace: PowerBIWorkspace
-    translator_cls: Type[DagsterPowerBITranslator]
+    translator: DagsterPowerBITranslator
     use_workspace_scan: bool
 
     @property
@@ -438,15 +441,13 @@ def fetch_state(self) -> PowerBIWorkspaceData:
             )
 
     def defs_from_state(self, state: PowerBIWorkspaceData) -> Definitions:
-        translator = self.translator_cls()
-
         all_external_data = [
             *state.dashboards_by_id.values(),
             *state.reports_by_id.values(),
             *state.semantic_models_by_id.values(),
         ]
         all_external_asset_specs = [
-            translator.get_asset_spec(
+            self.translator.get_asset_spec(
                 PowerBITranslatorData(
                     content_data=content,
                     workspace_data=state,
diff --git a/python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_asset_specs.py b/python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_asset_specs.py
index 555288ebf7498..32196a361cebd 100644
--- a/python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_asset_specs.py
+++ b/python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_asset_specs.py
@@ -99,7 +99,9 @@ def test_translator_custom_metadata(workspace_data_api_mocks: None, workspace_id
         workspace_id=workspace_id,
     )
     all_asset_specs = load_powerbi_asset_specs(
-        workspace=resource, dagster_powerbi_translator=MyCustomTranslator, use_workspace_scan=False
+        workspace=resource,
+        dagster_powerbi_translator=MyCustomTranslator(),
+        use_workspace_scan=False,
     )
     asset_spec = next(spec for spec in all_asset_specs)