Merge pull request #2903 from chaoss/dev

Release Ready

Author: sgoggins

.gitignore

@@ -10,6 +10,7 @@ augur_export_env.sh
 config.yml
 reports.yml
 *.pid
+*.sock
 
 node_modules/
 .idea/

augur/api/view/init.py

@@ -91,4 +91,4 @@ def write_settings(current_settings):
 # Initialize logging
 def init_logging():
     global logger
-    logger = AugurLogger("augur_view", reset_logfiles=True).get_logger()
+    logger = AugurLogger("augur_view", reset_logfiles=False).get_logger()

augur/application/cli/__init__.py

@@ -32,7 +32,7 @@ def new_func(ctx, *args, **kwargs):
                 You are not connected to the internet.\n \
                 Please connect to the internet to run Augur\n \
                 Consider setting http_proxy variables for limited access installations.")
-            sys.exit()        
+            sys.exit(-1)        
 
     return update_wrapper(new_func, function_internet_connection)
 
@@ -78,7 +78,7 @@ def new_func(ctx, *args, **kwargs):
                 print(f"\n\n{usage} command setup failed\nERROR: connecting to database\nHINT: The {incorrect_values} may be incorrectly specified in {location}\n")
 
             engine.dispose()
-            sys.exit()
+            sys.exit(-2)
 
     return update_wrapper(new_func, function_db_connection)
 

augur/application/cli/_multicommand.py

@@ -30,7 +30,6 @@ def get_command(self, ctx, name):
 
         # Check that the command exists before importing
         if not cmdfile.is_file():
-
             return
 
         # Prefer to raise exception instead of silcencing it

augur/application/cli/api.py

@@ -14,15 +14,16 @@
 
 from augur.application.db.session import DatabaseSession
 from augur.application.logs import AugurLogger
-from augur.application.cli import test_connection, test_db_connection, with_database
+from augur.application.cli import test_connection, test_db_connection, with_database, DatabaseContext
 from augur.application.cli._cli_util import _broadcast_signal_to_processes, raise_open_file_limit, clear_redis_caches, clear_rabbitmq_messages
 from augur.application.db.lib import get_value
 
-logger = AugurLogger("augur", reset_logfiles=True).get_logger()
+logger = AugurLogger("augur", reset_logfiles=False).get_logger()
 
 @click.group('api', short_help='Commands for controlling the backend API server')
-def cli():
-    pass
+@click.pass_context
+def cli(ctx):
+    ctx.obj = DatabaseContext()
 
 @cli.command("start")
 @click.option("--development", is_flag=True, default=False, help="Enable development mode")

augur/application/cli/backend.py

@@ -47,8 +47,8 @@ def cli(ctx):
 @click.pass_context
 def start(ctx, disable_collection, development, pidfile, port):
     """Start Augur's backend server."""
-    with open(pidfile, "w") as pidfile:
-        pidfile.write(str(os.getpid()))
+    with open(pidfile, "w") as pidfile_io:
+        pidfile_io.write(str(os.getpid()))
 
     try:
         if os.environ.get('AUGUR_DOCKER_DEPLOY') != "1":
@@ -63,6 +63,8 @@ def start(ctx, disable_collection, development, pidfile, port):
     if development:
         os.environ["AUGUR_DEV"] = "1"
         logger.info("Starting in development mode")
+    
+    os.environ["AUGUR_PIDFILE"] = pidfile
 
     try:
         gunicorn_location = os.getcwd() + "/augur/api/gunicorn_conf.py"
@@ -74,6 +76,11 @@ def start(ctx, disable_collection, development, pidfile, port):
     if not port:
         port = get_value("Server", "port")
 
+    os.environ["AUGUR_PORT"] = str(port)
+    
+    if disable_collection:
+        os.environ["AUGUR_DISABLE_COLLECTION"] = "1"
+    
     worker_vmem_cap = get_value("Celery", 'worker_process_vmem_cap')
 
     gunicorn_command = f"gunicorn -c {gunicorn_location} -b {host}:{port} augur.api.server:app --log-file gunicorn.log"
@@ -128,7 +135,7 @@ def start(ctx, disable_collection, development, pidfile, port):
         augur_collection_monitor.si().apply_async()
 
     else:
-        logger.info("Collection disabled")   
+        logger.info("Collection disabled")
 
     try:
         server.wait()
@@ -153,6 +160,8 @@ def start(ctx, disable_collection, development, pidfile, port):
                 cleanup_after_collection_halt(logger, ctx.obj.engine)
             except RedisConnectionError:
                 pass
+            
+    os.unlink(pidfile)
 
 def start_celery_worker_processes(vmem_cap_ratio, disable_collection=False):
 
@@ -224,6 +233,54 @@ def stop(ctx):
 
     augur_stop(signal.SIGTERM, logger, ctx.obj.engine)
 
+@cli.command('stop-collection-blocking')
+@test_connection
+@test_db_connection
+@with_database
+@click.pass_context
+def stop_collection(ctx):
+    """
+    Stop collection tasks if they are running, block until complete
+    """
+    processes = get_augur_processes()
+    
+    stopped = []
+    
+    p: psutil.Process
+    for p in processes:
+        if p.name() == "celery":
+            stopped.append(p)
+            p.terminate()
+    
+    if not len(stopped):
+        logger.info("No collection processes found")
+        return
+    
+    _, alive = psutil.wait_procs(stopped, 5,
+                                 lambda p: logger.info(f"STOPPED: {p.pid}"))
+    
+    killed = []
+    while True:
+        for i in range(len(alive)):
+            if alive[i].status() == psutil.STATUS_ZOMBIE:
+                logger.info(f"KILLING ZOMBIE: {alive[i].pid}")
+                alive[i].kill()
+                killed.append(i)
+            elif not alive[i].is_running():
+                logger.info(f"STOPPED: {p.pid}")
+                killed.append(i)
+        
+        for i in reversed(killed):
+            alive.pop(i)
+        
+        if not len(alive):
+            break
+        
+        logger.info(f"Waiting on [{', '.join(str(p.pid for p in alive))}]")
+        time.sleep(0.5)
+    
+    cleanup_after_collection_halt(logger, ctx.obj.engine)
+
 @cli.command('kill')
 @test_connection
 @test_db_connection
@@ -388,7 +445,7 @@ def processes():
     Outputs the name/PID of all Augur server & worker processes"""
     augur_processes = get_augur_processes()
     for process in augur_processes:
-        logger.info(f"Found process {process.pid}")
+        logger.info(f"Found process {process.pid} [{process.name()}] -> Parent: {process.parent().pid}")
 
 def get_augur_processes():
     augur_processes = []

augur/application/cli/collection.py

@@ -22,14 +22,15 @@
 from augur.application.db.session import DatabaseSession
 from augur.application.logs import AugurLogger
 from augur.application.db.lib import get_value
-from augur.application.cli import test_connection, test_db_connection, with_database
+from augur.application.cli import test_connection, test_db_connection, with_database, DatabaseContext
 from augur.application.cli._cli_util import _broadcast_signal_to_processes, raise_open_file_limit, clear_redis_caches, clear_rabbitmq_messages
 
-logger = AugurLogger("augur", reset_logfiles=True).get_logger()
+logger = AugurLogger("augur", reset_logfiles=False).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("--development", is_flag=True, default=False, help="Enable development mode, implies --disable-collection")

augur/application/cli/jumpstart.py

@@ -0,0 +1,98 @@
+import psutil
+import click
+import time
+import subprocess
+from pathlib import Path
+from datetime import datetime
+
+@click.group(invoke_without_command=True)
+@click.pass_context
+def cli(ctx):
+    if ctx.invoked_subcommand is None:
+        p = check_running()
+        if not p:
+            click.echo("Jumpstart is not running. Start it with: augur jumpstart run")
+            return
+
+        click.echo(f"Connecting to Jumpstart: [{p.pid}]")
+        
+        while p.is_running() and not len(p.connections("unix")):
+            # Waiting for app to open fd socket
+            time.sleep(0.1)
+        
+        if not p.is_running():
+            click.echo("Error: Jumpstart server exited abnormally")
+            return
+        
+        from jumpstart.tui import run_app
+        run_app(ctx=ctx)
+        
+def check_running(pidfile = ".jumpstart.pid") -> psutil.Process:
+    jumpidf = Path(pidfile)
+    
+    try:
+        jumpid, create_time = jumpidf.read_text().splitlines()
+        jumpp = psutil.Process(int(jumpid))
+        
+        if create_time != str(jumpp.create_time()):
+            # PID was reused, not the original
+            jumpidf.unlink()
+            return
+        
+        return jumpp
+    except (psutil.NoSuchProcess, FileNotFoundError):
+        return
+    except PermissionError:
+        click.echo(f"Permission denied while reading from or writing to pidfile [{str(jumpidf.resolve())}]")
+
+@cli.command("status")
+def get_status():
+    p = check_running()
+    
+    if not p:
+        click.echo("Jumpstart is not running")
+    else:
+        since = datetime.fromtimestamp(p.create_time()).astimezone()
+        delta = datetime.now().astimezone() - since
+        click.echo(f"Jumpstart is running at: [{p.pid}] since {since.strftime('%a %b %d, %Y %H:%M:%S %z:%Z')} [{delta}]")
+
+@cli.command("run")
+@click.pass_context
+def startup(ctx):
+    p = check_running()
+
+    if not p:
+        click.echo("Starting")
+        p = launch(ctx)
+    else:
+        click.echo(f"Jumpstart is already running [{p.pid}]")
+
+@cli.command("processID")
+def get_main_ID():
+    p = check_running()
+    
+    if p:
+        click.echo(p.pid)
+
+@cli.command("shutdown")
+def shutdown_server():
+    p = check_running()
+    
+    if not p:
+        click.echo("Jumpstart is not running")
+        return
+    
+    click.echo("Blocking on shutdown")
+    p.terminate()
+    p.wait()
+
+def launch(ctx, pidfile = ".jumpstart.pid", socketfile = "jumpstart.sock"):
+    service = subprocess.Popen(f"python -m jumpstart.jumpstart pidfile={pidfile} socketfile={socketfile}".split())
+    
+    # Popen object does not have create_time for some reason
+    ext_process = psutil.Process(service.pid)
+
+    with open(pidfile, "w") as file:
+        file.write(f"{ext_process.pid}\n{ext_process.create_time()}")
+    
+    return ext_process

augur/application/cli/tasks.py

@@ -17,7 +17,7 @@
 from augur.application.cli import test_connection, test_db_connection 
 from augur.application.cli.backend import clear_rabbitmq_messages, raise_open_file_limit
 
-logger = AugurLogger("augur", reset_logfiles=True).get_logger()
+logger = AugurLogger("augur", reset_logfiles=False).get_logger()
 
 @click.group('celery', short_help='Commands for controlling the backend API server & data collection workers')
 def cli():

augur/application/db/data_parse.py

@@ -283,21 +283,21 @@ def extract_pr_review_message_ref_data(comment: dict, augur_pr_review_id, github
     return pr_review_comment_message_ref
 
 
-def extract_pr_event_data(event: dict, pr_id: int, platform_id: int, repo_id: int, tool_source: str, tool_version: str, data_source: str) -> dict:
+def extract_pr_event_data(event: dict, pr_id: int, gh_src_id: int, platform_id: int, repo_id: int, tool_source: str, tool_version: str, data_source: str) -> dict:
 
     pr_event = {
         'pull_request_id': pr_id,
         'cntrb_id': event["cntrb_id"] if "cntrb_id" in event else None,
         'action': event['event'],
         'action_commit_hash': None,
         'created_at': event['created_at'],
-        'issue_event_src_id': int(event['issue']["id"]),
+        'issue_event_src_id': gh_src_id,
         'node_id': event['node_id'],
         'node_url': event['url'],
         'tool_source': tool_source,
         'tool_version': tool_version,
         'data_source': data_source,
-        'pr_platform_event_id': int(event['issue']["id"]),
+        'pr_platform_event_id': gh_src_id,
         'platform_id': platform_id,
         'repo_id': repo_id
     }