Present the _db_pool interface on RustConnectionPool

Make RustConnectionPool a drop-in for DatabasePool._db_pool so real
runInteraction traffic can flow through it (once make_pool is switched):

  - it now hands each function a DBAPI2-adapter connection (rust_dbapi.Connection)
    wrapping the pooled shim, so LoggingDatabaseConnection.cursor() yields a
    working LoggingTransaction and the engine's in_transaction / is_closed /
    set_autocommit operate on it;
  - the entry point is named `runWithConnection` (matching
    twisted.enterprise.adbapi.ConnectionPool, which database.py calls by that
    name), alongside `threadID` for the transaction-limit path. With
    `threadpool` and `running` already present, the pool covers the slice of the
    adbapi interface DatabasePool uses.

A new test drives a full transaction through the pool the way
DatabasePool.runWithConnection's inner_func does — engine.in_transaction check,
LoggingDatabaseConnection + LoggingTransaction, `?`→`$n` conversion, commit —
off the reactor thread, and gets the result back via the Deferred.

Still outstanding before make_pool can return it: `reconnect` on the connection
(transaction-limit / closed-connection paths) and the startup path
(make_conn / check_database).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

synapse/storage/rust_pool.py

@@ -15,26 +15,32 @@
 Synapse's transaction functions are synchronous and expect a DBAPI2 connection.
 This adapter lets them run unchanged against the Rust ``Connection`` / ``Cursor``
 shim: it owns a dedicated Twisted thread pool and, for each call, checks a
-connection out of the native Rust ``ConnectionPool``, runs the caller's function
-against it on a worker thread, then returns the connection to the pool and hands
-the result back to the reactor as a ``Deferred``.
+connection out of the native Rust ``ConnectionPool``, wraps it in the DBAPI2
+adapter (:mod:`synapse.storage.rust_dbapi`), runs the caller's function against
+it on a worker thread, then returns the connection to the pool and hands the
+result back to the reactor as a ``Deferred``.
 
-This is deliberately a thin *execution bridge*. Slotting it into
-``DatabasePool`` (so ``runInteraction`` flows through it) additionally needs
-engine-level support for the shim connection — ``in_transaction``,
-``is_connection_closed``, autocommit / isolation and ``reconnect`` — and is left
-to a follow-up.
+It presents the slice of ``twisted.enterprise.adbapi.ConnectionPool`` that
+``DatabasePool`` uses — ``runWithConnection``, ``threadID``, ``threadpool`` and
+``running`` — so it can stand in for ``_db_pool`` (paired with
+:class:`~synapse.storage.engines.RustPostgresEngine`, which drives the wrapped
+connection). What remains before ``make_pool`` can return it: ``reconnect`` on
+the connection (only used on the transaction-limit / closed-connection paths)
+and the startup path (``make_conn`` / ``check_database``).
 """
 
 import logging
+import threading
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 
 from typing_extensions import Concatenate, ParamSpec
 
+from twisted.internet import threads
 from twisted.python.threadpool import ThreadPool
 
 from synapse.logging.context import defer_to_threadpool
+from synapse.storage.rust_dbapi import Connection as DBAPI2Connection
 from synapse.synapse_rust.database import postgres
 
 if TYPE_CHECKING:
@@ -51,11 +57,12 @@ R = TypeVar("R")
 class RustConnectionPool:
     """Runs blocking database functions against pooled Rust connections.
 
-    Each :meth:`run_with_connection` call runs its function on a worker thread
-    with a connection checked out of the native Rust pool, and returns a
-    ``Deferred`` that fires on the reactor thread with the result (or an
-    errback if it raised). Log contexts are preserved across the hop, following
-    the same rules as :func:`synapse.logging.context.defer_to_threadpool`.
+    Each :meth:`runWithConnection` call runs its function on a worker thread
+    with a (DBAPI2-adapter) connection checked out of the native Rust pool, and
+    returns a ``Deferred`` that fires on the reactor thread with the result (or
+    an errback if it raised). Log contexts are preserved across the hop,
+    following the same rules as
+    :func:`synapse.logging.context.defer_to_threadpool`.
     """
 
     def __init__(
@@ -107,7 +114,7 @@ class RustConnectionPool:
         self.threadpool.stop()
         self._pool.close()
 
-    def run_with_connection(
+    def runWithConnection(  # noqa: N802 (implements adbapi's interface)
         self,
         func: Callable[Concatenate[Any, P], R],
         *args: P.args,
@@ -115,22 +122,38 @@ class RustConnectionPool:
     ) -> "Deferred[R]":
         """Run ``func(conn, *args, **kwargs)`` on a worker thread.
 
-        ``conn`` is a connection checked out of the Rust pool for the duration
-        of the call. The function is responsible for committing or rolling back
-        (as Synapse's ``new_transaction`` does); the connection is returned to
-        the pool afterwards regardless.
+        ``conn`` is a DBAPI2-adapter connection wrapping one checked out of the
+        Rust pool for the duration of the call. The function is responsible for
+        committing or rolling back (as Synapse's ``new_transaction`` does); the
+        connection is returned to the pool afterwards regardless.
+
+        Named to match ``twisted.enterprise.adbapi.ConnectionPool`` so this can
+        stand in for ``DatabasePool._db_pool``.
 
         Returns:
-            A ``Deferred`` firing with ``func``'s result, following the Synapse
-            logcontext rules (``yield`` / ``await`` it).
+            A raw ``Deferred`` — like ``adbapi.ConnectionPool.runWithConnection``
+            (and unlike a `make_deferred_yieldable`-wrapped one), because the
+            caller wraps it: ``DatabasePool.runWithConnection`` supplies the
+            single ``make_deferred_yieldable``, and the DB function sets up its
+            own ``LoggingContext``. Wrapping it here as well (double
+            ``make_deferred_yieldable``, plus a nested logcontext) breaks
+            logcontext handling when the awaiting request is cancelled.
         """
         if not self.running:
             raise RuntimeError("connection pool is not running")
 
-        return defer_to_threadpool(
+        return threads.deferToThreadPool(
             self._reactor, self.threadpool, self._run, func, args, kwargs
         )
 
+    def threadID(self) -> int:  # noqa: N802 (implements adbapi's interface)
+        """Identify the current worker thread (adbapi interface).
+
+        Used by ``DatabasePool`` only when a per-connection transaction limit is
+        configured, to count transactions per thread.
+        """
+        return threading.get_ident()
+
     def _run(
         self,
         func: Callable[..., R],
@@ -142,7 +165,7 @@ class RustConnectionPool:
         A checkout failure surfaces as the raised exception (→ errback) before
         there is any connection to release.
         """
-        conn = self._pool.connect()
+        conn = DBAPI2Connection(self._pool.connect())
         try:
             return func(conn, *args, **kwargs)
         finally:

tests/storage/test_rust_pool.py

@@ -26,6 +26,8 @@ from twisted.internet import reactor as _reactor
 from twisted.internet.defer import gatherResults, inlineCallbacks
 from twisted.trial import unittest as trial_unittest
 
+from synapse.storage.database import LoggingDatabaseConnection
+from synapse.storage.engines.postgres_rust import RustPostgresEngine
 from synapse.storage.rust_pool import RustConnectionPool
 
 from tests.unittest import skip_unless
@@ -81,11 +83,11 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         def txn(conn: Any) -> Any:
             cursor = conn.cursor()
             cursor.execute("SELECT 42::int")
-            row = cursor.fetch_one()
+            row = cursor.fetchone()
             conn.commit()
             return row
 
-        result = yield self.pool.run_with_connection(txn)
+        result = yield self.pool.runWithConnection(txn)
         self.assertEqual(result, (42,))
 
     @inlineCallbacks
@@ -93,7 +95,7 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         def txn(conn: Any, a: int, b: int, c: int = 0) -> int:
             return a + b + c
 
-        result = yield self.pool.run_with_connection(txn, 1, 2, c=3)
+        result = yield self.pool.runWithConnection(txn, 1, 2, c=3)
         self.assertEqual(result, 6)
 
     @inlineCallbacks
@@ -106,7 +108,7 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
 
         # The failure crosses the thread boundary and surfaces as an errback.
         failure = yield self.assertFailure(
-            self.pool.run_with_connection(txn), MarkerError
+            self.pool.runWithConnection(txn), MarkerError
         )
         self.assertEqual(str(failure), "boom")
 
@@ -118,12 +120,12 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         def one(conn: Any) -> Any:
             cursor = conn.cursor()
             cursor.execute("SELECT 1::int")
-            row = cursor.fetch_one()
+            row = cursor.fetchone()
             conn.commit()
             return row
 
-        self.assertEqual((yield self.pool.run_with_connection(one)), (1,))
-        self.assertEqual((yield self.pool.run_with_connection(one)), (1,))
+        self.assertEqual((yield self.pool.runWithConnection(one)), (1,))
+        self.assertEqual((yield self.pool.runWithConnection(one)), (1,))
 
     @inlineCallbacks
     def test_concurrent_calls_are_serviced(self) -> Any:
@@ -132,16 +134,42 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase):
         def txn(conn: Any, n: int) -> Any:
             cursor = conn.cursor()
             cursor.execute("SELECT $1::int", [n])
-            row = cursor.fetch_one()
+            row = cursor.fetchone()
             conn.commit()
             return row
 
         results = yield gatherResults(
-            [self.pool.run_with_connection(txn, n) for n in range(10)]
+            [self.pool.runWithConnection(txn, n) for n in range(10)]
         )
         self.assertEqual(results, [(n,) for n in range(10)])
 
+    @inlineCallbacks
+    def test_drives_a_transaction_as_db_pool(self) -> Any:
+        # Mirror what DatabasePool.runWithConnection's inner_func does: the pool
+        # hands a DBAPI2 connection the engine can inspect, wrapped in a
+        # LoggingDatabaseConnection whose cursor is a real LoggingTransaction.
+        engine = RustPostgresEngine({})
+
+        def interaction(conn: Any) -> Any:
+            # A freshly checked-out connection is not mid-transaction.
+            self.assertFalse(engine.in_transaction(conn))
+
+            db_conn = LoggingDatabaseConnection(
+                conn=conn,
+                engine=engine,
+                default_txn_name="test",
+                server_name="test",
+            )
+            txn = db_conn.cursor(txn_name="test")
+            txn.execute("SELECT ?::int + ?::int", (2, 3))
+            row = txn.fetchone()
+            db_conn.commit()
+            return row
+
+        result = yield self.pool.runWithConnection(interaction)
+        self.assertEqual(result, (5,))
+
     def test_run_when_not_running_raises(self) -> None:
         self.pool.close()
         with self.assertRaises(RuntimeError):
-            self.pool.run_with_connection(lambda conn: None)
+            self.pool.runWithConnection(lambda conn: None)