diff --git a/CHANGES.md b/CHANGES.md index 22b2f830..076a1385 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -18,6 +18,17 @@ columns need migrating to `NUMERIC` - Types: Added `numeric` and `numeric_array` to the reflected type map, where they previously resolved to an abstract type that could not be compiled +- Compiler: Fixed array indexes rendering as object keys, as in `arr['1']`. + Added support for array slices, and for columns and expressions as indexes +- Compiler: Fixed cached statements on `sa.JSON` and `ObjectArray` columns + reusing the subscript of an earlier statement, which returned values of + another key. Fixed object keys containing a quote +- BREAKING: Compiler: A statement with a literal subscript on an `sa.JSON`, + `ObjectArray` or `ARRAY` column can't run with `executemany()` on a cached + engine. For that, use `execution_options(compiled_cache=None)` +- BREAKING: Compiler: An integer subscript on a nested object renders as an + array position. Write the key as a string, as in `obj['n']['1']` +- Compiler: `bool`, `float` and `None` subscripts raise `CompileError` ## 2026/06/22 0.43.1 - Compiler: Fixed `AttributeError: 'CrateCompilerSA20' object has no attribute diff --git a/docs/working-with-types.rst b/docs/working-with-types.rst index 6bdbdfb7..afe8df63 100644 --- a/docs/working-with-types.rst +++ b/docs/working-with-types.rst @@ -203,6 +203,21 @@ all values of that field of all objects in that object array: >>> query.all() [([1, 2, 3],), (None,), (None,)] +Select an element of an array by its position, or a part of it by a slice. +Positions start at 1, and slices include both bounds. A position of 0 or less, +or past the end of the array, returns ``NULL``: + + >>> query = session.query(Character.more_details[1]['foo']).order_by(Character.name) + >>> query.all() + [(1,), (None,), (None,)] + + >>> query = session.query(Character.more_details['foo'][2:3]) + >>> query.filter_by(name='Arthur Dent').all() + [([2, 3],)] + +Positions and slices work the same way on ``sa.ARRAY`` columns. With +``sa.ARRAY(..., zero_indexes=True)``, positions start at 0, as in Python. + Geospatial types ================ diff --git a/src/sqlalchemy_cratedb/compiler.py b/src/sqlalchemy_cratedb/compiler.py index 3856dc68..8901e852 100644 --- a/src/sqlalchemy_cratedb/compiler.py +++ b/src/sqlalchemy_cratedb/compiler.py @@ -26,7 +26,15 @@ import sqlalchemy as sa from sqlalchemy.dialects.postgresql.base import RESERVED_WORDS as POSTGRESQL_RESERVED_WORDS from sqlalchemy.dialects.postgresql.base import PGCompiler -from sqlalchemy.sql import compiler +from sqlalchemy.sql import compiler, operators +from sqlalchemy.sql.elements import ( + BinaryExpression, + BindParameter, + Grouping, + Null, + Slice, + TypeCoerce, +) from sqlalchemy.types import String from .sa_version import SA_1_4, SA_VERSION @@ -335,6 +343,72 @@ def visit_JSONB(self, type_, **kw): return "OBJECT" +def _slice_bound(value): + if value is None: + return "" + if isinstance(value, bool) or not isinstance(value, int): + raise sa.exc.CompileError(f"CrateDB array slices take integer bounds, not {value!r}") + return str(value) + + +class _SubscriptType(sa.types.TypeEngine): + """ + Renders an array index, an array slice, or an object key as a literal. + + With `int_as_key`, an integer renders as an object key. + """ + + cache_ok = True + # Reject `None`, instead of rendering it as `NULL`. + should_evaluate_none = True + + def __init__(self, int_as_key=False): + self.int_as_key = int_as_key + + def literal_processor(self, dialect): + def process(value): + if isinstance(value, slice): + if value.step is not None: + raise sa.exc.CompileError("CrateDB array slices do not support a step") + return "%s:%s" % (_slice_bound(value.start), _slice_bound(value.stop)) + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise sa.exc.CompileError( + f"CrateDB subscripts take an integer index or a string key, not {value!r}" + ) + if isinstance(value, int) and not self.int_as_key: + return str(value) + return "'%s'" % str(value).replace("'", "''") + + return process + + +def _is_top_level_object(element): + """ + Whether `element` is an object, and not a value taken out of one, which + may be an array. CrateDB does not accept an array index on an object. + """ + while isinstance(element, Grouping): + element = element.element + if isinstance(element, BinaryExpression) and element.operator in ( + operators.getitem, + operators.json_getitem_op, + ): + return False + return isinstance(element.type, sa.types.JSON) + + +def _unwrap_bindparam(element): + """ + Return the bind parameter inside `element`, if it renders as a bare one. + """ + inner = element + while isinstance(inner, Grouping): + inner = inner.element + if isinstance(inner, TypeCoerce): + inner = inner.typed_expression + return inner if isinstance(inner, BindParameter) else element + + class CrateCompiler(compiler.SQLCompiler): def visit_typeclause(self, typeclause, **kw): """ @@ -345,11 +419,45 @@ def visit_typeclause(self, typeclause, **kw): visit_on_conflict_do_update = PGCompiler.visit_on_conflict_do_update _on_conflict_target = PGCompiler._on_conflict_target + def _render_subscript(self, binary, **kw): + left = self.process(binary.left, **kw) + index = _unwrap_bindparam(binary.right) + if isinstance(index, Slice): + return "%s[%s]" % (left, self._render_slice(index, **kw)) + if not isinstance(index, BindParameter): + # CrateDB accepts bind parameters inside an expression. + return "%s[%s]" % (left, self.process(index, **kw)) + + # CrateDB does not accept a bind parameter as an array index or object + # key, so render it as a literal. + type_ = _SubscriptType(int_as_key=_is_top_level_object(binary.left)) + if not index.required: + type_.literal_processor(self.dialect)(index.effective_value) + elif SA_VERSION < SA_1_4: + raise sa.exc.CompileError( + "A subscript taking its value on execution needs SQLAlchemy 1.4 or later" + ) + if index.required or getattr(self, "cache_key", None) is not None: + kw["literal_execute"] = True + else: + kw["literal_binds"] = True + index = sa.type_coerce(index, type_).typed_expression + return "%s[%s]" % (left, self.process(index, **kw)) + + def _render_slice(self, slice_, **kw): + # CrateDB accepts bind parameters as slice bounds. + if not isinstance(slice_.step, Null): + raise sa.exc.CompileError("CrateDB array slices do not support a step") + return ":".join( + "" if isinstance(bound, Null) else self.process(bound, **kw) + for bound in (slice_.start, slice_.stop) + ) + def visit_getitem_binary(self, binary, operator, **kw): - return "{0}['{1}']".format(self.process(binary.left, **kw), binary.right.value) + return self._render_subscript(binary, **kw) def visit_json_getitem_op_binary(self, binary, operator, _cast_applied=False, **kw): - return "{0}['{1}']".format(self.process(binary.left, **kw), binary.right.value) + return self._render_subscript(binary, **kw) def visit_any(self, element, **kw): return "%s%sANY (%s)" % ( diff --git a/tests/__init__.py b/tests/__init__.py index 5005cecf..1d1914e2 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -27,7 +27,7 @@ from .function_test import SqlAlchemyFunctionTest from .insert_from_select_test import SqlAlchemyInsertFromSelectTest from .match_test import SqlAlchemyMatchTest -from .query_caching import SqlAlchemyQueryCompilationCaching +from .query_caching_test import SqlAlchemyQueryCompilationCaching from .update_test import SqlAlchemyUpdateTest from .warnings_test import SqlAlchemyWarningsTest diff --git a/tests/array_test.py b/tests/array_test.py index 1b9332f9..7f5a9e54 100644 --- a/tests/array_test.py +++ b/tests/array_test.py @@ -27,7 +27,7 @@ from sqlalchemy.orm import Session from sqlalchemy.sql import operators -from sqlalchemy_cratedb import SA_1_4, SA_VERSION +from sqlalchemy_cratedb import SA_1_4, SA_VERSION, ObjectArray try: from sqlalchemy.orm import declarative_base @@ -57,10 +57,24 @@ class User(Base): self.User = User self.session = Session(bind=self.engine) + self.subscripts = sa.Table( + "subscripts", + self.metadata, + sa.Column("idx", sa.Integer), + sa.Column("arr0", sa.ARRAY(sa.Integer, zero_indexes=True)), + sa.Column("arr2d", sa.ARRAY(sa.Integer, dimensions=2)), + sa.Column("data_list", ObjectArray), + ) def assertSQL(self, expected_str, actual_expr): self.assertEqual(expected_str, str(actual_expr).replace("\n", "")) + def assertWhere(self, expected_where, clause): + s = self.session.query(self.subscripts.c.idx).filter(clause) + self.assertSQL( + "SELECT subscripts.idx AS subscripts_idx FROM subscripts WHERE " + expected_where, s + ) + @skipIf(SA_VERSION < SA_1_4, "`as_generic` not available with SQLAlchemy 1.3") def test_as_generic(self): t1 = sa.Table( @@ -116,6 +130,129 @@ def test_any_with_operator(self): s, ) + def test_index(self): + s = self.session.query(self.User.name).filter(self.User.scores[1] == 5) + self.assertSQL( + "SELECT users.name AS users_name FROM users WHERE users.scores[1] = %(param_1)s", s + ) + + def test_index_zero_indexes(self): + self.assertWhere("subscripts.arr0[1] = %(param_1)s", self.subscripts.c.arr0[0] == 5) + + def test_index_nested(self): + self.assertWhere("subscripts.arr2d[1][2] = %(param_1)s", self.subscripts.c.arr2d[1][2] == 5) + + def test_index_column(self): + t = self.subscripts + self.assertWhere( + "subscripts.arr0[(subscripts.idx + %(idx_1)s)] = %(param_1)s", t.c.arr0[t.c.idx] == 5 + ) + + def test_index_expression(self): + s = self.session.query(self.User.name).filter( + self.User.scores[sa.func.char_length(self.User.name) + 1] == 5 + ) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[(char_length(users.name) + %(char_length_1)s)] = %(param_1)s", + s, + ) + + def test_index_expression_bindparam(self): + s = self.session.query(self.User.name).filter( + self.User.scores[sa.func.char_length(self.User.name) + sa.bindparam("off")] == 5 + ) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[(char_length(users.name) + %(off)s)] = %(param_1)s", + s, + ) + + def test_index_zero_indexes_bindparam(self): + self.assertWhere( + "subscripts.arr0[(%(pos)s + %(param_1)s)] = %(param_2)s", + self.subscripts.c.arr0[sa.bindparam("pos")] == 5, + ) + + def test_index_type_coerce(self): + s = self.session.query(self.User.name).filter( + self.User.scores[sa.type_coerce(sa.literal(2), sa.Integer)] == 5 + ) + self.assertSQL( + "SELECT users.name AS users_name FROM users WHERE users.scores[2] = %(param_1)s", s + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no `literal_execute`") + def test_index_bindparam(self): + s = self.session.query(self.User.name).filter(self.User.scores[sa.bindparam("pos")] == 5) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[__[POSTCOMPILE_pos]] = %(param_1)s", + s, + ) + + @skipIf(SA_VERSION >= SA_1_4, "SQLAlchemy 1.4+ renders the parameter on execution") + def test_index_bindparam_sa13(self): + s = self.session.query(self.User.name).filter(self.User.scores[sa.bindparam("pos")] == 5) + with self.assertRaises(sa.exc.CompileError) as cm: + str(s) + self.assertEqual( + "A subscript taking its value on execution needs SQLAlchemy 1.4 or later", + str(cm.exception), + ) + + def test_index_invalid(self): + s = self.session.query(self.User.name).filter(self.User.scores[1.5] == 5) + with self.assertRaises(sa.exc.CompileError) as cm: + str(s) + self.assertEqual( + "CrateDB subscripts take an integer index or a string key, not 1.5", str(cm.exception) + ) + + def test_slice(self): + s = self.session.query(self.User.name).filter(self.User.scores[1:2] == [5]) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[%(scores_1)s:%(scores_2)s] = %(param_1)s", + s, + ) + + def test_slice_open(self): + s = self.session.query(self.User.name).filter(self.User.scores[:2] == [5]) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[:%(scores_1)s] = %(param_1)s", + s, + ) + s = self.session.query(self.User.name).filter(self.User.scores[2:] == [5]) + self.assertSQL( + "SELECT users.name AS users_name FROM users " + "WHERE users.scores[%(scores_1)s:] = %(param_1)s", + s, + ) + + def test_slice_step(self): + s = self.session.query(self.User.name).filter(self.User.scores[1:3:2] == [5]) + with self.assertRaises(sa.exc.CompileError) as cm: + str(s) + self.assertEqual("CrateDB array slices do not support a step", str(cm.exception)) + + def test_update_element(self): + stmt = sa.update(self.User.__table__).values({self.User.scores[1]: 99}) + self.assertSQL("UPDATE users SET scores[1] = %(param_1)s", stmt.compile(bind=self.engine)) + + def test_object_array_key(self): + self.assertWhere( + "%(param_1)s = ANY (subscripts.data_list['foo'])", + self.subscripts.c.data_list["foo"].any(1), + ) + + def test_object_array_index_key(self): + self.assertWhere( + "subscripts.data_list[1]['foo'] = %(param_1)s", + self.subscripts.c.data_list[1]["foo"] == 1, + ) + def test_multidimensional_arrays(self): t1 = sa.Table( "t", diff --git a/tests/dict_test.py b/tests/dict_test.py index 9eb97d78..37367bb1 100644 --- a/tests/dict_test.py +++ b/tests/dict_test.py @@ -92,6 +92,53 @@ def test_select_with_dict_column_where_clause_other_col(self): "SELECT mytable.name FROM mytable " + "WHERE mytable.data['x'] = mytable.name", s ) + def test_select_with_dict_column_quote_in_key(self): + mytable = self.mytable + self.assertSQL( + "SELECT mytable.data['it''s'] AS anon_1 FROM mytable", select(mytable.c.data["it's"]) + ) + + def test_select_with_dict_column_injection_in_key(self): + mytable = self.mytable + s = select(mytable.c.name).where(mytable.c.data["x'] = 1 OR TRUE OR mytable.data['x"] == 1) + self.assertSQL( + "SELECT mytable.name FROM mytable " + "WHERE mytable.data['x''] = 1 OR TRUE OR mytable.data[''x'] = %(param_1)s", + s, + ) + + def test_select_with_dict_column_integer_key(self): + # CrateDB does not accept an array index on an object. + mytable = self.mytable + self.assertSQL("SELECT mytable.data['1'] AS anon_1 FROM mytable", select(mytable.c.data[1])) + + def test_select_with_dict_column_invalid_key(self): + mytable = self.mytable + with self.assertRaises(sa.exc.CompileError) as cm: + select(mytable.c.data[True]).compile(bind=self.engine) + self.assertEqual( + "CrateDB subscripts take an integer index or a string key, not True", str(cm.exception) + ) + + def test_select_with_dict_column_nested_index(self): + # A value inside an object may be an array. + mytable = self.mytable + self.assertSQL( + "SELECT mytable.data['nums'][2] AS anon_1 FROM mytable", + select(mytable.c.data["nums"][2]), + ) + + def test_select_with_dict_column_nested_slice(self): + mytable = self.mytable + self.assertSQL( + "SELECT mytable.data['nums'][1:2] AS anon_1 FROM mytable", + select(mytable.c.data["nums"][1:2]), + ) + self.assertSQL( + "SELECT mytable.data['nums'][:2] AS anon_1 FROM mytable", + select(mytable.c.data["nums"][:2]), + ) + def test_update_with_dict_column(self): mytable = self.mytable stmt = ( diff --git a/tests/query_caching.py b/tests/query_caching.py deleted file mode 100644 index 35df4d02..00000000 --- a/tests/query_caching.py +++ /dev/null @@ -1,143 +0,0 @@ -# -*- 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 sqlalchemy_cratedb import ObjectArray, ObjectType -from sqlalchemy_cratedb.sa_version import SA_1_4, SA_VERSION - -try: - from sqlalchemy.orm import declarative_base -except ImportError: - from sqlalchemy.ext.declarative import declarative_base - -from tests.settings import crate_host - - -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(ObjectType) - 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_legacy(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. - - This variant uses direct value matching on the `OBJECT`s attribute. - """ - 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_object_multiple_select_modern(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. - - This variant uses comparator method matching on the `OBJECT`s attribute. - """ - self.setup_data() - Character = self.Character - - selectable = sa.select(Character).where(Character.data["x"].as_integer() == 1) - result = self.session.execute(selectable).scalar_one().data - self.assertEqual({"x": 1}, result) - - selectable = sa.select(Character).where(Character.data["y"].as_integer() == 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) diff --git a/tests/query_caching_test.py b/tests/query_caching_test.py new file mode 100644 index 00000000..1362785f --- /dev/null +++ b/tests/query_caching_test.py @@ -0,0 +1,304 @@ +# -*- 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 sqlalchemy_cratedb import ObjectArray, ObjectType +from sqlalchemy_cratedb.sa_version import SA_1_4, SA_VERSION + +try: + from sqlalchemy.orm import declarative_base +except ImportError: + from sqlalchemy.ext.declarative import declarative_base + +from tests.settings import crate_host + + +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() + # Without an `ObjectType` column, which would turn off caching for any + # statement that selects the whole row. + self.subscripts = sa.Table( + "subscripts", + self.metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("arr", sa.ARRAY(sa.Integer)), + sa.Column("objarr", ObjectArray), + sa.Column("js", sa.JSON), + ) + + def tearDown(self): + self.session.close() + self.metadata.drop_all(self.engine) + + 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(ObjectType) + 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;")) + + def setup_subscript_data(self): + """ + Insert one record into the `subscripts` table. + """ + self.metadata.drop_all(self.engine) + self.metadata.create_all(self.engine) + with self.engine.connect() as connection: + connection.execute( + sa.text( + "INSERT INTO testdrive.subscripts (id, arr, objarr, js) VALUES " + "(1, [10, 20, 30], [{foo='F1', bar='B1'}, {foo='F2', bar='B2'}], " + "{x='JX', y='JY', \"it's\"='Q'})" + ) + ) + connection.execute(sa.text("REFRESH TABLE testdrive.subscripts")) + + def execute_twice(self, first, second, first_params=None, second_params=None): + """ + Execute two statements that share a cache key, and return both results. + The second one must hit the cache, or the test proves nothing. + """ + results = [] + with self.engine.connect() as connection: + for statement, params in ((first, first_params), (second, second_params)): + result = connection.execute(statement, params or {}) + results.append(result.rowcount if statement.is_dml else result.scalar()) + self.assertIs(connection.dialect.CACHE_HIT, result.context.cache_hit) + return results + + @skipIf(SA_VERSION < SA_1_4, "On SA13, the 'ResultProxy' object has no attribute 'scalar_one'") + def test_object_multiple_select_legacy(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. + + This variant uses direct value matching on the `OBJECT`s attribute. + """ + 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_object_multiple_select_modern(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. + + This variant uses comparator method matching on the `OBJECT`s attribute. + """ + self.setup_data() + Character = self.Character + + selectable = sa.select(Character).where(Character.data["x"].as_integer() == 1) + result = self.session.execute(selectable).scalar_one().data + self.assertEqual({"x": 1}, result) + + selectable = sa.select(Character).where(Character.data["y"].as_integer() == 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 `ObjectType` column of the entity turns off caching for these + statements. `test_subscript_objectarray_key` covers the cached case. + """ + 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) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_index(self): + self.setup_subscript_data() + arr = self.subscripts.c.arr + self.assertEqual([10, 20], self.execute_twice(sa.select(arr[1]), sa.select(arr[2]))) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_index_bindparam(self): + self.setup_subscript_data() + statement = sa.select(self.subscripts.c.arr[sa.bindparam("pos")]) + self.assertEqual([10, 30], self.execute_twice(statement, statement, {"pos": 1}, {"pos": 3})) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_index_expression_bindparam(self): + self.setup_subscript_data() + t = self.subscripts + statement = sa.select(t.c.arr[t.c.id + sa.bindparam("off")]) + self.assertEqual([20, 30], self.execute_twice(statement, statement, {"off": 1}, {"off": 2})) + with self.engine.connect() as connection: + connection = connection.execution_options(compiled_cache=None) + self.assertEqual(20, connection.execute(statement, {"off": 1}).scalar()) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_executemany(self): + """ + SQLAlchemy does not support `literal_execute` with `executemany()`. + """ + self.setup_subscript_data() + t = self.subscripts + statement = sa.delete(t).where(t.c.js["x"] == sa.bindparam("value")) + parameters = [{"value": "JX"}, {"value": "unknown"}] + with self.engine.connect() as connection: + with self.assertRaises(sa.exc.StatementError) as cm: + connection.execute(statement, parameters) + self.assertIn("can't be used with executemany()", str(cm.exception)) + + connection.execution_options(compiled_cache=None).execute(statement, parameters) + connection.execute(sa.text("REFRESH TABLE testdrive.subscripts")) + self.assertEqual( + 0, connection.execute(sa.select(sa.func.count()).select_from(t)).scalar() + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has a different `select()` signature") + def test_subscript_expression_executemany(self): + self.setup_subscript_data() + t = self.subscripts + statement = sa.delete(t).where( + t.c.arr[t.c.id + sa.bindparam("off")] == sa.bindparam("value") + ) + with self.engine.connect() as connection: + connection.execute(statement, [{"off": 1, "value": 99}, {"off": 1, "value": 20}]) + connection.execute(sa.text("REFRESH TABLE testdrive.subscripts")) + self.assertEqual( + 0, connection.execute(sa.select(sa.func.count()).select_from(t)).scalar() + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_index_where(self): + self.setup_subscript_data() + t = self.subscripts + self.assertEqual( + [1, None], + self.execute_twice( + sa.select(t.c.id).where(t.c.arr[2] == 20), + sa.select(t.c.id).where(t.c.arr[3] == 20), + ), + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_slice(self): + self.setup_subscript_data() + arr = self.subscripts.c.arr + self.assertEqual( + [[10, 20], [20, 30]], self.execute_twice(sa.select(arr[1:2]), sa.select(arr[2:3])) + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_array_update(self): + self.setup_subscript_data() + t = self.subscripts + self.assertEqual( + [1, 1], + self.execute_twice( + sa.update(t).values({t.c.arr[1]: 99}), sa.update(t).values({t.c.arr[2]: 88}) + ), + ) + with self.engine.connect() as connection: + connection.execute(sa.text("REFRESH TABLE testdrive.subscripts")) + self.assertEqual([99, 88, 30], connection.execute(sa.select(t.c.arr)).scalar()) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_objectarray_key(self): + self.setup_subscript_data() + objarr = self.subscripts.c.objarr + self.assertEqual( + [["F1", "F2"], ["B1", "B2"]], + self.execute_twice(sa.select(objarr["foo"]), sa.select(objarr["bar"])), + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_objectarray_index_key(self): + self.setup_subscript_data() + objarr = self.subscripts.c.objarr + self.assertEqual( + ["F1", "F2"], + self.execute_twice(sa.select(objarr[1]["foo"]), sa.select(objarr[2]["foo"])), + ) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_json_key(self): + self.setup_subscript_data() + js = self.subscripts.c.js + self.assertEqual(["JX", "JY"], self.execute_twice(sa.select(js["x"]), sa.select(js["y"]))) + + @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 has no statement cache") + def test_subscript_json_key_with_quote(self): + self.setup_subscript_data() + js = self.subscripts.c.js + self.assertEqual(["JX", "Q"], self.execute_twice(sa.select(js["x"]), sa.select(js["it's"])))