Merge pull request hashicorp#310 from ScottSuarez/support-generation

Add support for gcs object generations
julian7 · Mar 1, 2021 · 8724413 · 8724413
2 parents bcfc71c + 5bb6635
commit 8724413
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 9 deletions.
diff --git a/get_gcs.go b/get_gcs.go
@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"cloud.google.com/go/storage"
@@ -22,7 +23,7 @@ func (g *GCSGetter) ClientMode(u *url.URL) (ClientMode, error) {
 	ctx := g.Context()
 
 	// Parse URL
-	bucket, object, err := g.parseURL(u)
+	bucket, object, _, err := g.parseURL(u)
 	if err != nil {
 		return 0, err
 	}
@@ -59,7 +60,7 @@ func (g *GCSGetter) Get(dst string, u *url.URL) error {
 	ctx := g.Context()
 
 	// Parse URL
-	bucket, object, err := g.parseURL(u)
+	bucket, object, _, err := g.parseURL(u)
 	if err != nil {
 		return err
 	}
@@ -105,7 +106,7 @@ func (g *GCSGetter) Get(dst string, u *url.URL) error {
 			}
 			objDst = filepath.Join(dst, objDst)
 			// Download the matching object.
-			err = g.getObject(ctx, client, objDst, bucket, obj.Name)
+			err = g.getObject(ctx, client, objDst, bucket, obj.Name, "")
 			if err != nil {
 				return err
 			}
@@ -118,7 +119,7 @@ func (g *GCSGetter) GetFile(dst string, u *url.URL) error {
 	ctx := g.Context()
 
 	// Parse URL
-	bucket, object, err := g.parseURL(u)
+	bucket, object, fragment, err := g.parseURL(u)
 	if err != nil {
 		return err
 	}
@@ -127,11 +128,21 @@ func (g *GCSGetter) GetFile(dst string, u *url.URL) error {
 	if err != nil {
 		return err
 	}
-	return g.getObject(ctx, client, dst, bucket, object)
+	return g.getObject(ctx, client, dst, bucket, object, fragment)
 }
 
-func (g *GCSGetter) getObject(ctx context.Context, client *storage.Client, dst, bucket, object string) error {
-	rc, err := client.Bucket(bucket).Object(object).NewReader(ctx)
+func (g *GCSGetter) getObject(ctx context.Context, client *storage.Client, dst, bucket, object, fragment string) error {
+	var rc *storage.Reader
+	var err error
+	if fragment != "" {
+		generation, err := strconv.ParseInt(fragment, 10, 64)
+		if err != nil {
+			return err
+		}
+		rc, err = client.Bucket(bucket).Object(object).Generation(generation).NewReader(ctx)
+	} else {
+		rc, err = client.Bucket(bucket).Object(object).NewReader(ctx)
+	}
 	if err != nil {
 		return err
 	}
@@ -145,7 +156,7 @@ func (g *GCSGetter) getObject(ctx context.Context, client *storage.Client, dst,
 	return copyReader(dst, rc, 0666, g.client.umask())
 }
 
-func (g *GCSGetter) parseURL(u *url.URL) (bucket, path string, err error) {
+func (g *GCSGetter) parseURL(u *url.URL) (bucket, path, fragment string, err error) {
 	if strings.Contains(u.Host, "googleapis.com") {
 		hostParts := strings.Split(u.Host, ".")
 		if len(hostParts) != 3 {
@@ -160,6 +171,7 @@ func (g *GCSGetter) parseURL(u *url.URL) (bucket, path string, err error) {
 		}
 		bucket = pathParts[3]
 		path = pathParts[4]
+		fragment = u.Fragment
 	}
 	return
 }
diff --git a/get_gcs_test.go b/get_gcs_test.go
@@ -1,6 +1,7 @@
 package getter
 
 import (
+	"io/ioutil"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -88,6 +89,47 @@ func TestGCSGetter_GetFile(t *testing.T) {
 	assertContents(t, dst, "# Main\n")
 }
 
+func TestGCSGetter_GetGenerationFile(t *testing.T) {
+	defer initGCPCredentials(t)()
+
+	g := new(GCSGetter)
+	dst := tempTestFile(t)
+	defer os.RemoveAll(filepath.Dir(dst))
+
+	// Download
+	err := g.GetFile(
+		dst, testURL("https://www.googleapis.com/storage/v1/go-getter-testcase-data/DO_NOT_DELETE/generation_test.txt#1614317688843055"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Verify contents are valid for this generation
+	content, err := ioutil.ReadFile(dst)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if string(content) != "a" {
+		t.Fatalf("expected file contents from generation to be `a` but got `%s`", content)
+	}
+
+	// Download
+	err = g.GetFile(
+		dst, testURL("https://www.googleapis.com/storage/v1/go-getter-testcase-data/DO_NOT_DELETE/generation_test.txt#1614317705239142"))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Verify contents are valid for this generation
+	content, err = ioutil.ReadFile(dst)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if string(content) != "b" {
+		t.Fatalf("expected file contents from generation to be `b` but got `%s`", content)
+	}
+
+}
+
 func TestGCSGetter_GetFile_notfound(t *testing.T) {
 	g := new(GCSGetter)
 	dst := tempTestFile(t)
@@ -180,7 +222,7 @@ func TestGCSGetter_Url(t *testing.T) {
 				t.Fatalf("expected forced protocol to be gcs")
 			}
 
-			bucket, path, err := g.parseURL(u)
+			bucket, path, _, err := g.parseURL(u)
 
 			if err != nil {
 				t.Fatalf("err: %s", err)