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)