From d85b94b417351180768090eabc2e562a3ac13dc4 Mon Sep 17 00:00:00 2001
From: ewezy <ziyi.ewe@gojek.com>
Date: Mon, 9 Sep 2024 15:43:44 +0800
Subject: [PATCH] Add s3 artifact service type as option in mlflow service

---
 api/pkg/client/mlflow/mlflow.go      | 13 ++++++++-----
 api/pkg/client/mlflow/mlflow_test.go | 10 ++++++++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/api/pkg/client/mlflow/mlflow.go b/api/pkg/client/mlflow/mlflow.go
index 5fce5396..7b05050f 100644
--- a/api/pkg/client/mlflow/mlflow.go
+++ b/api/pkg/client/mlflow/mlflow.go
@@ -7,8 +7,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"cloud.google.com/go/storage"
-
 	"github.com/caraml-dev/mlp/api/pkg/artifact"
 )
 
@@ -25,14 +23,20 @@ type mlflowService struct {
 
 func NewMlflowService(httpClient *http.Client, config Config) (Service, error) {
 	var artifactService artifact.Service
+	var err error
+	// TODO: To consider removing since it doesn't make sense to have nop artifact service
 	if config.ArtifactServiceType == "nop" {
 		artifactService = artifact.NewNopArtifactClient()
 	} else if config.ArtifactServiceType == "gcs" {
-		api, err := storage.NewClient(context.Background())
+		artifactService, err = artifact.NewGcsArtifactClient()
 		if err != nil {
 			return &mlflowService{}, fmt.Errorf("failed initializing gcs for mlflow delete package")
 		}
-		artifactService = artifact.NewGcsArtifactClient(api)
+	} else if config.ArtifactServiceType == "s3" {
+		artifactService, err = artifact.NewS3ArtifactClient()
+		if err != nil {
+			return &mlflowService{}, fmt.Errorf("failed initializing s3 for mlflow delete package")
+		}
 	} else {
 		return &mlflowService{}, fmt.Errorf("invalid artifact service type")
 	}
@@ -115,7 +119,6 @@ func (mfs *mlflowService) searchRunData(RunID string) (SearchRunResponse, error)
 }
 
 func (mfs *mlflowService) DeleteExperiment(ctx context.Context, ExperimentID string, deleteArtifact bool) error {
-
 	relatedRunData, err := mfs.searchRunsForExperiment(ExperimentID)
 	if err != nil {
 		return err
diff --git a/api/pkg/client/mlflow/mlflow_test.go b/api/pkg/client/mlflow/mlflow_test.go
index 619cfd26..187f6b00 100644
--- a/api/pkg/client/mlflow/mlflow_test.go
+++ b/api/pkg/client/mlflow/mlflow_test.go
@@ -246,6 +246,16 @@ func TestNewMlflowService(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:          "Mlflow Service with s3 Artifact",
+			artifactType:  "s3",
+			expectedError: nil,
+			expectedResult: &mlflowService{
+				API:             &httpClient,
+				Config:          Config{TrackingURL: "", ArtifactServiceType: "s3"},
+				ArtifactService: &artifact.S3ArtifactClient{},
+			},
+		},
 		{
 			name:          "Mlflow Service with nop Artifact",
 			artifactType:  "nop",