feat(db): wire schema registry and multi-db routing

Author: Sakeeb91

app/routes.py

@@ -473,8 +473,17 @@ async def get_schema(
     try:
         from db.schema import SchemaIntrospector
 
-        db_manager = await get_database()
-        introspector = SchemaIntrospector(db_manager.engine)
+        settings = get_settings()
+        if settings.multi_database.enabled:
+            from db.registry import get_database_registry
+
+            registry = await get_database_registry()
+            introspector = SchemaIntrospector(
+                engine_provider=registry.get_engine,
+            )
+        else:
+            db_manager = await get_database()
+            introspector = SchemaIntrospector(db_manager.engine)
         schema = await introspector.get_schema(database_id)
 
         return SchemaResponse(
@@ -502,7 +511,7 @@ async def register_schema(
     request: Request,
     response: Response,
     schema_request: SchemaRequest,
-) -> dict[str, str]:
+) -> dict[str, Any]:
     """
     Register a new database schema.
 
@@ -530,8 +539,96 @@ async def register_schema(
         dialect=schema_request.dialect,
     )
 
-    # TODO: Implement actual registration
-    return {"status": "registered", "database_id": schema_request.database_id}
+    settings = get_settings()
+
+    try:
+        from app.cache import cache_schema
+        from db.schema import SchemaIntrospector
+
+        registry = None
+
+        if settings.multi_database.enabled:
+            from db.dialects import SQLDialect
+            from db.registry import DatabaseConfig, get_database_registry
+
+            registry = await get_database_registry()
+            if registry.database_count >= settings.multi_database.max_databases:
+                raise ValidationException(
+                    message=(
+                        "Maximum database limit reached "
+                        f"({settings.multi_database.max_databases})"
+                    ),
+                    validation_errors=[
+                        {"field": "database_id", "error": "Registry is full"}
+                    ],
+                )
+
+            dialect: SQLDialect | None = None
+            if schema_request.dialect:
+                try:
+                    dialect = SQLDialect(schema_request.dialect.lower())
+                except ValueError:
+                    raise ValidationException(
+                        message=f"Invalid dialect: {schema_request.dialect}",
+                        validation_errors=[
+                            {
+                                "field": "dialect",
+                                "error": (
+                                    "Supported: "
+                                    f"{SQLDialect.supported_dialects()}"
+                                ),
+                            }
+                        ],
+                    ) from None
+
+            config = DatabaseConfig(
+                database_id=schema_request.database_id,
+                connection_string=schema_request.connection_string,
+                dialect=dialect,
+                pool_size=settings.multi_database.default_pool_size,
+                max_overflow=settings.multi_database.default_max_overflow,
+                pool_timeout=settings.multi_database.default_pool_timeout,
+                read_only=not settings.multi_database.allow_mutations,
+            )
+
+            registered = await registry.register_database(
+                config,
+                test_connection=settings.multi_database.require_connection_test,
+            )
+            introspector = SchemaIntrospector(registered.engine)
+        else:
+            db_manager = await get_database()
+            introspector = SchemaIntrospector(db_manager.engine)
+
+        schema = await introspector.get_schema(schema_request.database_id)
+        serialized = introspector.serialize_for_prompt(schema)
+        if registry is not None:
+            registry.update_schema(schema_request.database_id, schema)
+
+        if settings.cache.enabled:
+            await cache_schema(
+                schema_request.database_id,
+                {
+                    "schema_info": schema.to_dict(),
+                    "serialized_schema": serialized,
+                    "table_names": [table.name for table in schema.tables],
+                    "dialect": schema.dialect,
+                },
+            )
+
+        return {
+            "status": "registered",
+            "database_id": schema_request.database_id,
+            "table_count": len(schema.tables),
+            "dialect": schema.dialect,
+        }
+    except Exception as e:
+        logger.error(
+            "register_schema_error",
+            database_id=schema_request.database_id,
+            error=str(e),
+        )
+        raise
 
 
 # =============================================================================

app/text2sql_engine.py

@@ -18,6 +18,7 @@
 import json
 import re
 import time
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
@@ -44,6 +45,7 @@
 from app.resilience import CircuitBreaker, CircuitBreakerConfig, compute_backoff_seconds
 from db.connection import DatabaseManager
 from db.executor import QueryResult
+from db.registry import get_database_registry
 from db.schema import SchemaInfo, SchemaIntrospector
 
 logger = get_logger(__name__)
@@ -103,6 +105,16 @@ def to_dict(self) -> dict[str, Any]:
         }
 
 
+@dataclass(frozen=True)
+class DatabaseContext:
+    """Resolved database resources for a database id."""
+
+    database_id: str
+    engine: Any
+    dialect: str
+    session_provider: Callable[[], Any]
+
+
 @dataclass
 class ReasoningStep:
     """
@@ -927,6 +939,7 @@ async def validate_sql(
     async def _get_schema_context(self, database_id: str) -> SchemaContext:
         """Get or create schema context for database."""
         start_time = time.perf_counter()
+        db_context = await self._resolve_database_context(database_id)
         # Check cache
         if database_id in self._schema_cache:
             return self._schema_cache[database_id]
@@ -940,7 +953,7 @@ async def _get_schema_context(self, database_id: str) -> SchemaContext:
                 )
                 serialized = cached_payload.get("serialized_schema")
                 if not serialized:
-                    introspector = SchemaIntrospector(self._db_manager.engine)
+                    introspector = SchemaIntrospector(db_context.engine)
                     serialized = introspector.serialize_for_prompt(schema_info)
                 table_names = cached_payload.get("table_names") or [
                     table.name for table in schema_info.tables
@@ -962,7 +975,7 @@ async def _get_schema_context(self, database_id: str) -> SchemaContext:
                 return context
 
         # Create introspector and get schema
-        introspector = SchemaIntrospector(self._db_manager.engine)
+        introspector = SchemaIntrospector(db_context.engine)
         with self._tracing.create_span(
             "schema.introspect",
             attributes={"db.database_id": database_id},
@@ -1004,6 +1017,35 @@ async def _get_schema_context(self, database_id: str) -> SchemaContext:
         self._schema_cache[database_id] = context
         return context
 
+    async def _resolve_database_context(
+        self,
+        database_id: str | None,
+    ) -> DatabaseContext:
+        resolved_id = database_id or "default"
+        settings = get_settings()
+
+        if settings.multi_database.enabled:
+            registry = await get_database_registry()
+            registered = registry.get_database(resolved_id)
+            dialect = (
+                registered.config.dialect.value
+                if registered.config.dialect is not None
+                else registered.engine.dialect.name
+            )
+            return DatabaseContext(
+                database_id=resolved_id,
+                engine=registered.engine,
+                dialect=dialect,
+                session_provider=lambda: registry.session(resolved_id),
+            )
+
+        return DatabaseContext(
+            database_id=resolved_id,
+            engine=self._db_manager.engine,
+            dialect=self._db_manager.dialect,
+            session_provider=self._db_manager.session,
+        )
+
     async def _get_few_shot_examples(
         self,
         natural_query: str,
@@ -1295,20 +1337,21 @@ async def execute_sql(
         """
         from db.executor import SafeQueryExecutor
 
-        async with self._db_manager.session() as session:
+        db_context = await self._resolve_database_context(database_id)
+        async with db_context.session_provider() as session:
             executor = SafeQueryExecutor(
                 session=session,
                 allow_mutations=False,
                 max_rows=max_rows,
-                database_id=database_id or "default",
-                dialect=self._db_manager.dialect,
+                database_id=db_context.database_id,
+                dialect=db_context.dialect,
             )
 
             result: QueryResult = await executor.execute(sql)
 
         logger.info(
             "text2sql_execute_sql",
-            database_id=database_id or "default",
+            database_id=db_context.database_id,
             row_count=result.row_count,
         )
 

db/registry.py

@@ -36,6 +36,7 @@
     convert_url_to_async,
     get_dialect_adapter,
 )
+from db.schema import SchemaInfo
 
 logger = get_logger(__name__)
 
@@ -160,14 +161,23 @@ class RegisteredDatabase:
     session_factory: async_sessionmaker[AsyncSession]
     adapter: DialectAdapter
     health: DatabaseHealth
+    schema: SchemaInfo | None = None
+    schema_updated_at: datetime | None = None
     registered_at: datetime = field(default_factory=datetime.now)
 
     def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary."""
+        schema_metadata = None
+        if self.schema:
+            schema_metadata = {
+                "table_count": len(self.schema.tables),
+                "last_updated": self.schema.last_updated.isoformat(),
+            }
         return {
             **self.config.to_dict(),
             "registered_at": self.registered_at.isoformat(),
             "health": self.health.to_dict(),
+            "schema": schema_metadata,
         }
 
 
@@ -400,6 +410,18 @@ def get_adapter(self, database_id: str) -> DialectAdapter:
         """
         return self.get_database(database_id).adapter
 
+    def update_schema(self, database_id: str, schema: SchemaInfo) -> None:
+        """
+        Store schema metadata for a registered database.
+
+        Args:
+            database_id: Database identifier
+            schema: Extracted schema information
+        """
+        registered = self.get_database(database_id)
+        registered.schema = schema
+        registered.schema_updated_at = schema.last_updated
+
     @asynccontextmanager
     async def session(self, database_id: str) -> AsyncGenerator[AsyncSession, None]:
         """

db/schema.py

@@ -5,6 +5,7 @@
 database schema information for use in SQL generation prompts.
 """
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any
@@ -181,14 +182,24 @@ class SchemaIntrospector:
         prompt_text = introspector.serialize_for_prompt(schema)
     """
 
-    def __init__(self, engine: AsyncEngine) -> None:
+    def __init__(
+        self,
+        engine: AsyncEngine | None = None,
+        engine_provider: Callable[[str], AsyncEngine] | None = None,
+    ) -> None:
         """
         Initialize schema introspector.
 
         Args:
             engine: SQLAlchemy async engine
+            engine_provider: Callable that resolves an engine per database ID
         """
+        if engine is None and engine_provider is None:
+            raise ValueError(
+                "SchemaIntrospector requires an engine or engine_provider."
+            )
         self._engine = engine
+        self._engine_provider = engine_provider
         self._cache: dict[str, SchemaInfo] = {}
 
     async def get_schema(
@@ -214,9 +225,19 @@ async def get_schema(
         logger.info("introspecting_schema", database_id=database_id)
 
         try:
-            async with self._engine.connect() as conn:
+            engine = (
+                self._engine_provider(database_id)
+                if self._engine_provider
+                else self._engine
+            )
+            if engine is None:
+                raise ValueError("SchemaIntrospector engine is not configured.")
+            async with engine.connect() as conn:
                 schema = await self._extract_schema(
-                    conn, database_id, include_row_counts
+                    conn,
+                    database_id,
+                    include_row_counts,
+                    engine.dialect.name,
                 )
                 self._cache[database_id] = schema
                 return schema
@@ -237,6 +258,7 @@ async def _extract_schema(
         conn: AsyncConnection,
         database_id: str,
         include_row_counts: bool,
+        dialect_name: str,
     ) -> SchemaInfo:
         """Extract schema from database connection."""
 
@@ -319,12 +341,9 @@ def sync_inspect(connection: Any) -> dict[str, Any]:
                         error=str(e),
                     )
 
-        # Get dialect
-        dialect = self._engine.dialect.name
-
         return SchemaInfo(
             database_id=database_id,
-            dialect=dialect,
+            dialect=dialect_name,
             tables=tables,
         )