Tenant provisioning in the dataplane (#2694)

* add tenant provisioning to data plane

* minor typing update

* ensure tenant router included

* proper auth check

* update disabling logic

* validated basic provisioning

* use new kv store

backend/ee/danswer/main.py

@@ -34,6 +34,7 @@ from ee.danswer.server.query_history.api import router as query_history_router
 from ee.danswer.server.reporting.usage_export_api import router as usage_export_router
 from ee.danswer.server.saml import router as saml_router
 from ee.danswer.server.seeding import seed_db
+from ee.danswer.server.tenants.api import router as tenants_router
 from ee.danswer.server.token_rate_limits.api import (
     router as token_rate_limit_settings_router,
 )
@@ -79,6 +80,8 @@ def get_application() -> FastAPI:
 
     # RBAC / group access control
     include_router_with_global_prefix_prepended(application, user_group_router)
+    # Tenant management
+    include_router_with_global_prefix_prepended(application, tenants_router)
     # Analytics endpoints
     include_router_with_global_prefix_prepended(application, analytics_router)
     include_router_with_global_prefix_prepended(application, query_history_router)

backend/ee/danswer/server/tenants/api.py

@@ -0,0 +1,46 @@
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+
+from danswer.auth.users import control_plane_dep
+from danswer.configs.app_configs import MULTI_TENANT
+from danswer.db.engine import get_session_with_tenant
+from danswer.setup import setup_danswer
+from danswer.utils.logger import setup_logger
+from ee.danswer.server.tenants.models import CreateTenantRequest
+from ee.danswer.server.tenants.provisioning import ensure_schema_exists
+from ee.danswer.server.tenants.provisioning import run_alembic_migrations
+
+logger = setup_logger()
+router = APIRouter(prefix="/tenants")
+
+
+@router.post("/create")
+def create_tenant(
+    create_tenant_request: CreateTenantRequest, _: None = Depends(control_plane_dep)
+) -> dict[str, str]:
+    try:
+        tenant_id = create_tenant_request.tenant_id
+
+        if not MULTI_TENANT:
+            raise HTTPException(status_code=403, detail="Multi-tenancy is not enabled")
+
+        if not ensure_schema_exists(tenant_id):
+            logger.info(f"Created schema for tenant {tenant_id}")
+        else:
+            logger.info(f"Schema already exists for tenant {tenant_id}")
+
+        run_alembic_migrations(tenant_id)
+        with get_session_with_tenant(tenant_id) as db_session:
+            setup_danswer(db_session)
+
+        logger.info(f"Tenant {tenant_id} created successfully")
+        return {
+            "status": "success",
+            "message": f"Tenant {tenant_id} created successfully",
+        }
+    except Exception as e:
+        logger.exception(f"Failed to create tenant {tenant_id}: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to create tenant: {str(e)}"
+        )

backend/ee/danswer/server/tenants/models.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class CreateTenantRequest(BaseModel):
+    tenant_id: str
+    initial_admin_email: str

backend/ee/danswer/server/tenants/provisioning.py

@@ -0,0 +1,63 @@
+import os
+from types import SimpleNamespace
+
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+from sqlalchemy.schema import CreateSchema
+
+from alembic import command
+from alembic.config import Config
+from danswer.db.engine import build_connection_string
+from danswer.db.engine import get_sqlalchemy_engine
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def run_alembic_migrations(schema_name: str) -> None:
+    logger.info(f"Starting Alembic migrations for schema: {schema_name}")
+
+    try:
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        root_dir = os.path.abspath(os.path.join(current_dir, "..", "..", "..", ".."))
+        alembic_ini_path = os.path.join(root_dir, "alembic.ini")
+
+        # Configure Alembic
+        alembic_cfg = Config(alembic_ini_path)
+        alembic_cfg.set_main_option("sqlalchemy.url", build_connection_string())
+        alembic_cfg.set_main_option(
+            "script_location", os.path.join(root_dir, "alembic")
+        )
+
+        # Mimic command-line options by adding 'cmd_opts' to the config
+        alembic_cfg.cmd_opts = SimpleNamespace()  # type: ignore
+        alembic_cfg.cmd_opts.x = [f"schema={schema_name}"]  # type: ignore
+
+        # Run migrations programmatically
+        command.upgrade(alembic_cfg, "head")
+
+        # Run migrations programmatically
+        logger.info(
+            f"Alembic migrations completed successfully for schema: {schema_name}"
+        )
+
+    except Exception as e:
+        logger.exception(f"Alembic migration failed for schema {schema_name}: {str(e)}")
+        raise
+
+
+def ensure_schema_exists(tenant_id: str) -> bool:
+    with Session(get_sqlalchemy_engine()) as db_session:
+        with db_session.begin():
+            result = db_session.execute(
+                text(
+                    "SELECT schema_name FROM information_schema.schemata WHERE schema_name = :schema_name"
+                ),
+                {"schema_name": tenant_id},
+            )
+            schema_exists = result.scalar() is not None
+            if not schema_exists:
+                stmt = CreateSchema(tenant_id)
+                db_session.execute(stmt)
+                return True
+            return False