feat: test parameter stream binding

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<values=large_string>.
     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,

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)