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]: