diff --git a/docker/client.py b/docker/client.py index 771feb3c..6a187cb5 100644 --- a/docker/client.py +++ b/docker/client.py @@ -30,7 +30,7 @@ if not six.PY3: class APIError(requests.exceptions.HTTPError): def __init__(self, message, response, explanation=None): - super(APIError, self).__init__(message, response=response) + super(APIError, self).__init__(message, response) self.explanation = explanation @@ -81,7 +81,18 @@ class Client(requests.Session): try: response.raise_for_status() except requests.exceptions.HTTPError as e: - raise APIError(e, response=response, explanation=explanation) + raise APIError(e, response, explanation=explanation) + + def _stream_result(self, response): + self._raise_for_status(response) + for line in response.iter_lines(chunk_size=1): + # filter out keep-alive new lines + if line: + yield line + '\n' + + def _stream_result_socket(self, response): + self._raise_for_status(response) + return response.raw._fp.fp._sock def _result(self, response, json=False): self._raise_for_status(response) @@ -192,7 +203,7 @@ class Client(requests.Session): } def build(self, path=None, tag=None, quiet=False, fileobj=None, - nocache=False, rm=False): + nocache=False, rm=False, stream=False): remote = context = headers = None if path is None and fileobj is None: raise Exception("Either path or fileobj needs to be provided.") @@ -214,15 +225,19 @@ class Client(requests.Session): } if context is not None: headers = {'Content-Type': 'application/tar'} - res = self._result(self.post(u, context, params=params, - headers=headers, stream=True)) + response = self.post( + u, context, params=params, headers=headers, stream=stream) if context is not None: context.close() - srch = r'Successfully built ([0-9a-f]+)' - match = re.search(srch, res) - if not match: - return None, res - return match.group(1), res + if stream: + return self._stream_result(response) + else: + output = self._result(response) + srch = r'Successfully built ([0-9a-f]+)' + match = re.search(srch, output) + if not match: + return None, output + return match.group(1), output def commit(self, container, repository=None, tag=None, message=None, author=None, conf=None): @@ -410,7 +425,7 @@ class Client(requests.Session): return f_port - def pull(self, repository, tag=None): + def pull(self, repository, tag=None, stream=False): registry, repo_name = auth.resolve_repository_name(repository) if repo_name.count(":") == 1: repository, tag = repository.rsplit(":", 1) @@ -430,9 +445,30 @@ class Client(requests.Session): if authcfg: headers['X-Registry-Auth'] = auth.encode_header(authcfg) u = self._url("/images/create") - return self._result(self.post(u, params=params, headers=headers)) + response = self.post(u, params=params, headers=headers, stream=stream) - def push(self, repository): + if stream: + return self.stream_helper(response) + else: + return self._result(response) + + def stream_helper(self, response): + socket = self._stream_result_socket(response) + while True: + chunk = socket.recv(4096) + if chunk: + parts = chunk.strip().split('\r\n') + for i in range(len(parts)): + if i % 2 != 0: + yield parts[i] + '\n' + else: + size = int(parts[i], 16) + if size <= 0: + break + else: + break + + def push(self, repository, stream=False): registry, repo_name = auth.resolve_repository_name(repository) u = self._url("/images/{0}/push".format(repository)) headers = {} @@ -444,8 +480,17 @@ class Client(requests.Session): # for this specific registry as we can have an anon push if authcfg: headers['X-Registry-Auth'] = auth.encode_header(authcfg) - return self._result(self._post_json(u, None, headers=headers)) - return self._result(self._post_json(u, authcfg)) + if stream: + return self.stream_helper( + self._post_json(u, None, headers=headers, stream=True)) + else: + return self._result( + self._post_json(u, None, headers=headers, stream=False)) + if stream: + return self.stream_helper( + self._post_json(u, authcfg, stream=True)) + else: + return self._result(self._post_json(u, authcfg, stream=False)) def remove_container(self, container, v=False): if isinstance(container, dict): diff --git a/tests/integration_test.py b/tests/integration_test.py index fd77c930..4db21808 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -459,6 +459,28 @@ class TestPull(BaseTestCase): self.tmp_imgs.append('376968a23351') +class TestPullStream(BaseTestCase): + def runTest(self): + try: + self.client.remove_image('joffrey/test001') + self.client.remove_image('376968a23351') + except docker.APIError: + pass + info = self.client.info() + self.assertIn('Images', info) + img_count = info['Images'] + stream = self.client.pull('joffrey/test001', stream=True) + res = u'' + for chunk in stream: + res += chunk + self.assertEqual(type(res), six.text_type) + self.assertEqual(img_count + 2, self.client.info()['Images']) + img_info = self.client.inspect_image('joffrey/test001') + self.assertIn('id', img_info) + self.tmp_imgs.append('joffrey/test001') + self.tmp_imgs.append('376968a23351') + + class TestCommit(BaseTestCase): def runTest(self): container = self.client.create_container('busybox', ['touch', '/test']) @@ -529,6 +551,23 @@ class TestBuild(BaseTestCase): self.tmp_imgs.append(img) +class TestBuildStream(BaseTestCase): + def runTest(self): + script = io.BytesIO('\n'.join([ + 'FROM busybox', + 'MAINTAINER docker-py', + 'RUN mkdir -p /tmp/test', + 'EXPOSE 8080', + 'ADD https://dl.dropboxusercontent.com/u/20637798/silence.tar.gz' + ' /tmp/silence.tar.gz' + ]).encode('ascii')) + stream = self.client.build(fileobj=script, stream=True) + logs = '' + for chunk in stream: + logs += chunk + self.assertNotEqual(logs, '') + + class TestBuildFromStringIO(BaseTestCase): def runTest(self): if six.PY3: diff --git a/tests/test.py b/tests/test.py index dce465c1..4ba0a253 100644 --- a/tests/test.py +++ b/tests/test.py @@ -467,7 +467,21 @@ class DockerClientTest(unittest.TestCase): fake_request.assert_called_with( 'unix://var/run/docker.sock/v1.4/images/create', headers={}, - params={'tag': None, 'fromImage': 'joffrey/test001'} + params={'tag': None, 'fromImage': 'joffrey/test001'}, + stream=False + ) + + def test_pull_stream(self): + try: + self.client.pull('joffrey/test001', stream=True) + except Exception as e: + self.fail('Command should not raise exception: {0}'.format(e)) + + fake_request.assert_called_with( + 'unix://var/run/docker.sock/v1.4/images/create', + headers={}, + params={'tag': None, 'fromImage': 'joffrey/test001'}, + stream=True ) def test_commit(self): @@ -517,6 +531,20 @@ class DockerClientTest(unittest.TestCase): except Exception as e: self.fail('Command should not raise exception: {0}'.format(e)) + def test_build_container_stream(self): + script = io.BytesIO('\n'.join([ + 'FROM busybox', + 'MAINTAINER docker-py', + 'RUN mkdir -p /tmp/test', + 'EXPOSE 8080', + 'ADD https://dl.dropboxusercontent.com/u/20637798/silence.tar.gz' + ' /tmp/silence.tar.gz' + ]).encode('ascii')) + try: + self.client.build(fileobj=script, stream=True) + except Exception as e: + self.fail('Command should not raise exception: {0}'.format(e)) + ####################### ## PY SPECIFIC TESTS ## #######################