Repository: TobikoData/sqlmesh
Pull request #3075 (2024-08-30 20:45:03)

Fix: Support for unique key expressions for SCD2 by column models


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/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/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,