Skip to content

Commit f7be8e6

Browse files
Add leader-redirect-on-connect to dqlitedbapi.connect
Production dqlite use requires leader-redirect-on-connect — Raft leadership can flip between any two connections, and a SA pool's reconnect-after-pre-ping path otherwise surfaces ``SQLITE_IOERR_NOT_LEADER`` from the demoted ex-leader instead of recovering. Mirrors go-dqlite's ``database/sql`` driver layering (``client.NewLeaderConnector(store)`` is what every connect goes through there). Implementation: introduce ``_resolve_leader(seed, *, timeout)`` that wraps the seed in a single-node ``MemoryNodeStore`` and calls ``ClusterClient.find_leader()``, then ``_build_and_connect`` runs this BEFORE constructing the ``DqliteConnection``. The OPEN request now lands on the leader's actual address, not the seed. Both sync and async connect paths share the same ``_build_and_connect`` and inherit the fix transparently. Error translation: - ``ClusterError`` (no leader reachable from seed) → ``OperationalError("Failed to find leader from <seed>: ...")``. Distinct from the post-find-leader connect failure prefix (``"Failed to connect: ..."``) so log triage can tell the two apart. SA's pool retry-on-OperationalError handles both. - ``ClusterPolicyError`` (operator allowlist denied a redirect target) → ``InterfaceError("Cluster policy rejection during leader discovery; ...")``. Permanent config mismatch — SA's ``is_disconnect`` does NOT enter a retry loop. Tests: - ``test_leader_redirect_on_connect.py`` (9) — unit boundary at ``_build_and_connect`` and ``_resolve_leader``: happy paths, redirect, error translation arms, mid-flip race. - ``test_leader_redirect_live.py`` (4) — live-cluster pins: seed-as-leader, follower-redirects-to-leader, flip-then-reconnect, unreachable seed. Lockstep test updates: 4 existing tests that mock at the ``DqliteConnection`` level needed an additional ``_resolve_leader`` mock to bypass the new bootstrap step (previously the leader-discovery used a real network probe that the mocks didn't intercept). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 187e080 commit f7be8e6

8 files changed

Lines changed: 518 additions & 11 deletions

src/dqlitedbapi/connection.py

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
import dqliteclient.exceptions as _client_exc
1717
from dqliteclient import DqliteConnection, validate_positive_int_or_none
1818
from dqliteclient import connection as _client_conn_mod
19+
from dqliteclient.cluster import ClusterClient
1920
from dqliteclient.connection import parse_address as _client_parse_address
21+
from dqliteclient.node_store import MemoryNodeStore
2022
from dqlitedbapi import exceptions as _exc
2123
from dqlitedbapi.cursor import Cursor, _call_client
2224
from dqlitedbapi.exceptions import (
@@ -140,6 +142,30 @@ def _validate_close_timeout(close_timeout: float) -> None:
140142
)
141143

142144

145+
async def _resolve_leader(address: str, *, timeout: float) -> str:
146+
"""Resolve the cluster's current leader address from a seed.
147+
148+
Bootstraps from the user-supplied ``address`` (the URL host:port
149+
in the SA dialect's case) and uses :class:`ClusterClient` to
150+
follow the leader-redirect chain — same pattern go-dqlite's
151+
``database/sql`` driver implements via
152+
``client.NewLeaderConnector(store)``. Without this step,
153+
connecting to a demoted-leader address surfaces
154+
``SQLITE_IOERR_NOT_LEADER`` from the server even though the
155+
cluster has a healthy leader at a different address; the SA
156+
pool's reconnect-after-pre-ping path cannot recover.
157+
158+
Wraps the seed in a single-node :class:`MemoryNodeStore` and
159+
delegates to :meth:`ClusterClient.find_leader`. Returns the
160+
leader's address on success; raises the underlying
161+
``ClusterError`` / ``ClusterPolicyError`` for the surrounding
162+
error-translation arms in :func:`_build_and_connect` to handle.
163+
"""
164+
store = MemoryNodeStore([address])
165+
cluster = ClusterClient(store, timeout=timeout)
166+
return await cluster.find_leader()
167+
168+
143169
async def _build_and_connect(
144170
address: str,
145171
*,
@@ -152,14 +178,53 @@ async def _build_and_connect(
152178
) -> DqliteConnection:
153179
"""Build a DqliteConnection with the given governors and connect it.
154180
155-
Wraps the construct-then-connect sequence that both the sync and
156-
async Connection flavours execute under their respective locks. The
157-
``OperationalError`` message phrasing ("Failed to connect: ...") is
158-
intentionally verbatim so test assertions that match on the prefix
159-
continue to pass.
181+
Performs the dqlite production-grade connect sequence:
182+
183+
1. Resolve the current leader via :func:`_resolve_leader` (one
184+
round-trip against the seed; if the seed is the leader, the
185+
leader-info reply is its own address).
186+
2. Construct + connect a :class:`DqliteConnection` against the
187+
leader address.
188+
189+
Wraps the sequence that both the sync and async Connection
190+
flavours execute under their respective locks. The
191+
``OperationalError`` message phrasing ("Failed to connect: ...")
192+
is intentionally verbatim so test assertions that match on the
193+
prefix continue to pass.
194+
195+
Mirrors the canonical go-dqlite/driver layering — applications
196+
should not need to special-case leader-flips between
197+
connections; the dbapi handles the redirect transparently.
160198
"""
199+
try:
200+
leader_address = await _resolve_leader(address, timeout=timeout)
201+
except _client_exc.ClusterPolicyError as e:
202+
# Operator allowlist rejected a redirect target. Surface as
203+
# InterfaceError with the canonical prefix — symmetric with
204+
# the post-construct ClusterPolicyError arm below.
205+
raw_msg = getattr(e, "raw_message", None) or str(e)
206+
raise InterfaceError(
207+
f"Cluster policy rejection during leader discovery; {e}",
208+
code=None,
209+
raw_message=raw_msg,
210+
) from e
211+
except _client_exc.ClusterError as e:
212+
# All nodes in the seed's resolved store rejected the leader
213+
# query (no node is currently leader, all unreachable, etc.).
214+
# Surface as OperationalError so the SA pool's retry loop
215+
# classifies it correctly. Different from the post-construct
216+
# ClusterError arm only in the message prefix — operators
217+
# reading logs need to tell "couldn't find leader" from
218+
# "found leader but couldn't connect".
219+
raw_msg = getattr(e, "raw_message", None) or str(e)
220+
raise OperationalError(
221+
f"Failed to find leader from {address}: {e}",
222+
code=None,
223+
raw_message=raw_msg,
224+
) from e
225+
161226
conn = DqliteConnection(
162-
address,
227+
leader_address,
163228
database=database,
164229
timeout=timeout,
165230
max_total_rows=max_total_rows,

tests/conftest.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,20 @@
11
"""Pytest configuration for dqlite-dbapi tests."""
2+
3+
import sys
4+
from pathlib import Path
5+
6+
# Add python-dqlite-dev's testlib to sys.path so tests (in particular
7+
# the leader-redirect integration suite) can import shared utilities
8+
# from ``dqlitetestlib``. ``python-dqlite-dev`` is expected as a
9+
# sibling of this checkout — see ``python-dqlite-dev/testlib/README.md``.
10+
# The insertion is harmless when the sibling repo is absent.
11+
_TESTLIB = Path(__file__).resolve().parent.parent.parent / "python-dqlite-dev" / "testlib"
12+
if _TESTLIB.exists() and str(_TESTLIB) not in sys.path:
13+
sys.path.insert(0, str(_TESTLIB))
14+
15+
# Pytest 8+ requires ``pytest_plugins`` at the top-level conftest.
16+
# Only register the testlib's fixtures plugin when the path resolved
17+
# so consumers running unit tests without the sibling repo see no
18+
# difference.
19+
if _TESTLIB.exists():
20+
pytest_plugins = ["dqlitetestlib.fixtures"]
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
"""Live integration: leader-redirect-on-connect against a real cluster.
2+
3+
End-to-end coverage for the ``_resolve_leader`` step added to
4+
``_build_and_connect``. Mocked unit coverage is in
5+
``tests/test_leader_redirect_on_connect.py``; this file pins the
6+
behaviour against the live cluster so a regression in the wire
7+
layer or in ``ClusterClient.find_leader`` would surface here even
8+
if the unit tests stay green.
9+
10+
Tests use ``cluster_control`` from ``dqlitetestlib`` (bootstrapped
11+
in the top-level ``tests/conftest.py``) where they need
12+
deterministic leader manipulation. Each test that mutates cluster
13+
topology restores the original leader on its way out so subsequent
14+
tests in the session see a stable starting state.
15+
"""
16+
17+
from __future__ import annotations
18+
19+
import contextlib
20+
import os
21+
from typing import TYPE_CHECKING
22+
23+
import pytest
24+
25+
from dqlitedbapi import connect
26+
from dqlitedbapi.exceptions import OperationalError
27+
28+
if TYPE_CHECKING:
29+
from dqlitetestlib import TestClusterControl # type: ignore[import-not-found]
30+
31+
32+
def _node_addresses() -> list[str]:
33+
"""Read DQLITE_TEST_CLUSTER_NODES with the python-dqlite-dev
34+
default."""
35+
raw = os.environ.get(
36+
"DQLITE_TEST_CLUSTER_NODES",
37+
"localhost:9001,localhost:9002,localhost:9003",
38+
)
39+
return [s.strip() for s in raw.split(",") if s.strip()]
40+
41+
42+
# --- normal flows ---
43+
44+
45+
@pytest.mark.integration
46+
def test_connect_with_seed_as_leader_succeeds() -> None:
47+
"""Happy path: the seed address is the current leader.
48+
``_resolve_leader`` returns the seed verbatim and the OPEN
49+
runs against it."""
50+
seed = os.environ.get("DQLITE_TEST_CLUSTER", "localhost:9001")
51+
conn = connect(seed, timeout=5.0)
52+
try:
53+
cur = conn.cursor()
54+
cur.execute("SELECT 1")
55+
assert cur.fetchall() == [(1,)]
56+
finally:
57+
conn.close()
58+
59+
60+
@pytest.mark.integration
61+
def test_connect_via_follower_address_redirects_to_leader() -> None:
62+
"""The seed is a follower; ``_resolve_leader`` follows the
63+
redirect and OPEN runs against the actual leader. Without
64+
leader-redirect-on-connect this would fail with
65+
``SQLITE_IOERR_NOT_LEADER`` from the follower's OPEN handler.
66+
67+
Picks the non-leader nodes from ``DQLITE_TEST_CLUSTER_NODES``
68+
and connects through each — at least two of the three nodes
69+
are followers in the steady state, so this is a robust pin.
70+
"""
71+
import asyncio
72+
73+
from dqliteclient.cluster import ClusterClient
74+
from dqliteclient.node_store import MemoryNodeStore
75+
76+
addresses = _node_addresses()
77+
78+
async def _resolve() -> str:
79+
store = MemoryNodeStore(addresses)
80+
cluster = ClusterClient(store, timeout=5.0)
81+
return await cluster.find_leader()
82+
83+
leader_addr = asyncio.run(_resolve())
84+
follower_addrs = [a for a in addresses if a != leader_addr]
85+
assert follower_addrs, (
86+
f"expected at least one follower in {addresses!r}; leader is {leader_addr!r}"
87+
)
88+
89+
for follower in follower_addrs:
90+
conn = connect(follower, timeout=5.0)
91+
try:
92+
cur = conn.cursor()
93+
cur.execute("SELECT 1")
94+
assert cur.fetchall() == [(1,)]
95+
finally:
96+
conn.close()
97+
98+
99+
@pytest.mark.integration
100+
def test_connect_after_leader_flip_routes_to_new_leader(
101+
cluster_control: TestClusterControl,
102+
) -> None:
103+
"""Force a leader flip; a brand-new ``connect()`` against the
104+
OLD-leader address (now a follower) succeeds because
105+
``_resolve_leader`` follows the redirect to the new leader.
106+
Restores the original leader on the way out."""
107+
import asyncio
108+
109+
starting = asyncio.run(cluster_control.current_leader_node())
110+
seed = starting.address.replace("127.0.0.1", "localhost")
111+
112+
flip = asyncio.run(cluster_control.force_leader_flip())
113+
assert flip.target.node_id != starting.node_id
114+
115+
try:
116+
# The seed is now a follower (the demoted ex-leader).
117+
# ``_resolve_leader`` should follow the redirect and the
118+
# OPEN should reach the new leader.
119+
conn = connect(seed, timeout=5.0)
120+
try:
121+
cur = conn.cursor()
122+
cur.execute("SELECT 1")
123+
assert cur.fetchall() == [(1,)]
124+
finally:
125+
conn.close()
126+
finally:
127+
with contextlib.suppress(Exception):
128+
asyncio.run(cluster_control.transfer_leadership_to(starting.node_id))
129+
130+
131+
# --- failure flows ---
132+
133+
134+
@pytest.mark.integration
135+
def test_connect_to_unreachable_seed_raises_operational_error() -> None:
136+
"""Seed unreachable: ``_resolve_leader`` cannot reach any node
137+
in its 1-node store; surfaces as ``OperationalError`` with the
138+
canonical ``Failed to find leader from`` prefix."""
139+
# Pick a port we know nothing is listening on.
140+
with pytest.raises(OperationalError, match="Failed to find leader"):
141+
conn = connect("127.0.0.1:1", timeout=1.0)
142+
conn.connect() # explicit connect for clearer failure point

tests/test_async_close_race.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,13 @@ async def slow_execute(_sql: str, _params: object) -> tuple[int, int]:
3131
async def fake_query_raw_typed(_sql: str, _params: object) -> tuple[list, list, list, list]: # type: ignore[type-arg]
3232
return ([], [], [], [])
3333

34-
with patch("dqlitedbapi.connection.DqliteConnection") as MockDqliteConn:
34+
with (
35+
patch(
36+
"dqlitedbapi.connection._resolve_leader",
37+
new=AsyncMock(side_effect=lambda address, *, timeout: address),
38+
),
39+
patch("dqlitedbapi.connection.DqliteConnection") as MockDqliteConn,
40+
):
3541
mock_instance = AsyncMock()
3642
mock_instance.connect = AsyncMock()
3743
mock_instance.execute = slow_execute

tests/test_close_timeout_plumbing.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,13 @@ def test_propagates_to_underlying_dqlite_connection(self) -> None:
8080
the dbapi-level close_timeout so the actual wait_closed drain
8181
honours the caller's budget.
8282
"""
83-
with patch("dqlitedbapi.connection.DqliteConnection") as MockConn:
83+
with (
84+
patch(
85+
"dqlitedbapi.connection._resolve_leader",
86+
new=AsyncMock(side_effect=lambda address, *, timeout: address),
87+
),
88+
patch("dqlitedbapi.connection.DqliteConnection") as MockConn,
89+
):
8490
instance = AsyncMock()
8591
instance.connect = AsyncMock()
8692
MockConn.return_value = instance

tests/test_connection_closed_checks.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,8 +190,17 @@ def test_policy_rejection_surfaces_as_interface_error(
190190
async def _raise_policy(*args: object, **kwargs: object) -> None:
191191
raise _client_exc.ClusterPolicyError("not allowed")
192192

193-
# Patch the inner DqliteConnection.connect — this is what
194-
# ``_build_and_connect`` awaits after constructing the conn.
193+
# Patch ``_resolve_leader`` to short-circuit the leader-discovery
194+
# step so the test exercises the post-find-leader
195+
# ``DqliteConnection.connect`` arm. Then patch
196+
# ``DqliteConnection.connect`` to raise the policy error this
197+
# test pins.
198+
async def _identity_resolve(address: str, *, timeout: float) -> str:
199+
return address
200+
201+
monkeypatch.setattr(
202+
"dqlitedbapi.connection._resolve_leader", _identity_resolve, raising=True
203+
)
195204
monkeypatch.setattr("dqliteclient.DqliteConnection.connect", _raise_policy, raising=True)
196205

197206
conn = connect("localhost:19001", timeout=2.0)

0 commit comments

Comments
 (0)