server/db: fix closing DB sessions

Certain exception scenarios led to small disasters. Moved database
session management directly to router, since it's that sensitive.

1 files changed, 11 insertions, 5 deletions

diff --git a/server/szurubooru/rest/app.py b/server/szurubooru/rest/app.py
index eb163f0..0823e6e 100644
--- a/server/szurubooru/rest/app.py
+++ b/server/szurubooru/rest/app.py
@@ -3,6 +3,7 @@ import cgi
 import json
 import re
 from datetime import datetime
+from szurubooru import db
 from szurubooru.func import util
 from szurubooru.rest import errors, middleware, routes, context
 
@@ -86,13 +87,18 @@ def application(env, start_response):
                 'Requested path ' + ctx.url + ' was not found.')
 
         try:
-            for hook in middleware.pre_hooks:
-                hook(ctx)
+            ctx.session = db.session()
             try:
-                response = handler(ctx, match.groupdict())
-            finally:
-                for hook in middleware.post_hooks:
+                for hook in middleware.pre_hooks:
                     hook(ctx)
+                try:
+                    response = handler(ctx, match.groupdict())
+                finally:
+                    for hook in middleware.post_hooks:
+                        hook(ctx)
+            finally:
+                db.session.remove()
+
             start_response('200', [('content-type', 'application/json')])
             return (_dump_json(response).encode('utf-8'),)