Improve the demo feature of the project

services/demo_session/app/api/demo_sessions.py

@@ -20,6 +20,41 @@ logger = structlog.get_logger()
 route_builder = RouteBuilder('demo')
 
 
+async def _background_cloning_task(session_id: str, session_obj_id: UUID, base_tenant_id: str):
+    """Background task for orchestrated cloning - creates its own DB session"""
+    from app.core.database import db_manager
+    from app.models import DemoSession
+    from sqlalchemy import select
+
+    # Create new database session for background task
+    async with db_manager.session_factory() as db:
+        try:
+            # Get Redis client
+            redis = await get_redis()
+
+            # Fetch the session from the database
+            result = await db.execute(
+                select(DemoSession).where(DemoSession.id == session_obj_id)
+            )
+            session = result.scalar_one_or_none()
+
+            if not session:
+                logger.error("Session not found for cloning", session_id=session_id)
+                return
+
+            # Create session manager with new DB session
+            session_manager = DemoSessionManager(db, redis)
+            await session_manager.trigger_orchestrated_cloning(session, base_tenant_id)
+
+        except Exception as e:
+            logger.error(
+                "Background cloning failed",
+                session_id=session_id,
+                error=str(e),
+                exc_info=True
+            )
+
+
 @router.post(
     route_builder.build_base_route("sessions", include_tenant_prefix=False),
     response_model=DemoSessionResponse,
@@ -46,23 +81,20 @@ async def create_demo_session(
             user_agent=user_agent
         )
 
-        # Trigger async data cloning job
-        from app.services.k8s_job_cloner import K8sJobCloner
+        # Trigger async orchestrated cloning in background
         import asyncio
+        from app.core.config import settings
+        from app.models import DemoSession
 
-        job_cloner = K8sJobCloner()
+        # Get base tenant ID from config
+        demo_config = settings.DEMO_ACCOUNTS.get(request.demo_account_type, {})
+        base_tenant_id = demo_config.get("base_tenant_id", str(session.base_demo_tenant_id))
+
+        # Start cloning in background task with session ID (not session object)
         asyncio.create_task(
-            job_cloner.clone_tenant_data(
-                session.session_id,
-                "",
-                str(session.virtual_tenant_id),
-                request.demo_account_type
-            )
+            _background_cloning_task(session.session_id, session.id, base_tenant_id)
         )
 
-        await session_manager.mark_data_cloned(session.session_id)
-        await session_manager.mark_redis_populated(session.session_id)
-
         # Generate session token
         session_token = jwt.encode(
             {
@@ -110,6 +142,61 @@ async def get_session_info(
     return session.to_dict()
 
 
+@router.get(
+    route_builder.build_resource_detail_route("sessions", "session_id", include_tenant_prefix=False) + "/status",
+    response_model=dict
+)
+async def get_session_status(
+    session_id: str = Path(...),
+    db: AsyncSession = Depends(get_db),
+    redis: RedisClient = Depends(get_redis)
+):
+    """
+    Get demo session provisioning status
+
+    Returns current status of data cloning and readiness.
+    Use this endpoint for polling (recommended interval: 1-2 seconds).
+    """
+    session_manager = DemoSessionManager(db, redis)
+    status = await session_manager.get_session_status(session_id)
+
+    if not status:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    return status
+
+
+@router.post(
+    route_builder.build_resource_detail_route("sessions", "session_id", include_tenant_prefix=False) + "/retry",
+    response_model=dict
+)
+async def retry_session_cloning(
+    session_id: str = Path(...),
+    db: AsyncSession = Depends(get_db),
+    redis: RedisClient = Depends(get_redis)
+):
+    """
+    Retry failed cloning operations
+
+    Only available for sessions in "failed" or "partial" status.
+    """
+    try:
+        session_manager = DemoSessionManager(db, redis)
+        result = await session_manager.retry_failed_cloning(session_id)
+
+        return {
+            "message": "Cloning retry initiated",
+            "session_id": session_id,
+            "result": result
+        }
+
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:
+        logger.error("Failed to retry cloning", error=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.delete(
     route_builder.build_resource_detail_route("sessions", "session_id", include_tenant_prefix=False),
     response_model=dict
@@ -129,3 +216,24 @@ async def destroy_demo_session(
     except Exception as e:
         logger.error("Failed to destroy session", error=str(e))
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post(
+    route_builder.build_resource_detail_route("sessions", "session_id", include_tenant_prefix=False) + "/destroy",
+    response_model=dict
+)
+async def destroy_demo_session_post(
+    session_id: str = Path(...),
+    db: AsyncSession = Depends(get_db),
+    redis: RedisClient = Depends(get_redis)
+):
+    """Destroy demo session via POST (for frontend compatibility)"""
+    try:
+        session_manager = DemoSessionManager(db, redis)
+        await session_manager.destroy_session(session_id)
+
+        return {"message": "Session destroyed successfully", "session_id": session_id}
+
+    except Exception as e:
+        logger.error("Failed to destroy session", error=str(e))
+        raise HTTPException(status_code=500, detail=str(e))