Repository: G-Research/fasttrackml
Pull request #378 (2023-09-21 15:54:40)

Support multiple artifact storage backends

dsuhinin: > Just scanned it very quickly and it looks to me like we will create a new storage provider on every request to list artifacts. For the S3 provider, I believe this would mean recreating a client (and thus connections) every time.
> 
> I think we should instead register each storage provider with a URL scheme and instantiate them (lazily?) in the "factory" so they can be returned based on the scheme of the requested URI. So a proxy instead of a factory I guess.

emmm, I understand what you mean. in general even now there is no heavy lifting but let me think.

Patch:
diff --git a/pkg/api/mlflow/config/config.go b/pkg/api/mlflow/config/config.go
index d065973..1c0776b 100644
--- a/pkg/api/mlflow/config/config.go
+++ b/pkg/api/mlflow/config/config.go
@@ -18,7 +18,7 @@ type ServiceConfig struct {
 	ListenAddress         string
 	AuthUsername          string
 	AuthPassword          string
-	ArtifactRoot          string
+	DefaultArtifactRoot   string
 	S3EndpointURI         string
 	DatabaseURI           string
 	DatabaseReset         bool
@@ -34,7 +34,7 @@ func NewServiceConfig() *ServiceConfig {
 		ListenAddress:         viper.GetString("listen-address"),
 		AuthUsername:          viper.GetString("auth-username"),
 		AuthPassword:          viper.GetString("auth-password"),
-		ArtifactRoot:          viper.GetString("artifact-root"),
+		DefaultArtifactRoot:   viper.GetString("default-artifact-root"),
 		S3EndpointURI:         viper.GetString("s3-endpoint-uri"),
 		DatabaseURI:           viper.GetString("database-uri"),
 		DatabaseReset:         viper.GetBool("database-reset"),
@@ -57,18 +57,18 @@ func (c *ServiceConfig) Validate() error {
 
 // validateConfiguration validates service configuration for correctness.
 func (c *ServiceConfig) validateConfiguration() error {
-	// 1. validate ArtifactRoot configuration parameter for correctness and valid values.
-	parsed, err := url.Parse(c.ArtifactRoot)
+	// 1. validate DefaultArtifactRoot configuration parameter for correctness and valid values.
+	parsed, err := url.Parse(c.DefaultArtifactRoot)
 	if err != nil {
-		return eris.Wrap(err, "error parsing 'artifact-root' flag")
+		return eris.Wrap(err, "error parsing 'default-artifact-root' flag")
 	}
 
 	if parsed.User != nil || parsed.RawQuery != "" || parsed.RawFragment != "" {
-		return eris.New("incorrect format of 'artifact-root' flag")
+		return eris.New("incorrect format of 'default-artifact-root' flag")
 	}
 
 	if !slices.Contains([]string{"", "file", "s3"}, parsed.Scheme) {
-		return eris.New("unsupported schema of 'artifact-root' flag")
+		return eris.New("unsupported schema of 'default-artifact-root' flag")
 	}
 
 	return nil
@@ -76,9 +76,9 @@ func (c *ServiceConfig) validateConfiguration() error {
 
 // normalizeConfiguration normalizes service configuration parameters.
 func (c *ServiceConfig) normalizeConfiguration() error {
-	parsed, err := url.Parse(c.ArtifactRoot)
+	parsed, err := url.Parse(c.DefaultArtifactRoot)
 	if err != nil {
-		return eris.Wrap(err, "error parsing 'artifact-root' flag")
+		return eris.Wrap(err, "error parsing 'default-artifact-root' flag")
 	}
 	switch parsed.Scheme {
 	case "s3":
@@ -86,9 +86,9 @@ func (c *ServiceConfig) normalizeConfiguration() error {
 	case "", "file":
 		absoluteArtifactRoot, err := filepath.Abs(path.Join(parsed.Host, parsed.Path))
 		if err != nil {
-			return eris.Wrapf(err, "error getting absolute path for 'artifact-root': %s", c.ArtifactRoot)
+			return eris.Wrapf(err, "error getting absolute path for 'default-artifact-root': %s", c.DefaultArtifactRoot)
 		}
-		c.ArtifactRoot = absoluteArtifactRoot
+		c.DefaultArtifactRoot = absoluteArtifactRoot
 	}
 	return nil
 }
diff --git a/pkg/api/mlflow/service/artifact/service.go b/pkg/api/mlflow/service/artifact/service.go
index 64c5c8a..271e7db 100644
--- a/pkg/api/mlflow/service/artifact/service.go
+++ b/pkg/api/mlflow/service/artifact/service.go
@@ -11,15 +11,18 @@ import (
 
 // Service provides service layer to work with `artifact` business logic.
 type Service struct {
-	runRepository   repositories.RunRepositoryProvider
-	artifactStorage storage.Provider
+	runRepository          repositories.RunRepositoryProvider
+	artifactStorageFactory storage.ArtifactStorageFactoryProvider
 }
 
 // NewService creates new Service instance.
-func NewService(artifactStorage storage.Provider, runRepository repositories.RunRepositoryProvider) *Service {
+func NewService(
+	runRepository repositories.RunRepositoryProvider,
+	artifactStorageFactory storage.ArtifactStorageFactoryProvider,
+) *Service {
 	return &Service{
-		runRepository:   runRepository,
-		artifactStorage: artifactStorage,
+		runRepository:          runRepository,
+		artifactStorageFactory: artifactStorageFactory,
 	}
 }
 
@@ -36,9 +39,11 @@ func (s Service) ListArtifacts(
 		return "", nil, api.NewInternalError("unable to get artifact URI for run '%s'", req.GetRunID())
 	}
 
-	rootURI, artifacts, err := s.artifactStorage.List(
-		run.ArtifactURI, req.Path,
-	)
+	artifactStorage, err := s.artifactStorageFactory.GetStorage(run.ArtifactURI)
+	if err != nil {
+		return "", nil, api.NewInternalError("run with id '%s' has unsupported artifact storage", run.ID)
+	}
+	rootURI, artifacts, err := artifactStorage.List(run.ArtifactURI, req.Path)
 	if err != nil {
 		return "", nil, api.NewInternalError("error getting artifact list from storage")
 	}
diff --git a/pkg/api/mlflow/service/artifact/storage/local.go b/pkg/api/mlflow/service/artifact/storage/local.go
index fff3a95..c3f3d6a 100644
--- a/pkg/api/mlflow/service/artifact/storage/local.go
+++ b/pkg/api/mlflow/service/artifact/storage/local.go
@@ -11,6 +11,11 @@ import (
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/config"
 )
 
+// LocalStorageName is a file storage name.
+const (
+	LocalStorageName = "file"
+)
+
 // Local represents local file storage adapter to work with artifacts.
 type Local struct {
 	config *config.ServiceConfig
@@ -23,7 +28,7 @@ func NewLocal(config *config.ServiceConfig) (*Local, error) {
 	}, nil
 }
 
-// List implements Provider interface.
+// List implements ArtifactStorageProvider interface.
 func (s Local) List(artifactURI, path string) (string, []ArtifactObject, error) {
 	// 1. process search `prefix` parameter.
 	path, err := url.JoinPath(artifactURI, path)
@@ -50,5 +55,5 @@ func (s Local) List(artifactURI, path string) (string, []ArtifactObject, error)
 			IsDir: object.IsDir(),
 		}
 	}
-	return s.config.ArtifactRoot, artifactList, nil
+	return s.config.DefaultArtifactRoot, artifactList, nil
 }
diff --git a/pkg/api/mlflow/service/artifact/storage/noop.go b/pkg/api/mlflow/service/artifact/storage/noop.go
deleted file mode 100644
index cf5b563..0000000
--- a/pkg/api/mlflow/service/artifact/storage/noop.go
+++ /dev/null
@@ -1,14 +0,0 @@
-package storage
-
-// Noop represents empty storage adapter.
-type Noop struct{}
-
-// NewNoop creates new Noop instance.
-func NewNoop() *Noop {
-	return &Noop{}
-}
-
-// List implements Provider interface.
-func (s Noop) List(_, _ string) (string, []ArtifactObject, error) {
-	return "", make([]ArtifactObject, 0), nil
-}
diff --git a/pkg/api/mlflow/service/artifact/storage/s3.go b/pkg/api/mlflow/service/artifact/storage/s3.go
index 7683617..e6d1aaa 100644
--- a/pkg/api/mlflow/service/artifact/storage/s3.go
+++ b/pkg/api/mlflow/service/artifact/storage/s3.go
@@ -14,6 +14,11 @@ import (
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/config"
 )
 
+// S3StorageName is a s3 storage name.
+const (
+	S3StorageName = "s3"
+)
+
 // S3 represents S3 adapter to work with artifacts.
 type S3 struct {
 	client *s3.Client
@@ -55,7 +60,7 @@ func NewS3(config *config.ServiceConfig) (*S3, error) {
 	return &storage, nil
 }
 
-// List implements Provider interface.
+// List implements ArtifactStorageProvider interface.
 func (s S3) List(artifactURI, path string) (string, []ArtifactObject, error) {
 	bucket, prefix, err := ExtractS3BucketAndPrefix(artifactURI)
 	if err != nil {
diff --git a/pkg/api/mlflow/service/artifact/storage/storage.go b/pkg/api/mlflow/service/artifact/storage/storage.go
index 30e0ed7..18f86b1 100644
--- a/pkg/api/mlflow/service/artifact/storage/storage.go
+++ b/pkg/api/mlflow/service/artifact/storage/storage.go
@@ -2,6 +2,7 @@ package storage
 
 import (
 	"net/url"
+	"sync"
 
 	"github.com/rotisserie/eris"
 
@@ -30,27 +31,67 @@ func (o ArtifactObject) IsDirectory() bool {
 	return o.IsDir
 }
 
-// Provider provides and interface to work with artifact storage.
-type Provider interface {
+// ArtifactStorageProvider provides and interface to work with particular artifact storage.
+type ArtifactStorageProvider interface {
+	// List lists all artifact object under provided path.
 	List(artifactURI, path string) (string, []ArtifactObject, error)
 }
 
-// NewArtifactStorage creates new Artifact storage.
-func NewArtifactStorage(config *config.ServiceConfig) (Provider, error) {
-	if config.ArtifactRoot != "" {
-		u, err := url.Parse(config.ArtifactRoot)
+// ArtifactStorageFactoryProvider provides an interface to work with Artifact Storage.
+type ArtifactStorageFactoryProvider interface {
+	// GetStorage returns Artifact storage based on provided runArtifactPath.
+	GetStorage(runArtifactPath string) (ArtifactStorageProvider, error)
+}
+
+// ArtifactStorageFactory represents Artifact Storage .
+type ArtifactStorageFactory struct {
+	config      *config.ServiceConfig
+	storageList sync.Map
+}
+
+// NewArtifactStorageFactory creates new Artifact Storage Factory instance.
+func NewArtifactStorageFactory(config *config.ServiceConfig) (*ArtifactStorageFactory, error) {
+	return &ArtifactStorageFactory{
+		config:      config,
+		storageList: sync.Map{},
+	}, nil
+}
+
+// GetStorage returns Artifact storage based on provided runArtifactPath.
+func (s *ArtifactStorageFactory) GetStorage(runArtifactPath string) (ArtifactStorageProvider, error) {
+	u, err := url.Parse(runArtifactPath)
+	if err != nil {
+		return nil, eris.Wrap(err, "error parsing artifact root")
+	}
+
+	switch u.Scheme {
+	case S3StorageName:
+		if storage, ok := s.storageList.Load(S3StorageName); ok {
+			if localStorage, ok := storage.(*S3); ok {
+				return localStorage, nil
+			}
+			return nil, eris.New("storage is not s3 artifact storage")
+		}
+		storage, err := NewS3(s.config)
 		if err != nil {
-			return nil, eris.Wrap(err, "error parsing artifact root")
+			return nil, eris.Wrap(err, "error initializing s3 artifact storage")
 		}
-
-		switch u.Scheme {
-		case "s3":
-			return NewS3(config)
-		case "", "file":
-			return NewLocal(config)
-		default:
-			return nil, eris.Errorf("unsupported schema has been provided: %s", u.Scheme)
+		s.storageList.Store(S3StorageName, storage)
+		return storage, nil
+	case "", LocalStorageName:
+		if storage, ok := s.storageList.Load(LocalStorageName); ok {
+			if localStorage, ok := storage.(*Local); ok {
+				return localStorage, nil
+			}
+			return nil, eris.New("storage is not local artifact storage")
+		}
+		storage, err := NewLocal(s.config)
+		if err != nil {
+			return nil, eris.Wrap(err, "error initializing local artifact storage")
 		}
+		s.storageList.Store(LocalStorageName, storage)
+		return storage, nil
+	default:
+		return nil, eris.Errorf("unsupported schema has been provided: %s", u.Scheme)
 	}
-	return NewNoop(), nil
 }
diff --git a/pkg/api/mlflow/service/experiment/service.go b/pkg/api/mlflow/service/experiment/service.go
index 922eda5..bba0d16 100644
--- a/pkg/api/mlflow/service/experiment/service.go
+++ b/pkg/api/mlflow/service/experiment/service.go
@@ -74,7 +74,7 @@ func (s Service) CreateExperiment(
 	}
 
 	if experiment.ArtifactLocation == "" {
-		path, err := url.JoinPath(s.config.ArtifactRoot, fmt.Sprintf("%d", *experiment.ID))
+		path, err := url.JoinPath(s.config.DefaultArtifactRoot, fmt.Sprintf("%d", *experiment.ID))
 		if err != nil {
 			return nil, api.NewInternalError(
 				"error creating artifact_location for experiment'%s': %s", experiment.Name, err,
diff --git a/pkg/cmd/import.go b/pkg/cmd/import.go
index 8980f32..cc6c95c 100644
--- a/pkg/cmd/import.go
+++ b/pkg/cmd/import.go
@@ -41,14 +41,14 @@ func initDBs() (input, output database.DBProvider, err error) {
 	databasePoolMax := 20
 	databaseReset := false
 	databaseMigrate := false
-	artifactRoot := viper.GetString("artifact-root")
+	defaultArtifactRoot := viper.GetString("default-artifact-root")
 	input, err = database.MakeDBProvider(
 		viper.GetString("input-database-uri"),
 		databaseSlowThreshold,
 		databasePoolMax,
 		databaseReset,
 		databaseMigrate,
-		artifactRoot,
+		defaultArtifactRoot,
 	)
 	if err != nil {
 		return input, output, fmt.Errorf("error connecting to input DB: %w", err)
@@ -61,7 +61,7 @@ func initDBs() (input, output database.DBProvider, err error) {
 		databasePoolMax,
 		databaseReset,
 		databaseMigrate,
-		artifactRoot,
+		defaultArtifactRoot,
 	)
 	if err != nil {
 		return input, output, fmt.Errorf("error connecting to output DB: %w", err)
@@ -74,7 +74,7 @@ func init() {
 
 	ImportCmd.Flags().StringP("input-database-uri", "i", "", "Input Database URI (eg., sqlite://fasttrackml.db)")
 	ImportCmd.Flags().StringP("output-database-uri", "o", "", "Output Database URI (eg., postgres://user:psw@postgres:5432)")
-	ImportCmd.Flags().StringP("artifact-root", "a", "./artifacts", "Artifact Root")
+	ImportCmd.Flags().StringP("default-artifact-root", "a", "./artifacts", "Artifact Root")
 	ImportCmd.MarkFlagRequired("input-database-uri")
 	ImportCmd.MarkFlagRequired("output-database-uri")
 }
diff --git a/pkg/cmd/server.go b/pkg/cmd/server.go
index e2dbb8f..9adb37b 100644
--- a/pkg/cmd/server.go
+++ b/pkg/cmd/server.go
@@ -14,7 +14,6 @@ import (
 	"github.com/gofiber/fiber/v2/middleware/cors"
 	"github.com/gofiber/fiber/v2/middleware/logger"
 	"github.com/gofiber/fiber/v2/middleware/recover"
-	"github.com/rotisserie/eris"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -65,9 +64,9 @@ func serverCmd(cmd *cobra.Command, args []string) error {
 	aimAPI.AddRoutes(server.Group("/aim/api/"))
 	aimUI.AddRoutes(server)
 
-	storage, err := storage.NewArtifactStorage(mlflowConfig)
+	artifactStorageFactory, err := storage.NewArtifactStorageFactory(mlflowConfig)
 	if err != nil {
-		return eris.Wrap(err, "error initializing artifact storage")
+		return err
 	}
 
 	// 5. init `mlflow` api and ui routes.
@@ -86,8 +85,8 @@ func serverCmd(cmd *cobra.Command, args []string) error {
 				repositories.NewMetricRepository(db.GormDB()),
 			),
 			artifact.NewService(
-				storage,
 				repositories.NewRunRepository(db.GormDB()),
+				artifactStorageFactory,
 			),
 			experiment.NewService(
 				mlflowConfig,
@@ -131,7 +130,7 @@ func initDB(config *mlflowConfig.ServiceConfig) (database.DBProvider, error) {
 		config.DatabasePoolMax,
 		config.DatabaseReset,
 		config.DatabaseMigrate,
-		config.ArtifactRoot,
+		config.DefaultArtifactRoot,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error connecting to DB: %w", err)
@@ -208,7 +207,7 @@ func init() {
 	RootCmd.AddCommand(ServerCmd)
 
 	ServerCmd.Flags().StringP("listen-address", "a", "localhost:5000", "Address (host:post) to listen to")
-	ServerCmd.Flags().String("artifact-root", "./artifacts", "Artifact root")
+	ServerCmd.Flags().String("default-artifact-root", "./artifacts", "Default artifact root")
 	ServerCmd.Flags().String("s3-endpoint-uri", "", "S3 compatible storage base endpoint url")
 	ServerCmd.Flags().String("auth-username", "", "BasicAuth username")
 	ServerCmd.Flags().String("auth-password", "", "BasicAuth password")
diff --git a/pkg/database/db_factory.go b/pkg/database/db_factory.go
index 293bb32..0651bcc 100644
--- a/pkg/database/db_factory.go
+++ b/pkg/database/db_factory.go
@@ -10,7 +10,7 @@ import (
 
 // MakeDBProvider will create a DbProvider of the correct type from the parameters.
 func MakeDBProvider(
-	dsn string, slowThreshold time.Duration, poolMax int, reset bool, migrate bool, artifactRoot string,
+	dsn string, slowThreshold time.Duration, poolMax int, reset bool, migrate bool, defaultArtifactRoot string,
 ) (db DBProvider, err error) {
 	dsnURL, err := url.Parse(dsn)
 	if err != nil {
@@ -55,7 +55,7 @@ func MakeDBProvider(
 		return nil, err
 	}
 
-	if err := createDefaultExperiment(artifactRoot, db); err != nil {
+	if err := createDefaultExperiment(defaultArtifactRoot, db); err != nil {
 		db.Close()
 		return nil, err
 	}
diff --git a/pkg/database/db_instance.go b/pkg/database/db_instance.go
index 9b2c4c6..ce1a5d3 100644
--- a/pkg/database/db_instance.go
+++ b/pkg/database/db_instance.go
@@ -52,7 +52,7 @@ func (db *DBInstance) GormDB() *gorm.DB {
 }
 
 // createDefaultExperiment will create the default experiment if needed.
-func createDefaultExperiment(artifactRoot string, db DBProvider) error {
+func createDefaultExperiment(defaultArtifactRoot string, db DBProvider) error {
 	if tx := db.GormDB().First(&Experiment{}, 0); tx.Error != nil {
 		if errors.Is(tx.Error, gorm.ErrRecordNotFound) {
 			log.Info("Creating default experiment")
@@ -75,7 +75,7 @@ func createDefaultExperiment(artifactRoot string, db DBProvider) error {
 				return fmt.Errorf("error creating default experiment: %s", tx.Error)
 			}
 
-			exp.ArtifactLocation = fmt.Sprintf("%s/%d", strings.TrimRight(artifactRoot, "/"), *exp.ID)
+			exp.ArtifactLocation = fmt.Sprintf("%s/%d", strings.TrimRight(defaultArtifactRoot, "/"), *exp.ID)
 			if tx := db.GormDB().Model(&exp).Update("ArtifactLocation", exp.ArtifactLocation); tx.Error != nil {
 				return fmt.Errorf("error updating artifact_location for experiment '%s': %s", exp.Name, tx.Error)
 			}


Test patch:
diff --git a/pkg/api/mlflow/config/config_test.go b/pkg/api/mlflow/config/config_test.go
index 8bea946..567965e 100644
--- a/pkg/api/mlflow/config/config_test.go
+++ b/pkg/api/mlflow/config/config_test.go
@@ -15,21 +15,21 @@ func TestServiceConfig_Validate_Ok(t *testing.T) {
 		expectedConfig *ServiceConfig
 	}{
 		{
-			name: "ArtifactRootHasS3Prefix",
+			name: "DefaultArtifactRootHasS3Prefix",
 			providedConfig: &ServiceConfig{
-				ArtifactRoot: "s3://bucket_name",
+				DefaultArtifactRoot: "s3://bucket_name",
 			},
 			expectedConfig: &ServiceConfig{
-				ArtifactRoot: "s3://bucket_name",
+				DefaultArtifactRoot: "s3://bucket_name",
 			},
 		},
 		{
-			name: "ArtifactRootHasFilePrefixAndIsRelative",
+			name: "DefaultArtifactRootHasFilePrefixAndIsRelative",
 			providedConfig: &ServiceConfig{
-				ArtifactRoot: "file://path1/path2/path3",
+				DefaultArtifactRoot: "file://path1/path2/path3",
 			},
 			expectedConfig: &ServiceConfig{
-				ArtifactRoot: (func() string {
+				DefaultArtifactRoot: (func() string {
 					path, err := filepath.Abs("path1/path2/path3")
 					assert.Nil(t, err)
 					return path
@@ -37,30 +37,30 @@ func TestServiceConfig_Validate_Ok(t *testing.T) {
 			},
 		},
 		{
-			name: "ArtifactRootHasFilePrefixAndIsAbsolute",
+			name: "DefaultArtifactRootHasFilePrefixAndIsAbsolute",
 			providedConfig: &ServiceConfig{
-				ArtifactRoot: "file:///path1/path2/path3",
+				DefaultArtifactRoot: "file:///path1/path2/path3",
 			},
 			expectedConfig: &ServiceConfig{
-				ArtifactRoot: "/path1/path2/path3",
+				DefaultArtifactRoot: "/path1/path2/path3",
 			},
 		},
 		{
-			name: "ArtifactRootHasEmptyPrefixAndIsAbsolute",
+			name: "DefaultArtifactRootHasEmptyPrefixAndIsAbsolute",
 			providedConfig: &ServiceConfig{
-				ArtifactRoot: "/path1/path2/path3",
+				DefaultArtifactRoot: "/path1/path2/path3",
 			},
 			expectedConfig: &ServiceConfig{
-				ArtifactRoot: "/path1/path2/path3",
+				DefaultArtifactRoot: "/path1/path2/path3",
 			},
 		},
 		{
-			name: "ArtifactRootHasEmptyPrefixAndIsRelative",
+			name: "DefaultArtifactRootHasEmptyPrefixAndIsRelative",
 			providedConfig: &ServiceConfig{
-				ArtifactRoot: "path1/path2/path3",
+				DefaultArtifactRoot: "path1/path2/path3",
 			},
 			expectedConfig: &ServiceConfig{
-				ArtifactRoot: (func() string {
+				DefaultArtifactRoot: (func() string {
 					path, err := filepath.Abs("path1/path2/path3")
 					assert.Nil(t, err)
 					return path
@@ -72,7 +72,7 @@ func TestServiceConfig_Validate_Ok(t *testing.T) {
 	for _, tt := range testData {
 		t.Run(tt.name, func(t *testing.T) {
 			assert.Nil(t, tt.providedConfig.Validate())
-			assert.Equal(t, tt.providedConfig.ArtifactRoot, tt.expectedConfig.ArtifactRoot)
+			assert.Equal(t, tt.providedConfig.DefaultArtifactRoot, tt.expectedConfig.DefaultArtifactRoot)
 		})
 	}
 }
@@ -84,19 +84,19 @@ func TestServiceConfig_Validate_Error(t *testing.T) {
 		config *ServiceConfig
 	}{
 		{
-			name: "ArtifactRootHasIncorrectFormat",
+			name: "DefaultArtifactRootHasIncorrectFormat",
 			error: eris.New(
-				`error validating service configuration: error parsing 'artifact-root' flag: parse "incorrect_format_of_schema://something": first path segment in URL cannot contain colon`,
+				`error validating service configuration: error parsing 'default-artifact-root' flag: parse "incorrect_format_of_schema://something": first path segment in URL cannot contain colon`,
 			),
 			config: &ServiceConfig{
-				ArtifactRoot: "incorrect_format_of_schema://something",
+				DefaultArtifactRoot: "incorrect_format_of_schema://something",
 			},
 		},
 		{
-			name:  "ArtifactRootHasUnsupportedSchema",
-			error: eris.New("error validating service configuration: unsupported schema of 'artifact-root' flag"),
+			name:  "DefaultArtifactRootHasUnsupportedSchema",
+			error: eris.New("error validating service configuration: unsupported schema of 'default-artifact-root' flag"),
 			config: &ServiceConfig{
-				ArtifactRoot: "unsupported://something",
+				DefaultArtifactRoot: "unsupported://something",
 			},
 		},
 	}
diff --git a/pkg/api/mlflow/service/artifact/service_test.go b/pkg/api/mlflow/service/artifact/service_test.go
index f9c4707..4ffa0e3 100644
--- a/pkg/api/mlflow/service/artifact/service_test.go
+++ b/pkg/api/mlflow/service/artifact/service_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 func TestService_ListArtifacts_Ok(t *testing.T) {
-	artifactStorage := storage.MockProvider{}
+	artifactStorage := storage.MockArtifactStorageProvider{}
 	artifactStorage.On(
 		"List", "/artifact/uri", "",
 	).Return(
@@ -34,6 +34,11 @@ func TestService_ListArtifacts_Ok(t *testing.T) {
 		}, nil,
 	)
 
+	artifactStorageFactory := storage.MockArtifactStorageFactoryProvider{}
+	artifactStorageFactory.On(
+		"GetStorage", "/artifact/uri",
+	).Return(&artifactStorage, nil)
+
 	// init repository mocks.
 	runRepository := repositories.MockRunRepositoryProvider{}
 	runRepository.On(
@@ -46,7 +51,7 @@ func TestService_ListArtifacts_Ok(t *testing.T) {
 	}, nil)
 
 	// call service under testing.
-	service := NewService(&artifactStorage, &runRepository)
+	service := NewService(&runRepository, &artifactStorageFactory)
 	rootURI, artifacts, err := service.ListArtifacts(
 		context.TODO(),
 		&request.ListArtifactsRequest{
@@ -83,8 +88,8 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 			request: &request.ListArtifactsRequest{},
 			service: func() *Service {
 				return NewService(
-					&storage.MockProvider{},
 					&repositories.MockRunRepositoryProvider{},
+					&storage.MockArtifactStorageFactoryProvider{},
 				)
 			},
 		},
@@ -97,8 +102,8 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 			},
 			service: func() *Service {
 				return NewService(
-					&storage.MockProvider{},
 					&repositories.MockRunRepositoryProvider{},
+					&storage.MockArtifactStorageFactoryProvider{},
 				)
 			},
 		},
@@ -116,8 +121,8 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 					"id",
 				).Return(nil, errors.New("database error"))
 				return NewService(
-					&storage.MockProvider{},
 					&runRepository,
+					&storage.MockArtifactStorageFactoryProvider{},
 				)
 			},
 		},
@@ -128,13 +133,18 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 				RunID: "id",
 			},
 			service: func() *Service {
-				artifactStorage := storage.MockProvider{}
+				artifactStorage := storage.MockArtifactStorageProvider{}
 				artifactStorage.On(
 					"List", "/artifact/uri", "",
 				).Return(
 					"", nil, errors.New("storage error"),
 				)
 
+				artifactStorageFactory := storage.MockArtifactStorageFactoryProvider{}
+				artifactStorageFactory.On(
+					"GetStorage", "/artifact/uri",
+				).Return(&artifactStorage, nil)
+
 				runRepository := repositories.MockRunRepositoryProvider{}
 				runRepository.On(
 					"GetByID",
@@ -145,8 +155,8 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 					ArtifactURI: "/artifact/uri",
 				}, nil)
 				return NewService(
-					&artifactStorage,
 					&runRepository,
+					&artifactStorageFactory,
 				)
 			},
 		},
diff --git a/pkg/database/db_factory_test.go b/pkg/database/db_factory_test.go
index 8642eec..3252a08 100644
--- a/pkg/database/db_factory_test.go
+++ b/pkg/database/db_factory_test.go
@@ -36,9 +36,7 @@ func TestMakeDBProvider(t *testing.T) {
 
 			// expecting the global 'DB' not to be set
 			assert.Nil(t, DB)
-
-			err = db.Close()
-			assert.Nil(t, err)
+			assert.Nil(t, db.Close())
 		})
 	}
 }
diff --git a/tests/integration/golang/database/import_test.go b/tests/integration/golang/database/import_test.go
index 387cc92..eb8d35b 100644
--- a/tests/integration/golang/database/import_test.go
+++ b/tests/integration/golang/database/import_test.go
@@ -102,14 +102,14 @@ func (s *ImportTestSuite) SetupTest() {
 	databasePoolMax := 20
 	databaseReset := false
 	databaseMigrate := false
-	artifactRoot := "./artifacts"
+	defaultArtifactRoot := "./artifacts"
 	input, err := database.MakeDBProvider(
 		helpers.GetInputDatabaseUri(),
 		databaseSlowThreshold,
 		databasePoolMax,
 		databaseReset,
 		databaseMigrate,
-		artifactRoot,
+		defaultArtifactRoot,
 	)
 	assert.Nil(s.T(), err)
 	output, err := database.MakeDBProvider(
@@ -118,7 +118,7 @@ func (s *ImportTestSuite) SetupTest() {
 		databasePoolMax,
 		databaseReset,
 		databaseMigrate,
-		artifactRoot,
+		defaultArtifactRoot,
 	)
 	assert.Nil(s.T(), err)