feat: add GET /api/v1/me/workspaces & update Workspace class

CHANGELOG.md

@@ -24,6 +24,7 @@ These are the section headers that we use:
 - Added `POST /api/v1/records/{record_id}/suggestions` API endpoint to create a suggestion for a response associated to a record ([#3304](https://github.com/argilla-io/argilla/pull/3304)).
 - Added breaking simutaneously running tests within GitHub package worflows. ([#3354](https://github.com/argilla-io/argilla/pull/3354)).
 - Added `allowed_for_roles` Python decorator to check whether the current user has the required role to access the decorated function/method for `User` and `Workspace` ([#3383](https://github.com/argilla-io/argilla/pull/3383))
+- Added `GET /api/v1/me/workspaces` endpoint to list the workspaces of the current active user ([#3390](https://github.com/argilla-io/argilla/pull/3390))
 
 ### Changed
 

src/argilla/client/sdk/users/api.py

@@ -44,6 +44,33 @@ def whoami(client: AuthenticatedClient) -> UserModel:
     return UserModel(**response)
 
 
+# TODO(frascuchon): rename this to `whoami` and deprecate the current `whoami` function
+# in favor of this one, as this is just a patch.
+def whoami_httpx(client: httpx.Client) -> Response[Union[UserModel, ErrorMessage, HTTPValidationError]]:
+    """Sends a GET request to `/api/me` endpoint to get the current user information.
+
+    Args:
+        client: the authenticated Argilla client to be used to send the request to the API.
+
+    Returns:
+        A `Response` object containing a `parsed` attribute with the parsed response if
+        the request was successful, which is an instance of `UserModel`.
+    """
+    url = "/api/me"
+
+    response = client.get(url)
+
+    if response.status_code == 200:
+        parsed_response = UserModel(**response.json())
+        return Response(
+            status_code=response.status_code,
+            content=response.content,
+            headers=response.headers,
+            parsed=parsed_response,
+        )
+    return handle_response_error(response)
+
+
 def list_users(
     client: httpx.Client,
 ) -> Response[Union[List[UserModel], ErrorMessage, HTTPValidationError]]:

src/argilla/client/sdk/v1/users/api.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import List, Optional, Union
+from typing import List, Union
 from uuid import UUID
 
 import httpx

src/argilla/client/sdk/v1/workspaces/api.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Union
+from typing import List, Union
 from uuid import UUID
 
 import httpx
@@ -53,3 +53,31 @@ def get_workspace(
             parsed=parsed_response,
         )
     return handle_response_error(response)
+
+
+def list_workspaces_me(
+    client: httpx.Client,
+) -> Response[Union[List[WorkspaceModel], ErrorMessage, HTTPValidationError]]:
+    """Sends a GET request to `/api/v1/me/workspaces` endpoint to get the list of
+    workspaces the current user has access to.
+
+    Args:
+        client: the authenticated Argilla client to be used to send the request to the API.
+
+    Returns:
+        A `Response` object containing a `parsed` attribute with the parsed response if
+        the request was successful, which is a list of `WorkspaceModel`.
+    """
+    url = "/api/v1/me/workspaces"
+
+    response = client.get(url=url)
+
+    if response.status_code == 200:
+        parsed_response = [WorkspaceModel(**workspace) for workspace in response.json()["items"]]
+        return Response(
+            status_code=response.status_code,
+            content=response.content,
+            headers=response.headers,
+            parsed=parsed_response,
+        )
+    return handle_response_error(response)

src/argilla/client/users.py

@@ -26,13 +26,14 @@
 from argilla.client.sdk.users import api as users_api
 from argilla.client.sdk.users.models import UserCreateModel, UserModel, UserRole
 from argilla.client.sdk.v1.users import api as users_api_v1
-from argilla.client.sdk.v1.workspaces.models import WorkspaceModel
+from argilla.client.sdk.v1.workspaces import api as workspaces_api_v1
 from argilla.client.utils import allowed_for_roles
 
 if TYPE_CHECKING:
     import httpx
 
     from argilla.client.sdk.client import AuthenticatedClient
+    from argilla.client.sdk.v1.workspaces.models import WorkspaceModel
 
 
 class User:
@@ -60,7 +61,7 @@ class User:
         >>> from argilla import rg
         >>> user = rg.User.from_name("my-user") # or `User.from_id("...")`
         >>> print(user)
-        User(id='...', username='my-user', first_name='Luke', last_name="Skywalker', full_name='Luke Skywalker', role='annotator', workspaces=[WorkspaceModel(...), ...], api_key='...', inserted_at=datetime.datetime(2021, 8, 31, 10, 0, 0), updated_at=datetime.datetime(2021, 8, 31, 10, 0, 0))
+        User(id='...', username='my-user', role='annotator', first_name='Luke', last_name="Skywalker', full_name='Luke Skywalker', role='annotator', api_key='...', inserted_at=datetime.datetime(2021, 8, 31, 10, 0, 0), updated_at=datetime.datetime(2021, 8, 31, 10, 0, 0))
     """
 
     __client: "httpx.Client"
@@ -108,21 +109,23 @@ def __init__(
         raise Exception(error_msg)
 
     @property
-    @allowed_for_roles(roles=[UserRole.owner])
-    def workspaces(self) -> Optional[List[WorkspaceModel]]:
+    def workspaces(self) -> Optional[List["WorkspaceModel"]]:
         """Returns the workspace names the current user is linked to.
 
         Returns:
             A list of `WorkspaceModel` the current user is linked to.
         """
-        return users_api_v1.list_user_workspaces(self.__client, self.id).parsed
+        connected_user = users_api.whoami_httpx(self.__client).parsed
+        if connected_user.role == UserRole.owner:
+            return users_api_v1.list_user_workspaces(self.__client, self.id).parsed
+        return workspaces_api_v1.list_workspaces_me(self.__client).parsed
 
     def __repr__(self) -> str:
         return (
             f"User(id={self.id}, username={self.username}, role={self.role},"
-            f" workspaces={self.workspaces}, api_key={self.api_key},"
-            f" first_name={self.first_name}, last_name={self.last_name}, role={self.role},"
-            f" inserted_at={self.inserted_at}, updated_at={self.updated_at})"
+            f" api_key={self.api_key}, first_name={self.first_name},"
+            f" last_name={self.last_name}, inserted_at={self.inserted_at},"
+            f" updated_at={self.updated_at})"
         )
 
     @staticmethod

src/argilla/client/utils.py

@@ -20,7 +20,8 @@
 except ImportError:
     from typing_extensions import ParamSpec
 
-from argilla.client.api import ArgillaSingleton
+from argilla.client.api import active_client
+from argilla.client.sdk.users import api as users_api
 from argilla.client.sdk.users.models import UserRole
 
 _P = ParamSpec("_P")
@@ -43,10 +44,11 @@ def allowed_for_roles(roles: List[UserRole]) -> Callable[[Callable[_P, _R]], Cal
     def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
         @functools.wraps(func)
         def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
-            role = ArgillaSingleton.get().user.role
-            if role not in roles:
+            client = args[0].__client if hasattr(args[0], "__client") else active_client().http_client.httpx
+            user = users_api.whoami_httpx(client).parsed
+            if user.role not in roles:
                 raise PermissionError(
-                    f"User with role={role} is not allowed to call `{func.__name__}`."
+                    f"User with role={user.role} is not allowed to call `{func.__name__}`."
                     f" Only users with role={roles} are allowed to call this function."
                 )
             return func(*args, **kwargs)

src/argilla/client/workspaces.py

@@ -149,7 +149,7 @@ def add_user(self, user_id: str) -> None:
         except BaseClientError as e:
             raise RuntimeError(f"Error while adding user with id=`{user_id}` to workspace with id=`{self.id}`.") from e
 
-    @allowed_for_roles(roles=[UserRole.owner, UserRole.admin])
+    @allowed_for_roles(roles=[UserRole.owner])
     def delete_user(self, user_id: str) -> None:
         """Deletes an existing user from the workspace in Argilla. Note that the user
         will not be deleted from Argilla, but just from the workspace.
@@ -285,7 +285,7 @@ def from_name(cls, name: str) -> "Workspace":
         """
         client = cls.__active_client()
         try:
-            workspaces = workspaces_api.list_workspaces(client).parsed
+            workspaces = workspaces_api_v1.list_workspaces_me(client).parsed
         except Exception as e:
             raise RuntimeError("Error while retrieving the list of workspaces from Argilla.") from e
 
@@ -315,7 +315,7 @@ def list(cls) -> Iterator["Workspace"]:
         """
         client = cls.__active_client()
         try:
-            workspaces = workspaces_api.list_workspaces(client).parsed
+            workspaces = workspaces_api_v1.list_workspaces_me(client).parsed
             for ws in workspaces:
                 yield cls.__new_instance(client, ws)
         except Exception as e:

src/argilla/server/apis/v1/handlers/users.py

@@ -19,7 +19,7 @@
 
 from argilla.server.contexts import accounts
 from argilla.server.database import get_async_db
-from argilla.server.models import User, UserRole
+from argilla.server.models import User
 from argilla.server.policies import UserPolicyV1, authorize
 from argilla.server.schemas.v1.workspaces import Workspaces
 from argilla.server.security import auth

src/argilla/server/apis/v1/handlers/workspaces.py

@@ -19,10 +19,10 @@
 
 from argilla.server.contexts import accounts
 from argilla.server.database import get_async_db
+from argilla.server.models import User
 from argilla.server.policies import WorkspacePolicyV1, authorize
-from argilla.server.schemas.v1.workspaces import Workspace
+from argilla.server.schemas.v1.workspaces import Workspace, Workspaces
 from argilla.server.security import auth
-from argilla.server.security.model import User
 
 router = APIRouter(tags=["workspaces"])
 
@@ -44,3 +44,19 @@ async def get_workspace(
         )
 
     return workspace
+
+
+@router.get("/me/workspaces", response_model=Workspaces)
+async def list_workspaces_me(
+    *,
+    db: AsyncSession = Depends(get_async_db),
+    current_user: User = Security(auth.get_current_user),
+) -> Workspaces:
+    await authorize(current_user, WorkspacePolicyV1.list_workspaces_me)
+
+    if current_user.is_owner:
+        workspaces = await accounts.list_workspaces(db)
+    else:
+        workspaces = await accounts.list_workspaces_by_user_id(db, current_user.id)
+
+    return Workspaces(items=workspaces)

src/argilla/server/policies.py

@@ -99,6 +99,10 @@ async def is_allowed(actor: User) -> bool:
 
         return is_allowed
 
+    @classmethod
+    async def list_workspaces_me(cls, actor: User) -> bool:
+        return True
+
 
 class UserPolicy:
     @classmethod

tests/client/sdk/v1/users/test_api.py

@@ -13,7 +13,6 @@
 #  limitations under the License.
 
 from typing import TYPE_CHECKING
-from uuid import uuid4
 
 import pytest
 from argilla.client.api import ArgillaSingleton

tests/client/sdk/v1/workspaces/__init__.py

@@ -0,0 +1,13 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.

tests/client/sdk/v1/workspaces/test_api.py

@@ -0,0 +1,51 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import pytest
+from argilla.client.api import ArgillaSingleton
+from argilla.client.sdk.users.models import UserRole
+from argilla.client.sdk.v1.workspaces.api import get_workspace, list_workspaces_me
+from argilla.client.sdk.v1.workspaces.models import WorkspaceModel
+
+from tests.factories import UserFactory, WorkspaceFactory
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("role", [UserRole.admin, UserRole.owner])
+async def test_get_workspace(role: UserRole) -> None:
+    workspace = await WorkspaceFactory.create()
+    user = await UserFactory.create(role=role, workspaces=[workspace])
+
+    httpx_client = ArgillaSingleton.init(api_key=user.api_key).http_client.httpx
+
+    response = get_workspace(client=httpx_client, id=workspace.id)
+    assert response.status_code == 200
+    assert isinstance(response.parsed, WorkspaceModel)
+    assert response.parsed.id == workspace.id
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("role", [UserRole.owner, UserRole.admin, UserRole.annotator])
+async def test_list_workspaces_me(role: UserRole) -> None:
+    workspaces = await WorkspaceFactory.create_batch(size=5)
+    user = await UserFactory.create(role=role, workspaces=workspaces if role != UserRole.owner else [])
+
+    httpx_client = ArgillaSingleton.init(api_key=user.api_key).http_client.httpx
+
+    response = list_workspaces_me(client=httpx_client)
+    assert response.status_code == 200
+    assert isinstance(response.parsed, list)
+    assert len(response.parsed) > 0
+    assert isinstance(response.parsed[0], WorkspaceModel)
+    assert len(response.parsed) == len(workspaces)

tests/client/sdk/v1/workspaces/test_models.py

@@ -0,0 +1,20 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from argilla.client.sdk.v1.workspaces.models import WorkspaceModel as ClientSchema
+from argilla.server.schemas.v1.workspaces import Workspace as ServerSchema
+
+
+def test_workspace_schema(helpers) -> None:
+    assert helpers.are_compatible_api_schemas(ClientSchema.schema(), ServerSchema.schema())