Handle uploading chunked media by stream.
Reviewed in http://codereview.appspot.com/6486046/
diff --git a/apiclient/http.py b/apiclient/http.py
index 0c7a278..569815a 100644
--- a/apiclient/http.py
+++ b/apiclient/http.py
@@ -545,6 +545,46 @@
raise HttpError(resp, content, uri=self._uri)
+class _StreamSlice(object):
+ """Truncated stream.
+
+ Takes a stream and presents a stream that is a slice of the original stream.
+ This is used when uploading media in chunks. In later versions of Python a
+ stream can be passed to httplib in place of the string of data to send. The
+ problem is that httplib just blindly reads to the end of the stream. This
+ wrapper presents a virtual stream that only reads to the end of the chunk.
+ """
+
+ def __init__(self, stream, begin, chunksize):
+ """Constructor.
+
+ Args:
+ stream: (io.Base, file object), the stream to wrap.
+ begin: int, the seek position the chunk begins at.
+ chunksize: int, the size of the chunk.
+ """
+ self._stream = stream
+ self._begin = begin
+ self._chunksize = chunksize
+ self._stream.seek(begin)
+
+ def read(self, n=-1):
+ """Read n bytes.
+
+ Args:
+ n, int, the number of bytes to read.
+
+ Returns:
+ A string of length 'n', or less if EOF is reached.
+ """
+ # The data left available to read sits in [cur, end)
+ cur = self._stream.tell()
+ end = self._begin + self._chunksize
+ if n == -1 or cur + n > end:
+ n = end - cur
+ return self._stream.read(n)
+
+
class HttpRequest(object):
"""Encapsulates a single HTTP request."""
@@ -711,11 +751,13 @@
# conditions then use it as the body argument.
if self.resumable.has_stream() and sys.version_info[1] >= 6:
data = self.resumable.stream()
- data.seek(self.resumable_progress)
if self.resumable.chunksize() == -1:
- # Upload everything in a single chunk.
- chunk_end = self.resumable.size() - 1
+ data.seek(self.resumable_progress)
+ chunk_end = self.resumable.size() - self.resumable_progress - 1
else:
+ # Doing chunking with a stream, so wrap a slice of the stream.
+ data = _StreamSlice(data, self.resumable_progress,
+ self.resumable.chunksize())
chunk_end = min(
self.resumable_progress + self.resumable.chunksize() - 1,
self.resumable.size() - 1)
@@ -731,7 +773,10 @@
headers = {
'Content-Range': 'bytes %d-%d/%s' % (
- self.resumable_progress, chunk_end, size)
+ self.resumable_progress, chunk_end, size),
+ # Must set the content-length header here because httplib can't
+ # calculate the size when working with _StreamSlice.
+ 'Content-Length': str(chunk_end - self.resumable_progress + 1)
}
try:
resp, content = http.request(self.resumable_uri, 'PUT',
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
index ffa30c4..e9d5553 100644
--- a/tests/test_discovery.py
+++ b/tests/test_discovery.py
@@ -590,6 +590,37 @@
except ImportError:
pass
+
+ def test_media_io_base_stream_chunksize_resume(self):
+ self.http = HttpMock(datafile('zoo.json'), {'status': '200'})
+ zoo = build('zoo', 'v1', http=self.http)
+
+ try:
+ import io
+
+ # Set up a seekable stream and try to upload in chunks.
+ fd = io.BytesIO('0123456789')
+ media_upload = MediaIoBaseUpload(
+ fd=fd, mimetype='text/plain', chunksize=5, resumable=True)
+
+ request = zoo.animals().insert(media_body=media_upload, body=None)
+
+ # The single chunk fails, pull the content sent out of the exception.
+ http = HttpMockSequence([
+ ({'status': '200',
+ 'location': 'http://upload.example.com'}, ''),
+ ({'status': '400'}, 'echo_request_body'),
+ ])
+
+ try:
+ body = request.execute(http=http)
+ except HttpError, e:
+ self.assertEqual('01234', e.content)
+
+ except ImportError:
+ pass
+
+
def test_resumable_media_handle_uploads_of_unknown_size(self):
http = HttpMockSequence([
({'status': '200',
@@ -621,8 +652,10 @@
request = zoo.animals().insert(media_body=upload, body=None)
status, body = request.next_chunk(http=http)
- self.assertEqual(body, {'Content-Range': 'bytes 0-9/*'},
- 'Should be 10 out of * bytes.')
+ self.assertEqual(body, {
+ 'Content-Range': 'bytes 0-9/*',
+ 'Content-Length': '10',
+ })
def test_resumable_media_no_streaming_on_unsupported_platforms(self):
self.http = HttpMock(datafile('zoo.json'), {'status': '200'})
@@ -665,8 +698,10 @@
# This should not raise an exception because stream() shouldn't be called.
status, body = request.next_chunk(http=http)
- self.assertEqual(body, {'Content-Range': 'bytes 0-9/*'},
- 'Should be 10 out of * bytes.')
+ self.assertEqual(body, {
+ 'Content-Range': 'bytes 0-9/*',
+ 'Content-Length': '10'
+ })
sys.version_info = (2, 6, 5, 'final', 0)
@@ -701,7 +736,10 @@
request = zoo.animals().insert(media_body=upload, body=None)
status, body = request.next_chunk(http=http)
- self.assertEqual(body, {'Content-Range': 'bytes 0-13/14'})
+ self.assertEqual(body, {
+ 'Content-Range': 'bytes 0-13/14',
+ 'Content-Length': '14',
+ })
def test_resumable_media_handle_resume_of_upload_of_unknown_size(self):
http = HttpMockSequence([
diff --git a/tests/test_http.py b/tests/test_http.py
index 37d8e9c..3609f40 100644
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -36,13 +36,14 @@
from apiclient.http import HttpMock
from apiclient.http import HttpMockSequence
from apiclient.http import HttpRequest
-from apiclient.http import MediaFileUpload
-from apiclient.http import MediaUpload
-from apiclient.http import MediaInMemoryUpload
-from apiclient.http import MediaIoBaseUpload
-from apiclient.http import MediaIoBaseDownload
-from apiclient.http import set_user_agent
from apiclient.http import MAX_URI_LENGTH
+from apiclient.http import MediaFileUpload
+from apiclient.http import MediaInMemoryUpload
+from apiclient.http import MediaIoBaseDownload
+from apiclient.http import MediaIoBaseUpload
+from apiclient.http import MediaUpload
+from apiclient.http import _StreamSlice
+from apiclient.http import set_user_agent
from apiclient.model import JsonModel
from oauth2client.client import Credentials
@@ -745,6 +746,7 @@
'"Access Not Configured">')
self.assertEqual(expected, str(callbacks.exceptions['2']))
+
class TestRequestUriTooLong(unittest.TestCase):
def test_turn_get_into_post(self):
@@ -784,5 +786,26 @@
self.assertEqual(
'application/x-www-form-urlencoded', response['content-type'])
+
+class TestStreamSlice(unittest.TestCase):
+ """Test _StreamSlice."""
+
+ def setUp(self):
+ self.stream = StringIO.StringIO('0123456789')
+
+ def test_read(self):
+ s = _StreamSlice(self.stream, 0, 4)
+ self.assertEqual('', s.read(0))
+ self.assertEqual('0', s.read(1))
+ self.assertEqual('123', s.read())
+
+ def test_read_too_much(self):
+ s = _StreamSlice(self.stream, 1, 4)
+ self.assertEqual('1234', s.read(6))
+
+ def test_read_all(self):
+ s = _StreamSlice(self.stream, 2, 1)
+ self.assertEqual('2', s.read(-1))
+
if __name__ == '__main__':
unittest.main()