SQLAlchemy: Fix SQL statement caching for CrateDB's OBJECT type

CHANGES.txt

@@ -13,6 +13,8 @@ Unreleased
 
 - Allow handling datetime values tagged with time zone info when inserting or updating.
 
+- SQLAlchemy: Fix SQL statement caching for CrateDB's ``OBJECT`` type.
+
 
 2023/04/18 0.31.1
 =================

src/crate/client/sqlalchemy/tests/__init__.py

@@ -23,9 +23,10 @@
 from .dialect_test import SqlAlchemyDialectTest
 from .function_test import SqlAlchemyFunctionTest
 from .warnings_test import SqlAlchemyWarningsTest
+from .query_caching import SqlAlchemyQueryCompilationCaching
 
 
-def test_suite():
+def test_suite_unit():
     tests = TestSuite()
     tests.addTest(makeSuite(SqlAlchemyConnectionTest))
     tests.addTest(makeSuite(SqlAlchemyDictTypeTest))
@@ -42,3 +43,9 @@ def test_suite():
     tests.addTest(makeSuite(SqlAlchemyArrayTypeTest))
     tests.addTest(makeSuite(SqlAlchemyWarningsTest))
     return tests
+
+
+def test_suite_integration():
+    tests = TestSuite()
+    tests.addTest(makeSuite(SqlAlchemyQueryCompilationCaching))
+    return tests

src/crate/client/sqlalchemy/tests/query_caching.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8; -*-
+#
+# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
+# license agreements.  See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.  Crate licenses
+# this file to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.  You may
+# obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+# License for the specific language governing permissions and limitations
+# under the License.
+#
+# However, if you have executed another commercial license agreement
+# with Crate these terms will supersede the license and you may use the
+# software solely pursuant to the terms of the relevant commercial agreement.
+
+from __future__ import absolute_import
+from unittest import TestCase, skipIf
+
+import sqlalchemy as sa
+from sqlalchemy.orm import Session
+from sqlalchemy.sql.operators import eq
+
+from crate.client.sqlalchemy import SA_VERSION, SA_1_4
+from crate.testing.settings import crate_host
+
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+
+from crate.client.sqlalchemy.types import Object, ObjectArray
+
+
+class SqlAlchemyQueryCompilationCaching(TestCase):
+
+    def setUp(self):
+        self.engine = sa.create_engine(f"crate://{crate_host}")
+        self.metadata = sa.MetaData(schema="testdrive")
+        self.session = Session(bind=self.engine)
+        self.Character = self.setup_entity()
+
+    def setup_entity(self):
+        """
+        Define ORM entity.
+        """
+        Base = declarative_base(metadata=self.metadata)
+
+        class Character(Base):
+            __tablename__ = 'characters'
+            name = sa.Column(sa.String, primary_key=True)
+            age = sa.Column(sa.Integer)
+            data = sa.Column(Object)
+            data_list = sa.Column(ObjectArray)
+
+        return Character
+
+    def setup_data(self):
+        """
+        Insert two records into the `characters` table.
+        """
+        self.metadata.drop_all(self.engine)
+        self.metadata.create_all(self.engine)
+
+        Character = self.Character
+        char1 = Character(name='Trillian', data={'x': 1}, data_list=[{'foo': 1, 'bar': 10}])
+        char2 = Character(name='Slartibartfast', data={'y': 2}, data_list=[{'bar': 2}])
+        self.session.add(char1)
+        self.session.add(char2)
+        self.session.commit()
+        self.session.execute(sa.text("REFRESH TABLE testdrive.characters;"))
+
+    @skipIf(SA_VERSION < SA_1_4, "On SA13, the 'ResultProxy' object has no attribute 'scalar_one'")
+    def test_object_multiple_select(self):
+        """
+        The SQLAlchemy implementation of CrateDB's `OBJECT` type offers indexed
+        access to the instance's content in form of a dictionary. Thus, it must
+        not use `cache_ok = True` on its implementation, i.e. this part of the
+        compiled SQL clause must not be cached.
+
+        This test verifies that two subsequent `SELECT` statements are translated
+        well, and don't trip on incorrect SQL compiled statement caching.
+        """
+        self.setup_data()
+        Character = self.Character
+
+        selectable = sa.select(Character).where(Character.data['x'] == 1)
+        result = self.session.execute(selectable).scalar_one().data
+        self.assertEqual({"x": 1}, result)
+
+        selectable = sa.select(Character).where(Character.data['y'] == 2)
+        result = self.session.execute(selectable).scalar_one().data
+        self.assertEqual({"y": 2}, result)
+
+    @skipIf(SA_VERSION < SA_1_4, "On SA13, the 'ResultProxy' object has no attribute 'scalar_one'")
+    def test_objectarray_multiple_select(self):
+        """
+        The SQLAlchemy implementation of CrateDB's `ARRAY` type in form of the
+        `ObjectArray`, does *not* offer indexed access to the instance's content.
+        Thus, using `cache_ok = True` on that type should be sane, and not mess
+        up SQLAlchemy's SQL compiled statement caching.
+        """
+        self.setup_data()
+        Character = self.Character
+
+        selectable = sa.select(Character).where(Character.data_list['foo'].any(1, operator=eq))
+        result = self.session.execute(selectable).scalar_one().data
+        self.assertEqual({"x": 1}, result)
+
+        selectable = sa.select(Character).where(Character.data_list['bar'].any(2, operator=eq))
+        result = self.session.execute(selectable).scalar_one().data
+        self.assertEqual({"y": 2}, result)

src/crate/client/sqlalchemy/types.py

@@ -132,7 +132,7 @@ def __eq__(self, other):
 
 
 class _Craty(sqltypes.UserDefinedType):
-    cache_ok = True
+    cache_ok = False
 
     class Comparator(sqltypes.TypeEngine.Comparator):
 

src/crate/client/tests.py

@@ -54,7 +54,8 @@
     TestCrateJsonEncoder,
     TestDefaultSchemaHeader,
 )
-from .sqlalchemy.tests import test_suite as sqlalchemy_test_suite
+from .sqlalchemy.tests import test_suite_unit as sqlalchemy_test_suite_unit
+from .sqlalchemy.tests import test_suite_integration as sqlalchemy_test_suite_integration
 
 log = logging.getLogger('crate.testing.layer')
 ch = logging.StreamHandler()
@@ -344,7 +345,7 @@ def test_suite():
     suite.addTest(unittest.makeSuite(TestUsernameSentAsHeader))
     suite.addTest(unittest.makeSuite(TestCrateJsonEncoder))
     suite.addTest(unittest.makeSuite(TestDefaultSchemaHeader))
-    suite.addTest(sqlalchemy_test_suite())
+    suite.addTest(sqlalchemy_test_suite_unit())
     suite.addTest(doctest.DocTestSuite('crate.client.connection'))
     suite.addTest(doctest.DocTestSuite('crate.client.http'))
 
@@ -394,6 +395,7 @@ def test_suite():
         encoding='utf-8'
     )
     s.layer = ensure_cratedb_layer()
+    s.addTest(sqlalchemy_test_suite_integration())
     suite.addTest(s)
 
     return suite