Compile Merge On clause + Fix spacing

Author: rad-pat

databend_sqlalchemy/connector.py

@@ -163,7 +163,7 @@ def mogrify(self, query, parameters):
     def execute(self, operation, parameters=None):
         """Prepare and execute a database operation (query or command)."""
 
-        # ToDo - Fix this, which is preventing the execution of blank DDL sunch as CREATE INDEX statements which aren't currently supported
+        # ToDo - Fix this, which is preventing the execution of blank DDL such as CREATE INDEX statements which aren't currently supported
         # Seems hard to fix when statements are coming from metadata.create_all()
         if not operation:
             return

databend_sqlalchemy/databend_dialect.py

@@ -906,13 +906,17 @@ def visit_merge(self, merge, **kw):
             )
         elif isinstance(merge.source, Subquery):
             source = merge.source._compiler_dispatch(self, **source_kw)
+        else:
+            source = merge.source
+
+        merge_on = merge.on._compiler_dispatch(self, **kw)
 
         target_table = self.preparer.format_table(merge.target)
         return (
             f"MERGE INTO {target_table}\n"
             f" USING {source}\n"
-            f" ON {merge.on}\n"
-            f"{clauses if clauses else ''}"
+            f" ON {merge_on}\n"
+            f" {clauses if clauses else ''}"
         )
 
     def visit_when_merge_matched_update(self, merge_matched_update, **kw):
@@ -921,7 +925,7 @@ def visit_when_merge_matched_update(self, merge_matched_update, **kw):
             if merge_matched_update.predicate is not None
             else ""
         )
-        update_str = f"WHEN MATCHED{case_predicate} THEN\n" f"\tUPDATE"
+        update_str = f"WHEN MATCHED{case_predicate} THEN\n UPDATE"
         if not merge_matched_update.set:
             return f"{update_str} *"
 
@@ -950,7 +954,7 @@ def visit_when_merge_unmatched(self, merge_unmatched, **kw):
             if merge_unmatched.predicate is not None
             else ""
         )
-        insert_str = f"WHEN NOT MATCHED{case_predicate} THEN\n" f"\tINSERT"
+        insert_str = f"WHEN NOT MATCHED{case_predicate} THEN\n INSERT"
         if not merge_unmatched.set:
             return f"{insert_str} *"
 

databend_sqlalchemy/dml.py

@@ -76,6 +76,8 @@ class Merge(UpdateBase):
     __visit_name__ = "merge"
     _bind = None
 
+    inherit_cache = False
+
     def __init__(self, target, source, on):
         if not isinstance(source, (TableClause, Select, Subquery)):
             raise Exception(f"Invalid type for merge source: {source}")

tests/test_merge.py

@@ -11,15 +11,15 @@
 from sqlalchemy import types as sqltypes
 
 # from sqlalchemy.dialects.postgresql import insert
-from sqlalchemy.testing import config
+from sqlalchemy.testing import config, AssertsCompiledSQL
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing.assertions import assert_raises
 from sqlalchemy.testing.assertions import eq_
 
 from databend_sqlalchemy import Merge
 
 
-class MergeIntoTest(fixtures.TablesTest):
+class MergeIntoTest(fixtures.TablesTest, AssertsCompiledSQL):
     __backend__ = True
     run_define_tables = "each"
 
@@ -170,6 +170,11 @@ def test_when_not_matched_insert(self, connection):
         merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id)
         merge.when_not_matched_then_insert()
 
+        self.assert_compile(
+            merge,
+            'MERGE INTO "users" USING (SELECT users_xtra.id AS id, users_xtra.name AS name, users_xtra.login_email AS login_email FROM users_xtra) AS users_xtra ON "users".id = users_xtra.id WHEN NOT MATCHED THEN INSERT *',
+        )
+
         result = connection.execute(merge)
         eq_(
             connection.execute(users.select().order_by(users.c.id)).fetchall(),