From 520a4564ac6578922cdf65d79a2f7f9d6e1de901 Mon Sep 17 00:00:00 2001 From: Fredrik Fornwall Date: Fri, 18 Sep 2026 11:39:21 +0200 Subject: [PATCH] feat: test parameter stream binding Add tests for binding SQL parameters from Arrow streams. Signed-off-by: Fredrik Fornwall --- adbc_drivers_validation/tests/query.py | 79 +++++++++++++++-- tests/test_query.py | 112 +++++++++++++++++++++++++ 2 files changed, 185 insertions(+), 6 deletions(-) create mode 100644 tests/test_query.py diff --git a/adbc_drivers_validation/tests/query.py b/adbc_drivers_validation/tests/query.py index 25bd895..b6671d2 100644 --- a/adbc_drivers_validation/tests/query.py +++ b/adbc_drivers_validation/tests/query.py @@ -41,6 +41,7 @@ def generate_tests( # Only value types that drivers widely support dictionary-encoded; pandas # categoricals of strings produce dictionary. bind_dictionary_queries: set[str] = {"type/bind/string", "type/bind/large_string"}, + bind_stream_queries: set[str] = {"type/bind/string", "type/bind/large_string"}, ) -> None: """Parameterize the tests in this module for the given driver.""" if utils.generate_tests_by_marks(all_quirks, metafunc): @@ -79,13 +80,21 @@ def generate_tests( elif metafunc.definition.name == "test_query": if not isinstance(query.query, model.SelectQuery): continue - elif metafunc.definition.name == "test_query_bind_dictionary": + elif metafunc.definition.name in { + "test_query_bind_dictionary", + "test_query_bind_stream", + }: if not isinstance(query.query, model.SelectQuery): continue if query.query.bind_query(quirks) is None: continue - if query.name not in bind_dictionary_queries: - # There's no need to test every type with every encoding + bind_queries = ( + bind_dictionary_queries + if metafunc.definition.name == "test_query_bind_dictionary" + else bind_stream_queries + ) + if query.name not in bind_queries: + # There's no need to test every type with every bind variant continue combinations.append( @@ -182,9 +191,7 @@ def test_query( ) else: if bind: - # TODO: also test with stream - # TODO: also test with multiple batches in stream - # TODO: also test with empty stream/empty batch + # TODO: also test binding only empty inputs data = subquery.bind_data().combine_chunks().to_batches()[0] with conn.cursor() as cursor: cursor.adbc_statement.set_sql_query(bind) @@ -250,6 +257,66 @@ def test_query_bind_dictionary( compare.compare_tables(subquery.expected_result(), result, query.metadata()) + @pytest.mark.parametrize( + "batch_size,empty_batches", + [ + pytest.param(None, False, id="single_batch"), + pytest.param(2, False, id="multiple_batches"), + pytest.param(2, True, id="multiple_batches_with_empty"), + ], + ) + def test_query_bind_stream( + self, + driver: model.DriverQuirks, + conn: adbc_driver_manager.dbapi.Connection, + query: Query, + batch_size: int | None, + empty_batches: bool, + ) -> None: + """Bind a stream and consume all parameter batches in one execution.""" + subquery = query.query + assert isinstance(subquery, model.SelectQuery) + bind = subquery.bind_query(driver) + assert bind is not None + + data = subquery.bind_data().combine_chunks() + assert data.num_rows > 1, "Stream bind tests require multiple parameter rows" + batches = data.to_batches(max_chunksize=batch_size) + if empty_batches: + empty = batches[0].slice(0, 0) + # Empty batches before, between, and after data must not end the stream. + batches = [empty] + [part for batch in batches for part in (batch, empty)] + + select_bind = _uses_select_bind(driver, query) + with setup_connection(query, conn): + if not select_bind and driver.features.select_fixture_setup: + utils.retry_adbc_operation( + lambda: _setup_query(driver, conn, query), driver.is_retryable + ) + + with pyarrow.RecordBatchReader.from_batches(data.schema, batches) as stream: + with conn.cursor() as cursor: + if select_bind: + with driver.setup_statement(query, cursor): + cursor.adbc_statement.set_sql_query(bind) + cursor.adbc_statement.bind_stream(stream) + handle, _ = cursor.adbc_statement.execute_query() + with pyarrow.RecordBatchReader._import_from_c( + handle.address + ) as reader: + result = reader.read_all() + else: + cursor.adbc_statement.set_sql_query(bind) + cursor.adbc_statement.bind_stream(stream) + cursor.adbc_statement.execute_update() + + if not select_bind: + with conn.cursor() as cursor: + with driver.setup_statement(query, cursor): + result = execute_query_without_prepare(cursor, subquery.query()) + + compare.compare_tables(subquery.expected_result(), result, query.metadata()) + def test_execute_schema( self, driver: model.DriverQuirks, diff --git a/tests/test_query.py b/tests/test_query.py new file mode 100644 index 0000000..a169e54 --- /dev/null +++ b/tests/test_query.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026 ADBC Drivers Contributors +# +# Licensed 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. + +import typing +from unittest.mock import MagicMock + +import adbc_driver_manager +import pyarrow +import pytest + +from adbc_drivers_validation import model +from adbc_drivers_validation.tests import query as query_tests + + +@pytest.fixture +def driver() -> MagicMock: + driver = MagicMock(spec=model.DriverQuirks) + driver.name = "test" + driver.short_version = "1" + driver.features = model.DriverFeatures(statement_bind=True) + driver.query_set = model.base_query_set() + driver.bind_parameter.return_value = "?" + return driver + + +@pytest.mark.parametrize("variant", ["stream", "dictionary"]) +@pytest.mark.parametrize("custom_queries", [False, True]) +@pytest.mark.parametrize("bind_supported", [False, True]) +def test_generate_bind_variants( + driver: MagicMock, + variant: str, + custom_queries: bool, + bind_supported: bool, +) -> None: + driver.features.statement_bind = bind_supported + metafunc = MagicMock(spec=pytest.Metafunc) + metafunc.definition = MagicMock() + metafunc.definition.name = f"test_query_bind_{variant}" + metafunc.definition.iter_markers.return_value = [] + options = {} + expected = {"type/bind/string", "type/bind/large_string"} + if custom_queries: + expected = {"type/bind/int32"} + options[f"bind_{variant}_queries"] = expected + + query_tests.generate_tests([driver], metafunc, **options) + + metafunc.parametrize.assert_called_once() + params = metafunc.parametrize.call_args.args[1] + assert {param.values[1].name for param in params} == expected + for param in params: + skip_reasons = [mark.kwargs["reason"] for mark in param.marks] + assert skip_reasons == ([] if bind_supported else ["bind not supported"]) + + +@pytest.mark.parametrize("mode", ["insert", "select"]) +@pytest.mark.parametrize("fixture_setup", [False, True]) +@pytest.mark.parametrize( + "batch_size,empty_batches,batch_lengths", + [ + (None, False, [5]), + (2, False, [2, 2, 1]), + (2, True, [0, 2, 0, 2, 0, 1, 0]), + ], +) +def test_bind_stream( + driver: MagicMock, + monkeypatch: pytest.MonkeyPatch, + mode: typing.Literal["insert", "select"], + fixture_setup: bool, + batch_size: int | None, + empty_batches: bool, + batch_lengths: list[int], +) -> None: + driver.features.statement_bind_test_mode = mode + driver.features.select_fixture_setup = fixture_setup + query = driver.query_set.queries["type/bind/string"] + conn = MagicMock() + statement = conn.cursor.return_value.__enter__.return_value.adbc_statement + setup = MagicMock() + monkeypatch.setattr(query_tests, "_setup_query", setup) + + def bind_stream(stream: pyarrow.RecordBatchReader) -> None: + batches = list(stream) + assert [batch.num_rows for batch in batches] == batch_lengths + data = pyarrow.Table.from_batches(batches, schema=stream.schema) + handle = adbc_driver_manager.ArrowArrayStreamHandle() + data.to_reader()._export_to_c(handle.address) + statement.execute_query.return_value = handle, -1 + + statement.bind_stream.side_effect = bind_stream + + query_tests.TestQuery().test_query_bind_stream( + driver, conn, query, batch_size, empty_batches + ) + + statement.bind_stream.assert_called_once() + statement.bind.assert_not_called() + statement.execute_query.assert_called_once() + assert statement.execute_update.call_count == (1 if mode == "insert" else 0) + assert setup.call_count == (1 if mode == "insert" and fixture_setup else 0)