nopaque/app/corpora/events.py

from app import db, logger, socketio
from app.events import connected_sessions
from app.models import Corpus, User
from .cqi import CQiClient
from flask import current_app, request
from flask_login import current_user, login_required
import math


'''
' A dictionary containing lists of, with corpus ids associated, Socket.IO
' session ids (sid). {<corpus_id>: [<sid>, ...], ...}
'''
corpus_analysis_sessions = {}
'''
' A dictionary containing Socket.IO session id - CQi client pairs.
' {<sid>: CQiClient, ...}
'''
corpus_analysis_clients = {}


@socketio.on('corpus_analysis_init')
@login_required
def init_corpus_analysis(corpus_id):
    logger.warning("Initiating corpus analysis.")
    socketio.start_background_task(corpus_analysis_session_handler,
                                   current_app._get_current_object(),
                                   corpus_id, current_user.id, request.sid)


@socketio.on('corpus_analysis_query')
@login_required
def corpus_analysis_query(query):
    client = corpus_analysis_clients.get(request.sid)
    if client is None:
        socketio.emit('query', '[424]: Failed Dependency', room=request.sid)
        return
    corpus = client.corpora.get('CORPUS')
    try:
        results = corpus.query(query)
    except Exception as e:
        logger.warning(e)
        response = str(e)
        if response == "CQI_CQP_ERROR_GENERAL":
            response = {'code': 1}
            socketio.emit('corpus_analysis_query', response, room=request.sid)
    else:
        chunk_size = 100
        chunk_start = 0
        context = 100
        progress = 0
        while chunk_start <= results.size:
            chunk = results.dump_values(context=context,
                                        first_result=chunk_start,
                                        num_results=chunk_size)
            progress = ((chunk_start + chunk_size) / results.size) * 100
            progress = min(100, int(math.ceil(progress)))
            socketio.emit('corpus_analysis_query',
                          {'chunk': chunk, 'progress': progress,
                           'num_matches_total': results.size},
                          room=request.sid)
            chunk_start += chunk_size


@socketio.on('inspect_match')
@login_required
def inspect_match(message):
    client = corpus_analysis_clients.get(request.sid)
    if client is None:
        socketio.emit('query', '[424]: Failed Dependency',
                      room=request.sid)
        return
    # Get more context for given match CPOS
    match_context = client.get_sentences(message['cpos'],
                                         get_surrounding_s=True,
                                         l_r_s_context_additional_len=3)
    socketio.emit('match_context', match_context, room=request.sid)


def corpus_analysis_session_handler(app, corpus_id, user_id, session_id):
    with app.app_context():
        ''' Setup analysis session '''
        corpus = Corpus.query.get(corpus_id)
        user = User.query.get(user_id)
        if corpus is None:
            response = {'code': 404, 'msg': 'Not Found'}
            socketio.emit('corpus_analysis_init', response, room=session_id)
            return
        elif not (corpus.creator == user or user.is_administrator()):
            response = {'code': 403, 'msg': 'Forbidden'}
            socketio.emit('corpus_analysis_init', response, room=session_id)
            return
        while corpus.status != 'analysing':
            db.session.refresh(corpus)
            socketio.sleep(3)
        client = CQiClient('corpus_{}_analysis'.format(corpus_id))
        try:
            client.connect()
        except Exception:
            response = {'code': 500, 'msg': 'Internal Server Error'}
            socketio.emit('corpus_analysis_init', response, room=session_id)
            return
        corpus_analysis_clients[session_id] = client
        if corpus_id not in corpus_analysis_sessions:
            corpus_analysis_sessions[corpus_id] = [session_id]
        else:
            corpus_analysis_sessions[corpus_id].append(session_id)
        response = {'code': 201, 'msg': 'Created'}
        socketio.emit('corpus_analysis_init', response, room=session_id)
        ''' Observe analysis session '''
        while session_id in connected_sessions:
            socketio.sleep(3)
        ''' Teardown analysis session '''
        try:
            client.disconnect()
        except Exception:
            pass
        corpus_analysis_clients.pop(session_id, None)
        corpus_analysis_sessions[corpus_id].remove(session_id)
        if not corpus_analysis_sessions[corpus_id]:
            corpus_analysis_sessions.pop(corpus_id, None)
            corpus.status = 'stop analysis'
            db.session.commit()