Oracle MCP Server

import json import time from pathlib import Path import sys from typing import Dict, List, Set, Optional, Any from ..models import TableInfo, SchemaCache, SchemaManager as SchemaManagerProtocol class SchemaManager(SchemaManagerProtocol): def __init__(self, db_connector: Any, cache_path: Path): self.db_connector = db_connector # Base cache directory self.cache_base_path = cache_path # Actual cache file path will be set after we get the schema name self.cache_path = None self.cache: Optional[SchemaCache] = None self.cache_stats = { 'hits': 0, 'misses': 0, 'last_full_refresh': time.time() } self.object_cache = { 'plsql': {}, 'constraints': {}, 'indexes': {}, 'types': {}, 'related_tables': {} # Added cache for related tables } self.ttl = { 'plsql': 1800, # 30 minutes 'constraints': 3600, # 1 hour 'indexes': 3600, # 1 hour 'types': 3600, # 1 hour 'related_tables': 1800 # 30 minutes - relationships might change more frequently } async def _initialize_cache_path(self) -> None: """Initialize the cache file path using the schema name""" schema_name = await self.db_connector.get_effective_schema() # Create schema-specific cache file name self.cache_path = self.cache_base_path.parent / f"{schema_name.lower()}.json" async def build_schema_index(self) -> Dict[str, TableInfo]: """ Build a basic schema index with just table names. Detailed information will be loaded lazily when needed. """ all_table_names = await self.db_connector.get_all_table_names() print(f"Found {len(all_table_names)} tables in the database", file=sys.stderr) # Initialize empty table info for each table (lazy loading) schema_index = { table_name: TableInfo( table_name=table_name, columns=[], relationships={}, fully_loaded=False ) for table_name in all_table_names } return schema_index async def load_or_build_cache(self, force_rebuild: bool = False) -> SchemaCache: """Load the schema cache from disk or rebuild it if needed""" # Initialize cache path if not already set if self.cache_path is None: await self._initialize_cache_path() if not force_rebuild and self.cache_path.exists(): try: print(f"Opening existing index file for schema: {self.cache_path.stem}...", file=sys.stderr) with open(self.cache_path, 'r') as f: data = json.load(f) print("Loading index in memory...", file=sys.stderr) # Load main schema cache cache = SchemaCache( tables={k: TableInfo(**{**v, 'table_name': k}) for k, v in data['tables'].items()}, last_updated=data['last_updated'], all_table_names=set(data.get('all_table_names', [])) ) # Load additional object caches if they exist if 'object_cache' in data: self.object_cache = data['object_cache'] if 'cache_stats' in data: self.cache_stats = data['cache_stats'] return cache except (json.JSONDecodeError, KeyError) as e: print(f"Error loading cache: {e}", file=sys.stderr) # Fall through to rebuild # Build new cache tables = await self.build_schema_index() all_table_names = set(tables.keys()) print("Loading index in memory...", file=sys.stderr) cache = SchemaCache( tables=tables, last_updated=time.time(), all_table_names=all_table_names ) # Save to disk await self.save_cache(cache) return cache async def save_cache(self, cache: Optional[SchemaCache] = None) -> None: """Save the current cache to disk""" cache_to_save = cache or self.cache if not cache_to_save or not self.cache_path: return print(f"Saving updated index to disk for schema: {self.cache_path.stem}...", file=sys.stderr) self.cache_path.parent.mkdir(parents=True, exist_ok=True) with open(self.cache_path, 'w') as f: json.dump({ 'tables': {k: v.__dict__ for k, v in cache_to_save.tables.items()}, 'last_updated': cache_to_save.last_updated, 'all_table_names': list(cache_to_save.all_table_names), 'object_cache': self.object_cache, 'cache_stats': self.cache_stats }, f, indent=2) print("Index saved!", file=sys.stderr) async def get_schema_info(self, table_name: str) -> Optional[TableInfo]: """Get schema information for a specific table, loading it if necessary""" if not self.cache: self.cache = await self.load_or_build_cache() table_name = table_name.upper() # Check if we know about this table if table_name not in self.cache.all_table_names: return None # Check if we have the table in our cache if table_name not in self.cache.tables: self.cache.tables[table_name] = TableInfo( columns=[], relationships={}, fully_loaded=False ) # If the table isn't fully loaded, load it now if not self.cache.tables[table_name].fully_loaded: print(f"Lazily loading details for table {table_name}...", file=sys.stderr) table_details = await self.db_connector.load_table_details(table_name) if table_details: table_info = TableInfo( table_name=table_name, columns=table_details["columns"], relationships=table_details["relationships"], fully_loaded=True ) self.cache.tables[table_name] = table_info # Save the updated cache to disk await self.save_cache() else: # Table doesn't actually exist, remove it from our cache self.cache.tables.pop(table_name, None) self.cache.all_table_names.discard(table_name) await self.save_cache() return None return self.cache.tables.get(table_name) async def search_tables(self, search_term: str, limit: int = 20) -> List[str]: """ Search for table names matching the search term. First searches in cache, then falls back to database search if needed. """ if not self.cache: self.cache = await self.load_or_build_cache() search_term = search_term.upper() # First try exact/substring matches in cache matching_tables = [ table_name for table_name in self.cache.all_table_names if search_term in table_name ] # If we don't have enough results, search in the database if len(matching_tables) < limit: try: db_results = await self.db_connector.search_in_database(search_term, limit) # Add new tables to our cache new_tables = [table for table in db_results if table not in matching_tables] matching_tables.extend(new_tables) # Update cache with any new tables found if new_tables: self.cache.all_table_names.update(new_tables) await self.save_cache() except Exception as e: print(f"Error during database table search: {str(e)}", file=sys.stderr) # Return the first 'limit' matching tables return matching_tables[:limit] async def search_columns(self, search_term: str, limit: int = 50) -> Dict[str, List[Dict[str, Any]]]: """Search for columns matching the given pattern across all tables""" if not self.cache: await self.initialize() search_term = search_term.upper() result = {} # First check in cached tables to avoid database queries for already loaded tables for table_name, table_info in self.cache.tables.items(): if not table_info.fully_loaded: continue for column in table_info.columns: if search_term in column["name"].upper(): if table_name not in result: result[table_name] = [] result[table_name].append(column) # If we don't have enough results, search in uncached tables if len(result) < limit: uncached_tables = [ t for t in self.cache.all_table_names if t not in self.cache.tables or not self.cache.tables[t].fully_loaded ] if uncached_tables: try: # Search for columns in uncached tables using database connector db_results = await self.db_connector.search_columns_in_database(uncached_tables, search_term) # Merge database results with cache results for table_name, columns in db_results.items(): if table_name not in result: # Only add if not already in cache results result[table_name] = columns # Update cache with the new column information if table_name not in self.cache.tables: self.cache.tables[table_name] = TableInfo( columns=columns, relationships={}, fully_loaded=True ) await self.save_cache() except Exception as e: print(f"Error during database column search: {str(e)}", file=sys.stderr) return dict(list(result.items())[:limit]) async def initialize(self) -> None: """Initialize the database context and build initial cache""" self.cache = await self.load_or_build_cache() if not self.cache: raise RuntimeError("Failed to initialize schema cache") def is_cache_valid(self, cache_type: str, key: str) -> bool: """Check if a cached item is still valid based on TTL""" if (cache_type not in self.object_cache or key not in self.object_cache[cache_type] or 'timestamp' not in self.object_cache[cache_type][key]): return False return (time.time() - self.object_cache[cache_type][key]['timestamp']) < self.ttl[cache_type] def get_cache_stats(self) -> Dict[str, Any]: """Get cache statistics""" return { **self.cache_stats, 'size': { 'tables': len(self.cache.tables) if self.cache else 0, 'plsql': len(self.object_cache['plsql']), 'constraints': len(self.object_cache['constraints']), 'indexes': len(self.object_cache['indexes']), 'types': len(self.object_cache['types']) } } def update_cache(self, cache_type: str, key: str, data: Any) -> None: """Update cache with new data""" self.object_cache[cache_type][key] = { 'data': data, 'timestamp': time.time() }