refactor(GenerateOutputURI): move and preserve comments

kubeflow · Dec 16, 2023 · e2c34bf · e2c34bf
1 parent 0f4044c
commit e2c34bf
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 23 deletions.
diff --git a/backend/src/v2/driver/driver.go b/backend/src/v2/driver/driver.go
@@ -17,9 +17,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"path"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/golang/glog"
@@ -1062,8 +1060,9 @@ func provisionOutputs(pipelineRoot, taskName string, outputsSpec *pipelinespec.C
 		outputs.Artifacts[name] = &pipelinespec.ArtifactList{
 			Artifacts: []*pipelinespec.RuntimeArtifact{
 				{
-					// Split off any query strings first
-					Uri:      fmt.Sprintf("%s/%s", strings.TrimRight(strings.Split(pipelineRoot, "?")[0], "/"), path.Join(taskName, name)),
+					// Do not preserve the query string for output artifacts, as otherwise
+					// they'd appear in file and artifact names.
+					Uri:      metadata.GenerateOutputURI(pipelineRoot, []string{taskName, name}, false),
 					Type:     artifact.GetArtifactType(),
 					Metadata: artifact.GetMetadata(),
 				},

diff --git a/backend/src/v2/metadata/client.go b/backend/src/v2/metadata/client.go
@@ -260,17 +260,17 @@ func (e *Execution) FingerPrint() string {
 	return e.execution.GetCustomProperties()[keyCacheFingerPrint].GetStringValue()
 }
 
-// AppendToPipelineRoot appends the specified paths to the pipeline root.
-// It preserves the query part of the pipeline root by splitting it off and
-// appending it back to the full URI.
-func AppendToPipelineRoot(pipelineRoot string, paths []string) string {
-	// split the query part off the root, if it exists.
-	// we will append it back later to the full URI.
+// GenerateOutputURI appends the specified paths to the pipeline root.
+// It may be configured to preserve the query part of the pipeline root
+// by splitting it off and appending it back to the full URI.
+func GenerateOutputURI(pipelineRoot string, paths []string, preserveQueryString bool) string {
 	querySplit := strings.Split(pipelineRoot, "?")
 	query := ""
-	if len(querySplit) > 1 {
+	if len(querySplit) == 2 {
 		pipelineRoot = querySplit[0]
-		query = "?" + querySplit[1]
+		if preserveQueryString {
+			query = "?" + querySplit[1]
+		}
 	} else if len(querySplit) > 2 {
 		// this should never happen, but just in case.
 		glog.Warningf("Unexpected pipeline root: %v", pipelineRoot)
@@ -292,7 +292,7 @@ func (c *Client) GetPipeline(ctx context.Context, pipelineName, runID, namespace
 		keyNamespace:    stringValue(namespace),
 		keyResourceName: stringValue(runResource),
 		// pipeline root of this run
-		keyPipelineRoot: stringValue(AppendToPipelineRoot(pipelineRoot, []string{pipelineName, runID})),
+		keyPipelineRoot: stringValue(GenerateOutputURI(pipelineRoot, []string{pipelineName, runID}, true)),
 	}
 	runContext, err := c.getOrInsertContext(ctx, runID, pipelineRunContextType, metadata)
 	glog.Infof("Pipeline Run Context: %+v", runContext)

diff --git a/backend/src/v2/metadata/client_test.go b/backend/src/v2/metadata/client_test.go
@@ -214,24 +214,57 @@ func Test_GetPipelineConcurrently(t *testing.T) {
 	wg.Wait()
 }
 
-func Test_appendToPipelineRoot(t *testing.T) {
+func Test_GenerateOutputURI(t *testing.T) {
 	// Const define the artifact name
 	const (
 		pipelineName      = "my-pipeline-name"
 		runID             = "my-run-id"
 		pipelineRoot      = "minio://mlpipeline/v2/artifacts"
 		pipelineRootQuery = "?query=string&another=query"
 	)
-	// Test output uri generation with plain pipeline root, i.e. without any query strings
-	plainResult := metadata.AppendToPipelineRoot(pipelineRoot, []string{pipelineName, runID})
-	if plainResult != fmt.Sprintf("%s/%s/%s", pipelineRoot, pipelineName, runID) {
-		t.Errorf("Expected %s, got %s", fmt.Sprintf("%s/%s/%s", pipelineRoot, pipelineName, runID), plainResult)
+	tests := []struct {
+		name                string
+		queryString         string
+		paths               []string
+		preserveQueryString bool
+		want                string
+	}{
+		{
+			name:                "plain pipeline root without preserveQueryString",
+			queryString:         "",
+			paths:               []string{pipelineName, runID},
+			preserveQueryString: false,
+			want:                fmt.Sprintf("%s/%s/%s", pipelineRoot, pipelineName, runID),
+		},
+		{
+			name:                "plain pipeline root with preserveQueryString",
+			queryString:         "",
+			paths:               []string{pipelineName, runID},
+			preserveQueryString: true,
+			want:                fmt.Sprintf("%s/%s/%s", pipelineRoot, pipelineName, runID),
+		},
+		{
+			name:                "pipeline root with query string without preserveQueryString",
+			queryString:         pipelineRootQuery,
+			paths:               []string{pipelineName, runID},
+			preserveQueryString: false,
+			want:                fmt.Sprintf("%s/%s/%s", pipelineRoot, pipelineName, runID),
+		},
+		{
+			name:                "pipeline root with query string with preserveQueryString",
+			queryString:         pipelineRootQuery,
+			paths:               []string{pipelineName, runID},
+			preserveQueryString: true,
+			want:                fmt.Sprintf("%s/%s/%s%s", pipelineRoot, pipelineName, runID, pipelineRootQuery),
+		},
 	}
-
-	// Make sure query strings in the pipeline root are correctly preserved (added to the end)
-	queryResult := metadata.AppendToPipelineRoot(fmt.Sprintf("%s%s", pipelineRoot, pipelineRootQuery), []string{pipelineName, runID})
-	if queryResult != fmt.Sprintf("%s/%s/%s%s", pipelineRoot, pipelineName, runID, pipelineRootQuery) {
-		t.Errorf("Expected %s, got %s", fmt.Sprintf("%s/%s/%s%s", pipelineRoot, pipelineName, runID, pipelineRootQuery), queryResult)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := metadata.GenerateOutputURI(fmt.Sprintf("%s%s", pipelineRoot, tt.queryString), tt.paths, tt.preserveQueryString)
+			if diff := cmp.Diff(got, tt.want); diff != "" {
+				t.Errorf("GenerateOutputURI() = %v, want %v\nDiff (-want, +got)\n%s", got, tt.want, diff)
+			}
+		})
 	}
 }