feat(experiment): Enable deleting experiments by their alias (#20)

* add find by tags, datetime to pyaqueduct. * add pydantic validator calls for the userfacing APIs. * fix docstring being too long. * Update examples/tutorial_notebook.ipynb Co-authored-by: Stanislav Protasov <stanislav.protasov@gmail.com> * add remove experiment functionality with unit tests and updated tutorial. * rename method to . --------- Co-authored-by: Stanislav Protasov <stanislav.protasov@gmail.com>
AqueductHub · May 10, 2024 · 4449343 · 4449343 · github-actions · May 10, 2024
1 parent 65adae1
commit 4449343
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 28 deletions.
diff --git a/examples/tutorial_notebook.ipynb b/examples/tutorial_notebook.ipynb
@@ -312,6 +312,24 @@
     "# In this example we only retrieve maximum 20 experimets starting from the 10th experiment (inclusive) with the specified tags.\n",
     "experiments_list = api.find_experiments(limit=20, offset=10, tags=[\"tag1\", \"tag2\"])"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Removing Experiments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Experiments and their files can be removed from the database. Use it with caution\n",
+    "# as this operation is not revertible.\n",
+    "api.remove_experiment(alias=\"[experiment alias]\")"
+   ]
   }
  ],
  "metadata": {

diff --git a/pyaqueduct/api.py b/pyaqueduct/api.py
@@ -88,6 +88,17 @@ def get_experiment(self, alias: str) -> Experiment:
             created_at=experiment_data.created_at,
         )
 
+    @validate_call
+    def remove_experiment_by_alias(self, alias: str) -> None:
+        """Remove experiment from the database. Experiment's files will be also removed.
+
+        Args:
+            alias: Alias of the specified experiment.
+
+        """
+        experiment_data = self._client.get_experiment_by_alias(alias=alias)
+        self._client.remove_experiment(experiment_uuid=experiment_data.id)
+
     @validate_call
     def find_experiments(
         self,

diff --git a/pyaqueduct/client/client.py b/pyaqueduct/client/client.py
@@ -25,6 +25,7 @@
 from pyaqueduct.schemas.mutations import (
     add_tag_to_experiment_mutation,
     create_experiment_mutation,
+    remove_experiment_mutation,
     remove_tag_from_experiment_mutation,
     update_experiment_mutation,
 )
@@ -84,15 +85,15 @@ def create_experiment(
         self, title: str, description: str, tags: Optional[List[str]] = None
     ) -> ExperimentData:
         """
-        Create an experiment with title, description and list of tags
+        Create an experiment with title, description and list of tags.
 
         Args:
-        - title (str): Title of experiment
-        - description (str): Description of experiment
-        - tags (List[str]): List of tags to be assigned to experiment
+            title: Title of experiment.
+            description: Description of experiment.
+            tags: List of tags to be assigned to experiment.
 
         Returns:
-        Experiment: Experiment objects having all fields
+            Experiment object.
         """
         try:
             experiment = self._gql_client.execute(
@@ -118,15 +119,16 @@ def update_experiment(
         self, experiment_uuid: UUID, title: Optional[str] = None, description: Optional[str] = None
     ) -> ExperimentData:
         """
-        Update title or description or both for experiment
+        Update title or description or both for experiment.
 
         Args:
-        - experiment_uuid (UUID): UUID of experiment to be updated
-        - title (str): New title of experiment
-        - description (str): New description of experiment
+            experiment_uuid: UUID of experiment.
+            title: New title of experiment.
+            description: New description of experiment.
 
         Returns:
-        Experiment: Experiment object with updated fields
+            Experiment object.
+
         """
         if title and title == "":
             raise ValueError("Title cannot be an empty string")
@@ -171,15 +173,15 @@ def get_experiments(
         Get a list of experiments
 
         Args:
-        - limit: Pagination field, number of experiments to be fetched.
-        - offset: Pagination field, number of experiments to skip.
-        - title: Perform search on experiments through their title and alias.
-        - tags: Get experiments that have these tags.
-        - start_date: Start datetime to filter experiments (timezone aware).
-        - end_date: End datetime to filter experiments to (timezone aware).
+            limit: Pagination field, number of experiments to be fetched.
+            offset: Pagination field, number of experiments to skip.
+            title: Perform search on experiments through their title and alias.
+            tags: Get experiments that have these tags.
+            start_date: Start datetime to filter experiments (timezone aware).
+            end_date: End datetime to filter experiments to (timezone aware).
 
         Returns:
-            List of experiments with filters applied
+            List of experiments with filters applied.
 
         """
         if limit <= 0:
@@ -218,14 +220,14 @@ def get_experiments(
 
     def get_experiment(self, experiment_uuid: UUID) -> ExperimentData:
         """
-        Get an Experiment by ID or Alias
+        Get an Experiment by UUID.
 
         Args:
-        - id_ (str): "UUID | ALIAS" Unique identifier to be used
-        - value (str): Value of unique identifier to use for fetching experiment
+            experiment_uuid: UUID of experiment.
 
         Returns:
-        Experiment: Experiment object having all fields
+            Experiment object.
+
         """
         try:
             experiment = self._gql_client.execute(
@@ -249,14 +251,14 @@ def get_experiment(self, experiment_uuid: UUID) -> ExperimentData:
 
     def get_experiment_by_alias(self, alias: str) -> ExperimentData:
         """
-        Get an Experiment by ID or Alias
+        Get an Experiment by Alias.
 
         Args:
-        - id_ (str): "UUID | ALIAS" Unique identifier to be used
-        - value (str): Value of unique identifier to use for fetching experiment
+            alias: Experiment's alias.
 
         Returns:
-        Experiment: Experiment object having all fields
+            Updated experiment.
+
         """
         try:
             experiment = self._gql_client.execute(
@@ -282,11 +284,12 @@ def add_tag_to_experiment(self, experiment_uuid: UUID, tag: str) -> ExperimentDa
         Add a tag to an experiment
 
         Args:
-        - experiment_uuid (UUID): UUID of experiment to which tag is to be added
-        - tag (str): Tag to be added to experiment
+            experiment_uuid: UUID of experiment.
+            tag: Tag to be added to experiment.
 
         Returns:
-        Experiment: Experiment having tag added
+            Updated experiment.
+
         """
         try:
             experiment = self._gql_client.execute(
@@ -308,6 +311,28 @@ def add_tag_to_experiment(self, experiment_uuid: UUID, tag: str) -> ExperimentDa
         logging.info("Added tag %s to experiment <%s>", tag, experiment_obj.title)
         return experiment_obj
 
+    def remove_experiment(self, experiment_uuid: UUID) -> None:
+        """
+        Remove experiment from the database. It removes the experiments files as well.
+
+        Args:
+            experiment_uuid: UUID of experiment.
+
+        """
+        try:
+            self._gql_client.execute(
+                remove_experiment_mutation,
+                variable_values={"experimentId": str(experiment_uuid)},
+            )
+        except gql_exceptions.TransportServerError as error:
+            if error.code:
+                process_response_common(codes(error.code))
+            raise
+        except gql_exceptions.TransportQueryError as error:
+            raise RemoteOperationError(
+                error.errors if error.errors else "Unknown error occurred in the remote operation."
+            ) from error
+
     def remove_tag_from_experiment(self, experiment_uuid: UUID, tag: str) -> ExperimentData:
         """
         Remove a tag from an experiment

diff --git a/pyaqueduct/schemas/mutations.py b/pyaqueduct/schemas/mutations.py
@@ -1,4 +1,5 @@
 """Aqueduct GraphQL Mutation schemas"""
+
 from gql import gql
 
 create_experiment_mutation = gql(
@@ -120,3 +121,13 @@
     }
     """
 )
+
+remove_experiment_mutation = gql(
+    """
+    mutation RemoveTagFromExperiment (
+        $experimentId: UUID!
+        ) {
+        removeExperiment(experimentRemoveInput: {experimentId: $experimentId})
+    }
+    """
+)
diff --git a/tests/unittests/test_api.py b/tests/unittests/test_api.py
@@ -64,6 +64,35 @@ def patched_get_experiment(self, alias):
     assert experiment.created_at == expected_datetime
 
 
+def test_remove_experiment_by_alias(monkeypatch):
+    expected_id = uuid4()
+    expected_title = "test title"
+    expected_description = "test description"
+    expected_alias = "mock_alias"
+    expected_datetime = datetime.now()
+
+    def patched_remove_experiment(self, experiment_uuid):
+        assert experiment_uuid == expected_id
+
+    def patched_get_experiment(self, alias):
+        assert alias == expected_alias
+
+        return ExperimentData(
+            id=expected_id,
+            title=expected_title,
+            description=expected_description,
+            alias=alias,
+            created_at=expected_datetime,
+            updated_at=expected_datetime,
+        )
+
+    monkeypatch.setattr(AqueductClient, "get_experiment_by_alias", patched_get_experiment)
+    monkeypatch.setattr(AqueductClient, "remove_experiment", patched_remove_experiment)
+    api = API(url="http://test.com", timeout=1)
+
+    api.remove_experiment_by_alias(alias=expected_alias)
+
+
 def test_find_experiments(monkeypatch):
     ExperimentMockData = namedtuple(
         "ExperimentData",

diff --git a/tests/unittests/test_client.py b/tests/unittests/test_client.py
@@ -145,6 +145,20 @@ def test_add_tag_to_experiment(monkeypatch):
     assert tag_name in experiment.tags
 
 
+def test_remove_experiment(monkeypatch):
+    experiment_id = uuid4()
+
+    def patched_execute(self, query, variable_values, **kwargs):
+        if variable_values["experimentId"]:
+            assert variable_values["experimentId"] == str(experiment_id)
+
+    monkeypatch.setattr(SyncClientSession, "execute", patched_execute)
+
+    client = AqueductClient(url="http://test.com", timeout=1)
+
+    client.remove_experiment(experiment_uuid=experiment_id)
+
+
 def test_remove_tag_from_experiment(monkeypatch):
     experiment_id = uuid4()
     tag_name = "tag"
File	Stmts	Miss	Cover	Missing
pyaqueduct
experiment.py	62	1	98%	85
pyaqueduct/client
client.py	172	72	58%	44–50, 103–108, 134, 137, 148–153, 188, 202–207, 237–242, 263–280, 299–304, 327–332, 352–357, 380, 387–392, 421–422, 459–460
types.py	39	1	97%	20
TOTAL	333	74	78%