update driver to support staging compute

Author: Sneha-at

cmd/gce-pd-csi-driver/main.go

@@ -66,6 +66,7 @@ var (
 
 	maxConcurrentFormatAndMount = flag.Int("max-concurrent-format-and-mount", 1, "If set then format and mount operations are serialized on each node. This is stronger than max-concurrent-format as it includes fsck and other mount operations")
 	formatAndMountTimeout       = flag.Duration("format-and-mount-timeout", 1*time.Minute, "The maximum duration of a format and mount operation before another such operation will be started. Used only if --serialize-format-and-mount")
+	computeEnvironment          = flag.String("compute-environment", "prod", "Sets the compute environment")
 
 	fallbackRequisiteZonesFlag = flag.String("fallback-requisite-zones", "", "Comma separated list of requisite zones that will be used if there are not sufficient zones present in requisite topologies when provisioning a disk")
 
@@ -137,7 +138,7 @@ func handle() {
 	// Initialize requirements for the controller service
 	var controllerServer *driver.GCEControllerServer
 	if *runControllerService {
-		cloudProvider, err := gce.CreateCloudProvider(ctx, version, *cloudConfigFilePath, *computeEndpoint)
+		cloudProvider, err := gce.CreateCloudProvider(ctx, version, *cloudConfigFilePath, *computeEndpoint, *computeEnvironment)
 		if err != nil {
 			klog.Fatalf("Failed to get cloud provider: %v", err.Error())
 		}

pkg/gce-cloud-provider/compute/gce.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/url"
 	"os"
 	"runtime"
 	"time"
@@ -29,6 +30,7 @@ import (
 
 	"cloud.google.com/go/compute/metadata"
 	"golang.org/x/oauth2"
+	computealpha "google.golang.org/api/compute/v0.alpha"
 	computebeta "google.golang.org/api/compute/v0.beta"
 	"google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
@@ -46,8 +48,27 @@ const (
 	regionURITemplate = "projects/%s/regions/%s"
 
 	replicaZoneURITemplateSingleZone = "projects/%s/zones/%s" // {gce.projectID}/zones/{disk.Zone}
+	versionV1                        = "v1"
+	versionBeta                      = "beta"
+	versionAlpha                     = "alpha"
+	googleEnv                        = "googleapis"
 )
 
+var computeVersionMap = map[string]map[string]map[string]string{
+	googleEnv: {
+		"prod": {
+			versionV1:    "compute/v1/",
+			versionBeta:  "compute/beta/",
+			versionAlpha: "compute/alpha/",
+		},
+		"staging": {
+			versionV1:    "compute/staging_v1/",
+			versionBeta:  "compute/staging_beta/",
+			versionAlpha: "compute/staging_alpha/",
+		},
+	},
+}
+
 type CloudProvider struct {
 	service     *compute.Service
 	betaService *computebeta.Service
@@ -70,7 +91,7 @@ type ConfigGlobal struct {
 	Zone      string `gcfg:"zone"`
 }
 
-func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint string) (*CloudProvider, error) {
+func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint string, computeEnvironment string) (*CloudProvider, error) {
 	configFile, err := readConfig(configPath)
 	if err != nil {
 		return nil, err
@@ -85,15 +106,23 @@ func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath s
 		return nil, err
 	}
 
-	svc, err := createCloudService(ctx, vendorVersion, tokenSource, computeEndpoint)
+	svc, err := createCloudService(ctx, vendorVersion, tokenSource, computeEndpoint, computeEnvironment)
 	if err != nil {
 		return nil, err
 	}
+	klog.Infof("Compute endpoint for V1 version: %s", svc.BasePath)
 
-	betasvc, err := createBetaCloudService(ctx, vendorVersion, tokenSource, computeEndpoint)
+	betasvc, err := createBetaCloudService(ctx, vendorVersion, tokenSource, computeEndpoint, computeEnvironment)
 	if err != nil {
 		return nil, err
 	}
+	klog.Infof("Compute endpoint for Beta version: %s", betasvc.BasePath)
+
+	alphasvc, err := createAlphaCloudService(ctx, vendorVersion, tokenSource, computeEndpoint, computeEnvironment)
+	if err != nil {
+		return nil, err
+	}
+	klog.Infof("Compute endpoint for Alpha version: %s", alphasvc.BasePath)
 
 	project, zone, err := getProjectAndZone(configFile)
 	if err != nil {
@@ -156,16 +185,23 @@ func readConfig(configPath string) (*ConfigFile, error) {
 	return cfg, nil
 }
 
-func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string) (*computebeta.Service, error) {
-	client, err := newOauthClient(ctx, tokenSource)
+func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*computealpha.Service, error) {
+	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionAlpha)
+	if err != nil {
+		klog.Errorf("Failed to get compute endpoint: %s", err)
+	}
+	service, err := computealpha.NewService(ctx, computeOpts...)
 	if err != nil {
 		return nil, err
 	}
+	service.UserAgent = fmt.Sprintf("GCE CSI Driver/%s (%s %s)", vendorVersion, runtime.GOOS, runtime.GOARCH)
+	return service, nil
+}
 
-	computeOpts := []option.ClientOption{option.WithHTTPClient(client)}
-	if computeEndpoint != "" {
-		betaEndpoint := fmt.Sprintf("%s/compute/beta/", computeEndpoint)
-		computeOpts = append(computeOpts, option.WithEndpoint(betaEndpoint))
+func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*computebeta.Service, error) {
+	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionBeta)
+	if err != nil {
+		klog.Errorf("Failed to get compute endpoint: %s", err)
 	}
 	service, err := computebeta.NewService(ctx, computeOpts...)
 	if err != nil {
@@ -175,28 +211,39 @@ func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSour
 	return service, nil
 }
 
-func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string) (*compute.Service, error) {
-	svc, err := createCloudServiceWithDefaultServiceAccount(ctx, vendorVersion, tokenSource, computeEndpoint)
-	return svc, err
+func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*compute.Service, error) {
+	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionV1)
+	if err != nil {
+		klog.Errorf("Failed to get compute endpoint: %s", err)
+	}
+	service, err := compute.NewService(ctx, computeOpts...)
+	if err != nil {
+		return nil, err
+	}
+	service.UserAgent = fmt.Sprintf("GCE CSI Driver/%s (%s %s)", vendorVersion, runtime.GOOS, runtime.GOARCH)
+	return service, nil
 }
 
-func createCloudServiceWithDefaultServiceAccount(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string) (*compute.Service, error) {
+func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string, computeVersion string) ([]option.ClientOption, error) {
 	client, err := newOauthClient(ctx, tokenSource)
 	if err != nil {
 		return nil, err
 	}
-
+	computeEnvironmentSuffix, ok := computeVersionMap[googleEnv][computeEnvironment][computeVersion]
+	if !ok {
+		return nil, errors.New("Unable to fetch compute endpoint")
+	}
 	computeOpts := []option.ClientOption{option.WithHTTPClient(client)}
 	if computeEndpoint != "" {
-		v1Endpoint := fmt.Sprintf("%s/compute/v1/", computeEndpoint)
-		computeOpts = append(computeOpts, option.WithEndpoint(v1Endpoint))
-	}
-	service, err := compute.NewService(ctx, computeOpts...)
-	if err != nil {
-		return nil, err
+		endpoint := fmt.Sprintf("%s%s", computeEndpoint, computeEnvironmentSuffix)
+		klog.Infof("Got compute endpoint %s", endpoint)
+		_, err := url.ParseRequestURI(endpoint)
+		if err != nil {
+			klog.Fatalf("Error parsing compute endpoint %s", endpoint)
+		}
+		computeOpts = append(computeOpts, option.WithEndpoint(endpoint))
 	}
-	service.UserAgent = fmt.Sprintf("GCE CSI Driver/%s (%s %s)", vendorVersion, runtime.GOOS, runtime.GOARCH)
-	return service, nil
+	return computeOpts, nil
 }
 
 func newOauthClient(ctx context.Context, tokenSource oauth2.TokenSource) (*http.Client, error) {

pkg/gce-pd-csi-driver/controller.go

@@ -156,7 +156,7 @@ const (
 )
 
 var (
-	validResourceApiVersions = map[string]bool{"v1": true, "alpha": true, "beta": true}
+	validResourceApiVersions = map[string]bool{"v1": true, "alpha": true, "beta": true, "staging_v1": true, "staging_beta": true, "staging_alpha": true}
 )
 
 func isDiskReady(disk *gce.CloudDisk) (bool, error) {