fix: invalidate connection state after protocol errors

After a ConnectionError or ProtocolError, the connection's _protocol
and _db_id are now cleared so is_connected returns False. Previously,
broken connections still reported is_connected=True, causing the pool
to hand out dead connections.

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

Author: antoineleclair

src/dqliteclient/connection.py

@@ -5,7 +5,7 @@
 from contextlib import asynccontextmanager
 from typing import Any
 
-from dqliteclient.exceptions import ConnectionError
+from dqliteclient.exceptions import ConnectionError, ProtocolError
 from dqliteclient.protocol import DqliteProtocol
 
 
@@ -92,24 +92,41 @@ def _ensure_connected(self) -> tuple[DqliteProtocol, int]:
             raise ConnectionError("Not connected")
         return self._protocol, self._db_id
 
+    def _invalidate(self) -> None:
+        """Mark the connection as broken after an unrecoverable error."""
+        self._protocol = None
+        self._db_id = None
+
     async def execute(self, sql: str, params: list[Any] | None = None) -> tuple[int, int]:
         """Execute a SQL statement.
 
         Returns (last_insert_id, rows_affected).
         """
         protocol, db_id = self._ensure_connected()
-        return await protocol.exec_sql(db_id, sql, params)
+        try:
+            return await protocol.exec_sql(db_id, sql, params)
+        except (ConnectionError, ProtocolError):
+            self._invalidate()
+            raise
 
     async def fetch(self, sql: str, params: list[Any] | None = None) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts."""
         protocol, db_id = self._ensure_connected()
-        columns, rows = await protocol.query_sql(db_id, sql, params)
+        try:
+            columns, rows = await protocol.query_sql(db_id, sql, params)
+        except (ConnectionError, ProtocolError):
+            self._invalidate()
+            raise
         return [dict(zip(columns, row, strict=True)) for row in rows]
 
     async def fetchall(self, sql: str, params: list[Any] | None = None) -> list[list[Any]]:
         """Execute a query and return results as list of lists."""
         protocol, db_id = self._ensure_connected()
-        _, rows = await protocol.query_sql(db_id, sql, params)
+        try:
+            _, rows = await protocol.query_sql(db_id, sql, params)
+        except (ConnectionError, ProtocolError):
+            self._invalidate()
+            raise
         return rows
 
     async def fetchone(self, sql: str, params: list[Any] | None = None) -> dict[str, Any] | None:
@@ -120,7 +137,11 @@ async def fetchone(self, sql: str, params: list[Any] | None = None) -> dict[str,
     async def fetchval(self, sql: str, params: list[Any] | None = None) -> Any:
         """Execute a query and return the first column of the first row."""
         protocol, db_id = self._ensure_connected()
-        _, rows = await protocol.query_sql(db_id, sql, params)
+        try:
+            _, rows = await protocol.query_sql(db_id, sql, params)
+        except (ConnectionError, ProtocolError):
+            self._invalidate()
+            raise
         if rows and rows[0]:
             return rows[0][0]
         return None

tests/test_connection.py

@@ -142,3 +142,37 @@ async def mock_execute(sql: str, params=None):
         assert "ROLLBACK" in call_log
         # _in_transaction was cleaned up
         assert not conn._in_transaction
+
+    async def test_connection_invalidated_after_protocol_error(self) -> None:
+        """After a connection error, is_connected should return False."""
+        conn = DqliteConnection("localhost:9001")
+
+        mock_reader = AsyncMock()
+        mock_writer = MagicMock()
+        mock_writer.drain = AsyncMock()
+        mock_writer.close = MagicMock()
+        mock_writer.wait_closed = AsyncMock()
+
+        from dqlitewire.messages import DbResponse, WelcomeResponse
+
+        responses = [
+            WelcomeResponse(heartbeat_timeout=15000).encode(),
+            DbResponse(db_id=1).encode(),
+        ]
+        mock_reader.read.side_effect = responses
+
+        with patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)):
+            await conn.connect()
+
+        assert conn.is_connected
+
+        # Now make the reader return empty (connection closed)
+        mock_reader.read.side_effect = [b""]
+
+        from dqliteclient.exceptions import ConnectionError as DqliteConnectionError
+
+        with pytest.raises(DqliteConnectionError):
+            await conn.execute("SELECT 1")
+
+        # Connection should be invalidated
+        assert not conn.is_connected