Fix SQLAlchemy query cache collisions for bound parameters

src/datasets/packaged_modules/sql/sql.py

@@ -1,4 +1,5 @@
 import sys
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional, Union
 
@@ -9,7 +10,9 @@
 import datasets.config
 from datasets.builder import Key
 from datasets.features.features import require_storage_cast
+from datasets.fingerprint import generate_random_fingerprint
 from datasets.table import table_cast
+from datasets.utils.py_utils import map_nested
 
 
 if TYPE_CHECKING:
@@ -48,7 +51,7 @@ def create_config_id(
         custom_features: Optional[datasets.Features] = None,
     ) -> str:
         config_kwargs = config_kwargs.copy()
-        # We need to stringify the Selectable object to make its hash deterministic
+        # Compile Selectable objects to make their hash independent of SQLAlchemy object identities.
 
         # The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
         sql = config_kwargs["sql"]
@@ -58,8 +61,26 @@ def create_config_id(
 
                 if isinstance(sql, sqlalchemy.sql.Selectable):
                     engine = sqlalchemy.create_engine(config_kwargs["con"].split("://")[0] + "://")
-                    sql_str = str(sql.compile(dialect=engine.dialect))
-                    config_kwargs["sql"] = sql_str
+                    compiled = sql.compile(dialect=engine.dialect)
+                    # Effective compiled parameters include explicit overrides; don't hash them twice.
+                    params = config_kwargs.pop("params", None)
+                    params = params if isinstance(params, Mapping) else {}
+                    if any(
+                        bind.callable and bind.key not in params and name not in params
+                        for bind, name in compiled.bind_names.items()
+                    ):
+                        # Callable values are only known at execution time. Avoid reusing their cache
+                        # rather than evaluating them here and again when executing the query.
+                        sql_params = generate_random_fingerprint()
+                    else:
+                        # DBAPI binary buffers have the same contents as bytes but aren't picklable.
+                        sql_params = map_nested(
+                            lambda value: value.tobytes() if isinstance(value, memoryview) else value,
+                            compiled.construct_params(params=params) if params else compiled.params,
+                            map_tuple=True,
+                        )
+                    # The SQL contains placeholders, so its bound values must also identify the cache.
+                    config_kwargs["sql"] = (str(compiled), sql_params)
                 else:
                     raise TypeError(
                         f"Supported types for 'sql' are string and sqlalchemy.sql.Selectable but got {type(sql)}: {sql}"

tests/io/test_sql.py

@@ -1,6 +1,10 @@
 import contextlib
 import os
 import sqlite3
+from collections import UserDict
+from datetime import datetime
+from decimal import Decimal
+from unittest.mock import Mock, patch
 
 import pytest
 
@@ -14,6 +18,22 @@
 STRING_FROM_PANDAS = "large_string" if datasets.config.PANDAS_VERSION.major >= 3 else "string"
 
 
+@pytest.fixture
+def sql_items(tmp_path):
+    import sqlalchemy as sa
+
+    uri = "sqlite:///" + str(tmp_path / "items.db")
+    engine = sa.create_engine(uri)
+    items = sa.Table("items", sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True))
+    items.metadata.create_all(engine)
+    try:
+        with engine.begin() as connection:
+            connection.execute(items.insert(), [{"id": i} for i in range(1, 5)])
+        yield uri, items
+    finally:
+        engine.dispose()
+
+
 def _check_sql_dataset(dataset, expected_features):
     assert isinstance(dataset, Dataset)
     assert dataset.num_rows == 4
@@ -115,3 +135,179 @@ def test_dataset_to_sql_preserves_nullable_int(dtype, big, tmp_path, set_sqlalch
     with contextlib.closing(sqlite3.connect(output_sqlite_path)) as con:
         rows = con.execute("SELECT a, typeof(a) FROM dataset").fetchall()
     assert rows == [(big, "integer"), (None, "null"), (5, "integer")]
+
+
+@require_sqlalchemy
+@pytest.mark.parametrize(
+    "query_kind,reverse,chunksize",
+    [
+        ("where", False, None),
+        ("limit", False, None),
+        ("limit", True, None),
+        ("offset", False, None),
+        ("in", False, 2),
+        ("in", True, 2),
+        ("tuple_in", False, None),
+        ("literal_execute", False, None),
+    ],
+)
+def test_dataset_from_sql_bound_parameters(query_kind, reverse, chunksize, sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+    query = sa.select(items).order_by(items.c.id)
+    if query_kind == "where":
+        queries = [query.where(items.c.id == value) for value in (1, 2)]
+        expected_ids = [[1], [2]]
+    elif query_kind == "limit":
+        queries = [query.limit(value) for value in (1, 2)]
+        expected_ids = [[1], [1, 2]]
+    elif query_kind == "offset":
+        queries = [query.limit(2).offset(value) for value in (0, 1)]
+        expected_ids = [[1, 2], [2, 3]]
+    elif query_kind == "in":
+        queries = [query.where(items.c.id.in_(values)) for values in ([1], [2, 3])]
+        expected_ids = [[1], [2, 3]]
+    elif query_kind == "tuple_in":
+        queries = [query.where(sa.tuple_(items.c.id, items.c.id).in_(values)) for values in ([(1, 1)], [(2, 2)])]
+        expected_ids = [[1], [2]]
+    else:
+        queries = [query.where(items.c.id == sa.bindparam("id", value, literal_execute=True)) for value in (1, 2)]
+        expected_ids = [[1], [2]]
+    if reverse:
+        queries.reverse()
+        expected_ids.reverse()
+
+    # Different queries must not reuse one another's cached rows (#8877).
+    first = Dataset.from_sql(queries[0], uri, cache_dir=tmp_path / "cache", chunksize=chunksize)
+    second = Dataset.from_sql(queries[1], uri, cache_dir=tmp_path / "cache", chunksize=chunksize)
+    assert first.to_dict() == {"id": expected_ids[0]}
+    assert second.to_dict() == {"id": expected_ids[1]}
+    assert first.cache_files != second.cache_files
+    assert first._fingerprint != second._fingerprint
+
+
+@require_sqlalchemy
+def test_dataset_from_sql_equivalent_queries_reuse_cache(sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+
+    def read():
+        query = sa.select(items).where(items.c.id > 1).where(items.c.id.in_([2, 3])).order_by(items.c.id)
+        return Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache")
+
+    first = read()
+    # Rebuild the statement so SQLAlchemy object identities cannot determine the cache key.
+    with patch("datasets.packaged_modules.sql.sql.pd.read_sql", side_effect=AssertionError("cache miss")):
+        second = read()
+    assert first.to_dict() == second.to_dict() == {"id": [2, 3]}
+    assert first.cache_files == second.cache_files
+    assert first._fingerprint == second._fingerprint
+
+
+@require_sqlalchemy
+@pytest.mark.parametrize("query_kind", ["selectable", "string"])
+def test_dataset_from_sql_explicit_params(query_kind, sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+    query = (
+        sa.select(items).where(items.c.id == sa.bindparam("id", 4))
+        if query_kind == "selectable"
+        else "SELECT id FROM items WHERE id = :id"
+    )
+    # SQLAlchemy accepts parameter mappings beyond plain dictionaries.
+    params = UserDict(id=1) if query_kind == "selectable" else {"id": 1}
+    first = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache", params=params)
+    second = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache", params=type(params)(id=2))
+    assert first.to_dict() == {"id": [1]}
+    assert second.to_dict() == {"id": [2]}
+    assert first.cache_files != second.cache_files
+    assert params == {"id": 1}
+    if query_kind == "selectable":
+        # An overridden default must not prevent reuse of the effective query's cache.
+        query = sa.select(items).where(items.c.id == sa.bindparam("id", 3))
+        with patch("datasets.packaged_modules.sql.sql.pd.read_sql", side_effect=AssertionError("cache miss")):
+            repeated = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache", params=params)
+        assert repeated.to_dict() == first.to_dict()
+
+
+@require_sqlalchemy
+def test_dataset_from_sql_bound_and_explicit_params(sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+    datasets = [
+        Dataset.from_sql(
+            sa.select(items).where(items.c.id == value).where(items.c.id > sa.bindparam("minimum")),
+            uri,
+            cache_dir=tmp_path / "cache",
+            params={"minimum": 0},
+        )
+        for value in (1, 2)
+    ]
+    assert [dataset.to_dict() for dataset in datasets] == [{"id": [1]}, {"id": [2]}]
+
+
+@require_sqlalchemy
+def test_dataset_from_sql_callable_parameters(sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+    values = iter([1, 2])
+    query = sa.select(items).where(items.c.id == sa.bindparam("id", callable_=lambda: next(values)))
+    first = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache")
+    second = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache")
+    assert first.to_dict() == {"id": [1]}
+    assert second.to_dict() == {"id": [2]}
+    # Cache identity must not evaluate the callable in addition to SQL execution.
+    assert next(values, None) is None
+
+
+@require_sqlalchemy
+def test_dataset_from_sql_explicit_params_override_callable(sql_items, tmp_path):
+    import sqlalchemy as sa
+
+    uri, items = sql_items
+    value = Mock(side_effect=AssertionError("overridden callable was evaluated"))
+    query = sa.select(items).where(items.c.id == sa.bindparam("id", callable_=value))
+    dataset = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache", params={"id": 2})
+    assert dataset.to_dict() == {"id": [2]}
+    with patch("datasets.packaged_modules.sql.sql.pd.read_sql", side_effect=AssertionError("cache miss")):
+        repeated = Dataset.from_sql(query, uri, cache_dir=tmp_path / "cache", params={"id": 2})
+    assert repeated.to_dict() == {"id": [2]}
+
+
+@require_sqlalchemy
+@pytest.mark.parametrize(
+    "sql_type,values,explicit",
+    [
+        ("DateTime", [datetime(2024, 1, 1), datetime(2024, 1, 2)], False),
+        ("Numeric", [Decimal("1.25"), Decimal("2.50")], False),
+        ("LargeBinary", [memoryview(b"\x00\xff"), memoryview(b"\x00\xfe")], False),
+        ("LargeBinary", [memoryview(b"\x00\xff"), memoryview(b"\x00\xfe")], True),
+    ],
+)
+def test_dataset_from_sql_typed_bound_parameters(sql_type, values, explicit, tmp_path):
+    import sqlalchemy as sa
+
+    def read(value):
+        bound = sa.bindparam("value", None if explicit else value, type_=getattr(sa, sql_type)())
+        params = {"value": value} if explicit else None
+        return Dataset.from_sql(
+            sa.select(bound.label("value")),
+            "sqlite://",
+            cache_dir=tmp_path / "cache",
+            params=params,
+            coerce_float=False,
+        )
+
+    datasets = [read(value) for value in values]
+    expected = [bytes(value) if isinstance(value, memoryview) else value for value in values]
+    assert [dataset.to_dict() for dataset in datasets] == [{"value": [value]} for value in expected]
+    assert datasets[0].cache_files != datasets[1].cache_files
+    with patch("datasets.packaged_modules.sql.sql.pd.read_sql", side_effect=AssertionError("cache miss")):
+        # Buffer contents identify the query, independently of the buffer object or representation.
+        repeated = read(expected[0])
+    assert repeated.to_dict() == datasets[0].to_dict()

tests/packaged_modules/test_sql.py

@@ -1,9 +1,14 @@
+from functools import partial
+from uuid import UUID
+
 import pytest
 
-from datasets.builder import InvalidConfigName
+from datasets.builder import BuilderConfig, InvalidConfigName
 from datasets.data_files import DataFilesList
 from datasets.packaged_modules.sql.sql import SqlConfig
 
+from ..utils import require_sqlalchemy
+
 
 def test_config_raises_when_invalid_name() -> None:
     with pytest.raises(InvalidConfigName, match="Bad characters"):
@@ -14,3 +19,58 @@ def test_config_raises_when_invalid_name() -> None:
 def test_config_raises_when_invalid_data_files(data_files) -> None:
     with pytest.raises(ValueError, match="Expected a DataFilesDict"):
         _ = SqlConfig(name="name", data_files=data_files)
+
+
+@require_sqlalchemy
+@pytest.mark.parametrize("uri", ["postgresql://", "mysql://", "oracle://"])
+def test_config_id_bound_parameters(uri, monkeypatch):
+    import sqlalchemy as sa
+
+    # Exercise real dialect compilers without requiring external servers or DBAPI drivers.
+    monkeypatch.setattr(sa, "create_engine", partial(sa.create_mock_engine, executor=None))
+
+    def config_id(value):
+        # Nested native values need deterministic hashing, not SQL literal rendering.
+        payload = {"uuid": UUID(int=value), "nested": (None, "λ'?:%", memoryview(b"\x00\xff"))}
+        query = sa.select(sa.bindparam("value", payload, type_=sa.types.NullType()))
+        kwargs = {"sql": query, "con": uri}
+        return SqlConfig(**kwargs).create_config_id(kwargs)
+
+    assert config_id(1) != config_id(2)
+    assert config_id(1) == config_id(1)
+
+
+@require_sqlalchemy
+def test_config_id_expanding_parameters():
+    import sqlalchemy as sa
+
+    def config_id(values):
+        query = sa.select(sa.column("id")).where(sa.column("id").in_(values))
+        kwargs = {"sql": query, "con": "sqlite://"}
+        return SqlConfig(**kwargs).create_config_id(kwargs)
+
+    assert len({config_id(values) for values in [[], [1], [2], [1, 2]]}) == 4
+    assert config_id([]) == config_id([])
+
+
+@require_sqlalchemy
+def test_config_id_explicit_params_order():
+    import sqlalchemy as sa
+
+    def config_id(params):
+        query = sa.select(sa.bindparam("first"), sa.bindparam("second"))
+        kwargs = {"sql": query, "con": "sqlite://", "params": params}
+        result = SqlConfig(**kwargs).create_config_id(kwargs)
+        assert kwargs["sql"] is query
+        assert kwargs["params"] is params
+        return result
+
+    assert config_id({"first": 1, "second": 2}) == config_id({"second": 2, "first": 1})
+    assert config_id({"first": 1, "second": 2}) != config_id({"first": 2, "second": 1})
+
+
+@require_sqlalchemy
+def test_config_id_raw_sql_unchanged():
+    kwargs = {"sql": "SELECT id FROM items WHERE id = ?", "con": "sqlite://", "params": (1,)}
+    config = SqlConfig(**kwargs)
+    assert config.create_config_id(kwargs) == BuilderConfig.create_config_id(config, kwargs)