fix(server): use the correct name when downloading artifacts (#4579)

Signed-off-by: Daniel Herman <dherman@factset.com>

pkg/apis/workflow/v1alpha1/workflow_types.go

@@ -924,6 +924,22 @@ func (a *ArtifactLocation) GetType() ArtifactLocationType {
 
 }
 
+func (a *ArtifactLocation) GetKey() string {
+	if a.S3 != nil {
+		return a.S3.Key
+	}
+
+	if a.OSS != nil {
+		return a.OSS.Key
+	}
+
+	if a.GCS != nil {
+		return a.GCS.Key
+	}
+
+	return ""
+}
+
 type ArtifactRepositoryRef struct {
 	ConfigMap string `json:"configMap,omitempty" protobuf:"bytes,1,opt,name=configMap"`
 	Key       string `json:"key,omitempty" protobuf:"bytes,2,opt,name=key"`

server/artifacts/artifact_server.go

@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"os"
+	"path"
 	"strings"
 
 	log "github.com/sirupsen/logrus"
@@ -27,10 +28,15 @@ type ArtifactServer struct {
 	hydrator          hydrator.Interface
 	wfArchive         sqldb.WorkflowArchive
 	instanceIDService instanceid.Service
+	artDriverFactory  artifact.NewDriverFunc
 }
 
 func NewArtifactServer(authN auth.Gatekeeper, hydrator hydrator.Interface, wfArchive sqldb.WorkflowArchive, instanceIDService instanceid.Service) *ArtifactServer {
-	return &ArtifactServer{authN, hydrator, wfArchive, instanceIDService}
+	return newArtifactServer(authN, hydrator, wfArchive, instanceIDService, artifact.NewDriver)
+}
+
+func newArtifactServer(authN auth.Gatekeeper, hydrator hydrator.Interface, wfArchive sqldb.WorkflowArchive, instanceIDService instanceid.Service, artDriverFactory artifact.NewDriverFunc) *ArtifactServer {
+	return &ArtifactServer{authN, hydrator, wfArchive, instanceIDService, artDriverFactory}
 }
 
 func (a *ArtifactServer) GetArtifact(w http.ResponseWriter, r *http.Request) {
@@ -55,12 +61,15 @@ func (a *ArtifactServer) GetArtifact(w http.ResponseWriter, r *http.Request) {
 		a.serverInternalError(err, w)
 		return
 	}
-	data, err := a.getArtifact(ctx, wf, nodeId, artifactName)
+
+	data, filename, err := a.getArtifact(ctx, wf, nodeId, artifactName)
+
 	if err != nil {
 		a.serverInternalError(err, w)
 		return
 	}
-	w.Header().Add("Content-Disposition", fmt.Sprintf(`filename="%s.tgz"`, artifactName))
+
+	w.Header().Add("Content-Disposition", fmt.Sprintf(`filename="%s"`, filename))
 	a.ok(w, data)
 }
 
@@ -87,12 +96,14 @@ func (a *ArtifactServer) GetArtifactByUID(w http.ResponseWriter, r *http.Request
 		return
 	}
 
-	data, err := a.getArtifact(ctx, wf, nodeId, artifactName)
+	data, filename, err := a.getArtifact(ctx, wf, nodeId, artifactName)
+
 	if err != nil {
 		a.serverInternalError(err, w)
 		return
 	}
-	w.Header().Add("Content-Disposition", fmt.Sprintf(`filename="%s.tgz"`, artifactName))
+
+	w.Header().Add("Content-Disposition", fmt.Sprintf(`filename="%s"`, filename))
 	a.ok(w, data)
 }
 
@@ -125,37 +136,37 @@ func (a *ArtifactServer) serverInternalError(err error, w http.ResponseWriter) {
 	_, _ = w.Write([]byte(err.Error()))
 }
 
-func (a *ArtifactServer) getArtifact(ctx context.Context, wf *wfv1.Workflow, nodeId, artifactName string) ([]byte, error) {
+func (a *ArtifactServer) getArtifact(ctx context.Context, wf *wfv1.Workflow, nodeId, artifactName string) ([]byte, string, error) {
 	kubeClient := auth.GetKubeClient(ctx)
 
 	art := wf.Status.Nodes[nodeId].Outputs.GetArtifactByName(artifactName)
 	if art == nil {
-		return nil, fmt.Errorf("artifact not found")
+		return nil, "", fmt.Errorf("artifact not found")
 	}
 
-	driver, err := artifact.NewDriver(art, resources{kubeClient, wf.Namespace})
+	driver, err := a.artDriverFactory(art, resources{kubeClient, wf.Namespace})
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	tmp, err := ioutil.TempFile("/tmp", "artifact")
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
-	path := tmp.Name()
-	defer func() { _ = os.Remove(path) }()
+	tmpPath := tmp.Name()
+	defer func() { _ = os.Remove(tmpPath) }()
 
-	err = driver.Load(art, path)
+	err = driver.Load(art, tmpPath)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 
-	file, err := ioutil.ReadFile(path)
+	file, err := ioutil.ReadFile(tmpPath)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	log.WithFields(log.Fields{"size": len(file)}).Debug("Artifact file size")
 
-	return file, nil
+	return file, path.Base(art.GetKey()), nil
 }
 
 func (a *ArtifactServer) getWorkflowAndValidate(ctx context.Context, namespace string, workflowName string) (*wfv1.Workflow, error) {

server/artifacts/artifact_server_test.go

@@ -2,10 +2,15 @@ package artifacts
 
 import (
 	"context"
+	"fmt"
+	"io/ioutil"
 	"net/http"
 	"net/url"
 	"testing"
 
+	artifact "github.com/argoproj/argo/workflow/artifacts"
+	"github.com/argoproj/argo/workflow/artifacts/resource"
+
 	"github.com/stretchr/testify/assert"
 	testhttp "github.com/stretchr/testify/http"
 	"github.com/stretchr/testify/mock"
@@ -30,6 +35,19 @@ func mustParse(text string) *url.URL {
 	return u
 }
 
+type fakeArtifactDriver struct {
+	artifact.ArtifactDriver
+	data []byte
+}
+
+func (a *fakeArtifactDriver) Load(_ *wfv1.Artifact, path string) error {
+	return ioutil.WriteFile(path, a.data, 0666)
+}
+
+func (a *fakeArtifactDriver) Save(_ string, _ *wfv1.Artifact) error {
+	return fmt.Errorf("not implemented")
+}
+
 func newServer() *ArtifactServer {
 	gatekeeper := &authmocks.Gatekeeper{}
 	kube := kubefake.NewSimpleClientset()
@@ -44,10 +62,26 @@ func newServer() *ArtifactServer {
 					Outputs: &wfv1.Outputs{
 						Artifacts: wfv1.Artifacts{
 							{
-								Name: "my-artifact",
+								Name: "my-s3-artifact",
 								ArtifactLocation: wfv1.ArtifactLocation{
-									Raw: &wfv1.RawArtifact{
-										Data: "my-data",
+									S3: &wfv1.S3Artifact{
+										Key: "my-wf/my-node/my-s3-artifact.tgz",
+									},
+								},
+							},
+							{
+								Name: "my-gcs-artifact",
+								ArtifactLocation: wfv1.ArtifactLocation{
+									GCS: &wfv1.GCSArtifact{
+										Key: "my-wf/my-node/my-gcs-artifact",
+									},
+								},
+							},
+							{
+								Name: "my-oss-artifact",
+								ArtifactLocation: wfv1.ArtifactLocation{
+									GCS: &wfv1.GCSArtifact{
+										Key: "my-wf/my-node/my-oss-artifact.zip",
 									},
 								},
 							},
@@ -62,18 +96,46 @@ func newServer() *ArtifactServer {
 	gatekeeper.On("Context", mock.Anything).Return(ctx, nil)
 	a := &mocks.WorkflowArchive{}
 	a.On("GetWorkflow", "my-uuid").Return(wf, nil)
-	return NewArtifactServer(gatekeeper, hydratorfake.Noop, a, instanceid.NewService(instanceId))
+
+	fakeArtifactDriverFactory := func(_ *wfv1.Artifact, _ resource.Interface) (artifact.ArtifactDriver, error) {
+		return &fakeArtifactDriver{data: []byte("my-data")}, nil
+	}
+
+	return newArtifactServer(gatekeeper, hydratorfake.Noop, a, instanceid.NewService(instanceId), fakeArtifactDriverFactory)
 }
 
 func TestArtifactServer_GetArtifact(t *testing.T) {
 	s := newServer()
-	r := &http.Request{}
-	r.URL = mustParse("/artifacts/my-ns/my-wf/my-node/my-artifact")
-	w := &testhttp.TestResponseWriter{}
-	s.GetArtifact(w, r)
-	assert.Equal(t, 200, w.StatusCode)
-	assert.Equal(t, "filename=\"my-artifact.tgz\"", w.Header().Get("Content-Disposition"))
-	assert.Equal(t, "my-data", w.Output)
+
+	tests := []struct {
+		fileName     string
+		artifactName string
+	}{
+		{
+			fileName:     "my-s3-artifact.tgz",
+			artifactName: "my-s3-artifact",
+		},
+		{
+			fileName:     "my-gcs-artifact",
+			artifactName: "my-gcs-artifact",
+		},
+		{
+			fileName:     "my-oss-artifact.zip",
+			artifactName: "my-oss-artifact",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.artifactName, func(t *testing.T) {
+			r := &http.Request{}
+			r.URL = mustParse(fmt.Sprintf("/artifacts/my-ns/my-wf/my-node/%s", tt.artifactName))
+			w := &testhttp.TestResponseWriter{}
+			s.GetArtifact(w, r)
+			assert.Equal(t, 200, w.StatusCode)
+			assert.Equal(t, fmt.Sprintf(`filename="%s"`, tt.fileName), w.Header().Get("Content-Disposition"))
+			assert.Equal(t, "my-data", w.Output)
+		})
+	}
 }
 
 func TestArtifactServer_GetArtifactWithoutInstanceID(t *testing.T) {

workflow/artifacts/artifacts.go

@@ -27,6 +27,8 @@ type ArtifactDriver interface {
 
 var ErrUnsupportedDriver = fmt.Errorf("unsupported artifact driver")
 
+type NewDriverFunc func(art *wfv1.Artifact, ri resource.Interface) (ArtifactDriver, error)
+
 // NewDriver initializes an instance of an artifact driver
 func NewDriver(art *wfv1.Artifact, ri resource.Interface) (ArtifactDriver, error) {
 	if art.S3 != nil {