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:
machka-pasla
2025-08-20 15:49:12 +03:00
parent d2402fea77
commit f22e359684
2 changed files with 50 additions and 26 deletions
+2 -1
View File
@@ -165,8 +165,9 @@ async def start_command_handler(message: types.Message,
"registration_date": datetime.now(timezone.utc) "registration_date": datetime.now(timezone.utc)
} }
try: try:
db_user = await user_dal.create_user(session, user_data_to_create) db_user, created = await user_dal.create_user(session, user_data_to_create)
if created:
logging.info( logging.info(
f"New user {user_id} added to session. Referred by: {referred_by_user_id or 'N/A'}." f"New user {user_id} added to session. Referred by: {referred_by_user_id or 'N/A'}."
) )
+33 -10
View File
@@ -4,7 +4,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select from sqlalchemy.future import select
from sqlalchemy.orm import selectinload from sqlalchemy.orm import selectinload
from sqlalchemy import update, delete, func, and_ 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 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 ## 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: 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) # Use PostgreSQL upsert to avoid IntegrityError on concurrent inserts
session.add(new_user) stmt = (
await session.flush() pg_insert(User)
await session.refresh(new_user) .values(**user_data)
logging.info( .on_conflict_do_nothing(index_elements=[User.user_id])
f"New user {new_user.user_id} created in DAL. Referred by: {new_user.referred_by_id or 'N/A'}." .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( async def update_user(