| # -*- 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. |
| """Media helper functions and classes for Google Cloud Storage JSON API.""" |
| |
| from __future__ import absolute_import |
| |
| import copy |
| import cStringIO |
| import httplib |
| import logging |
| import socket |
| import types |
| import urlparse |
| |
| from apitools.base.py import exceptions as apitools_exceptions |
| import httplib2 |
| from httplib2 import parse_uri |
| |
| from gslib.cloud_api import BadRequestException |
| from gslib.progress_callback import ProgressCallbackWithBackoff |
| from gslib.util import SSL_TIMEOUT |
| from gslib.util import TRANSFER_BUFFER_SIZE |
| |
| |
| class BytesTransferredContainer(object): |
| """Container class for passing number of bytes transferred to lower layers. |
| |
| For resumed transfers or connection rebuilds in the middle of a transfer, we |
| need to rebuild the connection class with how much we've transferred so far. |
| For uploads, we don't know the total number of bytes uploaded until we've |
| queried the server, but we need to create the connection class to pass to |
| httplib2 before we can query the server. This container object allows us to |
| pass a reference into Upload/DownloadCallbackConnection. |
| """ |
| |
| def __init__(self): |
| self.__bytes_transferred = 0 |
| |
| @property |
| def bytes_transferred(self): |
| return self.__bytes_transferred |
| |
| @bytes_transferred.setter |
| def bytes_transferred(self, value): |
| self.__bytes_transferred = value |
| |
| |
| class UploadCallbackConnectionClassFactory(object): |
| """Creates a class that can override an httplib2 connection. |
| |
| This is used to provide progress callbacks and disable dumping the upload |
| payload during debug statements. It can later be used to provide on-the-fly |
| hash digestion during upload. |
| """ |
| |
| def __init__(self, bytes_uploaded_container, |
| buffer_size=TRANSFER_BUFFER_SIZE, |
| total_size=0, progress_callback=None): |
| self.bytes_uploaded_container = bytes_uploaded_container |
| self.buffer_size = buffer_size |
| self.total_size = total_size |
| self.progress_callback = progress_callback |
| |
| def GetConnectionClass(self): |
| """Returns a connection class that overrides send.""" |
| outer_bytes_uploaded_container = self.bytes_uploaded_container |
| outer_buffer_size = self.buffer_size |
| outer_total_size = self.total_size |
| outer_progress_callback = self.progress_callback |
| |
| class UploadCallbackConnection(httplib2.HTTPSConnectionWithTimeout): |
| """Connection class override for uploads.""" |
| bytes_uploaded_container = outer_bytes_uploaded_container |
| # After we instantiate this class, apitools will check with the server |
| # to find out how many bytes remain for a resumable upload. This allows |
| # us to update our progress once based on that number. |
| processed_initial_bytes = False |
| GCS_JSON_BUFFER_SIZE = outer_buffer_size |
| callback_processor = None |
| size = outer_total_size |
| |
| def __init__(self, *args, **kwargs): |
| kwargs['timeout'] = SSL_TIMEOUT |
| httplib2.HTTPSConnectionWithTimeout.__init__(self, *args, **kwargs) |
| |
| def send(self, data): |
| """Overrides HTTPConnection.send.""" |
| if not self.processed_initial_bytes: |
| self.processed_initial_bytes = True |
| if outer_progress_callback: |
| self.callback_processor = ProgressCallbackWithBackoff( |
| outer_total_size, outer_progress_callback) |
| self.callback_processor.Progress( |
| self.bytes_uploaded_container.bytes_transferred) |
| # httplib.HTTPConnection.send accepts either a string or a file-like |
| # object (anything that implements read()). |
| if isinstance(data, basestring): |
| full_buffer = cStringIO.StringIO(data) |
| else: |
| full_buffer = data |
| partial_buffer = full_buffer.read(self.GCS_JSON_BUFFER_SIZE) |
| while partial_buffer: |
| httplib2.HTTPSConnectionWithTimeout.send(self, partial_buffer) |
| send_length = len(partial_buffer) |
| if self.callback_processor: |
| # This is the only place where gsutil has control over making a |
| # callback, but here we can't differentiate the metadata bytes |
| # (such as headers and OAuth2 refreshes) sent during an upload |
| # from the actual upload bytes, so we will actually report |
| # slightly more bytes than desired to the callback handler. |
| # |
| # One considered/rejected alternative is to move the callbacks |
| # into the HashingFileUploadWrapper which only processes reads on |
| # the bytes. This has the disadvantages of being removed from |
| # where we actually send the bytes and unnecessarily |
| # multi-purposing that class. |
| self.callback_processor.Progress(send_length) |
| partial_buffer = full_buffer.read(self.GCS_JSON_BUFFER_SIZE) |
| |
| return UploadCallbackConnection |
| |
| |
| def WrapUploadHttpRequest(upload_http): |
| """Wraps upload_http so we only use our custom connection_type on PUTs. |
| |
| POSTs are used to refresh oauth tokens, and we don't want to process the |
| data sent in those requests. |
| |
| Args: |
| upload_http: httplib2.Http instance to wrap |
| """ |
| request_orig = upload_http.request |
| def NewRequest(uri, method='GET', body=None, headers=None, |
| redirections=httplib2.DEFAULT_MAX_REDIRECTS, |
| connection_type=None): |
| if method == 'PUT' or method == 'POST': |
| override_connection_type = connection_type |
| else: |
| override_connection_type = None |
| return request_orig(uri, method=method, body=body, |
| headers=headers, redirections=redirections, |
| connection_type=override_connection_type) |
| # Replace the request method with our own closure. |
| upload_http.request = NewRequest |
| |
| |
| class DownloadCallbackConnectionClassFactory(object): |
| """Creates a class that can override an httplib2 connection. |
| |
| This is used to provide progress callbacks, disable dumping the download |
| payload during debug statements, and provide on-the-fly hash digestion during |
| download. On-the-fly digestion is particularly important because httplib2 |
| will decompress gzipped content on-the-fly, thus this class provides our |
| only opportunity to calculate the correct hash for an object that has a |
| gzip hash in the cloud. |
| """ |
| |
| def __init__(self, bytes_downloaded_container, |
| buffer_size=TRANSFER_BUFFER_SIZE, total_size=0, |
| progress_callback=None, digesters=None): |
| self.buffer_size = buffer_size |
| self.total_size = total_size |
| self.progress_callback = progress_callback |
| self.digesters = digesters |
| self.bytes_downloaded_container = bytes_downloaded_container |
| |
| def GetConnectionClass(self): |
| """Returns a connection class that overrides getresponse.""" |
| |
| class DownloadCallbackConnection(httplib2.HTTPSConnectionWithTimeout): |
| """Connection class override for downloads.""" |
| outer_total_size = self.total_size |
| outer_digesters = self.digesters |
| outer_progress_callback = self.progress_callback |
| outer_bytes_downloaded_container = self.bytes_downloaded_container |
| processed_initial_bytes = False |
| callback_processor = None |
| |
| def __init__(self, *args, **kwargs): |
| kwargs['timeout'] = SSL_TIMEOUT |
| httplib2.HTTPSConnectionWithTimeout.__init__(self, *args, **kwargs) |
| |
| def getresponse(self, buffering=False): |
| """Wraps an HTTPResponse to perform callbacks and hashing. |
| |
| In this function, self is a DownloadCallbackConnection. |
| |
| Args: |
| buffering: Unused. This function uses a local buffer. |
| |
| Returns: |
| HTTPResponse object with wrapped read function. |
| """ |
| orig_response = httplib.HTTPConnection.getresponse(self) |
| if orig_response.status not in (httplib.OK, httplib.PARTIAL_CONTENT): |
| return orig_response |
| orig_read_func = orig_response.read |
| |
| def read(amt=None): # pylint: disable=invalid-name |
| """Overrides HTTPConnection.getresponse.read. |
| |
| This function only supports reads of TRANSFER_BUFFER_SIZE or smaller. |
| |
| Args: |
| amt: Integer n where 0 < n <= TRANSFER_BUFFER_SIZE. This is a |
| keyword argument to match the read function it overrides, |
| but it is required. |
| |
| Returns: |
| Data read from HTTPConnection. |
| """ |
| if not amt or amt > TRANSFER_BUFFER_SIZE: |
| raise BadRequestException( |
| 'Invalid HTTP read size %s during download, expected %s.' % |
| (amt, TRANSFER_BUFFER_SIZE)) |
| else: |
| amt = amt or TRANSFER_BUFFER_SIZE |
| |
| if not self.processed_initial_bytes: |
| self.processed_initial_bytes = True |
| if self.outer_progress_callback: |
| self.callback_processor = ProgressCallbackWithBackoff( |
| self.outer_total_size, self.outer_progress_callback) |
| self.callback_processor.Progress( |
| self.outer_bytes_downloaded_container.bytes_transferred) |
| |
| data = orig_read_func(amt) |
| read_length = len(data) |
| if self.callback_processor: |
| self.callback_processor.Progress(read_length) |
| if self.outer_digesters: |
| for alg in self.outer_digesters: |
| self.outer_digesters[alg].update(data) |
| return data |
| orig_response.read = read |
| |
| return orig_response |
| return DownloadCallbackConnection |
| |
| |
| def WrapDownloadHttpRequest(download_http): |
| """Overrides download request functions for an httplib2.Http object. |
| |
| Args: |
| download_http: httplib2.Http.object to wrap / override. |
| |
| Returns: |
| Wrapped / overridden httplib2.Http object. |
| """ |
| |
| # httplib2 has a bug https://code.google.com/p/httplib2/issues/detail?id=305 |
| # where custom connection_type is not respected after redirects. This |
| # function is copied from httplib2 and overrides the request function so that |
| # the connection_type is properly passed through. |
| # pylint: disable=protected-access,g-inconsistent-quotes,unused-variable |
| # pylint: disable=g-equals-none,g-doc-return-or-yield |
| # pylint: disable=g-short-docstring-punctuation,g-doc-args |
| # pylint: disable=too-many-statements |
| def OverrideRequest(self, conn, host, absolute_uri, request_uri, method, |
| body, headers, redirections, cachekey): |
| """Do the actual request using the connection object. |
| |
| Also follow one level of redirects if necessary. |
| """ |
| |
| auths = ([(auth.depth(request_uri), auth) for auth in self.authorizations |
| if auth.inscope(host, request_uri)]) |
| auth = auths and sorted(auths)[0][1] or None |
| if auth: |
| auth.request(method, request_uri, headers, body) |
| |
| (response, content) = self._conn_request(conn, request_uri, method, body, |
| headers) |
| |
| if auth: |
| if auth.response(response, body): |
| auth.request(method, request_uri, headers, body) |
| (response, content) = self._conn_request(conn, request_uri, method, |
| body, headers) |
| response._stale_digest = 1 |
| |
| if response.status == 401: |
| for authorization in self._auth_from_challenge( |
| host, request_uri, headers, response, content): |
| authorization.request(method, request_uri, headers, body) |
| (response, content) = self._conn_request(conn, request_uri, method, |
| body, headers) |
| if response.status != 401: |
| self.authorizations.append(authorization) |
| authorization.response(response, body) |
| break |
| |
| if (self.follow_all_redirects or (method in ["GET", "HEAD"]) |
| or response.status == 303): |
| if self.follow_redirects and response.status in [300, 301, 302, |
| 303, 307]: |
| # Pick out the location header and basically start from the beginning |
| # remembering first to strip the ETag header and decrement our 'depth' |
| if redirections: |
| if not response.has_key('location') and response.status != 300: |
| raise httplib2.RedirectMissingLocation( |
| "Redirected but the response is missing a Location: header.", |
| response, content) |
| # Fix-up relative redirects (which violate an RFC 2616 MUST) |
| if response.has_key('location'): |
| location = response['location'] |
| (scheme, authority, path, query, fragment) = parse_uri(location) |
| if authority == None: |
| response['location'] = urlparse.urljoin(absolute_uri, location) |
| if response.status == 301 and method in ["GET", "HEAD"]: |
| response['-x-permanent-redirect-url'] = response['location'] |
| if not response.has_key('content-location'): |
| response['content-location'] = absolute_uri |
| httplib2._updateCache(headers, response, content, self.cache, |
| cachekey) |
| if headers.has_key('if-none-match'): |
| del headers['if-none-match'] |
| if headers.has_key('if-modified-since'): |
| del headers['if-modified-since'] |
| if ('authorization' in headers and |
| not self.forward_authorization_headers): |
| del headers['authorization'] |
| if response.has_key('location'): |
| location = response['location'] |
| old_response = copy.deepcopy(response) |
| if not old_response.has_key('content-location'): |
| old_response['content-location'] = absolute_uri |
| redirect_method = method |
| if response.status in [302, 303]: |
| redirect_method = "GET" |
| body = None |
| (response, content) = self.request( |
| location, redirect_method, body=body, headers=headers, |
| redirections=redirections-1, |
| connection_type=conn.__class__) |
| response.previous = old_response |
| else: |
| raise httplib2.RedirectLimit( |
| "Redirected more times than redirection_limit allows.", |
| response, content) |
| elif response.status in [200, 203] and method in ["GET", "HEAD"]: |
| # Don't cache 206's since we aren't going to handle byte range |
| # requests |
| if not response.has_key('content-location'): |
| response['content-location'] = absolute_uri |
| httplib2._updateCache(headers, response, content, self.cache, |
| cachekey) |
| |
| return (response, content) |
| |
| # Wrap download_http so we do not use our custom connection_type |
| # on POSTS, which are used to refresh oauth tokens. We don't want to |
| # process the data received in those requests. |
| request_orig = download_http.request |
| def NewRequest(uri, method='GET', body=None, headers=None, |
| redirections=httplib2.DEFAULT_MAX_REDIRECTS, |
| connection_type=None): |
| if method == 'POST': |
| return request_orig(uri, method=method, body=body, |
| headers=headers, redirections=redirections, |
| connection_type=None) |
| else: |
| return request_orig(uri, method=method, body=body, |
| headers=headers, redirections=redirections, |
| connection_type=connection_type) |
| |
| # Replace the request methods with our own closures. |
| download_http._request = types.MethodType(OverrideRequest, download_http) |
| download_http.request = NewRequest |
| |
| return download_http |
| |
| |
| class HttpWithNoRetries(httplib2.Http): |
| """httplib2.Http variant that does not retry. |
| |
| httplib2 automatically retries requests according to httplib2.RETRIES, but |
| in certain cases httplib2 ignores the RETRIES value and forces a retry. |
| Because httplib2 does not handle the case where the underlying request body |
| is a stream, a retry may cause a non-idempotent write as the stream is |
| partially consumed and not reset before the retry occurs. |
| |
| Here we override _conn_request to disable retries unequivocally, so that |
| uploads may be retried at higher layers that properly handle stream request |
| bodies. |
| """ |
| |
| def _conn_request(self, conn, request_uri, method, body, headers): # pylint: disable=too-many-statements |
| |
| try: |
| if hasattr(conn, 'sock') and conn.sock is None: |
| conn.connect() |
| conn.request(method, request_uri, body, headers) |
| except socket.timeout: |
| raise |
| except socket.gaierror: |
| conn.close() |
| raise httplib2.ServerNotFoundError( |
| 'Unable to find the server at %s' % conn.host) |
| except httplib2.ssl_SSLError: |
| conn.close() |
| raise |
| except socket.error, e: |
| err = 0 |
| if hasattr(e, 'args'): |
| err = getattr(e, 'args')[0] |
| else: |
| err = e.errno |
| if err == httplib2.errno.ECONNREFUSED: # Connection refused |
| raise |
| except httplib.HTTPException: |
| conn.close() |
| raise |
| try: |
| response = conn.getresponse() |
| except (socket.error, httplib.HTTPException): |
| conn.close() |
| raise |
| else: |
| content = '' |
| if method == 'HEAD': |
| conn.close() |
| else: |
| content = response.read() |
| response = httplib2.Response(response) |
| if method != 'HEAD': |
| # pylint: disable=protected-access |
| content = httplib2._decompressContent(response, content) |
| return (response, content) |
| |
| |
| class HttpWithDownloadStream(httplib2.Http): |
| """httplib2.Http variant that only pushes bytes through a stream. |
| |
| httplib2 handles media by storing entire chunks of responses in memory, which |
| is undesirable particularly when multiple instances are used during |
| multi-threaded/multi-process copy. This class copies and then overrides some |
| httplib2 functions to use a streaming copy approach that uses small memory |
| buffers. |
| |
| Also disables httplib2 retries (for reasons stated in the HttpWithNoRetries |
| class doc). |
| """ |
| |
| def __init__(self, *args, **kwds): |
| self._stream = None |
| self._logger = logging.getLogger() |
| super(HttpWithDownloadStream, self).__init__(*args, **kwds) |
| |
| @property |
| def stream(self): |
| return self._stream |
| |
| @stream.setter |
| def stream(self, value): |
| self._stream = value |
| |
| def _conn_request(self, conn, request_uri, method, body, headers): # pylint: disable=too-many-statements |
| try: |
| if hasattr(conn, 'sock') and conn.sock is None: |
| conn.connect() |
| conn.request(method, request_uri, body, headers) |
| except socket.timeout: |
| raise |
| except socket.gaierror: |
| conn.close() |
| raise httplib2.ServerNotFoundError( |
| 'Unable to find the server at %s' % conn.host) |
| except httplib2.ssl_SSLError: |
| conn.close() |
| raise |
| except socket.error, e: |
| err = 0 |
| if hasattr(e, 'args'): |
| err = getattr(e, 'args')[0] |
| else: |
| err = e.errno |
| if err == httplib2.errno.ECONNREFUSED: # Connection refused |
| raise |
| except httplib.HTTPException: |
| # Just because the server closed the connection doesn't apparently mean |
| # that the server didn't send a response. |
| conn.close() |
| raise |
| try: |
| response = conn.getresponse() |
| except (socket.error, httplib.HTTPException): |
| conn.close() |
| raise |
| else: |
| content = '' |
| if method == 'HEAD': |
| conn.close() |
| response = httplib2.Response(response) |
| else: |
| if response.status in (httplib.OK, httplib.PARTIAL_CONTENT): |
| content_length = None |
| if hasattr(response, 'msg'): |
| content_length = response.getheader('content-length') |
| http_stream = response |
| bytes_read = 0 |
| while True: |
| new_data = http_stream.read(TRANSFER_BUFFER_SIZE) |
| if new_data: |
| if self.stream is None: |
| raise apitools_exceptions.InvalidUserInputError( |
| 'Cannot exercise HttpWithDownloadStream with no stream') |
| self.stream.write(new_data) |
| bytes_read += len(new_data) |
| else: |
| break |
| |
| if (content_length is not None and |
| long(bytes_read) != long(content_length)): |
| # The input stream terminated before we were able to read the |
| # entire contents, possibly due to a network condition. Set |
| # content-length to indicate how many bytes we actually read. |
| self._logger.log( |
| logging.DEBUG, 'Only got %s bytes out of content-length %s ' |
| 'for request URI %s. Resetting content-length to match ' |
| 'bytes read.', bytes_read, content_length, request_uri) |
| response.msg['content-length'] = str(bytes_read) |
| response = httplib2.Response(response) |
| else: |
| # We fall back to the current httplib2 behavior if we're |
| # not processing bytes (eg it's a redirect). |
| content = response.read() |
| response = httplib2.Response(response) |
| # pylint: disable=protected-access |
| content = httplib2._decompressContent(response, content) |
| return (response, content) |
