60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
"""
|
|
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
|
|
|