FEAT: TVP support for async queries

mssql_python/async_query/__init__.py

@@ -2,6 +2,9 @@
 
 The internal entry points are ``_AsyncConnection`` and ``_AsyncCursor``.
 Their methods retain DB-API names, but the classes are not stable public API.
+``_TableValuedParameter`` lazily exposes the native TVP constructor. The
+``SQL_*`` tokens are native setinputsizes() hints; server support is required
+for the corresponding SQL types, particularly JSON and VECTOR.
 
 Warning:
     Async query execution APIs are under active development and are not intended
@@ -25,6 +28,40 @@
     Warning,
 )
 
+SQL_MONEY: int
+SQL_SMALLMONEY: int
+SQL_XML: int
+SQL_JSON: int
+SQL_VECTOR: int
+
+_NATIVE_EXPORTS = {
+    "_TableValuedParameter": "TableValuedParameter",
+    "SQL_MONEY": "SQL_MONEY",
+    "SQL_SMALLMONEY": "SQL_SMALLMONEY",
+    "SQL_XML": "SQL_XML",
+    "SQL_JSON": "SQL_JSON",
+    "SQL_VECTOR": "SQL_VECTOR",
+}
+
+
+def __getattr__(name: str):
+    native_name = _NATIVE_EXPORTS.get(name)
+    if native_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    py_core = load_py_core()
+    try:
+        return getattr(py_core, native_name)
+    except AttributeError as error:
+        raise ImportError(
+            f"The installed mssql-python-rs dependency does not provide {native_name}; "
+            "install a version with this async feature."
+        ) from error
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_NATIVE_EXPORTS))
+
+
 __all__ = [
     "DataError",
     "DatabaseError",
@@ -35,6 +72,11 @@
     "NotSupportedError",
     "OperationalError",
     "ProgrammingError",
+    "SQL_MONEY",
+    "SQL_SMALLMONEY",
+    "SQL_XML",
+    "SQL_JSON",
+    "SQL_VECTOR",
     "Warning",
     "load_py_core",
 ]

mssql_python/async_query/async_cursor.py

@@ -7,7 +7,7 @@
 """
 
 import asyncio
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from contextlib import asynccontextmanager
 from typing import Any, Optional
 import uuid
@@ -142,7 +142,7 @@ async def execute(
     async def executemany(
         self,
         operation: str,
-        seq_of_parameters: Sequence[Sequence[Any]] | Sequence[Mapping[str, Any]],
+        seq_of_parameters: Iterable[Sequence[Any] | Mapping[str, Any]],
         *,
         use_prepare: bool = True,
     ) -> None:
@@ -164,6 +164,11 @@ async def fetchall(self) -> list[Row]:
         return await async_fetch.fetchall(self)
 
     async def nextset(self) -> bool:
+        """Discard remaining rows and advance to the next statement result.
+
+        Return True for another result, including DML results, or False at
+        batch end. Fetch counts are reset rather than carried between results.
+        """
         async with self._result_transition():
             try:
                 with translate_py_core_exceptions():
@@ -178,6 +183,13 @@ async def nextset(self) -> bool:
             return has_next
 
     async def close(self) -> None:
+        """Close this cursor; repeated calls are harmless.
+
+        Successful close clears description and resets rowcount to -1, unlike
+        the synchronous cursor, which retains its last metadata and count.
+        Rejected close preserves the cursor; accepted but interrupted cleanup
+        retires the wrapper without guaranteeing native cleanup has completed.
+        """
         logger.debug("AsyncCursor.close: starting")
         async with self._result_transition():
             with translate_py_core_exceptions():
@@ -203,17 +215,33 @@ def setinputsizes(self, sizes: Any) -> None:
         with translate_py_core_exceptions():
             self._py_core_async_cursor.setinputsizes(sizes)
 
+    @property
+    def closed(self) -> bool:
+        """Whether the wrapper is retired or its parent connection is closed.
+
+        This does not indicate completion of native cleanup or an in-flight
+        close operation. A rejected close leaves this value unchanged.
+        """
+        return self._closed or (self._connection is not None and self._connection.closed)
+
     @property
     def timeout(self) -> int:
         with translate_py_core_exceptions():
             return self._py_core_async_cursor.timeout
 
     @property
     def description(self) -> Any:
+        """Seven-item column descriptors, or None when no row result is published."""
         return self._description
 
     @property
     def rowcount(self) -> int:
+        """Fetched rows for the current SELECT, or the native affected-row count.
+
+        SELECT counts start at -1 and accumulate as rows are fetched; fetching
+        an empty result sets the count to zero. Unlike the synchronous cursor,
+        nextset() does not retain fetched counts from the previous result.
+        """
         if self._fetch_rowcount is not None:
             return self._fetch_rowcount
         with translate_py_core_exceptions():

mssql_python/async_query/async_execute.py

@@ -1,7 +1,7 @@
 """Asynchronous statement execution through mssql-py-core."""
 
 import asyncio
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
 from ..exceptions import OperationalError, ProgrammingError
@@ -85,29 +85,57 @@ async def execute(
 async def executemany(
     cursor: "_AsyncCursor",
     operation: str,
-    seq_of_parameters: Sequence[Sequence[Any]] | Sequence[Mapping[str, Any]],
+    seq_of_parameters: Iterable[Sequence[Any] | Mapping[str, Any]],
     *,
     use_prepare: bool = True,
 ) -> None:
-    """Execute a statement for every parameter row using the py-core async cursor."""
+    """Execute parameter rows from a synchronous iterable through py-core.
+
+    Py-core consumes and validates the iterable before dispatch; this is not
+    streaming execution. Asynchronous iterables are not supported. Exceptions
+    raised by the parameter iterable are propagated unchanged.
+    """
     cursor._check_closed()  # pyright: ignore[reportPrivateUsage]
-    batch_count = len(seq_of_parameters)
+    iteration_error: BaseException | None = None
+
+    def parameter_rows():
+        nonlocal iteration_error
+        try:
+            iterator = iter(seq_of_parameters)
+        except BaseException as error:
+            iteration_error = error
+            raise
+        while True:
+            try:
+                row = next(iterator)
+            except StopIteration:
+                return
+            except BaseException as error:
+                iteration_error = error
+                raise
+            yield row
+
     logger.debug(
-        "AsyncCursor.executemany: starting; batch_count=%d; use_prepare=%s",
-        batch_count,
+        "AsyncCursor.executemany: starting; use_prepare=%s",
         use_prepare,
     )
-    with translate_py_core_exceptions():
-        executemany_awaitable = _get_py_core_async_cursor(cursor).executemany(
-            operation,
-            seq_of_parameters,
-            use_prepare=use_prepare,
-        )
     try:
         with translate_py_core_exceptions():
-            await executemany_awaitable
-    except (Exception, asyncio.CancelledError) as error:
-        _reconcile_failed_execution(cursor, error)
+            executemany_awaitable = _get_py_core_async_cursor(cursor).executemany(
+                operation,
+                parameter_rows(),
+                use_prepare=use_prepare,
+            )
+        try:
+            with translate_py_core_exceptions():
+                await executemany_awaitable
+        except (Exception, asyncio.CancelledError) as error:
+            if iteration_error is None:
+                _reconcile_failed_execution(cursor, error)
+            raise
+    except BaseException:
+        if iteration_error is not None:
+            raise iteration_error from iteration_error.__cause__
         raise
     cursor._reset_fetch_tracking()  # pyright: ignore[reportPrivateUsage]
     cursor._clear_result_metadata()  # pyright: ignore[reportPrivateUsage]

tests/AsyncTest/test_001_async_query_native.py

@@ -1,3 +1,6 @@
+import datetime
+from decimal import Decimal
+
 import sys
 from importlib.machinery import ModuleSpec
 from types import ModuleType
@@ -6,6 +9,7 @@
 
 import mssql_python.async_query as async_query
 import mssql_python.async_query._native as native
+from mssql_python.constants import ConstantsDDBC
 
 
 def test_async_entry_points_are_internal():
@@ -62,3 +66,91 @@ def test_load_py_core_reports_missing_async_types(monkeypatch):
 
     with pytest.raises(ImportError, match="PyAsyncConnection, PyAsyncCursor"):
         native.load_py_core()
+
+
+@pytest.mark.parametrize(
+    "name", ("SQL_MONEY", "SQL_SMALLMONEY", "SQL_XML", "SQL_JSON", "SQL_VECTOR")
+)
+def test_async_type_hints_match_native_exports(name):
+    py_core = pytest.importorskip("mssql_py_core.mssql_py_core", exc_type=ImportError)
+    assert getattr(async_query, name) == getattr(py_core, name)
+    assert isinstance(getattr(async_query, name), int)
+    assert name in async_query.__all__
+    assert name in dir(async_query)
+
+
+def test_async_tvp_is_internal_native_constructor():
+    pytest.importorskip("mssql_py_core.mssql_py_core", exc_type=ImportError)
+    from mssql_python.async_query import _TableValuedParameter
+
+    assert _TableValuedParameter is native.load_py_core().TableValuedParameter
+    assert "_TableValuedParameter" not in async_query.__all__
+    assert "_TableValuedParameter" in dir(async_query)
+    assert not hasattr(async_query, "TableValuedParameter")
+    value = _TableValuedParameter("TestType", [(4, 0, 0)], [(1,)], schema="dbo")
+    assert (value.type_name, value.schema) == ("TestType", "dbo")
+    assert (value.column_count, value.row_count, value.is_null) == (1, 1, False)
+    assert _TableValuedParameter("dbo.TestType").is_null is True
+    with pytest.raises(ValueError, match="requires column definitions"):
+        _TableValuedParameter("dbo.TestType", rows=[(1,)])
+
+
+def test_async_tvp_accepts_numeric_and_temporal_column_metadata():
+    pytest.importorskip("mssql_py_core.mssql_py_core", exc_type=ImportError)
+    from mssql_python.async_query import _TableValuedParameter
+
+    columns = [
+        (ConstantsDDBC.SQL_DECIMAL.value, 12, 2),
+        (ConstantsDDBC.SQL_NUMERIC.value, 38, 10),
+        (ConstantsDDBC.SQL_TYPE_TIME.value, 0, 3),
+        (ConstantsDDBC.SQL_TYPE_TIMESTAMP.value, 0, 3),
+        (ConstantsDDBC.SQL_DATETIMEOFFSET.value, 0, 3),
+    ]
+    populated_row = (
+        Decimal("123.45"),
+        Decimal("123.4500000000"),
+        datetime.time(12, 34, 56, 123000),
+        datetime.datetime(2024, 1, 2, 12, 34, 56, 123000),
+        datetime.datetime(
+            2024,
+            1,
+            2,
+            12,
+            34,
+            56,
+            123000,
+            tzinfo=datetime.timezone(datetime.timedelta(hours=-5, minutes=-30)),
+        ),
+    )
+
+    populated = _TableValuedParameter("dbo.MetadataType", columns, [populated_row])
+    null_cells = _TableValuedParameter("dbo.MetadataType", columns, [(None,) * len(columns)])
+    empty = _TableValuedParameter("dbo.MetadataType", columns, [])
+
+    assert (populated.column_count, populated.row_count) == (len(columns), 1)
+    assert (null_cells.column_count, null_cells.row_count) == (len(columns), 1)
+    assert (empty.column_count, empty.row_count) == (len(columns), 0)
+
+
+def test_async_native_exports_load_only_when_requested(monkeypatch):
+    calls = []
+    py_core = ModuleType("mssql_py_core")
+    setattr(py_core, "SQL_JSON", 244)
+
+    def load():
+        calls.append(True)
+        return py_core
+
+    monkeypatch.setattr(async_query, "load_py_core", load)
+    assert "SQL_JSON" in dir(async_query)
+    assert not hasattr(async_query, "unknown_export")
+    assert calls == []
+    assert async_query.SQL_JSON == 244
+    assert calls == [True]
+
+
+@pytest.mark.parametrize("name", ("_TableValuedParameter", "SQL_VECTOR"))
+def test_async_native_exports_report_missing_feature(monkeypatch, name):
+    monkeypatch.setattr(async_query, "load_py_core", lambda: ModuleType("mssql_py_core"))
+    with pytest.raises(ImportError, match="installed mssql-python-rs dependency"):
+        getattr(async_query, name)

tests/AsyncTest/test_003_async_exceptions.py

@@ -304,6 +304,61 @@ async def test_executemany_integrity_error_reports_row_and_preserves_partial_pro
         await connection.close()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_prepare", (False, True))
+@pytest.mark.parametrize("batch", (False, True))
+@pytest.mark.parametrize("supplementary", (False, True))
+async def test_string_truncation_diagnostics_partial_progress_and_recovery(
+    async_cursor, use_prepare, batch, supplementary
+):
+    value = "\U0001f600" * 3 if supplementary else "abcdef"
+    await async_cursor.execute(
+        "CREATE TABLE #async_truncation (id INT, value NVARCHAR(5))", use_prepare=False
+    )
+    try:
+        with pytest.raises(
+            (public_exceptions.DataError, public_exceptions.OperationalError)
+        ) as raised:
+            if batch:
+                await async_cursor.executemany(
+                    "INSERT INTO #async_truncation VALUES (?, ?)",
+                    [(1, "ok"), (2, value), (3, "later")],
+                    use_prepare=use_prepare,
+                )
+            else:
+                await async_cursor.execute(
+                    "INSERT INTO #async_truncation VALUES (?, ?)",
+                    2,
+                    value,
+                    use_prepare=use_prepare,
+                )
+        error = cast(Any, raised.value)
+        native_decode_failure = (
+            supplementary
+            and isinstance(error, public_exceptions.OperationalError)
+            and "invalid utf-16: lone surrogate found" in str(error)
+            and not getattr(error, "sql_errors", None)
+        )
+        if not native_decode_failure:
+            assert isinstance(error, public_exceptions.DataError)
+            diagnostics = getattr(error, "sql_errors", [])
+            assert diagnostics
+            assert any(item["number"] in (8152, 2628) for item in diagnostics)
+        assert error.__cause__ is not None
+        await async_cursor.execute("SELECT id, value FROM #async_truncation ORDER BY id")
+        assert [tuple(row) for row in await async_cursor.fetchall()] == (
+            [(1, "ok")] if batch else []
+        )
+        await async_cursor.execute("INSERT INTO #async_truncation VALUES (4, N'new')")
+        assert async_cursor.rowcount == 1
+        await async_cursor.execute("SELECT value FROM #async_truncation WHERE id = 4")
+        assert tuple(await async_cursor.fetchone()) == ("new",)
+        if native_decode_failure:
+            pytest.xfail("py-core truncation diagnostic decoder rejects split UTF-16 surrogate")
+    finally:
+        await async_cursor.execute("DROP TABLE IF EXISTS #async_truncation", use_prepare=False)
+
+
 @pytest.mark.asyncio
 async def test_timeout_is_operational_error_and_cursor_is_reusable(async_connection_string):
     connection = await _AsyncConnection.connect(async_connection_string)

tests/AsyncTest/test_004_async_logging.py

@@ -94,7 +94,7 @@ async def test_default_logger_combines_python_and_py_core_operation_logs(
         "AsyncCursor.fetchmany: starting; requested_size=1",
         "AsyncCursor.fetchmany: completed; row_count=1; rowcount=2",
         "AsyncCursor.fetchall: completed; row_count=1; rowcount=1",
-        "AsyncCursor.executemany: starting; batch_count=2",
+        "AsyncCursor.executemany: starting; use_prepare=True",
         "AsyncCursor.executemany: completed; rowcount=-1",
     )
     for expected in expected_python_messages: