Move DELETE /mlflow/registered-models/delete endpoint. (#89)

Signed-off-by: Software Developer <7852635+dsuhinin@users.noreply.github.com>

magefiles/generate/endpoints.go

@@ -51,7 +51,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			// "createRegisteredModel",
 			"renameRegisteredModel",
 			"updateRegisteredModel",
-			// "deleteRegisteredModel",
+			"deleteRegisteredModel",
 			// "getRegisteredModel",
 			// "searchRegisteredModels",
 			"getLatestVersions",

mlflow_go/store/model_registry.py

@@ -3,6 +3,7 @@
 
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
 from mlflow.protos.model_registry_pb2 import (
+    DeleteRegisteredModel,
     GetLatestVersions,
     RenameRegisteredModel,
     UpdateRegisteredModel,
@@ -55,6 +56,10 @@ def rename_registered_model(self, name, new_name):
         )
         return RegisteredModel.from_proto(response.registered_model)
 
+    def delete_registered_model(self, name):
+        request = DeleteRegisteredModel(name=name)
+        self.service.call_endpoint(get_lib().ModelRegistryServiceDeleteRegisteredModel, request)
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/model_registry/service/model_versions.go

@@ -53,3 +53,13 @@ func (m *ModelRegistryService) RenameRegisteredModel(
 		RegisteredModel: registeredModel.ToProto(),
 	}, nil
 }
+
+func (m *ModelRegistryService) DeleteRegisteredModel(
+	ctx context.Context, input *protos.DeleteRegisteredModel,
+) (*protos.DeleteRegisteredModel_Response, *contract.Error) {
+	if err := m.store.DeleteRegisteredModel(ctx, input.GetName()); err != nil {
+		return nil, err
+	}
+
+	return &protos.DeleteRegisteredModel_Response{}, nil
+}

pkg/model_registry/store/sql/model_versions.go

@@ -199,3 +199,60 @@ func (m *ModelRegistrySQLStore) RenameRegisteredModel(
 
 	return registeredModel, nil
 }
+
+func (m *ModelRegistrySQLStore) DeleteRegisteredModel(ctx context.Context, name string) *contract.Error {
+	registeredModel, err := m.GetRegisteredModelByName(ctx, name)
+	if err != nil {
+		return err
+	}
+
+	if err := m.db.WithContext(ctx).Transaction(func(transaction *gorm.DB) error {
+		if err := transaction.Where(
+			"name = ?", registeredModel.Name,
+		).Delete(
+			models.ModelVersionTag{},
+		).Error; err != nil {
+			return err
+		}
+
+		if err := transaction.Where(
+			"name = ?", registeredModel.Name,
+		).Delete(
+			models.ModelVersion{},
+		).Error; err != nil {
+			return err
+		}
+
+		if err := transaction.Where(
+			"name = ?", registeredModel.Name,
+		).Delete(
+			models.RegisteredModelTag{},
+		).Error; err != nil {
+			return err
+		}
+
+		if err := transaction.Where(
+			"name = ?", registeredModel.Name,
+		).Delete(
+			models.RegisteredModelAlias{},
+		).Error; err != nil {
+			return err
+		}
+
+		if err := transaction.Where(
+			"name = ?", registeredModel.Name,
+		).Delete(
+			models.RegisteredModel{},
+		).Error; err != nil {
+			return err
+		}
+
+		return nil
+	}); err != nil {
+		return contract.NewError(
+			protos.ErrorCode_INTERNAL_ERROR, fmt.Sprintf("error deleting registered model: %v", err),
+		)
+	}
+
+	return nil
+}

pkg/model_registry/store/store.go

@@ -13,4 +13,5 @@ type ModelRegistryStore interface {
 	GetLatestVersions(ctx context.Context, name string, stages []string) ([]*protos.ModelVersion, *contract.Error)
 	UpdateRegisteredModel(ctx context.Context, name, description string) (*entities.RegisteredModel, *contract.Error)
 	RenameRegisteredModel(ctx context.Context, name, newName string) (*entities.RegisteredModel, *contract.Error)
+	DeleteRegisteredModel(ctx context.Context, name string) *contract.Error
 }