require_google_service decorator for calendar_tools
This commit is contained in:
Taylor Wilsdon
324
auth/service_decorator.py
Normal file
324
auth/service_decorator.py
Normal file
@@ -0,0 +1,324 @@
|
||||
import inspect
|
||||
import logging
|
||||
from functools import wraps
|
||||
from typing import Dict, List, Optional, Any, Callable, Union
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from auth.google_auth import get_authenticated_google_service, GoogleAuthenticationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Service configuration mapping
|
||||
SERVICE_CONFIGS = {
|
||||
"gmail": {"service": "gmail", "version": "v1"},
|
||||
"drive": {"service": "drive", "version": "v3"},
|
||||
"calendar": {"service": "calendar", "version": "v3"},
|
||||
"docs": {"service": "docs", "version": "v1"},
|
||||
"sheets": {"service": "sheets", "version": "v4"},
|
||||
"chat": {"service": "chat", "version": "v1"}
|
||||
}
|
||||
|
||||
# Scope group definitions for easy reference
|
||||
SCOPE_GROUPS = {
|
||||
# Gmail scopes
|
||||
"gmail_read": "https://www.googleapis.com/auth/gmail.readonly",
|
||||
"gmail_send": "https://www.googleapis.com/auth/gmail.send",
|
||||
"gmail_compose": "https://www.googleapis.com/auth/gmail.compose",
|
||||
"gmail_modify": "https://www.googleapis.com/auth/gmail.modify",
|
||||
"gmail_labels": "https://www.googleapis.com/auth/gmail.labels",
|
||||
|
||||
# Drive scopes
|
||||
"drive_read": "https://www.googleapis.com/auth/drive.readonly",
|
||||
"drive_file": "https://www.googleapis.com/auth/drive.file",
|
||||
|
||||
# Docs scopes
|
||||
"docs_read": "https://www.googleapis.com/auth/documents.readonly",
|
||||
"docs_write": "https://www.googleapis.com/auth/documents",
|
||||
|
||||
# Calendar scopes
|
||||
"calendar_read": "https://www.googleapis.com/auth/calendar.readonly",
|
||||
"calendar_events": "https://www.googleapis.com/auth/calendar.events",
|
||||
|
||||
# Sheets scopes
|
||||
"sheets_read": "https://www.googleapis.com/auth/spreadsheets.readonly",
|
||||
"sheets_write": "https://www.googleapis.com/auth/spreadsheets",
|
||||
|
||||
# Chat scopes
|
||||
"chat_read": "https://www.googleapis.com/auth/chat.messages.readonly",
|
||||
"chat_write": "https://www.googleapis.com/auth/chat.messages",
|
||||
"chat_spaces": "https://www.googleapis.com/auth/chat.spaces.readonly",
|
||||
}
|
||||
|
||||
# Service cache: {cache_key: (service, cached_time, user_email)}
|
||||
_service_cache: Dict[str, tuple[Any, datetime, str]] = {}
|
||||
_cache_ttl = timedelta(minutes=30) # Cache services for 30 minutes
|
||||
|
||||
|
||||
def _get_cache_key(user_email: str, service_name: str, version: str, scopes: List[str]) -> str:
|
||||
"""Generate a cache key for service instances."""
|
||||
sorted_scopes = sorted(scopes)
|
||||
return f"{user_email}:{service_name}:{version}:{':'.join(sorted_scopes)}"
|
||||
|
||||
|
||||
def _is_cache_valid(cached_time: datetime) -> bool:
|
||||
"""Check if cached service is still valid."""
|
||||
return datetime.now() - cached_time < _cache_ttl
|
||||
|
||||
|
||||
def _get_cached_service(cache_key: str) -> Optional[tuple[Any, str]]:
|
||||
"""Retrieve cached service if valid."""
|
||||
if cache_key in _service_cache:
|
||||
service, cached_time, user_email = _service_cache[cache_key]
|
||||
if _is_cache_valid(cached_time):
|
||||
logger.debug(f"Using cached service for key: {cache_key}")
|
||||
return service, user_email
|
||||
else:
|
||||
# Remove expired cache entry
|
||||
del _service_cache[cache_key]
|
||||
logger.debug(f"Removed expired cache entry: {cache_key}")
|
||||
return None
|
||||
|
||||
|
||||
def _cache_service(cache_key: str, service: Any, user_email: str) -> None:
|
||||
"""Cache a service instance."""
|
||||
_service_cache[cache_key] = (service, datetime.now(), user_email)
|
||||
logger.debug(f"Cached service for key: {cache_key}")
|
||||
|
||||
|
||||
def _resolve_scopes(scopes: Union[str, List[str]]) -> List[str]:
|
||||
"""Resolve scope names to actual scope URLs."""
|
||||
if isinstance(scopes, str):
|
||||
if scopes in SCOPE_GROUPS:
|
||||
return [SCOPE_GROUPS[scopes]]
|
||||
else:
|
||||
return [scopes]
|
||||
|
||||
resolved = []
|
||||
for scope in scopes:
|
||||
if scope in SCOPE_GROUPS:
|
||||
resolved.append(SCOPE_GROUPS[scope])
|
||||
else:
|
||||
resolved.append(scope)
|
||||
return resolved
|
||||
|
||||
|
||||
def require_google_service(
|
||||
service_type: str,
|
||||
scopes: Union[str, List[str]],
|
||||
version: Optional[str] = None,
|
||||
cache_enabled: bool = True
|
||||
):
|
||||
"""
|
||||
Decorator that automatically handles Google service authentication and injection.
|
||||
|
||||
Args:
|
||||
service_type: Type of Google service ("gmail", "drive", "calendar", etc.)
|
||||
scopes: Required scopes (can be scope group names or actual URLs)
|
||||
version: Service version (defaults to standard version for service type)
|
||||
cache_enabled: Whether to use service caching (default: True)
|
||||
|
||||
Usage:
|
||||
@require_google_service("gmail", "gmail_read")
|
||||
async def search_messages(service, user_google_email: str, query: str):
|
||||
# service parameter is automatically injected
|
||||
# Original authentication logic is handled automatically
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
# Extract user_google_email from function parameters
|
||||
sig = inspect.signature(func)
|
||||
param_names = list(sig.parameters.keys())
|
||||
|
||||
# Find user_google_email parameter
|
||||
user_google_email = None
|
||||
if 'user_google_email' in kwargs:
|
||||
user_google_email = kwargs['user_google_email']
|
||||
else:
|
||||
# Look for user_google_email in positional args
|
||||
try:
|
||||
user_email_index = param_names.index('user_google_email')
|
||||
if user_email_index < len(args):
|
||||
user_google_email = args[user_email_index]
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if not user_google_email:
|
||||
raise Exception("user_google_email parameter is required but not found")
|
||||
|
||||
# Get service configuration
|
||||
if service_type not in SERVICE_CONFIGS:
|
||||
raise Exception(f"Unknown service type: {service_type}")
|
||||
|
||||
config = SERVICE_CONFIGS[service_type]
|
||||
service_name = config["service"]
|
||||
service_version = version or config["version"]
|
||||
|
||||
# Resolve scopes
|
||||
resolved_scopes = _resolve_scopes(scopes)
|
||||
|
||||
# Check cache first if enabled
|
||||
service = None
|
||||
actual_user_email = user_google_email
|
||||
|
||||
if cache_enabled:
|
||||
cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
|
||||
cached_result = _get_cached_service(cache_key)
|
||||
if cached_result:
|
||||
service, actual_user_email = cached_result
|
||||
|
||||
# If not cached, authenticate
|
||||
if service is None:
|
||||
try:
|
||||
tool_name = func.__name__
|
||||
service, actual_user_email = await get_authenticated_google_service(
|
||||
service_name=service_name,
|
||||
version=service_version,
|
||||
tool_name=tool_name,
|
||||
user_google_email=user_google_email,
|
||||
required_scopes=resolved_scopes,
|
||||
)
|
||||
|
||||
# Cache the service if caching is enabled
|
||||
if cache_enabled:
|
||||
cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
|
||||
_cache_service(cache_key, service, actual_user_email)
|
||||
|
||||
except GoogleAuthenticationError as e:
|
||||
raise Exception(str(e))
|
||||
|
||||
# Inject service as first parameter
|
||||
if 'service' in param_names:
|
||||
kwargs['service'] = service
|
||||
else:
|
||||
# Insert service as first positional argument
|
||||
args = (service,) + args
|
||||
|
||||
# Call the original function
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
return decorator
|
||||
|
||||
|
||||
def require_multiple_services(service_configs: List[Dict[str, Any]]):
|
||||
"""
|
||||
Decorator for functions that need multiple Google services.
|
||||
|
||||
Args:
|
||||
service_configs: List of service configurations, each containing:
|
||||
- service_type: Type of service
|
||||
- scopes: Required scopes
|
||||
- param_name: Name to inject service as (e.g., 'drive_service', 'docs_service')
|
||||
- version: Optional version override
|
||||
|
||||
Usage:
|
||||
@require_multiple_services([
|
||||
{"service_type": "drive", "scopes": "drive_read", "param_name": "drive_service"},
|
||||
{"service_type": "docs", "scopes": "docs_read", "param_name": "docs_service"}
|
||||
])
|
||||
async def get_doc_with_metadata(drive_service, docs_service, user_google_email: str, doc_id: str):
|
||||
# Both services are automatically injected
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
# Extract user_google_email
|
||||
sig = inspect.signature(func)
|
||||
param_names = list(sig.parameters.keys())
|
||||
|
||||
user_google_email = None
|
||||
if 'user_google_email' in kwargs:
|
||||
user_google_email = kwargs['user_google_email']
|
||||
else:
|
||||
try:
|
||||
user_email_index = param_names.index('user_google_email')
|
||||
if user_email_index < len(args):
|
||||
user_google_email = args[user_email_index]
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if not user_google_email:
|
||||
raise Exception("user_google_email parameter is required but not found")
|
||||
|
||||
# Authenticate all services
|
||||
for config in service_configs:
|
||||
service_type = config["service_type"]
|
||||
scopes = config["scopes"]
|
||||
param_name = config["param_name"]
|
||||
version = config.get("version")
|
||||
|
||||
if service_type not in SERVICE_CONFIGS:
|
||||
raise Exception(f"Unknown service type: {service_type}")
|
||||
|
||||
service_config = SERVICE_CONFIGS[service_type]
|
||||
service_name = service_config["service"]
|
||||
service_version = version or service_config["version"]
|
||||
resolved_scopes = _resolve_scopes(scopes)
|
||||
|
||||
try:
|
||||
tool_name = func.__name__
|
||||
service, _ = await get_authenticated_google_service(
|
||||
service_name=service_name,
|
||||
version=service_version,
|
||||
tool_name=tool_name,
|
||||
user_google_email=user_google_email,
|
||||
required_scopes=resolved_scopes,
|
||||
)
|
||||
|
||||
# Inject service with specified parameter name
|
||||
kwargs[param_name] = service
|
||||
|
||||
except GoogleAuthenticationError as e:
|
||||
raise Exception(str(e))
|
||||
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
return decorator
|
||||
|
||||
|
||||
def clear_service_cache(user_email: Optional[str] = None) -> int:
|
||||
"""
|
||||
Clear service cache entries.
|
||||
|
||||
Args:
|
||||
user_email: If provided, only clear cache for this user. If None, clear all.
|
||||
|
||||
Returns:
|
||||
Number of cache entries cleared.
|
||||
"""
|
||||
global _service_cache
|
||||
|
||||
if user_email is None:
|
||||
count = len(_service_cache)
|
||||
_service_cache.clear()
|
||||
logger.info(f"Cleared all {count} service cache entries")
|
||||
return count
|
||||
|
||||
keys_to_remove = [key for key in _service_cache.keys() if key.startswith(f"{user_email}:")]
|
||||
for key in keys_to_remove:
|
||||
del _service_cache[key]
|
||||
|
||||
logger.info(f"Cleared {len(keys_to_remove)} service cache entries for user {user_email}")
|
||||
return len(keys_to_remove)
|
||||
|
||||
|
||||
def get_cache_stats() -> Dict[str, Any]:
|
||||
"""Get service cache statistics."""
|
||||
now = datetime.now()
|
||||
valid_entries = 0
|
||||
expired_entries = 0
|
||||
|
||||
for _, (_, cached_time, _) in _service_cache.items():
|
||||
if _is_cache_valid(cached_time):
|
||||
valid_entries += 1
|
||||
else:
|
||||
expired_entries += 1
|
||||
|
||||
return {
|
||||
"total_entries": len(_service_cache),
|
||||
"valid_entries": valid_entries,
|
||||
"expired_entries": expired_entries,
|
||||
"cache_ttl_minutes": _cache_ttl.total_seconds() / 60
|
||||
}
|
||||
Reference in New Issue
Block a user