[artifacts] Support fetching directories from artifacts_gcs_bucket.

Bug: 45302
Change-Id: I062061851146e8d9edeeb5c80a58f41ab4262a84
Reviewed-on: https://fuchsia-review.googlesource.com/c/infra/infra/+/372499
Commit-Queue: Ina Huh <ihuh@google.com>
Reviewed-by: Gary Miguel <garymm@google.com>
diff --git a/artifacts/artifacts.go b/artifacts/artifacts.go
index de3e464..b2dd072 100644
--- a/artifacts/artifacts.go
+++ b/artifacts/artifacts.go
@@ -7,7 +7,7 @@
 import (
 	"context"
 	"fmt"
-	"strings"
+	"path"
 
 	"cloud.google.com/go/storage"
 
@@ -38,10 +38,14 @@
 
 // GetBuildDir returns the BuildDirectory for the given build. bucket is the GCS bucket.
 // build is the BuildBucket build ID.
-func (c *Client) GetBuildDir(bucket, build string) *BuildDirectory {
+func (c *Client) GetBuildDir(bucket, build string) Directory {
 	bkt := c.client.Bucket(bucket)
+	root := ""
+	if build != "" {
+		root = path.Join("builds", build)
+	}
 	return &BuildDirectory{&directory{
 		bucket: bkt,
-		root:   strings.Join([]string{"builds", build}, "/"),
+		root:   root,
 	}}
 }
diff --git a/artifacts/directory.go b/artifacts/directory.go
index c1c8998..bcbf95a 100644
--- a/artifacts/directory.go
+++ b/artifacts/directory.go
@@ -7,9 +7,12 @@
 import (
 	"context"
 	"fmt"
+	"io"
+	"path"
 	"strings"
 
 	"cloud.google.com/go/storage"
+	"go.fuchsia.dev/fuchsia/tools/lib/osmisc"
 	"google.golang.org/api/iterator"
 )
 
@@ -17,6 +20,13 @@
 // test's combined stdout and stderr streams today.  This will change in the future.
 const DefaultTestOutputName = "stdout-stderr.txt"
 
+// Directory is an interface for interacting with a Cloud Storage "directory".
+type Directory interface {
+	Object(string) *storage.ObjectHandle
+	CopyFile(ctx context.Context, src, dest string) error
+	List(context.Context, string) ([]string, error)
+}
+
 // BuildDirectory represents a Fuchsia CI build's artifact directory. Refer to the
 // layout in doc.go for the layout of this directory. When amending the layout, prefer
 // adding convenience methods on this type to encourages all clients to create objects
@@ -41,16 +51,35 @@
 	root   string
 }
 
-// Object returns a handle to the given object within this directory. path is the path to
+// Object returns a handle to the given object within this directory. objPath is the path to
 // the object relative to this directory.
-func (d *directory) Object(path string) *storage.ObjectHandle {
-	path = fmt.Sprintf("%s/%s", d.root, path)
-	return d.bucket.Object(path)
+func (d *directory) Object(objPath string) *storage.ObjectHandle {
+	objPath = path.Join(d.root, objPath)
+	return d.bucket.Object(objPath)
 }
 
-// List lists all of the objects in this directory.
-func (d *directory) List(ctx context.Context) ([]string, error) {
-	prefix := strings.Join([]string{d.root}, "/")
+// CopyFile copies the object from src to dst, creating all the parent directories
+// of dest if they don't exist.
+func (d *directory) CopyFile(ctx context.Context, src, dest string) error {
+	obj := d.Object(src)
+	input, err := obj.NewReader(ctx)
+	if err != nil {
+		return err
+	}
+
+	output, err := osmisc.CreateFile(dest)
+	if err != nil {
+		return err
+	}
+
+	_, err = io.Copy(output, input)
+	return err
+}
+
+// List lists all of the objects in this directory with the given prefix. If prefix == "",
+// it will list everything in this directory.
+func (d *directory) List(ctx context.Context, prefix string) ([]string, error) {
+	prefix = path.Join(d.root, prefix)
 	iter := d.bucket.Objects(ctx, &storage.Query{
 		Prefix: prefix,
 	})
@@ -64,7 +93,7 @@
 		if err != nil {
 			return nil, err
 		}
-		items = append(items, strings.TrimPrefix(attrs.Name, prefix+"/"))
+		items = append(items, strings.TrimPrefix(attrs.Name, d.root+"/"))
 	}
 
 	return items, nil
diff --git a/cmd/artifacts/common.go b/cmd/artifacts/common.go
index 00c9aba..f4fa44a 100644
--- a/cmd/artifacts/common.go
+++ b/cmd/artifacts/common.go
@@ -18,14 +18,17 @@
 
 // The name of the recipe property passed to Fuchsia builders, which communicates which
 // Cloud storage bucket to upload artifacts to.
-const bucketPropertyName = "gcs_bucket"
+const bucketPropertyName = "artifact_gcs_bucket"
+
+// TODO(fxb/45302): Remove once no longer in use.
+const archiveBucketPropertyName = "gcs_bucket"
 
 // buildsClient sends RPCs to a BuildBucket server.
 type buildsClient interface {
 	GetBuild(context.Context, *buildbucketpb.GetBuildRequest, ...grpc.CallOption) (*buildbucketpb.Build, error)
 }
 
-func getStorageBucket(ctx context.Context, client buildsClient, build string) (string, error) {
+func getStorageBucket(ctx context.Context, client buildsClient, build string, withArchives bool) (string, error) {
 	buildID, err := strconv.ParseInt(build, 10, 64)
 	if err != nil {
 		return "", err
@@ -46,10 +49,14 @@
 		return "", fmt.Errorf("build %s not found", build)
 	}
 
+	propertyName := bucketPropertyName
+	if withArchives {
+		propertyName = archiveBucketPropertyName
+	}
 	wrapper := buildbucket.Build(*response)
-	property, ok := wrapper.OutputProperty(bucketPropertyName)
+	property, ok := wrapper.OutputProperty(propertyName)
 	if !ok {
-		return "", fmt.Errorf("no output property %q found for build %q", bucketPropertyName, build)
+		return "", fmt.Errorf("no output property %q found for build %q", propertyName, build)
 	}
 
 	return property.StringValue(), nil
diff --git a/cmd/artifacts/common_test.go b/cmd/artifacts/common_test.go
index 30bb362..e63e675 100644
--- a/cmd/artifacts/common_test.go
+++ b/cmd/artifacts/common_test.go
@@ -53,7 +53,7 @@
 					Output: &buildbucketpb.Build_Output{
 						Properties: &_struct.Struct{
 							Fields: map[string]*_struct.Value{
-								"gcs_bucket": &_struct.Value{
+								"artifact_gcs_bucket": {
 									Kind: &_struct.Value_StringValue{
 										StringValue: "the_bucket",
 									},
@@ -122,7 +122,7 @@
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			bucket, err := getStorageBucket(context.Background(), &tt.mock, tt.input)
+			bucket, err := getStorageBucket(context.Background(), &tt.mock, tt.input, false)
 			if err != nil != tt.expectErr {
 				if err == nil {
 					t.Error("wanted an error but got nil")
diff --git a/cmd/artifacts/copy.go b/cmd/artifacts/copy.go
index d4c6eaa..72e91f8 100644
--- a/cmd/artifacts/copy.go
+++ b/cmd/artifacts/copy.go
@@ -8,17 +8,22 @@
 	"context"
 	"flag"
 	"fmt"
-	"io"
 	"log"
-	"os"
+	"path/filepath"
+	"strings"
 
 	"github.com/google/subcommands"
 	"go.chromium.org/luci/auth/client/authcli"
 	"go.chromium.org/luci/hardcoded/chromeinfra"
 	"go.fuchsia.dev/infra/artifacts"
 	"go.fuchsia.dev/infra/buildbucket"
+	"golang.org/x/sync/errgroup"
 )
 
+type artifactsClient interface {
+	GetBuildDir(bucket, build string) artifacts.Directory
+}
+
 type CopyCommand struct {
 	authFlags authcli.Flags
 
@@ -28,6 +33,10 @@
 
 	// The local path to write the artifact to.
 	dest string
+
+	// Whether source should be relative to the Cloud Storage bucket's root directory.
+	// If false, source will be relative to the target build's Cloud Storage directory.
+	fromRoot bool
 }
 
 func (*CopyCommand) Name() string {
@@ -45,8 +54,10 @@
 func (cmd *CopyCommand) SetFlags(f *flag.FlagSet) {
 	cmd.authFlags.Register(flag.CommandLine, chromeinfra.DefaultAuthOptions())
 	f.StringVar(&cmd.build, "build", "", "the ID of the build that produced the artifacts")
-	f.StringVar(&cmd.source, "src", "", "The artifact to download from the build's Cloud Storage directory")
-	f.StringVar(&cmd.dest, "dst", "", "The local path to write the artifact to")
+	f.StringVar(&cmd.source, "src", "", "The artifact file or directory to download from the build's Cloud Storage directory")
+	f.StringVar(&cmd.dest, "dst", "", "The local path to write the artifact(s) to")
+	f.BoolVar(&cmd.fromRoot, "root", false, "Whether src is relative to the root directory of the Cloud Storage bucket."+
+		"If false, src will be taken as a relative path to the build-specific directory under the Cloud Storage bucket.")
 }
 
 func (cmd *CopyCommand) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
@@ -84,24 +95,46 @@
 	return cmd.execute(ctx, buildsCli, artifactsCli)
 }
 
-func (cmd *CopyCommand) execute(ctx context.Context, buildsCli buildsClient, artifactsCli *artifacts.Client) error {
-	bucket, err := getStorageBucket(ctx, buildsCli, cmd.build)
+func (cmd *CopyCommand) execute(ctx context.Context, buildsCli buildsClient, artifactsCli artifactsClient) error {
+	// If cmd.source is an archive, it will have an extension like .tgz or .tar.gz.
+	// Otherwise, we assume it to be a directory in the artifacts bucket.
+	// getStorageBucket will retrieve the appropriate bucket based on whether the source
+	// is an archive or not.
+	fetchArchive := filepath.Ext(cmd.source) != ""
+	bucket, err := getStorageBucket(ctx, buildsCli, cmd.build, fetchArchive)
 	if err != nil {
 		return err
 	}
 
-	dir := artifactsCli.GetBuildDir(bucket, cmd.build)
-	obj := dir.Object(cmd.source)
-	input, err := obj.NewReader(ctx)
+	build := cmd.build
+	if cmd.fromRoot {
+		build = ""
+	}
+	dir := artifactsCli.GetBuildDir(bucket, build)
+	objs, err := dir.List(ctx, cmd.source)
 	if err != nil {
 		return err
 	}
-
-	output, err := os.Create(cmd.dest)
-	if err != nil {
-		return err
+	var eg errgroup.Group
+	for _, obj := range objs {
+		obj := obj
+		eg.Go(func() error {
+			var dest string
+			if obj == cmd.source {
+				// cmd.source is a single artifact. Copy it directly to cmd.dest.
+				dest = cmd.dest
+			} else {
+				// obj is a path relative to the dir with cmd.source as its prefix.
+				// We want to get the relative path to cmd.source so that we can write
+				// it to the same relative path to cmd.dest.
+				relPath := strings.TrimPrefix(obj, cmd.source+"/")
+				dest = filepath.Join(cmd.dest, relPath)
+			}
+			if err := dir.CopyFile(ctx, obj, dest); err != nil {
+				return err
+			}
+			return nil
+		})
 	}
-
-	_, err = io.Copy(output, input)
-	return err
+	return eg.Wait()
 }
diff --git a/cmd/artifacts/copy_test.go b/cmd/artifacts/copy_test.go
new file mode 100644
index 0000000..15a06b1
--- /dev/null
+++ b/cmd/artifacts/copy_test.go
@@ -0,0 +1,153 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+package main
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+
+	"cloud.google.com/go/storage"
+
+	_struct "github.com/golang/protobuf/ptypes/struct"
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+	"go.fuchsia.dev/infra/artifacts"
+)
+
+type mockDirectory struct {
+	files       []string
+	copiedFiles []string
+	lock        sync.Mutex
+}
+
+func (d *mockDirectory) Object(path string) *storage.ObjectHandle {
+	return nil
+}
+
+func (d *mockDirectory) CopyFile(ctx context.Context, src, dest string) error {
+	for _, file := range d.files {
+		if file == src {
+			d.lock.Lock()
+			d.copiedFiles = append(d.copiedFiles, dest)
+			d.lock.Unlock()
+			return nil
+		}
+	}
+	return fmt.Errorf("source doesn't exist")
+}
+
+func (d *mockDirectory) List(ctx context.Context, prefix string) ([]string, error) {
+	var result []string
+	for _, file := range d.files {
+		if strings.HasPrefix(file, prefix) {
+			result = append(result, file)
+		}
+	}
+	return result, nil
+}
+
+type mockArtifactsClient struct {
+	files map[string][]string
+	dir   *mockDirectory
+}
+
+func (a *mockArtifactsClient) GetBuildDir(bucket, buildID string) artifacts.Directory {
+	a.dir = &mockDirectory{
+		files: a.files[bucket],
+	}
+	return a.dir
+}
+
+func TestExecute(t *testing.T) {
+	archiveBucket := []string{
+		"build-archive.tgz",
+		"packages.tar.gz",
+	}
+	artifactsBucket := []string{
+		"images/a",
+		"images/b/c",
+		"packages/a",
+		"packages/b/c",
+	}
+	storageContents := map[string][]string{
+		"gcs_bucket":       archiveBucket,
+		"artifacts_bucket": artifactsBucket,
+	}
+	buildID := "123"
+	buildsCli := &mockBuildsClient{
+		response: &buildbucketpb.Build{
+			Id: 123,
+			Output: &buildbucketpb.Build_Output{
+				Properties: &_struct.Struct{
+					Fields: map[string]*_struct.Value{
+						"gcs_bucket": {
+							Kind: &_struct.Value_StringValue{
+								StringValue: "gcs_bucket",
+							},
+						},
+						"artifact_gcs_bucket": {
+							Kind: &_struct.Value_StringValue{
+								StringValue: "artifacts_bucket",
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	tests := []struct {
+		name          string
+		source        string
+		dest          string
+		expectedFiles []string
+	}{
+		{
+			name:          "copy directory from the artifacts bucket",
+			source:        "packages",
+			dest:          "output/packages",
+			expectedFiles: []string{"output/packages/a", "output/packages/b/c"},
+		}, {
+			name:          "copy archive",
+			source:        "build-archive.tgz",
+			dest:          "output/build-archive.tgz",
+			expectedFiles: []string{"output/build-archive.tgz"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			artifactsCli := &mockArtifactsClient{
+				files: storageContents,
+			}
+			cmd := &CopyCommand{
+				build:  buildID,
+				source: tt.source,
+				dest:   tt.dest,
+			}
+			err := cmd.execute(context.Background(), buildsCli, artifactsCli)
+			if err != nil {
+				t.Errorf("unxpected error: %v", err)
+			}
+
+			compare := func(expected []string, actual []string) {
+				if len(expected) != len(actual) {
+					t.Errorf("expected:\n%+v\nbut got:\n%+v", expected, actual)
+				}
+				fileMap := map[string]bool{}
+				for _, file := range expected {
+					fileMap[file] = true
+				}
+				for _, file := range actual {
+					if _, ok := fileMap[file]; !ok {
+						t.Errorf("expected:\n%+v\nbut got:\n%+v", expected, actual)
+					}
+				}
+			}
+
+			compare(tt.expectedFiles, artifactsCli.dir.copiedFiles)
+		})
+	}
+}
diff --git a/cmd/artifacts/list.go b/cmd/artifacts/list.go
index 576bf48..8c39916 100644
--- a/cmd/artifacts/list.go
+++ b/cmd/artifacts/list.go
@@ -67,13 +67,18 @@
 }
 
 func (cmd *ListCommand) execute(ctx context.Context, buildsCli buildsClient, artifactsCli *artifacts.Client) error {
-	bucket, err := getStorageBucket(ctx, buildsCli, cmd.build)
+	// This is currently only called to check that the build's artifacts were uploaded to GCS.
+	// All builders that upload archives will also have uploaded artifacts to the artifacts bucket,
+	// so we can just list the contents of the archive bucket here, and it will mean that artifacts
+	// have been uploaded to both the archive and the artifacts buckets. Once support for fetching
+	// archives is removed, this will list the contents of the artifacts bucket.
+	bucket, err := getStorageBucket(ctx, buildsCli, cmd.build, true)
 	if err != nil {
 		return err
 	}
 
 	dir := artifactsCli.GetBuildDir(bucket, cmd.build)
-	items, err := dir.List(ctx)
+	items, err := dir.List(ctx, "")
 	if err != nil {
 		return err
 	}
diff --git a/cmd/artifacts/storetestoutputs.go b/cmd/artifacts/storetestoutputs.go
index 2390862..f014589 100644
--- a/cmd/artifacts/storetestoutputs.go
+++ b/cmd/artifacts/storetestoutputs.go
@@ -161,7 +161,7 @@
 		errs <- err
 		return
 	}
-	dir := cli.GetBuildDir(cmd.bucket, cmd.build)
+	dir := cli.GetBuildDir(cmd.bucket, cmd.build).(*artifacts.BuildDirectory)
 	for output := range outputs {
 		if err := cmd.upload(context.Background(), output, dir); err != nil {
 			errs <- err
diff --git a/go.sum b/go.sum
index 9c99053..3a32837 100644
--- a/go.sum
+++ b/go.sum
@@ -22,6 +22,7 @@
 cloud.google.com/go/storage v1.5.0 h1:RPUcBvDeYgQFMfQu1eBMq6piD1SXmLH+vK3qjewZPus=
 cloud.google.com/go/storage v1.5.0/go.mod h1:tpKbwo567HUNpVclU5sGELwQWBDZ8gh0ZeosJ0Rtdos=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
+github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/Microsoft/go-winio v0.4.14 h1:+hMXMk01us9KgxGb7ftKQt2Xpf5hH/yky+TDA+qxleU=