fix: wrap writer.drain and _read_data transport errors as DqliteConnectionError

Each protocol method awaited writer.drain() directly; ConnectionReset,
BrokenPipeError, OSError, and RuntimeError("Transport is closed") leaked
to callers raw. Application code catching DqliteConnectionError to retry
missed write-path transport errors entirely. _read_data had the same gap
on the read path: only TimeoutError was wrapped.

Add a _send() helper that wraps the drain and route all seven call sites
through it, and extend _read_data's except to the same transport-error
set. Parametrized tests cover the four known leaking types for both
paths and assert __cause__ is preserved.

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

Author: antoineleclair

src/dqliteclient/protocol.py

@@ -56,7 +56,7 @@ async def handshake(self, client_id: int | None = None) -> int:
         # Send protocol version + client registration together
         request = ClientRequest(client_id=client_id)
         self._writer.write(MessageEncoder().encode_handshake() + request.encode())
-        await self._writer.drain()
+        await self._send()
 
         # Read welcome response
         response = await self._read_response()
@@ -83,7 +83,7 @@ async def get_leader(self) -> tuple[int, str]:
         """
         request = LeaderRequest()
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -102,7 +102,7 @@ async def open_database(self, name: str, flags: int = 0, vfs: str = "") -> int:
         """
         request = OpenRequest(name=name, flags=flags, vfs=vfs)
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -121,7 +121,7 @@ async def prepare(self, db_id: int, sql: str) -> tuple[int, int]:
         """
         request = PrepareRequest(db_id=db_id, sql=sql)
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -137,7 +137,7 @@ async def finalize(self, db_id: int, stmt_id: int) -> None:
         """Finalize (close) a prepared statement."""
         request = FinalizeRequest(db_id=db_id, stmt_id=stmt_id)
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -159,7 +159,7 @@ async def exec_sql(
         """
         request = ExecSqlRequest(db_id=db_id, sql=sql, params=params if params is not None else [])
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -182,7 +182,7 @@ async def query_sql(
         """
         request = QuerySqlRequest(db_id=db_id, sql=sql, params=params if params is not None else [])
         self._writer.write(request.encode())
-        await self._writer.drain()
+        await self._send()
 
         response = await self._read_response()
 
@@ -214,12 +214,26 @@ async def query_sql(
 
         return column_names, all_rows
 
+    async def _send(self) -> None:
+        """Drain the writer, wrapping transport errors as DqliteConnectionError."""
+        try:
+            await self._writer.drain()
+        except (ConnectionError, OSError, RuntimeError) as e:
+            raise DqliteConnectionError(f"Write failed: {e}") from e
+
     async def _read_data(self) -> bytes:
-        """Read data from the stream with timeout."""
+        """Read data from the stream with timeout.
+
+        Transport errors (ConnectionResetError, BrokenPipeError, OSError,
+        RuntimeError("Transport is closed")) are wrapped in
+        DqliteConnectionError to match the write-path behaviour.
+        """
         try:
             data = await asyncio.wait_for(self._reader.read(4096), timeout=self._timeout)
         except TimeoutError:
             raise DqliteConnectionError(f"Server read timed out after {self._timeout}s") from None
+        except (ConnectionError, OSError, RuntimeError) as e:
+            raise DqliteConnectionError(f"Read failed: {e}") from e
         if not data:
             raise DqliteConnectionError("Connection closed by server")
         return data

tests/test_protocol.py

@@ -16,6 +16,59 @@ class TestDqliteProtocol:
     def protocol(self, mock_reader: AsyncMock, mock_writer: MagicMock) -> DqliteProtocol:
         return DqliteProtocol(mock_reader, mock_writer)
 
+    @pytest.mark.parametrize(
+        "err",
+        [
+            ConnectionResetError("peer reset"),
+            BrokenPipeError("pipe gone"),
+            OSError("generic os error"),
+            RuntimeError("Transport is closed"),
+        ],
+    )
+    async def test_writer_drain_errors_are_wrapped(
+        self,
+        protocol: DqliteProtocol,
+        mock_reader: AsyncMock,
+        mock_writer: MagicMock,
+        err: BaseException,
+    ) -> None:
+        """Transport errors raised by writer.drain() must surface as
+        DqliteConnectionError so callers catching the client's error
+        hierarchy don't miss them, with the original attached via __cause__.
+        """
+        from dqliteclient.exceptions import DqliteConnectionError
+
+        mock_writer.drain = AsyncMock(side_effect=err)
+
+        with pytest.raises(DqliteConnectionError) as exc_info:
+            await protocol.handshake()
+        assert exc_info.value.__cause__ is err
+
+    @pytest.mark.parametrize(
+        "err",
+        [
+            ConnectionResetError("peer reset mid-read"),
+            BrokenPipeError("pipe gone mid-read"),
+            OSError("generic os error mid-read"),
+            RuntimeError("Transport is closed mid-read"),
+        ],
+    )
+    async def test_reader_errors_are_wrapped(
+        self,
+        protocol: DqliteProtocol,
+        mock_reader: AsyncMock,
+        err: BaseException,
+    ) -> None:
+        """Transport errors on _read_data must also surface as
+        DqliteConnectionError, matching the write-path behaviour.
+        """
+        from dqliteclient.exceptions import DqliteConnectionError
+
+        mock_reader.read.side_effect = err
+        with pytest.raises(DqliteConnectionError) as exc_info:
+            await protocol._read_data()
+        assert exc_info.value.__cause__ is err
+
     async def test_handshake_success(
         self,
         protocol: DqliteProtocol,