feat: Allow to delete workspaces (#3358)

<!-- Thanks for your contribution! As part of our Community Growers
initiative 🌱, we're donating Justdiggit bunds in your name to reforest
sub-Saharan Africa. To claim your Community Growers certificate, please
contact David Berenstein in our Slack community or fill in this form
https://tally.so/r/n9XrxK once your PR has been merged. -->

# Description

This PR includes the missing functionality for deleting workspaces. A
workspace can be deleted only if no datasets are linked to it.
Otherwise, the operation will raise an error.


Closes #3260 

**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)
- [X] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Refactor (change restructuring the codebase without changing
functionality)
- [ ] Improvement (change adding some improvement to an existing
functionality)
- [ ] Documentation update



**Checklist**

- [ ] I added relevant documentation
- [x] follows the style guidelines of this project
- [x] I did a self-review of my code
- [ ] I made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [x] 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)
- [x] I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

---------

Co-authored-by: Alvaro Bartolome <[email protected]>

CHANGELOG.md

@@ -22,9 +22,10 @@ These are the section headers that we use:
 - Added `HuggingFaceDatasetMixin` for internal usage, to detach the `FeedbackDataset` integrations from the class itself, and use Mixins instead ([#3326](https://github.com/argilla-io/argilla/pull/3326)).
 - Added `GET /api/v1/records/{record_id}/suggestions` API endpoint to get the list of suggestions for the responses associated to a record ([#3304](https://github.com/argilla-io/argilla/pull/3304)).
 - Added `PUT /api/v1/records/{record_id}/suggestions` API endpoint to create or update a suggestion for a response associated to a record ([#3304](https://github.com/argilla-io/argilla/pull/3304) & [3391](https://github.com/argilla-io/argilla/pull/3391)).
-- Added breaking simutaneously running tests within GitHub package worflows. ([#3354](https://github.com/argilla-io/argilla/pull/3354)).
+- Added breaking simultaneously running tests within GitHub package workflows. ([#3354](https://github.com/argilla-io/argilla/pull/3354)).
 - Added `suggestions` attribute to `FeedbackRecord`, and allow adding and retrieving suggestions from the Python client ([#3370](https://github.com/argilla-io/argilla/pull/3370))
 - 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 API and Python Client support for workspace deletion (Closes [#3260](https://github.com/argilla-io/argilla/issues/3260))
 - 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

docs/_source/getting_started/installation/configurations/workspace_management.md

@@ -125,3 +125,23 @@ for user in users:
 workspace.add_user("<USER_ID>")
 workspace.delete_user("<USER_ID>")
 ```
+
+### Delete a `Workspace`
+
+#### Python client
+
+You can also delete a workspace using the python client.
+
+:::{note}
+To delete a workspace, no dataset can be linked to it. If workspace contains any dataset, deletion will fail.
+:::
+
+```python
+import argilla as rg
+
+rg.init(api_url="<ARGILLA_API_URL>", api_key="<ARGILLA_API_KEY>")
+
+workspace = rg.Workspace.from_name("new-workspace")
+
+workspace.delete()
+```

src/argilla/client/sdk/commons/errors.py

@@ -46,7 +46,7 @@ def __str__(self):
 
 
 class HttpResponseError(BaseClientError):
-    """Used for handle http errros other than defined in Argilla server"""
+    """Used for handle http errors other than defined in Argilla server"""
 
     def __init__(self, response: httpx.Response):
         self.status_code = response.status_code

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

@@ -55,6 +55,24 @@ def get_workspace(
     return handle_response_error(response)
 
 
+def delete_workspace(
+    client: httpx.Client, id: UUID
+) -> Response[Union[WorkspaceModel, ErrorMessage, HTTPValidationError]]:
+    url = f"/api/v1/workspaces/{id}"
+
+    response = client.delete(url=url)
+
+    if response.status_code == 200:
+        parsed_response = WorkspaceModel(**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_workspaces_me(
     client: httpx.Client,
 ) -> Response[Union[List[WorkspaceModel], ErrorMessage, HTTPValidationError]]:

src/argilla/client/workspaces.py

@@ -123,7 +123,7 @@ def __repr__(self) -> str:
         )
 
     @allowed_for_roles(roles=[UserRole.owner])
-    def add_user(self, user_id: str) -> None:
+    def add_user(self, user_id: UUID) -> None:
         """Adds an existing user to the workspace in Argilla.
 
         Args:
@@ -150,7 +150,7 @@ def add_user(self, user_id: str) -> None:
             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])
-    def delete_user(self, user_id: str) -> None:
+    def delete_user(self, user_id: UUID) -> 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.
 
@@ -182,6 +182,32 @@ def delete_user(self, user_id: str) -> None:
                 f"Error while deleting user with id=`{user_id}` from workspace with id=`{self.id}`."
             ) from e
 
+    @allowed_for_roles(roles=[UserRole.owner])
+    def delete(self) -> None:
+        """Deletes an existing workspace from Argilla. Note that the workspace
+        cannot have any linked dataset to be removed from Argilla. Otherwise an error will be raised.
+
+        Raises:
+            ValueError: if the workspace does not exists or some datasets are linked to it.
+            RuntimeError: if there was an unexpected error while deleting the user from the workspace.
+
+        Examples:
+            >>> from argilla import rg
+            >>> workspace = rg.Workspace.from_name("my-workspace")
+            >>> workspace.delete()
+        """
+        try:
+            workspaces_api_v1.delete_workspace(client=self.__client, id=self.id)
+        except NotFoundApiError as e:
+            raise ValueError(f"Workspace with id {self.id} doesn't exist in Argilla.") from e
+        except AlreadyExistsApiError as e:
+            # TODO: the already exists is to explicit for this context and should be generalized
+            raise ValueError(
+                f"Cannot delete workspace with id {self.id}. Some datasets are still linked to this workspace."
+            ) from e
+        except BaseClientError as e:
+            raise RuntimeError(f"Error while deleting workspace with id {self.id!r}.") from e
+
     @staticmethod
     def __active_client() -> "httpx.Client":
         """Returns the active Argilla `httpx.Client` instance."""

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

@@ -62,27 +62,6 @@ async def create_workspace(
     return Workspace.from_orm(workspace)
 
 
-# TODO: We can't do workspaces deletions right now. Workspaces are associated with datasets and used on
-# ElasticSearch indexes. Once that we have datasets on the database and we can check if the workspace doesn't have
-# any dataset then we can delete them.
-# @router.delete("/workspaces/{workspace_id}", response_model=Workspace, response_model_exclude_none=True)
-async def delete_workspace(
-    *,
-    db: AsyncSession = Depends(get_async_db),
-    workspace_id: UUID,
-    current_user: User = Security(auth.get_current_user),
-):
-    workspace = await accounts.get_workspace_by_id(db, workspace_id)
-    if not workspace:
-        raise EntityNotFoundError(name=str(workspace_id), type=Workspace)
-
-    await authorize(current_user, WorkspacePolicy.delete(workspace))
-
-    await accounts.delete_workspace(db, workspace)
-
-    return Workspace.from_orm(workspace)
-
-
 @router.get("/workspaces/{workspace_id}/users", response_model=List[User], response_model_exclude_none=True)
 async def list_workspace_users(
     *,

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

@@ -17,12 +17,13 @@
 from fastapi import APIRouter, Depends, HTTPException, Security, status
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from argilla.server.contexts import accounts
+from argilla.server.contexts import accounts, datasets
 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, Workspaces
 from argilla.server.security import auth
+from argilla.server.services.datasets import DatasetsService
 
 router = APIRouter(tags=["workspaces"])
 
@@ -46,6 +47,38 @@ async def get_workspace(
     return workspace
 
 
+@router.delete("/workspaces/{workspace_id}", response_model=Workspace)
+async def delete_workspace(
+    *,
+    db: AsyncSession = Depends(get_async_db),
+    datasets_service: DatasetsService = Depends(DatasetsService.get_instance),
+    workspace_id: UUID,
+    current_user: User = Security(auth.get_current_user),
+):
+    await authorize(current_user, WorkspacePolicyV1.delete)
+
+    workspace = await accounts.get_workspace_by_id(db, workspace_id)
+    if not workspace:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Workspace with id `{workspace_id}` not found",
+        )
+
+    if await datasets.list_datasets_by_workspace_id(db, workspace_id):
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail=f"Cannot delete the workspace {workspace_id}. This workspace has some feedback datasets linked",
+        )
+
+    if await datasets_service.list(current_user, workspaces=[workspace.name]):
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail=f"Cannot delete the workspace {workspace_id}. This workspace has some datasets linked",
+        )
+
+    return await accounts.delete_workspace(db, workspace)
+
+
 @router.get("/me/workspaces", response_model=Workspaces)
 async def list_workspaces_me(
     *,

src/argilla/server/contexts/datasets.py

@@ -80,7 +80,14 @@ async def list_datasets(db: "AsyncSession") -> List[Dataset]:
     return result.scalars().all()
 
 
-async def create_dataset(db: "AsyncSession", dataset_create: DatasetCreate) -> Dataset:
+async def list_datasets_by_workspace_id(db: "AsyncSession", workspace_id: UUID) -> List[Dataset]:
+    result = await db.execute(
+        select(Dataset).where(Dataset.workspace_id == workspace_id).order_by(Dataset.inserted_at.asc())
+    )
+    return result.scalars().all()
+
+
+async def create_dataset(db: "AsyncSession", dataset_create: DatasetCreate):
     return await Dataset.create(
         db,
         name=dataset_create.name,

src/argilla/server/policies.py

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

tests/client/test_workspaces.py

@@ -21,7 +21,12 @@
 from argilla.client.sdk.workspaces.models import WorkspaceUserModel
 from argilla.client.workspaces import Workspace
 
-from tests.factories import UserFactory, WorkspaceFactory, WorkspaceUserFactory
+from tests.factories import (
+    DatasetFactory,
+    UserFactory,
+    WorkspaceFactory,
+    WorkspaceUserFactory,
+)
 
 if TYPE_CHECKING:
     from argilla.server.models import User as ServerUser
@@ -197,3 +202,86 @@ async def test_print_workspace(owner: "ServerUser"):
         f"Workspace(id={workspace.id}, name={workspace.name}, "
         f"inserted_at={workspace.inserted_at}, updated_at={workspace.updated_at})"
     )
+
+
+def test_set_new_workspace(owner: "ServerUser"):
+    ArgillaSingleton.init(api_key=owner.api_key)
+    ws = Workspace.create("new-workspace")
+
+    ArgillaSingleton.get().set_workspace(ws.name)
+    assert ArgillaSingleton.get().get_workspace() == ws.name
+
+
+@pytest.mark.asyncio
+async def test_init_with_workspace(owner: "ServerUser"):
+    workspace = await WorkspaceFactory.create(name="test_workspace")
+
+    ArgillaSingleton.init(api_key=owner.api_key, workspace=workspace.name)
+
+    assert ArgillaSingleton.get().get_workspace() == workspace.name
+
+
+def test_set_workspace_with_missing_workspace(owner: "ServerUser"):
+    ArgillaSingleton.init(api_key=owner.api_key)
+    with pytest.raises(ValueError):
+        ArgillaSingleton.get().set_workspace("missing-workspace")
+
+
+def test_init_with_missing_workspace(owner: "ServerUser"):
+    with pytest.raises(ValueError):
+        ArgillaSingleton.init(api_key=owner.api_key, workspace="missing-workspace")
+
+
+@pytest.mark.asyncio
+async def test_delete_workspace(owner: "ServerUser"):
+    workspace = await WorkspaceFactory.create(name="test_workspace")
+
+    ArgillaSingleton.init(api_key=owner.api_key)
+
+    ws = Workspace.from_id(workspace.id)
+    ws.delete()
+
+    with pytest.raises(ValueError, match=rf"Workspace with id=`{ws.id}` doesn't exist in Argilla"):
+        Workspace.from_id(workspace.id)
+
+
+@pytest.mark.asyncio
+async def test_delete_non_existing_workspace(owner: "ServerUser"):
+    workspace = await WorkspaceFactory.create(name="test_workspace")
+
+    ArgillaSingleton.init(api_key=owner.api_key)
+
+    ws = Workspace.from_id(workspace.id)
+    ws.delete()
+
+    with pytest.raises(ValueError, match=rf"Workspace with id {ws.id} doesn't exist in Argilla."):
+        ws.delete()
+
+
+@pytest.mark.asyncio
+async def test_delete_workspace_with_linked_datasets(owner: "ServerUser"):
+    workspace = await WorkspaceFactory.create(name="test_workspace")
+    await DatasetFactory.create(workspace=workspace)
+
+    ArgillaSingleton.init(api_key=owner.api_key)
+
+    ws = Workspace.from_id(workspace.id)
+    with pytest.raises(
+        ValueError,
+        match=rf"Cannot delete workspace with id {ws.id}. Some datasets are still linked to this workspace.",
+    ):
+        ws.delete()
+
+
+@pytest.mark.asyncio
+async def test_delete_workspace_without_permissions():
+    workspace = await WorkspaceFactory.create(name="test_workspace")
+
+    user = await UserFactory.create(workspaces=[workspace])
+
+    ArgillaSingleton.init(api_key=user.api_key)
+
+    ws = Workspace.from_id(workspace.id)
+
+    with pytest.raises(PermissionError, match=rf"User with role={user.role.value} is not allowed to call `delete`"):
+        ws.delete()

tests/conftest.py

@@ -14,7 +14,7 @@
 import asyncio
 import contextlib
 import tempfile
-from typing import TYPE_CHECKING, AsyncGenerator, Dict, Generator
+from typing import TYPE_CHECKING, AsyncGenerator, Dict, Generator, Iterator
 
 import httpx
 import pytest
@@ -162,8 +162,8 @@ def elasticsearch_config():
     return {"hosts": settings.elasticsearch}
 
 
-@pytest.fixture(scope="session")
-def opensearch(elasticsearch_config):
+@pytest.fixture(scope="session", autouse=True)
+def opensearch(elasticsearch_config) -> Generator[OpenSearch, None, None]:
     client = OpenSearch(**elasticsearch_config)
     yield client