Skip to content

Commit e525615

Browse files
add the dbt node name
1 parent 86b532d commit e525615

File tree

10 files changed

+104
-45
lines changed

10 files changed

+104
-45
lines changed

sqlmesh/core/context.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1677,7 +1677,11 @@ def plan_builder(
16771677
end_override_per_model=max_interval_end_per_model,
16781678
console=self.console,
16791679
user_provided_flags=user_provided_flags,
1680-
selected_models=model_selector.expand_model_selections(select_models or "*"),
1680+
selected_models={
1681+
node_name
1682+
for model in model_selector.expand_model_selections(select_models or "*")
1683+
if (node_name := snapshots[model].node.node_name)
1684+
},
16811685
explain=explain or False,
16821686
ignore_cron=ignore_cron or False,
16831687
)

sqlmesh/core/node.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,7 @@ class _Node(PydanticModel):
199199
interval_unit_: t.Optional[IntervalUnit] = Field(alias="interval_unit", default=None)
200200
tags: t.List[str] = []
201201
stamp: t.Optional[str] = None
202+
node_name: t.Optional[str] = None # dbt node name
202203
_path: t.Optional[Path] = None
203204
_data_hash: t.Optional[str] = None
204205
_metadata_hash: t.Optional[str] = None

sqlmesh/core/scheduler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -812,7 +812,7 @@ def _run_or_audit(
812812
run_environment_statements=run_environment_statements,
813813
audit_only=audit_only,
814814
auto_restatement_triggers=auto_restatement_triggers,
815-
selected_models=selected_snapshots or {s.name for s in merged_intervals},
815+
selected_models={s.node.node_name for s in merged_intervals if s.node.node_name},
816816
)
817817

818818
return CompletionStatus.FAILURE if errors else CompletionStatus.SUCCESS

sqlmesh/dbt/adapter.py

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,6 @@ def graph(self) -> t.Any:
180180
}
181181
)
182182

183-
@property
184-
def selected_resources(self) -> t.List[str]:
185-
return []
186-
187183

188184
class ParsetimeAdapter(BaseAdapter):
189185
def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
@@ -284,13 +280,6 @@ def __init__(
284280
def graph(self) -> t.Any:
285281
return self.jinja_globals.get("flat_graph", super().graph)
286282

287-
@property
288-
def selected_resources(self) -> t.List[str]:
289-
selected_models = self.jinja_globals.get("selected_models")
290-
if selected_models:
291-
return [self._dbt_model_id(model) for model in sorted(selected_models)]
292-
return []
293-
294283
def get_relation(
295284
self, database: t.Optional[str], schema: str, identifier: str
296285
) -> t.Optional[BaseRelation]:
@@ -515,8 +504,3 @@ def _normalize(self, input_table: exp.Table) -> exp.Table:
515504
normalized_table.set("db", normalized_table.this)
516505
normalized_table.set("this", None)
517506
return normalized_table
518-
519-
def _dbt_model_id(self, sqlmesh_model_name: str) -> str:
520-
# Model prefix is needed to correspond to the key in the nodes within the dbt context variable
521-
parts = [part.strip('"') for part in sqlmesh_model_name.split(".")]
522-
return f"model.{parts[0]}.{parts[-1]}"

sqlmesh/dbt/builtin.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -540,7 +540,7 @@ def create_builtin_globals(
540540
"run_query": sql_execution.run_query,
541541
"statement": sql_execution.statement,
542542
"graph": adapter.graph,
543-
"selected_resources": adapter.selected_resources,
543+
"selected_resources": list(jinja_globals.get("selected_models") or []),
544544
}
545545
)
546546

sqlmesh/dbt/model.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -689,6 +689,7 @@ def to_sqlmesh(
689689
extract_dependencies_from_query=False,
690690
allow_partials=allow_partials,
691691
virtual_environment_mode=virtual_environment_mode,
692+
node_name=self.node_name,
692693
**optional_kwargs,
693694
**model_kwargs,
694695
)

sqlmesh/dbt/seed.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ def to_sqlmesh(
9292
audit_definitions=audit_definitions,
9393
virtual_environment_mode=virtual_environment_mode,
9494
start=self.start or context.sqlmesh_config.model_defaults.start,
95+
node_name=self.node_name,
9596
**kwargs,
9697
)
9798

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
"""Add 'node_name' property to node definition."""
2+
3+
4+
def migrate_schemas(state_sync, **kwargs): # type: ignore
5+
pass
6+
7+
8+
def migrate_rows(state_sync, **kwargs): # type: ignore
9+
pass

tests/dbt/test_model.py

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,3 +579,80 @@ def test_load_microbatch_with_ref_no_filter(
579579
context.render(microbatch_two_snapshot_fqn, start="2025-01-01", end="2025-01-10").sql()
580580
== 'SELECT "microbatch"."cola" AS "cola", "microbatch"."ds" AS "ds" FROM "local"."main"."microbatch" AS "microbatch"'
581581
)
582+
583+
584+
@pytest.mark.slow
585+
def test_node_name_populated_for_dbt_models(dbt_dummy_postgres_config: PostgresConfig) -> None:
586+
model_config = ModelConfig(
587+
name="test_model",
588+
package_name="test_package",
589+
sql="SELECT 1 as id",
590+
database="test_db",
591+
schema_="test_schema",
592+
alias="test_model",
593+
)
594+
595+
context = DbtContext()
596+
context.project_name = "test_project"
597+
context.target = dbt_dummy_postgres_config
598+
599+
# check after convert to SQLMesh model that node_name is populated correctly
600+
sqlmesh_model = model_config.to_sqlmesh(context)
601+
assert sqlmesh_model.node_name == "model.test_package.test_model"
602+
603+
604+
@pytest.mark.slow
605+
def test_load_model_dbt_node_name(tmp_path: Path) -> None:
606+
yaml = YAML()
607+
dbt_project_dir = tmp_path / "dbt"
608+
dbt_project_dir.mkdir()
609+
dbt_model_dir = dbt_project_dir / "models"
610+
dbt_model_dir.mkdir()
611+
612+
model_contents = "SELECT 1 as id, 'test' as name"
613+
model_file = dbt_model_dir / "simple_model.sql"
614+
with open(model_file, "w", encoding="utf-8") as f:
615+
f.write(model_contents)
616+
617+
dbt_project_config = {
618+
"name": "test_project",
619+
"version": "1.0.0",
620+
"config-version": 2,
621+
"profile": "test",
622+
"model-paths": ["models"],
623+
}
624+
dbt_project_file = dbt_project_dir / "dbt_project.yml"
625+
with open(dbt_project_file, "w", encoding="utf-8") as f:
626+
yaml.dump(dbt_project_config, f)
627+
628+
sqlmesh_config = {
629+
"model_defaults": {
630+
"start": "2025-01-01",
631+
}
632+
}
633+
sqlmesh_config_file = dbt_project_dir / "sqlmesh.yaml"
634+
with open(sqlmesh_config_file, "w", encoding="utf-8") as f:
635+
yaml.dump(sqlmesh_config, f)
636+
637+
dbt_data_dir = tmp_path / "dbt_data"
638+
dbt_data_dir.mkdir()
639+
dbt_data_file = dbt_data_dir / "local.db"
640+
dbt_profile_config = {
641+
"test": {
642+
"outputs": {"duckdb": {"type": "duckdb", "path": str(dbt_data_file)}},
643+
"target": "duckdb",
644+
}
645+
}
646+
db_profile_file = dbt_project_dir / "profiles.yml"
647+
with open(db_profile_file, "w", encoding="utf-8") as f:
648+
yaml.dump(dbt_profile_config, f)
649+
650+
context = Context(paths=dbt_project_dir)
651+
652+
# find the model by its sqlmesh fully qualified name
653+
model_fqn = '"local"."main"."simple_model"'
654+
assert model_fqn in context.snapshots
655+
656+
# Verify that node_name is the equivalent dbt one
657+
model = context.snapshots[model_fqn].model
658+
assert model.node_name == "model.test_project.simple_model"

tests/dbt/test_transformation.py

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
import typing as t
66
from pathlib import Path
77
from unittest.mock import patch
8-
from sqlmesh.dbt.adapter import RuntimeAdapter
9-
from sqlmesh.utils.jinja import JinjaMacroRegistry
108

119
from sqlmesh.dbt.util import DBT_VERSION
1210

@@ -2394,7 +2392,7 @@ def test_selected_resources_context_variable(
23942392
):
23952393
context = sushi_test_project.context
23962394

2397-
# should be empty list during parse time
2395+
# empty selected resources
23982396
direct_access = context.render("{{ selected_resources }}")
23992397
assert direct_access == "[]"
24002398

@@ -2420,32 +2418,16 @@ def test_selected_resources_context_variable(
24202418
result = context.render(test_condition)
24212419
assert result.strip() == "no_resources"
24222420

2423-
# Test 4: Test with runtime adapter (simulating runtime execution)
2424-
runtime_adapter = RuntimeAdapter(
2425-
engine_adapter=sushi_test_dbt_context.engine_adapter,
2426-
jinja_macros=JinjaMacroRegistry(),
2427-
jinja_globals={
2428-
"selected_models": {
2429-
'"jaffle_shop"."main"."customers"',
2430-
'"jaffle_shop"."main"."orders"',
2431-
'"jaffle_shop"."main"."items"',
2432-
},
2433-
},
2434-
)
2435-
2436-
# it should return correct selected resources in dbt format
2437-
selected_resources = runtime_adapter.selected_resources
2438-
assert len(selected_resources) == 3
2439-
assert "model.jaffle_shop.customers" in selected_resources
2440-
assert "model.jaffle_shop.orders" in selected_resources
2441-
assert "model.jaffle_shop.items" in selected_resources
2421+
# selected resources in dbt format
2422+
selected_resources = [
2423+
"model.jaffle_shop.customers",
2424+
"model.jaffle_shop.items",
2425+
"model.jaffle_shop.orders",
2426+
]
24422427

24432428
# check the jinja macros rendering
24442429
result = context.render("{{ selected_resources }}", selected_resources=selected_resources)
2445-
assert (
2446-
result
2447-
== "['model.jaffle_shop.customers', 'model.jaffle_shop.items', 'model.jaffle_shop.orders']"
2448-
)
2430+
assert result == selected_resources.__repr__()
24492431

24502432
result = context.render(test_jinja, selected_resources=selected_resources)
24512433
assert result.strip() == "3"

0 commit comments

Comments
 (0)