fix query sessions

uc-cdis · Jul 25, 2024 · bc90b49 · bc90b49
1 parent 2f972f5
commit bc90b49
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 45 deletions.
diff --git a/wts/blueprints/external_oidc.py b/wts/blueprints/external_oidc.py
@@ -124,23 +124,26 @@ def get_refresh_token_expirations(username, idps):
     Returns:
         dict: IdP to expiration of the most recent refresh token, or None if it's expired.
     """
-    now = int(time.time())
-    refresh_tokens = (
-        db.session.query(RefreshToken)
-        .filter_by(username=username)
-        .filter(RefreshToken.idp.in_(idps))
-        .order_by(RefreshToken.expires.asc())
-    )
-    if not refresh_tokens:
-        return {}
-    # the tokens are ordered by oldest to most recent, because we only want
-    # to return None if the most recent token is expired
-    expirations = {idp: None for idp in idps}
-    expirations.update(
-        {
-            t.idp: seconds_to_human_time(t.expires - now)
-            for t in refresh_tokens
-            if t.expires > now
-        }
-    )
-    return expirations
+    try:
+        now = int(time.time())
+        refresh_tokens = (
+            db.session.query(RefreshToken)
+            .filter_by(username=username)
+            .filter(RefreshToken.idp.in_(idps))
+            .order_by(RefreshToken.expires.asc())
+        )
+        if not refresh_tokens:
+            return {}
+        # the tokens are ordered by oldest to most recent, because we only want
+        # to return None if the most recent token is expired
+        expirations = {idp: None for idp in idps}
+        expirations.update(
+            {
+                t.idp: seconds_to_human_time(t.expires - now)
+                for t in refresh_tokens
+                if t.expires > now
+            }
+        )
+        return expirations
+    finally:
+        db.session.close()
diff --git a/wts/tokens.py b/wts/tokens.py
@@ -29,32 +29,41 @@ def get_data_for_fence_request(refresh_token):
 
 
 def get_access_token(requested_idp, expires=None):
-    if requested_idp not in flask.current_app.oauth2_clients:
-        raise UserError('Requested IdP "{}" is not configured'.format(requested_idp))
-    username = flask.g.user.username
-    flask.current_app.logger.info(
-        "Getting refresh token for user '{}', IdP '{}'".format(username, requested_idp)
-    )
-    refresh_token = (
-        db.session.query(RefreshToken)
-        .filter_by(username=username)
-        .filter_by(idp=requested_idp)
-        .order_by(RefreshToken.expires.desc())
-        .first()
-    )
-    now = int(time.time())
-    if not refresh_token:
-        raise AuthError("User doesn't have a refresh token")
-    if refresh_token.expires <= now:
-        raise AuthError("your refresh token is expired, please login again")
-    url, data, auth = get_data_for_fence_request(refresh_token)
     try:
-        r = httpx.post(url, data=data, auth=auth)
-    except Exception:
-        raise InternalError("Fail to reach fence")
-    if r.status_code != 200:
-        raise InternalError("Fail to get a access token from fence: {}".format(r.text))
-    return r.json()["access_token"]
+        if requested_idp not in flask.current_app.oauth2_clients:
+            raise UserError(
+                'Requested IdP "{}" is not configured'.format(requested_idp)
+            )
+        username = flask.g.user.username
+        flask.current_app.logger.info(
+            "Getting refresh token for user '{}', IdP '{}'".format(
+                username, requested_idp
+            )
+        )
+        refresh_token = (
+            db.session.query(RefreshToken)
+            .filter_by(username=username)
+            .filter_by(idp=requested_idp)
+            .order_by(RefreshToken.expires.desc())
+            .first()
+        )
+        now = int(time.time())
+        if not refresh_token:
+            raise AuthError("User doesn't have a refresh token")
+        if refresh_token.expires <= now:
+            raise AuthError("your refresh token is expired, please login again")
+        url, data, auth = get_data_for_fence_request(refresh_token)
+        try:
+            r = httpx.post(url, data=data, auth=auth)
+        except Exception:
+            raise InternalError("Fail to reach fence")
+        if r.status_code != 200:
+            raise InternalError(
+                "Fail to get a access token from fence: {}".format(r.text)
+            )
+        return r.json()["access_token"]
+    finally:
+        db.session.close()
 
 
 async def async_get_access_token(refresh_token, commons_hostname=None):