Add DEMO feature to the project

services/demo_session/app/services/cleanup_service.py

@@ -0,0 +1,147 @@
+"""
+Demo Cleanup Service
+Handles automatic cleanup of expired sessions
+"""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, update
+from datetime import datetime, timezone
+from typing import List
+import structlog
+
+from app.models import DemoSession, DemoSessionStatus
+from app.services.data_cloner import DemoDataCloner
+from app.core import RedisClient
+
+logger = structlog.get_logger()
+
+
+class DemoCleanupService:
+    """Handles cleanup of expired demo sessions"""
+
+    def __init__(self, db: AsyncSession, redis: RedisClient):
+        self.db = db
+        self.redis = redis
+        self.data_cloner = DemoDataCloner(db, redis)
+
+    async def cleanup_expired_sessions(self) -> dict:
+        """
+        Find and cleanup all expired sessions
+
+        Returns:
+            Cleanup statistics
+        """
+        logger.info("Starting demo session cleanup")
+
+        now = datetime.now(timezone.utc)
+
+        # Find expired sessions
+        result = await self.db.execute(
+            select(DemoSession).where(
+                DemoSession.status == DemoSessionStatus.ACTIVE,
+                DemoSession.expires_at < now
+            )
+        )
+        expired_sessions = result.scalars().all()
+
+        stats = {
+            "total_expired": len(expired_sessions),
+            "cleaned_up": 0,
+            "failed": 0,
+            "errors": []
+        }
+
+        for session in expired_sessions:
+            try:
+                # Mark as expired
+                session.status = DemoSessionStatus.EXPIRED
+                await self.db.commit()
+
+                # Delete session data
+                await self.data_cloner.delete_session_data(
+                    str(session.virtual_tenant_id),
+                    session.session_id
+                )
+
+                stats["cleaned_up"] += 1
+
+                logger.info(
+                    "Session cleaned up",
+                    session_id=session.session_id,
+                    age_minutes=(now - session.created_at).total_seconds() / 60
+                )
+
+            except Exception as e:
+                stats["failed"] += 1
+                stats["errors"].append({
+                    "session_id": session.session_id,
+                    "error": str(e)
+                })
+                logger.error(
+                    "Failed to cleanup session",
+                    session_id=session.session_id,
+                    error=str(e)
+                )
+
+        logger.info("Demo session cleanup completed", stats=stats)
+        return stats
+
+    async def cleanup_old_destroyed_sessions(self, days: int = 7) -> int:
+        """
+        Delete destroyed session records older than specified days
+
+        Args:
+            days: Number of days to keep destroyed sessions
+
+        Returns:
+            Number of deleted records
+        """
+        from datetime import timedelta
+
+        cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+
+        result = await self.db.execute(
+            select(DemoSession).where(
+                DemoSession.status == DemoSessionStatus.DESTROYED,
+                DemoSession.destroyed_at < cutoff_date
+            )
+        )
+        old_sessions = result.scalars().all()
+
+        for session in old_sessions:
+            await self.db.delete(session)
+
+        await self.db.commit()
+
+        logger.info(
+            "Old destroyed sessions deleted",
+            count=len(old_sessions),
+            older_than_days=days
+        )
+
+        return len(old_sessions)
+
+    async def get_cleanup_stats(self) -> dict:
+        """Get cleanup statistics"""
+        result = await self.db.execute(select(DemoSession))
+        all_sessions = result.scalars().all()
+
+        now = datetime.now(timezone.utc)
+
+        active_count = len([s for s in all_sessions if s.status == DemoSessionStatus.ACTIVE])
+        expired_count = len([s for s in all_sessions if s.status == DemoSessionStatus.EXPIRED])
+        destroyed_count = len([s for s in all_sessions if s.status == DemoSessionStatus.DESTROYED])
+
+        # Find sessions that should be expired but aren't marked yet
+        should_be_expired = len([
+            s for s in all_sessions
+            if s.status == DemoSessionStatus.ACTIVE and s.expires_at < now
+        ])
+
+        return {
+            "total_sessions": len(all_sessions),
+            "active_sessions": active_count,
+            "expired_sessions": expired_count,
+            "destroyed_sessions": destroyed_count,
+            "pending_cleanup": should_be_expired
+        }