Merge pull request #8 from level12/cli-commands

Add CLI commands to plugin

keg_storage/__init__.py

@@ -1,2 +1,2 @@
-from .plugin import Storage  # noqa
-from .backends import *  # noqa
+from keg_storage.plugin import Storage  # noqa
+from keg_storage.backends import *  # noqa

keg_storage/cli.py

@@ -0,0 +1,129 @@
+import click
+import functools
+from flask.cli import with_appcontext
+import humanize
+
+from keg_storage.backends.base import FileNotFoundInStorageError
+from keg_storage import utils
+
+
+def handle_not_found(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            func(*args, **kwargs)
+        except FileNotFoundInStorageError as e:
+            raise click.FileError(
+                e.filename,
+                hint="Not found in {}.".format(str(e.storage_type))
+            )
+
+    return wrapper
+
+
+@click.command('list')
+@click.option('--simple', is_flag=True)
+@click.argument('path', default='/')
+@click.pass_context
+def storage_list(ctx, path, simple):
+    objs = ctx.obj.data['storage'].list(path)
+
+    def fmt(item):
+        fmt_str = '{key}' if simple else '{date}\t{size}\t{key}'
+        keys = {
+            'date': humanize.naturaldate(item.last_modified),
+            'key': item.key,
+            'size': humanize.naturalsize(item.size, gnu=True),
+        }
+        return fmt_str.format(**keys)
+
+    lines = [fmt(item) for item in objs]
+    click.echo("\n".join(lines))
+
+
+@click.command('get')
+@click.argument('path')
+@click.argument('dest', default='')
+@click.pass_context
+@handle_not_found
+def storage_get(ctx, path, dest):
+    if dest == '':
+        dest = path.split('/')[-1]
+
+    ctx.obj.data['storage'].get(path, dest)
+    click.echo("Downloaded {path} to {dest}.".format(path=path, dest=dest))
+
+
+@click.command('put')
+@click.argument('path')
+@click.argument('key')
+@click.pass_context
+def storage_put(ctx, path, key):
+    ctx.obj.data['storage'].put(path, key)
+    click.echo("Uploaded {path} to {key}.".format(key=key, path=path))
+
+
+@click.command('delete')
+@click.argument('path')
+@click.pass_context
+@handle_not_found
+def storage_delete(ctx, path):
+    ctx.obj.data['storage'].delete(path)
+    click.echo("Deleted {path}.".format(path=path))
+
+
+@click.command('link_for')
+@click.argument('path')
+@click.option('--expiration', '-e', default=3600,
+              help="Expiration time (in seconds) of the link, defaults to 1 hours")
+@click.pass_context
+def storage_link_for(ctx, path, expiration):
+
+    try:
+        retval = ctx.obj.data['storage'].link_for(path, expiration)
+    except Exception as e:
+        click.echo(str(e))
+        ctx.abort()
+
+    click.echo("{data}".format(data=retval))
+
+
+@click.command('reencrypt')
+@click.argument('path')
+@click.pass_context
+@handle_not_found
+def storage_reencrypt(ctx, path):
+    old_key = click.prompt('Old Key', hide_input=True).encode('ascii')
+    new_key = click.prompt('New Key', hide_input=True).encode('ascii')
+
+    utils.reencrypt(ctx.obj.data['storage'], path, old_key, new_key)
+    click.echo('Re-encrypted {path}'.format(path=path))
+
+
+def add_cli_to_app(app, cli_group_name):
+    @app.cli.group(
+        cli_group_name,
+        help='Commands for working with remotely stored files.'
+    )
+    @click.option('--location')
+    @with_appcontext
+    @click.pass_context
+    def storage(ctx, location):
+        from flask import current_app
+        location = location or current_app.config.get('KEG_STORAGE_DEFAULT_LOCATION')
+        if not location:
+            click.echo('No location given and no default was configured.')
+            ctx.abort()
+        try:
+            ctx.obj.data['storage'] = current_app.storage.get_interface(location)
+        except KeyError:
+            click.echo('The location {} does not exist. '
+                       'Pass --location or change your configuration.'.format(location))
+            ctx.abort()
+
+    storage.add_command(storage_list)
+    storage.add_command(storage_get)
+    storage.add_command(storage_put)
+    storage.add_command(storage_delete)
+    storage.add_command(storage_link_for)
+    storage.add_command(storage_reencrypt)

keg_storage/plugin.py

@@ -1,24 +1,32 @@
-
 import collections
 
+from keg_storage import cli
+
 
 class Storage:
     """A proxy and management object for storage backends."""
 
     _interaces = None
 
-    def __init__(self, app=None):
+    def __init__(self, app=None, cli_group_name='storage'):
+        self.cli_group_name = cli_group_name
 
         if app:
             self.init_app(app)
 
     def init_app(self, app):
+        app.storage = self
+
         self._interfaces = collections.OrderedDict(
             (params['name'], interface(**params))
             for interface, params in app.config['STORAGE_PROFILES']
         )
-
         self.interface = next(iter(self._interfaces)) if self._interfaces else None
 
+        self.init_cli(app)
+
     def get_interface(self, interface=None):
         return self._interfaces[interface or self.interface]
+
+    def init_cli(self, app):
+        cli.add_cli_to_app(app, self.cli_group_name)

keg_storage/tests/conftest.py

@@ -0,0 +1,5 @@
+from keg_storage_ta.app import KegStorageTestApp
+
+
+def pytest_configure(config):
+    KegStorageTestApp.testing_prep()

keg_storage/tests/test_backends.py

@@ -1,4 +1,8 @@
+import click
+import pytest
+
 import keg_storage
+from keg_storage.cli import handle_not_found
 
 
 class TestStorageBackend:
@@ -21,3 +25,17 @@ def test_methods_not_implemented(self):
                 pass
             else:
                 raise AssertionError('Should have raised exception')
+
+
+class TestFileNotFoundException:
+    def test_click_wrapper(self):
+        s3 = keg_storage.S3Storage('bucket', 'key', 'secret', name='test')
+
+        @handle_not_found
+        def test_func():
+            raise keg_storage.FileNotFoundInStorageError(s3, 'foo')
+
+        with pytest.raises(click.FileError) as exc_info:
+            test_func()
+        assert exc_info.value.filename == 'foo'
+        assert exc_info.value.message == 'Not found in S3Storage.'

keg_storage/tests/test_cli.py

@@ -0,0 +1,167 @@
+import click.testing
+from blazeutils.containers import LazyDict
+from datetime import date
+from flask import current_app
+from keg.testing import CLIBase, invoke_command
+from mock import mock
+
+from keg_storage import FileNotFoundInStorageError
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestCLI(CLIBase):
+    def test_no_location(self, m_get_interface):
+        results = self.invoke('storage', 'get', 'foo/bar', exit_code=1)
+        assert results.output == 'No location given and no default was configured.\nAborted!\n'
+
+        self.invoke('storage', '--location', 'loc', 'get', 'foo/bar')
+        m_get_interface.assert_called_once_with('loc')
+
+    def test_bad_location(self, m_get_interface):
+        m_get_interface.side_effect = KeyError
+        results = self.invoke('storage', '--location', 'loc', 'get', 'foo/bar', exit_code=1)
+        assert results.output == 'The location loc does not exist. ' \
+            'Pass --location or change your configuration.\nAborted!\n'
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestList(CLIBase):
+    cmd_name = 'storage --location loc list'
+
+    def test_list(self, m_get_interface):
+        m_list = mock.MagicMock(
+            return_value=[LazyDict(last_modified=date(2017, 4, 1), key='foo/bar', size=1024 * 3)]
+        )
+        m_get_interface.return_value.list = m_list
+
+        results = self.invoke('foo')
+        assert results.output == 'Apr 01 2017\t3.0K\tfoo/bar\n'
+        m_list.assert_called_once_with('foo')
+
+    def test_list_simple(self, m_get_interface):
+        m_list = mock.MagicMock(
+            return_value=[LazyDict(last_modified=date(2017, 4, 1), key='foo/bar', size=1024 * 3)]
+        )
+        m_get_interface.return_value.list = m_list
+
+        results = self.invoke('--simple')
+        assert results.output == 'foo/bar\n'
+        m_list.assert_called_once_with('/')
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestGet(CLIBase):
+    cmd_name = 'storage --location loc get'
+
+    def test_get_default_dest(self, m_get_interface):
+        m_get = mock.MagicMock()
+        m_get_interface.return_value.get = m_get
+
+        results = self.invoke('foo/bar')
+        assert results.output == 'Downloaded foo/bar to bar.\n'
+        m_get.assert_called_once_with('foo/bar', 'bar')
+
+    def test_get_given_dest(self, m_get_interface):
+        m_get = mock.MagicMock()
+        m_get_interface.return_value.get = m_get
+
+        results = self.invoke('foo/bar', 'dest/path')
+        assert results.output == 'Downloaded foo/bar to dest/path.\n'
+        m_get.assert_called_once_with('foo/bar', 'dest/path')
+
+    def test_get_file_not_found(self, m_get_interface):
+        m_get_interface.return_value.get.side_effect = FileNotFoundInStorageError('abc', 'def')
+
+        results = self.invoke('foo/bar', exit_code=1)
+        assert results.output == 'Error: Could not open file def: Not found in abc.\n'
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestPut(CLIBase):
+    cmd_name = 'storage --location loc put'
+
+    def test_put(self, m_get_interface):
+        m_put = mock.MagicMock()
+        m_get_interface.return_value.put = m_put
+
+        results = self.invoke('foo/bar', 'baz/bar')
+        assert results.output == 'Uploaded foo/bar to baz/bar.\n'
+        m_put.assert_called_once_with('foo/bar', 'baz/bar')
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestDelete(CLIBase):
+    cmd_name = 'storage --location loc delete'
+
+    def test_delete(self, m_get_interface):
+        m_delete = mock.MagicMock()
+        m_get_interface.return_value.delete = m_delete
+
+        results = self.invoke('foo/bar')
+        assert results.output == 'Deleted foo/bar.\n'
+        m_delete.assert_called_once_with('foo/bar')
+
+    def test_delete_file_not_found(self, m_get_interface):
+        m_get_interface.return_value.delete.side_effect = FileNotFoundInStorageError('abc', 'def')
+
+        results = self.invoke('foo/bar', exit_code=1)
+        assert results.output == 'Error: Could not open file def: Not found in abc.\n'
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+class TestLinkFor(CLIBase):
+    cmd_name = 'storage --location loc link_for'
+
+    def test_link_for_default_expiration(self, m_get_interface):
+        m_link = mock.MagicMock(return_value='http://example.com/foo/bar')
+        m_get_interface.return_value.link_for = m_link
+
+        results = self.invoke('foo/bar')
+        assert results.output == 'http://example.com/foo/bar\n'
+        m_link.assert_called_once_with('foo/bar', 3600)
+
+    def test_link_for_expiration_given(self, m_get_interface):
+        m_link = mock.MagicMock(return_value='http://example.com/foo/bar')
+        m_get_interface.return_value.link_for = m_link
+
+        results = self.invoke('-e', '100', 'foo/bar')
+        assert results.output == 'http://example.com/foo/bar\n'
+        m_link.assert_called_once_with('foo/bar', 100)
+
+    def test_link_for_error(self, m_get_interface):
+        m_get_interface.return_value.link_for.side_effect = ValueError('Some error')
+
+        results = self.invoke('foo/bar', exit_code=1)
+        assert results.output == 'Some error\nAborted!\n'
+
+
+@mock.patch.object(current_app.storage, 'get_interface', autospec=True, spec_set=True)
+@mock.patch('keg_storage.utils.reencrypt', autospec=True, spec_set=True)
+class TestReencrypt(CLIBase):
+    class InputRunner(click.testing.CliRunner):
+        def isolation(self, **kwargs):
+            kwargs['input'] = 'abc\nxyz'
+            return super(self.__class__, self).isolation(**kwargs)
+
+    def test_reencrypt(self, m_reencrypt, m_get_interface):
+        m_get_interface.return_value = 'STORAGE'
+        app_cls = current_app._get_current_object().__class__
+        results = invoke_command(app_cls, 'storage', '--location', 'loc', 'reencrypt', 'foo',
+                                 runner=self.InputRunner())
+
+        assert results.output == 'Old Key: \nNew Key: \nRe-encrypted foo\n'
+        m_reencrypt.assert_called_once_with('STORAGE', 'foo', b'abc', b'xyz')
+
+    def test_file_not_found(self, m_reencrypt, m_get_interface):
+        m_get_interface.return_value = 'STORAGE'
+        m_reencrypt.side_effect = FileNotFoundInStorageError('abc', 'def')
+
+        app_cls = current_app._get_current_object().__class__
+        results = invoke_command(app_cls, 'storage', '--location', 'loc', 'reencrypt', 'foo',
+                                 runner=self.InputRunner(), exit_code=1)
+
+        assert results.output.splitlines() == [
+            'Old Key: ',
+            'New Key: ',
+            'Error: Could not open file def: Not found in abc.'
+        ]

keg_storage/tests/test_lib.py

@@ -1,14 +1,11 @@
-import keg_storage
-
+from mock import mock
 
-class MockApp:
-    pass
+import keg_storage
 
 
 class TestStorage:
-
     def test_app_to_init_call_init(self):
-        app = MockApp()
+        app = mock.MagicMock()
         app.config = {'STORAGE_PROFILES': [(keg_storage.backends.StorageBackend, {'name': 'test'})]}
         storage = keg_storage.Storage(app)
         assert 'test' in storage._interfaces