From 79249dd73334eb992a690b80975250934a8de4df Mon Sep 17 00:00:00 2001 From: Themis Valtinos <73662635+themisvaltinos@users.noreply.github.com> Date: Tue, 9 Sep 2025 15:51:11 +0300 Subject: [PATCH] Fix: Set dbt execute variable for parse time and run time accordingly --- sqlmesh/dbt/builtin.py | 2 +- tests/core/test_integration.py | 9 ++-- tests/dbt/test_transformation.py | 49 ++++++++++++++++++- .../dbt/sushi_test/macros/execute_test.sql | 9 ++++ .../sushi_test/models/execute_test_model.sql | 2 + 5 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 tests/fixtures/dbt/sushi_test/macros/execute_test.sql create mode 100644 tests/fixtures/dbt/sushi_test/models/execute_test_model.sql diff --git a/sqlmesh/dbt/builtin.py b/sqlmesh/dbt/builtin.py index 8d48c4a77a..fe32ac2e24 100644 --- a/sqlmesh/dbt/builtin.py +++ b/sqlmesh/dbt/builtin.py @@ -533,7 +533,7 @@ def create_builtin_globals( builtin_globals.update( { "adapter": adapter, - "execute": True, + "execute": isinstance(adapter, RuntimeAdapter), "load_relation": adapter.load_relation, "store_result": sql_execution.store_result, "load_result": sql_execution.load_result, diff --git a/tests/core/test_integration.py b/tests/core/test_integration.py index ca0789d262..d36ec97f3b 100644 --- a/tests/core/test_integration.py +++ b/tests/core/test_integration.py @@ -2068,6 +2068,7 @@ def test_dbt_is_incremental_table_is_missing(sushi_test_dbt_context: Context): assert context.engine_adapter.table_exists(snapshot.table_name()) +@pytest.mark.xdist_group("dbt_manifest") def test_model_attr(sushi_test_dbt_context: Context, assert_exp_eq): context = sushi_test_dbt_context model = context.get_model("sushi.top_waiters") @@ -2075,14 +2076,14 @@ def test_model_attr(sushi_test_dbt_context: Context, assert_exp_eq): model.render_query(), """ SELECT - CAST("waiter_id" AS INT) AS "waiter_id", - CAST("revenue" AS DOUBLE) AS "revenue", + CAST("waiter_revenue_by_day_v2"."waiter_id" AS INT) AS "waiter_id", + CAST("waiter_revenue_by_day_v2"."revenue" AS DOUBLE) AS "revenue", 3 AS "model_columns" FROM "memory"."sushi"."waiter_revenue_by_day_v2" AS "waiter_revenue_by_day_v2" WHERE - "ds" = ( + "waiter_revenue_by_day_v2"."ds" = ( SELECT - MAX("ds") + MAX("waiter_revenue_by_day_v2"."ds") AS "_col_0" FROM "memory"."sushi"."waiter_revenue_by_day_v2" AS "waiter_revenue_by_day_v2" ) ORDER BY diff --git a/tests/dbt/test_transformation.py b/tests/dbt/test_transformation.py index 89e4bca154..80e12d6701 100644 --- a/tests/dbt/test_transformation.py +++ b/tests/dbt/test_transformation.py @@ -6,6 +6,7 @@ from pathlib import Path from unittest.mock import patch +from sqlmesh.dbt.adapter import ParsetimeAdapter, RuntimeAdapter from sqlmesh.dbt.util import DBT_VERSION import pytest @@ -43,7 +44,7 @@ OnAdditiveChange, ) from sqlmesh.core.state_sync.db.snapshot import _snapshot_to_json -from sqlmesh.dbt.builtin import _relation_info_to_relation, Config +from sqlmesh.dbt.builtin import _relation_info_to_relation, Config, create_builtin_globals from sqlmesh.dbt.common import Dependencies from sqlmesh.dbt.column import ( ColumnConfig, @@ -64,7 +65,7 @@ ) from sqlmesh.dbt.test import TestConfig from sqlmesh.utils.errors import ConfigError, MacroEvalError, SQLMeshError -from sqlmesh.utils.jinja import MacroReference +from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroReference pytestmark = [pytest.mark.dbt, pytest.mark.slow] @@ -2352,3 +2353,47 @@ def test_dynamic_var_names_in_macro(sushi_test_project: Project): ) converted_model = model_config.to_sqlmesh(context) assert "dynamic_test_var" in converted_model.jinja_macros.global_objs["vars"] # type: ignore + + +@pytest.mark.xdist_group("dbt_manifest") +def test_execute_variable_parse_vs_runtime(sushi_test_dbt_context: Context): + execute_model = sushi_test_dbt_context.get_model('"memory"."sushi"."execute_test_model"') + parse_time_query = execute_model.render_query() + if parse_time_query: + parse_sql = parse_time_query.sql() + # should contain parse-time placeholders + assert "parse_time_placeholder" in parse_sql or "parse_time_context" in parse_sql + + runtime_query = execute_model.render_query_or_raise( + engine_adapter=sushi_test_dbt_context.engine_adapter + ) + runtime_sql = runtime_query.sql() + # At runtime, the macro should have executed the query and returned the result + assert "1" in runtime_sql or "runtime_context" in runtime_sql + + +@pytest.mark.xdist_group("dbt_manifest") +def test_execute_globals(mocker: MockerFixture): + # No engine adapter shoulde create parse adapter + parse_time_globals = create_builtin_globals( + jinja_macros=JinjaMacroRegistry(), + jinja_globals={}, + engine_adapter=None, + ) + + assert parse_time_globals["execute"] is False + assert parse_time_globals["flags"].WHICH == "parse" + assert isinstance(parse_time_globals["adapter"], ParsetimeAdapter) + + mock_engine_adapter = mocker.Mock() + + # Runtime globals should have execute=True and RuntimeAdapter + runtime_globals = create_builtin_globals( + jinja_macros=JinjaMacroRegistry(), + jinja_globals={}, + engine_adapter=mock_engine_adapter, + ) + + assert runtime_globals["execute"] is True + assert runtime_globals["flags"].WHICH == "run" + assert isinstance(runtime_globals["adapter"], RuntimeAdapter) diff --git a/tests/fixtures/dbt/sushi_test/macros/execute_test.sql b/tests/fixtures/dbt/sushi_test/macros/execute_test.sql new file mode 100644 index 0000000000..b587a13179 --- /dev/null +++ b/tests/fixtures/dbt/sushi_test/macros/execute_test.sql @@ -0,0 +1,9 @@ +{% macro runtime_sql() %} + {% if execute %} + {% set result = run_query("SELECT 1 as test_col") %} + {% set test_value = result.columns[0][0] %} + {{ return(test_value) }} + {% else %} + {{ return("parse_time_placeholder") }} + {% endif %} +{% endmacro %} diff --git a/tests/fixtures/dbt/sushi_test/models/execute_test_model.sql b/tests/fixtures/dbt/sushi_test/models/execute_test_model.sql new file mode 100644 index 0000000000..cff0e744ac --- /dev/null +++ b/tests/fixtures/dbt/sushi_test/models/execute_test_model.sql @@ -0,0 +1,2 @@ +SELECT + '{{ runtime_sql() }}' as runtime_result \ No newline at end of file