import os import re from contextlib import asynccontextmanager from datetime import datetime, timezone from typing import Any, Dict, List, Optional, TypeVar, Union from loguru import logger from surrealdb import AsyncSurreal, RecordID # type: ignore from surrealdb.data.types.table import Table # type: ignore T = TypeVar("T", Dict[str, Any], List[Dict[str, Any]]) # Bare SurrealDB table/relation identifier: no ':', whitespace, or query # syntax. Used to validate the parts of RELATE/UPSERT/UPDATE that name a # table or edge-relation and therefore can't be bound as a query parameter # (SurrealQL only allows binding record/table *values*, not identifiers in # that position). _IDENTIFIER_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") def _ensure_safe_identifier(value: str, kind: str) -> str: """Validate a table/relationship name before it is interpolated into a query.""" if not isinstance(value, str) or not _IDENTIFIER_RE.match(value): raise ValueError(f"Invalid {kind} name: {value!r}") return value def _get_env_or_default(name: str, default: str) -> str: value = os.getenv(name) return value if value else default def get_database_url() -> str: """Get database URL with backward compatibility""" surreal_url = os.getenv("SURREAL_URL") if surreal_url: return surreal_url # Fallback to old format - WebSocket URL format address = os.getenv("SURREAL_ADDRESS", "localhost") port = os.getenv("SURREAL_PORT", "8000") return f"ws://{address}/rpc:{port}" def get_database_password() -> str: """Get password with backward compatibility""" return os.getenv("SURREAL_PASSWORD") or os.getenv("SURREAL_PASS") or "root" def get_database_namespace() -> str: """Get configured SurrealDB namespace.""" return _get_env_or_default("SURREAL_NAMESPACE", "open_notebook") def get_database_name() -> str: """Get configured SurrealDB database name.""" return _get_env_or_default("SURREAL_DATABASE", "open_notebook") 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(): db = AsyncSurreal(get_database_url()) await db.signin( { "username": os.environ.get("SURREAL_USER"), "password": get_database_password(), } ) await db.use(get_database_namespace(), get_database_name()) try: yield db finally: await db.close() async def repo_query( query_str: str, vars: Optional[Dict[str, Any]] = None ) -> List[Dict[str, Any]]: """Execute a SurrealQL query and return the results""" async with db_connection() as connection: try: result = parse_record_ids(await connection.query(query_str, vars)) if isinstance(result, str): raise RuntimeError(result) return result except RuntimeError as e: # RuntimeError is raised for retriable transaction conflicts - log at debug to avoid noise logger.debug(str(e)) raise except Exception as e: logger.exception(e) raise async def repo_create(table: str, data: Dict[str, Any]) -> Dict[str, Any]: """Create a new record in the specified table""" # 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() as connection: result = parse_record_ids(await connection.insert(table, data)) # SurrealDB may return a string error message instead of the expected record if isinstance(result, str): raise RuntimeError(result) return result except RuntimeError as e: logger.error(str(e)) raise except Exception as e: logger.exception(e) raise RuntimeError("Failed to create record") async def repo_relate( source: Union[str, RecordID], relationship: str, target: Union[str, RecordID], data: Optional[Dict[str, Any]] = None, ) -> List[Dict[str, Any]]: """Create a relationship between two records with optional data""" if data is None: data = {} # relationship is an edge-table name, not a record value, so it can't be # bound as a query parameter; validate it against an identifier allowlist # instead of trusting the caller. source/target are always bound. _ensure_safe_identifier(relationship, "relationship") query = f"RELATE $source->{relationship}->$target CONTENT $data;" # logger.debug(f"Relate query: {query}") return await repo_query( query, { "source": ensure_record_id(source), "target": ensure_record_id(target), "data": data, }, ) async def repo_upsert( table: str, id: Optional[str], data: Dict[str, Any], add_timestamp: bool = False ) -> List[Dict[str, Any]]: """Create or update a record in the specified table""" data.pop("id", None) if add_timestamp: data["updated"] = datetime.now(timezone.utc) _ensure_safe_identifier(table, "table") target: Union[RecordID, Table] = ensure_record_id(id) if id else Table(table) query = "UPSERT $target MERGE $data;" return await repo_query(query, {"target": target, "data": data}) async def repo_update( table: str, id: Union[str, RecordID], data: Dict[str, Any] ) -> List[Dict[str, Any]]: """Update an existing record by table and id""" # If id already contains the table name, use it as is try: _ensure_safe_identifier(table, "table") if isinstance(id, RecordID): record_id: RecordID = id elif ":" in id and id.startswith(f"{table}:"): record_id = ensure_record_id(id) else: record_id = RecordID(table, id) data.pop("id", None) if "created" in data and isinstance(data["created"], str): data["created"] = datetime.fromisoformat(data["created"]) data["updated"] = datetime.now(timezone.utc) query = "UPDATE $target MERGE $data;" # logger.debug(f"Update query: {query}") result = await repo_query(query, {"target": record_id, "data": data}) # if isinstance(result, list): # return [_return_data(item) for item in result] return parse_record_ids(result) except Exception as e: raise RuntimeError(f"Failed to update record: {str(e)}") async def repo_delete(record_id: Union[str, RecordID]): """Delete a record by record id""" try: async with db_connection() as connection: return await connection.delete(ensure_record_id(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 ) -> List[Dict[str, Any]]: """Create a new record in the specified table""" try: async with db_connection() as connection: result = parse_record_ids(await connection.insert(table, data)) # SurrealDB may return a string error message instead of the expected records if isinstance(result, str): raise RuntimeError(result) return result except RuntimeError as e: if ignore_duplicates and "already contains" in str(e): return [] # Log transaction conflicts at debug level (they are expected during concurrent operations) error_str = str(e).lower() if "transaction" in error_str or "conflict" in error_str: logger.debug(str(e)) else: logger.error(str(e)) raise except Exception as e: if ignore_duplicates and "already contains" in str(e): return [] logger.exception(e) raise RuntimeError("Failed to create record")