SQLAlchemy: Ignore SQL "FOR UPDATE" clause, CrateDB does not support it

amotl · amotl · commit c1ac6790753c · 2023-09-29T01:16:46.000-07:00
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -17,6 +17,7 @@ Unreleased
   constraints
 - DBAPI: Properly raise ``IntegrityError`` exceptions instead of
   ``ProgrammingError``, when CrateDB raises a ``DuplicateKeyException``.
+- SQLAlchemy: Ignore SQL's ``FOR UPDATE`` clause. Thanks, @surister.
 
 .. _urllib3 v2.0 migration guide: https://urllib3.readthedocs.io/en/latest/v2-migration-guide.html
 .. _urllib3 v2.0 roadmap: https://urllib3.readthedocs.io/en/stable/v2-roadmap.html
diff --git a/src/crate/client/sqlalchemy/compiler.py b/src/crate/client/sqlalchemy/compiler.py
@@ -309,3 +309,10 @@ def limit_clause(self, select, **kw):
         Generate OFFSET / LIMIT clause, PostgreSQL-compatible.
         """
         return PGCompiler.limit_clause(self, select, **kw)
+
+    def for_update_clause(self, select, **kw):
+        # CrateDB does not support the `INSERT ... FOR UPDATE` clause.
+        # See https://github.com/crate/crate-python/issues/577.
+        warnings.warn("CrateDB does not support the 'INSERT ... FOR UPDATE' clause, "
+                      "it will be omitted when generating SQL statements.")
+        return ''
diff --git a/src/crate/client/sqlalchemy/tests/compiler_test.py b/src/crate/client/sqlalchemy/tests/compiler_test.py
@@ -43,7 +43,7 @@
 from crate.testing.settings import crate_host
 
 
-class SqlAlchemyCompilerTest(ParametrizedTestCase):
+class SqlAlchemyCompilerTest(ParametrizedTestCase, ExtraAssertions):
 
     def setUp(self):
         self.crate_engine = sa.create_engine('crate://')
@@ -257,6 +257,34 @@ def test_insert_manyvalues(self):
             mock.call(mock.ANY, 'INSERT INTO mytable (name) VALUES (?)', ('foo_4', ), None),
         ])
 
+    def test_for_update(self):
+        """
+        Verify the `CrateCompiler.for_update_clause` method to
+        omit the clause, since CrateDB does not support it.
+        """
+
+        with warnings.catch_warnings(record=True) as w:
+
+            # By default, warnings from a loop will only be emitted once.
+            # This scenario tests exactly this behaviour, to verify logs
+            # don't get flooded.
+            warnings.simplefilter("once")
+
+            selectable = self.mytable.select().with_for_update()
+            _ = str(selectable.compile(bind=self.crate_engine))
+
+            selectable = self.mytable.select().with_for_update()
+            statement = str(selectable.compile(bind=self.crate_engine))
+
+        # Verify SQL statement.
+        self.assertEqual(statement, "SELECT mytable.name, mytable.data \nFROM mytable")
+
+        # Verify if corresponding warning is emitted, once.
+        self.assertEqual(len(w), 1)
+        self.assertIsSubclass(w[-1].category, UserWarning)
+        self.assertIn("CrateDB does not support the 'INSERT ... FOR UPDATE' clause, "
+                      "it will be omitted when generating SQL statements.", str(w[-1].message))
+
 
 FakeCursor = MagicMock(name='FakeCursor', spec=Cursor)
 
