# -*- 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 -> gs://prefix-mid -> 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)