Client.add_image() accepts image data as iterable.

Fixes bug 773562

Allow the image data to be passed to glance.Client.add_add_image()
via an iterator, in addition to the previously accepted images types
(file-like or string).

Change-Id: I1f90b9875f5610d478a5ed123fbf431501f1c2b6

glance/common/client.py

@@ -19,6 +19,7 @@
 # http://code.activestate.com/recipes/
 #   577548-https-httplib-client-connection-with-certificate-v/
 
+import collections
 import functools
 import httplib
 import logging
@@ -36,6 +37,10 @@
 from glance.common import exception
 
 
+# common chunk size for get and put
+CHUNKSIZE = 65536
+
+
 def handle_unauthorized(func):
     """
     Wrap a function to re-authenticate and retry.
@@ -77,21 +82,20 @@ class ImageBodyIterator(object):
     tuple from `glance.client.Client.get_image`
     """
 
-    CHUNKSIZE = 65536
-
-    def __init__(self, response):
+    def __init__(self, source):
         """
-        Constructs the object from an HTTPResponse object
+        Constructs the object from a readable image source
+        (such as an HTTPResponse or file-like object)
         """
-        self.response = response
+        self.source = source
 
     def __iter__(self):
         """
         Exposes an iterator over the chunks of data in the
         image file.
         """
         while True:
-            chunk = self.response.read(ImageBodyIterator.CHUNKSIZE)
+            chunk = self.source.read(CHUNKSIZE)
             if chunk:
                 yield chunk
             else:
@@ -144,7 +148,6 @@ class BaseClient(object):
 
     """A base client class"""
 
-    CHUNKSIZE = 65536
     DEFAULT_PORT = 80
     DEFAULT_DOC_ROOT = None
 
@@ -355,15 +358,16 @@ def _do_request(self, method, url, body, headers):
 
         :param method: HTTP method ("GET", "POST", "PUT", etc...)
         :param url: urlparse.ParsedResult object with URL information
-        :param body: string of data to send, or None (default)
+        :param body: data to send (as string, filelike or iterable),
+                     or None (default)
         :param headers: mapping of key/value pairs to add as headers
 
         :note
 
         If the body param has a read attribute, and method is either
         POST or PUT, this method will automatically conduct a chunked-transfer
-        encoding and use the body as a file object, transferring chunks
-        of data using the connection's send() method. This allows large
+        encoding and use the body as a file object or iterable, transferring
+        chunks of data using the connection's send() method. This allows large
         objects to be transferred efficiently without buffering the entire
         body in memory.
         """
@@ -381,10 +385,26 @@ def _do_request(self, method, url, body, headers):
 
             c = connection_type(url.hostname, url.port, **self.connect_kwargs)
 
+            def _pushing(method):
+                return method.lower() in ('post', 'put')
+
+            def _simple(body):
+                return body is None or isinstance(body, basestring)
+
+            def _filelike(body):
+                return hasattr(body, 'read')
+
+            def _iterable(body):
+                return isinstance(body, collections.Iterable)
+
             # Do a simple request or a chunked request, depending
-            # on whether the body param is a file-like object and
+            # on whether the body param is file-like or iterable and
             # the method is PUT or POST
-            if hasattr(body, 'read') and method.lower() in ('post', 'put'):
+            #
+            if not _pushing(method) or _simple(body):
+                # Simple request...
+                c.request(method, path, body, headers)
+            elif _filelike(body) or _iterable(body):
                 # Chunk it, baby...
                 c.putrequest(method, path)
 
@@ -393,14 +413,14 @@ def _do_request(self, method, url, body, headers):
                 c.putheader('Transfer-Encoding', 'chunked')
                 c.endheaders()
 
-                chunk = body.read(self.CHUNKSIZE)
-                while chunk:
+                iter = body if _iterable(body) else ImageBodyIterator(body)
+
+                for chunk in iter:
                     c.send('%x\r\n%s\r\n' % (len(chunk), chunk))
-                    chunk = body.read(self.CHUNKSIZE)
                 c.send('0\r\n\r\n')
             else:
-                # Simple request...
-                c.request(method, path, body, headers)
+                raise TypeError('Unsupported image type: %s' % body.__class__)
+
             res = c.getresponse()
             status_code = self.get_status_code(res)
             if status_code in self.OK_RESPONSE_CODES:

glance/tests/stubs.py

@@ -155,7 +155,7 @@ def fake_get_connection_type(client):
             return FakeRegistryConnection
 
     def fake_image_iter(self):
-        for i in self.response.app_iter:
+        for i in self.source.app_iter:
             yield i
 
     stubs.Set(glance.common.client.BaseClient, 'get_connection_type',

glance/tests/unit/test_clients.py

@@ -1813,6 +1813,60 @@ def test_add_image_with_image_data_as_file(self):
         for k, v in fixture.items():
             self.assertEquals(v, new_meta[k])
 
+    def _add_image_as_iterable(self):
+        fixture = {'name': 'fake public image',
+                   'is_public': True,
+                   'disk_format': 'vhd',
+                   'container_format': 'ovf',
+                   'size': 10 * 65536,
+                   'properties': {'distro': 'Ubuntu 10.04 LTS'},
+                  }
+
+        class Zeros:
+            def __init__(self, chunks):
+                self.chunks = chunks
+                self.zeros = open('/dev/zero', 'rb')
+
+            def __iter__(self):
+                while self.chunks > 0:
+                    self.chunks -= 1
+                    chunk = self.zeros.read(65536)
+                    yield chunk
+
+        new_image = self.client.add_image(fixture, Zeros(10))
+        new_image_id = new_image['id']
+
+        new_meta, new_image_chunks = self.client.get_image(new_image_id)
+
+        return (fixture, new_meta, new_image_chunks)
+
+    def _verify_image_iterable(self, fixture, meta, chunks):
+        image_data_len = 0
+        for image_chunk in chunks:
+            image_data_len += len(image_chunk)
+        self.assertEquals(10 * 65536, image_data_len)
+
+        for k, v in fixture.items():
+            self.assertEquals(v, meta[k])
+
+    def test_add_image_with_image_data_as_iterable(self):
+        """Tests we can add image by passing image data as an iterable"""
+        fixture, new_meta, new_chunks = self._add_image_as_iterable()
+
+        self._verify_image_iterable(fixture, new_meta, new_chunks)
+
+    def test_roundtrip_image_with_image_data_as_iterable(self):
+        """Tests we can roundtrip image as an iterable"""
+        fixture, new_meta, new_chunks = self._add_image_as_iterable()
+
+        # duplicate directly from iterable returned from get
+        dup_image = self.client.add_image(fixture, new_chunks)
+        dup_image_id = dup_image['id']
+
+        roundtrip_meta, roundtrip_chunks = self.client.get_image(dup_image_id)
+
+        self._verify_image_iterable(fixture, roundtrip_meta, roundtrip_chunks)
+
     def test_add_image_with_image_data_as_string_and_no_size(self):
         """Tests add image by passing image data as string w/ no size attr"""
         fixture = {'name': 'fake public image',