Correct such logic in store.get() when chunk_size param provided

The change corrected the logic of image content reading function in
get() for some store drivers. if the chunk_size param provided from
outside, it means client want to read partial image content from
storage, current implementation is incomplete and has some problem.

The change didn't try to make drivers fully support random access from
backend with chunk_size (and offset) param, that stuff will be
implementated by a dedicated patch, this one only to make default logic
be correct.

Change-Id: I76fd3746a2c449e257ea2a19c2dc6118be207eec
Signed-off-by: Zhi Yan Liu <zhiyanl@cn.ibm.com>

9 files changed, 51 insertions, 16 deletions

diff --git a/glance_store/_drivers/filesystem.py b/glance_store/_drivers/filesystem.py
index 5460102..ae2b29e 100644
--- a/glance_store/_drivers/filesystem.py
+++ b/glance_store/_drivers/filesystem.py
@@ -116,10 +116,12 @@ class ChunkedFile(object):
     something that can iterate over a large file
     """
 
-    def __init__(self, filepath, offset=0, chunk_size=None, partial=False):
-        self.partial = partial
+    def __init__(self, filepath, offset=0, chunk_size=4096,
+                 partial_length=None):
         self.filepath = filepath
         self.chunk_size = chunk_size
+        self.partial_length = partial_length
+        self.partial = self.partial_length is not None
         self.fp = open(self.filepath, 'rb')
         if offset:
             self.fp.seek(offset)
@@ -129,12 +131,19 @@ class ChunkedFile(object):
         try:
             if self.fp:
                 while True:
-                    chunk = self.fp.read(self.chunk_size)
+                    if self.partial:
+                        size = min(self.chunk_size, self.partial_length)
+                    else:
+                        size = self.chunk_size
+
+                    chunk = self.fp.read(size)
                     if chunk:
                         yield chunk
 
                         if self.partial:
-                            break
+                            self.partial_length -= len(chunk)
+                            if self.partial_length <= 0:
+                                break
                     else:
                         break
         finally:
@@ -451,8 +460,8 @@ class Store(glance_store.driver.Store):
         LOG.debug(msg)
         return (ChunkedFile(filepath,
                             offset=offset,
-                            chunk_size=chunk_size or self.READ_CHUNKSIZE,
-                            partial=chunk_size is not None),
+                            chunk_size=self.READ_CHUNKSIZE,
+                            partial_length=chunk_size),
                 chunk_size or filesize)
 
     def get_size(self, location, context=None):
diff --git a/glance_store/_drivers/http.py b/glance_store/_drivers/http.py
index eef578d..a2073de 100644
--- a/glance_store/_drivers/http.py
+++ b/glance_store/_drivers/http.py
@@ -133,8 +133,7 @@ class Store(glance_store.driver.Store):
             LOG.error(reason)
             raise exceptions.RemoteServiceUnavailable()
 
-        cs = chunk_size or self.READ_CHUNKSIZE
-        iterator = http_response_iterator(conn, resp, cs)
+        iterator = http_response_iterator(conn, resp, self.READ_CHUNKSIZE)
 
         class ResponseIndexable(glance_store.Indexable):
             def another(self):
diff --git a/glance_store/_drivers/rbd.py b/glance_store/_drivers/rbd.py
index 09d2468..767a4ff 100644
--- a/glance_store/_drivers/rbd.py
+++ b/glance_store/_drivers/rbd.py
@@ -220,9 +220,8 @@ class Store(driver.Store):
         :raises `glance_store.exceptions.NotFound` if image does not exist
         """
         loc = location.store_location
-        return (ImageIterator(loc.pool, loc.image, loc.snapshot,
-                              self, chunk_size=chunk_size),
-                chunk_size or self.get_size(location))
+        return (ImageIterator(loc.pool, loc.image, loc.snapshot, self),
+                self.get_size(location))
 
     def get_size(self, location, context=None):
         """
diff --git a/glance_store/_drivers/s3.py b/glance_store/_drivers/s3.py
index 2fcb550..1f82813 100644
--- a/glance_store/_drivers/s3.py
+++ b/glance_store/_drivers/s3.py
@@ -378,7 +378,7 @@ class Store(glance_store.driver.Store):
         :raises `glance_store.exceptions.NotFound` if image does not exist
         """
         key = self._retrieve_key(location)
-        cs = chunk_size or self.READ_CHUNKSIZE
+        cs = self.READ_CHUNKSIZE
         key.BufferSize = cs
 
         class ChunkedIndexable(glance_store.Indexable):
diff --git a/glance_store/_drivers/sheepdog.py b/glance_store/_drivers/sheepdog.py
index 597f857..18a5b29 100644
--- a/glance_store/_drivers/sheepdog.py
+++ b/glance_store/_drivers/sheepdog.py
@@ -226,7 +226,7 @@ class Store(glance_store.driver.Store):
 
         loc = location.store_location
         image = SheepdogImage(self.addr, self.port, loc.image,
-                              chunk_size or self.READ_CHUNKSIZE)
+                              self.READ_CHUNKSIZE)
         if not image.exist():
             raise exceptions.NotFound(_("Sheepdog image %s does not exist")
                                       % image.name)
diff --git a/tests/unit/test_http_store.py b/tests/unit/test_http_store.py
index b133e94..5204f61 100644
--- a/tests/unit/test_http_store.py
+++ b/tests/unit/test_http_store.py
@@ -60,6 +60,12 @@ class TestHttpStore(base.StoreBaseTest,
         chunks = [c for c in image_file]
         self.assertEqual(expected_returns, chunks)
 
+    def test_http_partial_get(self):
+        uri = "http://netloc/path/to/file.tar.gz"
+        loc = location.get_location_from_uri(uri, conf=self.conf)
+        self.assertRaises(exceptions.StoreRandomGetNotSupported,
+                          self.store.get, loc, chunk_size=1)
+
     def test_http_get_redirect(self):
         # Add two layers of redirects to the response stack, which will
         # return the default 200 OK with the expected data after resolving
diff --git a/tests/unit/test_rbd_store.py b/tests/unit/test_rbd_store.py
index 917b7ff..dca51e2 100644
--- a/tests/unit/test_rbd_store.py
+++ b/tests/unit/test_rbd_store.py
@@ -173,9 +173,9 @@ class TestStore(base.StoreBaseTest,
             with mock.patch.object(rbd_store.rbd.Image, 'write') as write:
                 ret = self.store.add('fake_image_id', self.data_iter, 0)
 
-        self.assertTrue(resize.called)
-        self.assertTrue(write.called)
-        self.assertEqual(ret[1], self.data_len)
+                self.assertTrue(resize.called)
+                self.assertTrue(write.called)
+                self.assertEqual(ret[1], self.data_len)
 
     @mock.patch.object(MockRBD.Image, '__enter__')
     @mock.patch.object(rbd_store.Store, '_create_image')
@@ -283,6 +283,12 @@ class TestStore(base.StoreBaseTest,
 
             self.called_commands_expected = ['remove']
 
+    def test_get_partial_image(self):
+        loc = Location('test_rbd_store', rbd_store.StoreLocation, self.conf,
+                       store_specs=self.store_specs)
+        self.assertRaises(exceptions.StoreRandomGetNotSupported,
+                          self.store.get, loc, chunk_size=1)
+
     def tearDown(self):
         self.assertEqual(self.called_commands_actual,
                          self.called_commands_expected)
diff --git a/tests/unit/test_s3_store.py b/tests/unit/test_s3_store.py
index 39c0ba5..ed46a0a 100644
--- a/tests/unit/test_s3_store.py
+++ b/tests/unit/test_s3_store.py
@@ -301,6 +301,14 @@ class TestStore(base.StoreBaseTest,
             data += chunk
         self.assertEqual(expected_data, data)
 
+    def test_partial_get(self):
+        """Test a "normal" retrieval of an image in chunks."""
+        loc = location.get_location_from_uri(
+            "s3://user:key@auth_address/glance/%s" % FAKE_UUID,
+            conf=self.conf)
+        self.assertRaises(exceptions.StoreRandomGetNotSupported,
+                          self.store.get, loc, chunk_size=1)
+
     def test_get_calling_format_path(self):
         """Test a "normal" retrieval of an image in chunks."""
         self.config(s3_store_bucket_url_format='path')
diff --git a/tests/unit/test_sheepdog_store.py b/tests/unit/test_sheepdog_store.py
index 4e4d917..41608fb 100644
--- a/tests/unit/test_sheepdog_store.py
+++ b/tests/unit/test_sheepdog_store.py
@@ -19,6 +19,8 @@ import mock
 from oslo_concurrency import processutils
 
 from glance_store._drivers import sheepdog
+from glance_store import exceptions
+from glance_store import location
 from glance_store.tests import base
 from tests.unit import test_store_capabilities
 
@@ -53,3 +55,9 @@ class TestSheepdogStore(base.StoreBaseTest,
             data = StringIO.StringIO('xx')
             self.store.add('fake_image_id', data, 2)
             self.assertEqual(called_commands, ['list -r', 'create', 'write'])
+
+    def test_partial_get(self):
+        loc = location.Location('test_sheepdog_store', sheepdog.StoreLocation,
+                                self.conf, store_specs={'image': 'fake_image'})
+        self.assertRaises(exceptions.StoreRandomGetNotSupported,
+                          self.store.get, loc, chunk_size=1)