cjs3/cjs/database.py

#!/usr/bin/env python3

# external
import click

import utils.io 
from utils.io import read_json, get_valid_filename

import utils.fs
from utils.fs import remove_file, file_exists

import utils.cli
from utils.cli import echo_title, echo_field

import core.database


def init_database_skel(root):
    create_dir(root)
    create_dir(f'{root}/input')
    create_dir(f'{root}/output')
    create_dir(f'{root}/log')

# BEGIN command-line interface

@click.group()
@click.pass_context
def database(ctx):
    """
    Manage automation databases.\f

    Args:
        ctx: Click command-line interface context.
    """

    # ensure database directory exists
    database_dir = '%s/database' % ctx.obj['data_dir']
    utils.fs.create_dir(database_dir)

    # store database_dir to context object
    ctx.obj['database_dir'] = database_dir

@database.command()
@click.argument('database_name')
@click.pass_context
def new(ctx, database_name):
    '''
    Create a new database.\f
    
    Args:
        ctx: Click command-line interface context.
    '''

    try:
        # try to register the new database
        core.database.new(database_name)
        #new_dir = f'{ctx.obj['database_dir']}/{get_valid_filename(database_name)}'
        #init_database_skel(new_dir)
    except Exception as e:
        # TODO: better exception handling
        click.echo("Error: Unable to create '%s': %s" % (db, e))
        ctx.exit()
    
    # if database was created successfully, initialize

@database.command(name='list')
@click.pass_context
def list_databases(ctx):
    '''
    Display available databases.\f
    
    Args:
        ctx: Click command-line interface context.
    '''

    echo_title('Available Databases', fg='yellow', bold=True)
    databases = core.database.list_databases()
    for i, db in enumerate(databases, start=1):
        echo_field(db.name, db.root, fg='yellow', bold=True)
        

@database.command()
@click.argument('database_name')
@click.pass_context
def delete(ctx, db):
    '''
    Remove a database.\f
    
    Args:
        ctx: Click command-line interface context.
        src: Source directory containing input file(s)
    '''

    if click.confirm("Warning: This action cannot be undone!\nDelete database? (%s)" % database):
        try:
            core.database.delete(db)
        except:
            click.echo('Error: Unable to create new database.')
            ctx.exit(1) # TODO: change exit code
    else:
        click.echo('No changes were made.')
    click.echo('database successfully deleted.')

@database.command()
@click.argument('database')
@click.argument('input_file')
@click.pass_context
def upload(ctx, db, input_file):
    '''
    Add data to a database.\f

    Args:
        ctx: Click command-line interface context.
        src: Source directory containing input file(s)
    '''

    click.echo('Please wait, this may take awhile...')
    
    try:
        data = read_json(input_file)
    except:
        click.echo("Error: Unable to read '%s'." % input_file)
        ctx.exit(1)

    with click.progressbar(data.items()) as bar:
        for slug, evidence in bar:
            # break apart data into separate files
            comparisons = evidence.pop('comparisons')

            core.database.upload(db, evidence, comparisons)

@database.command()
@click.argument('database')
@click.pass_context
def clear(ctx, db):
    '''
    Clear jobs.\f

    Args:
        ctx: Click command-line interface context.
    '''
    
    if click.confirm("Warning: This action cannot be undone!\nClear the Job table? (%s)" % database):
        try:
            core.database.clear(db)
        except:
            click.echo("Error: Unable to clear jobs.")
            ctx.exit(1)
        click.echo('Job table successfully cleared.')
    else:
        click.echo('No changes were made.')


@database.command()
@click.argument('database')
@click.argument('batch_name', type=str)
@click.pass_context
def stage(ctx, db, batch_name):
    '''
    Generate jobs.\f
    
    Args:
        ctx: Click command-line interface context.
    '''

    if click.confirm("Warning: This may cause overlapping jobs!\nLoad the Job table? (%s)" % database):
        core.database.stage(db, batch_name)

    else:
        click.echo('No changes were made.')

@database.command()
@click.argument('database')
def status(db):
    '''
    Display information about a database.\f
    
    Args:
        ctx: Click command-line interface context.
        src: Source directory containing input file(s)
    '''
  
    from core.database import database_status
    s = core.database.status(db)

    if s['num_jobs'] > 0:
        plabel('Job Tracker', '', 'yellow')
        plabel('  Total', s['num_jobs'], 'yellow')
        plabel('  Pending', (s['num_pending']/s['num_jobs'])*100, 'yellow')
        plabel('  Processing', (s['num_processing']/s['num_jobs'])*100, 'yellow')
        plabel('  Complete', (s['num_complete']/s['num_jobs'])*100, 'yellow')
        click.echo('')
        evidence, comparison = s['next_job']
        
        plabel('Next Job', '%s and %s' % (evidence, comparison), 'yellow')
    else:
        click.echo('No job currently staged.')
    

@database.command()
@click.argument('database')
@click.pass_context
def clone(ctx, db):
    '''
    Copy a database.\f
    '''

    click.echo('New database:')
    name = click.prompt("Name (e.g. 'my_database')", type=str)
    while True:
        path = click.prompt("Location (e.g. 'path/to/database.db')", type=str)
        if not file_exists(path):
            break
        click.echo("'%s' already exists." % path)
    try:
        core.database.copy(db, path)
        core.database.new(name, path)
    except Exception as e:
        click.echo(e)
        raise
        click.echo('Error: Unable to create new database.')
        ctx.exit(1) # TODO: change exit code

# END command-line interface