[FEATURE-BRANCH] refactor: improve API v1 error handling

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