Repository: TobikoData/sqlmesh
Pull request #3074 (2024-08-30 18:24:38)

Fix: Prepopulate the model's depends on in-memory cache for converted dbt models
With this update, SQLMesh will no longer attempt to render queries of dbt models to extract additional dependencies that were not expressed using the `ref` macro.

As a result, this update speeds up the loading of dbt projects, as we no longer rely on unoptimized rendered queries to construct the set of dependencies needed to build the model DAG.


Patch:
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 65f94b60..7e6f0489 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1455,11 +1455,12 @@ class EngineAdapter:
                 )
             row_value_check = exp.or_(*row_check_conditions)
             unique_key_conditions = []
-            for col in unique_key:
-                t_col = col.copy()
-                t_col.this.set("this", f"t_{col.name}")
+            for key in unique_key:
+                t_key = key.copy()
+                for col in t_key.find_all(exp.Column):
+                    col.this.set("this", f"t_{col.name}")
                 unique_key_conditions.extend(
-                    [t_col.is_(exp.Null()).not_(), col.is_(exp.Null()).not_()]
+                    [t_key.is_(exp.Null()).not_(), key.is_(exp.Null()).not_()]
                 )
             unique_key_check = exp.and_(*unique_key_conditions)
             # unique_key_check is saying "if the row is updated"
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index aec5d703..526ad9af 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -72,7 +72,7 @@ from sqlmesh.core.state_sync.common import transactional
 from sqlmesh.utils import major_minor, random_id, unique
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import TimeLike, now, now_timestamp, time_like_to_str, to_timestamp
-from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.errors import ConflictingPlanError, SQLMeshError
 
 logger = logging.getLogger(__name__)
 
@@ -286,7 +286,7 @@ class EngineAdapterStateSync(StateSync):
             }
             if not existing_environment.expired:
                 if environment.previous_plan_id != existing_environment.plan_id:
-                    raise SQLMeshError(
+                    raise ConflictingPlanError(
                         f"Plan '{environment.plan_id}' is no longer valid for the target environment '{environment.name}'. "
                         f"Expected previous plan ID: '{environment.previous_plan_id}', actual previous plan ID: '{existing_environment.plan_id}'. "
                         "Please recreate the plan and try again"
diff --git a/sqlmesh/dbt/model.py b/sqlmesh/dbt/model.py
index 2cf0c412..2f1b2428 100644
--- a/sqlmesh/dbt/model.py
+++ b/sqlmesh/dbt/model.py
@@ -422,7 +422,7 @@ class ModelConfig(BaseModelConfig):
                     "target_lag": self.target_lag,
                 }
 
-        return create_sql_model(
+        model = create_sql_model(
             self.canonical_name(context),
             query,
             dialect=model_dialect,
@@ -431,6 +431,12 @@ class ModelConfig(BaseModelConfig):
             **optional_kwargs,
             **model_kwargs,
         )
+        # Prepopulate the _full_depends_on cache with dependencies sourced directly from the manifest.
+        # This ensures that we bypass query rendering that would otherwise be required to extract additional
+        # dependencies from the model's SQL.
+        # Note: any table dependencies that are not referenced using the `ref` macro will not be included.
+        model._full_depends_on = model.depends_on_
+        return model
 
     def _dbt_max_partition_blob(self) -> t.Optional[str]:
         """Returns a SQL blob which declares the _dbt_max_partition variable. Only applicable to BigQuery."""
diff --git a/sqlmesh/integrations/github/cicd/command.py b/sqlmesh/integrations/github/cicd/command.py
index e5fc207b..1a4ac413 100644
--- a/sqlmesh/integrations/github/cicd/command.py
+++ b/sqlmesh/integrations/github/cicd/command.py
@@ -12,7 +12,7 @@ from sqlmesh.integrations.github.cicd.controller import (
     GithubController,
     TestFailure,
 )
-from sqlmesh.utils.errors import CICDBotError, PlanError
+from sqlmesh.utils.errors import CICDBotError, ConflictingPlanError, PlanError
 
 logger = logging.getLogger(__name__)
 
@@ -147,11 +147,23 @@ def _deploy_production(controller: GithubController) -> bool:
         controller.try_merge_pr()
         controller.try_invalidate_pr_environment()
         return True
+    except ConflictingPlanError as e:
+        controller.update_prod_environment_check(
+            status=GithubCheckStatus.COMPLETED,
+            conclusion=GithubCheckConclusion.SKIPPED,
+            skip_reason=str(e),
+        )
+        return False
     except PlanError:
         controller.update_prod_environment_check(
             status=GithubCheckStatus.COMPLETED, conclusion=GithubCheckConclusion.ACTION_REQUIRED
         )
         return False
+    except Exception:
+        controller.update_prod_environment_check(
+            status=GithubCheckStatus.COMPLETED, conclusion=GithubCheckConclusion.FAILURE
+        )
+        return False
 
 
 @github.command()
diff --git a/sqlmesh/utils/errors.py b/sqlmesh/utils/errors.py
index d2267aa6..010e6ffc 100644
--- a/sqlmesh/utils/errors.py
+++ b/sqlmesh/utils/errors.py
@@ -47,6 +47,10 @@ class UncategorizedPlanError(PlanError):
     pass
 
 
+class ConflictingPlanError(PlanError):
+    pass
+
+
 class MissingContextException(Exception):
     pass
 


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index 42d69c16..b95d36e2 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -1719,6 +1719,198 @@ SELECT CAST("id" AS INT) AS "id", CAST("name" AS VARCHAR) AS "name", CAST("price
     )
 
 
+def test_scd_type_2_by_column_composite_key(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    adapter.scd_type_2_by_column(
+        target_table="target",
+        source_table=t.cast(exp.Select, parse_one("SELECT id_a, id_b, name, price FROM source")),
+        unique_key=[exp.func("CONCAT", exp.column("id_a"), exp.column("id_b"))],
+        valid_from_col=exp.column("test_VALID_from", quoted=True),
+        valid_to_col=exp.column("test_valid_to", quoted=True),
+        check_columns=[exp.column("name"), exp.column("price")],
+        columns_to_types={
+            "id_a": exp.DataType.build("VARCHAR"),
+            "id_b": exp.DataType.build("VARCHAR"),
+            "name": exp.DataType.build("VARCHAR"),
+            "price": exp.DataType.build("DOUBLE"),
+            "test_VALID_from": exp.DataType.build("TIMESTAMP"),
+            "test_valid_to": exp.DataType.build("TIMESTAMP"),
+        },
+        execution_time=datetime(2020, 1, 1, 0, 0, 0),
+    )
+
+    assert (
+        adapter.cursor.execute.call_args[0][0]
+        == parse_one(
+            """
+CREATE OR REPLACE TABLE "target" AS
+WITH "source" AS (
+  SELECT DISTINCT ON (CONCAT("id_a", "id_b"))
+    TRUE AS "_exists",
+    "id_a",
+    "id_b",
+    "name",
+    "price",
+  FROM (
+    SELECT
+      "id_a",
+      "id_b",
+      "name",
+      "price"
+    FROM "source"
+  ) AS "raw_source"
+), "static" AS (
+  SELECT
+    "id_a",
+    "id_b",
+    "name",
+    "price",
+    "test_VALID_from",
+    "test_valid_to",
+    TRUE AS "_exists"
+  FROM "target"
+  WHERE
+    NOT "test_valid_to" IS NULL
+), "latest" AS (
+  SELECT
+    "id_a",
+    "id_b",
+    "name",
+    "price",
+    "test_VALID_from",
+    "test_valid_to",
+    TRUE AS "_exists"
+  FROM "target"
+  WHERE
+    "test_valid_to" IS NULL
+), "deleted" AS (
+  SELECT
+    "static"."id_a",
+    "static"."id_b",
+    "static"."name",
+    "static"."price",
+    "static"."test_VALID_from",
+    "static"."test_valid_to"
+  FROM "static"
+  LEFT JOIN "latest"
+    ON CONCAT("static"."id_a", "static"."id_b") = CONCAT("latest"."id_a", "latest"."id_b")
+  WHERE
+    "latest"."test_valid_to" IS NULL
+), "latest_deleted" AS (
+  SELECT
+    TRUE AS "_exists",
+    CONCAT("id_a", "id_b") AS "_key0",
+    MAX("test_valid_to") AS "test_valid_to"
+  FROM "deleted"
+  GROUP BY
+    CONCAT("id_a", "id_b")
+), "joined" AS (
+  SELECT
+    "source"."_exists" AS "_exists",
+    "latest"."id_a" AS "t_id_a",
+    "latest"."id_b" AS "t_id_b",
+    "latest"."name" AS "t_name",
+    "latest"."price" AS "t_price",
+    "latest"."test_VALID_from" AS "t_test_VALID_from",
+    "latest"."test_valid_to" AS "t_test_valid_to",
+    "source"."id_a" AS "id_a",
+    "source"."id_b" AS "id_b",
+    "source"."name" AS "name",
+    "source"."price" AS "price"
+  FROM "latest"
+  LEFT JOIN "source"
+    ON CONCAT("latest"."id_a", "latest"."id_b") = CONCAT("source"."id_a", "source"."id_b")
+  UNION ALL
+  SELECT
+    "source"."_exists" AS "_exists",
+    "latest"."id_a" AS "t_id_a",
+    "latest"."id_b" AS "t_id_b",
+    "latest"."name" AS "t_name",
+    "latest"."price" AS "t_price",
+    "latest"."test_VALID_from" AS "t_test_VALID_from",
+    "latest"."test_valid_to" AS "t_test_valid_to",
+    "source"."id_a" AS "id_a",
+    "source"."id_b" AS "id_b",
+    "source"."name" AS "name",
+    "source"."price" AS "price"
+  FROM "latest"
+  RIGHT JOIN "source"
+    ON CONCAT("latest"."id_a", "latest"."id_b") = CONCAT("source"."id_a", "source"."id_b")
+  WHERE
+    "latest"."_exists" IS NULL
+), "updated_rows" AS (
+  SELECT
+    COALESCE("joined"."t_id_a", "joined"."id_a") AS "id_a",
+    COALESCE("joined"."t_id_b", "joined"."id_b") AS "id_b",
+    COALESCE("joined"."t_name", "joined"."name") AS "name",
+    COALESCE("joined"."t_price", "joined"."price") AS "price",
+    COALESCE("t_test_VALID_from", CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AS "test_VALID_from",
+    CASE
+      WHEN "joined"."_exists" IS NULL
+      OR (
+        (
+          NOT CONCAT("t_id_a", "t_id_b") IS NULL AND NOT CONCAT("id_a", "id_b") IS NULL
+        )
+        AND (
+          "name" <> "t_name"
+          OR (
+            "t_name" IS NULL AND NOT "name" IS NULL
+          )
+          OR (
+            NOT "t_name" IS NULL AND "name" IS NULL
+          )
+          OR "price" <> "t_price"
+          OR (
+            "t_price" IS NULL AND NOT "price" IS NULL
+          )
+          OR (
+            NOT "t_price" IS NULL AND "price" IS NULL
+          )
+        )
+      )
+      THEN CAST('2020-01-01 00:00:00' AS TIMESTAMP)
+      ELSE "t_test_valid_to"
+    END AS "test_valid_to"
+  FROM "joined"
+  LEFT JOIN "latest_deleted"
+    ON CONCAT("joined"."id_a", "joined"."id_b") = "latest_deleted"."_key0"
+), "inserted_rows" AS (
+  SELECT
+    "id_a",
+    "id_b",
+    "name",
+    "price",
+    CAST('2020-01-01 00:00:00' AS TIMESTAMP) AS "test_VALID_from",
+    CAST(NULL AS TIMESTAMP) AS "test_valid_to"
+  FROM "joined"
+  WHERE
+    (
+      NOT CONCAT("t_id_a", "t_id_b") IS NULL AND NOT CONCAT("id_a", "id_b") IS NULL
+    )
+    AND (
+      "name" <> "t_name"
+      OR (
+        "t_name" IS NULL AND NOT "name" IS NULL
+      )
+      OR (
+        NOT "t_name" IS NULL AND "name" IS NULL
+      )
+      OR "price" <> "t_price"
+      OR (
+        "t_price" IS NULL AND NOT "price" IS NULL
+      )
+      OR (
+        NOT "t_price" IS NULL AND "price" IS NULL
+      )
+    )
+)
+SELECT CAST("id_a" AS VARCHAR) AS "id_a", CAST("id_b" AS VARCHAR) AS "id_b", CAST("name" AS VARCHAR) AS "name", CAST("price" AS DOUBLE) AS "price", CAST("test_VALID_from" AS TIMESTAMP) AS "test_VALID_from", CAST("test_valid_to" AS TIMESTAMP) AS "test_valid_to" FROM (SELECT "id_a", "id_b", "name", "price", "test_VALID_from", "test_valid_to" FROM "static" UNION ALL SELECT "id_a", "id_b", "name", "price", "test_VALID_from", "test_valid_to" FROM "updated_rows" UNION ALL SELECT "id_a", "id_b", "name", "price", "test_VALID_from", "test_valid_to" FROM "inserted_rows") AS "_subquery"
+    """
+        ).sql()
+    )
+
+
 def test_scd_type_2_truncate(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
diff --git a/tests/dbt/test_config.py b/tests/dbt/test_config.py
index 57e39071..0495d9e0 100644
--- a/tests/dbt/test_config.py
+++ b/tests/dbt/test_config.py
@@ -789,3 +789,22 @@ def test_variable_override():
         variables={"yet_another_var": 2, "start": "2021-01-01"},
     )
     assert project.packages["sushi"].variables["yet_another_var"] == 2
+
+
+def test_depends_on(assert_exp_eq, sushi_test_project):
+    # Case 1: using an undefined variable without a default value
+    context = sushi_test_project.context
+
+    model_config = ModelConfig(
+        alias="sushi.test",
+        sql="SELECT * FROM {{ ref('waiter_revenue_by_day') }} JOIN other_table",
+        dependencies=Dependencies(refs=["waiter_revenue_by_day"]),
+    )
+
+    sqlmesh_model = model_config.to_sqlmesh(context)
+    assert sqlmesh_model.depends_on_ == {'"memory"."sushi"."waiter_revenue_by_day_v2"'}
+    assert sqlmesh_model.depends_on == {'"memory"."sushi"."waiter_revenue_by_day_v2"'}
+    assert sqlmesh_model.full_depends_on == {'"memory"."sushi"."waiter_revenue_by_day_v2"'}
+
+    # Make sure the query wasn't rendered
+    assert not sqlmesh_model._query_renderer._cache
diff --git a/tests/integrations/github/cicd/test_github_commands.py b/tests/integrations/github/cicd/test_github_commands.py
index 97684dbe..222118c2 100644
--- a/tests/integrations/github/cicd/test_github_commands.py
+++ b/tests/integrations/github/cicd/test_github_commands.py
@@ -16,7 +16,7 @@ from sqlmesh.integrations.github.cicd.controller import (
     GithubCheckConclusion,
     GithubCheckStatus,
 )
-from sqlmesh.utils.errors import PlanError, TestError
+from sqlmesh.utils.errors import ConflictingPlanError, PlanError, TestError
 
 pytest_plugins = ["tests.integrations.github.cicd.fixtures"]
 pytestmark = [
@@ -788,7 +788,7 @@ def test_pr_update_failure(
         )
 
 
-def test_prod_update_failure(
+def make_test_prod_update_failure_case(
     github_client,
     make_controller,
     make_mock_check_run,
@@ -796,6 +796,8 @@ def test_prod_update_failure(
     make_pull_request_review,
     tmp_path: pathlib.Path,
     mocker: MockerFixture,
+    to_raise_on_prod_plan: Exception,
+    expect_prod_sync_conclusion: GithubCheckConclusion,
 ):
     """
     Scenario:
@@ -842,7 +844,7 @@ def test_prod_update_failure(
 
     def raise_on_prod_plan(plan: Plan):
         if plan.environment.name == c.PROD:
-            raise PlanError("Failed to update Prod environment")
+            raise to_raise_on_prod_plan
 
     controller._context.apply = mocker.MagicMock(side_effect=lambda plan: raise_on_prod_plan(plan))
 
@@ -875,7 +877,7 @@ def test_prod_update_failure(
     assert GithubCheckStatus(prod_checks_runs[0]["status"]).is_queued
     assert GithubCheckStatus(prod_checks_runs[1]["status"]).is_in_progress
     assert GithubCheckStatus(prod_checks_runs[2]["status"]).is_completed
-    assert GithubCheckConclusion(prod_checks_runs[2]["conclusion"]).is_action_required
+    assert GithubCheckConclusion(prod_checks_runs[2]["conclusion"]) == expect_prod_sync_conclusion
 
     assert "SQLMesh - Run Unit Tests" in controller._check_run_mapping
     test_checks_runs = controller._check_run_mapping["SQLMesh - Run Unit Tests"].all_kwargs
@@ -919,10 +921,106 @@ def test_prod_update_failure(
         output = f.read()
         assert (
             output
-            == "run_unit_tests=success\nhas_required_approval=success\ncreated_pr_environment=true\npr_environment_name=hello_world_2\npr_environment_synced=success\nprod_plan_preview=success\nprod_environment_synced=action_required\n"
+            == f"run_unit_tests=success\nhas_required_approval=success\ncreated_pr_environment=true\npr_environment_name=hello_world_2\npr_environment_synced=success\nprod_plan_preview=success\nprod_environment_synced={expect_prod_sync_conclusion.value}\n"
         )
 
 
+def test_prod_update_failure(
+    github_client,
+    make_controller,
+    make_mock_check_run,
+    make_mock_issue_comment,
+    make_pull_request_review,
+    tmp_path: pathlib.Path,
+    mocker: MockerFixture,
+):
+    """
+    Scenario:
+    - PR is not merged
+    - PR has been approved by a required reviewer
+    - Tests passed
+    - PR Merge Method defined
+    - Delete environment is enabled
+    - Prod environment update failed
+    """
+
+    make_test_prod_update_failure_case(
+        github_client=github_client,
+        make_controller=make_controller,
+        make_mock_check_run=make_mock_check_run,
+        make_mock_issue_comment=make_mock_issue_comment,
+        make_pull_request_review=make_pull_request_review,
+        tmp_path=tmp_path,
+        mocker=mocker,
+        to_raise_on_prod_plan=PlanError("Failed to update Prod environment"),
+        expect_prod_sync_conclusion=GithubCheckConclusion.ACTION_REQUIRED,
+    )
+
+
+def test_prod_update_conflict(
+    github_client,
+    make_controller,
+    make_mock_check_run,
+    make_mock_issue_comment,
+    make_pull_request_review,
+    tmp_path: pathlib.Path,
+    mocker: MockerFixture,
+):
+    """
+    Scenario:
+    - PR is not merged
+    - PR has been approved by a required reviewer
+    - Tests passed
+    - PR Merge Method defined
+    - Delete environment is enabled
+    - Prod environment update conflicted
+    """
+
+    make_test_prod_update_failure_case(
+        github_client=github_client,
+        make_controller=make_controller,
+        make_mock_check_run=make_mock_check_run,
+        make_mock_issue_comment=make_mock_issue_comment,
+        make_pull_request_review=make_pull_request_review,
+        tmp_path=tmp_path,
+        mocker=mocker,
+        to_raise_on_prod_plan=ConflictingPlanError("Plan a conflicts with plan b"),
+        expect_prod_sync_conclusion=GithubCheckConclusion.SKIPPED,
+    )
+
+
+def test_prod_update_exception(
+    github_client,
+    make_controller,
+    make_mock_check_run,
+    make_mock_issue_comment,
+    make_pull_request_review,
+    tmp_path: pathlib.Path,
+    mocker: MockerFixture,
+):
+    """
+    Scenario:
+    - PR is not merged
+    - PR has been approved by a required reviewer
+    - Tests passed
+    - PR Merge Method defined
+    - Delete environment is enabled
+    - Prod environment update fails with an unknown exception
+    """
+
+    make_test_prod_update_failure_case(
+        github_client=github_client,
+        make_controller=make_controller,
+        make_mock_check_run=make_mock_check_run,
+        make_mock_issue_comment=make_mock_issue_comment,
+        make_pull_request_review=make_pull_request_review,
+        tmp_path=tmp_path,
+        mocker=mocker,
+        to_raise_on_prod_plan=RuntimeError("boom"),
+        expect_prod_sync_conclusion=GithubCheckConclusion.FAILURE,
+    )
+
+
 def test_comment_command_invalid(
     github_client,
     make_controller,