Repository: TobikoData/sqlmesh
Pull request #3607 (2025-01-09 19:02:17)

Feat(redshift): Add native Merge support for Redshift
This update adds native `MERGE` support for Incremental by Unique Key models in Redshift (fixes: #3316)

This required adjustments to the generated statement to accommodate for the lack of alias support for the target table and limitations of the when matched clause [Redshift MERGE Documentation](https://docs.aws.amazon.com/redshift/latest/dg/r_MERGE.html)


erindru: This theoretically looks good but are you able to show the Redshift integration tests still passing by commenting out the branch filter for the cloud tests in [.circleci/continue_config.yml](https://github.com/TobikoData/sqlmesh/blob/main/.circleci/continue_config.yml#L316)? (and re-instating it prior to merge)

LogicalMergeMixin had less syntax limitations so I wonder if "upgrading" to Redshift's limited version of MERGE will break things
themisvaltinos: > This theoretically looks good but are you able to show the Redshift integration tests still passing by commenting out the branch filter for the cloud tests in [.circleci/continue_config.yml](https://github.com/TobikoData/sqlmesh/blob/main/.circleci/continue_config.yml#L316)? (and re-instating it prior to merge)

Thanks for this suggestion I tested it and it appears to be working as expected

> LogicalMergeMixin had less syntax limitations so I wonder if "upgrading" to Redshift's limited version of MERGE will break things

I was wondering of this as well, but given that logical merge does not support `when_matched` or `merge_filter` I was thinking it should be fine in that aspect. But can you think of any other potential risks we should consider?
erindru: > I tested it and it appears to be working as expected
Thanks, looks good!

> But can you think of any other potential risks we should consider?
Every change carries risk :) but in this case, our tests pass so I think we are good. If this change breaks someones workflow, its an opportunity to improve our test suite

Patch:
diff --git a/docs/concepts/models/model_kinds.md b/docs/concepts/models/model_kinds.md
index dc15de8a..cc7de46c 100644
--- a/docs/concepts/models/model_kinds.md
+++ b/docs/concepts/models/model_kinds.md
@@ -349,9 +349,12 @@ MODEL (
 * BigQuery
 * Databricks
 * Postgres
+* Redshift
 * Snowflake
 * Spark
 
+Redshift supports only the `UPDATE` or `DELETE` actions for the `WHEN MATCHED` clause and does not allow multiple `WHEN MATCHED` expressions. For further information, refer to the [Redshift documentation](https://docs.aws.amazon.com/redshift/latest/dg/r_MERGE.html#r_MERGE-parameters).
+
 ### Merge Filter Expression
 
 The `MERGE` statement typically induces a full table scan of the existing table, which can be problematic with large data volumes. 
diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py
index cfe40708..badb472b 100644
--- a/sqlmesh/core/engine_adapter/redshift.py
+++ b/sqlmesh/core/engine_adapter/redshift.py
@@ -7,10 +7,10 @@ import pandas as pd
 from sqlglot import exp
 
 from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.engine_adapter.base import MERGE_SOURCE_ALIAS, MERGE_TARGET_ALIAS
 from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
-    LogicalMergeMixin,
     NonTransactionalTruncateMixin,
     VarcharSizeWorkaroundMixin,
     RowDiffMixin,
@@ -23,10 +23,11 @@ from sqlmesh.core.engine_adapter.shared import (
     set_catalog,
 )
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import SchemaName, TableName
-    from sqlmesh.core.engine_adapter.base import QueryOrDF
+    from sqlmesh.core.engine_adapter.base import QueryOrDF, Query
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +35,6 @@ logger = logging.getLogger(__name__)
 @set_catalog()
 class RedshiftEngineAdapter(
     BasePostgresEngineAdapter,
-    LogicalMergeMixin,
     GetCurrentCatalogFromFunctionMixin,
     NonTransactionalTruncateMixin,
     VarcharSizeWorkaroundMixin,
@@ -329,6 +329,45 @@ class RedshiftEngineAdapter(
             for row in df.itertuples()
         ]
 
+    def _merge(
+        self,
+        target_table: TableName,
+        query: Query,
+        on: exp.Expression,
+        whens: exp.Whens,
+    ) -> None:
+        # Redshift does not support table aliases in the target table of a MERGE statement.
+        # So we must use the actual table name instead of an alias, as we do with the source table.
+        def resolve_target_table(expression: exp.Expression) -> exp.Expression:
+            if (
+                isinstance(expression, exp.Column)
+                and expression.table.upper() == MERGE_TARGET_ALIAS
+            ):
+                expression.set("table", exp.to_table(target_table))
+            return expression
+
+        # Ensure that there is exactly one "WHEN MATCHED" and one "WHEN NOT MATCHED" clause.
+        # Since Redshift does not support multiple "WHEN MATCHED" clauses.
+        if (
+            len(whens.expressions) != 2
+            or whens.expressions[0].args["matched"] == whens.expressions[1].args["matched"]
+        ):
+            raise SQLMeshError(
+                "Redshift only supports a single WHEN MATCHED and WHEN NOT MATCHED clause"
+            )
+
+        using = exp.alias_(
+            exp.Subquery(this=query), alias=MERGE_SOURCE_ALIAS, copy=False, table=True
+        )
+        self.execute(
+            exp.Merge(
+                this=target_table,
+                using=using,
+                on=on.transform(resolve_target_table),
+                whens=whens.transform(resolve_target_table),
+            )
+        )
+
     def _normalize_decimal_value(self, expr: exp.Expression, precision: int) -> exp.Expression:
         # Redshift is finicky. It truncates when the data is already in a table, but rounds when the data is generated as part of a SELECT.
         #


Test patch:
diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py
index 1560916d..adca20c7 100644
--- a/tests/core/engine_adapter/test_redshift.py
+++ b/tests/core/engine_adapter/test_redshift.py
@@ -8,6 +8,7 @@ from sqlglot import expressions as exp
 from sqlglot import parse_one
 
 from sqlmesh.core.engine_adapter import RedshiftEngineAdapter
+from sqlmesh.utils.errors import SQLMeshError
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.engine, pytest.mark.redshift]
@@ -322,3 +323,84 @@ def test_alter_table_drop_column_cascade(adapter: t.Callable):
     assert to_sql_calls(adapter) == [
         'ALTER TABLE "test_table" DROP COLUMN "test_column" CASCADE',
     ]
+
+
+def test_merge(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(RedshiftEngineAdapter)
+
+    adapter.merge(
+        target_table=exp.to_table("target_table_name"),
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
+        columns_to_types={
+            "ID": exp.DataType.build("int"),
+            "ts": exp.DataType.build("timestamp"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("ID", quoted=True)],
+    )
+
+    # Test additional predicates in the merge_filter
+    adapter.merge(
+        target_table=exp.to_table("target_table_name"),
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
+        columns_to_types={
+            "ID": exp.DataType.build("int"),
+            "ts": exp.DataType.build("timestamp"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("ID", quoted=True)],
+        merge_filter=exp.and_(
+            exp.and_(exp.column("ID", "__merge_source__") > 0),
+            exp.column("ts", "__merge_target__") < exp.column("ts", "__merge_source__"),
+        ),
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        'MERGE INTO "target_table_name" USING (SELECT "ID", "ts", "val" FROM "source") AS "__MERGE_SOURCE__" ON "target_table_name"."ID" = "__MERGE_SOURCE__"."ID" WHEN MATCHED THEN UPDATE SET "ID" = "__merge_source__"."ID", "ts" = "__merge_source__"."ts", "val" = "__merge_source__"."val" WHEN NOT MATCHED THEN INSERT ("ID", "ts", "val") VALUES ("__merge_source__"."ID", "__merge_source__"."ts", "__merge_source__"."val")',
+        'MERGE INTO "target_table_name" USING (SELECT "ID", "ts", "val" FROM "source") AS "__MERGE_SOURCE__" ON ("__merge_source__"."ID" > 0 AND "target_table_name"."ts" < "__merge_source__"."ts") AND "target_table_name"."ID" = "__MERGE_SOURCE__"."ID" WHEN MATCHED THEN UPDATE SET "ID" = "__merge_source__"."ID", "ts" = "__merge_source__"."ts", "val" = "__merge_source__"."val" WHEN NOT MATCHED THEN INSERT ("ID", "ts", "val") VALUES ("__merge_source__"."ID", "__merge_source__"."ts", "__merge_source__"."val")',
+    ]
+
+
+def test_merge_when_matched_error(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(RedshiftEngineAdapter)
+
+    with pytest.raises(
+        SQLMeshError,
+        match=r".*Redshift only supports a single WHEN MATCHED and WHEN NOT MATCHED clause*",
+    ):
+        adapter.merge(
+            target_table=exp.to_table("target_table_name"),
+            source_table=t.cast(exp.Select, parse_one('SELECT "ID", val FROM source')),
+            columns_to_types={
+                "ID": exp.DataType.build("int"),
+                "val": exp.DataType.build("int"),
+            },
+            unique_key=[exp.to_identifier("ID", quoted=True)],
+            when_matched=exp.Whens(
+                expressions=[
+                    exp.When(
+                        matched=True,
+                        condition=exp.column("ID", "__MERGE_SOURCE__").eq(exp.Literal.number(1)),
+                        then=exp.Update(
+                            expressions=[
+                                exp.column("val", "__MERGE_TARGET__").eq(
+                                    exp.column("val", "__MERGE_SOURCE__")
+                                ),
+                            ],
+                        ),
+                    ),
+                    exp.When(
+                        matched=True,
+                        source=False,
+                        then=exp.Update(
+                            expressions=[
+                                exp.column("val", "__MERGE_TARGET__").eq(
+                                    exp.column("val", "__MERGE_SOURCE__")
+                                ),
+                            ],
+                        ),
+                    ),
+                ]
+            ),
+        )