fix: invalidate connection on CancelledError to prevent decoder corruption

When a protocol read is cancelled mid-flight, the MessageDecoder may
contain partial data from an incomplete message, or the TCP buffer
may contain an orphaned response. Either way, the connection is
irrecoverably corrupted — subsequent reads would decode garbage or
return wrong responses.

Add `except BaseException: self._invalidate(); raise` to execute(),
fetch(), fetchall(), and fetchval() so that CancelledError (and other
BaseException subclasses) properly invalidate the connection rather
than leaving it in a poisoned state.

Closes #085

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

Author: antoineleclair

src/dqliteclient/connection.py

@@ -151,6 +151,9 @@ async def execute(self, sql: str, params: Sequence[Any] | None = None) -> tuple[
             if e.code in _LEADER_ERROR_CODES:
                 self._invalidate()
             raise
+        except BaseException:
+            self._invalidate()
+            raise
 
     async def fetch(self, sql: str, params: Sequence[Any] | None = None) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts."""
@@ -164,6 +167,9 @@ async def fetch(self, sql: str, params: Sequence[Any] | None = None) -> list[dic
             if e.code in _LEADER_ERROR_CODES:
                 self._invalidate()
             raise
+        except BaseException:
+            self._invalidate()
+            raise
         return [dict(zip(columns, row, strict=True)) for row in rows]
 
     async def fetchall(self, sql: str, params: Sequence[Any] | None = None) -> list[list[Any]]:
@@ -178,6 +184,9 @@ async def fetchall(self, sql: str, params: Sequence[Any] | None = None) -> list[
             if e.code in _LEADER_ERROR_CODES:
                 self._invalidate()
             raise
+        except BaseException:
+            self._invalidate()
+            raise
         return rows
 
     async def fetchone(
@@ -199,6 +208,9 @@ async def fetchval(self, sql: str, params: Sequence[Any] | None = None) -> Any:
             if e.code in _LEADER_ERROR_CODES:
                 self._invalidate()
             raise
+        except BaseException:
+            self._invalidate()
+            raise
         if rows and rows[0]:
             return rows[0][0]
         return None

tests/test_connection.py

@@ -226,6 +226,56 @@ async def mock_execute(sql: str, params=None):
         # _in_transaction was cleaned up
         assert not conn._in_transaction
 
+    async def test_cancellation_invalidates_connection(self) -> None:
+        """CancelledError during a query must invalidate the connection."""
+        import asyncio
+
+        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
+
+        # Make the reader hang forever (will be cancelled)
+        read_entered = asyncio.Event()
+
+        async def hanging_read(*args, **kwargs):
+            read_entered.set()
+            await asyncio.sleep(100)
+
+        mock_reader.read.side_effect = hanging_read
+
+        async def do_execute():
+            await conn.execute("INSERT INTO t VALUES (1)")
+
+        task = asyncio.create_task(do_execute())
+        await read_entered.wait()
+        task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        # Connection must be invalidated — the decoder may have partial data
+        assert not conn.is_connected, (
+            "Connection should be invalidated after CancelledError to prevent "
+            "decoder corruption from partial reads"
+        )
+
     async def test_connection_invalidated_after_protocol_error(self) -> None:
         """After a connection error, is_connected should return False."""
         conn = DqliteConnection("localhost:9001")