Skip to content

Commit 9ab65c3

Browse files
maximearmstrongmaximearmstrong
committed
Make PowerBITranslatorData a container class
1 parent e018d8f commit 9ab65c3

File tree

6 files changed

+63
-40
lines changed

6 files changed

+63
-40
lines changed

‎examples/docs_snippets/docs_snippets/integrations/power-bi/customize-power-bi-asset-defs.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
class MyCustomPowerBITranslator(DagsterPowerBITranslator):
2424
def get_asset_spec(self, data: PowerBIContentData) -> dg.AssetSpec:
2525
# We create the default asset spec using super()
26-
default_spec = super().get_asset_spec(data)
26+
default_spec = super().get_asset_spec(data) # type: ignore
2727
# We customize the team owner tag for all assets,
2828
# and we customize the asset key prefix only for dashboards.
2929
return default_spec.replace_attributes(

‎examples/project_atproto_dashboard/project_atproto_dashboard/dashboard/definitions.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ class CustomDagsterPowerBITranslator(DagsterPowerBITranslator):
2121
def get_report_spec(self, data: PowerBIContentData) -> dg.AssetSpec:
2222
return (
2323
super()
24-
.get_report_spec(data)
24+
.get_report_spec(data) # type: ignore
2525
.replace_attributes(
2626
group_name="reporting",
2727
)
@@ -33,7 +33,7 @@ def get_semantic_model_spec(self, data: PowerBIContentData) -> dg.AssetSpec:
3333
]
3434
return (
3535
super()
36-
.get_semantic_model_spec(data)
36+
.get_semantic_model_spec(data) # type: ignore
3737
.replace_attributes(
3838
group_name="reporting",
3939
deps=upsteam_table_deps,

‎python_modules/libraries/dagster-powerbi/dagster_powerbi/resource.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -448,8 +448,7 @@ def defs_from_state(self, state: PowerBIWorkspaceData) -> Definitions:
448448
all_external_asset_specs = [
449449
translator.get_asset_spec(
450450
PowerBITranslatorData(
451-
content_type=content.content_type,
452-
properties=content.properties,
451+
content_data=content,
453452
workspace_data=state,
454453
)
455454
)

‎python_modules/libraries/dagster-powerbi/dagster_powerbi/translator.py‎

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import re
22
import urllib.parse
33
from enum import Enum
4-
from typing import Any, Dict, List, Literal, Optional, Sequence
4+
from typing import Any, Dict, List, Literal, NamedTuple, Optional, Sequence
55

66
from dagster import (
77
UrlMetadataValue,
@@ -78,15 +78,22 @@ class PowerBIContentData:
7878
properties: Dict[str, Any]
7979

8080

81-
@whitelist_for_serdes
82-
@record
83-
class PowerBITranslatorData(PowerBIContentData):
81+
class PowerBITranslatorData(NamedTuple):
8482
"""A record representing a piece of content in PowerBI and the PowerBI workspace data.
8583
Includes the content's type and data as returned from the API.
8684
"""
8785

86+
content_data: "PowerBIContentData"
8887
workspace_data: "PowerBIWorkspaceData"
8988

89+
@property
90+
def content_type(self) -> PowerBIContentType:
91+
return self.content_data.content_type
92+
93+
@property
94+
def properties(self) -> Dict[str, Any]:
95+
return self.content_data.properties
96+
9097

9198
@whitelist_for_serdes
9299
@record
@@ -165,7 +172,7 @@ class DagsterPowerBITranslator:
165172
Subclass this class to implement custom logic for each type of PowerBI content.
166173
"""
167174

168-
def get_asset_spec(self, data: PowerBIContentData) -> AssetSpec:
175+
def get_asset_spec(self, data: PowerBITranslatorData) -> AssetSpec:
169176
data = check.inst(data, PowerBITranslatorData)
170177
if data.content_type == PowerBIContentType.DASHBOARD:
171178
return self.get_dashboard_spec(data)
@@ -182,11 +189,11 @@ def get_asset_spec(self, data: PowerBIContentData) -> AssetSpec:
182189
breaking_version="1.10",
183190
additional_warn_text="Use `DagsterPowerBITranslator.get_asset_spec(...).key` instead",
184191
)
185-
def get_dashboard_asset_key(self, data: PowerBIContentData) -> AssetKey:
192+
def get_dashboard_asset_key(self, data: PowerBITranslatorData) -> AssetKey:
186193
data = check.inst(data, PowerBITranslatorData)
187194
return self.get_dashboard_spec(data).key
188195

189-
def get_dashboard_spec(self, data: PowerBIContentData) -> AssetSpec:
196+
def get_dashboard_spec(self, data: PowerBITranslatorData) -> AssetSpec:
190197
data = check.inst(data, PowerBITranslatorData)
191198
dashboard_id = data.properties["id"]
192199
tile_report_ids = [
@@ -195,8 +202,7 @@ def get_dashboard_spec(self, data: PowerBIContentData) -> AssetSpec:
195202
report_keys = [
196203
self.get_report_spec(
197204
PowerBITranslatorData(
198-
content_type=data.workspace_data.reports_by_id[report_id].content_type,
199-
properties=data.workspace_data.reports_by_id[report_id].properties,
205+
content_data=data.workspace_data.reports_by_id[report_id],
200206
workspace_data=data.workspace_data,
201207
)
202208
).key
@@ -224,11 +230,11 @@ def get_dashboard_spec(self, data: PowerBIContentData) -> AssetSpec:
224230
breaking_version="1.10",
225231
additional_warn_text="Use `DagsterPowerBITranslator.get_asset_spec(...).key` instead",
226232
)
227-
def get_report_asset_key(self, data: PowerBIContentData) -> AssetKey:
233+
def get_report_asset_key(self, data: PowerBITranslatorData) -> AssetKey:
228234
data = check.inst(data, PowerBITranslatorData)
229235
return self.get_report_spec(data).key
230236

231-
def get_report_spec(self, data: PowerBIContentData) -> AssetSpec:
237+
def get_report_spec(self, data: PowerBITranslatorData) -> AssetSpec:
232238
data = check.inst(data, PowerBITranslatorData)
233239
report_id = data.properties["id"]
234240
dataset_id = data.properties.get("datasetId")
@@ -238,8 +244,7 @@ def get_report_spec(self, data: PowerBIContentData) -> AssetSpec:
238244
dataset_key = (
239245
self.get_semantic_model_spec(
240246
PowerBITranslatorData(
241-
content_type=dataset_data.content_type,
242-
properties=dataset_data.properties,
247+
content_data=dataset_data,
243248
workspace_data=data.workspace_data,
244249
)
245250
).key
@@ -266,19 +271,18 @@ def get_report_spec(self, data: PowerBIContentData) -> AssetSpec:
266271
breaking_version="1.10",
267272
additional_warn_text="Use `DagsterPowerBITranslator.get_asset_spec(...).key` instead",
268273
)
269-
def get_semantic_model_asset_key(self, data: PowerBIContentData) -> AssetKey:
274+
def get_semantic_model_asset_key(self, data: PowerBITranslatorData) -> AssetKey:
270275
data = check.inst(data, PowerBITranslatorData)
271276
return self.get_semantic_model_spec(data).key
272277

273-
def get_semantic_model_spec(self, data: PowerBIContentData) -> AssetSpec:
278+
def get_semantic_model_spec(self, data: PowerBITranslatorData) -> AssetSpec:
274279
data = check.inst(data, PowerBITranslatorData)
275280
dataset_id = data.properties["id"]
276281
source_ids = data.properties.get("sources", [])
277282
source_keys = [
278283
self.get_data_source_spec(
279284
PowerBITranslatorData(
280-
content_type=data.workspace_data.data_sources_by_id[source_id].content_type,
281-
properties=data.workspace_data.data_sources_by_id[source_id].properties,
285+
content_data=data.workspace_data.data_sources_by_id[source_id],
282286
workspace_data=data.workspace_data,
283287
)
284288
).key
@@ -328,11 +332,11 @@ def get_semantic_model_spec(self, data: PowerBIContentData) -> AssetSpec:
328332
breaking_version="1.10",
329333
additional_warn_text="Use `DagsterPowerBITranslator.get_asset_spec(...).key` instead",
330334
)
331-
def get_data_source_asset_key(self, data: PowerBIContentData) -> AssetKey:
335+
def get_data_source_asset_key(self, data: PowerBITranslatorData) -> AssetKey:
332336
data = check.inst(data, PowerBITranslatorData)
333337
return self.get_data_source_spec(data).key
334338

335-
def get_data_source_spec(self, data: PowerBIContentData) -> AssetSpec:
339+
def get_data_source_spec(self, data: PowerBITranslatorData) -> AssetSpec:
336340
data = check.inst(data, PowerBITranslatorData)
337341
connection_name = (
338342
data.properties["connectionDetails"].get("path")

‎python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_asset_specs.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from dagster import materialize
77
from dagster._config.field_utils import EnvVar
88
from dagster._core.code_pointer import CodePointer
9+
from dagster._core.definitions.asset_spec import AssetSpec
910
from dagster._core.definitions.assets import AssetsDefinition
1011
from dagster._core.definitions.decorators.asset_decorator import asset
1112
from dagster._core.definitions.definitions_class import Definitions
@@ -24,6 +25,7 @@
2425
from dagster_powerbi import PowerBIWorkspace
2526
from dagster_powerbi.assets import build_semantic_model_refresh_asset_definition
2627
from dagster_powerbi.resource import BASE_API_URL, PowerBIToken, load_powerbi_asset_specs
28+
from dagster_powerbi.translator import DagsterPowerBITranslator, PowerBIContentData
2729

2830
from dagster_powerbi_tests.conftest import SAMPLE_SEMANTIC_MODEL
2931

@@ -82,6 +84,31 @@ def test_translator_dashboard_spec(workspace_data_api_mocks: None, workspace_id:
8284
assert semantic_model_asset.key.path == ["semantic_model", "Sales_Returns_Sample_v201912"]
8385

8486

87+
class MyCustomTranslator(DagsterPowerBITranslator):
88+
def get_asset_spec(self, data: PowerBIContentData) -> AssetSpec:
89+
default_spec = super().get_asset_spec(data) # type: ignore
90+
return default_spec.replace_attributes(
91+
key=default_spec.key.with_prefix("prefix"),
92+
).merge_attributes(metadata={"custom": "metadata"})
93+
94+
95+
def test_translator_custom_metadata(workspace_data_api_mocks: None, workspace_id: str) -> None:
96+
fake_token = uuid.uuid4().hex
97+
resource = PowerBIWorkspace(
98+
credentials=PowerBIToken(api_token=fake_token),
99+
workspace_id=workspace_id,
100+
)
101+
all_asset_specs = load_powerbi_asset_specs(
102+
workspace=resource, dagster_powerbi_translator=MyCustomTranslator, use_workspace_scan=False
103+
)
104+
asset_spec = next(spec for spec in all_asset_specs)
105+
106+
assert "custom" in asset_spec.metadata
107+
assert asset_spec.metadata["custom"] == "metadata"
108+
assert asset_spec.key.path == ["prefix", "dashboard", "Sales_Returns_Sample_v201912"]
109+
assert "dagster/kind/powerbi" in asset_spec.tags
110+
111+
85112
@lazy_definitions
86113
def state_derived_defs_two_workspaces() -> Definitions:
87114
resource = PowerBIWorkspace(

‎python_modules/libraries/dagster-powerbi/dagster_powerbi_tests/test_translator.py‎

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@ def test_translator_dashboard_spec(workspace_data: PowerBIWorkspaceData) -> None
1818
translator = DagsterPowerBITranslator()
1919
asset_spec = translator.get_asset_spec(
2020
PowerBITranslatorData(
21-
content_type=dashboard.content_type,
22-
properties=dashboard.properties,
21+
content_data=dashboard,
2322
workspace_data=workspace_data,
2423
)
2524
)
@@ -46,8 +45,7 @@ def test_translator_report_spec(workspace_data: PowerBIWorkspaceData) -> None:
4645
translator = DagsterPowerBITranslator()
4746
asset_spec = translator.get_asset_spec(
4847
PowerBITranslatorData(
49-
content_type=report.content_type,
50-
properties=report.properties,
48+
content_data=report,
5149
workspace_data=workspace_data,
5250
)
5351
)
@@ -75,8 +73,7 @@ def test_translator_semantic_model(workspace_data: PowerBIWorkspaceData) -> None
7573
translator = DagsterPowerBITranslator()
7674
asset_spec = translator.get_asset_spec(
7775
PowerBITranslatorData(
78-
content_type=semantic_model.content_type,
79-
properties=semantic_model.properties,
76+
content_data=semantic_model,
8077
workspace_data=workspace_data,
8178
)
8279
)
@@ -116,8 +113,7 @@ def test_translator_semantic_model_many_tables(second_workspace_data: PowerBIWor
116113
translator = DagsterPowerBITranslator()
117114
asset_spec = translator.get_asset_spec(
118115
PowerBITranslatorData(
119-
content_type=semantic_model.content_type,
120-
properties=semantic_model.properties,
116+
content_data=semantic_model,
121117
workspace_data=second_workspace_data,
122118
)
123119
)
@@ -147,21 +143,19 @@ def test_translator_semantic_model_many_tables(second_workspace_data: PowerBIWor
147143

148144
class MyCustomTranslator(DagsterPowerBITranslator):
149145
def get_asset_spec(self, data: PowerBIContentData) -> AssetSpec:
150-
default_spec = super().get_asset_spec(data)
146+
default_spec = super().get_asset_spec(data) # type: ignore
151147
return default_spec.replace_attributes(
152148
key=default_spec.key.with_prefix("prefix"),
153-
metadata={**default_spec.metadata, "custom": "metadata"},
154-
)
149+
).merge_attributes(metadata={"custom": "metadata"})
155150

156151

157152
def test_translator_custom_metadata(workspace_data: PowerBIWorkspaceData) -> None:
158153
dashboard = next(iter(workspace_data.dashboards_by_id.values()))
159154

160155
translator = MyCustomTranslator()
161156
asset_spec = translator.get_asset_spec(
162-
PowerBITranslatorData(
163-
content_type=dashboard.content_type,
164-
properties=dashboard.properties,
157+
PowerBITranslatorData( # type: ignore
158+
content_data=dashboard,
165159
workspace_data=workspace_data,
166160
)
167161
)
@@ -191,8 +185,7 @@ def test_translator_report_spec_no_dataset(workspace_data: PowerBIWorkspaceData)
191185
translator = DagsterPowerBITranslator()
192186
asset_spec = translator.get_asset_spec(
193187
PowerBITranslatorData(
194-
content_type=report_no_dataset.content_type,
195-
properties=report_no_dataset.properties,
188+
content_data=report_no_dataset,
196189
workspace_data=workspace_data,
197190
)
198191
)

0 commit comments

Comments
 (0)