nopaque/app/main/views.py

from flask import (abort, current_app, flash, redirect, request,
                   render_template, url_for, send_from_directory)
from flask_login import current_user, login_required
from . import main
from .forms import CreateCorpusForm
from .. import db
from ..models import Corpus, Job
import os
import logging


@main.route('/')
def index():
    return render_template('main/index.html.j2', title='Opaque')


@main.route('/corpora/<int:corpus_id>')
@login_required
def corpus(corpus_id):
    corpus = current_user.corpora.filter_by(id=corpus_id).first()
    if not corpus:
        print('Corpus not found.')
        abort(404)

    dir = os.path.join(current_app.config['OPAQUE_STORAGE_DIRECTORY'],
                       str(current_user.id),
                       'corpora',
                       str(corpus.id))
    files = {}
    for file in os.listdir(dir):
        files[file] = {}
        files[file]['path'] = os.path.join(file)

    return render_template('main/corpora/corpus.html.j2',
                           corpus=corpus,
                           files=files,
                           title='Corpus')


@main.route('/corpora/<int:corpus_id>/download')
@login_required
def corpus_download(corpus_id):
    file = request.args.get('file')
    corpus = current_user.corpora.filter_by(id=corpus_id).first()
    if not file or not corpus:
        print('File not found.')
        abort(404)
    dir = os.path.join(current_app.config['OPAQUE_STORAGE_DIRECTORY'],
                       str(current_user.id),
                       'corpora',
                       str(corpus.id))
    return send_from_directory(as_attachment=True,
                               directory=dir,
                               filename=file)


@main.route('/dashboard', methods=['GET', 'POST'])
@login_required
def dashboard():
    create_corpus_form = CreateCorpusForm()

    if create_corpus_form.validate_on_submit():
        app = current_app._get_current_object()
        corpus = Corpus(creator=current_user._get_current_object(),
                        description=create_corpus_form.description.data,
                        title=create_corpus_form.title.data)
        db.session.add(corpus)
        db.session.commit()

        dir = os.path.join(app.config['OPAQUE_STORAGE_DIRECTORY'],
                           str(corpus.user_id),
                           'corpora',
                           str(corpus.id))

        try:
            os.makedirs(dir)
        except OSError:
            flash('OSError!')
        else:
            for file in create_corpus_form.files.data:
                file.save(os.path.join(dir, file.filename))
            flash('Corpus created!')
        return redirect(url_for('main.dashboard'))

    return render_template(
        'main/dashboard.html.j2',
        title='Dashboard',
        create_corpus_form=create_corpus_form
    )


@main.route('/jobs/<int:job_id>')
@login_required
def job(job_id):
    job = current_user.jobs.filter_by(id=job_id).first()
    if not job:
        print('Job not found.')
        abort(404)

    dir = os.path.join(current_app.config['OPAQUE_STORAGE_DIRECTORY'],
                       str(current_user.id),
                       'jobs',
                       str(job.id))
    files = {}
    for file in os.listdir(dir):
        if file == 'output':
            continue
        files[file] = {}
        files[file]['path'] = os.path.join(file)
        if job.status == 'complete':
            files[file]['results'] = {}
            results_dir = os.path.join(dir, 'output', file)
            for result in os.listdir(results_dir):
                result_type = result.rsplit(".", 1)[1]
                files[file]['results'][result_type] = {}
                files[file]['results'][result_type]['path'] = os.path.join(
                    'output', files[file]['path'], result
                )

    return render_template('main/jobs/job.html.j2',
                           files=files,
                           job=job,
                           title='Job')


@main.route('/jobs/<int:job_id>/download')
@login_required
def job_download(job_id):
    file = request.args.get('file')
    job = current_user.jobs.filter_by(id=job_id).first()
    if not file or not job:
        print('File not found.')
        abort(404)
    dir = os.path.join(current_app.config['OPAQUE_STORAGE_DIRECTORY'],
                       str(current_user.id),
                       'jobs',
                       str(job.id))
    return send_from_directory(as_attachment=True,
                               directory=dir,
                               filename=file)


@main.route('/jobs/<int:job_id>/delete')
@login_required
def delete_job(job_id):
    logger = logging.getLogger(__name__)
    job = Job.query.filter_by(id=job_id).first()
    job.flag_for_stop()
    logger.warning('Job status: {}'.format(job.status))
    return redirect(url_for('main.dashboard'))