fixed handler

Author: shreyas-goenka

bundle/direct/dresources/all_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/apps"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/database"
+	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -60,6 +61,18 @@ var testConfig map[string]any = map[string]any{
 			Name: "main.myschema.my_synced_table",
 		},
 	},
+	"models": &resources.MlflowModel{
+		CreateModelRequest: ml.CreateModelRequest{
+			Name:        "my_mlflow_model",
+			Description: "my_mlflow_model_description",
+			Tags: []ml.ModelTag{
+				{
+					Key:   "k1",
+					Value: "v1",
+				},
+			},
+		},
+	},
 }
 
 type prepareWorkspace func(client *databricks.WorkspaceClient) error

libs/testserver/fake_workspace.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/dashboards"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
@@ -60,21 +61,22 @@ type FakeWorkspace struct {
 	repoIdByPath map[string]int64
 
 	// normally, ids are not sequential, but we make them sequential for deterministic diff
-	nextJobId       int64
-	nextJobRunId    int64
-	Jobs            map[int64]jobs.Job
-	JobRuns         map[int64]jobs.Run
-	JobPermissions  map[string][]jobs.JobAccessControlRequest
-	Pipelines       map[string]pipelines.GetPipelineResponse
-	PipelineUpdates map[string]bool
-	Monitors        map[string]catalog.MonitorInfo
-	Apps            map[string]apps.App
-	Schemas         map[string]catalog.SchemaInfo
-	SchemasGrants   map[string][]catalog.PrivilegeAssignment
-	Volumes         map[string]catalog.VolumeInfo
-	Dashboards      map[string]dashboards.Dashboard
-	SqlWarehouses   map[string]sql.GetWarehouseResponse
-	Alerts          map[string]sql.AlertV2
+	nextJobId           int64
+	nextJobRunId        int64
+	Jobs                map[int64]jobs.Job
+	JobRuns             map[int64]jobs.Run
+	JobPermissions      map[string][]jobs.JobAccessControlRequest
+	Pipelines           map[string]pipelines.GetPipelineResponse
+	PipelineUpdates     map[string]bool
+	Monitors            map[string]catalog.MonitorInfo
+	Apps                map[string]apps.App
+	Schemas             map[string]catalog.SchemaInfo
+	SchemasGrants       map[string][]catalog.PrivilegeAssignment
+	Volumes             map[string]catalog.VolumeInfo
+	Dashboards          map[string]dashboards.Dashboard
+	SqlWarehouses       map[string]sql.GetWarehouseResponse
+	Alerts              map[string]sql.AlertV2
+	ModelRegistryModels map[string]ml.Model
 
 	Acls map[string][]workspace.AclItem
 
@@ -172,6 +174,7 @@ func NewFakeWorkspace(url, token string) *FakeWorkspace {
 		DatabaseCatalogs:     map[string]database.DatabaseCatalog{},
 		SyncedDatabaseTables: map[string]database.SyncedDatabaseTable{},
 		Alerts:               map[string]sql.AlertV2{},
+		ModelRegistryModels:  map[string]ml.Model{},
 	}
 }
 

libs/testserver/handlers.go

@@ -521,4 +521,21 @@ func AddDefaultHandlers(server *Server) {
 	server.Handle("GET", "/api/2.0/permissions/jobs/{job_id}", func(req Request) any {
 		return req.Workspace.JobsGetPermissions(req, req.Vars["job_id"])
 	})
+
+	// Model registry models.
+	server.Handle("POST", "/api/2.0/mlflow/registered-models/create", func(req Request) any {
+		return req.Workspace.ModelRegistryCreateModel(req)
+	})
+
+	server.Handle("GET", "/api/2.0/mlflow/databricks/registered-models/get", func(req Request) any {
+		return req.Workspace.ModelRegistryGetModel(req)
+	})
+
+	server.Handle("PATCH", "/api/2.0/mlflow/registered-models/update", func(req Request) any {
+		return req.Workspace.ModelRegistryUpdateModel(req)
+	})
+
+	server.Handle("DELETE", "/api/2.0/mlflow/registered-models/delete", func(req Request) any {
+		return MapDelete(req.Workspace, req.Workspace.ModelRegistryModels, req.URL.Query().Get("name"))
+	})
 }

libs/testserver/models.go

@@ -0,0 +1,90 @@
+package testserver
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go/service/ml"
+)
+
+func (s *FakeWorkspace) ModelRegistryCreateModel(req Request) any {
+	defer s.LockUnlock()()
+
+	var request ml.CreateModelRequest
+	if err := json.Unmarshal(req.Body, &request); err != nil {
+		return Response{
+			StatusCode: 400,
+			Body:       map[string]string{"message": fmt.Sprintf("Failed to parse request: %s", err)},
+		}
+	}
+
+	// Create the model
+	model := ml.Model{
+		Name:        request.Name,
+		Description: request.Description,
+		Tags:        request.Tags,
+	}
+
+	s.ModelRegistryModels[request.Name] = model
+
+	return Response{
+		Body: ml.CreateModelResponse{
+			RegisteredModel: &model,
+		},
+	}
+}
+
+func (s *FakeWorkspace) ModelRegistryUpdateModel(req Request) any {
+	defer s.LockUnlock()()
+
+	var request ml.UpdateModelRequest
+	if err := json.Unmarshal(req.Body, &request); err != nil {
+		return Response{
+			StatusCode: 400,
+			Body:       map[string]string{"message": fmt.Sprintf("Failed to parse request: %s", err)},
+		}
+	}
+
+	existingModel, ok := s.ModelRegistryModels[request.Name]
+	if !ok {
+		return Response{
+			StatusCode: 404,
+			Body:       map[string]string{"message": fmt.Sprintf("Model not found: %v", request.Name)},
+		}
+	}
+
+	// Update the model
+	existingModel.Description = request.Description
+	s.ModelRegistryModels[request.Name] = existingModel
+
+	return Response{
+		Body: ml.UpdateModelResponse{
+			RegisteredModel: &existingModel,
+		},
+	}
+}
+
+func (s *FakeWorkspace) ModelRegistryGetModel(req Request) any {
+	defer s.LockUnlock()()
+
+	name := req.URL.Query().Get("name")
+
+	model, ok := s.ModelRegistryModels[name]
+	if !ok {
+		return Response{
+			StatusCode: 404,
+			Body:       map[string]string{"message": fmt.Sprintf("Model not found: %v", name)},
+		}
+	}
+
+	return Response{
+		Body: ml.GetModelResponse{
+			RegisteredModelDatabricks: &ml.ModelDatabricks{
+				Name:            model.Name,
+				Description:     model.Description,
+				Tags:            model.Tags,
+				ForceSendFields: model.ForceSendFields,
+			},
+		},
+	}
+}