1 files changed, 9 insertions, 7 deletions
diff --git a/gn_auth/auth/authorisation/users/masquerade/views.py b/gn_auth/auth/authorisation/users/masquerade/views.py
index 8b897f2..12a8c97 100644
--- a/gn_auth/auth/authorisation/users/masquerade/views.py
+++ b/gn_auth/auth/authorisation/users/masquerade/views.py
@@ -1,14 +1,14 @@
 """Endpoints for user masquerade"""
 from dataclasses import asdict
 from uuid import UUID
-from functools import partial
 
-from flask import request, jsonify, Response, Blueprint
+from flask import request, jsonify, Response, Blueprint, current_app
 
 from gn_auth.auth.errors import InvalidData
+from gn_auth.auth.authorisation.resources.groups.models import user_group
 
+from ....db import sqlite3 as db
 from ...checks import require_json
-from ....db.sqlite3 import with_db_connection
 from ....authentication.users import user_by_id
 from ....authentication.oauth2.resource_server import require_oauth
 
@@ -21,13 +21,13 @@ masq = Blueprint("masquerade", __name__)
 @require_json
 def masquerade() -> Response:
     """Masquerade as a particular user."""
-    with require_oauth.acquire("profile user masquerade") as token:
+    with (require_oauth.acquire("profile user masquerade") as token,
+          db.connection(current_app.config["AUTH_DB"]) as conn):
         masqueradee_id = UUID(request.json["masquerade_as"])#type: ignore[index]
         if masqueradee_id == token.user.user_id:
             raise InvalidData("You are not allowed to masquerade as yourself.")
 
-        masq_user = with_db_connection(partial(
-            user_by_id, user_id=masqueradee_id))
+        masq_user = user_by_id(conn, user_id=masqueradee_id)
 
         def __masq__(conn):
             new_token = masquerade_as(conn, original_token=token, masqueradee=masq_user)
@@ -39,6 +39,8 @@ def masquerade() -> Response:
             },
             "masquerade_as": {
                 "user": asdict(masq_user),
-                "token": with_db_connection(__masq__)
+                "token": __masq__(conn),
+                **(user_group(conn, masq_user).maybe(# type: ignore[misc]
+                    {}, lambda grp: {"group": grp}))
             }
         })