REFACTOR ALL APIs fix 1

services/auth/app/repositories/onboarding_repository.py

@@ -51,17 +51,26 @@ class OnboardingRepository:
             return None
 
     async def upsert_user_step(
-        self, 
-        user_id: str, 
-        step_name: str, 
-        completed: bool, 
-        step_data: Dict[str, Any] = None
+        self,
+        user_id: str,
+        step_name: str,
+        completed: bool,
+        step_data: Dict[str, Any] = None,
+        auto_commit: bool = True
     ) -> UserOnboardingProgress:
-        """Insert or update a user's onboarding step"""
+        """Insert or update a user's onboarding step
+
+        Args:
+            user_id: User ID
+            step_name: Name of the step
+            completed: Whether the step is completed
+            step_data: Additional data for the step
+            auto_commit: Whether to auto-commit (set to False when used within UnitOfWork)
+        """
         try:
             completed_at = datetime.now(timezone.utc) if completed else None
             step_data = step_data or {}
-            
+
             # Use PostgreSQL UPSERT (INSERT ... ON CONFLICT ... DO UPDATE)
             stmt = insert(UserOnboardingProgress).values(
                 user_id=user_id,
@@ -71,7 +80,7 @@ class OnboardingRepository:
                 step_data=step_data,
                 updated_at=datetime.now(timezone.utc)
             )
-            
+
             # On conflict, update the existing record
             stmt = stmt.on_conflict_do_update(
                 index_elements=['user_id', 'step_name'],
@@ -82,17 +91,24 @@ class OnboardingRepository:
                     updated_at=stmt.excluded.updated_at
                 )
             )
-            
+
             # Return the updated record
             stmt = stmt.returning(UserOnboardingProgress)
             result = await self.db.execute(stmt)
-            await self.db.commit()
-            
+
+            # Only commit if auto_commit is True (not within a UnitOfWork)
+            if auto_commit:
+                await self.db.commit()
+            else:
+                # Flush to ensure the statement is executed
+                await self.db.flush()
+
             return result.scalars().first()
-            
+
         except Exception as e:
             logger.error(f"Error upserting step {step_name} for user {user_id}: {e}")
-            await self.db.rollback()
+            if auto_commit:
+                await self.db.rollback()
             raise
 
     async def get_user_summary(self, user_id: str) -> Optional[UserOnboardingSummary]:

services/auth/app/services/auth_service.py

@@ -108,16 +108,15 @@ class EnhancedAuthService:
                     }
                     
                     await token_repo.create_token(token_data)
-                    
-                    # Commit transaction
-                    await uow.commit()
 
-                    # Store subscription plan selection in onboarding progress for later retrieval
+                    # Store subscription plan selection in onboarding progress BEFORE committing
+                    # This ensures it's part of the same transaction
                     if user_data.subscription_plan or user_data.use_trial or user_data.payment_method_id:
                         try:
                             from app.repositories.onboarding_repository import OnboardingRepository
                             from app.models.onboarding import UserOnboardingProgress
 
+                            # Use upsert_user_step instead of save_step_data to avoid double commits
                             onboarding_repo = OnboardingRepository(db_session)
                             plan_data = {
                                 "subscription_plan": user_data.subscription_plan or "starter",
@@ -126,17 +125,29 @@ class EnhancedAuthService:
                                 "saved_at": datetime.now(timezone.utc).isoformat()
                             }
 
-                            await onboarding_repo.save_step_data(
-                                str(new_user.id),
-                                "user_registered",
-                                plan_data
+                            # Create the onboarding step record with plan data
+                            # Note: We use completed=True to mark user_registered as complete
+                            # auto_commit=False to let UnitOfWork handle the commit
+                            await onboarding_repo.upsert_user_step(
+                                user_id=str(new_user.id),
+                                step_name="user_registered",
+                                completed=True,
+                                step_data=plan_data,
+                                auto_commit=False
                             )
 
                             logger.info("Subscription plan saved to onboarding progress",
                                        user_id=new_user.id,
                                        plan=user_data.subscription_plan)
                         except Exception as e:
-                            logger.warning("Failed to save subscription plan to onboarding progress", error=str(e))
+                            logger.error("Failed to save subscription plan to onboarding progress",
+                                        user_id=new_user.id,
+                                        error=str(e))
+                            # Re-raise to ensure registration fails if onboarding data can't be saved
+                            raise
+
+                    # Commit transaction (includes user, tokens, and onboarding data)
+                    await uow.commit()
 
                     # Publish registration event (non-blocking)
                     try: