Skip to content

Commit d9b666e

Browse files
committed
Move key scanning functions from redis_query_engine.py to misc.py
- Move get_all_keys(), scan_keys(), and scan_all_keys() to misc.py - Better organization: key scanning is general Redis operation, not search-specific - Keep redis_query_engine.py focused on vector search and indexing operations
1 parent 057a3f7 commit d9b666e

File tree

2 files changed

+100
-100
lines changed

2 files changed

+100
-100
lines changed

src/tools/misc.py

Lines changed: 100 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,4 +94,103 @@ async def rename(old_key: str, new_key: str) -> Dict[str, Any]:
9494
}
9595

9696
except RedisError as e:
97-
return {"error": str(e)}
97+
return {"error": str(e)}
98+
99+
100+
@mcp.tool()
101+
async def get_all_keys(pattern: str = "*") -> list:
102+
"""
103+
Retrieve all keys matching a pattern from the Redis database using the KEYS command.
104+
105+
Note: The KEYS command is blocking and can impact performance on large databases.
106+
For production use with large datasets, consider using SCAN instead.
107+
108+
Args:
109+
pattern: Pattern to match keys against (default is "*" for all keys).
110+
Common patterns: "user:*", "cache:*", "*:123", etc.
111+
112+
Returns:
113+
A list of keys matching the pattern or an error message.
114+
"""
115+
try:
116+
r = RedisConnectionManager.get_connection()
117+
keys = r.keys(pattern)
118+
# Convert bytes to strings if needed
119+
return [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
120+
except RedisError as e:
121+
return f"Error retrieving keys with pattern '{pattern}': {str(e)}"
122+
123+
124+
@mcp.tool()
125+
async def scan_keys(pattern: str = "*", count: int = 100, cursor: int = 0) -> dict:
126+
"""
127+
Scan keys in the Redis database using the SCAN command (non-blocking, production-safe).
128+
129+
The SCAN command iterates through the keyspace in small chunks, making it safe to use
130+
on large databases without blocking other operations.
131+
132+
Args:
133+
pattern: Pattern to match keys against (default is "*" for all keys).
134+
Common patterns: "user:*", "cache:*", "*:123", etc.
135+
count: Hint for the number of keys to return per iteration (default 100).
136+
Redis may return more or fewer keys than this hint.
137+
cursor: The cursor position to start scanning from (0 to start from beginning).
138+
139+
Returns:
140+
A dictionary containing:
141+
- 'cursor': Next cursor position (0 means scan is complete)
142+
- 'keys': List of keys found in this iteration
143+
- 'total_scanned': Number of keys returned in this batch
144+
Or an error message if something goes wrong.
145+
"""
146+
try:
147+
r = RedisConnectionManager.get_connection()
148+
cursor, keys = r.scan(cursor=cursor, match=pattern, count=count)
149+
150+
# Convert bytes to strings if needed
151+
decoded_keys = [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
152+
153+
return {
154+
'cursor': cursor,
155+
'keys': decoded_keys,
156+
'total_scanned': len(decoded_keys),
157+
'scan_complete': cursor == 0
158+
}
159+
except RedisError as e:
160+
return f"Error scanning keys with pattern '{pattern}': {str(e)}"
161+
162+
163+
@mcp.tool()
164+
async def scan_all_keys(pattern: str = "*", batch_size: int = 100) -> list:
165+
"""
166+
Scan and return ALL keys matching a pattern using multiple SCAN iterations.
167+
168+
This function automatically handles the SCAN cursor iteration to collect all matching keys.
169+
It's safer than KEYS * for large databases but will still collect all results in memory.
170+
171+
Args:
172+
pattern: Pattern to match keys against (default is "*" for all keys).
173+
batch_size: Number of keys to scan per iteration (default 100).
174+
175+
Returns:
176+
A list of all keys matching the pattern or an error message.
177+
"""
178+
try:
179+
r = RedisConnectionManager.get_connection()
180+
all_keys = []
181+
cursor = 0
182+
183+
while True:
184+
cursor, keys = r.scan(cursor=cursor, match=pattern, count=batch_size)
185+
186+
# Convert bytes to strings if needed and add to results
187+
decoded_keys = [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
188+
all_keys.extend(decoded_keys)
189+
190+
# Break when scan is complete (cursor returns to 0)
191+
if cursor == 0:
192+
break
193+
194+
return all_keys
195+
except RedisError as e:
196+
return f"Error scanning all keys with pattern '{pattern}': {str(e)}"

src/tools/redis_query_engine.py

Lines changed: 0 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -136,102 +136,3 @@ async def vector_search_hash(query_vector: list,
136136
return [doc.__dict__ for doc in results.docs]
137137
except RedisError as e:
138138
return f"Error performing vector search on index '{index_name}': {str(e)}"
139-
140-
141-
@mcp.tool()
142-
async def get_all_keys(pattern: str = "*") -> list:
143-
"""
144-
Retrieve all keys matching a pattern from the Redis database using the KEYS command.
145-
146-
Note: The KEYS command is blocking and can impact performance on large databases.
147-
For production use with large datasets, consider using SCAN instead.
148-
149-
Args:
150-
pattern: Pattern to match keys against (default is "*" for all keys).
151-
Common patterns: "user:*", "cache:*", "*:123", etc.
152-
153-
Returns:
154-
A list of keys matching the pattern or an error message.
155-
"""
156-
try:
157-
r = RedisConnectionManager.get_connection()
158-
keys = r.keys(pattern)
159-
# Convert bytes to strings if needed
160-
return [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
161-
except RedisError as e:
162-
return f"Error retrieving keys with pattern '{pattern}': {str(e)}"
163-
164-
165-
@mcp.tool()
166-
async def scan_keys(pattern: str = "*", count: int = 100, cursor: int = 0) -> dict:
167-
"""
168-
Scan keys in the Redis database using the SCAN command (non-blocking, production-safe).
169-
170-
The SCAN command iterates through the keyspace in small chunks, making it safe to use
171-
on large databases without blocking other operations.
172-
173-
Args:
174-
pattern: Pattern to match keys against (default is "*" for all keys).
175-
Common patterns: "user:*", "cache:*", "*:123", etc.
176-
count: Hint for the number of keys to return per iteration (default 100).
177-
Redis may return more or fewer keys than this hint.
178-
cursor: The cursor position to start scanning from (0 to start from beginning).
179-
180-
Returns:
181-
A dictionary containing:
182-
- 'cursor': Next cursor position (0 means scan is complete)
183-
- 'keys': List of keys found in this iteration
184-
- 'total_scanned': Number of keys returned in this batch
185-
Or an error message if something goes wrong.
186-
"""
187-
try:
188-
r = RedisConnectionManager.get_connection()
189-
cursor, keys = r.scan(cursor=cursor, match=pattern, count=count)
190-
191-
# Convert bytes to strings if needed
192-
decoded_keys = [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
193-
194-
return {
195-
'cursor': cursor,
196-
'keys': decoded_keys,
197-
'total_scanned': len(decoded_keys),
198-
'scan_complete': cursor == 0
199-
}
200-
except RedisError as e:
201-
return f"Error scanning keys with pattern '{pattern}': {str(e)}"
202-
203-
204-
@mcp.tool()
205-
async def scan_all_keys(pattern: str = "*", batch_size: int = 100) -> list:
206-
"""
207-
Scan and return ALL keys matching a pattern using multiple SCAN iterations.
208-
209-
This function automatically handles the SCAN cursor iteration to collect all matching keys.
210-
It's safer than KEYS * for large databases but will still collect all results in memory.
211-
212-
Args:
213-
pattern: Pattern to match keys against (default is "*" for all keys).
214-
batch_size: Number of keys to scan per iteration (default 100).
215-
216-
Returns:
217-
A list of all keys matching the pattern or an error message.
218-
"""
219-
try:
220-
r = RedisConnectionManager.get_connection()
221-
all_keys = []
222-
cursor = 0
223-
224-
while True:
225-
cursor, keys = r.scan(cursor=cursor, match=pattern, count=batch_size)
226-
227-
# Convert bytes to strings if needed and add to results
228-
decoded_keys = [key.decode('utf-8') if isinstance(key, bytes) else key for key in keys]
229-
all_keys.extend(decoded_keys)
230-
231-
# Break when scan is complete (cursor returns to 0)
232-
if cursor == 0:
233-
break
234-
235-
return all_keys
236-
except RedisError as e:
237-
return f"Error scanning all keys with pattern '{pattern}': {str(e)}"

0 commit comments

Comments
 (0)