feat(sqlserver): add statement-level NOSCAN execution option

Author: vistart

src/rhosocial/activerecord/backend/impl/sqlserver/__init__.py

@@ -5,6 +5,7 @@
 from .config import SQLServerConnectionConfig
 from .collation import SQLServerCollation
 from .dialect import SQLServerDialect
+from .options import SQLServerExecutionOptions
 from .transaction import SQLServerTransactionManager
 from .adapters import (
     SQLServerUUIDAdapter,
@@ -66,6 +67,7 @@
 __all__ = [
     "SQLServerBackend",
     "SQLServerConnectionConfig",
+    "SQLServerExecutionOptions",
     "SQLServerDialect",
     "SQLServerCollation",
     "SQLServerTransactionManager",

src/rhosocial/activerecord/backend/impl/sqlserver/async_backend.py

@@ -9,7 +9,8 @@
 import asyncio
 import logging
 import time
-from typing import Any, Dict, List, Optional, Set, Tuple
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Dict, List, Optional, Set, Tuple
 
 from rhosocial.activerecord.backend.base import AsyncStorageBackend
 from rhosocial.activerecord.backend.base.operations import AsyncSQLOperationsMixin
@@ -329,6 +330,21 @@ async def execute(
 
         return await self._execute_statement(sql, params, options)
 
+    @asynccontextmanager
+    async def _cursor_noscan(self, cursor: Any, noscan: Optional[bool]) -> AsyncIterator[None]:
+        if noscan is None:
+            yield
+            return
+
+        impl = cursor._impl
+        run_operation = cursor._run_operation
+        previous = await run_operation(getattr, impl, "noscan")
+        try:
+            await run_operation(setattr, impl, "noscan", noscan)
+            yield
+        finally:
+            await run_operation(setattr, impl, "noscan", previous)
+
     async def _execute_statement(
         self, sql: str, params: Optional[Tuple], options, start_time: Optional[float] = None
     ) -> QueryResult:
@@ -366,7 +382,8 @@ async def _execute_statement(
                 if identity_table is not None:
                     await self.set_identity_insert(identity_table, True)
                 try:
-                    await cursor.execute(final_sql, final_params or ())
+                    async with self._cursor_noscan(cursor, getattr(options, "noscan", None)):
+                        await cursor.execute(final_sql, final_params or ())
                 finally:
                     if identity_table is not None:
                         await self.set_identity_insert(identity_table, False)

src/rhosocial/activerecord/backend/impl/sqlserver/backend.py

@@ -10,6 +10,7 @@
 import logging
 import re
 import time
+from contextvars import ContextVar
 from typing import Any, Dict, List, Optional, Tuple, Type
 
 import pyodbc
@@ -24,6 +25,7 @@
     OperationalError,
     QueryError,
 )
+from rhosocial.activerecord.backend.options import ExecutionOptions
 from rhosocial.activerecord.backend.result import QueryResult
 from rhosocial.activerecord.backend.introspection.backend_mixin import IntrospectorBackendMixin
 from rhosocial.activerecord.backend.explain import SyncExplainBackendMixin
@@ -161,6 +163,9 @@ def __init__(self, **kwargs):
         self._connection = None
         self._transaction_manager = None
         self._default_suggestions_cache = None
+        self._execution_noscan: ContextVar[Optional[bool]] = ContextVar(
+            "sqlserver_execution_noscan", default=None
+        )
 
         self._register_sqlserver_adapters()
 
@@ -358,7 +363,7 @@ def execute(
             cursor.execute(f"SET {set_clause} ON")
             cursor.close()
             try:
-                return super().execute(statement, params, options=options)
+                return self._execute_with_options(statement, params, options)
             finally:
                 cursor = self._get_cursor()
                 cursor.execute(f"SET {set_clause} OFF")
@@ -368,7 +373,7 @@ def execute(
         if identity_table is not None:
             self.set_identity_insert(identity_table, True)
         try:
-            result = super().execute(sql, params, options=options)
+            result = self._execute_with_options(sql, params, options)
         finally:
             if identity_table is not None:
                 self.set_identity_insert(identity_table, False)
@@ -380,6 +385,27 @@ def execute(
 
         return result
 
+    def _execute_with_options(
+        self, sql: str, params: Optional[Tuple], options: ExecutionOptions
+    ) -> QueryResult:
+        token = self._execution_noscan.set(getattr(options, "noscan", None))
+        try:
+            return super().execute(sql, params, options=options)
+        finally:
+            self._execution_noscan.reset(token)
+
+    def _execute_query(self, cursor: Any, sql: str, params: Optional[Tuple]) -> Any:
+        noscan = self._execution_noscan.get()
+        if noscan is None:
+            return super()._execute_query(cursor, sql, params)
+
+        previous = cursor.noscan
+        try:
+            cursor.noscan = noscan
+            return super()._execute_query(cursor, sql, params)
+        finally:
+            cursor.noscan = previous
+
     def execute_many(self, sql: str, params_list: List[Tuple]) -> QueryResult:
         """Execute the same SQL statement multiple times with different parameters.
         

src/rhosocial/activerecord/backend/impl/sqlserver/options.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from rhosocial.activerecord.backend.options import ExecutionOptions
+
+
+@dataclass
+class SQLServerExecutionOptions(ExecutionOptions):
+    noscan: Optional[bool] = None

tests/rhosocial/activerecord_sqlserver_test/feature/backend/dialect/test_sqlserver_sql_graph.py

@@ -34,8 +34,10 @@
     SQLServerEdgeConstraint,
     SQLServerGraphTableKind,
 )
+from rhosocial.activerecord.backend.impl.sqlserver.options import SQLServerExecutionOptions
 from rhosocial.activerecord.backend.impl.sqlserver.protocols import SQLServerGraphSupport
 from rhosocial.activerecord.testsuite.utils import requires_protocol, skip_test_if_protocol_unsupported
+from rhosocial.activerecord.backend.options import StatementType
 
 
 @pytest.fixture
@@ -419,6 +421,7 @@ def test_shortest_path_and_edge_constraint(self, sqlserver_backend, maximum, dir
                 f"SELECT LAST_VALUE([Person{segment_count + 1}].ID) WITHIN GROUP (GRAPH PATH) AS [endpoint_id] "
                 f"FROM {', '.join(tables)} WHERE {match_sql} AND [Person1].ID = {dialect.p()}",
                 params + (start_id,),
+                options=SQLServerExecutionOptions(stmt_type=StatementType.DQL, noscan=True),
             )
             assert result.data is not None
             actual = {row["endpoint_id"] for row in result.data}

tests/rhosocial/activerecord_sqlserver_test/feature/backend/test_execution_options.py

@@ -0,0 +1,232 @@
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from rhosocial.activerecord.backend.base import StorageBackend
+from rhosocial.activerecord.backend.errors import DatabaseError
+from rhosocial.activerecord.backend.impl.sqlserver import (
+    AsyncSQLServerBackend,
+    SQLServerBackend,
+    SQLServerExecutionOptions,
+)
+from rhosocial.activerecord.backend.impl.sqlserver.config import SQLServerConnectionConfig
+from rhosocial.activerecord.backend.options import ExecutionOptions, StatementType
+
+
+SQL = "SELECT {fn UCASE(?)} AS [value], '?' AS [literal]"
+PARAMS = ("{not an ODBC escape}",)
+OPTION_CASES = [
+    pytest.param(SQLServerExecutionOptions(StatementType.DQL, noscan=True), id="true"),
+    pytest.param(SQLServerExecutionOptions(StatementType.DQL, noscan=False), id="false"),
+    pytest.param(SQLServerExecutionOptions(StatementType.DQL, noscan=None), id="none"),
+    pytest.param(ExecutionOptions(StatementType.DQL), id="core"),
+]
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _prepare_read_committed_snapshot():
+    return None
+
+
+@pytest.fixture(autouse=True)
+def block_connections(monkeypatch):
+    monkeypatch.setattr(
+        SQLServerBackend, "connect", MagicMock(side_effect=AssertionError("Unexpected database connection"))
+    )
+    monkeypatch.setattr(
+        AsyncSQLServerBackend, "connect", AsyncMock(side_effect=AssertionError("Unexpected database connection"))
+    )
+
+
+class NativeCursor:
+    def __init__(self, noscan):
+        self._noscan = noscan
+        self.accesses = []
+
+    @property
+    def noscan(self):
+        self.accesses.append(("get",))
+        return self._noscan
+
+    @noscan.setter
+    def noscan(self, value):
+        self.accesses.append(("set", value))
+        self._noscan = value
+
+
+class AsyncCursor:
+    def __init__(self, previous, error=None):
+        self._impl = NativeCursor(previous)
+        self.noscan = "wrapper must not change"
+        self.operations = []
+        self.executions = []
+        self.error = error
+        self.description = None
+        self.rowcount = 0
+        self.lastrowid = None
+        self.exit_value = None
+
+    async def _run_operation(self, operation, *args):
+        self.operations.append((operation, args))
+        return operation(*args)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.exit_value = self._impl._noscan
+
+    async def execute(self, sql, params):
+        self.executions.append((sql, params, self._impl._noscan, self.noscan))
+        if self.error is not None:
+            raise self.error
+
+
+def make_backend(backend_class):
+    config = SQLServerConnectionConfig(host="localhost", database="master", username="sa", password="")
+    backend = backend_class(connection_config=config)
+    backend._connection = object()
+    return backend
+
+
+def expected_accesses(noscan, previous):
+    return [] if noscan is None else [("get",), ("set", noscan), ("set", previous)]
+
+
+def test_execution_options_default():
+    options = SQLServerExecutionOptions(StatementType.DQL)
+    assert isinstance(options, ExecutionOptions)
+    assert options.noscan is None
+
+
+@pytest.mark.parametrize("options", OPTION_CASES)
+@pytest.mark.parametrize("fails", [False, True], ids=["success", "error"])
+@pytest.mark.parametrize("outer", [None, True, False])
+def test_sync_execution_restores_cursor_and_context(monkeypatch, options, fails, outer):
+    backend = make_backend(SQLServerBackend)
+    assert backend._execution_noscan.get() is None
+    noscan = getattr(options, "noscan", None)
+    previous = noscan is not True
+    cursor = NativeCursor(previous)
+    error = RuntimeError("execution failed")
+    executions = []
+    delegated = []
+
+    def execute(sql, params):
+        executions.append((sql, params, cursor._noscan))
+        if fails:
+            raise error
+
+    cursor.execute = execute
+
+    def core_execute(self, sql, params, *, options):
+        delegated.append((sql, params, options, self._execution_noscan.get()))
+        return self._execute_query(cursor, sql, params)
+
+    monkeypatch.setattr(StorageBackend, "execute", core_execute)
+    token = backend._execution_noscan.set(outer)
+    try:
+        if fails:
+            with pytest.raises(RuntimeError, match="execution failed") as caught:
+                backend.execute(SQL, PARAMS, options=options)
+            assert caught.value is error
+        else:
+            assert backend.execute(SQL, PARAMS, options=options) is cursor
+        assert backend._execution_noscan.get() is outer
+        assert cursor._noscan is previous
+        assert cursor.accesses == expected_accesses(noscan, previous)
+        assert delegated == [(SQL, PARAMS, options, noscan)]
+        assert delegated[0][2] is options
+        assert executions == [(SQL, PARAMS, previous if noscan is None else noscan)]
+    finally:
+        backend._execution_noscan.reset(token)
+
+    assert backend._execution_noscan.get() is None
+    fails = False
+    cursor.accesses.clear()
+    assert backend.execute(SQL, PARAMS) is cursor
+    assert backend._execution_noscan.get() is None
+    assert executions[-1] == (SQL, PARAMS, previous)
+    assert len(executions) == 2
+    assert type(delegated[-1][2]) is ExecutionOptions
+    assert delegated[-1][3] is None
+    assert cursor.accesses == []
+
+
+@pytest.mark.parametrize("options", OPTION_CASES)
+@pytest.mark.parametrize("fails", [False, True], ids=["success", "error"])
+@pytest.mark.parametrize("previous", [False, True])
+@pytest.mark.asyncio
+async def test_async_statement_uses_native_noscan(monkeypatch, options, fails, previous):
+    backend = make_backend(AsyncSQLServerBackend)
+    noscan = getattr(options, "noscan", None)
+    cursor = AsyncCursor(previous, RuntimeError("execution failed") if fails else None)
+    monkeypatch.setattr(backend, "_get_cursor", AsyncMock(return_value=cursor))
+    monkeypatch.setattr(backend, "_handle_auto_commit_if_needed", AsyncMock())
+
+    if fails:
+        with pytest.raises(DatabaseError, match="execution failed"):
+            await backend.execute(SQL, PARAMS, options=options)
+    else:
+        result = await backend.execute(SQL, PARAMS, options=options)
+        assert result.affected_rows == 0
+        assert result.data is None
+
+    assert cursor.executions == [
+        (SQL, PARAMS, previous if noscan is None else noscan, "wrapper must not change")
+    ]
+    assert cursor._impl._noscan is previous
+    assert cursor.exit_value is previous
+    assert cursor.noscan == "wrapper must not change"
+    assert cursor._impl.accesses == expected_accesses(noscan, previous)
+    assert cursor.operations == ([] if noscan is None else [
+        (getattr, (cursor._impl, "noscan")),
+        (setattr, (cursor._impl, "noscan", noscan)),
+        (setattr, (cursor._impl, "noscan", previous)),
+    ])
+
+    cursor.error = None
+    cursor.operations.clear()
+    cursor._impl.accesses.clear()
+    await backend.execute(SQL, PARAMS)
+    assert cursor.executions[-1] == (SQL, PARAMS, previous, "wrapper must not change")
+    assert len(cursor.executions) == 2
+    assert cursor.operations == []
+    assert cursor._impl.accesses == []
+
+
+@pytest.mark.parametrize("noscan", [True, False, None])
+@pytest.mark.parametrize("fails", [False, True], ids=["success", "error"])
+@pytest.mark.asyncio
+async def test_async_cursor_noscan_context_restores_native_value(noscan, fails):
+    backend = make_backend(AsyncSQLServerBackend)
+    previous = noscan is not True
+    cursor = AsyncCursor(previous)
+    error = RuntimeError("context failed")
+
+    async def use_context():
+        async with backend._cursor_noscan(cursor, noscan):
+            assert cursor._impl._noscan is (previous if noscan is None else noscan)
+            assert cursor.noscan == "wrapper must not change"
+            if fails:
+                raise error
+
+    if fails:
+        with pytest.raises(RuntimeError, match="context failed") as caught:
+            await use_context()
+        assert caught.value is error
+    else:
+        await use_context()
+    assert cursor._impl._noscan is previous
+    assert cursor.noscan == "wrapper must not change"
+    assert cursor._impl.accesses == expected_accesses(noscan, previous)
+
+
+@pytest.mark.asyncio
+async def test_async_cursor_noscan_none_does_not_require_native_cursor():
+    backend = make_backend(AsyncSQLServerBackend)
+    cursor = SimpleNamespace()
+    async with backend._cursor_noscan(cursor, None):
+        assert vars(cursor) == {}
+    assert vars(cursor) == {}