OLD | NEW |
| (Empty) |
1 # -*- coding: utf-8 -*- | |
2 # Copyright 2014 Google Inc. All Rights Reserved. | |
3 # | |
4 # Licensed under the Apache License, Version 2.0 (the "License"); | |
5 # you may not use this file except in compliance with the License. | |
6 # You may obtain a copy of the License at | |
7 # | |
8 # http://www.apache.org/licenses/LICENSE-2.0 | |
9 # | |
10 # Unless required by applicable law or agreed to in writing, software | |
11 # distributed under the License is distributed on an "AS IS" BASIS, | |
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
13 # See the License for the specific language governing permissions and | |
14 # limitations under the License. | |
15 """Shell tab completion.""" | |
16 | |
17 import itertools | |
18 import json | |
19 import threading | |
20 import time | |
21 | |
22 import boto | |
23 | |
24 from boto.gs.acl import CannedACLStrings | |
25 from gslib.storage_url import IsFileUrlString | |
26 from gslib.storage_url import StorageUrlFromString | |
27 from gslib.storage_url import StripOneSlash | |
28 from gslib.util import GetTabCompletionCacheFilename | |
29 from gslib.util import GetTabCompletionLogFilename | |
30 from gslib.wildcard_iterator import CreateWildcardIterator | |
31 | |
32 TAB_COMPLETE_CACHE_TTL = 15 | |
33 | |
34 _TAB_COMPLETE_MAX_RESULTS = 1000 | |
35 | |
36 _TIMEOUT_WARNING = """ | |
37 Tab completion aborted (took >%ss), you may complete the command manually. | |
38 The timeout can be adjusted in the gsutil configuration file. | |
39 """.rstrip() | |
40 | |
41 | |
42 class CompleterType(object): | |
43 CLOUD_BUCKET = 'cloud_bucket' | |
44 CLOUD_OBJECT = 'cloud_object' | |
45 CLOUD_OR_LOCAL_OBJECT = 'cloud_or_local_object' | |
46 LOCAL_OBJECT = 'local_object' | |
47 LOCAL_OBJECT_OR_CANNED_ACL = 'local_object_or_canned_acl' | |
48 NO_OP = 'no_op' | |
49 | |
50 | |
51 class LocalObjectCompleter(object): | |
52 """Completer object for local files.""" | |
53 | |
54 def __init__(self): | |
55 # This is only safe to import if argcomplete is present in the install | |
56 # (which happens for Cloud SDK installs), so import on usage, not on load. | |
57 # pylint: disable=g-import-not-at-top | |
58 from argcomplete.completers import FilesCompleter | |
59 self.files_completer = FilesCompleter() | |
60 | |
61 def __call__(self, prefix, **kwargs): | |
62 return self.files_completer(prefix, **kwargs) | |
63 | |
64 | |
65 class LocalObjectOrCannedACLCompleter(object): | |
66 """Completer object for local files and canned ACLs. | |
67 | |
68 Currently, only Google Cloud Storage canned ACL names are supported. | |
69 """ | |
70 | |
71 def __init__(self): | |
72 self.local_object_completer = LocalObjectCompleter() | |
73 | |
74 def __call__(self, prefix, **kwargs): | |
75 local_objects = self.local_object_completer(prefix, **kwargs) | |
76 canned_acls = [acl for acl in CannedACLStrings if acl.startswith(prefix)] | |
77 return local_objects + canned_acls | |
78 | |
79 | |
80 class TabCompletionCache(object): | |
81 """Cache for tab completion results.""" | |
82 | |
83 def __init__(self, prefix, results, timestamp, partial_results): | |
84 self.prefix = prefix | |
85 self.results = results | |
86 self.timestamp = timestamp | |
87 self.partial_results = partial_results | |
88 | |
89 @staticmethod | |
90 def LoadFromFile(filename): | |
91 """Instantiates the cache from a file. | |
92 | |
93 Args: | |
94 filename: The file to load. | |
95 Returns: | |
96 TabCompletionCache instance with loaded data or an empty cache | |
97 if the file cannot be loaded | |
98 """ | |
99 try: | |
100 with open(filename, 'r') as fp: | |
101 cache_dict = json.loads(fp.read()) | |
102 prefix = cache_dict['prefix'] | |
103 results = cache_dict['results'] | |
104 timestamp = cache_dict['timestamp'] | |
105 partial_results = cache_dict['partial-results'] | |
106 except Exception: # pylint: disable=broad-except | |
107 # Guarding against incompatible format changes in the cache file. | |
108 # Erring on the side of not breaking tab-completion in case of cache | |
109 # issues. | |
110 prefix = None | |
111 results = [] | |
112 timestamp = 0 | |
113 partial_results = False | |
114 | |
115 return TabCompletionCache(prefix, results, timestamp, partial_results) | |
116 | |
117 def GetCachedResults(self, prefix): | |
118 """Returns the cached results for prefix or None if not in cache.""" | |
119 current_time = time.time() | |
120 if current_time - self.timestamp >= TAB_COMPLETE_CACHE_TTL: | |
121 return None | |
122 | |
123 results = None | |
124 | |
125 if prefix == self.prefix: | |
126 results = self.results | |
127 elif (not self.partial_results and prefix.startswith(self.prefix) | |
128 and prefix.count('/') == self.prefix.count('/')): | |
129 results = [x for x in self.results if x.startswith(prefix)] | |
130 | |
131 if results is not None: | |
132 # Update cache timestamp to make sure the cache entry does not expire if | |
133 # the user is performing multiple completions in a single | |
134 # bucket/subdirectory since we can answer these requests from the cache. | |
135 # e.g. gs://prefix<tab> -> gs://prefix-mid<tab> -> gs://prefix-mid-suffix | |
136 self.timestamp = time.time() | |
137 return results | |
138 | |
139 def UpdateCache(self, prefix, results, partial_results): | |
140 """Updates the in-memory cache with the results for the given prefix.""" | |
141 self.prefix = prefix | |
142 self.results = results | |
143 self.partial_results = partial_results | |
144 self.timestamp = time.time() | |
145 | |
146 def WriteToFile(self, filename): | |
147 """Writes out the cache to the given file.""" | |
148 json_str = json.dumps({ | |
149 'prefix': self.prefix, | |
150 'results': self.results, | |
151 'partial-results': self.partial_results, | |
152 'timestamp': self.timestamp, | |
153 }) | |
154 | |
155 try: | |
156 with open(filename, 'w') as fp: | |
157 fp.write(json_str) | |
158 except IOError: | |
159 pass | |
160 | |
161 | |
162 class CloudListingRequestThread(threading.Thread): | |
163 """Thread that performs a listing request for the given URL string.""" | |
164 | |
165 def __init__(self, wildcard_url_str, gsutil_api): | |
166 """Instantiates Cloud listing request thread. | |
167 | |
168 Args: | |
169 wildcard_url_str: The URL to list. | |
170 gsutil_api: gsutil Cloud API instance to use. | |
171 """ | |
172 super(CloudListingRequestThread, self).__init__() | |
173 self.daemon = True | |
174 self._wildcard_url_str = wildcard_url_str | |
175 self._gsutil_api = gsutil_api | |
176 self.results = None | |
177 | |
178 def run(self): | |
179 it = CreateWildcardIterator( | |
180 self._wildcard_url_str, self._gsutil_api).IterAll( | |
181 bucket_listing_fields=['name']) | |
182 self.results = [ | |
183 str(c) for c in itertools.islice(it, _TAB_COMPLETE_MAX_RESULTS)] | |
184 | |
185 | |
186 class TimeoutError(Exception): | |
187 pass | |
188 | |
189 | |
190 class CloudObjectCompleter(object): | |
191 """Completer object for Cloud URLs.""" | |
192 | |
193 def __init__(self, gsutil_api, bucket_only=False): | |
194 """Instantiates completer for Cloud URLs. | |
195 | |
196 Args: | |
197 gsutil_api: gsutil Cloud API instance to use. | |
198 bucket_only: Whether the completer should only match buckets. | |
199 """ | |
200 self._gsutil_api = gsutil_api | |
201 self._bucket_only = bucket_only | |
202 | |
203 def _PerformCloudListing(self, wildcard_url, timeout): | |
204 """Perform a remote listing request for the given wildcard URL. | |
205 | |
206 Args: | |
207 wildcard_url: The wildcard URL to list. | |
208 timeout: Time limit for the request. | |
209 Returns: | |
210 Cloud resources matching the given wildcard URL. | |
211 Raises: | |
212 TimeoutError: If the listing does not finish within the timeout. | |
213 """ | |
214 request_thread = CloudListingRequestThread(wildcard_url, self._gsutil_api) | |
215 request_thread.start() | |
216 request_thread.join(timeout) | |
217 | |
218 if request_thread.is_alive(): | |
219 # This is only safe to import if argcomplete is present in the install | |
220 # (which happens for Cloud SDK installs), so import on usage, not on load. | |
221 # pylint: disable=g-import-not-at-top | |
222 import argcomplete | |
223 argcomplete.warn(_TIMEOUT_WARNING % timeout) | |
224 raise TimeoutError() | |
225 | |
226 results = request_thread.results | |
227 | |
228 return results | |
229 | |
230 def __call__(self, prefix, **kwargs): | |
231 if not prefix: | |
232 prefix = 'gs://' | |
233 elif IsFileUrlString(prefix): | |
234 return [] | |
235 | |
236 wildcard_url = prefix + '*' | |
237 url = StorageUrlFromString(wildcard_url) | |
238 if self._bucket_only and not url.IsBucket(): | |
239 return [] | |
240 | |
241 timeout = boto.config.getint('GSUtil', 'tab_completion_timeout', 5) | |
242 if timeout == 0: | |
243 return [] | |
244 | |
245 start_time = time.time() | |
246 | |
247 cache = TabCompletionCache.LoadFromFile(GetTabCompletionCacheFilename()) | |
248 cached_results = cache.GetCachedResults(prefix) | |
249 | |
250 timing_log_entry_type = '' | |
251 if cached_results is not None: | |
252 results = cached_results | |
253 timing_log_entry_type = ' (from cache)' | |
254 else: | |
255 try: | |
256 results = self._PerformCloudListing(wildcard_url, timeout) | |
257 if self._bucket_only and len(results) == 1: | |
258 results = [StripOneSlash(results[0])] | |
259 partial_results = (len(results) == _TAB_COMPLETE_MAX_RESULTS) | |
260 cache.UpdateCache(prefix, results, partial_results) | |
261 except TimeoutError: | |
262 timing_log_entry_type = ' (request timeout)' | |
263 results = [] | |
264 | |
265 cache.WriteToFile(GetTabCompletionCacheFilename()) | |
266 | |
267 end_time = time.time() | |
268 num_results = len(results) | |
269 elapsed_seconds = end_time - start_time | |
270 _WriteTimingLog( | |
271 '%s results%s in %.2fs, %.2f results/second for prefix: %s\n' % | |
272 (num_results, timing_log_entry_type, elapsed_seconds, | |
273 num_results / elapsed_seconds, prefix)) | |
274 | |
275 return results | |
276 | |
277 | |
278 class CloudOrLocalObjectCompleter(object): | |
279 """Completer object for Cloud URLs or local files. | |
280 | |
281 Invokes the Cloud object completer if the input looks like a Cloud URL and | |
282 falls back to local file completer otherwise. | |
283 """ | |
284 | |
285 def __init__(self, gsutil_api): | |
286 self.cloud_object_completer = CloudObjectCompleter(gsutil_api) | |
287 self.local_object_completer = LocalObjectCompleter() | |
288 | |
289 def __call__(self, prefix, **kwargs): | |
290 if IsFileUrlString(prefix): | |
291 completer = self.local_object_completer | |
292 else: | |
293 completer = self.cloud_object_completer | |
294 return completer(prefix, **kwargs) | |
295 | |
296 | |
297 class NoOpCompleter(object): | |
298 """Completer that always returns 0 results.""" | |
299 | |
300 def __call__(self, unused_prefix, **unused_kwargs): | |
301 return [] | |
302 | |
303 | |
304 def MakeCompleter(completer_type, gsutil_api): | |
305 """Create a completer instance of the given type. | |
306 | |
307 Args: | |
308 completer_type: The type of completer to create. | |
309 gsutil_api: gsutil Cloud API instance to use. | |
310 Returns: | |
311 A completer instance. | |
312 Raises: | |
313 RuntimeError: if completer type is not supported. | |
314 """ | |
315 if completer_type == CompleterType.CLOUD_OR_LOCAL_OBJECT: | |
316 return CloudOrLocalObjectCompleter(gsutil_api) | |
317 elif completer_type == CompleterType.LOCAL_OBJECT: | |
318 return LocalObjectCompleter() | |
319 elif completer_type == CompleterType.LOCAL_OBJECT_OR_CANNED_ACL: | |
320 return LocalObjectOrCannedACLCompleter() | |
321 elif completer_type == CompleterType.CLOUD_BUCKET: | |
322 return CloudObjectCompleter(gsutil_api, bucket_only=True) | |
323 elif completer_type == CompleterType.CLOUD_OBJECT: | |
324 return CloudObjectCompleter(gsutil_api) | |
325 elif completer_type == CompleterType.NO_OP: | |
326 return NoOpCompleter() | |
327 else: | |
328 raise RuntimeError( | |
329 'Unknown completer "%s"' % completer_type) | |
330 | |
331 | |
332 def _WriteTimingLog(message): | |
333 """Write an entry to the tab completion timing log, if it's enabled.""" | |
334 if boto.config.getbool('GSUtil', 'tab_completion_time_logs', False): | |
335 with open(GetTabCompletionLogFilename(), 'ab') as fp: | |
336 fp.write(message) | |
337 | |
OLD | NEW |