1# -*- coding: utf-8 -*- # 2# Copyright 2020 Google LLC. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15"""Utilities for AI Platform operations commands.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import unicode_literals 20 21from googlecloudsdk.core import log 22 23 24def WaitForOpMaybe(operations_client, op, op_ref, asynchronous=False, 25 log_method=None, message=None, kind=None): 26 """Waits for an operation if asynchronous flag is off. 27 28 Args: 29 operations_client: api_lib.ai.operations.OperationsClient, the client via 30 which to poll. 31 op: Cloud AI Platform operation, the operation to poll. 32 op_ref: The operation reference to the operation resource. It's the result 33 by calling resources.REGISTRY.Parse 34 asynchronous: bool, whether to wait for the operation or return immediately 35 log_method: Logging method used for synchronous operation. If None, no log 36 message: str, the message to display while waiting for the operation. 37 kind: str, the resource kind (instance, cluster, project, etc.), which will 38 be passed to logging function. 39 40 Returns: 41 The result of the operation if asynchronous is true, or the Operation 42 message otherwise 43 """ 44 logging_function = { 45 'create': log.CreatedResource, 46 'delete': log.DeletedResource, 47 'update': log.UpdatedResource, 48 } 49 if asynchronous: 50 if logging_function.get(log_method) is not None: 51 logging_function[log_method](op.name, kind=kind) 52 return op 53 return operations_client.WaitForOperation( 54 op, op_ref, message=message).response 55