Pipelines CLI Fixes (#298)

This PR makes the following changes to the Pipelines CLI:

* A successful spec deployment prints the URL to the pipeline.
* If the provided spec file does not have an ID, one is generated and the spec file is updated with the new ID.
* Adds a new exception eater to print cleaner error messages from the pipelines API
* Adds exception handing for specs with invalid json
* A validation check for pipeline id: only -, _ and alphanumeric characters are allowed

Testing:
* Unit Tests
* Manually verified deployment URLs, new exception formatting, validation of pipeline ids

Author: null-sleep

databricks_cli/pipelines/cli.py

@@ -21,16 +21,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from json import loads as json_loads
 import os
+import uuid
+import json
+import string
+
+try:
+    from urlparse import urlparse, urljoin
+except ImportError:
+    from urllib.parse import urlparse, urljoin
 
 import click
 
 from databricks_cli.click_types import PipelineSpecClickType, PipelineIdClickType
-from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, pretty_format
 from databricks_cli.version import print_version_callback, version
 from databricks_cli.pipelines.api import PipelinesApi
 from databricks_cli.configure.config import provide_api_client, profile_option, debug_option
+from databricks_cli.utils import pipelines_exception_eater, CONTEXT_SETTINGS, pretty_format, \
+    error_and_quit
+
+try:
+    json_parse_exception = json.decoder.JSONDecodeError
+except AttributeError:  # Python 2
+    json_parse_exception = ValueError
+
+PIPELINE_ID_PERMITTED_CHARACTERS = set(string.ascii_letters + string.digits + '-_')
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
@@ -39,7 +54,7 @@
 @click.option('--spec', default=None, type=PipelineSpecClickType(), help=PipelineSpecClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions
+@pipelines_exception_eater
 @provide_api_client
 def deploy_cli(api_client, spec_arg, spec):
     """
@@ -59,8 +74,19 @@ def deploy_cli(api_client, spec_arg, spec):
         raise RuntimeError('The spec should be provided either by an option or argument')
     src = spec_arg if bool(spec_arg) else spec
     spec_obj = _read_spec(src)
+    if 'id' not in spec_obj:
+        pipeline_id = str(uuid.uuid4())
+        click.echo("Updating spec at {} with id: {}".format(src, pipeline_id))
+        spec_obj['id'] = pipeline_id
+        _write_spec(src, spec_obj)
+    _validate_pipeline_id(spec_obj['id'])
     PipelinesApi(api_client).deploy(spec_obj)
 
+    pipeline_id = spec_obj['id']
+    base_url = "{0.scheme}://{0.netloc}/".format(urlparse(api_client.url))
+    pipeline_url = urljoin(base_url, "#joblist/pipelines/{}".format(pipeline_id))
+    click.echo("Pipeline successfully deployed: {}".format(pipeline_url))
+
 
 @click.command(context_settings=CONTEXT_SETTINGS,
                short_help='Stops a delta pipeline and deletes its associated Databricks resources')
@@ -70,7 +96,7 @@ def deploy_cli(api_client, spec_arg, spec):
               help=PipelineIdClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions
+@pipelines_exception_eater
 @provide_api_client
 def delete_cli(api_client, spec_arg, spec, pipeline_id):
     """
@@ -91,6 +117,7 @@ def delete_cli(api_client, spec_arg, spec, pipeline_id):
     """
     pipeline_id = _get_pipeline_id(spec_arg=spec_arg, spec=spec, pipeline_id=pipeline_id)
     PipelinesApi(api_client).delete(pipeline_id)
+    click.echo("Pipeline {} deleted".format(pipeline_id))
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
@@ -101,7 +128,7 @@ def delete_cli(api_client, spec_arg, spec, pipeline_id):
               help=PipelineIdClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions
+@pipelines_exception_eater
 @provide_api_client
 def get_cli(api_client, spec_arg, spec, pipeline_id):
     """
@@ -131,7 +158,7 @@ def get_cli(api_client, spec_arg, spec, pipeline_id):
               help=PipelineIdClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions
+@pipelines_exception_eater
 @provide_api_client
 def reset_cli(api_client, spec_arg, spec, pipeline_id):
     """
@@ -152,6 +179,7 @@ def reset_cli(api_client, spec_arg, spec, pipeline_id):
     """
     pipeline_id = _get_pipeline_id(spec_arg=spec_arg, spec=spec, pipeline_id=pipeline_id)
     PipelinesApi(api_client).reset(pipeline_id)
+    click.echo("Reset triggered for pipeline {}".format(pipeline_id))
 
 
 def _read_spec(src):
@@ -161,13 +189,25 @@ def _read_spec(src):
     """
     extension = os.path.splitext(src)[1]
     if extension.lower() == '.json':
-        with open(src, 'r') as f:
-            json = f.read()
-        return json_loads(json)
+        try:
+            with open(src, 'r') as f:
+                data = f.read()
+            return json.loads(data)
+        except json_parse_exception as e:
+            error_and_quit("Invalid JSON provided in spec\n{}".format(e))
     else:
         raise RuntimeError('The provided file extension for the spec is not supported')
 
 
+def _write_spec(src, spec):
+    """
+    Writes the spec at src as JSON.
+    """
+    data = json.dumps(spec, indent=2) + '\n'
+    with open(src, 'w') as f:
+        f.write(data)
+
+
 def _get_pipeline_id(spec_arg, spec, pipeline_id):
     """
     Ensures that the user has either specified a spec (either through argument or option) or a
@@ -180,9 +220,22 @@ def _get_pipeline_id(spec_arg, spec, pipeline_id):
     if bool(spec_arg) or bool(spec):
         src = spec_arg if bool(spec_arg) else spec
         pipeline_id = _read_spec(src)["id"]
+    _validate_pipeline_id(pipeline_id)
     return pipeline_id
 
 
+def _validate_pipeline_id(pipeline_id):
+    """
+    Checks if the pipeline_id only contain -, _ and alphanumeric characters
+    """
+    if len(pipeline_id) == 0:
+        error_and_quit(u'Empty pipeline id provided')
+    if not set(pipeline_id) <= PIPELINE_ID_PERMITTED_CHARACTERS:
+        message = u'Pipeline id {} has invalid character(s)\n'.format(pipeline_id)
+        message += u'Valid characters are: _ - a-z A-Z 0-9'
+        error_and_quit(message)
+
+
 @click.group(context_settings=CONTEXT_SETTINGS,
              short_help='Utility to interact with the Databricks Delta Pipelines.')
 @click.option('--version', '-v', is_flag=True, callback=print_version_callback,

databricks_cli/utils.py

@@ -46,9 +46,41 @@ def decorator(*args, **kwargs):
                                'reconfigure with ``dbfs configure``')
             else:
                 error_and_quit(exception.response.content)
-        except Exception as exception: # noqa
+        except Exception as exception:  # noqa
             if not DEBUG_MODE:
                 error_and_quit('{}: {}'.format(type(exception).__name__, str(exception)))
+
+    decorator.__doc__ = function.__doc__
+    return decorator
+
+
+def pipelines_exception_eater(function):
+    """
+    Formats error messages from the pipelines API while keeping the existing
+    behavior of eat_exception
+    """
+
+    @six.wraps(function)
+    def decorator(*args, **kwargs):
+        try:
+            return function(*args, **kwargs)
+        except HTTPError as exception:  # noqa
+            if exception.response.status_code == 401:
+                error_and_quit('Your authentication information may be incorrect. Please '
+                               + 'reconfigure with ``dbfs configure``')
+            else:
+                try:
+                    exp_context = json_loads(exception.response.content.decode('utf-8'))
+                    message = exception.response.content
+                    if 'error_code' in exp_context and 'message' in exp_context:
+                        message = exp_context['error_code'] + '\n' + exp_context['message']
+                    error_and_quit(message)
+                except Exception:  # noqa
+                    error_and_quit(exception.response.content)
+        except Exception as exception:  # noqa
+            if not DEBUG_MODE:
+                error_and_quit('{}: {}'.format(type(exception).__name__, str(exception)))
+
     decorator.__doc__ = function.__doc__
     return decorator
 
@@ -58,7 +90,7 @@ def error_and_quit(message):
     context_object = ctx.ensure_object(ContextObject)
     if context_object.debug_mode:
         traceback.print_exc()
-    click.echo('Error: {}'.format(message))
+    click.echo(u'Error: {}'.format(message))
     sys.exit(1)
 
 

tests/pipelines/test_cli.py

@@ -24,8 +24,10 @@
 # pylint:disable=redefined-outer-name
 
 import json
+
 import mock
 import pytest
+from click import Context, Command
 from click.testing import CliRunner
 
 import databricks_cli.pipelines.cli as cli
@@ -43,6 +45,15 @@ def pipelines_api_mock():
         yield _pipelines_api_mock
 
 
+@pytest.fixture()
+def click_ctx():
+    """
+    A dummy Click context to allow testing of methods that raise exceptions. Fixes Click capturing
+    actual exceptions and raising its own `RuntimeError: There is no active click context`.
+    """
+    return Context(Command('cmd'))
+
+
 @provide_conf
 def test_deploy_cli_spec_arg(pipelines_api_mock, tmpdir):
     path = tmpdir.join('/spec.json').strpath
@@ -124,6 +135,31 @@ def test_delete_cli_incorrect_parameters(pipelines_api_mock, tmpdir):
     assert pipelines_api_mock.delete.call_count == 0
 
 
+@provide_conf
+def test_deploy_spec_updated_with_id_if_pipeline_id_not_in_spec(tmpdir):
+    path = tmpdir.join('/spec.json').strpath
+    spec_with_no_id = '{"name": "no id"}'
+    with open(path, 'w') as f:
+        f.write(spec_with_no_id)
+    runner = CliRunner()
+    runner.invoke(cli.deploy_cli, ['--spec', path])
+    with open(path, 'r') as f:
+        updated_spec = json.loads(f.read())
+    assert 'id' in updated_spec
+
+
+@provide_conf
+def test_deploy_spec_pipeline_id_is_not_changed_if_provided_in_spec(tmpdir):
+    path = tmpdir.join('/spec.json').strpath
+    with open(path, 'w') as f:
+        f.write(DEPLOY_SPEC)
+    runner = CliRunner()
+    runner.invoke(cli.deploy_cli, ['--spec', path])
+    with open(path, 'r') as f:
+        spec = json.loads(f.read())
+    assert spec['id'] == '123'
+
+
 @provide_conf
 def test_deploy_delete_cli_incorrect_spec_extension(pipelines_api_mock, tmpdir):
     path = tmpdir.join('/spec.wrong_ext').strpath
@@ -244,3 +280,15 @@ def test_get_cli_no_id(pipelines_api_mock):
     result = runner.invoke(cli.get_cli, [])
     assert result.exit_code == 1
     assert pipelines_api_mock.get.call_count == 0
+
+
+def test_validate_pipeline_id(click_ctx):
+    empty_pipeline_id = ''
+    pipeline_id_with_unicode = b'pipeline_id-\xe2\x9d\x8c-123'.decode('utf-8')
+    invalid_pipline_ids = ['pipeline_id-?-123', 'pipeline_id-\\-\'-123', 'pipeline_id-/-123',
+                           pipeline_id_with_unicode, empty_pipeline_id]
+    with click_ctx:
+        for pipline_id in invalid_pipline_ids:
+            with pytest.raises(SystemExit):
+                cli._validate_pipeline_id(pipline_id)
+    assert cli._validate_pipeline_id('pipeline_id-ac345cd1') is None

tests/test_utils.py

@@ -33,6 +33,7 @@ def test_eat_exceptions_normal_case():
     """
     If no exceptions, this wrapper should do nothing.
     """
+
     @utils.eat_exceptions
     def test_function(x):
         return x
@@ -50,11 +51,73 @@ def test_function():
             resp = Response()
             resp.status_code = 401
             raise HTTPError(response=resp)
+
+        test_function()
+        assert error_and_quit_mock.call_count == 1
+        assert 'Your authentication information' in error_and_quit_mock.call_args[0][0]
+
+
+def test_pipelines_exception_eater_normal_case():
+    """
+    If no exceptions, this wrapper should do nothing.
+    """
+
+    @utils.pipelines_exception_eater
+    def test_function(x):
+        return x
+
+    assert test_function(1) == 1
+
+
+def test_pipelines_exception_eater_http_error_401():
+    """
+    If wrapped function returns 401 HTTPError, then print special error message.
+    """
+    with mock.patch('databricks_cli.utils.error_and_quit') as error_and_quit_mock:
+        @utils.pipelines_exception_eater
+        def test_function():
+            resp = Response()
+            resp.status_code = 401
+            raise HTTPError(response=resp)
+
         test_function()
         assert error_and_quit_mock.call_count == 1
         assert 'Your authentication information' in error_and_quit_mock.call_args[0][0]
 
 
+def test_pipelines_exception_eater_non_401_http_error():
+    """
+    If wrapped function returns a non 401 HTTPError, then try to parse json response
+    to print a formatted error message.
+    """
+    with mock.patch('databricks_cli.utils.error_and_quit') as error_and_quit_mock:
+        @utils.pipelines_exception_eater
+        def test_function(content):
+            resp = Response()
+            resp.status_code = 400
+            resp._content_consumed = True
+            resp._content = content
+            raise HTTPError(response=resp)
+
+        test_function(content=b'{"error_code":"TEST_ERROR_CODE","message":"test message"}')
+        assert error_and_quit_mock.call_count == 1
+        assert error_and_quit_mock.call_args[0][0] == 'TEST_ERROR_CODE\ntest message'
+        test_function(content=b'{"message":"test message"}')
+        assert error_and_quit_mock.call_count == 2
+        assert error_and_quit_mock.call_args[0][0] == b'{"message":"test message"}'
+
+
+def test_pipelines_exception_eater_non_http_error_exceptions():
+    with mock.patch('databricks_cli.utils.error_and_quit') as error_and_quit_mock:
+        @utils.pipelines_exception_eater
+        def test_function():
+            raise ValueError('value error test message')
+
+        test_function()
+        assert error_and_quit_mock.call_count == 1
+        assert error_and_quit_mock.call_args[0][0] == 'ValueError: value error test message'
+
+
 def test_json_cli_base_both_args():
     with pytest.raises(RuntimeError):
         utils.json_cli_base('a', 'b', mock.Mock())