surreal-mcp

from contextlib import asynccontextmanager from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union from loguru import logger from surrealdb import RecordID # type: ignore from .connection_pool import ( close_connection_pool, get_pool_stats, override_db_connection, pooled_db_connection, ) T = TypeVar("T", Dict[str, Any], List[Dict[str, Any]]) def parse_record_ids(obj: Any) -> Any: """Recursively parse and convert RecordIDs into strings.""" if isinstance(obj, dict): return {k: parse_record_ids(v) for k, v in obj.items()} elif isinstance(obj, list): return [parse_record_ids(item) for item in obj] elif isinstance(obj, RecordID): return str(obj) return obj def ensure_record_id(value: Union[str, RecordID]) -> RecordID: """Ensure a value is a RecordID.""" if isinstance(value, RecordID): return value return RecordID.parse(value) @asynccontextmanager async def db_connection( namespace: Optional[str] = None, database: Optional[str] = None, ): """ Database connection context manager using connection pool or override connection. If namespace and database are provided, creates a one-off connection with those values. Otherwise, uses the pooled connection with default namespace/database from env vars. Args: namespace: Optional namespace override database: Optional database override Yields: An authenticated AsyncSurreal connection """ if namespace is not None and database is not None: # Use override connection for custom namespace/database async with override_db_connection(namespace, database) as connection: yield connection else: # Use pooled connection for default namespace/database async with pooled_db_connection() as connection: yield connection async def repo_query( query_str: str, vars: Optional[Dict[str, Any]] = None, namespace: Optional[str] = None, database: Optional[str] = None, ) -> List[Dict[str, Any]]: """Execute a SurrealQL query and return the results. Args: query_str: The SurrealQL query to execute vars: Optional variables for the query namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The query results as a list of dictionaries """ async with db_connection(namespace, database) as connection: try: result = parse_record_ids(await connection.query(query_str, vars)) if isinstance(result, str): raise RuntimeError(result) return result except Exception as e: logger.error(f"Query: {query_str[:200]} vars: {vars}") logger.exception(e) raise async def repo_create( table: str, data: Dict[str, Any], namespace: Optional[str] = None, database: Optional[str] = None, ) -> Dict[str, Any]: """Create a new record in the specified table. Args: table: The table to create the record in data: The record data namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The created record """ # Remove 'id' attribute if it exists in data data.pop("id", None) data["created"] = datetime.now(timezone.utc) data["updated"] = datetime.now(timezone.utc) try: async with db_connection(namespace, database) as connection: return parse_record_ids(await connection.insert(table, data)) except Exception as e: logger.exception(e) raise RuntimeError("Failed to create record") async def repo_relate( source: str, relationship: str, target: str, data: Optional[Dict[str, Any]] = None, namespace: Optional[str] = None, database: Optional[str] = None, ) -> List[Dict[str, Any]]: """Create a relationship between two records with optional data. Args: source: The source record ID relationship: The relationship/edge name target: The target record ID data: Optional data to store on the relationship namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The created relationship records """ if data is None: data = {} query = f"RELATE {source}->{relationship}->{target} CONTENT $data;" return await repo_query( query, {"data": data}, namespace=namespace, database=database, ) async def repo_upsert( table: str, id: Optional[str], data: Dict[str, Any], add_timestamp: bool = False, namespace: Optional[str] = None, database: Optional[str] = None, ) -> List[Dict[str, Any]]: """Create or update a record in the specified table. Args: table: The table name id: Optional record ID (if provided, upserts that specific record) data: The record data to upsert add_timestamp: Whether to add/update the 'updated' timestamp namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The upserted record(s) """ data.pop("id", None) if add_timestamp: data["updated"] = datetime.now(timezone.utc) query = f"UPSERT {id if id else table} MERGE $data;" return await repo_query(query, {"data": data}, namespace=namespace, database=database) async def repo_update( table: str, id: str, data: Dict[str, Any], namespace: Optional[str] = None, database: Optional[str] = None, ) -> List[Dict[str, Any]]: """Update an existing record by table and id. Args: table: The table name id: The record ID data: The data to update namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The updated record(s) """ # If id already contains the table name, use it as is try: if isinstance(id, RecordID) or (":" in id and id.startswith(f"{table}:")): record_id = id else: record_id = f"{table}:{id}" data["updated"] = datetime.now(timezone.utc) query = f"UPDATE {record_id} MERGE $data;" result = await repo_query(query, {"data": data}, namespace=namespace, database=database) return parse_record_ids(result) except Exception as e: raise RuntimeError(f"Failed to update record: {str(e)}") async def repo_get_news_by_jota_id(jota_id: str) -> Dict[str, Any]: try: results = await repo_query( "SELECT * omit embedding FROM news where jota_id=$jota_id", {"jota_id": jota_id}, ) return parse_record_ids(results) except Exception as e: logger.exception(e) raise RuntimeError(f"Failed to fetch record: {str(e)}") async def repo_delete( record_id: Union[str, RecordID], namespace: Optional[str] = None, database: Optional[str] = None, ): """Delete a record by record id. Args: record_id: The record ID to delete namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The deletion result """ try: async with db_connection(namespace, database) as connection: return await connection.delete(record_id) except Exception as e: logger.exception(e) raise RuntimeError(f"Failed to delete record: {str(e)}") async def repo_insert( table: str, data: List[Dict[str, Any]], ignore_duplicates: bool = False, namespace: Optional[str] = None, database: Optional[str] = None, ) -> List[Dict[str, Any]]: """Insert multiple records into a table. Args: table: The table to insert into data: List of records to insert ignore_duplicates: Whether to ignore duplicate key errors namespace: Optional namespace override (uses env var if not provided) database: Optional database override (uses env var if not provided) Returns: The inserted records """ try: async with db_connection(namespace, database) as connection: return parse_record_ids(await connection.insert(table, data)) except Exception as e: if ignore_duplicates and "already contains" in str(e): return [] logger.exception(e) raise RuntimeError("Failed to create record") # Connection pool utilities async def get_connection_pool_stats() -> Dict[str, Any]: """Get connection pool statistics for monitoring.""" return await get_pool_stats() async def close_database_pool(): """Close the database connection pool. Call this on application shutdown.""" await close_connection_pool()