diff --git a/glanceclient/common/http.py b/glanceclient/common/http.py index 5a888dd..f54d650 100644 --- a/glanceclient/common/http.py +++ b/glanceclient/common/http.py @@ -14,6 +14,8 @@ # under the License. import copy +import errno +import hashlib import httplib import logging import posixpath @@ -435,17 +437,55 @@ class VerifiedHTTPSConnection(HTTPSConnection): class ResponseBodyIterator(object): - """A class that acts as an iterator over an HTTP response.""" + """ + A class that acts as an iterator over an HTTP response. + + This class will also check response body integrity when iterating over + the instance and if a checksum was supplied using `set_checksum` method, + else by default the class will not do any integrity check. + """ def __init__(self, resp): - self.resp = resp + self._resp = resp + self._checksum = None + self._size = int(resp.getheader('content-length', 0)) + self._end_reached = False + + def set_checksum(self, checksum): + """ + Set checksum to check against when iterating over this instance. + + :raise: AttributeError if iterator is already consumed. + """ + if self._end_reached: + raise AttributeError("Can't set checksum for an already consumed" + " iterator") + self._checksum = checksum + + def __len__(self): + return int(self._size) def __iter__(self): + md5sum = hashlib.md5() while True: - yield self.next() + try: + chunk = self.next() + except StopIteration: + self._end_reached = True + # NOTE(mouad): Check image integrity when the end of response + # body is reached. + md5sum = md5sum.hexdigest() + if self._checksum is not None and md5sum != self._checksum: + raise IOError(errno.EPIPE, + 'Corrupted image. Checksum was %s ' + 'expected %s' % (md5sum, self._checksum)) + raise + else: + yield chunk + md5sum.update(chunk) def next(self): - chunk = self.resp.read(CHUNKSIZE) + chunk = self._resp.read(CHUNKSIZE) if chunk: return chunk else: diff --git a/glanceclient/common/progressbar.py b/glanceclient/common/progressbar.py new file mode 100644 index 0000000..63d4d8d --- /dev/null +++ b/glanceclient/common/progressbar.py @@ -0,0 +1,87 @@ +# Copyright 2013 OpenStack LLC. +# All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +import sys + + +class _ProgressBarBase(object): + """ + Base abstract class used by specific class wrapper to show a progress bar + when the wrapped object are consumed. + + :param wrapped: Object to wrap that hold data to be consumed. + :param totalsize: The total size of the data in the wrapped object. + + :note: The progress will be displayed only if sys.stdout is a tty. + """ + + def __init__(self, wrapped, totalsize): + self._wrapped = wrapped + self._totalsize = float(totalsize) + self._show_progress = sys.stdout.isatty() and self._totalsize != 0 + self._percent = 0 + + def _display_progress_bar(self, size_read): + if self._show_progress: + self._percent += size_read / self._totalsize + # Output something like this: [==========> ] 49% + sys.stdout.write('\r[{0:<30}] {1:.0%}'.format( + '=' * int(round(self._percent * 29)) + '>', self._percent + )) + sys.stdout.flush() + + def __getattr__(self, attr): + # Forward other attribute access to the wrapped object. + return getattr(self._wrapped, attr) + + +class VerboseFileWrapper(_ProgressBarBase): + """ + A file wrapper that show and advance a progress bar whenever file's read + method is called. + """ + + def read(self, *args, **kwargs): + data = self._wrapped.read(*args, **kwargs) + if data: + self._display_progress_bar(len(data)) + else: + # Break to a new line from the progress bar for incoming output. + sys.stdout.write('\n') + return data + + +class VerboseIteratorWrapper(_ProgressBarBase): + """ + An iterator wrapper that show and advance a progress bar whenever + data is consumed from the iterator. + + :note: Use only with iterator that yield strings. + """ + + def __iter__(self): + return self + + def next(self): + try: + data = self._wrapped.next() + # NOTE(mouad): Assuming that data is a string b/c otherwise calling + # len function will not make any sense. + self._display_progress_bar(len(data)) + return data + except StopIteration: + # Break to a new line from the progress bar for incoming output. + sys.stdout.write('\n') + raise diff --git a/glanceclient/common/utils.py b/glanceclient/common/utils.py index c0e65db..82360df 100644 --- a/glanceclient/common/utils.py +++ b/glanceclient/common/utils.py @@ -14,7 +14,6 @@ # under the License. import errno -import hashlib import os import sys import uuid @@ -161,23 +160,6 @@ def save_image(data, path): image.close() -def integrity_iter(iter, checksum): - """ - Check image data integrity. - - :raises: IOError - """ - md5sum = hashlib.md5() - for chunk in iter: - yield chunk - md5sum.update(chunk) - md5sum = md5sum.hexdigest() - if md5sum != checksum: - raise IOError(errno.EPIPE, - 'Corrupt image download. Checksum was %s expected %s' % - (md5sum, checksum)) - - def make_size_human_readable(size): suffix = ['B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB'] base = 1024.0 @@ -214,3 +196,31 @@ def exception_to_str(exc): error = ("Caught '%(exception)s' exception." % {"exception": exc.__class__.__name__}) return strutils.safe_encode(error, errors='ignore') + + +def get_file_size(file_obj): + """ + Analyze file-like object and attempt to determine its size. + + :param file_obj: file-like object. + :retval The file's size or None if it cannot be determined. + """ + if hasattr(file_obj, 'seek') and hasattr(file_obj, 'tell'): + try: + curr = file_obj.tell() + file_obj.seek(0, os.SEEK_END) + size = file_obj.tell() + file_obj.seek(curr) + return size + except IOError, e: + if e.errno == errno.ESPIPE: + # Illegal seek. This means the file object + # is a pipe (e.g the user is trying + # to pipe image data to the client, + # echo testdata | bin/glance add blah...), or + # that file object is empty, or that a file-like + # object which doesn't support 'seek/tell' has + # been supplied. + return + else: + raise diff --git a/glanceclient/v1/images.py b/glanceclient/v1/images.py index 1c09cbb..fbb406d 100644 --- a/glanceclient/v1/images.py +++ b/glanceclient/v1/images.py @@ -14,9 +14,7 @@ # under the License. import copy -import errno import json -import os import urllib from glanceclient.common import base @@ -129,9 +127,8 @@ class ImageManager(base.Manager): % urllib.quote(image_id)) checksum = resp.getheader('x-image-meta-checksum', None) if do_checksum and checksum is not None: - return utils.integrity_iter(body, checksum) - else: - return body + body.set_checksum(checksum) + return body def list(self, **kwargs): """Get a list of images. @@ -228,35 +225,6 @@ class ImageManager(base.Manager): """Delete an image.""" self._delete("/v1/images/%s" % base.getid(image)) - def _get_file_size(self, obj): - """Analyze file-like object and attempt to determine its size. - - :param obj: file-like object, typically redirected from stdin. - :retval The file's size or None if it cannot be determined. - """ - # For large images, we need to supply the size of the - # image file. See LP Bugs #827660 and #845788. - if hasattr(obj, 'seek') and hasattr(obj, 'tell'): - try: - obj.seek(0, os.SEEK_END) - obj_size = obj.tell() - obj.seek(0) - return obj_size - except IOError as e: - if e.errno == errno.ESPIPE: - # Illegal seek. This means the user is trying - # to pipe image data to the client, e.g. - # echo testdata | bin/glance add blah..., or - # that stdin is empty, or that a file-like - # object which doesn't support 'seek/tell' has - # been supplied. - return None - else: - raise - else: - # Cannot determine size of input image - return None - def create(self, **kwargs): """Create an image @@ -264,7 +232,7 @@ class ImageManager(base.Manager): """ image_data = kwargs.pop('data', None) if image_data is not None: - image_size = self._get_file_size(image_data) + image_size = utils.get_file_size(image_data) if image_size is not None: kwargs.setdefault('size', image_size) @@ -293,7 +261,7 @@ class ImageManager(base.Manager): """ image_data = kwargs.pop('data', None) if image_data is not None: - image_size = self._get_file_size(image_data) + image_size = utils.get_file_size(image_data) if image_size is not None: kwargs.setdefault('size', image_size) diff --git a/glanceclient/v1/shell.py b/glanceclient/v1/shell.py index 41bfe49..31c507b 100644 --- a/glanceclient/v1/shell.py +++ b/glanceclient/v1/shell.py @@ -25,6 +25,7 @@ else: from glanceclient import exc from glanceclient.common import utils +from glanceclient.common import progressbar from glanceclient.openstack.common import strutils import glanceclient.v1.images @@ -165,10 +166,14 @@ def do_image_show(gc, args): 'If this is not specified the image data will be ' 'written to stdout.') @utils.arg('image', metavar='', help='Name or ID of image to download.') +@utils.arg('--progress', action='store_true', default=False, + help='Show download progress bar.') def do_image_download(gc, args): """Download a specific image.""" image = utils.find_resource(gc.images, args.image) body = image.data() + if args.progress: + body = progressbar.VerboseIteratorWrapper(body, len(body)) utils.save_image(body, args.file) @@ -219,6 +224,8 @@ def do_image_download(gc, args): "May be used multiple times.")) @utils.arg('--human-readable', action='store_true', default=False, help='Print image size in a human-friendly format.') +@utils.arg('--progress', action='store_true', default=False, + help='Show upload progress bar.') def do_image_create(gc, args): """Create a new image.""" # Filter out None values @@ -241,6 +248,12 @@ def do_image_create(gc, args): _set_data_field(fields, args) + if args.progress: + filesize = utils.get_file_size(fields['data']) + fields['data'] = progressbar.VerboseFileWrapper( + fields['data'], filesize + ) + image = gc.images.create(**fields) _image_show(image, args.human_readable) @@ -287,6 +300,8 @@ def do_image_create(gc, args): "those properties not referenced are preserved.")) @utils.arg('--human-readable', action='store_true', default=False, help='Print image size in a human-friendly format.') +@utils.arg('--progress', action='store_true', default=False, + help='Show upload progress bar.') def do_image_update(gc, args): """Update a specific image.""" # Filter out None values @@ -311,6 +326,12 @@ def do_image_update(gc, args): if image.status == 'queued': _set_data_field(fields, args) + if args.progress: + filesize = utils.get_file_size(fields['data']) + fields['data'] = progressbar.VerboseFileWrapper( + fields['data'], filesize + ) + image = gc.images.update(image, purge_props=args.purge_props, **fields) _image_show(image, args.human_readable) diff --git a/glanceclient/v2/images.py b/glanceclient/v2/images.py index fe24b0a..8cb60e1 100644 --- a/glanceclient/v2/images.py +++ b/glanceclient/v2/images.py @@ -84,9 +84,8 @@ class Controller(object): resp, body = self.http_client.raw_request('GET', url) checksum = resp.getheader('content-md5', None) if do_checksum and checksum is not None: - return utils.integrity_iter(body, checksum) - else: - return body + body.set_checksum(checksum) + return body def upload(self, image_id, image_data): """ @@ -114,7 +113,7 @@ class Controller(object): try: setattr(image, key, value) except warlock.InvalidOperation, e: - raise TypeError(unicode(message)) + raise TypeError(utils.exception_to_str(e)) resp, body = self.http_client.json_request('POST', url, body=image) #NOTE(esheffield): remove 'self' for now until we have an elegant diff --git a/glanceclient/v2/shell.py b/glanceclient/v2/shell.py index ff94154..87705ba 100644 --- a/glanceclient/v2/shell.py +++ b/glanceclient/v2/shell.py @@ -13,6 +13,7 @@ # License for the specific language governing permissions and limitations # under the License. +from glanceclient.common import progressbar from glanceclient.common import utils from glanceclient import exc @@ -125,9 +126,13 @@ def do_explain(gc, args): 'If this is not specified the image data will be ' 'written to stdout.') @utils.arg('id', metavar='', help='ID of image to download.') +@utils.arg('--progress', action='store_true', default=False, + help='Show download progress bar.') def do_image_download(gc, args): """Download a specific image.""" body = gc.images.data(args.id) + if args.progress: + body = progressbar.VerboseIteratorWrapper(body, len(body)) utils.save_image(body, args.file) diff --git a/tests/test_http.py b/tests/test_http.py index e6cd770..f25227a 100644 --- a/tests/test_http.py +++ b/tests/test_http.py @@ -13,6 +13,7 @@ # License for the specific language governing permissions and limitations # under the License. +import errno import httplib import socket import StringIO @@ -241,8 +242,42 @@ class TestHostResolutionError(testtools.TestCase): class TestResponseBodyIterator(testtools.TestCase): + def test_iter_default_chunk_size_64k(self): resp = utils.FakeResponse({}, StringIO.StringIO('X' * 98304)) iterator = http.ResponseBodyIterator(resp) chunks = list(iterator) self.assertEqual(chunks, ['X' * 65536, 'X' * 32768]) + + def test_integrity_check_with_correct_checksum(self): + resp = utils.FakeResponse({}, StringIO.StringIO('CCC')) + body = http.ResponseBodyIterator(resp) + body.set_checksum('defb99e69a9f1f6e06f15006b1f166ae') + list(body) + + def test_integrity_check_with_wrong_checksum(self): + resp = utils.FakeResponse({}, StringIO.StringIO('BB')) + body = http.ResponseBodyIterator(resp) + body.set_checksum('wrong') + try: + list(body) + self.fail('integrity checked passed with wrong checksum') + except IOError as e: + self.assertEqual(errno.EPIPE, e.errno) + + def test_set_checksum_in_consumed_iterator(self): + resp = utils.FakeResponse({}, StringIO.StringIO('CCC')) + body = http.ResponseBodyIterator(resp) + list(body) + # Setting checksum for an already consumed iterator should raise an + # AttributeError. + self.assertRaises( + AttributeError, body.set_checksum, + 'defb99e69a9f1f6e06f15006b1f166ae') + + def test_body_size(self): + size = 1000000007 + resp = utils.FakeResponse( + {'content-length': str(size)}, StringIO.StringIO('BB')) + body = http.ResponseBodyIterator(resp) + self.assertEqual(len(body), size) diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py new file mode 100644 index 0000000..891f15e --- /dev/null +++ b/tests/test_progressbar.py @@ -0,0 +1,59 @@ +# Copyright 2013 OpenStack LLC. +# All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +import sys +import StringIO + +import testtools + +from glanceclient.common import progressbar +from tests import utils as test_utils + + +class TestProgressBarWrapper(testtools.TestCase): + + def test_iter_iterator_display_progress_bar(self): + size = 100 + iterator = iter('X' * 100) + saved_stdout = sys.stdout + try: + sys.stdout = output = test_utils.FakeTTYStdout() + # Consume iterator. + data = list(progressbar.VerboseIteratorWrapper(iterator, size)) + self.assertEqual(data, ['X'] * 100) + self.assertEqual( + output.getvalue().strip(), + '[%s>] 100%%' % ('=' * 29) + ) + finally: + sys.stdout = saved_stdout + + def test_iter_file_display_progress_bar(self): + size = 98304 + file_obj = StringIO.StringIO('X' * size) + saved_stdout = sys.stdout + try: + sys.stdout = output = test_utils.FakeTTYStdout() + file_obj = progressbar.VerboseFileWrapper(file_obj, size) + chunksize = 1024 + chunk = file_obj.read(chunksize) + while chunk: + chunk = file_obj.read(chunksize) + self.assertEqual( + output.getvalue().strip(), + '[%s>] 100%%' % ('=' * 29) + ) + finally: + sys.stdout = saved_stdout diff --git a/tests/test_utils.py b/tests/test_utils.py index d47c7bb..fbcf8fb 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -13,39 +13,16 @@ # License for the specific language governing permissions and limitations # under the License. -import errno -import testtools import sys import StringIO +import testtools + from glanceclient.common import utils class TestUtils(testtools.TestCase): - def test_integrity_iter_without_checksum(self): - try: - data = ''.join([f for f in utils.integrity_iter('A', None)]) - self.fail('integrity checked passed without checksum.') - except IOError as e: - self.assertEqual(errno.EPIPE, e.errno) - msg = 'was 7fc56270e7a70fa81a5935b72eacbe29 expected None' - self.assertTrue(msg in str(e)) - - def test_integrity_iter_with_wrong_checksum(self): - try: - data = ''.join([f for f in utils.integrity_iter('BB', 'wrong')]) - self.fail('integrity checked passed with wrong checksum') - except IOError as e: - self.assertEqual(errno.EPIPE, e.errno) - msg = 'was 9d3d9048db16a7eee539e93e3618cbe7 expected wrong' - self.assertTrue('expected wrong' in str(e)) - - def test_integrity_iter_with_checksum(self): - fixture = 'CCC' - checksum = 'defb99e69a9f1f6e06f15006b1f166ae' - data = ''.join([f for f in utils.integrity_iter(fixture, checksum)]) - def test_make_size_human_readable(self): self.assertEqual("106B", utils.make_size_human_readable(106)) self.assertEqual("1000kB", utils.make_size_human_readable(1024000)) @@ -53,6 +30,27 @@ class TestUtils(testtools.TestCase): self.assertEqual("1.4GB", utils.make_size_human_readable(1476395008)) self.assertEqual("9.3MB", utils.make_size_human_readable(9761280)) + def test_get_new_file_size(self): + size = 98304 + file_obj = StringIO.StringIO('X' * size) + try: + self.assertEqual(utils.get_file_size(file_obj), size) + # Check that get_file_size didn't change original file position. + self.assertEqual(file_obj.tell(), 0) + finally: + file_obj.close() + + def test_get_consumed_file_size(self): + size, consumed = 98304, 304 + file_obj = StringIO.StringIO('X' * size) + file_obj.seek(consumed) + try: + self.assertEqual(utils.get_file_size(file_obj), size) + # Check that get_file_size didn't change original file position. + self.assertEqual(file_obj.tell(), consumed) + finally: + file_obj.close() + def test_prettytable(self): class Struct: def __init__(self, **entries): diff --git a/tests/utils.py b/tests/utils.py index 0202a42..3d85951 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -33,8 +33,9 @@ class FakeAPI(object): def raw_request(self, *args, **kwargs): fixture = self._request(*args, **kwargs) - body_iter = http.ResponseBodyIterator(StringIO.StringIO(fixture[1])) - return FakeResponse(fixture[0]), body_iter + resp = FakeResponse(fixture[0], StringIO.StringIO(fixture[1])) + body_iter = http.ResponseBodyIterator(resp) + return resp, body_iter def json_request(self, *args, **kwargs): fixture = self._request(*args, **kwargs) @@ -95,3 +96,17 @@ class TestResponse(requests.Response): @property def text(self): return self._text + + +class FakeTTYStdout(StringIO.StringIO): + """A Fake stdout that try to emulate a TTY device as much as possible.""" + + def isatty(self): + return True + + def write(self, data): + # When a CR (carriage return) is found reset file. + if data.startswith('\r'): + self.seek(0) + data = data[1:] + return StringIO.StringIO.write(self, data) diff --git a/tests/v1/test_shell.py b/tests/v1/test_shell.py index 9d937be..9de0d62 100644 --- a/tests/v1/test_shell.py +++ b/tests/v1/test_shell.py @@ -365,7 +365,8 @@ class ShellStdinHandlingTests(testtools.TestCase): property={}, purge_props=False, human_readable=False, - file=None + file=None, + progress=False ) ) diff --git a/tests/v2/test_shell_v2.py b/tests/v2/test_shell_v2.py index 1826e0c..d4ef1a2 100644 --- a/tests/v2/test_shell_v2.py +++ b/tests/v2/test_shell_v2.py @@ -15,11 +15,16 @@ # under the License. # vim: tabstop=4 shiftwidth=4 softtabstop=4 +import StringIO + import mock import testtools +from glanceclient.common import http +from glanceclient.common import progressbar from glanceclient.common import utils from glanceclient.v2 import shell as test_shell +from tests import utils as test_utils class ShellV2Test(testtools.TestCase): @@ -109,15 +114,32 @@ class ShellV2Test(testtools.TestCase): self.gc.schemas.get.assert_called_once_with('test') def test_image_download(self): - args = self._make_args({'id': 'pass', 'file': 'test'}) + args = self._make_args( + {'id': 'pass', 'file': 'test', 'progress': False}) with mock.patch.object(self.gc.images, 'data') as mocked_data: - mocked_data.return_value = 'test_passed' - + resp = test_utils.FakeResponse({}, StringIO.StringIO('CCC')) + ret = mocked_data.return_value = http.ResponseBodyIterator(resp) test_shell.do_image_download(self.gc, args) mocked_data.assert_called_once_with('pass') - utils.save_image.assert_called_once_with('test_passed', 'test') + utils.save_image.assert_called_once_with(ret, 'test') + + def test_image_download_with_progressbar(self): + args = self._make_args( + {'id': 'pass', 'file': 'test', 'progress': True}) + + with mock.patch.object(self.gc.images, 'data') as mocked_data: + resp = test_utils.FakeResponse({}, StringIO.StringIO('CCC')) + mocked_data.return_value = http.ResponseBodyIterator(resp) + test_shell.do_image_download(self.gc, args) + + mocked_data.assert_called_once_with('pass') + utils.save_image.assert_called_once_with(mock.ANY, 'test') + self.assertIsInstance( + utils.save_image.call_args[0][0], + progressbar.VerboseIteratorWrapper + ) def test_do_image_delete(self): args = self._make_args({'id': 'pass', 'file': 'test'})