fix: preserve healthy connections in pool when user code raises exceptions

The pool's acquire() context manager previously closed the connection
unconditionally on any exception — including application errors like
ValueError, KeyError, or constraint-violation OperationalError. The
connection was perfectly healthy in these cases but got destroyed,
forcing a full TCP + dqlite handshake + leader discovery to replace it.

Now the except handler checks conn.is_connected before deciding:
- If True (healthy): return the connection to the pool for reuse
- If False (broken, already invalidated by execute/fetch handlers):
  close and discard

This matches asyncpg's behavior, where connections are only discarded
when they are actually in an error state.

Closes #090

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -123,25 +123,36 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
         try:
             yield conn
         except BaseException:
-            # On error (including cancellation), close connection
-            with contextlib.suppress(BaseException):
-                await conn.close()
-            self._in_use.discard(conn)
-            self._size -= 1
-            raise
-        else:
             self._in_use.discard(conn)
-            # Return to pool
-            if self._closed:
-                await conn.close()
-                self._size -= 1
-            else:
+            if conn.is_connected and not self._closed:
+                # Connection is healthy — user code raised a non-connection error.
+                # Return it to the pool instead of destroying it.
                 try:
                     self._pool.put_nowait(conn)
                 except asyncio.QueueFull:
-                    # Pool full, close connection
                     await conn.close()
                     self._size -= 1
+            else:
+                # Connection is broken (invalidated by execute/fetch error handlers).
+                with contextlib.suppress(BaseException):
+                    await conn.close()
+                self._size -= 1
+            raise
+        else:
+            self._in_use.discard(conn)
+            await self._release(conn)
+
+    async def _release(self, conn: DqliteConnection) -> None:
+        """Return a connection to the pool or close it."""
+        if self._closed:
+            await conn.close()
+            self._size -= 1
+        else:
+            try:
+                self._pool.put_nowait(conn)
+            except asyncio.QueueFull:
+                await conn.close()
+                self._size -= 1
 
     async def execute(self, sql: str, params: Sequence[Any] | None = None) -> tuple[int, int]:
         """Execute a SQL statement using a pooled connection."""

tests/test_pool.py

@@ -132,8 +132,8 @@ async def fail_then_succeed(**kwargs):
 
         await pool.close()
 
-    async def test_cancellation_does_not_leak_connection(self) -> None:
-        """Cancelling a task that holds a connection should clean it up."""
+    async def test_cancellation_returns_healthy_connection_to_pool(self) -> None:
+        """Cancelling a task holding a healthy connection should return it to the pool."""
         import asyncio
 
         pool = ConnectionPool(["localhost:9001"], max_size=1)
@@ -152,18 +152,21 @@ async def test_cancellation_does_not_leak_connection(self) -> None:
         async def hold_connection():
             async with pool.acquire() as _:
                 acquired.set()
-                await asyncio.sleep(10)  # Hold forever
+                await asyncio.sleep(10)  # Hold forever — will be cancelled
 
         task = asyncio.create_task(hold_connection())
-        await acquired.wait()  # Deterministic: wait until connection is acquired
+        await acquired.wait()
 
         task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
             await task
 
-        # Connection should have been closed, size decremented
-        mock_conn.close.assert_called()
-        assert pool._size < initial_size
+        # Connection is still healthy (CancelledError was in user code, not in
+        # a protocol operation), so it should be returned to the pool
+        assert pool._size == initial_size
+        assert pool._pool.qsize() == 1
+
+        await pool.close()
 
     async def test_acquire_timeout_when_pool_exhausted(self) -> None:
         """acquire() should timeout, not block forever, when pool is exhausted."""
@@ -252,7 +255,9 @@ async def holder():
                     assert conn1 is mock_conn1
                     holder_acquired.set()
                     await holder_release.wait()
-                    raise RuntimeError("simulated failure")
+                    # Simulate a real connection failure (broken connection)
+                    conn1.is_connected = False
+                    raise DqliteConnectionError("connection lost")
 
             holder_task = asyncio.create_task(holder())
             await holder_acquired.wait()
@@ -273,9 +278,9 @@ async def waiter():
             # check (_size == max_size) and enter the queue.get() wait
             await asyncio.sleep(0.1)
 
-            # Now release the holder — it fails, closing conn and decrementing _size
+            # Now release the holder — connection breaks, _size decrements
             holder_release.set()
-            with contextlib.suppress(RuntimeError):
+            with contextlib.suppress(DqliteConnectionError):
                 await holder_task
             # At this point _size == 0. The waiter is blocked on queue.get().
             # With the fix, the waiter should wake up and create a new connection.
@@ -384,6 +389,59 @@ async def test_close_then_return_closes_connection(self) -> None:
         mock_conn.close.assert_called()
 
 
+    async def test_user_exception_preserves_healthy_connection(self) -> None:
+        """A user-code exception should not destroy a healthy connection."""
+        pool = ConnectionPool(["localhost:9001"], max_size=2)
+
+        mock_conn = MagicMock()
+        mock_conn.is_connected = True
+        mock_conn.connect = AsyncMock()
+        mock_conn.close = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=mock_conn):
+            await pool.initialize()
+
+        assert pool._size == 1
+
+        # User code raises a non-connection error
+        with pytest.raises(ValueError, match="application error"):
+            async with pool.acquire():
+                raise ValueError("application error")
+
+        # Connection should NOT have been closed — it's healthy
+        mock_conn.close.assert_not_called()
+        # Pool size should be unchanged
+        assert pool._size == 1
+        # Connection should be back in the pool queue
+        assert pool._pool.qsize() == 1
+
+        await pool.close()
+
+    async def test_broken_connection_discarded_on_exception(self) -> None:
+        """A broken connection should be discarded even if user code raised."""
+        pool = ConnectionPool(["localhost:9001"], max_size=2)
+
+        mock_conn = MagicMock()
+        mock_conn.is_connected = True
+        mock_conn.connect = AsyncMock()
+        mock_conn.close = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=mock_conn):
+            await pool.initialize()
+
+        # Simulate a connection that becomes broken during use
+        with pytest.raises(ValueError, match="user error"):
+            async with pool.acquire() as conn:
+                conn.is_connected = False  # Mark as broken (simulates invalidation)
+                raise ValueError("user error")
+
+        # Broken connection SHOULD have been closed and discarded
+        mock_conn.close.assert_called()
+        assert pool._size == 0
+
+        await pool.close()
+
+
 class TestConnectionPoolIntegration:
     """Integration tests requiring mocked connections."""