[API] Adding APIs to delete specific versions of feature-store objects (

#811)

mlrun/api/api/endpoints/feature_sets.py

@@ -86,10 +86,17 @@ def get_feature_set(
 
 
 @router.delete("/projects/{project}/feature-sets/{name}")
+@router.delete("/projects/{project}/feature-sets/{name}/references/{reference}")
 def delete_feature_set(
-    project: str, name: str, db_session: Session = Depends(deps.get_db_session),
+    project: str,
+    name: str,
+    reference: str = None,
+    db_session: Session = Depends(deps.get_db_session),
 ):
-    get_db().delete_feature_set(db_session, project, name)
+    tag = uid = None
+    if reference:
+        tag, uid = parse_reference(reference)
+    get_db().delete_feature_set(db_session, project, name, tag, uid)
     return Response(status_code=HTTPStatus.NO_CONTENT.value)
 
 
@@ -231,8 +238,15 @@ def patch_feature_vector(
 
 
 @router.delete("/projects/{project}/feature-vectors/{name}")
+@router.delete("/projects/{project}/feature-vectors/{name}/references/{reference}")
 def delete_feature_vector(
-    project: str, name: str, db_session: Session = Depends(deps.get_db_session),
+    project: str,
+    name: str,
+    reference: str = None,
+    db_session: Session = Depends(deps.get_db_session),
 ):
-    get_db().delete_feature_vector(db_session, project, name)
+    tag = uid = None
+    if reference:
+        tag, uid = parse_reference(reference)
+    get_db().delete_feature_vector(db_session, project, name, tag, uid)
     return Response(status_code=HTTPStatus.NO_CONTENT.value)

mlrun/api/db/base.py

@@ -303,7 +303,7 @@ def patch_feature_set(
         pass
 
     @abstractmethod
-    def delete_feature_set(self, session, project, name):
+    def delete_feature_set(self, session, project, name, tag=None, uid=None):
         pass
 
     @abstractmethod
@@ -358,7 +358,9 @@ def patch_feature_vector(
         pass
 
     @abstractmethod
-    def delete_feature_vector(self, session, project, name):
+    def delete_feature_vector(
+        self, session, project, name, tag=None, uid=None,
+    ):
         pass
 
     def list_artifact_tags(self, session, project):

mlrun/api/db/filedb/db.py

@@ -242,7 +242,7 @@ def patch_feature_set(
     ):
         raise NotImplementedError()
 
-    def delete_feature_set(self, session, project, name):
+    def delete_feature_set(self, session, project, name, tag=None, uid=None):
         raise NotImplementedError()
 
     def create_feature_vector(
@@ -291,7 +291,7 @@ def patch_feature_vector(
     ):
         raise NotImplementedError()
 
-    def delete_feature_vector(self, session, project, name):
+    def delete_feature_vector(self, session, project, name, tag=None, uid=None):
         raise NotImplementedError()
 
     def list_artifact_tags(self, session, project):

mlrun/api/db/sqldb/db.py

@@ -1482,9 +1482,38 @@ def patch_feature_set(
             always_overwrite=True,
         )
 
-    def delete_feature_set(self, session, project, name):
-        self._delete(session, FeatureSet.Tag, project=project, obj_name=name)
-        self._delete(session, FeatureSet, project=project, name=name)
+    def _delete_feature_store_object(self, session, cls, project, name, tag, uid):
+        if tag and uid:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Both uid and tag specified when deleting an object."
+            )
+
+        object_id = None
+        if uid:
+            object_record = self._query(
+                session, cls, project=project, name=name, uid=uid
+            ).one_or_none()
+            if object_record is None:
+                return
+            object_id = object_record.id
+        elif tag:
+            tag_record = self._query(
+                session, cls.Tag, project=project, name=tag, obj_name=name
+            ).one_or_none()
+            if tag_record is None:
+                return
+            object_id = tag_record.obj_id
+
+        if object_id:
+            self._delete(session, cls, id=object_id)
+            self._delete(session, cls.Tag, obj_id=object_id)
+        else:
+            # If we got here, neither tag nor uid were provided - delete all references by name.
+            self._delete(session, cls, project=project, name=name)
+            self._delete(session, cls.Tag, project=project, obj_name=name)
+
+    def delete_feature_set(self, session, project, name, tag=None, uid=None):
+        self._delete_feature_store_object(session, FeatureSet, project, name, tag, uid)
 
     def create_feature_vector(
         self, session, project, feature_vector: schemas.FeatureVector, versioned=True
@@ -1669,9 +1698,10 @@ def patch_feature_vector(
             always_overwrite=True,
         )
 
-    def delete_feature_vector(self, session, project, name):
-        self._delete(session, FeatureVector.Tag, project=project, obj_name=name)
-        self._delete(session, FeatureVector, project=project, name=name)
+    def delete_feature_vector(self, session, project, name, tag=None, uid=None):
+        self._delete_feature_store_object(
+            session, FeatureVector, project, name, tag, uid
+        )
 
     def _resolve_tag(self, session, cls, project, name):
         ids = []

mlrun/db/base.py

@@ -232,7 +232,7 @@ def patch_feature_set(
         pass
 
     @abstractmethod
-    def delete_feature_set(self, name, project=""):
+    def delete_feature_set(self, name, project="", tag=None, uid=None):
         pass
 
     @abstractmethod
@@ -286,7 +286,7 @@ def patch_feature_vector(
         pass
 
     @abstractmethod
-    def delete_feature_vector(self, name, project=""):
+    def delete_feature_vector(self, name, project="", tag=None, uid=None):
         pass
 
     @abstractmethod

mlrun/db/filedb.py

@@ -566,7 +566,7 @@ def patch_feature_set(
     ):
         raise NotImplementedError()
 
-    def delete_feature_set(self, name, project=""):
+    def delete_feature_set(self, name, project="", tag=None, uid=None):
         raise NotImplementedError()
 
     def create_feature_vector(self, feature_vector, project="", versioned=True) -> dict:
@@ -603,7 +603,7 @@ def patch_feature_vector(
     ):
         raise NotImplementedError()
 
-    def delete_feature_vector(self, name, project=""):
+    def delete_feature_vector(self, name, project="", tag=None, uid=None):
         raise NotImplementedError()
 
     def list_pipelines(

mlrun/db/httpdb.py

@@ -1371,10 +1371,19 @@ def patch_feature_set(
             headers=headers,
         )
 
-    def delete_feature_set(self, name, project=""):
-        """ Delete a :py:class:`~mlrun.feature_store.FeatureSet` object from the DB. """
+    def delete_feature_set(self, name, project="", tag=None, uid=None):
+        """ Delete a :py:class:`~mlrun.feature_store.FeatureSet` object from the DB.
+        If ``tag`` or ``uid`` are specified, then just the version referenced by them will be deleted. Using both
+        is not allowed.
+        If none are specified, then all instances of the object whose name is ``name`` will be deleted.
+        """
         project = project or default_project
         path = f"projects/{project}/feature-sets/{name}"
+
+        if tag or uid:
+            reference = self._resolve_reference(tag, uid)
+            path = path + f"/references/{reference}"
+
         error_message = f"Failed deleting feature-set {name}"
         self.api_call("DELETE", path, error_message)
 
@@ -1541,11 +1550,18 @@ def patch_feature_vector(
             headers=headers,
         )
 
-    def delete_feature_vector(self, name, project=""):
-        """ Delete a :py:class:`~mlrun.feature_store.FeatureVector` object from the DB. """
-
+    def delete_feature_vector(self, name, project="", tag=None, uid=None):
+        """ Delete a :py:class:`~mlrun.feature_store.FeatureVector` object from the DB.
+        If ``tag`` or ``uid`` are specified, then just the version referenced by them will be deleted. Using both
+        is not allowed.
+        If none are specified, then all instances of the object whose name is ``name`` will be deleted.
+        """
         project = project or default_project
         path = f"projects/{project}/feature-vectors/{name}"
+        if tag or uid:
+            reference = self._resolve_reference(tag, uid)
+            path = path + f"/references/{reference}"
+
         error_message = f"Failed deleting feature-vector {name}"
         self.api_call("DELETE", path, error_message)
 

mlrun/db/sqldb.py

@@ -320,9 +320,9 @@ def patch_feature_set(
             patch_mode,
         )
 
-    def delete_feature_set(self, name, project=""):
+    def delete_feature_set(self, name, project="", tag=None, uid=None):
         return self._transform_db_error(
-            self.db.delete_feature_set, self.session, project, name
+            self.db.delete_feature_set, self.session, project, name, tag, uid
         )
 
     def create_feature_vector(self, feature_vector, project="", versioned=True):
@@ -393,9 +393,9 @@ def patch_feature_vector(
             patch_mode,
         )
 
-    def delete_feature_vector(self, name, project=""):
+    def delete_feature_vector(self, name, project="", tag=None, uid=None):
         return self._transform_db_error(
-            self.db.delete_feature_vector, self.session, project, name,
+            self.db.delete_feature_vector, self.session, project, name, tag, uid
         )
 
     def list_pipelines(

tests/api/api/feature_store/test_feature_sets.py

@@ -272,6 +272,54 @@ def test_feature_set_delete(db: Session, client: TestClient) -> None:
     _list_and_assert_objects(client, "feature_sets", project_name, None, count - 2)
 
 
+def test_feature_set_delete_version(db: Session, client: TestClient) -> None:
+    project_name = f"prj-{uuid4().hex}"
+
+    name = "feature_set"
+    feature_set = _generate_feature_set(name)
+
+    count = 5
+    uids = {}
+    for i in range(count):
+        # Store different copies of the feature set with different uids and tags
+        feature_set["metadata"]["extra_metadata"] = i * 100
+        tag = f"tag{i}"
+        result = _store_and_assert_feature_set(
+            client, project_name, name, f"tag{i}", feature_set
+        )
+        uids[result["metadata"]["uid"]] = tag
+
+    _list_and_assert_objects(
+        client, "feature_sets", project_name, f"name={name}", count
+    )
+
+    delete_by_tag = True
+    objects_left = count
+    for uid, tag in uids.items():
+        reference = tag if delete_by_tag else uid
+        delete_by_tag = not delete_by_tag
+
+        response = client.delete(
+            f"/api/projects/{project_name}/feature-sets/{name}/references/{reference}"
+        )
+        assert response.status_code == HTTPStatus.NO_CONTENT.value
+        objects_left = objects_left - 1
+        _list_and_assert_objects(
+            client, "feature_sets", project_name, f"name={name}", objects_left
+        )
+
+    for i in range(count):
+        feature_set["metadata"]["extra_metadata"] = i * 100
+        _store_and_assert_feature_set(
+            client, project_name, name, f"tag{i}", feature_set
+        )
+
+    # Now delete by name
+    response = client.delete(f"/api/projects/{project_name}/feature-sets/{name}")
+    assert response.status_code == HTTPStatus.NO_CONTENT.value
+    _list_and_assert_objects(client, "feature_sets", project_name, f"name={name}", 0)
+
+
 def test_feature_set_create_failure_already_exists(
     db: Session, client: TestClient
 ) -> None:

tests/api/api/feature_store/test_feature_vectors.py

@@ -281,6 +281,54 @@ def test_feature_vector_delete(db: Session, client: TestClient) -> None:
     _list_and_assert_objects(client, "feature_vectors", project_name, None, count - 2)
 
 
+def test_feature_vector_delete_version(db: Session, client: TestClient) -> None:
+    project_name = f"prj-{uuid4().hex}"
+
+    name = "feature_vector"
+    feature_vector = _generate_feature_vector(name)
+
+    count = 5
+    uids = {}
+    for i in range(count):
+        # Store different copies of the feature set with different uids and tags
+        feature_vector["metadata"]["extra_metadata"] = i * 100
+        tag = f"tag{i}"
+        result = _assert_store_feature_vector(
+            client, project_name, name, f"tag{i}", feature_vector
+        )
+        uids[result["metadata"]["uid"]] = tag
+
+    _list_and_assert_objects(
+        client, "feature_vectors", project_name, f"name={name}", count
+    )
+
+    delete_by_tag = True
+    objects_left = count
+    for uid, tag in uids.items():
+        reference = tag if delete_by_tag else uid
+        delete_by_tag = not delete_by_tag
+
+        response = client.delete(
+            f"/api/projects/{project_name}/feature-vectors/{name}/references/{reference}"
+        )
+        assert response.status_code == HTTPStatus.NO_CONTENT.value
+        objects_left = objects_left - 1
+        _list_and_assert_objects(
+            client, "feature_vectors", project_name, f"name={name}", objects_left
+        )
+
+    for i in range(count):
+        feature_vector["metadata"]["extra_metadata"] = i * 100
+        _assert_store_feature_vector(
+            client, project_name, name, f"tag{i}", feature_vector
+        )
+
+    # Now delete by name
+    response = client.delete(f"/api/projects/{project_name}/feature-vectors/{name}")
+    assert response.status_code == HTTPStatus.NO_CONTENT.value
+    _list_and_assert_objects(client, "feature_vectors", project_name, f"name={name}", 0)
+
+
 def test_unversioned_feature_vector_actions(db: Session, client: TestClient) -> None:
     project_name = f"prj-{uuid4().hex}"
     name = "feature_vector1"