Move GET /mlflow/model-versions/get-download-uri endpoint.

magefiles/generate/endpoints.go

@@ -61,7 +61,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			"deleteModelVersion",
 			"getModelVersion",
 			// "searchModelVersions",
-			// "getModelVersionDownloadUri",
+			"getModelVersionDownloadUri",
 			"setRegisteredModelTag",
 			"setModelVersionTag",
 			"deleteRegisteredModelTag",

magefiles/generate/validations.go

@@ -72,4 +72,5 @@ var validations = map[string]string{
 	"SetModelVersionTag_Value":           "omitempty,max=5000,truncate=5000",
 	"SetModelVersionTag_Version":         "stringAsInteger",
 	"GetModelVersion_Version":            "stringAsInteger",
+	"GetModelVersionDownloadUri_Version": "stringAsInteger",
 }

mlflow_go_backend/store/model_registry.py

@@ -12,6 +12,7 @@
     GetLatestVersions,
     GetModelVersion,
     GetModelVersionByAlias,
+    GetModelVersionDownloadUri,
     GetRegisteredModel,
     RenameRegisteredModel,
     SetModelVersionTag,
@@ -168,6 +169,13 @@ def set_model_version_tag(self, name, version, tag):
         request = SetModelVersionTag(name=name, version=str(version), key=tag.key, value=tag.value)
         self.service.call_endpoint(get_lib().ModelRegistryServiceSetModelVersionTag, request)
 
+    def get_model_version_download_uri(self, name, version):
+        request = GetModelVersionDownloadUri(name=name, version=str(version))
+        response = self.service.call_endpoint(
+            get_lib().ModelRegistryServiceGetModelVersionDownloadUri, request
+        )
+        return response.artifact_uri
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/model_registry/service/model_versions.go

@@ -9,6 +9,7 @@ import (
 	"github.com/mlflow/mlflow-go-backend/pkg/contract"
 	"github.com/mlflow/mlflow-go-backend/pkg/model_registry/store/sql/models"
 	"github.com/mlflow/mlflow-go-backend/pkg/protos"
+	"github.com/mlflow/mlflow-go-backend/pkg/utils"
 )
 
 func (m *ModelRegistryService) GetLatestVersions(
@@ -142,3 +143,17 @@ func (m *ModelRegistryService) SetModelVersionTag(
 
 	return &protos.SetModelVersionTag_Response{}, nil
 }
+
+//nolint:revive,stylecheck
+func (m *ModelRegistryService) GetModelVersionDownloadUri(
+	ctx context.Context, input *protos.GetModelVersionDownloadUri,
+) (*protos.GetModelVersionDownloadUri_Response, *contract.Error) {
+	artifactURI, err := m.store.GetModelVersionDownloadURI(ctx, input.GetName(), input.GetVersion())
+	if err != nil {
+		return nil, err
+	}
+
+	return &protos.GetModelVersionDownloadUri_Response{
+		ArtifactUri: utils.PtrTo(artifactURI),
+	}, nil
+}

pkg/model_registry/store/sql/model_versions.go

@@ -385,3 +385,18 @@ func (m *ModelRegistrySQLStore) SetModelVersionTag(
 
 	return nil
 }
+
+func (m *ModelRegistrySQLStore) GetModelVersionDownloadURI(
+	ctx context.Context, name, version string,
+) (string, *contract.Error) {
+	modelVersion, err := m.GetModelVersion(ctx, name, version, false)
+	if err != nil {
+		return "", err
+	}
+
+	if modelVersion.StorageLocation != "" {
+		return modelVersion.StorageLocation, nil
+	}
+
+	return modelVersion.Source, nil
+}

pkg/model_registry/store/store.go

@@ -26,6 +26,7 @@ type ModelVersionStore interface {
 	DeleteModelVersionTag(ctx context.Context, name, version, key string) *contract.Error
 	GetModelVersionByAlias(ctx context.Context, name, alias string) (*entities.ModelVersion, *contract.Error)
 	SetModelVersionTag(ctx context.Context, name, version, key, value string) *contract.Error
+	GetModelVersionDownloadURI(ctx context.Context, name, version string) (string, *contract.Error)
 }
 
 type RegisteredModelStore interface {