BigQuery Validator

"""Caching module for MCP BigQuery server.""" import hashlib import json import time from typing import Any, Optional, Union from .config import get_config from .constants import CACHE_KEY_PREFIX from .logging_config import get_logger logger = get_logger(__name__) class CacheEntry: """A cache entry with value and metadata.""" def __init__(self, value: Any, ttl: int): self.value = value self.created_at = time.time() self.ttl = ttl self.access_count = 0 self.last_accessed = self.created_at def is_expired(self) -> bool: """Check if the cache entry has expired.""" return time.time() - self.created_at > self.ttl def access(self) -> Any: """Access the cache entry and update metadata.""" self.access_count += 1 self.last_accessed = time.time() return self.value class Cache: """In-memory cache implementation.""" def __init__(self, max_size: int = 1000, default_ttl: int = 300): """ Initialize cache. Args: max_size: Maximum number of entries in cache default_ttl: Default time-to-live in seconds """ self.max_size = max_size self.default_ttl = default_ttl self._cache: dict[str, CacheEntry] = {} self._stats = {"hits": 0, "misses": 0, "evictions": 0, "expirations": 0} def _make_key(self, prefix: str, *args: Any, **kwargs: Any) -> str: """ Create a cache key from prefix and arguments. Args: prefix: Key prefix *args: Positional arguments **kwargs: Keyword arguments Returns: Cache key string """ key_data = {"prefix": prefix, "args": args, "kwargs": sorted(kwargs.items())} key_str = json.dumps(key_data, sort_keys=True, default=str) return hashlib.md5(key_str.encode()).hexdigest() def get(self, key: str) -> Optional[Any]: """ Get value from cache. Args: key: Cache key Returns: Cached value or None if not found/expired """ if key in self._cache: entry = self._cache[key] if entry.is_expired(): del self._cache[key] self._stats["expirations"] += 1 self._stats["misses"] += 1 logger.debug(f"Cache miss (expired): {key}") return None self._stats["hits"] += 1 logger.debug(f"Cache hit: {key}") return entry.access() self._stats["misses"] += 1 logger.debug(f"Cache miss: {key}") return None def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None: """ Set value in cache. Args: key: Cache key value: Value to cache ttl: Time-to-live in seconds (uses default if None) """ if ttl is None: ttl = self.default_ttl # Check if we need to evict entries if len(self._cache) >= self.max_size: self._evict_lru() self._cache[key] = CacheEntry(value, ttl) logger.debug(f"Cache set: {key} (TTL: {ttl}s)") def delete(self, key: str) -> bool: """ Delete value from cache. Args: key: Cache key Returns: True if key was deleted, False if not found """ if key in self._cache: del self._cache[key] logger.debug(f"Cache delete: {key}") return True return False def clear(self) -> None: """Clear all cache entries.""" self._cache.clear() logger.info("Cache cleared") def _evict_lru(self) -> None: """Evict least recently used entry.""" if not self._cache: return # Find LRU entry lru_key = min(self._cache.keys(), key=lambda k: self._cache[k].last_accessed) del self._cache[lru_key] self._stats["evictions"] += 1 logger.debug(f"Cache eviction (LRU): {lru_key}") def cleanup_expired(self) -> int: """ Remove all expired entries. Returns: Number of entries removed """ expired_keys = [k for k, v in self._cache.items() if v.is_expired()] for key in expired_keys: del self._cache[key] if expired_keys: logger.info(f"Cleaned up {len(expired_keys)} expired cache entries") return len(expired_keys) def get_stats(self) -> dict[str, Any]: """ Get cache statistics. Returns: Dictionary with cache statistics """ total = self._stats["hits"] + self._stats["misses"] hit_rate = self._stats["hits"] / total if total > 0 else 0 return { "size": len(self._cache), "max_size": self.max_size, "hits": self._stats["hits"], "misses": self._stats["misses"], "hit_rate": hit_rate, "evictions": self._stats["evictions"], "expirations": self._stats["expirations"], } class BigQueryClientCache: """Cache for BigQuery client instances.""" def __init__(self): self._clients: dict[str, Any] = {} self._client_locks: dict[str, Any] = {} def get_client(self, project_id: Optional[str] = None, location: Optional[str] = None) -> Any: """ Get or create a BigQuery client. Args: project_id: GCP project ID location: BigQuery location Returns: BigQuery client instance """ from google.cloud import bigquery from .bigquery_client import get_bigquery_client key = f"{project_id or 'default'}:{location or 'default'}" if key not in self._clients: logger.info(f"Creating new BigQuery client for {key}") self._clients[key] = get_bigquery_client(project_id, location) else: logger.debug(f"Reusing BigQuery client for {key}") return self._clients[key] def clear(self) -> None: """Clear all cached clients.""" self._clients.clear() logger.info("BigQuery client cache cleared") # Global cache instances _query_cache: Optional[Cache] = None _schema_cache: Optional[Cache] = None _client_cache: Optional[BigQueryClientCache] = None def get_query_cache() -> Cache: """Get the global query cache instance.""" global _query_cache if _query_cache is None: config = get_config() _query_cache = Cache(max_size=1000, default_ttl=config.cache_ttl) return _query_cache def get_schema_cache() -> Cache: """Get the global schema cache instance.""" global _schema_cache if _schema_cache is None: config = get_config() _schema_cache = Cache( max_size=500, default_ttl=config.cache_ttl * 2 # Schema changes less frequently ) return _schema_cache def get_client_cache() -> BigQueryClientCache: """Get the global BigQuery client cache instance.""" global _client_cache if _client_cache is None: _client_cache = BigQueryClientCache() return _client_cache def clear_all_caches() -> None: """Clear all cache instances.""" if _query_cache: _query_cache.clear() if _schema_cache: _schema_cache.clear() if _client_cache: _client_cache.clear() logger.info("All caches cleared") # Cache decorators def cache_query_result(ttl: Optional[int] = None): """ Decorator to cache query results. Args: ttl: Time-to-live in seconds """ import functools def decorator(func): @functools.wraps(func) async def wrapper(*args, **kwargs): config = get_config() if not config.cache_enabled: return await func(*args, **kwargs) cache = get_query_cache() cache_key = cache._make_key(CACHE_KEY_PREFIX["query"], *args, **kwargs) # Check cache result = cache.get(cache_key) if result is not None: return result # Execute function result = await func(*args, **kwargs) # Cache result cache.set(cache_key, result, ttl) return result return wrapper return decorator def cache_schema_info(ttl: Optional[int] = None): """ Decorator to cache schema information. Args: ttl: Time-to-live in seconds """ import functools def decorator(func): @functools.wraps(func) async def wrapper(*args, **kwargs): config = get_config() if not config.cache_enabled: return await func(*args, **kwargs) cache = get_schema_cache() cache_key = cache._make_key(CACHE_KEY_PREFIX["schema"], *args, **kwargs) # Check cache result = cache.get(cache_key) if result is not None: return result # Execute function result = await func(*args, **kwargs) # Cache result cache.set(cache_key, result, ttl) return result return wrapper return decorator