finish up db rollback checks

- move catch_db_error to utils
- tidy catch/propagate errors in prepare, get_current_user

jupyterhub/app.py

@@ -5,9 +5,6 @@
 import asyncio
 import atexit
 import binascii
-import functools
-import inspect
-import json
 import logging
 import os
 import re
@@ -94,6 +91,7 @@ from .pagination import Pagination
 from .proxy import Proxy, ConfigurableHTTPProxy
 from .traitlets import URLPrefix, Command, EntryPointType, Callable
 from .utils import (
+    catch_db_error,
     maybe_future,
     url_path_join,
     print_stacks,
@@ -1561,23 +1559,6 @@ class JupyterHub(Application):
         if os.path.exists(path) and not os.access(path, os.W_OK):
             self.log.error("%s cannot edit %s", user, path)
 
-    def catch_db_error(f):
-        """Catch and rollback database errors"""
-
-        @functools.wraps(f)
-        async def catching(self, *args, **kwargs):
-            try:
-                r = f(self, *args, **kwargs)
-                if inspect.isawaitable(r):
-                    r = await r
-            except SQLAlchemyError:
-                self.log.exception("Rolling back session due to database error")
-                self.db.rollback()
-            else:
-                return r
-
-        return catching
-
     def init_secrets(self):
         trait_name = 'cookie_secret'
         trait = self.traits()[trait_name]
@@ -2236,7 +2217,7 @@ class JupyterHub(Application):
         await self._add_tokens(self.service_tokens, kind='service')
         await self._add_tokens(self.api_tokens, kind='user')
 
-        self.purge_expired_tokens()
+        await self.purge_expired_tokens()
         # purge expired tokens hourly
         # we don't need to be prompt about this
         # because expired tokens cannot be used anyway

jupyterhub/handlers/base.py

@@ -85,12 +85,14 @@ class BaseHandler(RequestHandler):
         self.expanded_scopes = set()
         try:
             await self.get_current_user()
-        except SQLAlchemyError:
+        except Exception as e:
-            self.log.exception("Rolling back session due to database error")
+            # ensure get_current_user is never called again for this handler,
-            self.db.rollback()
+            # since it failed
-        except Exception:
-            self.log.exception("Failed to get current user")
             self._jupyterhub_user = None
+            self.log.exception("Failed to get current user")
+            if isinstance(e, SQLAlchemyError):
+                self.log.error("Rolling back session due to database error")
+                self.db.rollback()
         self._resolve_roles_and_scopes()
         return await maybe_future(super().prepare())
 
@@ -414,12 +416,11 @@ class BaseHandler(RequestHandler):
                 if user and isinstance(user, User):
                     user = await self.refresh_auth(user)
                 self._jupyterhub_user = user
-            except Exception as e:
+            except Exception:
-                if isinstance(e, SQLAlchemyError):
-                    raise SQLAlchemyError()
                 # don't let errors here raise more than once
                 self._jupyterhub_user = None
-                self.log.exception("Error getting current user")
+                # but still raise, which will get handled in .prepare()
+                raise
         return self._jupyterhub_user
 
     def _resolve_roles_and_scopes(self):

jupyterhub/utils.py

@@ -4,9 +4,9 @@
 import asyncio
 import concurrent.futures
 import errno
+import functools
 import hashlib
 import inspect
-import os
 import random
 import secrets
 import socket
@@ -22,6 +22,7 @@ from hmac import compare_digest
 from operator import itemgetter
 
 from async_generator import aclosing
+from sqlalchemy.exc import SQLAlchemyError
 from tornado import ioloop
 from tornado import web
 from tornado.httpclient import AsyncHTTPClient
@@ -635,3 +636,21 @@ def get_accepted_mimetype(accept_header, choices=None):
         else:
             return mime
     return None
+
+
+def catch_db_error(f):
+    """Catch and rollback database errors"""
+
+    @functools.wraps(f)
+    async def catching(self, *args, **kwargs):
+        try:
+            r = f(self, *args, **kwargs)
+            if inspect.isawaitable(r):
+                r = await r
+        except SQLAlchemyError:
+            self.log.exception("Rolling back session due to database error")
+            self.db.rollback()
+        else:
+            return r
+
+    return catching