| # -*- 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) |
| |