Use separate connection pool for result storage (#276)

The py-key-value library's RedisStore requires `decode_responses=True`, 
but Docket's internal Redis usage expects bytes for performance. This
creates two connection pools - one for each mode.

Also fixes a bug where the `result_storage` parameter was accepted but
silently ignored, and works around a py-key-value issue where usernames
in Redis URLs weren't being passed to the client
(github.com/strawgate/py-key-value/issues/254).

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: chrisguidry

src/docket/_redis.py

@@ -34,24 +34,29 @@ def get_memory_server(url: str) -> "FakeServer | None":
     return _memory_servers.get(url)
 
 
-async def connection_pool_from_url(url: str) -> ConnectionPool:
+async def connection_pool_from_url(
+    url: str, decode_responses: bool = False
+) -> ConnectionPool:
     """Create a Redis connection pool from a URL.
 
     Handles both real Redis (redis://) and in-memory fakeredis (memory://).
     This is the only place in the codebase that imports fakeredis.
 
     Args:
         url: Redis URL (redis://...) or memory:// for in-memory backend
+        decode_responses: If True, decode Redis responses from bytes to strings
 
     Returns:
         A ConnectionPool ready for use with Redis clients
     """
     if url.startswith("memory://"):
-        return await _memory_connection_pool(url)
-    return ConnectionPool.from_url(url)
+        return await _memory_connection_pool(url, decode_responses)
+    return ConnectionPool.from_url(url, decode_responses=decode_responses)
 
 
-async def _memory_connection_pool(url: str) -> ConnectionPool:
+async def _memory_connection_pool(
+    url: str, decode_responses: bool = False
+) -> ConnectionPool:
     """Create a connection pool for a memory:// URL using fakeredis."""
     global _memory_servers
 
@@ -63,16 +68,28 @@ async def _memory_connection_pool(url: str) -> ConnectionPool:
     # Fast path: server already exists
     server = _memory_servers.get(url)
     if server is not None:
-        return ConnectionPool(connection_class=FakeConnection, server=server)
+        return ConnectionPool(
+            connection_class=FakeConnection,
+            server=server,
+            decode_responses=decode_responses,
+        )
 
     async with _memory_servers_lock:
         server = _memory_servers.get(url)
         if server is not None:  # pragma: no cover
-            return ConnectionPool(connection_class=FakeConnection, server=server)
+            return ConnectionPool(
+                connection_class=FakeConnection,
+                server=server,
+                decode_responses=decode_responses,
+            )
 
         server = FakeServer()
         _memory_servers[url] = server
-        return ConnectionPool(connection_class=FakeConnection, server=server)
+        return ConnectionPool(
+            connection_class=FakeConnection,
+            server=server,
+            decode_responses=decode_responses,
+        )
 
 
 # ------------------------------------------------------------------------------

src/docket/docket.py

@@ -46,7 +46,6 @@
 )
 from key_value.aio.protocols.key_value import AsyncKeyValue
 from key_value.aio.stores.redis import RedisStore
-from key_value.aio.stores.memory import MemoryStore
 
 from .instrumentation import (
     TASKS_ADDED,
@@ -192,14 +191,7 @@ def __init__(
         self.execution_ttl = execution_ttl
         self.enable_internal_instrumentation = enable_internal_instrumentation
         self._cancel_task_script = None
-
-        self.result_storage: AsyncKeyValue
-        if url.startswith("memory://"):
-            self.result_storage = MemoryStore()
-        else:
-            self.result_storage = RedisStore(
-                url=url, default_collection=self.results_collection
-            )
+        self._user_result_storage = result_storage
 
         from .tasks import standard_tasks
 
@@ -240,6 +232,23 @@ async def __aenter__(self) -> Self:
         # Connect the strike list to Redis and start monitoring
         await self.strike_list.connect()
 
+        # Initialize result storage
+        # We use a separate connection pool for result storage because RedisStore
+        # requires decode_responses=True while Docket's internal Redis usage
+        # expects bytes (decode_responses=False). We also pass a client to
+        # RedisStore to work around a py-key-value bug that ignores username
+        # in URLs (github.com/strawgate/py-key-value/issues/254).
+        if self._user_result_storage is not None:
+            self.result_storage: AsyncKeyValue = self._user_result_storage
+        else:
+            self._result_storage_pool = await connection_pool_from_url(
+                self.url, decode_responses=True
+            )
+            result_client = Redis(connection_pool=self._result_storage_pool)
+            self.result_storage = RedisStore(
+                client=result_client, default_collection=self.results_collection
+            )
+
         if isinstance(self.result_storage, BaseContextManagerStore):
             await self.result_storage.__aenter__()
         else:
@@ -255,6 +264,11 @@ async def __aexit__(
         if isinstance(self.result_storage, BaseContextManagerStore):
             await self.result_storage.__aexit__(exc_type, exc_value, traceback)
 
+        # Close the result storage pool if we created it
+        if hasattr(self, "_result_storage_pool"):
+            await asyncio.shield(self._result_storage_pool.disconnect())
+            del self._result_storage_pool
+
         # Close the strike list (stops monitoring and disconnects)
         await self.strike_list.close()
         del self.strike_list

tests/test_results.py

@@ -498,3 +498,45 @@ async def waits_forever() -> int:
 
     event.set()
     await worker_task
+
+
+async def test_result_storage_uses_provided_or_default(redis_url: str):
+    """Test that result_storage uses your store if provided, RedisStore if not."""
+    from unittest.mock import AsyncMock, MagicMock
+    from urllib.parse import urlparse
+
+    from key_value.aio.protocols.key_value import AsyncKeyValue
+    from key_value.aio.stores.redis import RedisStore
+
+    # If you give us one, it's yours
+    custom_storage = MagicMock(spec=AsyncKeyValue)
+    custom_storage.setup = AsyncMock()
+    async with Docket(
+        name="test-custom-storage",
+        url=redis_url,
+        result_storage=custom_storage,
+    ) as docket:
+        assert docket.result_storage is custom_storage
+
+    # If you don't, it's a RedisStore pointing at the same server
+    async with Docket(name="test-default-storage", url=redis_url) as docket:
+        assert isinstance(docket.result_storage, RedisStore)
+
+        # Verify it's connected to the same Redis
+        result_client = docket.result_storage._client  # type: ignore[attr-defined]
+        pool_kwargs: dict[str, Any] = result_client.connection_pool.connection_kwargs  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+        if redis_url.startswith("memory://"):  # pragma: no cover
+            # For memory://, just verify it has a server (fakeredis)
+            assert "server" in pool_kwargs
+        else:
+            # For real Redis, verify host/port/db match
+            parsed = urlparse(redis_url)
+            assert pool_kwargs.get("host") == (parsed.hostname or "localhost")
+            assert pool_kwargs.get("port") == (parsed.port or 6379)
+            expected_db = (
+                int(parsed.path.lstrip("/"))
+                if parsed.path and parsed.path != "/"
+                else 0
+            )
+            assert pool_kwargs.get("db") == expected_db