From b0f909c0720f22005b691db5bbba311152a410f3 Mon Sep 17 00:00:00 2001 From: James Saryerwinnie Date: Thu, 3 Jan 2013 14:07:16 -0800 Subject: [PATCH] Use a custom sender for glacier archive upload This sender will reset a filelike object so that retries work properly. Fixes #1189. --- boto/glacier/layer1.py | 18 ++++++++++++++---- boto/glacier/utils.py | 13 +++++++++++++ tests/unit/glacier/test_layer1.py | 24 ++++++++++++++++++++++-- 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/boto/glacier/layer1.py b/boto/glacier/layer1.py index b0090628c4..e5a3963034 100644 --- a/boto/glacier/layer1.py +++ b/boto/glacier/layer1.py @@ -29,6 +29,7 @@ from boto.connection import AWSAuthConnection from .exceptions import UnexpectedHTTPResponseError from .response import GlacierResponse +from .utils import ResettingFileSender class Layer1(AWSAuthConnection): @@ -66,7 +67,7 @@ def _required_auth_capability(self): def make_request(self, verb, resource, headers=None, data='', ok_responses=(200,), params=None, - response_headers=None): + sender=None, response_headers=None): if headers is None: headers = {} headers['x-amz-glacier-version'] = self.Version @@ -74,6 +75,7 @@ def make_request(self, verb, resource, headers=None, response = AWSAuthConnection.make_request(self, verb, uri, params=params, headers=headers, + sender=sender, data=data) if response.status in ok_responses: return GlacierResponse(response, response_headers) @@ -420,7 +422,7 @@ def upload_archive(self, vault_name, archive, uri = 'vaults/%s/archives' % vault_name try: content_length = str(len(archive)) - except TypeError: + except (TypeError, AttributeError): # If a file like object is provided, try to retrieve # the file size via fstat. content_length = str(os.fstat(archive.fileno()).st_size) @@ -429,9 +431,17 @@ def upload_archive(self, vault_name, archive, 'Content-Length': content_length} if description: headers['x-amz-archive-description'] = description + if self._is_file_like(archive): + sender = ResettingFileSender(archive) + else: + sender = None return self.make_request('POST', uri, headers=headers, - data=archive, ok_responses=(201,), - response_headers=response_headers) + sender=sender, + data=archive, ok_responses=(201,), + response_headers=response_headers) + + def _is_file_like(self, archive): + return hasattr(archive, 'seek') and hasattr(archive, 'tell') def delete_archive(self, vault_name, archive_id): """ diff --git a/boto/glacier/utils.py b/boto/glacier/utils.py index e6a9baf65a..5f26f6f38f 100644 --- a/boto/glacier/utils.py +++ b/boto/glacier/utils.py @@ -112,3 +112,16 @@ def compute_hashes_from_fileobj(fileobj, chunk_size=1024 * 1024): def bytes_to_hex(str_as_bytes): return ''.join(["%02x" % ord(x) for x in str_as_bytes]).strip() + + +class ResettingFileSender(object): + def __init__(self, archive): + self._archive = archive + self._starting_offset = archive.tell() + + def __call__(self, connection, method, path, body, headers): + try: + connection.request(method, path, self._archive, headers) + return connection.getresponse() + finally: + self._archive.seek(self._starting_offset) diff --git a/tests/unit/glacier/test_layer1.py b/tests/unit/glacier/test_layer1.py index 7d7b6fc273..4cabbeae50 100644 --- a/tests/unit/glacier/test_layer1.py +++ b/tests/unit/glacier/test_layer1.py @@ -1,7 +1,9 @@ -from tests.unit import AWSMockServiceTestCase -from boto.glacier.layer1 import Layer1 import json import copy +import tempfile + +from tests.unit import AWSMockServiceTestCase +from boto.glacier.layer1 import Layer1 class GlacierLayer1ConnectionBase(AWSMockServiceTestCase): @@ -76,3 +78,21 @@ def test_get_archive_output(self): response = self.service_connection.get_job_output(self.vault_name, 'example-job-id') self.assertEqual(self.job_content, response.read()) + + +class GlacierUploadArchiveResets(GlacierLayer1ConnectionBase): + def test_upload_archive(self): + fake_data = tempfile.NamedTemporaryFile() + fake_data.write('foobarbaz') + # First seek to a non zero offset. + fake_data.seek(2) + self.set_http_response(status_code=201) + # Simulate reading the request body when we send the request. + self.service_connection.connection.request.side_effect = \ + lambda *args: fake_data.read() + self.service_connection.upload_archive('vault_name', fake_data, 'linear_hash', + 'tree_hash') + # Verify that we seek back to the original offset after making + # a request. This ensures that if we need to resend the request we're + # back at the correct location within the file. + self.assertEqual(fake_data.tell(), 2)