validated basic provisioning

danswer-ai · Oct 6, 2024 · a7b0a0d · a7b0a0d
1 parent 3a144d1
commit a7b0a0d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 18 deletions.
diff --git a/backend/danswer/db/engine.py b/backend/danswer/db/engine.py
@@ -117,7 +117,7 @@ def get_db_current_time(db_session: Session) -> datetime:
 
 
 # Regular expression to validate schema names to prevent SQL injection
-SCHEMA_NAME_REGEX = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+SCHEMA_NAME_REGEX = re.compile(r"^[a-zA-Z0-9_-]+$")
 
 
 def is_valid_schema_name(name: str) -> bool:
@@ -281,6 +281,7 @@ def get_session_with_tenant(tenant_id: str | None = None) -> Session:
         tenant_id = current_tenant_id.get()
 
     if not is_valid_schema_name(tenant_id):
+        logger.error(f"Invalid tenant ID: {tenant_id}")
         raise Exception("Invalid tenant ID")
 
     engine = SqlEngine.get_engine()

diff --git a/backend/ee/danswer/server/tenants/api.py b/backend/ee/danswer/server/tenants/api.py
@@ -19,19 +19,28 @@
 def create_tenant(
     create_tenant_request: CreateTenantRequest, _: None = Depends(control_plane_dep)
 ) -> dict[str, str]:
-    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"}
+    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)}"
+        )
diff --git a/backend/ee/danswer/server/tenants/provisioning.py b/backend/ee/danswer/server/tenants/provisioning.py
@@ -19,12 +19,15 @@ def run_alembic_migrations(schema_name: str) -> None:
 
     try:
         current_dir = os.path.dirname(os.path.abspath(__file__))
-        root_dir = os.path.abspath(os.path.join(current_dir, "..", "..", ".."))
+        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