You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

224 lines
5.7 KiB

#!/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