bakery-ia/services/demo_session/app/api/demo_sessions.py

"""
Demo Sessions API - Atomic CRUD operations on DemoSession model
"""

from fastapi import APIRouter, Depends, HTTPException, Path, Query, Request
from typing import Optional
from uuid import UUID
from datetime import datetime, timezone
import structlog
import jwt

from app.api.schemas import DemoSessionCreate, DemoSessionResponse
from app.services import DemoSessionManager
from app.core import get_db
from app.core.redis_wrapper import get_redis, DemoRedisWrapper
from sqlalchemy.ext.asyncio import AsyncSession
from shared.routing import RouteBuilder

router = APIRouter(tags=["demo-sessions"])
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
    from app.core.redis_wrapper import get_redis

    # 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
            )
            # Attempt to update session status to failed if possible
            try:
                from app.core.database import db_manager
                from app.models import DemoSession
                from sqlalchemy import select, update

                # Try to update the session directly in DB to mark it as failed
                async with db_manager.session_factory() as update_db:
                    from app.models import DemoSessionStatus
                    update_result = await update_db.execute(
                        update(DemoSession)
                        .where(DemoSession.id == session_obj_id)
                        .values(status=DemoSessionStatus.FAILED, cloning_completed_at=datetime.now(timezone.utc))
                    )
                    await update_db.commit()
            except Exception as update_error:
                logger.error(
                    "Failed to update session status to FAILED after background task error",
                    session_id=session_id,
                    error=str(update_error)
                )


def _handle_task_result(task, session_id: str):
    """Handle the result of the background cloning task"""
    try:
        # This will raise the exception if the task failed
        task.result()
    except Exception as e:
        logger.error(
            "Background cloning task failed with exception",
            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,
    status_code=201
)
async def create_demo_session(
    request: DemoSessionCreate,
    http_request: Request,
    db: AsyncSession = Depends(get_db),
    redis: DemoRedisWrapper = Depends(get_redis)
):
    """Create a new isolated demo session (ATOMIC)"""
    logger.info("Creating demo session", demo_account_type=request.demo_account_type)

    try:
        ip_address = request.ip_address or http_request.client.host
        user_agent = request.user_agent or http_request.headers.get("user-agent", "")

        session_manager = DemoSessionManager(db, redis)
        session = await session_manager.create_session(
            demo_account_type=request.demo_account_type,
            subscription_tier=request.subscription_tier,
            user_id=request.user_id,
            ip_address=ip_address,
            user_agent=user_agent
        )

        # Trigger async orchestrated cloning in background
        import asyncio
        from app.core.config import settings
        from app.models import DemoSession

        # 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)
        # Store task reference in case we need to track it
        task = asyncio.create_task(
            _background_cloning_task(session.session_id, session.id, base_tenant_id)
        )

        # Add error handling for the task to prevent silent failures
        task.add_done_callback(lambda t: _handle_task_result(t, session.session_id))

        # Generate session token
        session_token = jwt.encode(
            {
                "session_id": session.session_id,
                "virtual_tenant_id": str(session.virtual_tenant_id),
                "demo_account_type": request.demo_account_type,
                "exp": session.expires_at.timestamp()
            },
            settings.JWT_SECRET_KEY,
            algorithm=settings.JWT_ALGORITHM
        )

        return {
            "session_id": session.session_id,
            "virtual_tenant_id": str(session.virtual_tenant_id),
            "demo_account_type": session.demo_account_type,
            "status": session.status.value,
            "created_at": session.created_at,
            "expires_at": session.expires_at,
            "demo_config": session.session_metadata.get("demo_config", {}),
            "session_token": session_token
        }

    except Exception as e:
        logger.error("Failed to create demo session", error=str(e))
        raise HTTPException(status_code=500, detail=f"Failed to create demo session: {str(e)}")


@router.get(
    route_builder.build_resource_detail_route("sessions", "session_id", include_tenant_prefix=False),
    response_model=dict
)
async def get_session_info(
    session_id: str = Path(...),
    db: AsyncSession = Depends(get_db),
    redis: DemoRedisWrapper = Depends(get_redis)
):
    """Get demo session information (ATOMIC READ)"""
    session_manager = DemoSessionManager(db, redis)
    session = await session_manager.get_session(session_id)

    if not session:
        raise HTTPException(status_code=404, detail="Session not found")

    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: DemoRedisWrapper = 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: DemoRedisWrapper = 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
)
async def destroy_demo_session(
    session_id: str = Path(...),
    db: AsyncSession = Depends(get_db),
    redis: DemoRedisWrapper = Depends(get_redis)
):
    """Destroy demo session and cleanup resources (ATOMIC DELETE)"""
    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))


@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: DemoRedisWrapper = 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))