Merge pull request #43 Fix false cache hit from LuckySting/fix/false-cache-hit-by-query-name

Author: rekby

test/test_core.py

@@ -83,6 +83,19 @@ def test_sa_crud(self, connection):
             (5, "c"),
         ]
 
+    def test_cached_query(self, connection_no_trans: sa.Connection, connection: sa.Connection):
+        table = self.tables.test
+
+        with connection_no_trans.begin() as transaction:
+            connection_no_trans.execute(sa.insert(table).values([{"id": 1, "text": "foo"}]))
+            connection_no_trans.execute(sa.insert(table).values([{"id": 2, "text": None}]))
+            connection_no_trans.execute(sa.insert(table).values([{"id": 3, "text": "bar"}]))
+            transaction.commit()
+
+        result = connection.execute(sa.select(table)).fetchall()
+
+        assert result == [(1, "foo"), (2, None), (3, "bar")]
+
     def test_sa_crud_with_add_declare(self):
         engine = sa.create_engine(config.db_url, _add_declare_for_yql_stmt_vars=True)
         with engine.connect() as connection:

ydb_sqlalchemy/dbapi/cursor.py

@@ -1,6 +1,7 @@
 import collections.abc
 import dataclasses
 import functools
+import hashlib
 import itertools
 import logging
 import posixpath
@@ -139,10 +140,29 @@ def _get_ydb_query(self, operation: YdbQuery) -> Union[ydb.DataQuery, str]:
         pragma = ""
         if self.root_directory:
             pragma = f'PRAGMA TablePathPrefix = "{self.root_directory}";\n'
+
+        yql_with_pragma = pragma + operation.yql_text
+
         if operation.is_ddl or not operation.parameters_types:
-            return pragma + operation.yql_text
+            return yql_with_pragma
 
-        return ydb.DataQuery(pragma + operation.yql_text, operation.parameters_types)
+        return self._make_data_query(yql_with_pragma, operation.parameters_types)
+
+    def _make_data_query(
+        self,
+        yql_text: str,
+        parameters_types: Dict[str, Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]],
+    ) -> ydb.DataQuery:
+        """
+        ydb.DataQuery uses hashed SQL text as cache key, which may cause issues if parameters change type within
+        the same session, so we include parameter types to the key to prevent false positive cache hit.
+        """
+
+        sorted_parameters = sorted(parameters_types.items())  # dict keys are unique, so the sorting is stable
+
+        yql_with_params = yql_text + "".join([k + str(v) for k, v in sorted_parameters])
+        name = hashlib.sha256(yql_with_params.encode("utf-8")).hexdigest()
+        return ydb.DataQuery(yql_text, parameters_types, name=name)
 
     @_handle_ydb_errors
     def _execute_dml(