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