import json
from typing import Union, Mapping, List, TYPE_CHECKING
from redis.exceptions import RedisError
from src.common.connection import RedisConnectionManager
from src.common.server import mcp
# Define JsonType for type checking to match redis-py definition
# Use object as runtime type to avoid issubclass() issues with Any in Python 3.10
if TYPE_CHECKING:
JsonType = Union[
str, int, float, bool, None, Mapping[str, "JsonType"], List["JsonType"]
]
else:
# Use object at runtime to avoid MCP framework issubclass() issues
JsonType = object
@mcp.tool()
async def json_set(
name: str, path: str, value: JsonType, expire_seconds: int = None
) -> str:
"""Set a JSON value in Redis at a given path with an optional expiration time.
Args:
name: The Redis key where the JSON document is stored.
path: The JSON path where the value should be set.
value: The JSON value to store.
expire_seconds: Optional; time in seconds after which the key should expire.
Returns:
A success message or an error message.
"""
try:
r = RedisConnectionManager.get_connection()
r.json().set(name, path, value)
if expire_seconds is not None:
r.expire(name, expire_seconds)
return f"JSON value set at path '{path}' in '{name}'." + (
f" Expires in {expire_seconds} seconds." if expire_seconds else ""
)
except RedisError as e:
return f"Error setting JSON value at path '{path}' in '{name}': {str(e)}"
@mcp.tool()
async def json_get(name: str, path: str = "$") -> str:
"""Retrieve a JSON value from Redis at a given path.
Args:
name: The Redis key where the JSON document is stored.
path: The JSON path to retrieve (default: root '$').
Returns:
The retrieved JSON value or an error message.
"""
try:
r = RedisConnectionManager.get_connection()
value = r.json().get(name, path)
if value is not None:
# Convert the value to JSON string for consistent return type
return json.dumps(value, ensure_ascii=False, indent=2)
else:
return f"No data found at path '{path}' in '{name}'."
except RedisError as e:
return f"Error retrieving JSON value at path '{path}' in '{name}': {str(e)}"
@mcp.tool()
async def json_del(name: str, path: str = "$") -> str:
"""Delete a JSON value from Redis at a given path.
Args:
name: The Redis key where the JSON document is stored.
path: The JSON path to delete (default: root '$').
Returns:
A success message or an error message.
"""
try:
r = RedisConnectionManager.get_connection()
deleted = r.json().delete(name, path)
return (
f"Deleted JSON value at path '{path}' in '{name}'."
if deleted
else f"No JSON value found at path '{path}' in '{name}'."
)
except RedisError as e:
return f"Error deleting JSON value at path '{path}' in '{name}': {str(e)}"
