Allow reading from object body on download

Currently, get_object returns a generator. This allows access to the
object's data in multiples of 'resp_chunk_size'.

This patch adds a read function to also allow accessing the object data
using read(size).

This allows, for example, the consumer of an object (where no byte range
has been specified) to read up to certain boundaries while streaming to
a new Large Object with segments of a specified size.

Reading and chunking can be safely mixed.

Related-Bug: 1367925
Co-Authored-By: Clay Gerrard <clay.gerrard@gmail.com>
Change-Id: I1cfb67f45afc7015fd896f1a89bebae048871769

swiftclient/client.py

@@ -138,6 +138,37 @@ def encode_meta_headers(headers):
     return ret
 
 
+class _ObjectBody(object):
+    """
+    Readable and iterable object body response wrapper.
+    """
+
+    def __init__(self, resp, chunk_size):
+        """
+        Wrap the underlying response
+
+        :param resp: the response to wrap
+        :param chunk_size: number of bytes to return each iteration/next call
+        """
+        self.resp = resp
+        self.chunk_size = chunk_size
+
+    def read(self, length=None):
+        return self.resp.read(length)
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        buf = self.resp.read(self.chunk_size)
+        if not buf:
+            raise StopIteration()
+        return buf
+
+    def __next__(self):
+        return self.next()
+
+
 class HTTPConnection(object):
     def __init__(self, url, proxy=None, cacert=None, insecure=False,
                  ssl_compression=False, default_user_agent=None):
@@ -874,13 +905,7 @@ def get_object(url, token, container, name, http_conn=None,
                               http_reason=resp.reason,
                               http_response_content=body)
     if resp_chunk_size:
-
-        def _object_body():
-            buf = resp.read(resp_chunk_size)
-            while buf:
-                yield buf
-                buf = resp.read(resp_chunk_size)
-        object_body = _object_body()
+        object_body = _ObjectBody(resp, resp_chunk_size)
     else:
         object_body = resp.read()
     http_log(('%s%s' % (url.replace(parsed.path, ''), path), method,),

tests/functional/test_swiftclient.py

@@ -16,7 +16,6 @@
 import os
 import testtools
 import time
-import types
 from io import BytesIO
 
 from six.moves import configparser
@@ -256,8 +255,24 @@ class TestFunctional(testtools.TestCase):
         hdrs, body = self.conn.get_object(
             self.containername, self.objectname,
             resp_chunk_size=10)
-        self.assertTrue(isinstance(body, types.GeneratorType))
-        self.assertEqual(self.test_data, b''.join(body))
+        downloaded_contents = b''
+        while True:
+            try:
+                chunk = next(body)
+            except StopIteration:
+                break
+            downloaded_contents += chunk
+        self.assertEqual(self.test_data, downloaded_contents)
+
+        # Download in chunks, should also work with read
+        hdrs, body = self.conn.get_object(
+            self.containername, self.objectname,
+            resp_chunk_size=10)
+        num_bytes = 5
+        downloaded_contents = body.read(num_bytes)
+        self.assertEqual(num_bytes, len(downloaded_contents))
+        downloaded_contents += body.read()
+        self.assertEqual(self.test_data, downloaded_contents)
 
     def test_post_account(self):
         self.conn.post_account({'x-account-meta-data': 'Something'})

tests/unit/test_swiftclient.py

@@ -655,6 +655,41 @@ class TestGetObject(MockHttpTest):
             }),
         ])
 
+    def test_chunk_size_read_method(self):
+        conn = c.Connection('http://auth.url/', 'some_user', 'some_key')
+        with mock.patch('swiftclient.client.get_auth_1_0') as mock_get_auth:
+            mock_get_auth.return_value = ('http://auth.url/', 'tToken')
+            c.http_connection = self.fake_http_connection(200, body='abcde')
+            __, resp = conn.get_object('asdf', 'asdf', resp_chunk_size=3)
+            self.assertTrue(hasattr(resp, 'read'))
+            self.assertEquals(resp.read(3), 'abc')
+            self.assertEquals(resp.read(None), 'de')
+            self.assertEquals(resp.read(), '')
+
+    def test_chunk_size_iter(self):
+        conn = c.Connection('http://auth.url/', 'some_user', 'some_key')
+        with mock.patch('swiftclient.client.get_auth_1_0') as mock_get_auth:
+            mock_get_auth.return_value = ('http://auth.url/', 'tToken')
+            c.http_connection = self.fake_http_connection(200, body='abcde')
+            __, resp = conn.get_object('asdf', 'asdf', resp_chunk_size=3)
+            self.assertTrue(hasattr(resp, 'next'))
+            self.assertEquals(next(resp), 'abc')
+            self.assertEquals(next(resp), 'de')
+            self.assertRaises(StopIteration, next, resp)
+
+    def test_chunk_size_read_and_iter(self):
+        conn = c.Connection('http://auth.url/', 'some_user', 'some_key')
+        with mock.patch('swiftclient.client.get_auth_1_0') as mock_get_auth:
+            mock_get_auth.return_value = ('http://auth.url/', 'tToken')
+            c.http_connection = self.fake_http_connection(200, body='abcdef')
+            __, resp = conn.get_object('asdf', 'asdf', resp_chunk_size=2)
+            self.assertTrue(hasattr(resp, 'read'))
+            self.assertEquals(resp.read(3), 'abc')
+            self.assertEquals(next(resp), 'de')
+            self.assertEquals(resp.read(), 'f')
+            self.assertRaises(StopIteration, next, resp)
+            self.assertEquals(resp.read(), '')
+
 
 class TestHeadObject(MockHttpTest):
 

tests/unit/utils.py

@@ -154,7 +154,10 @@ def fake_http_connect(*code_iter, **kwargs):
                     sleep(0.1)
                     return ' '
             rv = self.body[:amt]
-            self.body = self.body[amt:]
+            if amt is not None:
+                self.body = self.body[amt:]
+            else:
+                self.body = ''
             return rv
 
         def send(self, amt=None):