Refactor user creation logic in DAL to support race-safe inserts and return creation status
- Updated the create_user function to use PostgreSQL upsert for concurrent user creation, preventing IntegrityError. - Modified the function to return a tuple containing the user object and a boolean indicating if the user was newly created. - Adjusted the start_command_handler to log user registration only if a new user was created, improving logging clarity.
This commit is contained in:
+33
-10
@@ -4,7 +4,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
from sqlalchemy.orm import selectinload
|
||||
from sqlalchemy import update, delete, func, and_
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
from ..models import User, Subscription
|
||||
|
||||
@@ -33,19 +34,41 @@ async def get_user_by_panel_uuid(
|
||||
## Removed unused generic get_user helper to keep DAL explicit and simple
|
||||
|
||||
|
||||
async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> User:
|
||||
async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple[User, bool]:
|
||||
"""Create a user if not exists in a race-safe way.
|
||||
|
||||
Returns a tuple of (user, created_flag).
|
||||
"""
|
||||
|
||||
if "registration_date" not in user_data:
|
||||
user_data["registration_date"] = datetime.now()
|
||||
user_data["registration_date"] = datetime.now(timezone.utc)
|
||||
|
||||
new_user = User(**user_data)
|
||||
session.add(new_user)
|
||||
await session.flush()
|
||||
await session.refresh(new_user)
|
||||
logging.info(
|
||||
f"New user {new_user.user_id} created in DAL. Referred by: {new_user.referred_by_id or 'N/A'}."
|
||||
# Use PostgreSQL upsert to avoid IntegrityError on concurrent inserts
|
||||
stmt = (
|
||||
pg_insert(User)
|
||||
.values(**user_data)
|
||||
.on_conflict_do_nothing(index_elements=[User.user_id])
|
||||
.returning(User.user_id)
|
||||
)
|
||||
return new_user
|
||||
|
||||
result = await session.execute(stmt)
|
||||
inserted_row = result.first()
|
||||
created = inserted_row is not None
|
||||
|
||||
# Fetch the user (inserted just now or pre-existing)
|
||||
user_id: int = user_data["user_id"]
|
||||
user = await get_user_by_id(session, user_id)
|
||||
|
||||
if created and user is not None:
|
||||
logging.info(
|
||||
f"New user {user.user_id} created in DAL. Referred by: {user.referred_by_id or 'N/A'}."
|
||||
)
|
||||
elif user is not None:
|
||||
logging.info(
|
||||
f"User {user.user_id} already exists in DAL. Proceeding without creation."
|
||||
)
|
||||
|
||||
return user, created
|
||||
|
||||
|
||||
async def update_user(
|
||||
|
||||
Reference in New Issue
Block a user