1 files changed, 33 insertions, 1 deletions
diff --git a/server/user_session.py b/server/user_session.py
index bb4179d..d8cdf1d 100644
--- a/server/user_session.py
+++ b/server/user_session.py
@@ -9,6 +9,7 @@ from . import prolog_session
 from . import problems
 import db
 from errors.session import NoSuchSession, AuthenticationFailed
+import psycopg2.extras
 
 __all__ = ['get_session_by_id', 'get_or_create_session', 'UserSession']
 
@@ -49,7 +50,7 @@ class UserSession(object):
     def get_prolog(self):
         with self._access_lock:
             if self.prolog_session is None:
-                self.prolog_session = prolog_session.PrologSession()  # lazy init
+                self.prolog_session = prolog_session.PrologSession(self)  # lazy init
             return self.prolog_session
 
     def get_problem_data(self, language, problem_group, problem):
@@ -79,6 +80,37 @@ class UserSession(object):
             conn.commit()
             db.return_connection(conn)
 
+    def update_solution(self, problem_id, trace, solution):
+        if (trace is None) and (solution is None):
+            return
+        conn = db.get_connection()
+        try:
+            cur = conn.cursor()
+            try:
+                # TODO: convert to upsert with postgresql 9.5 to eliminate the small window where it's possible for more than one concurrent insert to execute
+                cur.execute('select id, trace, solution content from solution where codeq_user_id = %s and problem_id = %s for update', (self.uid, problem_id))
+                row = cur.fetchone()
+                if row:
+                    if row[1]:
+                        new_trace = row[1]
+                        if trace:
+                            new_trace.extend(trace)
+                    else:
+                        new_trace = trace
+                    new_solution = row[2] if solution is None else solution
+                    cur.execute('update solution set content = %s, trace = %s where id = %s', (new_solution, psycopg2.extras.Json(new_trace), row[0]))
+                else:
+                    # this is the first entry
+                    cur.execute('insert into solution (done, content, problem_id, codeq_user_id, trace) values (%s, %s, %s, %s, %s)', (False, solution, problem_id, self.uid, psycopg2.extras.Json(trace)))
+            finally:
+                cur.close()
+            conn.commit()
+        except:
+            conn.rollback()
+            raise
+        finally:
+            db.return_connection(conn)
+
     def __del__(self):
         #  no locking needed if GC is removing us, as there cannot be any concurrent access by definition
         if hasattr(self, 'prolog_session') and (self.prolog_session is not None):