from flask import current_app, request from flask_login import current_user from .cqi import CQiClient from .. import db, socketio from ..decorators import socketio_login_required from ..events import connected_sessions from ..models import Corpus, User import math ''' ' A dictionary containing lists of, with corpus ids associated, Socket.IO ' session ids (sid). {: [, ...], ...} ''' pj_corpus_analysis_sessions = {} ''' ' A dictionary containing Socket.IO session id - CQi client pairs. ' {: CQiClient, ...} ''' pj_corpus_analysis_clients = {} @socketio.on('pj_corpus_analysis_init') @socketio_login_required def pj_init_corpus_analysis(corpus_id): socketio.start_background_task(pj_corpus_analysis_session_handler, current_app._get_current_object(), corpus_id, current_user.id, request.sid) @socketio.on('pj_corpus_analysis_query') @socketio_login_required def pj_corpus_analysis_query(query): client = pj_corpus_analysis_clients.get(request.sid) if client is None: response = {'code': 404, 'msg': 'Failed Dependency'} socketio.emit('pj_corpus_analysis_query', response, room=request.sid) return corpus = client.corpora.get('CORPUS') try: results = corpus.query(query) except Exception as e: response = {'code': 1, 'msg': str(e)} socketio.emit('pj_corpus_analysis_query', response, room=request.sid) else: response = {'code': 200, 'msg': 'OK', 'data': {'num_matches': results.size}} socketio.emit('pj_corpus_analysis_query', response, room=request.sid) chunk_size = 100 chunk_start = 0 context = 100 progress = 0 while chunk_start <= results.size: chunk = results.export(context=context, offset=chunk_start, cutoff=chunk_size) if (results.size == 0): progress = 100 else: progress = ((chunk_start + chunk_size) / results.size) * 100 progress = min(100, int(math.ceil(progress))) socketio.emit('pj_corpus_analysis_query_results', {'chunk': chunk, 'progress': progress}, room=request.sid) chunk_start += chunk_size chunk_size = 250 def pj_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('pj_corpus_analysis_init', response, room=session_id) return elif not (corpus.creator == user or user.is_administrator()): response = {'code': 403, 'msg': 'Forbidden'} socketio.emit('pj_corpus_analysis_init', response, room=session_id) return while corpus.status != 'analysing': db.session.refresh(corpus) socketio.sleep(3) client = CQiClient('pj_corpus_{}_analysis'.format(corpus_id)) try: client.connect() except Exception: response = {'code': 500, 'msg': 'Internal Server Error'} socketio.emit('pj_corpus_analysis_init', response, room=session_id) return pj_corpus_analysis_clients[session_id] = client if corpus_id not in pj_corpus_analysis_sessions: pj_corpus_analysis_sessions[corpus_id] = [session_id] else: pj_corpus_analysis_sessions[corpus_id].append(session_id) response = {'code': 200, 'msg': 'OK'} socketio.emit('pj_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 pj_corpus_analysis_clients.pop(session_id, None) pj_corpus_analysis_sessions[corpus_id].remove(session_id) if not pj_corpus_analysis_sessions[corpus_id]: pj_corpus_analysis_sessions.pop(corpus_id, None) corpus.status = 'stop analysis' db.session.commit()