YDB MCP

import asyncio import logging import re from typing import Optional, Tuple from urllib.parse import urlparse import ydb logger = logging.getLogger(__name__) class YDBConnection: """Manages YDB connection with async support.""" def __init__(self, connection_string: str, database: str | None = None): """Initialize YDB connection. Args: connection_string: YDB connection string database: Optional database path. If not provided, will be extracted from connection_string if present """ self.connection_string = connection_string self.driver: Optional[ydb.Driver] = None self.session_pool: Optional[ydb.aio.QuerySessionPool] = None self._database = database self.last_error: str | None = None def _parse_endpoint_and_database(self) -> Tuple[str, str]: """Parse endpoint and database from connection string. Returns: Tuple of (endpoint, database) Raises: RuntimeError: If no database is specified either in connection string or explicitly """ # Parse the URL connection_string = self.connection_string if not connection_string.startswith(("grpc://", "grpcs://")): # If no scheme, assume grpc:// and parse as host:port if "/" in connection_string: host_port, path = connection_string.split("/", 1) connection_string = f"grpc://{host_port}/{path}" else: connection_string = f"grpc://{connection_string}" parsed = urlparse(connection_string) # Extract endpoint (scheme + netloc) endpoint = f"{parsed.scheme}://{parsed.netloc}" # Extract database path database = self._database if not database: if parsed.path: database = parsed.path # Remove query parameters if present if "?" in database: database = database.split("?")[0] # Ensure database starts with / if database and not database.startswith("/"): database = f"/{database}" # Raise error if no database specified if not database: raise RuntimeError("Database not specified in connection string or explicitly") return endpoint, database async def connect(self) -> Tuple[ydb.Driver, ydb.aio.QuerySessionPool]: """Connect to YDB and setup session pool asynchronously. Returns: Tuple of (driver, session_pool) Raises: RuntimeError: If connection fails """ try: endpoint, database = self._parse_endpoint_and_database() logger.info(f"Connecting to YDB endpoint: {endpoint}, database: {database}") # Create driver with direct parameters instead of config self.driver = ydb.aio.Driver( endpoint=endpoint, database=database, credentials=ydb.credentials.AnonymousCredentials(), ) # Wait for driver to be ready with timeout try: await asyncio.wait_for(self.driver.wait(), timeout=10.0) except asyncio.TimeoutError: self.last_error = "Connection timeout" raise RuntimeError("YDB driver connection timeout after 10 seconds") # Check if we connected successfully if not self.driver.discovery_debug_details().startswith("Resolved endpoints"): debug_details = self.driver.discovery_debug_details() self.last_error = f"Driver not ready: {debug_details}" raise RuntimeError(f"YDB driver failed to connect: {debug_details}") logger.info("Connected to YDB successfully") # Create session pool self.session_pool = ydb.aio.QuerySessionPool(self.driver) return self.driver, self.session_pool except Exception as e: self.last_error = str(e) logger.error(f"Failed to connect to YDB: {e}") raise RuntimeError(f"Failed to connect to YDB: {e}") async def close(self) -> None: """Close YDB connection.""" logger.info("Closing YDB connection") if self.session_pool: await asyncio.get_event_loop().run_in_executor(None, self.session_pool.stop) self.session_pool = None if self.driver: await asyncio.get_event_loop().run_in_executor(None, self.driver.stop) self.driver = None logger.info("YDB connection closed") def _extract_database_path(self, connection_string: Optional[str] = None) -> str: """Extract database path from connection string. Args: connection_string: YDB connection string, or None to use the instance's connection string Returns: Database path """ # Use instance connection string if none provided if connection_string is None: connection_string = self.connection_string # Handle connection string with query parameters if "?" in connection_string: connection_string = connection_string.split("?")[0] # Extract path using regex match = re.match(r"^(?:[^:]+://[^/]+)?(/.*)?$", connection_string) return match.group(1) if match and match.group(1) else "/"