Fix mypy errors, ruff formatting, and add requested tests

Fix QueryOrDF import path, replace reduce with exp.and_ for type
safety, add 3-column composite key test and expression-based
composite key test via insert_overwrite_by_partition.

Signed-off-by: Hashwanth S <s.hashwanth531@gmail.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Hashwanth S <s.hashwanth531@gmail.com>

Author: hashwnath

sqlmesh/core/engine_adapter/mysql.py

@@ -2,8 +2,6 @@
 
 import logging
 import typing as t
-from functools import reduce
-
 from sqlglot import exp, parse_one
 
 from sqlmesh.core.dialect import to_schema
@@ -22,7 +20,8 @@
 )
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import QueryOrDF, SchemaName, TableName
+    from sqlmesh.core._typing import SchemaName, TableName
+    from sqlmesh.core.engine_adapter._typing import QueryOrDF
 
 logger = logging.getLogger(__name__)
 
@@ -198,7 +197,12 @@ def _replace_by_key(
     ) -> None:
         if len(key) <= 1:
             return super()._replace_by_key(
-                target_table, source_table, target_columns_to_types, key, is_unique_key, source_columns
+                target_table,
+                source_table,
+                target_columns_to_types,
+                key,
+                is_unique_key,
+                source_columns,
             )
 
         if target_columns_to_types is None:
@@ -222,14 +226,13 @@ def _replace_by_key(
             try:
                 # Build a JOIN-based DELETE instead of using CONCAT_WS.
                 # CONCAT_WS prevents MySQL/MariaDB from using indexes, causing full table scans.
-                on_condition = reduce(
-                    lambda a, b: exp.And(this=a, expression=b),
-                    [
+                on_condition = exp.and_(
+                    *[
                         self._qualify_columns(k, target_alias).eq(
                             self._qualify_columns(k, temp_alias)
                         )
                         for k in key
-                    ],
+                    ]
                 )
 
                 target_table_aliased = exp.to_table(target_table).as_(target_alias, quoted=True)

tests/core/engine_adapter/test_mysql.py

@@ -92,9 +92,7 @@ def test_replace_by_key_composite_uses_join_delete(
 ):
     """Composite key DELETE uses JOIN instead of CONCAT_WS to allow index usage."""
     adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
-    temp_table_mock = mocker.patch(
-        "sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table"
-    )
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table")
     temp_table_mock.return_value = exp.to_table("temporary")
 
     adapter.merge(
@@ -135,14 +133,81 @@ def test_replace_by_key_composite_uses_join_delete(
     )
 
 
+def test_replace_by_key_three_column_composite_key(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    """3-column composite key matching the original issue scenario (#5711)."""
+    adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table")
+    temp_table_mock.return_value = exp.to_table("temporary")
+
+    adapter.merge(
+        target_table="target",
+        source_table=t.cast(exp.Select, parse_one("SELECT id, region, ts, val FROM source")),
+        target_columns_to_types={
+            "id": exp.DataType(this=exp.DataType.Type.INT),
+            "region": exp.DataType(this=exp.DataType.Type.VARCHAR),
+            "ts": exp.DataType(this=exp.DataType.Type.TIMESTAMP),
+            "val": exp.DataType(this=exp.DataType.Type.INT),
+        },
+        unique_key=[parse_one("id"), parse_one("region"), parse_one("ts")],
+    )
+
+    sql_calls = to_sql_calls(adapter)
+
+    assert any("CONCAT_WS" in s for s in sql_calls) is False
+    assert any("INNER JOIN" in s for s in sql_calls) is True
+
+    adapter.cursor.execute.assert_has_calls(
+        [
+            call(
+                "DELETE `_target` FROM `target` AS `_target` INNER JOIN `temporary` AS `_temp` ON `_target`.`id` = `_temp`.`id` AND `_target`.`region` = `_temp`.`region` AND `_target`.`ts` = `_temp`.`ts`"
+            ),
+        ]
+    )
+
+
+def test_replace_by_key_expression_based_composite_key(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    """Expression-based composite keys (e.g. DATE_TRUNC + column) via insert_overwrite_by_partition."""
+    adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table")
+    temp_table_mock.return_value = exp.to_table("temporary")
+
+    adapter.insert_overwrite_by_partition(
+        table_name="target",
+        query_or_df=t.cast(exp.Select, parse_one("SELECT id, ts, val FROM source")),
+        partitioned_by=[
+            parse_one("DATE_TRUNC('month', ts)"),
+            parse_one("id"),
+        ],
+        target_columns_to_types={
+            "id": exp.DataType(this=exp.DataType.Type.INT),
+            "ts": exp.DataType(this=exp.DataType.Type.TIMESTAMP),
+            "val": exp.DataType(this=exp.DataType.Type.INT),
+        },
+    )
+
+    sql_calls = to_sql_calls(adapter)
+
+    assert any("CONCAT_WS" in s for s in sql_calls) is False
+    assert any("INNER JOIN" in s for s in sql_calls) is True
+
+    # DATE_TRUNC transpiles to STR_TO_DATE in MySQL; both key parts should be qualified
+    delete_sql = [s for s in sql_calls if "DELETE" in s][0]
+    assert "`_target`.`ts`" in delete_sql
+    assert "`_temp`.`ts`" in delete_sql
+    assert "`_target`.`id`" in delete_sql
+    assert "`_temp`.`id`" in delete_sql
+
+
 def test_replace_by_key_single_key_uses_in(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
 ):
     """Single key DELETE still uses the IN-based approach (indexes work fine for single column)."""
     adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
-    temp_table_mock = mocker.patch(
-        "sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table"
-    )
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._get_temp_table")
     temp_table_mock.return_value = exp.to_table("temporary")
 
     adapter.merge(