nopaque/app/api/jobs.py

from apifairy import authenticate, response
from apifairy.decorators import body, other_responses
from flask import abort, Blueprint
from werkzeug.exceptions import InternalServerError
from app import db, hashids
from app.models import Job, JobInput, JobStatus, TesseractOCRPipelineModel
from .schemas import EmptySchema, JobSchema, SpaCyNLPPipelineJobSchema, TesseractOCRPipelineJobSchema, TesseractOCRPipelineModelSchema
from .auth import auth_error_responses, token_auth


bp = Blueprint('jobs', __name__)
job_schema = JobSchema()
jobs_schema = JobSchema(many=True)
spacy_nlp_pipeline_job_schema = SpaCyNLPPipelineJobSchema()
tesseract_ocr_pipeline_job_schema = TesseractOCRPipelineJobSchema()
tesseract_ocr_pipeline_model_schema = TesseractOCRPipelineModelSchema()
tesseract_ocr_pipeline_models_schema = TesseractOCRPipelineModelSchema(many=True)


@bp.route('', methods=['GET'])
@authenticate(token_auth, role='Administrator')
@response(jobs_schema)
@other_responses(auth_error_responses)
def get_jobs():
    """Get all jobs"""
    return Job.query.all()


@bp.route('/tesseract-ocr-pipeline', methods=['POST'])
@authenticate(token_auth)
@body(tesseract_ocr_pipeline_job_schema, location='form')
@response(job_schema)
@other_responses({**auth_error_responses, InternalServerError.code: InternalServerError.description})
def create_tesseract_ocr_pipeline_job(args):
    """Create a new Tesseract OCR Pipeline job"""
    current_user = token_auth.current_user()
    try:
        job = Job.create(
            title=args['title'],
            description=args['description'],
            service='tesseract-ocr-pipeline',
            service_args={
                'model': hashids.decode(args['model_id']),
                'binarization': args['binarization']
            },
            service_version=args['service_version'],
            user=current_user
        )
    except OSError:
        abort(500)
    try:
        JobInput.create(args['pdf'], job=job)
    except OSError:
        abort(500)
    job.status = JobStatus.SUBMITTED
    db.session.commit()
    return job, 201


@bp.route('/tesseract-ocr-pipeline/models', methods=['GET'])
@authenticate(token_auth)
@response(tesseract_ocr_pipeline_models_schema)
@other_responses(auth_error_responses)
def get_tesseract_ocr_models():
    """Get all Tesseract OCR Models"""
    return TesseractOCRPipelineModel.query.all()


@bp.route('/<hashid:job_id>', methods=['DELETE'])
@authenticate(token_auth)
@response(EmptySchema, status_code=204)
@other_responses(auth_error_responses)
def delete_job(job_id):
    """Delete a job by id"""
    current_user = token_auth.current_user()
    job = Job.query.get(job_id)
    if job is None:
        abort(404)
    if not (job.user == current_user or current_user.is_administrator()):
        abort(403)
    try:
        job.delete()
    except OSError as e:
        abort(500)
    db.session.commit()
    return {}, 204


@bp.route('/<hashid:job_id>', methods=['GET'])
@authenticate(token_auth)
@response(job_schema)
@other_responses(auth_error_responses)
def get_job(job_id):
    """Get a job by id"""
    current_user = token_auth.current_user()
    job = Job.query.get(job_id)
    if job is None:
        abort(404)
    if not (job.user == current_user or current_user.is_administrator()):
        abort(403)
    return job