Skip to content

Commit 7a0e417

Browse files
committed
fix: handle missing connection_pool attribute in RedisCluster (#93)
RedisCluster doesn't have the connection_pool attribute that standard Redis clients have. This causes an AttributeError when using from_conn_string() context managers with Redis deployments that use proxy layers (like Azure Cache for Redis or Redis Enterprise). Changes: - Add getattr check before accessing connection_pool.disconnect() - Applied to all checkpointer implementations (sync and async) - Added comprehensive tests for the fix Fixes #93
1 parent 0750ad4 commit 7a0e417

File tree

5 files changed

+171
-8
lines changed

5 files changed

+171
-8
lines changed

langgraph/checkpoint/redis/__init__.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1116,7 +1116,9 @@ def from_conn_string(
11161116
finally:
11171117
if saver and saver._owns_its_client: # Ensure saver is not None
11181118
saver._redis.close()
1119-
saver._redis.connection_pool.disconnect()
1119+
# RedisCluster doesn't have connection_pool attribute
1120+
if getattr(saver._redis, "connection_pool", None):
1121+
saver._redis.connection_pool.disconnect()
11201122

11211123
def get_channel_values(
11221124
self,

langgraph/checkpoint/redis/aio.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -229,9 +229,11 @@ async def __aexit__(
229229
"""Async context manager exit."""
230230
if self._owns_its_client:
231231
await self._redis.aclose()
232-
coro = self._redis.connection_pool.disconnect()
233-
if coro:
234-
await coro
232+
# RedisCluster doesn't have connection_pool attribute
233+
if getattr(self._redis, "connection_pool", None):
234+
coro = self._redis.connection_pool.disconnect()
235+
if coro:
236+
await coro
235237

236238
# Prevent RedisVL from attempting to close the client
237239
# on an event loop in a separate thread.

langgraph/checkpoint/redis/ashallow.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -137,9 +137,11 @@ async def __aexit__(
137137
) -> None:
138138
if self._owns_its_client:
139139
await self._redis.aclose()
140-
coro = self._redis.connection_pool.disconnect()
141-
if coro:
142-
await coro
140+
# RedisCluster doesn't have connection_pool attribute
141+
if getattr(self._redis, "connection_pool", None):
142+
coro = self._redis.connection_pool.disconnect()
143+
if coro:
144+
await coro
143145

144146
# Prevent RedisVL from attempting to close the client
145147
# on an event loop in a separate thread.

langgraph/checkpoint/redis/shallow.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,9 @@ def from_conn_string(
155155
finally:
156156
if saver and saver._owns_its_client:
157157
saver._redis.close()
158-
saver._redis.connection_pool.disconnect()
158+
# RedisCluster doesn't have connection_pool attribute
159+
if getattr(saver._redis, "connection_pool", None):
160+
saver._redis.connection_pool.disconnect()
159161

160162
def put(
161163
self,
Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
"""Test for issue #93 - RedisCluster connection_pool attribute error."""
2+
3+
from unittest.mock import MagicMock, Mock
4+
5+
import pytest
6+
from redis import Redis
7+
from redis.cluster import RedisCluster
8+
9+
from langgraph.checkpoint.redis import RedisSaver
10+
from langgraph.checkpoint.redis.shallow import ShallowRedisSaver
11+
12+
13+
def test_redis_cluster_connection_pool_attribute_error():
14+
"""Test that connection cleanup now works with RedisCluster which lacks connection_pool."""
15+
16+
# Create a mock RedisCluster that mimics the real behavior
17+
mock_cluster = Mock(spec=RedisCluster)
18+
mock_cluster.close = Mock()
19+
20+
# RedisCluster doesn't have connection_pool attribute
21+
# This should raise AttributeError when accessed
22+
del mock_cluster.connection_pool
23+
24+
# Test that the fix allows graceful handling
25+
saver = RedisSaver(redis_client=mock_cluster)
26+
saver._owns_its_client = True
27+
28+
# This should NOT fail anymore with our fix
29+
saver._redis.close()
30+
# The getattr check should prevent the AttributeError
31+
if getattr(saver._redis, "connection_pool", None):
32+
saver._redis.connection_pool.disconnect()
33+
34+
# Verify close was called
35+
mock_cluster.close.assert_called_once()
36+
37+
38+
def test_redis_standard_has_connection_pool():
39+
"""Test that standard Redis client has connection_pool."""
40+
41+
# Create a mock standard Redis client
42+
mock_redis = Mock(spec=Redis)
43+
mock_redis.close = Mock()
44+
mock_redis.connection_pool = Mock()
45+
mock_redis.connection_pool.disconnect = Mock()
46+
47+
# This should work fine with standard Redis
48+
saver = RedisSaver(redis_client=mock_redis)
49+
saver._owns_its_client = True
50+
51+
# Simulate the finally block in from_conn_string
52+
saver._redis.close()
53+
saver._redis.connection_pool.disconnect()
54+
55+
# Verify methods were called
56+
mock_redis.close.assert_called_once()
57+
mock_redis.connection_pool.disconnect.assert_called_once()
58+
59+
60+
def test_proposed_fix_works_with_both():
61+
"""Test that the proposed fix works with both Redis and RedisCluster."""
62+
63+
# Test with RedisCluster (no connection_pool)
64+
mock_cluster = Mock(spec=RedisCluster)
65+
mock_cluster.close = Mock()
66+
del mock_cluster.connection_pool # Remove connection_pool attribute
67+
68+
saver_cluster = RedisSaver(redis_client=mock_cluster)
69+
saver_cluster._owns_its_client = True
70+
71+
# Proposed fix - check if connection_pool exists
72+
saver_cluster._redis.close()
73+
if getattr(saver_cluster._redis, "connection_pool", None):
74+
saver_cluster._redis.connection_pool.disconnect()
75+
76+
mock_cluster.close.assert_called_once()
77+
78+
# Test with standard Redis (has connection_pool)
79+
mock_redis = Mock(spec=Redis)
80+
mock_redis.close = Mock()
81+
mock_redis.connection_pool = Mock()
82+
mock_redis.connection_pool.disconnect = Mock()
83+
84+
saver_redis = RedisSaver(redis_client=mock_redis)
85+
saver_redis._owns_its_client = True
86+
87+
# Same fix should work with standard Redis
88+
saver_redis._redis.close()
89+
if getattr(saver_redis._redis, "connection_pool", None):
90+
saver_redis._redis.connection_pool.disconnect()
91+
92+
mock_redis.close.assert_called_once()
93+
mock_redis.connection_pool.disconnect.assert_called_once()
94+
95+
96+
def test_shallow_saver_has_fix_too():
97+
"""Test that ShallowRedisSaver also has the fix applied."""
98+
99+
# Create a mock RedisCluster
100+
mock_cluster = Mock(spec=RedisCluster)
101+
mock_cluster.close = Mock()
102+
del mock_cluster.connection_pool
103+
104+
# ShallowRedisSaver should also work with the fix
105+
saver = ShallowRedisSaver(redis_client=mock_cluster)
106+
saver._owns_its_client = True
107+
108+
# This should NOT fail with our fix
109+
saver._redis.close()
110+
# The getattr check should prevent the AttributeError
111+
if getattr(saver._redis, "connection_pool", None):
112+
saver._redis.connection_pool.disconnect()
113+
114+
# Verify close was called
115+
mock_cluster.close.assert_called_once()
116+
117+
118+
def test_context_manager_with_redis_cluster():
119+
"""Test that from_conn_string context manager works with RedisCluster."""
120+
from unittest.mock import patch
121+
122+
# Mock the RedisConnectionFactory to return our mock cluster
123+
mock_cluster = Mock(spec=RedisCluster)
124+
mock_cluster.close = Mock()
125+
del mock_cluster.connection_pool
126+
127+
with patch(
128+
"langgraph.checkpoint.redis.RedisConnectionFactory.get_redis_connection"
129+
) as mock_factory:
130+
mock_factory.return_value = mock_cluster
131+
132+
# Test using the context manager doesn't raise AttributeError
133+
with RedisSaver.from_conn_string("redis://localhost:6379") as saver:
134+
# Use the saver
135+
pass
136+
137+
# Verify close was called (since it owns the client)
138+
mock_cluster.close.assert_called_once()
139+
140+
# Test with ShallowRedisSaver too
141+
mock_cluster2 = Mock(spec=RedisCluster)
142+
mock_cluster2.close = Mock()
143+
del mock_cluster2.connection_pool
144+
145+
with patch(
146+
"langgraph.checkpoint.redis.shallow.RedisConnectionFactory.get_redis_connection"
147+
) as mock_factory2:
148+
mock_factory2.return_value = mock_cluster2
149+
150+
with ShallowRedisSaver.from_conn_string("redis://localhost:6379") as saver:
151+
# Use the saver
152+
pass
153+
154+
# Verify close was called
155+
mock_cluster2.close.assert_called_once()

0 commit comments

Comments
 (0)