Skip to content

Commit 7e7199f

Browse files
themisvaltinosthemisvaltinos
committed
refactor graph to the manifest; make test more deterministic
1 parent 69d658b commit 7e7199f

File tree

4 files changed

+62
-47
lines changed

4 files changed

+62
-47
lines changed

‎sqlmesh/dbt/context.py‎

Lines changed: 1 addition & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
from sqlmesh.dbt.manifest import ManifestHelper
1212
from sqlmesh.dbt.target import TargetConfig
1313
from sqlmesh.utils import AttributeDict
14-
from sqlmesh.utils.conversions import serializable
1514
from sqlmesh.utils.errors import ConfigError, SQLMeshError
1615
from sqlmesh.utils.jinja import (
1716
JinjaGlobalAttribute,
@@ -196,49 +195,6 @@ def refs(self) -> t.Dict[str, t.Union[ModelConfig, SeedConfig]]:
196195
self._refs[f"{config_name}_v{model.version}"] = model
197196
return self._refs
198197

199-
@property
200-
def flat_graph(self) -> t.Dict[str, t.Any]:
201-
if self._manifest is None:
202-
return {
203-
"exposures": {},
204-
"groups": {},
205-
"metrics": {},
206-
"nodes": {},
207-
"sources": {},
208-
"semantic_models": {},
209-
"saved_queries": {},
210-
}
211-
212-
manifest = self._manifest._manifest
213-
return {
214-
"exposures": {
215-
k: serializable(v.to_dict(omit_none=False))
216-
for k, v in getattr(manifest, "exposures", {}).items()
217-
},
218-
"groups": {
219-
k: serializable(v.to_dict(omit_none=False))
220-
for k, v in getattr(manifest, "groups", {}).items()
221-
},
222-
"metrics": {
223-
k: serializable(v.to_dict(omit_none=False))
224-
for k, v in getattr(manifest, "metrics", {}).items()
225-
},
226-
"nodes": {
227-
k: serializable(v.to_dict(omit_none=False)) for k, v in manifest.nodes.items()
228-
},
229-
"sources": {
230-
k: serializable(v.to_dict(omit_none=False)) for k, v in manifest.sources.items()
231-
},
232-
"semantic_models": {
233-
k: serializable(v.to_dict(omit_none=False))
234-
for k, v in getattr(manifest, "semantic_models", {}).items()
235-
},
236-
"saved_queries": {
237-
k: serializable(v.to_dict(omit_none=False))
238-
for k, v in getattr(manifest, "saved_queries", {}).items()
239-
},
240-
}
241-
242198
@property
243199
def target(self) -> TargetConfig:
244200
if not self._target:
@@ -288,7 +244,7 @@ def jinja_globals(self) -> t.Dict[str, JinjaGlobalAttribute]:
288244
output["dialect"] = self.sqlmesh_config.dialect
289245
# Pass flat graph structure like dbt
290246
if self._manifest is not None:
291-
output["flat_graph"] = AttributeDict(self.flat_graph)
247+
output["flat_graph"] = AttributeDict(self.manifest.flat_graph)
292248
return output
293249

294250
def context_for_dependencies(self, dependencies: Dependencies) -> DbtContext:

‎sqlmesh/dbt/manifest.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212

1313
from dbt import constants as dbt_constants, flags
1414

15+
from sqlmesh.utils.conversions import serializable
16+
1517
# Override the file name to prevent dbt commands from invalidating the cache.
1618
dbt_constants.PARTIAL_PARSE_FILE_NAME = "sqlmesh_partial_parse.msgpack"
1719

@@ -155,6 +157,38 @@ def all_macros(self) -> t.Dict[str, t.Dict[str, MacroInfo]]:
155157
result[package_name][macro_name] = macro_config.info
156158
return result
157159

160+
@property
161+
def flat_graph(self) -> t.Dict[str, t.Any]:
162+
return {
163+
"exposures": {
164+
k: serializable(v.to_dict(omit_none=False))
165+
for k, v in getattr(self._manifest, "exposures", {}).items()
166+
},
167+
"groups": {
168+
k: serializable(v.to_dict(omit_none=False))
169+
for k, v in getattr(self._manifest, "groups", {}).items()
170+
},
171+
"metrics": {
172+
k: serializable(v.to_dict(omit_none=False))
173+
for k, v in getattr(self._manifest, "metrics", {}).items()
174+
},
175+
"nodes": {
176+
k: serializable(v.to_dict(omit_none=False)) for k, v in self._manifest.nodes.items()
177+
},
178+
"sources": {
179+
k: serializable(v.to_dict(omit_none=False))
180+
for k, v in self._manifest.sources.items()
181+
},
182+
"semantic_models": {
183+
k: serializable(v.to_dict(omit_none=False))
184+
for k, v in getattr(self._manifest, "semantic_models", {}).items()
185+
},
186+
"saved_queries": {
187+
k: serializable(v.to_dict(omit_none=False))
188+
for k, v in getattr(self._manifest, "saved_queries", {}).items()
189+
},
190+
}
191+
158192
def _load_all(self) -> None:
159193
if self._is_loaded:
160194
return

‎sqlmesh/utils/jinja.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,9 @@ def add_globals(self, globals: t.Dict[str, JinjaGlobalAttribute]) -> None:
363363
Args:
364364
globals: The global objects that should be added.
365365
"""
366+
# Keep the registry lightweight when the graph is not needed
367+
if not "graph" in self.packages:
368+
globals.pop("flat_graph", None)
366369
self.global_objs.update(**self._validate_global_objs(globals))
367370

368371
def build_macro(self, reference: MacroReference, **kwargs: t.Any) -> t.Optional[t.Callable]:

‎tests/dbt/test_transformation.py‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1641,10 +1641,32 @@ def test_on_run_start_end():
16411641
"CREATE OR REPLACE TABLE schema_table_snapshots__dev AS SELECT 'snapshots__dev' AS schema",
16421642
"CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
16431643
"DROP TABLE to_be_executed_last",
1644-
"CREATE OR REPLACE TABLE graph_table AS SELECT 'model.sushi.simple_model_a' AS unique_id, 'table' AS materialized UNION ALL SELECT 'model.sushi.waiters' AS unique_id, 'ephemeral' AS materialized UNION ALL SELECT 'model.sushi.simple_model_b' AS unique_id, 'table' AS materialized UNION ALL SELECT 'model.sushi.waiter_as_customer_by_day' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.top_waiters' AS unique_id, 'view' AS materialized UNION ALL SELECT 'model.customers.customers' AS unique_id, 'view' AS materialized UNION ALL SELECT 'model.customers.customer_revenue_by_day' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.waiter_revenue_by_day.v1' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.waiter_revenue_by_day.v2' AS unique_id, 'incremental' AS materialized",
16451644
]
1645+
assert sorted(rendered_after_all[:-1]) == sorted(expected_statements)
16461646

1647-
assert sorted(rendered_after_all) == sorted(expected_statements)
1647+
# Assert the models with their materialisations are present in the rendered graph_table statement
1648+
graph_table_stmt = rendered_after_all[-1]
1649+
assert "'model.sushi.simple_model_a' AS unique_id, 'table' AS materialized" in graph_table_stmt
1650+
assert "'model.sushi.waiters' AS unique_id, 'ephemeral' AS materialized" in graph_table_stmt
1651+
assert "'model.sushi.simple_model_b' AS unique_id, 'table' AS materialized" in graph_table_stmt
1652+
assert (
1653+
"'model.sushi.waiter_as_customer_by_day' AS unique_id, 'incremental' AS materialized"
1654+
in graph_table_stmt
1655+
)
1656+
assert "'model.sushi.top_waiters' AS unique_id, 'view' AS materialized" in graph_table_stmt
1657+
assert "'model.customers.customers' AS unique_id, 'view' AS materialized" in graph_table_stmt
1658+
assert (
1659+
"'model.customers.customer_revenue_by_day' AS unique_id, 'incremental' AS materialized"
1660+
in graph_table_stmt
1661+
)
1662+
assert (
1663+
"'model.sushi.waiter_revenue_by_day.v1' AS unique_id, 'incremental' AS materialized"
1664+
in graph_table_stmt
1665+
)
1666+
assert (
1667+
"'model.sushi.waiter_revenue_by_day.v2' AS unique_id, 'incremental' AS materialized"
1668+
in graph_table_stmt
1669+
)
16481670

16491671
# Nested dbt_packages on run start / on run end
16501672
packaged_environment_statements = sushi_context._environment_statements[1]

0 commit comments

Comments
 (0)