Repository: nchammas/flintrock (python)

Issue:
Enforce CLI option dependencies in a user-friendly way
If you are launching clusters into EC2 (i.e. when `--provider` is set to `ec2`), then `--ec2-region` and `--ec2-ami` are required.

At the moment we don't enforce this dependency directly in the CLI, partly because [Click doesn't directly support option dependencies like this](https://github.com/mitsuhiko/click/issues/257#issuecomment-99302880). So if you forget to specify `--ec2-ami`, it gets caught somewhere in the bowels of Flintrock as opposed to at the front gate, and some ugly error gets barfed out.

This is a user experience problem, and it will get worse when we add support for other providers like GCE since there will be more opportunities for users to forget to specify required options (hah!).

We should fix this.

Pull request description:
Check CLI option dependencies cleanly
Option dependencies that could not be checked directly using Click are now checked using a few new utility functions. This means users get nice, fast error messages when they don't provide the right options to Flintrock, rather than getting some ugly barf from EC2 or worse--having Flintrock launch an unusable cluster.

For example, before this PR:

```
$ flintrock launch nick
Launching 2 instances...
Exception: Error reading SSH protocol banner
Traceback (most recent call last):
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1724, in _check_banner
    buf = self.packetizer.readline(timeout)
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/packet.py", line 326, in readline
    buf += self._read_timeout(timeout)
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/packet.py", line 481, in _read_timeout
    raise EOFError()
EOFError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1594, in run
    self._check_banner()
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1728, in _check_banner
    raise SSHException('Error reading SSH protocol banner' + str(e))
paramiko.ssh_exception.SSHException: Error reading SSH protocol banner

Exception: Error reading SSH protocol banner[Errno 54] Connection reset by peer
Traceback (most recent call last):
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1724, in _check_banner
    buf = self.packetizer.readline(timeout)
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/packet.py", line 326, in readline
    buf += self._read_timeout(timeout)
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/packet.py", line 479, in _read_timeout
    x = self.__socket.recv(128)
ConnectionResetError: [Errno 54] Connection reset by peer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1594, in run
    self._check_banner()
  File ".../flintrock/venv/lib/python3.5/site-packages/paramiko/transport.py", line 1728, in _check_banner
    raise SSHException('Error reading SSH protocol banner' + str(e))
paramiko.ssh_exception.SSHException: Error reading SSH protocol banner[Errno 54] Connection reset by peer

At least one node raised an error: Error reading SSH protocol banner
Do you want to terminate the 2 instances created by this operation? [Y/n]: 
Terminating instances...
At least one node raised an error: Error reading SSH protocol banner
```

And with this PR:

```
Error: Missing option "--ec2-key-name" is required by "--provider ec2".
```

Before this PR:

```
$ flintrock launch nick
EC2ResponseError: 400 Bad Request
<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>InvalidAMIID.Malformed</Code><Message>Invalid id: "None" (expecting "ami-...")</Message></Error></Errors><RequestID>275a403d-da32-4bce-81d7-e02f2e66fb18</RequestID></Response>
```

With this PR:

```
Error: Missing option "--ec2-ami" is required by "--provider ec2".
```

Fixes #5.

Patch:
diff --git a/.travis.yml b/.travis.yml
index e80feaa..0b32032 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -6,3 +6,4 @@ install:
   - "pip install -r requirements/developer.pip"
 script:
   - "py.test ./tests/test_static.py"
+  - "py.test ./tests/test_flintrock.py"
diff --git a/flintrock/ec2.py b/flintrock/ec2.py
index 4a7ad15..ba5e3fc 100644
--- a/flintrock/ec2.py
+++ b/flintrock/ec2.py
@@ -15,6 +15,7 @@ import click
 from .core import FlintrockCluster
 from .core import provision_cluster
 from .exceptions import (
+    Error,
     ClusterNotFound,
     ClusterAlreadyExists,
     ClusterInvalidState,
@@ -402,7 +403,7 @@ def launch(
             region=region)
     except boto.exception.EC2ResponseError as e:
         if e.error_code == 'InvalidAMIID.NotFound':
-            raise Exception(
+            raise Error(
                 "Error: Could not find {ami} in region {region}.".format(
                     ami=ami,
                     region=region))
diff --git a/flintrock/exceptions.py b/flintrock/exceptions.py
index 23b42c8..1767ffc 100644
--- a/flintrock/exceptions.py
+++ b/flintrock/exceptions.py
@@ -1,3 +1,7 @@
+class NothingToDo(Exception):
+    pass
+
+
 class UsageError(Exception):
     pass
 
@@ -9,19 +13,19 @@ class UnsupportedProviderError(UsageError):
         self.provider = provider
 
 
-class NothingToDo(Exception):
+class Error(Exception):
     pass
 
 
-class ClusterNotFound(Exception):
+class ClusterNotFound(Error):
     pass
 
 
-class ClusterAlreadyExists(Exception):
+class ClusterAlreadyExists(Error):
     pass
 
 
-class ClusterInvalidState(Exception):
+class ClusterInvalidState(Error):
     def __init__(self, *, attempted_command: str, state: str):
         super().__init__(
             "Cluster is in state '{s}'. Cannot execute {c}.".format(
@@ -31,7 +35,7 @@ class ClusterInvalidState(Exception):
         self.state = state
 
 
-class NodeError(Exception):
+class NodeError(Error):
     def __init__(self, error: str):
         super().__init__(
             "At least one node raised an error: " + error)
diff --git a/flintrock/flintrock.py b/flintrock/flintrock.py
index a5e74fa..c66473c 100644
--- a/flintrock/flintrock.py
+++ b/flintrock/flintrock.py
@@ -4,6 +4,7 @@ import errno
 import sys
 import shutil
 import textwrap
+import traceback
 import warnings
 
 # External modules
@@ -15,7 +16,8 @@ from . import ec2
 from .exceptions import (
     UsageError,
     UnsupportedProviderError,
-    NothingToDo)
+    NothingToDo,
+    Error)
 from flintrock import __version__
 from .services import HDFS, Spark  # TODO: Remove this dependency.
 
@@ -33,6 +35,96 @@ def format_message(*, message: str, indent: int=4, wrap: int=70):
         prefix=' ' * indent)
 
 
+def option_name_to_variable_name(option: str):
+    """
+    Convert an option name like `--ec2-user` to the Python name it gets mapped to,
+    like `ec2_user`.
+    """
+    return option.replace('--', '', 1).replace('-', '_')
+
+
+def variable_name_to_option_name(variable: str):
+    """
+    Convert a variable name like `ec2_user` to the Click option name it gets mapped to,
+    like `--ec2-user`.
+    """
+    return '--' + variable.replace('_', '-')
+
+
+def option_requires(
+        *,
+        option: str,
+        conditional_value=None,
+        requires_all: list=[],
+        requires_any: list=[],
+        scope: dict):
+    """
+    Raise an exception if an option's requirements are not met. If conditional_value
+    is not None, then only check the requirements if option is set to that value.
+
+    requires_all: Every option in this list must be defined.
+    requires_any: At least one option in this list must be defined.
+
+    This function looks for values by converting the option names to their
+    corresponding variable names (e.g. --option-a becomes option_a) and looking them
+    up in the provided scope.
+    """
+    if (conditional_value is None or
+            scope[option_name_to_variable_name(option)] == conditional_value):
+        if requires_all:
+            for required_option in requires_all:
+                required_name = option_name_to_variable_name(required_option)
+                if required_name not in scope or scope[required_name] is None:
+                    raise UsageError(
+                        'Error: Missing option "{missing_option}" is required by '
+                        '"{option}{space}{conditional_value}".'.format(
+                            missing_option=required_option,
+                            option=option,
+                            space=' ' if conditional_value is not None else '',
+                            conditional_value=conditional_value if conditional_value is not None else ''))
+        if requires_any:
+            for required_option in requires_any:
+                required_name = option_name_to_variable_name(required_option)
+                if required_name in scope and scope[required_name] is not None:
+                    break
+            else:
+                raise UsageError(
+                    'Error: "{option}{space}{conditional_value}" requires at least '
+                    'one of the following options to be set: {at_least}'.format(
+                        option=option,
+                        space=' ' if conditional_value is not None else '',
+                        conditional_value=conditional_value if conditional_value is not None else '',
+                        at_least=', '.join(['"' + ra + '"' for ra in requires_any])))
+
+
+def mutually_exclusive(*, options: list, scope: dict):
+    """
+    Raise an exception if more than one of the provided options is specified.
+
+    This function looks for values by converting the option names to their
+    corresponding variable names (e.g. --option-a becomes option_a) and looking them
+    up in the provided scope.
+    """
+    mutually_exclusive_names = [option_name_to_variable_name(o) for o in options]
+
+    used_options = set()
+    for name, value in scope.items():
+        if name in mutually_exclusive_names and scope[name]:  # is not None:
+            used_options.add(name)
+
+    if len(used_options) > 1:
+        bad_option1 = used_options.pop()
+        bad_option2 = used_options.pop()
+        raise UsageError(
+            'Error: "{option1}" and "{option2}" are mutually exclusive.\n'
+            '  {option1}: {value1}\n'
+            '  {option2}: {value2}'.format(
+                option1=variable_name_to_option_name(bad_option1),
+                value1=scope[bad_option1],
+                option2=variable_name_to_option_name(bad_option2),
+                value2=scope[bad_option2]))
+
+
 def get_config_file() -> str:
     """
     Get the path to Flintrock's default configuration file.
@@ -130,34 +222,50 @@ def launch(
     """
     Launch a new cluster.
     """
+    provider = cli_context.obj['provider']
     services = []
 
+    option_requires(
+        option='--install-hdfs',
+        requires_all=['--hdfs-version'],
+        scope=locals())
+    option_requires(
+        option='--install-spark',
+        requires_any=[
+            '--spark-version',
+            '--spark-git-commit'],
+        scope=locals())
+    mutually_exclusive(
+        options=[
+            '--spark-version',
+            '--spark-git-commit'],
+        scope=locals())
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=[
+            '--ec2-key-name',
+            '--ec2-identity-file',
+            '--ec2-instance-type',
+            '--ec2-region',
+            '--ec2-ami',
+            '--ec2-user'],
+        scope=locals())
+
     if install_hdfs:
-        if not hdfs_version:
-            raise UsageError(
-                "Error: Cannot install HDFS. Missing option \"--hdfs-version\".")
         hdfs = HDFS(version=hdfs_version)
         services += [hdfs]
     if install_spark:
-        if ((not spark_version and not spark_git_commit) or
-                (spark_version and spark_git_commit)):
-            # TODO: API for capturing option dependencies like this.
-            raise UsageError(
-                "Error: Cannot install Spark. Exactly one of \"--spark-version\" or "
-                "\"--spark-git-commit\" must be specified.\n"
-                "--spark-version: " + spark_version + "\n"
-                "--spark-git-commit: " + spark_git_commit)
-        else:
-            if spark_version:
-                spark = Spark(version=spark_version)
-            elif spark_git_commit:
-                print("Warning: Building Spark takes a long time. "
-                      "e.g. 15-20 minutes on an m3.xlarge instance on EC2.")
-                spark = Spark(git_commit=spark_git_commit,
-                              git_repository=spark_git_repository)
-            services += [spark]
-
-    if cli_context.obj['provider'] == 'ec2':
+        if spark_version:
+            spark = Spark(version=spark_version)
+        elif spark_git_commit:
+            print("Warning: Building Spark takes a long time. "
+                  "e.g. 15-20 minutes on an m3.xlarge instance on EC2.")
+            spark = Spark(git_commit=spark_git_commit,
+                          git_repository=spark_git_repository)
+        services += [spark]
+
+    if provider == 'ec2':
         return ec2.launch(
             cluster_name=cluster_name,
             num_slaves=num_slaves,
@@ -179,7 +287,7 @@ def launch(
             ebs_optimized=ec2_ebs_optimized,
             instance_initiated_shutdown_behavior=ec2_instance_initiated_shutdown_behavior)
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
 
 @cli.command()
@@ -191,12 +299,20 @@ def destroy(cli_context, cluster_name, assume_yes, ec2_region):
     """
     Destroy a cluster.
     """
-    if cli_context.obj['provider'] == 'ec2':
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=['--ec2-region'],
+        scope=locals())
+
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     if not assume_yes:
         cluster.print()
@@ -226,20 +342,27 @@ def describe(
     The output of this command is both human- and machine-friendly. Full cluster
     descriptions are output in YAML.
     """
+    provider = cli_context.obj['provider']
     search_area = ""
 
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=['--ec2-region'],
+        scope=locals())
+
     if cluster_name:
         cluster_names = [cluster_name]
     else:
         cluster_names = []
 
-    if cli_context.obj['provider'] == 'ec2':
+    if provider == 'ec2':
         search_area = "in region {r}".format(r=ec2_region)
         clusters = ec2.get_clusters(
             cluster_names=cluster_names,
             region=ec2_region)
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     if cluster_name:
         cluster = clusters[0]
@@ -277,14 +400,25 @@ def login(cli_context, cluster_name, ec2_region, ec2_identity_file, ec2_user):
     """
     Login to the master of an existing cluster.
     """
-    if cli_context.obj['provider'] == 'ec2':
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=[
+            '--ec2-region',
+            '--ec2-identity-file',
+            '--ec2-user'],
+        scope=locals())
+
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
         user = ec2_user
         identity_file = ec2_identity_file
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     # TODO: Check that master up first and error out cleanly if not
     #       via ClusterInvalidState.
@@ -304,14 +438,25 @@ def start(cli_context, cluster_name, ec2_region, ec2_identity_file, ec2_user):
     """
     Start an existing, stopped cluster.
     """
-    if cli_context.obj['provider'] == 'ec2':
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=[
+            '--ec2-region',
+            '--ec2-identity-file',
+            '--ec2-user'],
+        scope=locals())
+
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
         user = ec2_user
         identity_file = ec2_identity_file
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     cluster.start_check()
     print("Starting {c}...".format(c=cluster_name))
@@ -327,12 +472,20 @@ def stop(cli_context, cluster_name, ec2_region, assume_yes):
     """
     Stop an existing, running cluster.
     """
-    if cli_context.obj['provider'] == 'ec2':
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=['--ec2-region'],
+        scope=locals())
+
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     cluster.stop_check()
 
@@ -376,14 +529,25 @@ def run_command(
     Flintrock will return a non-zero code if any of the cluster nodes raises an error
     while running the command.
     """
-    if cli_context.obj['provider'] == 'ec2':
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=[
+            '--ec2-region',
+            '--ec2-identity-file',
+            '--ec2-user'],
+        scope=locals())
+
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
         user = ec2_user
         identity_file = ec2_identity_file
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     cluster.run_command_check()
 
@@ -431,19 +595,30 @@ def copy_file(
 
     Flintrock will return a non-zero code if any of the cluster nodes raises an error.
     """
+    provider = cli_context.obj['provider']
+
+    option_requires(
+        option='--provider',
+        conditional_value='ec2',
+        requires_all=[
+            '--ec2-region',
+            '--ec2-identity-file',
+            '--ec2-user'],
+        scope=locals())
+
     # We assume POSIX for the remote path since Flintrock
     # only supports clusters running CentOS / Amazon Linux.
     if not posixpath.basename(remote_path):
         remote_path = posixpath.join(remote_path, os.path.basename(local_path))
 
-    if cli_context.obj['provider'] == 'ec2':
+    if provider == 'ec2':
         cluster = ec2.get_cluster(
             cluster_name=cluster_name,
             region=ec2_region)
         user = ec2_user
         identity_file = ec2_identity_file
     else:
-        raise UnsupportedProviderError(cli_context.obj['provider'])
+        raise UnsupportedProviderError(provider)
 
     cluster.copy_file_check()
 
@@ -594,5 +769,9 @@ def main() -> int:
         print(e, file=sys.stderr)
         return 2
     except Exception as e:
+        if not isinstance(e, Error):
+            # This not one of our custom exceptions, so print
+            # a traceback to help the user debug.
+            traceback.print_tb(e.__traceback__, file=sys.stderr)
         print(e, file=sys.stderr)
         return 1


Test patch:
diff --git a/tests/test_flintrock.py b/tests/test_flintrock.py
new file mode 100644
index 0000000..09db5c5
--- /dev/null
+++ b/tests/test_flintrock.py
@@ -0,0 +1,128 @@
+# External modules
+import pytest
+
+# Flintrock modules
+from flintrock.exceptions import (
+    UsageError
+)
+from flintrock.flintrock import (
+    option_name_to_variable_name,
+    variable_name_to_option_name,
+    option_requires,
+    mutually_exclusive
+)
+
+
+def test_option_name_to_variable_name_conversions():
+    test_cases = [
+        ('--ec2-user', 'ec2_user'),
+        ('--provider', 'provider'),
+        ('--spark-git-commit', 'spark_git_commit')
+    ]
+
+    for option_name, variable_name in test_cases:
+        assert option_name_to_variable_name(option_name) == variable_name
+        assert variable_name_to_option_name(variable_name) == option_name
+        assert option_name == variable_name_to_option_name(
+            option_name_to_variable_name(option_name))
+        assert variable_name == option_name_to_variable_name(
+            variable_name_to_option_name(variable_name))
+
+
+def test_option_requires():
+    unset_option = None
+    set_option = '와 짠이다'
+
+    option_requires(
+        option='--some-option',
+        requires_all=['--set_option'],
+        scope=locals()
+    )
+
+    option_requires(
+        option='--some-option',
+        requires_any=[
+            '--set_option',
+            '--unset-option'],
+        scope=locals()
+    )
+
+    with pytest.raises(UsageError):
+        option_requires(
+            option='--some-option',
+            requires_all=[
+                '--set-option',
+                '--unset-option'],
+            scope=locals()
+        )
+
+    with pytest.raises(UsageError):
+        option_requires(
+            option='--some-option',
+            requires_any=[
+                '--unset-option'],
+            scope=locals()
+        )
+
+
+def test_option_requires_conditional_value():
+    with pytest.raises(Exception):
+        option_requires(
+            option='--some-option',
+            conditional_value='magic',
+            requires_any=[
+                '--set_option',
+                '--unset-option'],
+            scope=locals()
+        )
+
+    some_option = 'not magic'
+    option_requires(
+        option='--some-option',
+        conditional_value='magic',
+        requires_any=[
+            '--set_option',
+            '--unset-option'],
+        scope=locals()
+    )
+
+    some_option = 'magic'
+    with pytest.raises(UsageError):
+        option_requires(
+            option='--some-option',
+            conditional_value='magic',
+            requires_any=[
+                '--set_option',
+                '--unset-option'],
+            scope=locals()
+        )
+
+    some_option = ''
+    with pytest.raises(UsageError):
+        option_requires(
+            option='--some-option',
+            conditional_value='',
+            requires_any=[
+                '--set_option',
+                '--unset-option'],
+            scope=locals()
+        )
+
+
+def test_mutually_exclusive():
+    option1 = 'yes'
+    option2 = None
+
+    mutually_exclusive(
+        options=[
+            '--option1',
+            '--option2'],
+        scope=locals())
+
+    option2 = 'no'
+    with pytest.raises(UsageError):
+        mutually_exclusive(
+            options=[
+                '--option1',
+                '--option2'],
+            scope=locals())