async def update(
        self,
        db: AsyncSession,
        *,
        db_obj: User,
        obj_in: UserUpdate | dict[str, Any],
        current_user: User,
    ) -> User:
        """Update an object.

        Args:
            db (AsyncSession): The database session.
            db_obj (User): The object to update.
            obj_in (UserUpdate | dict[str, Any]): The updated object.
            current_user (User): The current user.

        Returns:
            User: The updated user.
        """
        # Ensure the organization_id is not removed
        if isinstance(obj_in, dict):
            update_data = obj_in
            if "organization_id" in update_data and update_data["organization_id"] is None:
                raise ValueError("User must have an organization")
        else:
            update_data = obj_in.model_dump(exclude_unset=True)
            if "organization_id" in update_data and update_data["organization_id"] is None:
                raise ValueError("User must have an organization")

        updated_user = await super().update(db, db_obj=db_obj, obj_in=update_data)

        # Explicitly load the organization after update
        stmt = (
            select(User).where(User.id == updated_user.id).options(selectinload(User.organization))
        )
        result = await db.execute(stmt)
        return result.scalar_one_or_none() or updated_user
Update an object.

        Args:
            db (AsyncSession): The database session.
            db_obj (User): The object to update.
            obj_in (UserUpdate | dict[str, Any]): The updated object.
            current_user (User): The current user.

        Returns:
            User: The updated user.
        
backend/airweave/crud/crud_user.py