feat(Datasets): restrict dataset deletion only to creators and super-…

…users (#1713)

* feat(Datasets): restrict dataset deletion only to creators and super-users

* tests: add tests for datasets deletion check

* chore: improve deletion error message

Co-authored-by: Daniel Vila Suero <daniel@recogn.ai>

* Format long strings

Co-authored-by: Daniel Vila Suero <daniel@recogn.ai>
(cherry picked from commit c5ab270)

src/rubrix/server/services/datasets.py

@@ -141,7 +141,13 @@ def delete(self, user: User, dataset: ServiceDataset):
             as_dataset_class=None,
         )
         if found:
-            self.__dao__.delete_dataset(dataset)
+            if user.is_superuser() or user.username == dataset.created_by:
+                self.__dao__.delete_dataset(dataset)
+            else:
+                raise ForbiddenOperationError(
+                    f"You don't have the necessary permissions to delete this dataset. "
+                    "Only dataset creators or administrators can delete datasets"
+                )
 
     def update(
         self,

tests/datasets/test_datasets.py

@@ -3,6 +3,7 @@
 import rubrix as rb
 from rubrix import TextClassificationSettings, TokenClassificationSettings
 from rubrix.client import api
+from rubrix.client.sdk.commons.errors import ForbiddenApiError
 
 
 @pytest.mark.parametrize(
@@ -37,3 +38,17 @@ def test_settings_workflow(mocked_client, settings_, wrong_settings):
 
     with pytest.raises(ValueError, match="Task type mismatch"):
         rb.configure_dataset(dataset, wrong_settings)
+
+
+def test_delete_dataset_by_non_creator(mocked_client):
+    try:
+        dataset = "test_delete_dataset_by_non_creator"
+        rb.delete(dataset)
+        rb.configure_dataset(
+            dataset, settings=TextClassificationSettings(label_schema={"A", "B", "C"})
+        )
+        mocked_client.change_current_user("mock-user")
+        with pytest.raises(ForbiddenApiError):
+            rb.delete(dataset)
+    finally:
+        mocked_client.reset_default_user()

tests/helpers.py

@@ -3,21 +3,46 @@
 from fastapi import FastAPI
 from starlette.testclient import TestClient
 
-from rubrix._constants import API_KEY_HEADER_NAME
+from rubrix._constants import API_KEY_HEADER_NAME, RUBRIX_WORKSPACE_HEADER_NAME
 from rubrix.client.api import active_api
 from rubrix.server.security import auth
 from rubrix.server.security.auth_provider.local.settings import settings
+from rubrix.server.security.auth_provider.local.users.model import UserInDB
 
 
 class SecuredClient:
     def __init__(self, client: TestClient):
         self._client = client
         self._header = {API_KEY_HEADER_NAME: settings.default_apikey}
+        self._current_user = None
 
     @property
     def fastpi_app(self) -> FastAPI:
         return self._client.app
 
+    def change_current_user(self, username):
+        default_user = auth.users.__dao__.__users__["rubrix"]
+        new_user = UserInDB(
+            username=username,
+            hashed_password=username,  # Even if required, we can ignore it
+            api_key=username,
+            workspaces=["rubrix"],  # The default workspace
+        )
+
+        auth.users.__dao__.__users__[username] = new_user
+        rb_api = active_api()
+        rb_api._user = new_user
+        rb_api.set_workspace(default_user.username)
+        rb_api.client.token = new_user.api_key
+
+    def reset_default_user(self):
+        default_user = auth.users.__dao__.__users__["rubrix"]
+
+        rb_api = active_api()
+        rb_api._user = default_user
+        rb_api.client.token = default_user.api_key
+        rb_api.client.headers.pop(RUBRIX_WORKSPACE_HEADER_NAME)
+
     def add_workspaces_to_rubrix_user(self, workspaces: List[str]):
         rubrix_user = auth.users.__dao__.__users__["rubrix"]
         rubrix_user.workspaces.extend(workspaces or [])