cleanup_anonymous.py

#!/usr/bin/env python3
"""
Cleanup script for anonymous session data
Runs on app startup to clear stale anonymous sessions
"""

import os
import sys
from datetime import datetime, timedelta
from pymongo import MongoClient
import logging

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def cleanup_anonymous_sessions():
    """Remove all anonymous session data from MongoDB"""
    try:
        # Get MongoDB connection from environment
        mongodb_url = os.environ.get('MONGODB_URL') or os.environ.get('MONGODB_URI')
        if not mongodb_url:
            logger.warning("No MongoDB URL found in environment. Skipping cleanup.")
            return True
        
        # Connect to MongoDB
        client = MongoClient(mongodb_url)
        db_name = os.environ.get('MONGODB_DATABASE', 'Atlas')
        db = client[db_name]
        
        # Clear anonymous sessions from chat_sessions collection
        # Anonymous sessions have user_id starting with 'anon_' or are older than 6 hours
        if 'chat_sessions' in db.list_collection_names():
            # Remove sessions with anonymous user IDs
            result = db.chat_sessions.delete_many({
                '$or': [
                    {'user_id': {'$regex': '^anon_'}},
                    {'session_data.anonymous': True}
                ]
            })
            logger.info(f"Removed {result.deleted_count} anonymous chat sessions")
        
        # Clear any temporary anonymous user records if they exist
        if 'users' in db.list_collection_names():
            # Remove any accidentally created anonymous user records
            result = db.users.delete_many({
                '$or': [
                    {'email': {'$regex': '^anon_'}},
                    {'is_anonymous': True}
                ]
            })
            if result.deleted_count > 0:
                logger.info(f"Removed {result.deleted_count} anonymous user records")
        
        # Clear expired sessions (older than 24 hours)
        cutoff_time = datetime.utcnow() - timedelta(hours=24)
        if 'sessions' in db.list_collection_names():
            result = db.sessions.delete_many({
                'expiry': {'$lt': cutoff_time}
            })
            if result.deleted_count > 0:
                logger.info(f"Removed {result.deleted_count} expired sessions")
        
        logger.info("Anonymous session cleanup completed successfully")
        return True
        
    except Exception as e:
        logger.error(f"Error during anonymous session cleanup: {e}")
        # Don't fail the app startup if cleanup fails
        return False

if __name__ == "__main__":
    # Load environment variables if .env exists
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass
    
    success = cleanup_anonymous_sessions()
    sys.exit(0 if success else 1)