@@ -17,16 +17,24 @@ limitations under the License.
1717package gce
1818
1919import (
20+ "context"
2021 "crypto/rand"
2122 "encoding/hex"
23+ "errors"
24+ "fmt"
2225 "os"
2326 "strings"
24- "time"
2527
28+ "cloud.google.com/go/iam"
29+ "cloud.google.com/go/storage"
2630 "k8s.io/klog/v2"
2731 "sigs.k8s.io/kubetest2/pkg/exec"
2832)
2933
34+ const (
35+ defaultRegion = "us-central1"
36+ )
37+
3038func GCSBucketName (projectID , prefix string ) string {
3139 var s string
3240 if jobID := os .Getenv ("PROW_JOB_ID" ); len (jobID ) >= 4 {
@@ -41,57 +49,58 @@ func GCSBucketName(projectID, prefix string) string {
4149}
4250
4351func EnsureGCSBucket (bucketPath , projectID string , public bool ) error {
44- lsArgs := []string {
45- "gsutil" , "ls" , "-b" ,
46- }
47- if projectID != "" {
48- lsArgs = append (lsArgs , "-p" , projectID )
52+ // TODO: Detect the GCP region used
53+ return EnsureGCSBucketWithRegion (bucketPath , projectID , defaultRegion , public )
54+ }
55+
56+ func EnsureGCSBucketWithRegion (bucketPath , projectID string , region string , public bool ) error {
57+ ctx := context .Background ()
58+ client , err := storage .NewClient (ctx )
59+ if err != nil {
60+ return fmt .Errorf ("failed to create storage client: %w" , err )
4961 }
50- lsArgs = append (lsArgs , bucketPath )
62+ defer client .Close ()
63+
64+ // Extract bucket name from gs:// path if the bucket's URI is provided
65+ bucketName := strings .TrimPrefix (bucketPath , "gs://" )
66+ bucketName = strings .TrimSuffix (bucketName , "/" )
5167
52- klog .Info (strings .Join (lsArgs , " " ))
53- cmd := exec .Command (lsArgs [0 ], lsArgs [1 :]... )
68+ bucket := client .Bucket (bucketName )
5469
55- output , err := exec .CombinedOutputLines (cmd )
70+ // Check if bucket exists
71+ klog .Infof ("Checking if bucket %s exists" , bucketName )
72+ _ , err = bucket .Attrs (ctx )
5673 if err == nil {
74+ klog .Infof ("Bucket %s already exists" , bucketName )
5775 return nil
58- } else if len (output ) != 1 || ! strings .Contains (output [0 ], "BucketNotFound" ) {
59- klog .Info (output )
60- return err
6176 }
6277
63- mbArgs := []string {
64- "gsutil" , "mb" ,
65- }
66- if projectID != "" {
67- mbArgs = append (mbArgs , "-p" , projectID )
78+ // If error is not "bucket doesn't exist", return error
79+ if errors .Is (err , storage .ErrBucketNotExist ) {
80+ return fmt .Errorf ("error checking bucket: %w" , err )
6881 }
69- mbArgs = append (mbArgs , bucketPath )
7082
71- klog .Info (strings .Join (mbArgs , " " ))
72- cmd = exec .Command (mbArgs [0 ], mbArgs [1 :]... )
73-
74- exec .InheritOutput (cmd )
75- err = cmd .Run ()
76- if err != nil {
77- return err
83+ // Create the bucket
84+ klog .Infof ("Creating bucket %s in project %s" , bucketName , projectID )
85+ if err := bucket .Create (ctx , projectID , nil ); err != nil {
86+ return fmt .Errorf ("failed to create bucket: %w" , err )
7887 }
7988
8089 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 ()
90+ klog .Infof ("Making bucket %s public" , bucketName )
91+ policy , err := bucket .IAM ().Policy (ctx )
9192 if err != nil {
92- return err
93+ return fmt .Errorf ("failed to get bucket IAM policy: %w" , err )
94+ }
95+
96+ // Add allUsers as objectViewer
97+ policy .Add (iam .AllUsers , "roles/storage.objectViewer" )
98+
99+ if err := bucket .IAM ().SetPolicy (ctx , policy ); err != nil {
100+ return fmt .Errorf ("failed to set bucket IAM policy: %w" , err )
93101 }
94102 }
103+
95104 return nil
96105}
97106
0 commit comments