Add chunk_size parameter to data downloading methods (export, get_archive, save)1352-data_stream_control

Signed-off-by: Joffrey F <joffrey@docker.com>

8 files changed, 53 insertions, 18 deletions

diff --git a/docker/api/client.py b/docker/api/client.py
index e69d143..bddab61 100644
--- a/docker/api/client.py
+++ b/docker/api/client.py
@@ -350,10 +350,10 @@ class APIClient(
                 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):
diff --git a/docker/api/container.py b/docker/api/container.py
index 962d8cb..e986cf2 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -3,6 +3,7 @@ from datetime import datetime
 
 from .. import errors
 from .. import utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 from ..types import (
     ContainerConfig, EndpointConfig, HostConfig, NetworkingConfig
 )
@@ -643,12 +644,15 @@ class ContainerApiMixin(object):
         )
 
     @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
@@ -660,10 +664,10 @@ class ContainerApiMixin(object):
         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.
@@ -671,6 +675,9 @@ class ContainerApiMixin(object):
         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
@@ -688,7 +695,7 @@ class ContainerApiMixin(object):
         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
         )
 
diff --git a/docker/api/image.py b/docker/api/image.py
index fa832a3..3ebca32 100644
--- a/docker/api/image.py
+++ b/docker/api/image.py
@@ -4,6 +4,7 @@ import os
 import six
 
 from .. import auth, errors, utils
+from ..constants import DEFAULT_DATA_CHUNK_SIZE
 
 log = logging.getLogger(__name__)
 
@@ -11,12 +12,15 @@ 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 @@ class ImageApiMixin(object):
             >>> 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):
diff --git a/docker/constants.py b/docker/constants.py
index 9ab3673..7565a76 100644
--- a/docker/constants.py
+++ b/docker/constants.py
@@ -17,3 +17,4 @@ IS_WINDOWS_PLATFORM = (sys.platform == 'win32')
 
 DEFAULT_USER_AGENT = "docker-sdk-python/{0}".format(version)
 DEFAULT_NUM_POOLS = 25
+DEFAULT_DATA_CHUNK_SIZE = 1024 * 2048
diff --git a/docker/models/containers.py b/docker/models/containers.py
index 107a020..b6e34dd 100644
--- a/docker/models/containers.py
+++ b/docker/models/containers.py
@@ -3,6 +3,7 @@ import ntpath
 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,10 +182,15 @@ class Container(Model):
             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
 
@@ -192,15 +198,18 @@ class Container(Model):
             :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 @@ class Container(Model):
             :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):
         """
diff --git a/docker/models/images.py b/docker/models/images.py
index 0f3c71a..d604f7c 100644
--- a/docker/models/images.py
+++ b/docker/models/images.py
@@ -4,6 +4,7 @@ import re
 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 @@ class Image(Model):
         """
         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 @@ class Image(Model):
             >>>   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):
         """
diff --git a/tests/unit/models_containers_test.py b/tests/unit/models_containers_test.py
index f79f5d5..2b0b499 100644
--- a/tests/unit/models_containers_test.py
+++ b/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 @@ class ContainerTest(unittest.TestCase):
         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()
diff --git a/tests/unit/models_images_test.py b/tests/unit/models_images_test.py
index dacd72b..6783279 100644
--- a/tests/unit/models_images_test.py
+++ b/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 @@ class ImageTest(unittest.TestCase):
         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()