From 0cd4be055ee56109898f4fe9c111412ed2ec33f2 Mon Sep 17 00:00:00 2001
From: Patrick Jentsch <p.jentsch@uni-bielefeld.de>
Date: Mon, 18 Nov 2019 11:08:33 +0100
Subject: [PATCH] Change corpus analysis socket behavior

---
 app/corpora/events.py | 48 +++++++++++++++++++++++--------------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/app/corpora/events.py b/app/corpora/events.py
index c25aadbe..88217dac 100644
--- a/app/corpora/events.py
+++ b/app/corpora/events.py
@@ -26,23 +26,36 @@ def init_corpus_analysis(corpus_id):
     if corpus is None:
         socketio.emit('init_corpus_analysis', '[ERROR 404]: Not Found',
                       room=request.sid)
-    if not (corpus.creator == current_user or current_user.is_administrator()):
+    elif not (corpus.creator == current_user
+              or current_user.is_administrator()):
         socketio.emit('init_corpus_analysis', '[ERROR 403]: Forbidden',
                       room=request.sid)
-    if corpus_id not in analysis_sessions:
-        analysis_sessions[corpus_id] = [request.sid]
     else:
-        analysis_sessions[corpus_id].append(request.sid)
-    socketio.start_background_task(observe_corpus_analysis_connection,
-                                   current_app._get_current_object(),
-                                   corpus_id, request.sid)
+        if corpus_id not in analysis_sessions:
+            analysis_sessions[corpus_id] = [request.sid]
+        else:
+            analysis_sessions[corpus_id].append(request.sid)
+        while corpus.status != 'analysing':
+            db.session.refresh(corpus)
+            socketio.sleep(3)
+        analysis_clients[request.sid] = CQiClient(
+            host='{}_analysis_container{}'.format(corpus.creator.username,
+                                                  corpus.id))
+        analysis_clients[request.sid].ctrl_connect('opaque', 'opaque')
+        socketio.emit('init_corpus_analysis', 'Ready', room=request.sid)
+        socketio.start_background_task(observe_corpus_analysis_connection,
+                                       current_app._get_current_object(),
+                                       corpus_id, request.sid)
 
 
 @socketio.on('query')
 @login_required
 def recv_query(message):
-    logger.warning(message)
-    analysis_client = analysis_clients[request.sid]
+    analysis_client = analysis_clients.get(request.sid)
+    if analysis_client is None:
+        socketio.emit('query', '[ERROR 424]: Failed Dependency',
+                      room=request.sid)
+        return
     """ Prepare and execute a query """
     corpus = 'CORPUS'
     query = message['query']
@@ -85,23 +98,14 @@ def recv_query(message):
 
 def observe_corpus_analysis_connection(app, corpus_id, session_id):
     with app.app_context():
-        corpus = Corpus.query.get(corpus_id)
-        if corpus is None:
-            ''' TODO: Handle this '''
-            return
-        while corpus.status != 'analysing':
-            db.session.refresh(corpus)
-            socketio.sleep(3)
-        analysis_client = CQiClient(host='{}_analysis_container{}'.format(corpus.creator.username, corpus.id), port=4877)
-        analysis_client.ctrl_connect('opaque', 'opaque')
-        analysis_clients[session_id] = analysis_client
-        socketio.emit('init_corpus_analysis', 'Ready', room=session_id)
         while session_id in connected_sessions:
             socketio.sleep(3)
-        analysis_client.ctrl_bye()
-        analysis_clients.pop(session_id, None)
+        analysis_client = analysis_clients.pop(session_id, None)
+        if analysis_client is not None:
+            analysis_client.ctrl_bye()
         analysis_sessions[corpus_id].remove(session_id)
         if not analysis_sessions[corpus_id]:
             analysis_sessions.pop(corpus_id, None)
+            corpus = Corpus.query.get(corpus_id)
             corpus.status = 'stop analysis'
             db.session.commit()