Forward trust_server_heartbeat into the leader-discovery probe

_query_leader constructed a fresh DqliteProtocol without threading the
flag through, so operators who opted into a widened heartbeat window
for the main query path got the default tight timeout during leader
discovery. Thread trust_server_heartbeat through find_leader and
_query_leader; connect() forwards the value at call time.

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

Author: antoineleclair

src/dqliteclient/cluster.py

@@ -78,10 +78,13 @@ def _check_redirect(self, address: str) -> None:
         if not self._redirect_policy(address):
             raise ClusterError(f"Leader redirect to {address!r} rejected by redirect_policy")
 
-    async def find_leader(self) -> str:
+    async def find_leader(self, *, trust_server_heartbeat: bool = False) -> str:
         """Find the current cluster leader.
 
-        Returns the leader address.
+        Returns the leader address. ``trust_server_heartbeat`` is forwarded
+        to each probe protocol so operators who opted into a widened
+        heartbeat window for the main query path get the same semantics
+        during leader discovery.
         """
         nodes = await self._node_store.get_nodes()
 
@@ -109,7 +112,11 @@ async def find_leader(self) -> str:
         for node in nodes:
             try:
                 leader_address = await asyncio.wait_for(
-                    self._query_leader(node.address), timeout=self._timeout
+                    self._query_leader(
+                        node.address,
+                        trust_server_heartbeat=trust_server_heartbeat,
+                    ),
+                    timeout=self._timeout,
                 )
                 if leader_address:
                     # Only leader_address values that did NOT come from
@@ -133,7 +140,9 @@ async def find_leader(self) -> str:
 
         raise ClusterError(f"Could not find leader. Errors: {'; '.join(errors)}") from last_exc
 
-    async def _query_leader(self, address: str) -> str | None:
+    async def _query_leader(
+        self, address: str, *, trust_server_heartbeat: bool = False
+    ) -> str | None:
         """Query a node for the current leader."""
         host, port = _parse_address(address)
 
@@ -146,7 +155,12 @@ async def _query_leader(self, address: str) -> str | None:
             return None
 
         try:
-            protocol = DqliteProtocol(reader, writer, timeout=self._timeout)
+            protocol = DqliteProtocol(
+                reader,
+                writer,
+                timeout=self._timeout,
+                trust_server_heartbeat=trust_server_heartbeat,
+            )
             await protocol.handshake()
             node_id, leader_addr = await protocol.get_leader()
 
@@ -198,7 +212,9 @@ async def try_connect() -> DqliteConnection:
             attempt = attempt_counter[0]
             leader: str | None = None
             try:
-                leader = await self.find_leader()
+                leader = await self.find_leader(
+                    trust_server_heartbeat=trust_server_heartbeat,
+                )
                 conn = DqliteConnection(
                     leader,
                     database=database,

tests/test_cluster.py

@@ -52,10 +52,12 @@ async def test_find_leader_success(self) -> None:
 
         from dqlitewire.messages import LeaderResponse, WelcomeResponse
 
-        # First call for handshake, second for leader query
+        # Upstream raft_leader sets id and address atomically: a voter
+        # that IS the leader returns its own id AND its own address
+        # (never (nonzero, "")).
         responses = [
             WelcomeResponse(heartbeat_timeout=15000).encode(),
-            LeaderResponse(node_id=1, address="").encode(),  # Empty = this node is leader
+            LeaderResponse(node_id=1, address="localhost:9001").encode(),
         ]
         mock_reader.read.side_effect = responses
 
@@ -217,7 +219,7 @@ async def test_find_leader_skips_node_with_bad_handshake(self) -> None:
         responses = [
             b"\x00" * 64,
             WelcomeResponse(heartbeat_timeout=15000).encode(),
-            LeaderResponse(node_id=2, address="").encode(),
+            LeaderResponse(node_id=2, address="localhost:9002").encode(),
         ]
         mock_reader.read.side_effect = responses
 
@@ -251,7 +253,7 @@ async def hang_forever():
 
         responses = [
             WelcomeResponse(heartbeat_timeout=15000).encode(),
-            LeaderResponse(node_id=1, address="").encode(),
+            LeaderResponse(node_id=1, address="localhost:9001").encode(),
         ]
         mock_reader.read.side_effect = responses
 
@@ -276,7 +278,7 @@ async def test_find_leader_propagates_programming_bugs(self) -> None:
         store = MemoryNodeStore(["localhost:9001", "localhost:9002"])
         client = ClusterClient(store, timeout=0.5)
 
-        async def buggy_query(_address: str) -> str | None:
+        async def buggy_query(_address: str, **_kwargs: object) -> str | None:
             raise TypeError("programmer mistake")
 
         with (
@@ -297,7 +299,7 @@ async def test_find_leader_transport_error_chains_cause(self) -> None:
 
         boom = DqliteConnectionError("handshake failed")
 
-        async def failing_query(_address: str) -> str | None:
+        async def failing_query(_address: str, **_kwargs: object) -> str | None:
             raise boom
 
         with (
@@ -318,7 +320,7 @@ async def test_find_leader_randomizes_node_order(self) -> None:
 
         first_probed: list[str] = []
 
-        async def track(address: str) -> str | None:
+        async def track(address: str, **_kwargs: object) -> str | None:
             first_probed.append(address)
             raise DqliteConnectionError("not leader")
 
@@ -350,7 +352,7 @@ async def test_find_leader_probes_voters_before_non_voters(self) -> None:
 
         order: list[str] = []
 
-        async def track(address: str) -> str | None:
+        async def track(address: str, **_kwargs: object) -> str | None:
             order.append(address)
             return None  # no leader known — keep probing
 
@@ -383,7 +385,7 @@ async def test_connect_does_not_retry_plain_sql_errors(self) -> None:
 
         call_count = 0
 
-        async def always_sql_error() -> str:
+        async def always_sql_error(**_kwargs: object) -> str:
             nonlocal call_count
             call_count += 1
             raise OperationalError(1, "some sql error")
@@ -431,7 +433,7 @@ async def test_max_attempts_override_honored(self) -> None:
 
         call_count = [0]
 
-        async def fake_find_leader() -> str:
+        async def fake_find_leader(**_kwargs: object) -> str:
             call_count[0] += 1
             raise DqliteConnectionError("unreachable")
 
@@ -457,7 +459,7 @@ async def test_failed_attempts_logged(self, caplog: pytest.LogCaptureFixture) ->
         store = MemoryNodeStore(["localhost:1"])  # unreachable
         client = ClusterClient(store, timeout=0.1)
 
-        async def fake_find_leader() -> str:
+        async def fake_find_leader(**_kwargs: object) -> str:
             raise DqliteConnectionError("simulated")
 
         client.find_leader = fake_find_leader  # type: ignore[method-assign]
@@ -472,3 +474,68 @@ async def fake_find_leader() -> str:
             f"Expected 2 per-attempt log lines, got {len(attempt_logs)}: "
             f"{[r.message for r in attempt_logs]}"
         )
+
+
+class TestQueryLeaderTrustsHeartbeat:
+    """_query_leader forwards the trust_server_heartbeat flag."""
+
+    async def test_flag_propagates_to_probe_protocol(self) -> None:
+        store = MemoryNodeStore(["localhost:9001"])
+        client = ClusterClient(store, timeout=1.0)
+
+        mock_reader = AsyncMock()
+        mock_writer = MagicMock()
+        mock_writer.drain = AsyncMock()
+        mock_writer.close = MagicMock()
+        mock_writer.wait_closed = AsyncMock()
+
+        captured: dict[str, object] = {}
+
+        class FakeProto:
+            def __init__(self, *args: object, **kwargs: object) -> None:
+                captured.update(kwargs)
+
+            async def handshake(self) -> None:
+                pass
+
+            async def get_leader(self) -> tuple[int, str]:
+                return (1, "localhost:9001")
+
+        with (
+            patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)),
+            patch("dqliteclient.cluster.DqliteProtocol", FakeProto),
+        ):
+            await client._query_leader("localhost:9001", trust_server_heartbeat=True)
+
+        assert captured.get("trust_server_heartbeat") is True
+
+    async def test_flag_default_false(self) -> None:
+        store = MemoryNodeStore(["localhost:9001"])
+        client = ClusterClient(store, timeout=1.0)
+
+        mock_reader = AsyncMock()
+        mock_writer = MagicMock()
+        mock_writer.drain = AsyncMock()
+        mock_writer.close = MagicMock()
+        mock_writer.wait_closed = AsyncMock()
+
+        captured: dict[str, object] = {}
+
+        class FakeProto:
+            def __init__(self, *args: object, **kwargs: object) -> None:
+                captured.update(kwargs)
+
+            async def handshake(self) -> None:
+                pass
+
+            async def get_leader(self) -> tuple[int, str]:
+                return (1, "localhost:9001")
+
+        with (
+            patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)),
+            patch("dqliteclient.cluster.DqliteProtocol", FakeProto),
+        ):
+            await client._query_leader("localhost:9001")
+
+        assert captured.get("trust_server_heartbeat") is False
+