Merge pull request #1906 from docker/1352-data_stream_control

Add chunk_size parameter to data downloading methods

Author: shin-

docker/api/client.py

@@ -350,10 +350,10 @@ def _multiplexed_response_stream_helper(self, response):
                 break
             yield data
 
-    def _stream_raw_result(self, response):
-        ''' Stream result for TTY-enabled container '''
+    def _stream_raw_result(self, response, chunk_size=1, decode=True):
+        ''' Stream result for TTY-enabled container and raw binary data'''
         self._raise_for_status(response)
-        for out in response.iter_content(chunk_size=1, decode_unicode=True):
+        for out in response.iter_content(chunk_size, decode):
             yield out
 
     def _read_from_socket(self, response, stream, tty=False):

docker/api/container.py

@@ -3,6 +3,7 @@
 
 from .. import errors
 from .. import utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 from ..types import (
     ContainerConfig, EndpointConfig, HostConfig, NetworkingConfig
 )
@@ -645,12 +646,15 @@ def diff(self, container):
         )
 
     @utils.check_resource('container')
-    def export(self, container):
+    def export(self, container, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Export the contents of a filesystem as a tar archive.
 
         Args:
             container (str): The container to export
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (generator): The archived filesystem data stream
@@ -662,17 +666,20 @@ def export(self, container):
         res = self._get(
             self._url("/containers/{0}/export", container), stream=True
         )
-        return self._stream_raw_result(res)
+        return self._stream_raw_result(res, chunk_size, False)
 
     @utils.check_resource('container')
-    def get_archive(self, container, path):
+    def get_archive(self, container, path, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Retrieve a file or folder from a container in the form of a tar
         archive.
 
         Args:
             container (str): The container where the file is located
             path (str): Path to the file or folder to retrieve
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (tuple): First element is a raw tar data stream. Second element is
@@ -690,7 +697,7 @@ def get_archive(self, container, path):
         self._raise_for_status(res)
         encoded_stat = res.headers.get('x-docker-container-path-stat')
         return (
-            self._stream_raw_result(res),
+            self._stream_raw_result(res, chunk_size, False),
             utils.decode_json_header(encoded_stat) if encoded_stat else None
         )
 

docker/api/image.py

@@ -4,19 +4,23 @@
 import six
 
 from .. import auth, errors, utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 
 log = logging.getLogger(__name__)
 
 
 class ImageApiMixin(object):
 
     @utils.check_resource('image')
-    def get_image(self, image):
+    def get_image(self, image, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Get a tarball of an image. Similar to the ``docker save`` command.
 
         Args:
             image (str): Image name to get
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (generator): A stream of raw archive data.
@@ -34,7 +38,7 @@ def get_image(self, image):
             >>> f.close()
         """
         res = self._get(self._url("/images/{0}/get", image), stream=True)
-        return self._stream_raw_result(res)
+        return self._stream_raw_result(res, chunk_size, False)
 
     @utils.check_resource('image')
     def history(self, image):

docker/constants.py

@@ -17,3 +17,4 @@
 
 DEFAULT_USER_AGENT = "docker-sdk-python/{0}".format(version)
 DEFAULT_NUM_POOLS = 25
+DEFAULT_DATA_CHUNK_SIZE = 1024 * 2048

docker/models/containers.py

@@ -3,6 +3,7 @@
 from collections import namedtuple
 
 from ..api import APIClient
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 from ..errors import (ContainerError, ImageNotFound,
                       create_unexpected_kwargs_error)
 from ..types import HostConfig
@@ -181,26 +182,34 @@ def exec_run(self, cmd, stdout=True, stderr=True, stdin=False, tty=False,
             exec_output
         )
 
-    def export(self):
+    def export(self, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Export the contents of the container's filesystem as a tar archive.
 
+        Args:
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
+
         Returns:
             (str): The filesystem tar archive
 
         Raises:
             :py:class:`docker.errors.APIError`
                 If the server returns an error.
         """
-        return self.client.api.export(self.id)
+        return self.client.api.export(self.id, chunk_size)
 
-    def get_archive(self, path):
+    def get_archive(self, path, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Retrieve a file or folder from the container in the form of a tar
         archive.
 
         Args:
             path (str): Path to the file or folder to retrieve
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
 
         Returns:
             (tuple): First element is a raw tar data stream. Second element is
@@ -210,7 +219,7 @@ def get_archive(self, path):
             :py:class:`docker.errors.APIError`
                 If the server returns an error.
         """
-        return self.client.api.get_archive(self.id, path)
+        return self.client.api.get_archive(self.id, path, chunk_size)
 
     def kill(self, signal=None):
         """

docker/models/images.py

@@ -4,6 +4,7 @@
 import six
 
 from ..api import APIClient
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 from ..errors import BuildError, ImageLoadError
 from ..utils import parse_repository_tag
 from ..utils.json_stream import json_stream
@@ -58,10 +59,15 @@ def history(self):
         """
         return self.client.api.history(self.id)
 
-    def save(self):
+    def save(self, chunk_size=DEFAULT_DATA_CHUNK_SIZE):
         """
         Get a tarball of an image. Similar to the ``docker save`` command.
 
+        Args:
+            chunk_size (int): The number of bytes returned by each iteration
+                of the generator. If ``None``, data will be streamed as it is
+                received. Default: 2 MB
+
         Returns:
             (generator): A stream of raw archive data.
 
@@ -77,7 +83,7 @@ def save(self):
             >>>   f.write(chunk)
             >>> f.close()
         """
-        return self.client.api.get_image(self.id)
+        return self.client.api.get_image(self.id, chunk_size)
 
     def tag(self, repository, tag=None, **kwargs):
         """

tests/unit/models_containers_test.py

@@ -1,4 +1,5 @@
 import docker
+from docker.constants import DEFAULT_DATA_CHUNK_SIZE
 from docker.models.containers import Container, _create_container_args
 from docker.models.images import Image
 import unittest
@@ -422,13 +423,17 @@ def test_export(self):
         client = make_fake_client()
         container = client.containers.get(FAKE_CONTAINER_ID)
         container.export()
-        client.api.export.assert_called_with(FAKE_CONTAINER_ID)
+        client.api.export.assert_called_with(
+            FAKE_CONTAINER_ID, DEFAULT_DATA_CHUNK_SIZE
+        )
 
     def test_get_archive(self):
         client = make_fake_client()
         container = client.containers.get(FAKE_CONTAINER_ID)
         container.get_archive('foo')
-        client.api.get_archive.assert_called_with(FAKE_CONTAINER_ID, 'foo')
+        client.api.get_archive.assert_called_with(
+            FAKE_CONTAINER_ID, 'foo', DEFAULT_DATA_CHUNK_SIZE
+        )
 
     def test_image(self):
         client = make_fake_client()

tests/unit/models_images_test.py

@@ -1,3 +1,4 @@
+from docker.constants import DEFAULT_DATA_CHUNK_SIZE
 from docker.models.images import Image
 import unittest
 
@@ -116,7 +117,9 @@ def test_save(self):
         client = make_fake_client()
         image = client.images.get(FAKE_IMAGE_ID)
         image.save()
-        client.api.get_image.assert_called_with(FAKE_IMAGE_ID)
+        client.api.get_image.assert_called_with(
+            FAKE_IMAGE_ID, DEFAULT_DATA_CHUNK_SIZE
+        )
 
     def test_tag(self):
         client = make_fake_client()