206 lines
6.2 KiB
Python
206 lines
6.2 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2020 Google LLC. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Customized versions of runners in subprocess.
|
|
|
|
Some of this is just for python 2 support and can be simplified.
|
|
"""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import json
|
|
import os.path
|
|
import subprocess
|
|
import threading
|
|
from googlecloudsdk.api_lib.compute import utils
|
|
from googlecloudsdk.command_lib.code import json_stream
|
|
from googlecloudsdk.core import config
|
|
from googlecloudsdk.core.updater import update_manager
|
|
from googlecloudsdk.core.util import files as file_utils
|
|
import six
|
|
|
|
|
|
def _FindOrInstallComponent(component_name):
|
|
"""Finds the path to a component or install it.
|
|
|
|
Args:
|
|
component_name: Name of the component.
|
|
|
|
Returns:
|
|
Path to the component. Returns None if the component can't be found.
|
|
"""
|
|
if (config.Paths().sdk_root and
|
|
update_manager.UpdateManager.EnsureInstalledAndRestart([component_name])):
|
|
return os.path.join(config.Paths().sdk_root, 'bin', component_name)
|
|
|
|
return None
|
|
|
|
|
|
def GetGcloudPreferredExecutable(exe):
|
|
"""Finds the path to an executable, preferring the gcloud packaged version.
|
|
|
|
Args:
|
|
exe: Name of the executable.
|
|
|
|
Returns:
|
|
Path to the executable.
|
|
Raises:
|
|
EnvironmentError: The executable can't be found.
|
|
"""
|
|
path = _FindOrInstallComponent(exe) or file_utils.FindExecutableOnPath(exe)
|
|
if not path:
|
|
raise EnvironmentError('Unable to locate %s.' % exe)
|
|
return path
|
|
|
|
|
|
class _TimeoutThread(object):
|
|
"""A context manager based on threading.Timer.
|
|
|
|
Pass a function to call after the given time has passed. If you exit before
|
|
the timer fires, nothing happens. If you exit after we've had to call the
|
|
timer function, we raise TimeoutError at exit time.
|
|
"""
|
|
|
|
def __init__(self,
|
|
func,
|
|
timeout_sec,
|
|
error_format='Task ran for more than {timeout_sec} seconds'):
|
|
self.func = func
|
|
self.timeout_sec = timeout_sec
|
|
self.error_format = error_format
|
|
self.timer = None
|
|
|
|
def __enter__(self):
|
|
self.Reset()
|
|
return self
|
|
|
|
def Reset(self):
|
|
if self.timer is not None:
|
|
self.timer.cancel()
|
|
self.timer = threading.Timer(self.timeout_sec, self.func)
|
|
self.timer.start()
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
timed_out = self.timer.finished.is_set()
|
|
self.timer.cancel()
|
|
|
|
if timed_out:
|
|
raise utils.TimeoutError(
|
|
self.error_format.format(timeout_sec=self.timeout_sec))
|
|
|
|
|
|
def Run(cmd, timeout_sec, show_output=True, inpt=None):
|
|
"""Run command and optionally send the output to /dev/null or nul."""
|
|
with file_utils.FileWriter(os.devnull) as devnull:
|
|
stdout = devnull
|
|
stderr = devnull
|
|
stdin = None
|
|
if show_output:
|
|
stdout = None
|
|
stderr = None
|
|
if inpt:
|
|
stdin = subprocess.PIPE
|
|
# [py3 port] Should be able to use subprocess.run (etc) with 'timeout' param
|
|
# here and below. We're only using the Popen API in order to have a process
|
|
# to give to _TimeoutThread.
|
|
p = subprocess.Popen(cmd, stdout=stdout, stderr=stderr, stdin=stdin)
|
|
with _TimeoutThread(p.kill, timeout_sec):
|
|
if inpt:
|
|
p.communicate(six.ensure_binary(inpt))
|
|
else:
|
|
p.wait()
|
|
if p.returncode != 0:
|
|
raise subprocess.CalledProcessError(p.returncode, cmd)
|
|
|
|
|
|
def _GetStdout(cmd, timeout_sec, show_stderr=True):
|
|
p = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=None if show_stderr else subprocess.PIPE)
|
|
with _TimeoutThread(p.kill, timeout_sec):
|
|
stdout, _ = p.communicate()
|
|
if p.returncode != 0:
|
|
raise subprocess.CalledProcessError(p.returncode, cmd)
|
|
return six.ensure_text(stdout)
|
|
|
|
|
|
def GetOutputLines(cmd, timeout_sec, show_stderr=True, strip_output=False):
|
|
"""Run command and get its stdout as a list of lines.
|
|
|
|
Args:
|
|
cmd: List of executable and arg strings.
|
|
timeout_sec: Command will be killed if it exceeds this.
|
|
show_stderr: False to suppress stderr from the command.
|
|
strip_output: Strip head/tail whitespace before splitting into lines.
|
|
|
|
Returns:
|
|
List of lines (without newlines).
|
|
"""
|
|
stdout = _GetStdout(cmd, timeout_sec, show_stderr=show_stderr)
|
|
if strip_output:
|
|
stdout = stdout.strip()
|
|
lines = stdout.splitlines()
|
|
return lines
|
|
|
|
|
|
def GetOutputJson(cmd, timeout_sec, show_stderr=True):
|
|
"""Run command and get its JSON stdout as a parsed dict.
|
|
|
|
Args:
|
|
cmd: List of executable and arg strings.
|
|
timeout_sec: Command will be killed if it exceeds this.
|
|
show_stderr: False to suppress stderr from the command.
|
|
|
|
Returns:
|
|
Parsed JSON.
|
|
"""
|
|
stdout = _GetStdout(cmd, timeout_sec, show_stderr=show_stderr)
|
|
return json.loads(stdout.strip())
|
|
|
|
|
|
def StreamOutputJson(cmd, event_timeout_sec, show_stderr=True):
|
|
"""Run command and get its output streamed as an iterable of dicts.
|
|
|
|
Args:
|
|
cmd: List of executable and arg strings.
|
|
event_timeout_sec: Command will be killed if we don't get a JSON line for
|
|
this long. (This is not the same as timeout_sec above).
|
|
show_stderr: False to suppress stderr from the command.
|
|
|
|
Yields:
|
|
Parsed JSON.
|
|
|
|
Raises:
|
|
CalledProcessError: cmd returned with a non-zero exit code.
|
|
TimeoutError: cmd has timed out.
|
|
"""
|
|
p = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=None if show_stderr else subprocess.PIPE)
|
|
with _TimeoutThread(
|
|
p.kill,
|
|
event_timeout_sec,
|
|
error_format='No subprocess output for {timeout_sec} seconds') as timer:
|
|
for obj in json_stream.ReadJsonStream(p.stdout):
|
|
timer.Reset()
|
|
yield obj
|
|
p.wait()
|
|
if p.returncode != 0:
|
|
raise subprocess.CalledProcessError(p.returncode, cmd)
|