Add --overwrite_virtualenv flag to submit and update_virtualenv (#416)

streamparse/cli/common.py

@@ -120,6 +120,15 @@ def add_override_name(parser):
                              'duplicate the topology file.')
 
 
+def add_overwrite_virtualenv(parser):
+    """ Add --overwrite_virtualenv option to parser """
+    parser.add_argument('--overwrite_virtualenv',
+                        help='Create the virtualenv even if it already exists.'
+                             ' This is useful when you have changed your '
+                             'virtualenv_flags.',
+                        action='store_true')
+
+
 def add_pattern(parser):
     """ Add --pattern option to parser """
     parser.add_argument('--pattern',

streamparse/cli/submit.py

@@ -23,9 +23,9 @@
                     get_topology_from_file, nimbus_storm_version,
                     set_topology_serializer, ssh_tunnel, warn)
 from .common import (add_ackers, add_config, add_debug, add_environment,
-                     add_name, add_options, add_override_name, add_pool_size,
-                     add_requirements, add_timeout, add_wait, add_workers,
-                     resolve_options)
+                     add_name, add_options, add_override_name,
+                     add_overwrite_virtualenv, add_pool_size, add_requirements,
+                     add_timeout, add_wait, add_workers, resolve_options)
 from .jar import jar_for_deploy
 from .kill import _kill_topology
 from .list import _list_topologies
@@ -143,7 +143,8 @@ def _upload_jar(nimbus_client, local_path):
 def submit_topology(name=None, env_name=None, options=None, force=False,
                     wait=None, simple_jar=True, override_name=None,
                     requirements_paths=None, local_jar_path=None,
-                    remote_jar_path=None, timeout=None, config_file=None):
+                    remote_jar_path=None, timeout=None, config_file=None,
+                    overwrite_virtualenv=False):
     """Submit a topology to a remote Storm cluster."""
     config = get_config()
     name, topology_file = get_topology_definition(name, config_file=config_file)
@@ -178,7 +179,8 @@ def submit_topology(name=None, env_name=None, options=None, force=False,
             create_or_update_virtualenvs(env_name, name, options,
                                          virtualenv_name=virtualenv_name,
                                          requirements_paths=requirements_paths,
-                                         config_file=config_file)
+                                         config_file=config_file,
+                                         overwrite_virtualenv=overwrite_virtualenv)
         streamparse_run_path = '/'.join([env.virtualenv_root, virtualenv_name,
                                          'bin', 'streamparse_run'])
         # Update python paths in bolts
@@ -259,6 +261,7 @@ def subparser_hook(subparsers):
     add_name(subparser)
     add_options(subparser)
     add_override_name(subparser)
+    add_overwrite_virtualenv(subparser)
     add_pool_size(subparser)
     add_requirements(subparser)
     subparser.add_argument('-R', '--remote_jar_path',
@@ -288,4 +291,5 @@ def main(args):
                     local_jar_path=args.local_jar_path,
                     remote_jar_path=args.remote_jar_path,
                     timeout=args.timeout,
-                    config_file=args.config)
+                    config_file=args.config,
+                    overwrite_virtualenv=args.overwrite_virtualenv)

streamparse/cli/update_virtualenv.py

@@ -13,7 +13,8 @@
 from fabric.contrib.files import exists
 from six import string_types
 
-from .common import (add_config, add_environment, add_name, add_options, add_override_name,
+from .common import (add_config, add_environment, add_name, add_options,
+                     add_override_name, add_overwrite_virtualenv,
                      add_pool_size, add_requirements, resolve_options)
 from ..util import (activate_env, die, get_config, get_env_config,
                     get_topology_definition, get_topology_from_file)
@@ -23,9 +24,13 @@
 def _create_or_update_virtualenv(virtualenv_root,
                                  virtualenv_name,
                                  requirements_paths,
-                                 virtualenv_flags=None):
+                                 virtualenv_flags=None,
+                                 overwrite_virtualenv=False):
     with show('output'):
         virtualenv_path = '/'.join((virtualenv_root, virtualenv_name))
+        if overwrite_virtualenv:
+            puts("Removing virtualenv if it exists in {}".format(virtualenv_root))
+            run('rm -rf {}'.format(virtualenv_path))
         if not exists(virtualenv_path):
             if virtualenv_flags is None:
                 virtualenv_flags = ''
@@ -51,8 +56,9 @@ def _create_or_update_virtualenv(virtualenv_root,
             run("rm {}".format(temp_req))
 
 
-def create_or_update_virtualenvs(env_name, topology_name, options, virtualenv_name=None,
-                                 requirements_paths=None, config_file=None):
+def create_or_update_virtualenvs(env_name, topology_name, options,
+                                 virtualenv_name=None, requirements_paths=None,
+                                 config_file=None, overwrite_virtualenv=False):
     """Create or update virtualenvs on remote servers.
 
     Assumes that virtualenv is on the path of the remote server(s).
@@ -63,6 +69,8 @@ def create_or_update_virtualenvs(env_name, topology_name, options, virtualenv_na
                           though the topology file has a different name.
     :param requirements_paths: a list of paths to requirements files to use to
                                create virtualenv
+    :param overwrite_virtualenv: Force the creation of a fresh virtualenv, even
+                                 if it already exists.
     """
     config = get_config()
     topology_name, topology_file = get_topology_definition(topology_name, config_file=config_file)
@@ -99,7 +107,8 @@ def create_or_update_virtualenvs(env_name, topology_name, options, virtualenv_na
     execute(_create_or_update_virtualenv, env.virtualenv_root, virtualenv_name,
             requirements_paths,
             virtualenv_flags=options.get('virtualenv_flags'),
-            hosts=env.storm_workers)
+            hosts=env.storm_workers,
+            overwrite_virtualenv=overwrite_virtualenv)
 
 
 def subparser_hook(subparsers):
@@ -110,6 +119,7 @@ def subparser_hook(subparsers):
     subparser.set_defaults(func=main)
     add_config(subparser)
     add_environment(subparser)
+    add_overwrite_virtualenv(subparser)
     add_name(subparser)
     add_options(subparser)
     add_override_name(subparser)
@@ -123,4 +133,5 @@ def main(args):
     create_or_update_virtualenvs(args.environment, args.name, args.options,
                                  virtualenv_name=args.override_name,
                                  requirements_paths=args.requirements,
-                                 config_file=args.config)
+                                 config_file=args.config,
+                                 overwrite_virtualenv=args.overwrite_virtualenv)