1 files changed, 77 insertions, 3 deletions

diff --git a/lib/api/ml/mlflow/registered_models.rb b/lib/api/ml/mlflow/registered_models.rb
index 18b705ad214..a68a2767a74 100644
--- a/lib/api/ml/mlflow/registered_models.rb
+++ b/lib/api/ml/mlflow/registered_models.rb
@@ -11,8 +11,9 @@ module API
 
         before do
           check_api_read!
-          check_api_write! unless request.get? || request.head?
           check_api_model_registry_read!
+          check_api_write! if route.settings.dig(:api, :write)
+          check_api_model_registry_write! if route.settings.dig(:model_registry, :write)
         end
 
         resource 'registered-models' do
@@ -20,6 +21,8 @@ module API
             success Entities::Ml::Mlflow::RegisteredModel
             detail 'MLFlow Registered Models map to GitLab Models. https://mlflow.org/docs/2.6.0/rest-api.html#create-registeredmodel'
           end
+          route_setting :api, write: true
+          route_setting :model_registry, write: true
           params do
             requires :name, type: String,
               desc: 'Register models under this name.'
@@ -60,6 +63,8 @@ module API
             success Entities::Ml::Mlflow::RegisteredModel
             detail 'https://mlflow.org/docs/2.6.0/rest-api.html#update-registeredmodel'
           end
+          route_setting :api, write: true
+          route_setting :model_registry, write: true
           params do
             # The name param is actually required, however it is listed as optional here
             # we can send a custom error response required by MLFlow
@@ -74,7 +79,7 @@ module API
           end
 
           desc 'Fetch the latest Model Version for the given Registered Model Name' do
-            success Entities::Ml::Mlflow::ModelVersions::Types::ModelVersion
+            success Entities::Ml::Mlflow::ModelVersion
             detail 'https://mlflow.org/docs/2.6.0/rest-api.html#get-latest-modelversions'
           end
           params do
@@ -86,9 +91,78 @@ module API
           post 'get-latest-versions', urgency: :low do
             model = find_model(user_project, params[:name])
 
-            present [model.latest_version], with: Entities::Ml::Mlflow::ModelVersions::Types::ModelVersion,
+            present [model.latest_version], with: Entities::Ml::Mlflow::ModelVersion,
               root: :model_versions
           end
+
+          desc 'Delete a Registered Model by Name' do
+            success Entities::Ml::Mlflow::RegisteredModel
+            detail 'https://mlflow.org/docs/2.6.0/rest-api.html#delete-registeredmodel'
+          end
+          route_setting :api, write: true
+          route_setting :model_registry, write: true
+          params do
+            # The name param is actually required, however it is listed as optional here
+            # we can send a custom error response required by MLFlow
+            optional :name, type: String,
+              desc: 'Registered model unique name identifier, in reference to the project'
+          end
+          delete 'delete', urgency: :low do
+            resource_not_found! unless params[:name]
+
+            model = ::Ml::FindModelService.new(user_project, params[:name]).execute
+
+            resource_not_found! unless model
+
+            if ::Ml::DestroyModelService.new(model, current_user).execute
+              present({})
+            else
+              render_api_error!('Model could not be deleted', 400)
+            end
+          end
+
+          desc 'Search Registered Models within a project' do
+            success Entities::Ml::Mlflow::RegisteredModel
+            detail 'https://mlflow.org/docs/2.6.0/rest-api.html#search-registeredmodels'
+          end
+          params do
+            optional :filter,
+              type: String,
+              desc: "Filter to search models. must be in the format `name='value'`. Only filtering by name is supported"
+            optional :max_results,
+              type: Integer,
+              desc: 'Maximum number of models desired. Default is 200. Max threshold is 1000.',
+              default: 200
+            optional :order_by,
+              type: String,
+              desc: 'Order criteria. Can be by name or last_updated_timestamp, with optional DESC or ASC (default)' \
+                    'Valid examples: `name`, `name DESC`, `last_updated_timestamp DESC`' \
+                    'Sorting by model metadata is not supported.',
+              default: 'name ASC'
+            optional :page_token,
+              type: String,
+              desc: 'Token for pagination'
+          end
+          get 'search', urgency: :low do
+            max_results = [params[:max_results], 1000].min
+
+            finder_params = model_order_params(params)
+            filter_params = model_filter_params(params)
+
+            if !params[:filter].nil? && !filter_params.key?(:name)
+              invalid_parameter!("Invalid attribute key specified. Valid keys are '{'name'}'")
+            end
+
+            finder = ::Projects::Ml::ModelFinder.new(user_project, finder_params.merge(filter_params))
+            paginator = finder.execute.keyset_paginate(cursor: params[:page_token], per_page: max_results)
+
+            result = {
+              registered_models: paginator.records,
+              next_page_token: paginator.cursor_for_next_page
+            }
+
+            present result, with: Entities::Ml::Mlflow::ListRegisteredModels
+          end
         end
       end
     end