@@ -17,12 +17,15 @@ limitations under the License.
1717package gce
1818
1919import (
20+ "context"
2021 "crypto/rand"
2122 "encoding/hex"
23+ "fmt"
2224 "os"
2325 "strings"
24- "time"
2526
27+ "cloud.google.com/go/iam"
28+ "cloud.google.com/go/storage"
2629 "k8s.io/klog/v2"
2730 "sigs.k8s.io/kubetest2/pkg/exec"
2831)
@@ -41,57 +44,58 @@ func GCSBucketName(projectID, prefix string) string {
4144}
4245
4346func EnsureGCSBucket (bucketPath , projectID string , public bool ) error {
44- lsArgs := []string {
45- "gsutil" , "ls" , "-b" ,
46- }
47- if projectID != "" {
48- lsArgs = append (lsArgs , "-p" , projectID )
47+ // TODO: Detect the GCP region used
48+ return EnsureGCSBucketWithRegion (bucketPath , projectID , "us-central1" , public )
49+ }
50+
51+ func EnsureGCSBucketWithRegion (bucketPath , projectID string , region string , public bool ) error {
52+ ctx := context .Background ()
53+ client , err := storage .NewClient (ctx )
54+ if err != nil {
55+ return fmt .Errorf ("failed to create storage client: %w" , err )
4956 }
50- lsArgs = append ( lsArgs , bucketPath )
57+ defer client . Close ( )
5158
52- klog .Info (strings .Join (lsArgs , " " ))
53- cmd := exec .Command (lsArgs [0 ], lsArgs [1 :]... )
59+ // Extract bucket name from gs:// path if the bucket's URI is provided
60+ bucketName := strings .TrimPrefix (bucketPath , "gs://" )
61+ bucketName = strings .TrimSuffix (bucketName , "/" )
5462
55- output , err := exec .CombinedOutputLines (cmd )
63+ bucket := client .Bucket (bucketName )
64+
65+ // Check if bucket exists
66+ klog .Infof ("Checking if bucket %s exists" , bucketName )
67+ _ , err = bucket .Attrs (ctx )
5668 if err == nil {
69+ klog .Infof ("Bucket %s already exists" , bucketName )
5770 return nil
58- } else if len (output ) != 1 || ! strings .Contains (output [0 ], "BucketNotFound" ) {
59- klog .Info (output )
60- return err
6171 }
6272
63- mbArgs := []string {
64- "gsutil" , "mb" ,
73+ // If error is not "bucket doesn't exist", return error
74+ if err != storage .ErrBucketNotExist {
75+ return fmt .Errorf ("error checking bucket: %w" , err )
6576 }
66- if projectID != "" {
67- mbArgs = append (mbArgs , "-p" , projectID )
68- }
69- mbArgs = append (mbArgs , bucketPath )
70-
71- klog .Info (strings .Join (mbArgs , " " ))
72- cmd = exec .Command (mbArgs [0 ], mbArgs [1 :]... )
7377
74- exec . InheritOutput ( cmd )
75- err = cmd . Run ( )
76- if err != nil {
77- return err
78+ // Create the bucket
79+ klog . Infof ( "Creating bucket %s in project %s" , bucketName , projectID )
80+ if err := bucket . Create ( ctx , projectID , nil ); err != nil {
81+ return fmt . Errorf ( "failed to create bucket: %w" , err )
7882 }
7983
8084 if public {
81- iamArgs := []string {
82- "gsutil" , "iam" , "ch" , "allUsers:objectViewer" ,
83- }
84- iamArgs = append (iamArgs , bucketPath )
85- klog .Info (strings .Join (iamArgs , " " ))
86- // GCS APIs are strongly consistent but this should help with flakes
87- time .Sleep (10 * time .Second )
88- cmd = exec .Command (iamArgs [0 ], iamArgs [1 :]... )
89- exec .InheritOutput (cmd )
90- err = cmd .Run ()
85+ klog .Infof ("Making bucket %s public" , bucketName )
86+ policy , err := bucket .IAM ().Policy (ctx )
9187 if err != nil {
92- return err
88+ return fmt .Errorf ("failed to get bucket IAM policy: %w" , err )
89+ }
90+
91+ // Add allUsers as objectViewer
92+ policy .Add (iam .AllUsers , "roles/storage.objectViewer" )
93+
94+ if err := bucket .IAM ().SetPolicy (ctx , policy ); err != nil {
95+ return fmt .Errorf ("failed to set bucket IAM policy: %w" , err )
9396 }
9497 }
98+
9599 return nil
96100}
97101
0 commit comments