refactor: rename internal classes and update related tests for clarity

Author: cofin

sqlspec/extensions/litestar/plugin.py

@@ -71,11 +71,16 @@
 CORRELATION_STATE_KEY = "sqlspec_correlation_id"
 
 __all__ = (
+    "CORRELATION_STATE_KEY",
     "DEFAULT_COMMIT_MODE",
     "DEFAULT_CONNECTION_KEY",
+    "DEFAULT_CORRELATION_HEADER",
     "DEFAULT_POOL_KEY",
     "DEFAULT_SESSION_KEY",
+    "TRACE_CONTEXT_FALLBACK_HEADERS",
     "CommitMode",
+    "CorrelationMiddleware",
+    "PluginConfigState",
     "SQLSpecPlugin",
 )
 
@@ -117,7 +122,7 @@ def _build_correlation_headers(*, primary: str, configured: list[str], auto_trac
     return tuple(_dedupe_headers(header_order))
 
 
-class _CorrelationMiddleware:
+class CorrelationMiddleware:
     __slots__ = ("_app", "_headers")
 
     def __init__(self, app: "ASGIApp", *, headers: tuple[str, ...]) -> None:
@@ -153,7 +158,7 @@ async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> No
 
 
 @dataclass
-class _PluginConfigState:
+class PluginConfigState:
     """Internal state for each database configuration."""
 
     config: "DatabaseConfigProtocol[Any, Any, Any]"
@@ -219,7 +224,7 @@ def __init__(self, sqlspec: SQLSpec, *, loader: "SQLFileLoader | None" = None) -
         """
         self._sqlspec = sqlspec
 
-        self._plugin_configs: list[_PluginConfigState] = []
+        self._plugin_configs: list[PluginConfigState] = []
         for cfg in self._sqlspec.configs.values():
             config_union = cast(
                 "SyncDatabaseConfig[Any, Any, Any] | NoPoolSyncConfig[Any, Any] | AsyncDatabaseConfig[Any, Any, Any] | NoPoolAsyncConfig[Any, Any]",
@@ -276,9 +281,9 @@ def _create_config_state(
         self,
         config: "SyncDatabaseConfig[Any, Any, Any] | NoPoolSyncConfig[Any, Any] | AsyncDatabaseConfig[Any, Any, Any] | NoPoolAsyncConfig[Any, Any]",
         settings: "dict[str, Any]",
-    ) -> _PluginConfigState:
+    ) -> PluginConfigState:
         """Create plugin state with handlers for the given configuration."""
-        state = _PluginConfigState(
+        state = PluginConfigState(
             config=config,
             connection_key=settings["connection_key"],
             pool_key=settings["pool_key"],
@@ -296,7 +301,7 @@ def _create_config_state(
             self._setup_handlers(state)
         return state
 
-    def _setup_handlers(self, state: _PluginConfigState) -> None:
+    def _setup_handlers(self, state: PluginConfigState) -> None:
         """Setup handlers for the plugin state."""
         connection_key = state.connection_key
         pool_key = state.pool_key
@@ -403,7 +408,7 @@ def store_sqlspec_in_state() -> None:
                 app_config.type_decoders = decoders_list
 
         if self._correlation_headers:
-            middleware = DefineMiddleware(_CorrelationMiddleware, headers=self._correlation_headers)
+            middleware = DefineMiddleware(CorrelationMiddleware, headers=self._correlation_headers)
             existing_middleware = list(app_config.middleware or [])
             existing_middleware.append(middleware)
             app_config.middleware = existing_middleware
@@ -579,7 +584,7 @@ def provide_request_connection(
 
     def _get_plugin_state(
         self, key: "str | SyncConfigT | AsyncConfigT | type[SyncConfigT | AsyncConfigT]"
-    ) -> _PluginConfigState:
+    ) -> PluginConfigState:
         """Get plugin state for a configuration by key."""
         if isinstance(key, str):
             for state in self._plugin_configs:

tests/unit/test_adapters/test_spanner/test_config.py

@@ -53,8 +53,7 @@ def test_driver_features_defaults() -> None:
 
 
 def test_provide_connection_batch_and_snapshot() -> None:
-    """Ensure provide_connection selects snapshot vs batch correctly."""
-    batch_obj = object()
+    """Ensure provide_connection selects snapshot vs transaction correctly."""
     snap_obj = object()
 
     class _Ctx:
@@ -67,39 +66,88 @@ def __enter__(self):
         def __exit__(self, *_):
             return False
 
+    class _Txn:
+        _transaction_id = "test-txn-id"
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *_):
+            return False
+
+        def commit(self):
+            pass
+
+        def rollback(self):
+            pass
+
+    class _Session:
+        def create(self):
+            pass
+
+        def delete(self):
+            pass
+
+        def transaction(self):
+            return _Txn()
+
     class _DB:
-        def batch(self):
-            return _Ctx(batch_obj)
+        def session(self):
+            return _Session()
 
-        def snapshot(self):
+        def snapshot(self, multi_use: bool = False):
             return _Ctx(snap_obj)
 
     config = SpannerSyncConfig(pool_config={"project": "p", "instance_id": "i", "database_id": "d"})
     config.get_database = lambda: _DB()  # type: ignore[assignment]
 
     with config.provide_connection(transaction=True) as conn:
-        assert conn is batch_obj
+        assert isinstance(conn, _Txn)
 
     with config.provide_connection(transaction=False) as conn:
         assert conn is snap_obj
 
 
 def test_provide_session_uses_batch_when_transaction_requested() -> None:
-    """Driver should receive batch connection when transaction=True."""
-    batch_obj = object()
+    """Driver should receive transaction connection when transaction=True."""
+
+    class _Txn:
+        _transaction_id = "test-txn-id"
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *_):
+            return False
+
+        def commit(self):
+            pass
+
+        def rollback(self):
+            pass
+
+    class _Session:
+        def create(self):
+            pass
+
+        def delete(self):
+            pass
+
+        def transaction(self):
+            return _Txn()
 
     class _Ctx:
         def __enter__(self):
-            return batch_obj
+            return object()
 
         def __exit__(self, *_):
             return False
 
     class _DB:
-        def batch(self):
-            return _Ctx()
+        def session(self):
+            return _Session()
 
-        def snapshot(self):
+        def snapshot(self, multi_use: bool = False):
             return _Ctx()
 
     config = SpannerSyncConfig(pool_config={"project": "p", "instance_id": "i", "database_id": "d"})
@@ -108,30 +156,54 @@ def snapshot(self):
 
     with config.provide_session(transaction=True) as driver:
         assert isinstance(driver, _DummyDriver)
-        assert driver.connection is batch_obj
+        assert isinstance(driver.connection, _Txn)
 
 
 def test_provide_write_session_alias() -> None:
-    """provide_write_session should always give a batch-backed driver."""
-    batch_obj = object()
+    """provide_write_session should always give a transaction-backed driver."""
+
+    class _Txn:
+        _transaction_id = "test-txn-id"
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *_):
+            return False
+
+        def commit(self):
+            pass
+
+        def rollback(self):
+            pass
+
+    class _Session:
+        def create(self):
+            pass
+
+        def delete(self):
+            pass
+
+        def transaction(self):
+            return _Txn()
 
     class _Ctx:
         def __enter__(self):
-            return batch_obj
+            return object()
 
         def __exit__(self, *_):
             return False
 
     class _DB:
-        def batch(self):
-            return _Ctx()
+        def session(self):
+            return _Session()
 
-        def snapshot(self):
+        def snapshot(self, multi_use: bool = False):
             return _Ctx()
 
     config = SpannerSyncConfig(pool_config={"project": "p", "instance_id": "i", "database_id": "d"})
     config.get_database = lambda: _DB()  # type: ignore[assignment]
     config.driver_type = _DummyDriver  # type: ignore[assignment,misc]
 
     with config.provide_write_session() as driver:
-        assert driver.connection is batch_obj
+        assert isinstance(driver.connection, _Txn)