From eedc00358f5dd9c46b9aa759d503ee2c2218a849 Mon Sep 17 00:00:00 2001
From: Charles Zaloom <38677807+czaloom@users.noreply.github.com>
Date: Wed, 10 Jul 2024 10:48:30 -0500
Subject: [PATCH] Allow Retrying of Deletion (#664)

---
 api/valor_api/backend/core/dataset.py   | 25 ++++++++---
 api/valor_api/backend/core/model.py     | 21 +++++----
 integration_tests/client/test_client.py | 59 ++++++++++++++++++++++++-
 3 files changed, 88 insertions(+), 17 deletions(-)

diff --git a/api/valor_api/backend/core/dataset.py b/api/valor_api/backend/core/dataset.py
index b418fc581..779a516c0 100644
--- a/api/valor_api/backend/core/dataset.py
+++ b/api/valor_api/backend/core/dataset.py
@@ -505,16 +505,27 @@ def delete_dataset(
     name : str
         The name of the dataset.
     """
-    dataset = fetch_dataset(db, name=name)
-    set_dataset_status(db, name, enums.TableStatus.DELETING)
+    if core.count_active_evaluations(db=db, dataset_names=[name]):
+        raise exceptions.EvaluationRunningError(dataset_name=name)
 
-    core.delete_evaluations(db=db, dataset_names=[name])
-    core.delete_dataset_predictions(db, dataset)
-    core.delete_groundtruths(db, dataset)
-    core.delete_dataset_annotations(db, dataset)
-    core.delete_datums(db, dataset)
+    dataset = (
+        db.query(models.Dataset)
+        .where(models.Dataset.name == name)
+        .one_or_none()
+    )
+    if not dataset:
+        raise exceptions.DatasetDoesNotExistError(name)
 
     try:
+        dataset.status = enums.TableStatus.DELETING
+        db.commit()
+
+        core.delete_evaluations(db=db, dataset_names=[name])
+        core.delete_dataset_predictions(db, dataset)
+        core.delete_groundtruths(db, dataset)
+        core.delete_dataset_annotations(db, dataset)
+        core.delete_datums(db, dataset)
+
         db.execute(
             delete(models.Dataset).where(models.Dataset.id == dataset.id)
         )
diff --git a/api/valor_api/backend/core/model.py b/api/valor_api/backend/core/model.py
index d7a415b00..6fee3b5d8 100644
--- a/api/valor_api/backend/core/model.py
+++ b/api/valor_api/backend/core/model.py
@@ -377,21 +377,24 @@ def delete_model(
     name : str
         The name of the model.
     """
-    model = fetch_model(db, name=name)
+    if core.count_active_evaluations(db=db, model_names=[name]):
+        raise exceptions.EvaluationRunningError(model_name=name)
+
+    model = (
+        db.query(models.Model).where(models.Model.name == name).one_or_none()
+    )
+    if not model:
+        raise exceptions.ModelDoesNotExistError(name)
 
-    # set status
     try:
+        # set status
         model.status = ModelStatus.DELETING
         db.commit()
-    except Exception as e:
-        db.rollback()
-        raise e
 
-    core.delete_evaluations(db=db, model_names=[name])
-    core.delete_model_predictions(db=db, model=model)
-    core.delete_model_annotations(db=db, model=model)
+        core.delete_evaluations(db=db, model_names=[name])
+        core.delete_model_predictions(db=db, model=model)
+        core.delete_model_annotations(db=db, model=model)
 
-    try:
         db.execute(delete(models.Model).where(models.Model.id == model.id))
         db.commit()
     except IntegrityError as e:
diff --git a/integration_tests/client/test_client.py b/integration_tests/client/test_client.py
index e40f57826..8f2c9c111 100644
--- a/integration_tests/client/test_client.py
+++ b/integration_tests/client/test_client.py
@@ -6,6 +6,7 @@
 
 import pytest
 import requests
+from sqlalchemy.orm import Session
 
 from valor import (
     Annotation,
@@ -18,7 +19,11 @@
     Prediction,
 )
 from valor.client import connect
-from valor.exceptions import ClientException
+from valor.exceptions import (
+    ClientException,
+    DatasetAlreadyExistsError,
+    ModelAlreadyExistsError,
+)
 from valor.schemas import And, Filter
 
 
@@ -301,3 +306,55 @@ def test_get_datums(
     requests_method = getattr(requests, "get")
     resp = requests_method("http://localhost:8000/data")
     assert resp.headers["content-range"] == "items 0-0/1"
+
+
+def test_delete_zombie_dataset(
+    db: Session, client: Client, created_dataset: Dataset
+):
+
+    dataset_name = created_dataset.name
+    assert isinstance(dataset_name, str)
+
+    # set deletion status in row (this simulates the zombie deletion)
+    from valor_api import enums as backend_enums
+    from valor_api.backend import models
+
+    row = (
+        db.query(models.Dataset)
+        .where(models.Dataset.name == dataset_name)
+        .one_or_none()
+    )
+    assert row
+    row.status = backend_enums.TableStatus.DELETING
+    db.commit()
+
+    with pytest.raises(DatasetAlreadyExistsError):
+        Dataset.create(name=dataset_name)
+
+    client.delete_dataset(name=dataset_name)
+
+
+def test_delete_zombie_model(
+    db: Session, client: Client, created_model: Model
+):
+
+    model_name = created_model.name
+    assert isinstance(model_name, str)
+
+    # set deletion status in row (this simulates the zombie deletion)
+    from valor_api import enums as backend_enums
+    from valor_api.backend import models
+
+    row = (
+        db.query(models.Model)
+        .where(models.Model.name == model_name)
+        .one_or_none()
+    )
+    assert row
+    row.status = backend_enums.ModelStatus.DELETING
+    db.commit()
+
+    with pytest.raises(ModelAlreadyExistsError):
+        Model.create(name=model_name)
+
+    client.delete_model(name=model_name)