[FEATURE-BRANCH] refactor: improve API v1 error handling (#4887)

# Description

This is the feature branch for all the issues related with improving
error handling.

The code changes should improve the code without modifying API v1
behavior at all. This means that error responses should be the same.

Closes #4871 

**Type of change**

(Please delete options that are not relevant. Remember to title the PR
according to the type of change)

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [x] Refactor (change restructuring the codebase without changing
functionality)
- [ ] Improvement (change adding some improvement to an existing
functionality)
- [ ] Documentation update

**How Has This Been Tested**

(Please describe the tests that you ran to verify your changes. And
ideally, reference `tests`)

- [x] Modifying and improving already existent tests.

**Checklist**

- [ ] I added relevant documentation
- [ ] follows the style guidelines of this project
- [ ] I did a self-review of my code
- [ ] I made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I filled out [the contributor form](https://tally.so/r/n9XrxK)
(see text above)
- [ ] I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

Author: jfcalvo

argilla-server/src/argilla_server/apis/errors/v1/exception_handlers.py

@@ -19,8 +19,50 @@
 
 
 def add_exception_handlers(app: FastAPI):
+    @app.exception_handler(errors.NotFoundError)
+    async def not_found_error_exception_handler(request, exc):
+        return JSONResponse(
+            status_code=status.HTTP_404_NOT_FOUND,
+            # TODO: Once we move to v2.0 we can remove the content using detail attribute
+            # and use the new one using code and message.
+            # content={"code": exc.code, "message": exc.message},
+            content={"detail": exc.message},
+        )
+
+    @app.exception_handler(errors.NotUniqueError)
+    async def not_unique_error_exception_handler(request, exc):
+        return JSONResponse(
+            status_code=status.HTTP_409_CONFLICT,
+            # TODO: Once we move to v2.0 we can remove the content using detail attribute
+            # and use the new one using code and message.
+            # content={"code": exc.code, "message": exc.message},
+            content={"detail": exc.message},
+        )
+
     @app.exception_handler(errors.UnprocessableEntityError)
     async def unprocessable_entity_error_exception_handler(request, exc):
+        return JSONResponse(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            # TODO: Once we move to v2.0 we can remove the content using detail attribute
+            # and use the new one using code and message.
+            # content={"code": exc.code, "message": exc.message},
+            content={"detail": exc.message},
+        )
+
+    # TODO: This is a temporary exception handler for ValueError exceptions.
+    # This is because we are using ValueError exceptions in some places and we want to
+    # return a 422 status code instead of a 500 status code.
+    # This exception handler should be removed once we move to v2.0 and we use UnprocessableEntityError.
+    @app.exception_handler(ValueError)
+    async def value_error_exception_handler(request, exc):
+        return JSONResponse(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            content={"detail": str(exc)},
+        )
+
+    # TODO: Once we move to v2.0 we can remove this exception handler and use UnprocessableEntityError
+    @app.exception_handler(errors.MissingVectorError)
+    async def missing_vector_error_exception_handler(request, exc):
         return JSONResponse(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             content={"code": exc.code, "message": exc.message},

argilla-server/src/argilla_server/apis/v0/handlers/users.py

@@ -19,24 +19,26 @@
 from fastapi import APIRouter, Depends, HTTPException, Request, Security, status
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from argilla_server import models, telemetry
+from argilla_server import telemetry
 from argilla_server.contexts import accounts
 from argilla_server.database import get_async_db
 from argilla_server.errors import EntityAlreadyExistsError, EntityNotFoundError
 from argilla_server.errors.future import NotUniqueError
+from argilla_server.models import User
 from argilla_server.policies import UserPolicy, authorize
 from argilla_server.pydantic_v1 import parse_obj_as
-from argilla_server.schemas.v0.users import User, UserCreate
+from argilla_server.schemas.v0.users import User as UserSchema
+from argilla_server.schemas.v0.users import UserCreate
 from argilla_server.security import auth
 
 router = APIRouter(tags=["users"])
 
 
-@router.get("/me", response_model=User, response_model_exclude_none=True, operation_id="whoami")
+@router.get("/me", response_model=UserSchema, response_model_exclude_none=True, operation_id="whoami")
 async def whoami(
     request: Request,
     db: AsyncSession = Depends(get_async_db),
-    current_user: models.User = Security(auth.get_current_user),
+    current_user: User = Security(auth.get_current_user),
 ):
     """
     User info endpoint
@@ -56,7 +58,7 @@ async def whoami(
 
     await telemetry.track_login(request, current_user)
 
-    user = User.from_orm(current_user)
+    user = UserSchema.from_orm(current_user)
     # TODO: The current client checks if a user can work on a specific workspace
     #  by using workspaces info returning in `/api/me`.
     #  Returning all workspaces from the `/api/me` for owner users keeps the
@@ -71,23 +73,25 @@ async def whoami(
     return user
 
 
-@router.get("/users", response_model=List[User], response_model_exclude_none=True)
+@router.get("/users", response_model=List[UserSchema], response_model_exclude_none=True)
 async def list_users(
-    *, db: AsyncSession = Depends(get_async_db), current_user: models.User = Security(auth.get_current_user)
+    *,
+    db: AsyncSession = Depends(get_async_db),
+    current_user: User = Security(auth.get_current_user),
 ):
     await authorize(current_user, UserPolicy.list)
 
     users = await accounts.list_users(db)
 
-    return parse_obj_as(List[User], users)
+    return parse_obj_as(List[UserSchema], users)
 
 
-@router.post("/users", response_model=User, response_model_exclude_none=True)
+@router.post("/users", response_model=UserSchema, response_model_exclude_none=True)
 async def create_user(
     *,
     db: AsyncSession = Depends(get_async_db),
     user_create: UserCreate,
-    current_user: models.User = Security(auth.get_current_user),
+    current_user: User = Security(auth.get_current_user),
 ):
     await authorize(current_user, UserPolicy.create)
 
@@ -96,32 +100,32 @@ async def create_user(
 
         telemetry.track_user_created(user)
     except NotUniqueError:
-        raise EntityAlreadyExistsError(name=user_create.username, type=User)
+        raise EntityAlreadyExistsError(name=user_create.username, type=UserSchema)
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(e))
 
     await user.awaitable_attrs.workspaces
 
-    return User.from_orm(user)
+    return user
 
 
-@router.delete("/users/{user_id}", response_model=User, response_model_exclude_none=True)
+@router.delete("/users/{user_id}", response_model=UserSchema, response_model_exclude_none=True)
 async def delete_user(
     *,
     db: AsyncSession = Depends(get_async_db),
     user_id: UUID,
-    current_user: models.User = Security(auth.get_current_user),
+    current_user: User = Security(auth.get_current_user),
 ):
-    user = await accounts.get_user_by_id(db, user_id)
+    user = await User.get(db, user_id)
     if not user:
         # TODO: Forcing here user_id to be an string.
         # Not casting it is causing a `Object of type UUID is not JSON serializable`.
         # Possible solution redefining JSONEncoder.default here:
         # https://github.com/jazzband/django-push-notifications/issues/586
-        raise EntityNotFoundError(name=str(user_id), type=User)
+        raise EntityNotFoundError(name=str(user_id), type=UserSchema)
 
     await authorize(current_user, UserPolicy.delete(user))
 
     await accounts.delete_user(db, user)
 
-    return User.from_orm(user)
+    return user

argilla-server/src/argilla_server/apis/v0/handlers/workspaces.py

@@ -22,16 +22,18 @@
 from argilla_server.database import get_async_db
 from argilla_server.errors import EntityAlreadyExistsError, EntityNotFoundError
 from argilla_server.errors.future import NotUniqueError
+from argilla_server.models import User, Workspace, WorkspaceUser
 from argilla_server.policies import WorkspacePolicy, WorkspaceUserPolicy, authorize
 from argilla_server.pydantic_v1 import parse_obj_as
-from argilla_server.schemas.v0.users import User
-from argilla_server.schemas.v0.workspaces import Workspace, WorkspaceCreate
+from argilla_server.schemas.v0.users import User as UserSchema
+from argilla_server.schemas.v0.workspaces import Workspace as WorkspaceSchema
+from argilla_server.schemas.v0.workspaces import WorkspaceCreate
 from argilla_server.security import auth
 
 router = APIRouter(tags=["workspaces"])
 
 
-@router.post("/workspaces", response_model=Workspace, response_model_exclude_none=True)
+@router.post("/workspaces", response_model=WorkspaceSchema, response_model_exclude_none=True)
 async def create_workspace(
     *,
     db: AsyncSession = Depends(get_async_db),
@@ -45,10 +47,10 @@ async def create_workspace(
     except NotUniqueError:
         raise EntityAlreadyExistsError(name=workspace_create.name, type=Workspace)
 
-    return Workspace.from_orm(workspace)
+    return workspace
 
 
-@router.get("/workspaces/{workspace_id}/users", response_model=List[User], response_model_exclude_none=True)
+@router.get("/workspaces/{workspace_id}/users", response_model=List[UserSchema], response_model_exclude_none=True)
 async def list_workspace_users(
     *,
     db: AsyncSession = Depends(get_async_db),
@@ -57,17 +59,18 @@ async def list_workspace_users(
 ):
     await authorize(current_user, WorkspaceUserPolicy.list(workspace_id))
 
-    workspace = await accounts.get_workspace_by_id(db, workspace_id)
+    workspace = await Workspace.get(db, workspace_id)
     if not workspace:
         raise EntityNotFoundError(name=str(workspace_id), type=Workspace)
 
     await workspace.awaitable_attrs.users
     for user in workspace.users:
         await user.awaitable_attrs.workspaces
-    return parse_obj_as(List[User], workspace.users)
 
+    return parse_obj_as(List[UserSchema], workspace.users)
 
-@router.post("/workspaces/{workspace_id}/users/{user_id}", response_model=User, response_model_exclude_none=True)
+
+@router.post("/workspaces/{workspace_id}/users/{user_id}", response_model=UserSchema, response_model_exclude_none=True)
 async def create_workspace_user(
     *,
     db: AsyncSession = Depends(get_async_db),
@@ -77,40 +80,42 @@ async def create_workspace_user(
 ):
     await authorize(current_user, WorkspaceUserPolicy.create)
 
-    workspace = await accounts.get_workspace_by_id(db, workspace_id)
+    workspace = await Workspace.get(db, workspace_id)
     if not workspace:
         raise EntityNotFoundError(name=str(workspace_id), type=Workspace)
 
-    user = await accounts.get_user_by_id(db, user_id)
+    user = await User.get(db, user_id)
     if not user:
-        raise EntityNotFoundError(name=str(user_id), type=User)
+        raise EntityNotFoundError(name=str(user_id), type=UserSchema)
 
     try:
         workspace_user = await accounts.create_workspace_user(db, {"workspace_id": workspace_id, "user_id": user_id})
     except NotUniqueError:
-        raise EntityAlreadyExistsError(name=str(user_id), type=User)
+        raise EntityAlreadyExistsError(name=str(user_id), type=UserSchema)
 
     await db.refresh(user, attribute_names=["workspaces"])
 
-    return User.from_orm(workspace_user.user)
+    return workspace_user.user
 
 
-@router.delete("/workspaces/{workspace_id}/users/{user_id}", response_model=User, response_model_exclude_none=True)
+@router.delete(
+    "/workspaces/{workspace_id}/users/{user_id}", response_model=UserSchema, response_model_exclude_none=True
+)
 async def delete_workspace_user(
     *,
     db: AsyncSession = Depends(get_async_db),
     workspace_id: UUID,
     user_id: UUID,
     current_user: User = Security(auth.get_current_user),
 ):
-    workspace_user = await accounts.get_workspace_user_by_workspace_id_and_user_id(db, workspace_id, user_id)
+    workspace_user = await WorkspaceUser.get_by(db, workspace_id=workspace_id, user_id=user_id)
     if not workspace_user:
-        raise EntityNotFoundError(name=str(user_id), type=User)
+        raise EntityNotFoundError(name=str(user_id), type=UserSchema)
 
     await authorize(current_user, WorkspaceUserPolicy.delete(workspace_user))
 
     user = await workspace_user.awaitable_attrs.user
     await accounts.delete_workspace_user(db, workspace_user)
     await db.refresh(user, attribute_names=["workspaces"])
 
-    return User.from_orm(user)
+    return user