[catapult] - Copy Telemetry's gsutilz over to third_party.

third_party/gsutil/gslib/tab_complete.py

+# -*- coding: utf-8 -*-
+# Copyright 2014 Google Inc. 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.
+"""Shell tab completion."""
+
+import itertools
+import json
+import threading
+import time
+
+import boto
+
+from boto.gs.acl import CannedACLStrings
+from gslib.storage_url import IsFileUrlString
+from gslib.storage_url import StorageUrlFromString
+from gslib.storage_url import StripOneSlash
+from gslib.util import GetTabCompletionCacheFilename
+from gslib.util import GetTabCompletionLogFilename
+from gslib.wildcard_iterator import CreateWildcardIterator
+
+TAB_COMPLETE_CACHE_TTL = 15
+
+_TAB_COMPLETE_MAX_RESULTS = 1000
+
+_TIMEOUT_WARNING = """
+Tab completion aborted (took >%ss), you may complete the command manually.
+The timeout can be adjusted in the gsutil configuration file.
+""".rstrip()
+
+
+class CompleterType(object):
+  CLOUD_BUCKET = 'cloud_bucket'
+  CLOUD_OBJECT = 'cloud_object'
+  CLOUD_OR_LOCAL_OBJECT = 'cloud_or_local_object'
+  LOCAL_OBJECT = 'local_object'
+  LOCAL_OBJECT_OR_CANNED_ACL = 'local_object_or_canned_acl'
+  NO_OP = 'no_op'
+
+
+class LocalObjectCompleter(object):
+  """Completer object for local files."""
+
+  def __init__(self):
+    # This is only safe to import if argcomplete is present in the install
+    # (which happens for Cloud SDK installs), so import on usage, not on load.
+    # pylint: disable=g-import-not-at-top
+    from argcomplete.completers import FilesCompleter
+    self.files_completer = FilesCompleter()
+
+  def __call__(self, prefix, **kwargs):
+    return self.files_completer(prefix, **kwargs)
+
+
+class LocalObjectOrCannedACLCompleter(object):
+  """Completer object for local files and canned ACLs.
+
+  Currently, only Google Cloud Storage canned ACL names are supported.
+  """
+
+  def __init__(self):
+    self.local_object_completer = LocalObjectCompleter()
+
+  def __call__(self, prefix, **kwargs):
+    local_objects = self.local_object_completer(prefix, **kwargs)
+    canned_acls = [acl for acl in CannedACLStrings if acl.startswith(prefix)]
+    return local_objects + canned_acls
+
+
+class TabCompletionCache(object):
+  """Cache for tab completion results."""
+
+  def __init__(self, prefix, results, timestamp, partial_results):
+    self.prefix = prefix
+    self.results = results
+    self.timestamp = timestamp
+    self.partial_results = partial_results
+
+  @staticmethod
+  def LoadFromFile(filename):
+    """Instantiates the cache from a file.
+
+    Args:
+      filename: The file to load.
+    Returns:
+      TabCompletionCache instance with loaded data or an empty cache
+          if the file cannot be loaded
+    """
+    try:
+      with open(filename, 'r') as fp:
+        cache_dict = json.loads(fp.read())
+        prefix = cache_dict['prefix']
+        results = cache_dict['results']
+        timestamp = cache_dict['timestamp']
+        partial_results = cache_dict['partial-results']
+    except Exception:  # pylint: disable=broad-except
+      # Guarding against incompatible format changes in the cache file.
+      # Erring on the side of not breaking tab-completion in case of cache
+      # issues.
+      prefix = None
+      results = []
+      timestamp = 0
+      partial_results = False
+
+    return TabCompletionCache(prefix, results, timestamp, partial_results)
+
+  def GetCachedResults(self, prefix):
+    """Returns the cached results for prefix or None if not in cache."""
+    current_time = time.time()
+    if current_time - self.timestamp >= TAB_COMPLETE_CACHE_TTL:
+      return None
+
+    results = None
+
+    if prefix == self.prefix:
+      results = self.results
+    elif (not self.partial_results and prefix.startswith(self.prefix)
+          and prefix.count('/') == self.prefix.count('/')):
+      results = [x for x in self.results if x.startswith(prefix)]
+
+    if results is not None:
+      # Update cache timestamp to make sure the cache entry does not expire if
+      # the user is performing multiple completions in a single
+      # bucket/subdirectory since we can answer these requests from the cache.
+      # e.g. gs://prefix<tab> -> gs://prefix-mid<tab> -> gs://prefix-mid-suffix
+      self.timestamp = time.time()
+      return results
+
+  def UpdateCache(self, prefix, results, partial_results):
+    """Updates the in-memory cache with the results for the given prefix."""
+    self.prefix = prefix
+    self.results = results
+    self.partial_results = partial_results
+    self.timestamp = time.time()
+
+  def WriteToFile(self, filename):
+    """Writes out the cache to the given file."""
+    json_str = json.dumps({
+        'prefix': self.prefix,
+        'results': self.results,
+        'partial-results': self.partial_results,
+        'timestamp': self.timestamp,
+    })
+
+    try:
+      with open(filename, 'w') as fp:
+        fp.write(json_str)
+    except IOError:
+      pass
+
+
+class CloudListingRequestThread(threading.Thread):
+  """Thread that performs a listing request for the given URL string."""
+
+  def __init__(self, wildcard_url_str, gsutil_api):
+    """Instantiates Cloud listing request thread.
+
+    Args:
+      wildcard_url_str: The URL to list.
+      gsutil_api: gsutil Cloud API instance to use.
+    """
+    super(CloudListingRequestThread, self).__init__()
+    self.daemon = True
+    self._wildcard_url_str = wildcard_url_str
+    self._gsutil_api = gsutil_api
+    self.results = None
+
+  def run(self):
+    it = CreateWildcardIterator(
+        self._wildcard_url_str, self._gsutil_api).IterAll(
+            bucket_listing_fields=['name'])
+    self.results = [
+        str(c) for c in itertools.islice(it, _TAB_COMPLETE_MAX_RESULTS)]
+
+
+class TimeoutError(Exception):
+  pass
+
+
+class CloudObjectCompleter(object):
+  """Completer object for Cloud URLs."""
+
+  def __init__(self, gsutil_api, bucket_only=False):
+    """Instantiates completer for Cloud URLs.
+
+    Args:
+      gsutil_api: gsutil Cloud API instance to use.
+      bucket_only: Whether the completer should only match buckets.
+    """
+    self._gsutil_api = gsutil_api
+    self._bucket_only = bucket_only
+
+  def _PerformCloudListing(self, wildcard_url, timeout):
+    """Perform a remote listing request for the given wildcard URL.
+
+    Args:
+      wildcard_url: The wildcard URL to list.
+      timeout: Time limit for the request.
+    Returns:
+      Cloud resources matching the given wildcard URL.
+    Raises:
+      TimeoutError: If the listing does not finish within the timeout.
+    """
+    request_thread = CloudListingRequestThread(wildcard_url, self._gsutil_api)
+    request_thread.start()
+    request_thread.join(timeout)
+
+    if request_thread.is_alive():
+      # This is only safe to import if argcomplete is present in the install
+      # (which happens for Cloud SDK installs), so import on usage, not on load.
+      # pylint: disable=g-import-not-at-top
+      import argcomplete
+      argcomplete.warn(_TIMEOUT_WARNING % timeout)
+      raise TimeoutError()
+
+    results = request_thread.results
+
+    return results
+
+  def __call__(self, prefix, **kwargs):
+    if not prefix:
+      prefix = 'gs://'
+    elif IsFileUrlString(prefix):
+      return []
+
+    wildcard_url = prefix + '*'
+    url = StorageUrlFromString(wildcard_url)
+    if self._bucket_only and not url.IsBucket():
+      return []
+
+    timeout = boto.config.getint('GSUtil', 'tab_completion_timeout', 5)
+    if timeout == 0:
+      return []
+
+    start_time = time.time()
+
+    cache = TabCompletionCache.LoadFromFile(GetTabCompletionCacheFilename())
+    cached_results = cache.GetCachedResults(prefix)
+
+    timing_log_entry_type = ''
+    if cached_results is not None:
+      results = cached_results
+      timing_log_entry_type = ' (from cache)'
+    else:
+      try:
+        results = self._PerformCloudListing(wildcard_url, timeout)
+        if self._bucket_only and len(results) == 1:
+          results = [StripOneSlash(results[0])]
+        partial_results = (len(results) == _TAB_COMPLETE_MAX_RESULTS)
+        cache.UpdateCache(prefix, results, partial_results)
+      except TimeoutError:
+        timing_log_entry_type = ' (request timeout)'
+        results = []
+
+    cache.WriteToFile(GetTabCompletionCacheFilename())
+
+    end_time = time.time()
+    num_results = len(results)
+    elapsed_seconds = end_time - start_time
+    _WriteTimingLog(
+        '%s results%s in %.2fs, %.2f results/second for prefix: %s\n' %
+        (num_results, timing_log_entry_type, elapsed_seconds,
+         num_results / elapsed_seconds, prefix))
+
+    return results
+
+
+class CloudOrLocalObjectCompleter(object):
+  """Completer object for Cloud URLs or local files.
+
+  Invokes the Cloud object completer if the input looks like a Cloud URL and
+  falls back to local file completer otherwise.
+  """
+
+  def __init__(self, gsutil_api):
+    self.cloud_object_completer = CloudObjectCompleter(gsutil_api)
+    self.local_object_completer = LocalObjectCompleter()
+
+  def __call__(self, prefix, **kwargs):
+    if IsFileUrlString(prefix):
+      completer = self.local_object_completer
+    else:
+      completer = self.cloud_object_completer
+    return completer(prefix, **kwargs)
+
+
+class NoOpCompleter(object):
+  """Completer that always returns 0 results."""
+
+  def __call__(self, unused_prefix, **unused_kwargs):
+    return []
+
+
+def MakeCompleter(completer_type, gsutil_api):
+  """Create a completer instance of the given type.
+
+  Args:
+    completer_type: The type of completer to create.
+    gsutil_api: gsutil Cloud API instance to use.
+  Returns:
+    A completer instance.
+  Raises:
+    RuntimeError: if completer type is not supported.
+  """
+  if completer_type == CompleterType.CLOUD_OR_LOCAL_OBJECT:
+    return CloudOrLocalObjectCompleter(gsutil_api)
+  elif completer_type == CompleterType.LOCAL_OBJECT:
+    return LocalObjectCompleter()
+  elif completer_type == CompleterType.LOCAL_OBJECT_OR_CANNED_ACL:
+    return LocalObjectOrCannedACLCompleter()
+  elif completer_type == CompleterType.CLOUD_BUCKET:
+    return CloudObjectCompleter(gsutil_api, bucket_only=True)
+  elif completer_type == CompleterType.CLOUD_OBJECT:
+    return CloudObjectCompleter(gsutil_api)
+  elif completer_type == CompleterType.NO_OP:
+    return NoOpCompleter()
+  else:
+    raise RuntimeError(
+        'Unknown completer "%s"' % completer_type)
+
+
+def _WriteTimingLog(message):
+  """Write an entry to the tab completion timing log, if it's enabled."""
+  if boto.config.getbool('GSUtil', 'tab_completion_time_logs', False):
+    with open(GetTabCompletionLogFilename(), 'ab') as fp:
+      fp.write(message)
+