From f5f0f99678a7e9e5f613293f3f0056534718f2af Mon Sep 17 00:00:00 2001 From: saidsurucu Date: Fri, 11 Jul 2025 21:19:16 +0300 Subject: [PATCH] add redis --- asgi_app.py | 52 ++++- mcp_auth_http_simple.py | 166 ++++++++++---- pyproject.toml | 1 + redis_session_store.py | 464 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 636 insertions(+), 47 deletions(-) create mode 100644 redis_session_store.py diff --git a/asgi_app.py b/asgi_app.py index 5f14602..94387b4 100644 --- a/asgi_app.py +++ b/asgi_app.py @@ -299,9 +299,9 @@ async def root(): async def oauth_authorization_server(): """OAuth 2.0 Authorization Server Metadata proxy to Clerk - MCP Auth Toolkit standard location""" return JSONResponse({ - "issuer": CLERK_ISSUER, - "authorization_endpoint": f"{BASE_URL}/auth/login", - "token_endpoint": f"{BASE_URL}/auth/callback", + "issuer": BASE_URL, + "authorization_endpoint": "https://yargimcp.com/mcp-callback", + "token_endpoint": f"{BASE_URL}/token", "jwks_uri": f"{CLERK_ISSUER}/.well-known/jwks.json", "response_types_supported": ["code"], "grant_types_supported": ["authorization_code", "refresh_token"], @@ -312,18 +312,54 @@ async def oauth_authorization_server(): "claims_supported": ["sub", "iss", "aud", "exp", "iat", "email", "name"], "code_challenge_methods_supported": ["S256"], "service_documentation": f"{BASE_URL}/mcp", - "registration_endpoint": f"{BASE_URL}/auth/register", + "registration_endpoint": f"{BASE_URL}/register", "resource_documentation": f"{BASE_URL}/mcp" }) +# Claude AI MCP specific endpoint format +@app.get("/.well-known/oauth-authorization-server/mcp") +async def oauth_authorization_server_mcp_suffix(): + """OAuth 2.0 Authorization Server Metadata - Claude AI MCP specific format""" + return JSONResponse({ + "issuer": BASE_URL, + "authorization_endpoint": "https://yargimcp.com/mcp-callback", + "token_endpoint": f"{BASE_URL}/token", + "jwks_uri": f"{CLERK_ISSUER}/.well-known/jwks.json", + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code", "refresh_token"], + "token_endpoint_auth_methods_supported": ["client_secret_basic", "none"], + "scopes_supported": ["read", "search", "openid", "profile", "email"], + "subject_types_supported": ["public"], + "id_token_signing_alg_values_supported": ["RS256"], + "claims_supported": ["sub", "iss", "aud", "exp", "iat", "email", "name"], + "code_challenge_methods_supported": ["S256"], + "service_documentation": f"{BASE_URL}/mcp", + "registration_endpoint": f"{BASE_URL}/register", + "resource_documentation": f"{BASE_URL}/mcp" + }) + +@app.get("/.well-known/oauth-protected-resource/mcp") +async def oauth_protected_resource_mcp_suffix(): + """OAuth 2.0 Protected Resource Metadata - Claude AI MCP specific format""" + return JSONResponse({ + "resource": BASE_URL, + "authorization_servers": [ + BASE_URL + ], + "scopes_supported": ["read", "search"], + "bearer_methods_supported": ["header"], + "resource_documentation": f"{BASE_URL}/mcp", + "resource_policy_uri": f"{BASE_URL}/privacy" + }) + # Keep root level for compatibility with some MCP clients @app.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_root(): """OAuth 2.0 Authorization Server Metadata proxy to Clerk - root level for compatibility""" return JSONResponse({ - "issuer": CLERK_ISSUER, - "authorization_endpoint": f"{BASE_URL}/auth/login", - "token_endpoint": f"{BASE_URL}/auth/callback", + "issuer": BASE_URL, + "authorization_endpoint": "https://yargimcp.com/mcp-callback", + "token_endpoint": f"{BASE_URL}/token", "jwks_uri": f"{CLERK_ISSUER}/.well-known/jwks.json", "response_types_supported": ["code"], "grant_types_supported": ["authorization_code", "refresh_token"], @@ -334,7 +370,7 @@ async def oauth_authorization_server_root(): "claims_supported": ["sub", "iss", "aud", "exp", "iat", "email", "name"], "code_challenge_methods_supported": ["S256"], "service_documentation": f"{BASE_URL}/mcp", - "registration_endpoint": f"{BASE_URL}/auth/register", + "registration_endpoint": f"{BASE_URL}/register", "resource_documentation": f"{BASE_URL}/mcp" }) diff --git a/mcp_auth_http_simple.py b/mcp_auth_http_simple.py index ee8eafa..4311ef9 100644 --- a/mcp_auth_http_simple.py +++ b/mcp_auth_http_simple.py @@ -1,5 +1,6 @@ """ Simplified MCP OAuth HTTP adapter - only Clerk JWT based authentication +Uses Redis for authorization code storage to support multi-machine deployment """ import os @@ -10,6 +11,9 @@ from urllib.parse import urlencode, quote from fastapi import APIRouter, Request, Query, HTTPException from fastapi.responses import RedirectResponse, JSONResponse +# Import Redis session store +from redis_session_store import get_redis_store + # Try to import Clerk SDK try: from clerk_backend_api import Clerk @@ -26,14 +30,50 @@ router = APIRouter() BASE_URL = os.getenv("BASE_URL", "https://api.yargimcp.com") CLERK_DOMAIN = os.getenv("CLERK_DOMAIN", "accounts.yargimcp.com") +# Initialize Redis store +redis_store = None + +def get_redis_session_store(): + """Get Redis store instance with lazy initialization.""" + global redis_store + if redis_store is None: + try: + import concurrent.futures + import functools + + # Use thread pool with timeout to prevent hanging + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(get_redis_store) + try: + # 5 second timeout for Redis initialization + redis_store = future.result(timeout=5.0) + if redis_store: + logger.info("Redis session store initialized for OAuth handler") + else: + logger.warning("Redis store initialization returned None") + except concurrent.futures.TimeoutError: + logger.error("Redis initialization timed out after 5 seconds") + redis_store = None + future.cancel() # Try to cancel the hanging operation + + except Exception as e: + logger.error(f"Failed to initialize Redis store: {e}") + redis_store = None + + if redis_store is None: + # Fall back to in-memory storage with warning + logger.warning("Falling back to in-memory storage - multi-machine deployment will not work") + + return redis_store + @router.get("/.well-known/oauth-authorization-server") async def get_oauth_metadata(): """OAuth 2.0 Authorization Server Metadata (RFC 8414)""" return JSONResponse({ "issuer": BASE_URL, - "authorization_endpoint": f"{BASE_URL}/auth/login", - "token_endpoint": f"{BASE_URL}/auth/callback", - "registration_endpoint": f"{BASE_URL}/auth/register", + "authorization_endpoint": "https://yargimcp.com/mcp-callback", + "token_endpoint": f"{BASE_URL}/token", + "registration_endpoint": f"{BASE_URL}/register", "response_types_supported": ["code"], "grant_types_supported": ["authorization_code"], "code_challenge_methods_supported": ["S256"], @@ -165,24 +205,38 @@ async def oauth_callback( # Generate authorization code auth_code = f"clerk_auth_{os.urandom(16).hex()}" - # Store code with JWT token mapping (in production, use proper storage) - # For now, we'll use a simple in-memory storage + # Prepare code data import time code_data = { "user_id": user_id, "session_id": session_id, "real_jwt_token": real_jwt_token, "user_authenticated": user_authenticated, - "created_at": time.time(), - "expires_at": time.time() + 300 # 5 minutes expiry + "client_id": client_id, + "redirect_uri": redirect_uri, + "scope": scope or "read search" } - # Store in module-level dict (in production, use Redis or database) - if not hasattr(oauth_callback, '_code_storage'): - oauth_callback._code_storage = {} - oauth_callback._code_storage[auth_code] = code_data - - logger.info(f"Stored authorization code with real JWT token") + # Try to store in Redis, fall back to in-memory if Redis unavailable + store = get_redis_session_store() + if store: + # Store in Redis with automatic expiration + success = store.set_oauth_code(auth_code, code_data) + if success: + logger.info(f"Stored authorization code {auth_code[:10]}... in Redis with real JWT token") + else: + logger.error(f"Failed to store authorization code in Redis, falling back to in-memory") + # Fall back to in-memory storage + if not hasattr(oauth_callback, '_code_storage'): + oauth_callback._code_storage = {} + oauth_callback._code_storage[auth_code] = code_data + else: + # Fall back to in-memory storage + logger.warning("Redis not available, using in-memory storage") + if not hasattr(oauth_callback, '_code_storage'): + oauth_callback._code_storage = {} + oauth_callback._code_storage[auth_code] = code_data + logger.info(f"Stored authorization code in memory (fallback)") # Redirect back to client with authorization code redirect_params = { @@ -272,15 +326,25 @@ async def oauth_callback_post(request: Request): content={"error": "invalid_grant", "error_description": "Invalid authorization code"} ) - # TODO: In production, validate code against stored session - # For now, we'll return a placeholder response - - # Retrieve stored JWT token using authorization code + # Retrieve stored JWT token using authorization code from Redis or in-memory fallback stored_code_data = None - # Get stored code data from authorization flow - if hasattr(oauth_callback, '_code_storage'): + # Try to get from Redis first, then fall back to in-memory + store = get_redis_session_store() + if store: + stored_code_data = store.get_oauth_code(code, delete_after_use=True) + if stored_code_data: + logger.info(f"Retrieved authorization code {code[:10]}... from Redis") + else: + logger.warning(f"Authorization code {code[:10]}... not found in Redis") + + # Fall back to in-memory storage if Redis unavailable or code not found + if not stored_code_data and hasattr(oauth_callback, '_code_storage'): stored_code_data = oauth_callback._code_storage.get(code) + if stored_code_data: + # Clean up in-memory storage + oauth_callback._code_storage.pop(code, None) + logger.info(f"Retrieved authorization code {code[:10]}... from in-memory storage") if not stored_code_data: logger.error(f"No stored data found for authorization code: {code}") @@ -289,13 +353,11 @@ async def oauth_callback_post(request: Request): content={"error": "invalid_grant", "error_description": "Authorization code not found or expired"} ) - # Check if code is expired + # Note: Redis TTL handles expiration automatically, but check for manual expiration for in-memory fallback import time - if time.time() > stored_code_data.get("expires_at", 0): + expires_at = stored_code_data.get("expires_at", 0) + if expires_at and time.time() > expires_at: logger.error(f"Authorization code expired: {code}") - # Clean up expired code - if hasattr(oauth_callback, '_code_storage'): - oauth_callback._code_storage.pop(code, None) return JSONResponse( status_code=400, content={"error": "invalid_grant", "error_description": "Authorization code expired"} @@ -306,7 +368,7 @@ async def oauth_callback_post(request: Request): if real_jwt_token: logger.info("Returning real Clerk JWT token") - # Clean up used code + # Note: Code already deleted from Redis, clean up in-memory fallback if used if hasattr(oauth_callback, '_code_storage'): oauth_callback._code_storage.pop(code, None) @@ -334,6 +396,26 @@ async def oauth_callback_post(request: Request): content={"error": "server_error", "error_description": str(e)} ) +@router.post("/register") +async def register_client(request: Request): + """Dynamic Client Registration (RFC 7591)""" + + data = await request.json() + logger.info(f"Client registration request: {data}") + + # Simple dynamic registration - accept any client + client_id = f"mcp-client-{os.urandom(8).hex()}" + + return JSONResponse({ + "client_id": client_id, + "client_secret": None, # Public client + "redirect_uris": data.get("redirect_uris", []), + "grant_types": ["authorization_code"], + "response_types": ["code"], + "client_name": data.get("client_name", "MCP Client"), + "token_endpoint_auth_method": "none" + }) + @router.post("/token") async def token_endpoint(request: Request): """OAuth 2.1 Token Endpoint - exchanges code for Clerk JWT""" @@ -369,17 +451,25 @@ async def token_endpoint(request: Request): content={"error": "invalid_grant", "error_description": "Invalid authorization code"} ) - # In a real implementation, you would: - # 1. Validate the code against stored session - # 2. Extract user info from the session - # 3. Return the actual Clerk JWT token - - # Retrieve stored JWT token using authorization code + # Retrieve stored JWT token using authorization code from Redis or in-memory fallback stored_code_data = None - # Get stored code data from authorization flow - if hasattr(oauth_callback, '_code_storage'): + # Try to get from Redis first, then fall back to in-memory + store = get_redis_session_store() + if store: + stored_code_data = store.get_oauth_code(code, delete_after_use=True) + if stored_code_data: + logger.info(f"Retrieved authorization code {code[:10]}... from Redis (/token endpoint)") + else: + logger.warning(f"Authorization code {code[:10]}... not found in Redis (/token endpoint)") + + # Fall back to in-memory storage if Redis unavailable or code not found + if not stored_code_data and hasattr(oauth_callback, '_code_storage'): stored_code_data = oauth_callback._code_storage.get(code) + if stored_code_data: + # Clean up in-memory storage + oauth_callback._code_storage.pop(code, None) + logger.info(f"Retrieved authorization code {code[:10]}... from in-memory storage (/token endpoint)") if not stored_code_data: logger.error(f"No stored data found for authorization code: {code}") @@ -388,13 +478,11 @@ async def token_endpoint(request: Request): content={"error": "invalid_grant", "error_description": "Authorization code not found or expired"} ) - # Check if code is expired + # Note: Redis TTL handles expiration automatically, but check for manual expiration for in-memory fallback import time - if time.time() > stored_code_data.get("expires_at", 0): + expires_at = stored_code_data.get("expires_at", 0) + if expires_at and time.time() > expires_at: logger.error(f"Authorization code expired: {code}") - # Clean up expired code - if hasattr(oauth_callback, '_code_storage'): - oauth_callback._code_storage.pop(code, None) return JSONResponse( status_code=400, content={"error": "invalid_grant", "error_description": "Authorization code expired"} @@ -405,7 +493,7 @@ async def token_endpoint(request: Request): if real_jwt_token: logger.info("Returning real Clerk JWT token from /token endpoint") - # Clean up used code + # Note: Code already deleted from Redis, clean up in-memory fallback if used if hasattr(oauth_callback, '_code_storage'): oauth_callback._code_storage.pop(code, None) diff --git a/pyproject.toml b/pyproject.toml index 2a717b7..fc7ae00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ production = [ saas = [ "clerk-backend-api>=3.0.0", "stripe>=9.1.0", + "upstash-redis>=1.1.0", ] [project.scripts] diff --git a/redis_session_store.py b/redis_session_store.py new file mode 100644 index 0000000..0eb1142 --- /dev/null +++ b/redis_session_store.py @@ -0,0 +1,464 @@ +""" +Redis Session Store for OAuth Authorization Codes and User Sessions + +This module provides Redis-based storage for OAuth authorization codes and user sessions, +enabling multi-machine deployment support by replacing in-memory storage. + +Uses Upstash Redis via REST API for serverless-friendly operation. +""" + +import os +import json +import time +import logging +from typing import Optional, Dict, Any, Union +from datetime import datetime, timedelta + +logger = logging.getLogger(__name__) + +try: + from upstash_redis import Redis + UPSTASH_AVAILABLE = True +except ImportError: + UPSTASH_AVAILABLE = False + Redis = None + +# Use standard Python exceptions for Redis connection errors +import socket +from requests.exceptions import ConnectionError as RequestsConnectionError, Timeout as RequestsTimeout + +class RedisSessionStore: + """ + Redis-based session store for OAuth flows and user sessions. + + Uses Upstash Redis REST API for connection-free operation suitable for + multi-instance deployments on platforms like Fly.io. + """ + + def __init__(self): + """Initialize Redis connection using environment variables.""" + if not UPSTASH_AVAILABLE: + raise ImportError("upstash-redis package is required. Install with: pip install upstash-redis") + + # Initialize Upstash Redis client from environment with optimized connection settings + try: + # Get Upstash Redis configuration + redis_url = os.getenv("UPSTASH_REDIS_REST_URL") + redis_token = os.getenv("UPSTASH_REDIS_REST_TOKEN") + + if not redis_url or not redis_token: + raise ValueError("UPSTASH_REDIS_REST_URL and UPSTASH_REDIS_REST_TOKEN must be set") + + logger.info(f"Connecting to Upstash Redis at {redis_url[:30]}...") + + # Initialize with explicit configuration for better SSL handling + self.redis = Redis( + url=redis_url, + token=redis_token + ) + + logger.info("Upstash Redis client created") + + # Skip connection test during initialization to prevent server hang + # Connection will be tested during first actual operation + logger.info("Redis client initialized - connection will be tested on first use") + + except Exception as e: + logger.error(f"Failed to initialize Upstash Redis: {e}") + raise + + # TTL values (in seconds) + self.oauth_code_ttl = int(os.getenv("OAUTH_CODE_TTL", "600")) # 10 minutes + self.session_ttl = int(os.getenv("SESSION_TTL", "3600")) # 1 hour + + def _serialize_data(self, data: Dict[str, Any]) -> Dict[str, str]: + """Convert data to Redis-compatible string format.""" + serialized = {} + for key, value in data.items(): + if isinstance(value, (dict, list)): + serialized[key] = json.dumps(value) + elif isinstance(value, (int, float)): + serialized[key] = str(value) + elif isinstance(value, bool): + serialized[key] = "true" if value else "false" + else: + serialized[key] = str(value) + return serialized + + def _deserialize_data(self, data: Dict[str, str]) -> Dict[str, Any]: + """Convert Redis string data back to original types.""" + if not data: + return {} + + deserialized = {} + for key, value in data.items(): + if not isinstance(value, str): + deserialized[key] = value + continue + + # Try to deserialize JSON + if value.startswith(('[', '{')): + try: + deserialized[key] = json.loads(value) + continue + except json.JSONDecodeError: + pass + + # Try to convert numbers + if value.isdigit(): + deserialized[key] = int(value) + continue + + if value.replace('.', '').isdigit(): + try: + deserialized[key] = float(value) + continue + except ValueError: + pass + + # Handle booleans + if value in ("true", "false"): + deserialized[key] = value == "true" + continue + + # Keep as string + deserialized[key] = value + + return deserialized + + # OAuth Authorization Code Methods + + def set_oauth_code(self, code: str, data: Dict[str, Any]) -> bool: + """ + Store OAuth authorization code with automatic expiration. + + Args: + code: Authorization code string + data: Code data including user_id, client_id, etc. + + Returns: + True if stored successfully, False otherwise + """ + try: + key = f"oauth:code:{code}" + + # Add timestamp for debugging + data_with_timestamp = data.copy() + data_with_timestamp.update({ + "created_at": time.time(), + "expires_at": time.time() + self.oauth_code_ttl + }) + + # Serialize and store - Upstash Redis doesn't support mapping parameter + serialized_data = self._serialize_data(data_with_timestamp) + + # Use individual hset calls for each field with retry logic + max_retries = 3 + for attempt in range(max_retries): + try: + # Clear any existing data first + self.redis.delete(key) + + # Set all fields in a pipeline-like manner + for field, value in serialized_data.items(): + self.redis.hset(key, field, value) + + # Set expiration + self.redis.expire(key, self.oauth_code_ttl) + + logger.info(f"Stored OAuth code {code[:10]}... with TTL {self.oauth_code_ttl}s (attempt {attempt + 1})") + return True + + except (RequestsConnectionError, RequestsTimeout, OSError, socket.error) as e: + logger.warning(f"Redis connection error on attempt {attempt + 1}: {e}") + if attempt == max_retries - 1: + raise # Re-raise on final attempt + time.sleep(0.5 * (attempt + 1)) # Exponential backoff + + except Exception as e: + logger.error(f"Failed to store OAuth code {code[:10]}... after {max_retries} attempts: {e}") + return False + + def get_oauth_code(self, code: str, delete_after_use: bool = True) -> Optional[Dict[str, Any]]: + """ + Retrieve OAuth authorization code data. + + Args: + code: Authorization code string + delete_after_use: If True, delete the code after retrieval (one-time use) + + Returns: + Code data dictionary or None if not found/expired + """ + max_retries = 3 + for attempt in range(max_retries): + try: + key = f"oauth:code:{code}" + + # Get all hash fields with retry + data = self.redis.hgetall(key) + + if not data: + logger.warning(f"OAuth code {code[:10]}... not found or expired (attempt {attempt + 1})") + return None + + # Deserialize data + deserialized_data = self._deserialize_data(data) + + # Check manual expiration (in case Redis TTL failed) + expires_at = deserialized_data.get("expires_at", 0) + if expires_at and time.time() > expires_at: + logger.warning(f"OAuth code {code[:10]}... manually expired") + try: + self.redis.delete(key) + except Exception as del_error: + logger.warning(f"Failed to delete expired code: {del_error}") + return None + + # Delete after use for security (one-time use) + if delete_after_use: + try: + self.redis.delete(key) + logger.info(f"Retrieved and deleted OAuth code {code[:10]}... (attempt {attempt + 1})") + except Exception as del_error: + logger.warning(f"Failed to delete code after use: {del_error}") + # Continue anyway since we got the data + else: + logger.info(f"Retrieved OAuth code {code[:10]}... (not deleted, attempt {attempt + 1})") + + return deserialized_data + + except (RequestsConnectionError, RequestsTimeout, OSError, socket.error) as e: + logger.warning(f"Redis connection error on retrieval attempt {attempt + 1}: {e}") + if attempt == max_retries - 1: + logger.error(f"Failed to retrieve OAuth code {code[:10]}... after {max_retries} attempts: {e}") + return None + time.sleep(0.5 * (attempt + 1)) # Exponential backoff + + except Exception as e: + logger.error(f"Failed to retrieve OAuth code {code[:10]}... on attempt {attempt + 1}: {e}") + if attempt == max_retries - 1: + return None + time.sleep(0.5 * (attempt + 1)) + + return None + + # User Session Methods + + def set_session(self, session_id: str, user_data: Dict[str, Any]) -> bool: + """ + Store user session data with sliding expiration. + + Args: + session_id: Unique session identifier + user_data: User session data (user_id, email, scopes, etc.) + + Returns: + True if stored successfully, False otherwise + """ + try: + key = f"session:{session_id}" + + # Add session metadata + session_data = user_data.copy() + session_data.update({ + "session_id": session_id, + "created_at": time.time(), + "last_accessed": time.time() + }) + + # Serialize and store - Upstash Redis doesn't support mapping parameter + serialized_data = self._serialize_data(session_data) + + # Use individual hset calls for each field (Upstash compatibility) + for field, value in serialized_data.items(): + self.redis.hset(key, field, value) + self.redis.expire(key, self.session_ttl) + + logger.info(f"Stored session {session_id[:10]}... with TTL {self.session_ttl}s") + return True + + except Exception as e: + logger.error(f"Failed to store session {session_id[:10]}...: {e}") + return False + + def get_session(self, session_id: str, refresh_ttl: bool = True) -> Optional[Dict[str, Any]]: + """ + Retrieve user session data. + + Args: + session_id: Session identifier + refresh_ttl: If True, extend session TTL on access + + Returns: + Session data dictionary or None if not found/expired + """ + try: + key = f"session:{session_id}" + + # Get session data + data = self.redis.hgetall(key) + + if not data: + logger.warning(f"Session {session_id[:10]}... not found or expired") + return None + + # Deserialize data + session_data = self._deserialize_data(data) + + # Update last accessed time and refresh TTL + if refresh_ttl: + session_data["last_accessed"] = time.time() + self.redis.hset(key, "last_accessed", str(time.time())) + self.redis.expire(key, self.session_ttl) + logger.debug(f"Refreshed session {session_id[:10]}... TTL") + + return session_data + + except Exception as e: + logger.error(f"Failed to retrieve session {session_id[:10]}...: {e}") + return None + + def delete_session(self, session_id: str) -> bool: + """ + Delete user session (logout). + + Args: + session_id: Session identifier + + Returns: + True if deleted successfully, False otherwise + """ + try: + key = f"session:{session_id}" + result = self.redis.delete(key) + + if result: + logger.info(f"Deleted session {session_id[:10]}...") + return True + else: + logger.warning(f"Session {session_id[:10]}... not found for deletion") + return False + + except Exception as e: + logger.error(f"Failed to delete session {session_id[:10]}...: {e}") + return False + + # Health Check Methods + + def health_check(self) -> Dict[str, Any]: + """ + Perform Redis health check. + + Returns: + Health status dictionary + """ + try: + # Test basic operations + test_key = f"health:check:{int(time.time())}" + test_value = {"timestamp": time.time(), "test": True} + + # Test set - Use individual hset calls for Upstash compatibility + serialized_test = self._serialize_data(test_value) + for field, value in serialized_test.items(): + self.redis.hset(test_key, field, value) + + # Test get + retrieved = self.redis.hgetall(test_key) + + # Test delete + self.redis.delete(test_key) + + return { + "status": "healthy", + "redis_connected": True, + "operations_working": bool(retrieved), + "timestamp": datetime.utcnow().isoformat() + } + + except Exception as e: + logger.error(f"Redis health check failed: {e}") + return { + "status": "unhealthy", + "redis_connected": False, + "error": str(e), + "timestamp": datetime.utcnow().isoformat() + } + + def get_stats(self) -> Dict[str, Any]: + """ + Get Redis usage statistics. + + Returns: + Statistics dictionary + """ + try: + # Get basic info (not all Upstash plans support INFO command) + stats = { + "oauth_codes_pattern": "oauth:code:*", + "sessions_pattern": "session:*", + "timestamp": datetime.utcnow().isoformat() + } + + try: + # Try to get counts (may fail on some Upstash plans) + oauth_keys = self.redis.keys("oauth:code:*") + session_keys = self.redis.keys("session:*") + + stats.update({ + "active_oauth_codes": len(oauth_keys) if oauth_keys else 0, + "active_sessions": len(session_keys) if session_keys else 0 + }) + except Exception as e: + logger.warning(f"Could not get detailed stats: {e}") + stats["warning"] = "Detailed stats not available on this Redis plan" + + return stats + + except Exception as e: + logger.error(f"Failed to get Redis stats: {e}") + return {"error": str(e), "timestamp": datetime.utcnow().isoformat()} + +# Global instance for easy importing +redis_store = None + +def get_redis_store() -> Optional[RedisSessionStore]: + """ + Get global Redis store instance (singleton pattern). + + Returns: + RedisSessionStore instance or None if initialization fails + """ + global redis_store + + if redis_store is None: + try: + logger.info("Initializing Redis store...") + redis_store = RedisSessionStore() + logger.info("Redis store initialized successfully") + except Exception as e: + logger.error(f"Failed to initialize Redis store: {e}") + redis_store = None + + return redis_store + +def init_redis_store() -> RedisSessionStore: + """ + Initialize Redis store and perform health check. + + Returns: + RedisSessionStore instance + + Raises: + Exception if Redis is not available or unhealthy + """ + store = get_redis_store() + + # Perform health check + health = store.health_check() + + if health["status"] != "healthy": + raise Exception(f"Redis health check failed: {health}") + + logger.info("Redis session store initialized and healthy") + return store \ No newline at end of file