Add source

shared/database/user_helpers.py

@@ -0,0 +1,59 @@
+"""
+Helper functions for user management
+"""
+from sqlalchemy.ext.asyncio import AsyncSession
+from shared.database.models import User
+from datetime import datetime
+from sqlalchemy.exc import IntegrityError
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+async def ensure_user_exists(user_id: int, db: AsyncSession) -> User:
+    """
+    Ensure user exists in database, create if not exists.
+    
+    Args:
+        user_id: User ID
+        db: Database session
+        
+    Returns:
+        User object (existing or newly created)
+    """
+    # Check if user exists
+    user = await db.get(User, user_id)
+    if user:
+        return user
+    
+    # User doesn't exist, create it
+    user = User(
+        user_id=user_id,
+        username=None,
+        first_name=None,
+        last_name=None,
+        is_admin=False,
+        is_blocked=False,
+        created_at=datetime.utcnow()
+    )
+    
+    try:
+        db.add(user)
+        await db.commit()
+        logger.info(f"Automatically created user {user_id} in database")
+        return user
+    except IntegrityError:
+        # User was created by another request/thread, fetch it
+        await db.rollback()
+        user = await db.get(User, user_id)
+        if user:
+            logger.debug(f"User {user_id} was created concurrently, using existing record")
+            return user
+        else:
+            logger.error(f"Failed to create or fetch user {user_id}")
+            raise Exception(f"Failed to ensure user {user_id} exists")
+    except Exception as e:
+        await db.rollback()
+        logger.error(f"Error ensuring user {user_id} exists: {e}", exc_info=True)
+        raise
+