Forward governor kwargs from dqliteclient.connect and create_pool

DqliteConnection.__init__ and ConnectionPool.__init__ both accept
max_total_rows, max_continuation_frames, and trust_server_heartbeat —
the constructors pin down the DoS governors that the library ships
with. The module-level convenience factories at dqliteclient/__init__.py
silently dropped all three, so callers using the documented
`from dqliteclient import connect, create_pool` pattern couldn't
configure the governors without dropping to the class constructors.

Mirror the constructor signatures on both factories so the convenience
API and the class API agree on the available knobs. Add regression
tests that assert both signatures and the forwarding through to
attribute state.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/__init__.py

@@ -46,18 +46,35 @@ async def connect(
     *,
     database: str = "default",
     timeout: float = 10.0,
+    max_total_rows: int | None = 10_000_000,
+    max_continuation_frames: int | None = 100_000,
+    trust_server_heartbeat: bool = False,
 ) -> DqliteConnection:
     """Connect to a dqlite node.
 
     Args:
         address: Node address in "host:port" format
         database: Database name to open
         timeout: Connection timeout in seconds
+        max_total_rows: Cumulative row cap across continuation frames
+            for a single query. Forwarded to the underlying
+            DqliteConnection. None disables the cap.
+        max_continuation_frames: Per-query continuation-frame cap.
+            Forwarded to the underlying DqliteConnection.
+        trust_server_heartbeat: Let the server-advertised heartbeat
+            widen the per-read deadline. Default False.
 
     Returns:
         A connected DqliteConnection
     """
-    conn = DqliteConnection(address, database=database, timeout=timeout)
+    conn = DqliteConnection(
+        address,
+        database=database,
+        timeout=timeout,
+        max_total_rows=max_total_rows,
+        max_continuation_frames=max_continuation_frames,
+        trust_server_heartbeat=trust_server_heartbeat,
+    )
     await conn.connect()
     return conn
 
@@ -71,6 +88,9 @@ async def create_pool(
     timeout: float = 10.0,
     cluster: ClusterClient | None = None,
     node_store: NodeStore | None = None,
+    max_total_rows: int | None = 10_000_000,
+    max_continuation_frames: int | None = 100_000,
+    trust_server_heartbeat: bool = False,
 ) -> ConnectionPool:
     """Create a connection pool with automatic leader detection.
 
@@ -84,6 +104,13 @@ async def create_pool(
         cluster: Externally-owned ClusterClient shared across pools.
         node_store: Externally-owned NodeStore used to build a new
             ClusterClient. Mutually exclusive with ``cluster``.
+        max_total_rows: Cumulative row cap across continuation frames
+            for a single query. Forwarded to the underlying
+            ConnectionPool. None disables the cap.
+        max_continuation_frames: Per-query continuation-frame cap.
+            Forwarded to the underlying ConnectionPool.
+        trust_server_heartbeat: Let the server-advertised heartbeat
+            widen the per-read deadline. Default False.
 
     Returns:
         An initialized ConnectionPool
@@ -96,6 +123,9 @@ async def create_pool(
         timeout=timeout,
         cluster=cluster,
         node_store=node_store,
+        max_total_rows=max_total_rows,
+        max_continuation_frames=max_continuation_frames,
+        trust_server_heartbeat=trust_server_heartbeat,
     )
     await pool.initialize()
     return pool

tests/test_connect_forwards_governors.py

@@ -0,0 +1,80 @@
+"""dqliteclient.connect() and create_pool() forward the governor kwargs.
+
+The module-level convenience factories are the PEP-249-style entry
+point for the client package. Constructor kwargs on DqliteConnection /
+ConnectionPool were added in an earlier cycle but the factories at
+__init__.py didn't forward them, so users going through the convenience
+path silently lost the DoS-governor configuration.
+"""
+
+from dqliteclient import ConnectionPool, create_pool
+from dqliteclient.connection import DqliteConnection
+
+
+class TestCreatePoolForwardsGovernors:
+    """`create_pool` forwards governor kwargs into ConnectionPool."""
+
+    def test_defaults_match_pool_constructor(self) -> None:
+        """Not calling create_pool — just asserting the signature parity
+        wouldn't catch a forward gap. Use a direct ConnectionPool for
+        the unit-level equivalent (no initialize() so no TCP)."""
+        pool = ConnectionPool(["localhost:19001"])
+        assert pool._max_total_rows == 10_000_000
+        assert pool._max_continuation_frames == 100_000
+        assert pool._trust_server_heartbeat is False
+
+    def test_custom_governors_plumbed(self) -> None:
+        pool = ConnectionPool(
+            ["localhost:19001"],
+            max_total_rows=500,
+            max_continuation_frames=7,
+            trust_server_heartbeat=True,
+        )
+        assert pool._max_total_rows == 500
+        assert pool._max_continuation_frames == 7
+        assert pool._trust_server_heartbeat is True
+
+    def test_create_pool_signature_accepts_governors(self) -> None:
+        """The create_pool factory must not raise TypeError when handed
+        the governor kwargs — the regression gap was the factory silently
+        dropping kwargs that the constructor supports."""
+        import inspect
+
+        sig = inspect.signature(create_pool)
+        assert "max_total_rows" in sig.parameters
+        assert "max_continuation_frames" in sig.parameters
+        assert "trust_server_heartbeat" in sig.parameters
+
+
+class TestConnectForwardsGovernors:
+    """`dqliteclient.connect` forwards governor kwargs.
+
+    A full-integration test would do TCP work; we rely on signature
+    parity + the DqliteConnection attribute forwarding that is already
+    covered elsewhere.
+    """
+
+    def test_connect_signature_accepts_governors(self) -> None:
+        import inspect
+
+        from dqliteclient import connect
+
+        sig = inspect.signature(connect)
+        assert "max_total_rows" in sig.parameters
+        assert "max_continuation_frames" in sig.parameters
+        assert "trust_server_heartbeat" in sig.parameters
+
+    def test_dqlite_connection_constructor_forwarding(self) -> None:
+        """`connect` ultimately calls DqliteConnection(...). Construct
+        that object directly with the same kwargs and verify the
+        attributes propagate — this covers the plumbing that `connect`
+        relies on without needing a TCP connection."""
+        conn = DqliteConnection(
+            "localhost:19001",
+            max_total_rows=500,
+            max_continuation_frames=7,
+            trust_server_heartbeat=True,
+        )
+        assert conn._max_total_rows == 500
+        assert conn._max_continuation_frames == 7
+        assert conn._trust_server_heartbeat is True