Merge pull request #2721 from chaoss/cli-engine-cleanup

Handling database engines in the CLI

augur/application/cli/__init__.py

@@ -8,6 +8,7 @@
 import json
 
 from augur.application.db.engine import DatabaseEngine
+from augur.application.db import get_engine, dispose_database_engine
 from sqlalchemy.exc import OperationalError 
 
 
@@ -72,6 +73,22 @@ def new_func(ctx, *args, **kwargs):
 
     return update_wrapper(new_func, function_db_connection)
 
+
+class DatabaseContext():
+    def __init__(self):
+        self.engine = None
+
+def with_database(f):
+    @click.pass_context
+    def new_func(ctx, *args, **kwargs):
+        ctx.obj.engine = get_engine()
+        try:
+            return ctx.invoke(f, *args, **kwargs)
+        finally:
+            dispose_database_engine()
+    return new_func
+
+
 # def pass_application(f):
 #     @click.pass_context
 #     def new_func(ctx, *args, **kwargs):

augur/application/cli/backend.py

@@ -23,25 +23,27 @@
 from augur.application.db.session import DatabaseSession
 from augur.application.logs import AugurLogger
 from augur.application.config import AugurConfig
-from augur.application.cli import test_connection, test_db_connection 
+from augur.application.cli import test_connection, test_db_connection, with_database, DatabaseContext
 import sqlalchemy as s
 
 
 logger = AugurLogger("augur", reset_logfiles=True).get_logger()
 
 
-
 @click.group('server', short_help='Commands for controlling the backend API server & data collection workers')
-def cli():
-    pass
+@click.pass_context
+def cli(ctx):
+    ctx.obj = DatabaseContext()
 
 @cli.command("start")
 @click.option("--disable-collection", is_flag=True, default=False, help="Turns off data collection workers")
 @click.option("--development", is_flag=True, default=False, help="Enable development mode, implies --disable-collection")
 @click.option('--port')
 @test_connection
 @test_db_connection
-def start(disable_collection, development, port):
+@with_database
+@click.pass_context
+def start(ctx, disable_collection, development, port):
     """Start Augur's backend server."""
 
     try:
@@ -63,7 +65,7 @@ def start(disable_collection, development, port):
     except FileNotFoundError:
         logger.error("\n\nPlease run augur commands in the root directory\n\n")
 
-    with DatabaseSession(logger) as db_session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as db_session:
         config = AugurConfig(logger, db_session)
         host = config.get_value("Server", "host")
 
@@ -85,7 +87,7 @@ def start(disable_collection, development, port):
             logger.info("Deleting old task schedule")
             os.remove("celerybeat-schedule.db")
 
-    with DatabaseSession(logger) as db_session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as db_session:
         config = AugurConfig(logger, db_session)
         log_level = config.get_value("Logging", "log_level")
         celery_beat_process = None
@@ -94,7 +96,7 @@ def start(disable_collection, development, port):
 
     if not disable_collection:
 
-        with DatabaseSession(logger) as session:
+        with DatabaseSession(logger, engine=ctx.obj.engine) as session:
 
             clean_collection_status(session)
             assign_orphan_repos_to_default_user(session)
@@ -132,7 +134,7 @@ def start(disable_collection, development, port):
         if not disable_collection:
 
             try:
-                cleanup_after_collection_halt(logger)
+                cleanup_after_collection_halt(logger, ctx.obj.engine)
             except RedisConnectionError:
                 pass
 
@@ -194,24 +196,32 @@ def determine_worker_processes(ratio,maximum):
 
 
 @cli.command('stop')
-def stop():
+@test_connection
+@test_db_connection
+@with_database
+@click.pass_context
+def stop(ctx):
     """
     Sends SIGTERM to all Augur server & worker processes
     """
     logger = logging.getLogger("augur.cli")
 
-    augur_stop(signal.SIGTERM, logger)
+    augur_stop(signal.SIGTERM, logger, ctx.obj.engine)
 
 @cli.command('kill')
-def kill():
+@test_connection
+@test_db_connection
+@with_database
+@click.pass_context
+def kill(ctx):
     """
     Sends SIGKILL to all Augur server & worker processes
     """
     logger = logging.getLogger("augur.cli")
-    augur_stop(signal.SIGKILL, logger)
+    augur_stop(signal.SIGKILL, logger, ctx.obj.engine)
 
 
-def augur_stop(signal, logger):
+def augur_stop(signal, logger, engine):
     """
     Stops augur with the given signal, 
     and cleans up collection if it was running
@@ -224,13 +234,13 @@ def augur_stop(signal, logger):
     _broadcast_signal_to_processes(augur_processes, broadcast_signal=signal, given_logger=logger)
 
     if "celery" in process_names:
-        cleanup_after_collection_halt(logger)
+        cleanup_after_collection_halt(logger, engine)
 
 
-def cleanup_after_collection_halt(logger):
+def cleanup_after_collection_halt(logger, engine):
     clear_redis_caches()
     connection_string = ""
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=engine) as session:
         config = AugurConfig(logger, session)
         connection_string = config.get_section("RabbitMQ")['connection_string']
 

augur/application/cli/config.py

@@ -9,7 +9,7 @@
 
 from augur.application.db.session import DatabaseSession
 from augur.application.config import AugurConfig
-from augur.application.cli import test_connection, test_db_connection 
+from augur.application.cli import DatabaseContext, test_connection, test_db_connection, with_database
 from augur.util.inspect_without_import import get_phase_names_without_import
 ROOT_AUGUR_DIRECTORY = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))))
 
@@ -18,8 +18,9 @@
 ENVVAR_PREFIX = "AUGUR_"
 
 @click.group('config', short_help='Generate an augur.config.json')
-def cli():
-    pass
+@click.pass_context
+def cli(ctx):
+    ctx.obj = DatabaseContext()
 
 @cli.command('init')
 @click.option('--github-api-key', help="GitHub API key for data collection from the GitHub API", envvar=ENVVAR_PREFIX + 'GITHUB_API_KEY')
@@ -29,7 +30,9 @@ def cli():
 @click.option('--rabbitmq-conn-string', help="String to connect to rabbitmq broker", envvar=ENVVAR_PREFIX + 'RABBITMQ_CONN_STRING')
 @test_connection
 @test_db_connection
-def init_config(github_api_key, facade_repo_directory, gitlab_api_key, redis_conn_string, rabbitmq_conn_string):
+@with_database
+@click.pass_context
+def init_config(ctx, github_api_key, facade_repo_directory, gitlab_api_key, redis_conn_string, rabbitmq_conn_string):
 
     if not github_api_key:
 
@@ -59,7 +62,7 @@ def init_config(github_api_key, facade_repo_directory, gitlab_api_key, redis_con
     keys["github_api_key"] = github_api_key
     keys["gitlab_api_key"] = gitlab_api_key
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as session:
 
         config = AugurConfig(logger, session)
 
@@ -104,9 +107,11 @@ def init_config(github_api_key, facade_repo_directory, gitlab_api_key, redis_con
 @click.option('--file', required=True)
 @test_connection
 @test_db_connection
-def load_config(file):
+@with_database
+@click.pass_context
+def load_config(ctx, file):
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as session:
         config = AugurConfig(logger, session)
 
         print("WARNING: This will override your current config")
@@ -127,9 +132,11 @@ def load_config(file):
 @click.option('--file', required=True)
 @test_connection
 @test_db_connection
-def add_section(section_name, file):
+@with_database
+@click.pass_context
+def add_section(ctx, section_name, file):
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as session:
         config = AugurConfig(logger, session)
 
         if config.is_section_in_config(section_name):
@@ -156,9 +163,11 @@ def add_section(section_name, file):
 @click.option('--data-type', required=True)
 @test_connection
 @test_db_connection
-def config_set(section, setting, value, data_type):
+@with_database
+@click.pass_context
+def config_set(ctx, section, setting, value, data_type):
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as session:
         config = AugurConfig(logger, session)
 
         if data_type not in config.accepted_types:
@@ -180,9 +189,11 @@ def config_set(section, setting, value, data_type):
 @click.option('--setting')
 @test_connection
 @test_db_connection
-def config_get(section, setting):
+@with_database
+@click.pass_context
+def config_get(ctx, section, setting):
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, engine=ctx.obj.engine) as session:
         config = AugurConfig(logger, session)
 
         if setting:
@@ -210,9 +221,11 @@ def config_get(section, setting):
 @cli.command('clear')
 @test_connection
 @test_db_connection
-def clear_config():
+@with_database
+@click.pass_context
+def clear_config(ctx):
 
-    with DatabaseSession(logger) as session:
+    with DatabaseSession(logger, ctx.obj.engine) as session:
         config = AugurConfig(logger, session)
 
         if not config.empty():