feat(proxy): introduce datastore abstraction

cmd/proxy.go

@@ -1,63 +1,41 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"net/http"
-	"time"
 
 	"github.com/golang/glog"
-	"github.com/kubeflow/model-registry/internal/mlmdtypes"
+	"github.com/kubeflow/model-registry/internal/datastore"
 	"github.com/kubeflow/model-registry/internal/server/openapi"
-	"github.com/kubeflow/model-registry/pkg/core"
 	"github.com/spf13/cobra"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
 )
 
 // proxyCmd represents the proxy command
 var proxyCmd = &cobra.Command{
 	Use:   "proxy",
-	Short: "Starts the ml-metadata go OpenAPI proxy",
-	Long: `This command launches the ml-metadata go OpenAPI proxy server.
+	Short: "Starts the go OpenAPI proxy server to connect to a metadata store",
+	Long: `This command launches the go OpenAPI proxy server.
 
-The server connects to a mlmd CPP server. It supports options to customize the
-hostname and port where it listens.'`,
+The server connects to a metadata store, currently only MLMD is supported. It supports options to customize the
+hostname and port where it listens.`,
 	RunE: runProxyServer,
 }
 
 func runProxyServer(cmd *cobra.Command, args []string) error {
 	glog.Infof("proxy server started at %s:%v", cfg.Hostname, cfg.Port)
 
-	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Second*30)
-	defer cancel()
-
-	mlmdAddr := fmt.Sprintf("%s:%d", proxyCfg.MLMDHostname, proxyCfg.MLMDPort)
-	glog.Infof("connecting to MLMD server %s..", mlmdAddr)
-	conn, err := grpc.DialContext( // nolint:staticcheck
-		ctxTimeout,
-		mlmdAddr,
-		grpc.WithReturnConnectionError(), // nolint:staticcheck
-		grpc.WithBlock(),                 // nolint:staticcheck
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-	)
+	ds, dsTeardownF, err := datastore.NewDatastore(proxyCfg.DatastoreType, proxyCfg.DatastoreHostname, proxyCfg.DatastorePort)
 	if err != nil {
-		return fmt.Errorf("error dialing connection to mlmd server %s: %v", mlmdAddr, err)
+		return fmt.Errorf("error creating datastore: %w", err)
 	}
-	defer conn.Close()
-	glog.Infof("connected to MLMD server")
 
-	mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
-	_, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig)
-	if err != nil {
-		return fmt.Errorf("error creating MLMD types: %v", err)
-	}
-	service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
-	if err != nil {
-		return fmt.Errorf("error creating core service: %v", err)
-	}
+	defer func() {
+		if err := dsTeardownF(); err != nil {
+			glog.Errorf("error during cleanup: %w", err)
+		}
+	}()
 
-	ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(service)
+	ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(ds)
 	ModelRegistryServiceAPIController := openapi.NewModelRegistryServiceAPIController(ModelRegistryServiceAPIService)
 
 	router := openapi.NewRouter(ModelRegistryServiceAPIController)
@@ -72,16 +50,29 @@ func init() {
 	proxyCmd.Flags().StringVarP(&cfg.Hostname, "hostname", "n", cfg.Hostname, "Proxy server listen hostname")
 	proxyCmd.Flags().IntVarP(&cfg.Port, "port", "p", cfg.Port, "Proxy server listen port")
 
-	proxyCmd.Flags().StringVar(&proxyCfg.MLMDHostname, "mlmd-hostname", proxyCfg.MLMDHostname, "MLMD hostname")
-	proxyCmd.Flags().IntVar(&proxyCfg.MLMDPort, "mlmd-port", proxyCfg.MLMDPort, "MLMD port")
+	proxyCmd.Flags().StringVar(&proxyCfg.DatastoreHostname, "mlmd-hostname", proxyCfg.DatastoreHostname, "MLMD hostname")
+	if err := proxyCmd.Flags().MarkDeprecated("mlmd-hostname", "please use --datastore-hostname instead"); err != nil {
+		glog.Errorf("error marking flag as deprecated: %v", err)
+	}
+
+	proxyCmd.Flags().IntVar(&proxyCfg.DatastorePort, "mlmd-port", proxyCfg.DatastorePort, "MLMD port")
+	if err := proxyCmd.Flags().MarkDeprecated("mlmd-port", "please use --datastore-port instead"); err != nil {
+		glog.Errorf("error marking flag as deprecated: %v", err)
+	}
+
+	proxyCmd.Flags().StringVar(&proxyCfg.DatastoreHostname, "datastore-hostname", proxyCfg.DatastoreHostname, "Datastore hostname")
+	proxyCmd.Flags().IntVar(&proxyCfg.DatastorePort, "datastore-port", proxyCfg.DatastorePort, "Datastore port")
+	proxyCmd.Flags().StringVar(&proxyCfg.DatastoreType, "datastore-type", proxyCfg.DatastoreType, "Datastore type")
 }
 
 type ProxyConfig struct {
-	MLMDHostname string
-	MLMDPort     int
+	DatastoreHostname string
+	DatastorePort     int
+	DatastoreType     string
 }
 
 var proxyCfg = ProxyConfig{
-	MLMDHostname: "localhost",
-	MLMDPort:     9090,
+	DatastoreHostname: "localhost",
+	DatastorePort:     9090,
+	DatastoreType:     "mlmd",
 }

docker-compose-local.yaml

@@ -13,7 +13,7 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
-    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080"]
+    command: ["proxy", "--hostname", "0.0.0.0", "--datastore-hostname", "mlmd-server", "--datastore-port", "8080", "--datastore-type", "mlmd"]
     container_name: model-registry
     ports:
       - "8080:8080"

docker-compose.yaml

@@ -11,7 +11,7 @@ services:
       - ./test/config/ml-metadata:/tmp/shared
   model-registry:
     image: docker.io/kubeflow/model-registry:latest
-    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080"]
+    command: ["proxy", "--hostname", "0.0.0.0", "--datastore-hostname", "mlmd-server", "--datastore-port", "8080", "--datastore-type", "mlmd"]
     container_name: model-registry
     ports:
       - "8080:8080"

internal/datastore/builder.go

@@ -0,0 +1,34 @@
+package datastore
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/kubeflow/model-registry/pkg/api"
+)
+
+var (
+	ErrCreatingDatastore    = errors.New("error creating datastore")
+	ErrUnsupportedDatastore = errors.New("unsupported datastore type")
+)
+
+type Builder interface {
+	Build() (api.ModelRegistryApi, error)
+	Teardown() error
+}
+
+func NewDatastore(dsType string, dsHostname string, dsPort int) (api.ModelRegistryApi, func() error, error) {
+	switch dsType {
+	case "mlmd":
+		mlmd := NewMLMDService(dsHostname, dsPort)
+
+		svc, err := mlmd.Build()
+		if err != nil {
+			return nil, nil, fmt.Errorf("%w: %w", ErrCreatingDatastore, err)
+		}
+
+		return svc, mlmd.Teardown, nil
+	default:
+		return nil, nil, fmt.Errorf("%w: %s", ErrUnsupportedDatastore, dsType)
+	}
+}

internal/datastore/mlmd.go

@@ -0,0 +1,83 @@
+package datastore
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/golang/glog"
+	"github.com/kubeflow/model-registry/internal/mlmdtypes"
+	"github.com/kubeflow/model-registry/pkg/api"
+	"github.com/kubeflow/model-registry/pkg/core"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+var (
+	ErrMLMDConnectionStart = errors.New("error dialing connection to mlmd server")
+	ErrMLMDTypeCreation    = errors.New("error creating MLMD types")
+	ErrMLMDCoreCreation    = errors.New("error creating core service")
+	ErrMLMDConnectionClose = errors.New("error closing connection to mlmd server")
+)
+
+type MLMDService struct {
+	Hostname       string
+	Port           int
+	GRPCConnection *grpc.ClientConn
+}
+
+func NewMLMDService(hostname string, port int) *MLMDService {
+	return &MLMDService{
+		Hostname: hostname,
+		Port:     port,
+	}
+}
+
+func (s *MLMDService) Build() (api.ModelRegistryApi, error) {
+	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Second*30)
+
+	defer cancel()
+
+	mlmdAddr := fmt.Sprintf("%s:%d", s.Hostname, s.Port)
+
+	glog.Infof("connecting to MLMD server %s..", mlmdAddr)
+
+	conn, err := grpc.DialContext( // nolint:staticcheck
+		ctxTimeout,
+		mlmdAddr,
+		grpc.WithReturnConnectionError(), // nolint:staticcheck
+		grpc.WithBlock(),                 // nolint:staticcheck
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("%w %s: %w", ErrMLMDConnectionStart, mlmdAddr, err)
+	}
+
+	s.GRPCConnection = conn
+
+	glog.Infof("connected to MLMD server")
+
+	mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
+
+	if _, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig); err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrMLMDTypeCreation, err)
+	}
+
+	service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrMLMDCoreCreation, err)
+	}
+
+	return service, nil
+}
+
+func (s *MLMDService) Teardown() error {
+	glog.Infof("closing connection to MLMD server")
+
+	if err := s.GRPCConnection.Close(); err != nil {
+		return fmt.Errorf("%w: %w", ErrMLMDConnectionClose, err)
+	}
+
+	return nil
+}

manifests/kustomize/base/model-registry-configmap.yaml

@@ -9,3 +9,4 @@ data:
   MODEL_REGISTRY_REST_SERVICE_PORT: "8080"
   MODEL_REGISTRY_GRPC_SERVICE_HOST: "model-registry-service"
   MODEL_REGISTRY_GRPC_SERVICE_PORT: "9090"
+  MODEL_REGISTRY_DATA_STORE_TYPE: "mlmd"

manifests/kustomize/base/model-registry-deployment.yaml

@@ -21,8 +21,9 @@ spec:
           args:
             - --hostname=0.0.0.0
             - --port=8080
-            - --mlmd-hostname=localhost
-            - --mlmd-port=9090
+            - --datastore-hostname=localhost
+            - --datastore-port=9090
+            - --datastore-type=mlmd
           command:
             - /model-registry
             - proxy

manifests/kustomize/overlays/db/patches/model-registry-deployment.yaml

@@ -20,8 +20,9 @@ spec:
           args:
             - --hostname=0.0.0.0
             - --port=$(MODEL_REGISTRY_REST_SERVICE_PORT)
-            - --mlmd-hostname=localhost
-            - --mlmd-port=$(MODEL_REGISTRY_GRPC_SERVICE_PORT)
+            - --datastore-hostname=localhost
+            - --datastore-port=$(MODEL_REGISTRY_GRPC_SERVICE_PORT)
+            - --datastore-type=$(MODEL_REGISTRY_DATA_STORE_TYPE)
         - name: grpc-container
           # Remove existing environment variables
           env: