wait.go

Documentation: github.com/GoogleCloudPlatform/k8s-config-connector/pkg/gcp

     1  // Copyright 2022 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package gcp
    16  
    17  import (
    18  	"fmt"
    19  	"time"
    20  
    21  	"google.golang.org/api/bigtableadmin/v2"
    22  	"google.golang.org/api/cloudasset/v1"
    23  	resourcemanager "google.golang.org/api/cloudresourcemanager/v1"
    24  	"google.golang.org/api/compute/v1"
    25  	container "google.golang.org/api/container/v1beta1"
    26  	"google.golang.org/api/redis/v1"
    27  	"google.golang.org/api/servicenetworking/v1"
    28  	"google.golang.org/api/spanner/v1"
    29  	sqladmin "google.golang.org/api/sqladmin/v1beta4"
    30  	"k8s.io/apimachinery/pkg/util/wait"
    31  )
    32  
    33  type AssetInventoryWaitCallback func(operation *cloudasset.Operation) error
    34  type BigtableWaitCallback func(operation *bigtableadmin.Operation) error
    35  type RedisWaitCallback func(operation *redis.Operation) error
    36  type SpannerWaitCallback func(operation *spanner.Operation) error
    37  type SqlWaitCallback func(operation *sqladmin.Operation) error
    38  type ResourceManagerCallback func(operation *resourcemanager.Operation) error
    39  type ComputeWaitCallback func(operation *compute.Operation) error
    40  type ContainerWaitCallback func(operation *container.Operation) error
    41  type ServiceNetworkingWaitCallback func(operation *servicenetworking.Operation) error
    42  
    43  func WaitForAssetInventoryOperationDefaultTimeout(assetClient *cloudasset.Service,
    44  	operation *cloudasset.Operation, projectNum string, callback AssetInventoryWaitCallback) (*cloudasset.Operation, error) {
    45  	return WaitForAssetInventoryOperation(assetClient, operation, projectNum, 10*time.Second, 30*time.Minute, callback)
    46  }
    47  
    48  // The project number should be the project number for the credentials which are making this request
    49  func WaitForAssetInventoryOperation(assetClient *cloudasset.Service, operation *cloudasset.Operation, projectNum string,
    50  	interval, timeout time.Duration, callback AssetInventoryWaitCallback) (*cloudasset.Operation, error) {
    51  	err := wait.PollImmediate(interval, timeout, func() (done bool, err error) {
    52  		request := assetClient.Operations.Get(operation.Name)
    53  		// CAI needs this value to use end user credentials instead of a service account
    54  		//   https://cloud.google.com/asset-inventory/docs/faq
    55  		// when the export is run on a project, the project number is in the operation name, however, when the operation
    56  		// is run on an organization or folder it is not in the name.
    57  		request.Header().Add("X-Goog-User-Project", projectNum)
    58  		newOp, err := request.Do()
    59  		if err != nil {
    60  			return false, fmt.Errorf("error getting operation %v: %v", operation.Name,
    61  				err)
    62  		}
    63  		operation = newOp
    64  		if callback != nil {
    65  			if err = callback(operation); err != nil {
    66  				return false, fmt.Errorf("error returned by wait callback: %v", err)
    67  			}
    68  		}
    69  		return operation.Done, nil
    70  	})
    71  	return operation, err
    72  }
    73  
    74  func WaitForResourceManagerOperationDefaultTimeout(rmClient *resourcemanager.Service, operation *resourcemanager.Operation,
    75  	callback ResourceManagerCallback) (*resourcemanager.Operation, error) {
    76  	return WaitForResourceManagerOperation(rmClient, operation, 10*time.Second, 10*time.Minute, callback)
    77  }
    78  
    79  func WaitForResourceManagerOperation(rmClient *resourcemanager.Service, operation *resourcemanager.Operation,
    80  	interval, timeout time.Duration, callback ResourceManagerCallback) (*resourcemanager.Operation, error) {
    81  	err := wait.PollImmediate(interval, timeout, func() (done bool, err error) {
    82  		newOp, err := rmClient.Operations.Get(operation.Name).Do()
    83  		if err != nil {
    84  			return false, fmt.Errorf("error getting operation %v: %v", operation.Name,
    85  				err)
    86  		}
    87  		operation = newOp
    88  		if callback != nil {
    89  			if err = callback(operation); err != nil {
    90  				return false, fmt.Errorf("error returned by wait callback: %v", err)
    91  			}
    92  		}
    93  		return operation.Done, nil
    94  	})
    95  	return operation, err
    96  }
    97  
    98  func WaitForComputeOperationDefaultTimeout(computeClient *compute.Service, operation *compute.Operation,
    99  	projectID string, callback ComputeWaitCallback) (*compute.Operation, error) {
   100  	return WaitForComputeOperation(computeClient, operation, projectID, 30*time.Second, 30*time.Minute, callback)
   101  }
   102  
   103  func WaitForComputeOperation(computeClient *compute.Service, operation *compute.Operation,
   104  	projectID string, interval, timeout time.Duration, callback ComputeWaitCallback) (*compute.Operation, error) {
   105  	err := wait.PollImmediate(interval, timeout, func() (done bool, err error) {
   106  		var newOp *compute.Operation
   107  		if operation.Zone != "" {
   108  			zone := FullResourceNameToShortName(operation.Zone)
   109  			newOp, err = computeClient.ZoneOperations.Get(projectID, zone, operation.Name).Do()
   110  		} else if operation.Region != "" {
   111  			region := FullResourceNameToShortName(operation.Region)
   112  			newOp, err = computeClient.RegionOperations.Get(projectID, region, operation.Name).Do()
   113  		} else {
   114  			newOp, err = computeClient.GlobalOperations.Get(projectID, operation.Name).Do()
   115  		}
   116  		if err != nil {
   117  			return false, fmt.Errorf("error getting operation %v/%v: %v", projectID, operation.Name,
   118  				err)
   119  		}
   120  		operation = newOp
   121  		if callback != nil {
   122  			if err = callback(operation); err != nil {
   123  				return false, fmt.Errorf("error returned by wait callback: %v", err)
   124  			}
   125  		}
   126  		return operation.Status == "DONE", nil
   127  	})
   128  	return operation, err
   129  }
   130
View as plain text