Skip to content

Commit ea94470

Browse files
razvanrazvan
committed
refactor opa authorizer to cache resolved roles
1 parent a4e345f commit ea94470

File tree

4 files changed

+197
-128
lines changed

4 files changed

+197
-128
lines changed

‎superset/stackable/opa-authorizer/opa_authorizer/opa_manager.py‎

Lines changed: 125 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,17 @@
33
"""
44

55
import logging
6+
from typing import Optional
67

7-
from typing import List, Optional
8-
from cachetools import cachedmethod, TTLCache
9-
from flask import current_app, g
10-
from flask_appbuilder.security.sqla.models import (
11-
Role,
12-
User,
13-
)
148
import requests
9+
from cachetools import TTLCache, cachedmethod
10+
from flask import current_app, g
11+
from flask_appbuilder.security.sqla.models import Role, User
12+
from overrides import override
1513
from superset.security import SupersetSecurityManager
1614

15+
log = logging.getLogger(__name__)
16+
1717

1818
class OpaSupersetSecurityManager(SupersetSecurityManager):
1919
"""
@@ -29,20 +29,39 @@ class OpaSupersetSecurityManager(SupersetSecurityManager):
2929
AUTH_USER_REGISTRATION_ROLE_DEFAULT = "Public"
3030

3131
def __init__(self, appbuilder):
32-
self.appbuilder = appbuilder
33-
super().__init__(self.appbuilder)
34-
config = self.appbuilder.get_app.config
35-
self.opa_cache = self.opa_cache = TTLCache(
32+
super().__init__(appbuilder)
33+
34+
config = appbuilder.get_app.config
35+
36+
self.role_cache: TTLCache[str, set[Role]] = TTLCache(
3637
maxsize=config.get(
3738
"AUTH_OPA_CACHE_MAXSIZE", self.AUTH_OPA_CACHE_MAXSIZE_DEFAULT
3839
),
3940
ttl=config.get(
4041
"AUTH_OPA_CACHE_TTL_IN_SEC", self.AUTH_OPA_CACHE_TTL_IN_SEC_DEFAULT
4142
),
4243
)
44+
45+
self.auth_opa_url: str = config.get(
46+
"AUTH_OPA_REQUEST_URL", self.AUTH_OPA_REQUEST_URL_DEFAULT
47+
)
48+
self.auth_opa_package: str = config.get(
49+
"AUTH_OPA_PACKAGE", self.AUTH_OPA_PACKAGE_DEFAULT
50+
)
51+
self.auth_opa_rule: str = config.get(
52+
"AUTH_OPA_RULE", self.AUTH_OPA_RULE_DEFAULT
53+
)
54+
self.auth_opa_request_timeout: int = current_app.config.get(
55+
"AUTH_OPA_REQUEST_TIMEOUT", self.AUTH_OPA_REQUEST_TIMEOUT_DEFAULT
56+
)
57+
# Cannot name this "auth_auth_user_registration_role" because it clashes with some super() property constraints
58+
self.user_registration_role: str = config.get(
59+
"AUTH_USER_REGISTRATION_ROLE", self.AUTH_USER_REGISTRATION_ROLE_DEFAULT
60+
)
4361
self.opa_session = requests.Session()
4462

45-
def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
63+
@override
64+
def get_user_roles(self, user: Optional[User] = None) -> list[Role]:
4665
"""
4766
Retrieves a user's roles from an Open Policy Agent instance updating the
4867
user-role mapping in Superset's database in the process.
@@ -52,96 +71,132 @@ def get_user_roles(self, user: Optional[User] = None) -> List[Role]:
5271
if not user:
5372
user = g.user
5473

55-
default_role = self.resolve_role(
56-
current_app.config.get(
57-
"AUTH_USER_REGISTRATION_ROLE", self.AUTH_USER_REGISTRATION_ROLE_DEFAULT
58-
)
59-
)
60-
61-
opa_role_names = self.get_opa_user_roles(user.username)
62-
logging.debug(
63-
"OPA returned roles for user %s: %s", user.get_full_name(), opa_role_names
64-
)
74+
resolved_opa_roles = self.resolve_opa_roles(user.username)
6575

66-
opa_roles = set(map(self.resolve_role, opa_role_names))
67-
logging.debug(
68-
"Resolved OPA Roles in database for user %s: %s",
69-
user.get_full_name(),
70-
opa_roles,
71-
)
72-
# Ensure that in case of a bad or no response from OPA each user will have
73-
# at least one role.
74-
if opa_roles == {None} or opa_roles == set():
75-
opa_roles = {default_role}
76-
77-
if set(user.roles) != opa_roles:
78-
logging.debug(
79-
"Found a diff in roles for user %s: %s (Superset) and %s (OPA).",
80-
user.get_full_name(),
81-
user.roles,
82-
opa_roles,
76+
if not resolved_opa_roles:
77+
log.error(
78+
f"No OPA roles for user [{user.username}], defaulting to role AUTH_USER_REGISTRATION_ROLE"
79+
)
80+
default_role = self.resolve_role(self.user_registration_role)
81+
if not default_role:
82+
log.error(
83+
f"Failed to resolve default role name {self.user_registration_role} for user [{user.username}]. User will have no roles."
84+
)
85+
return []
86+
else:
87+
log.info(
88+
f"User [{user.username}] will only have default role [{self.user_registration_role}]"
89+
)
90+
return [default_role]
91+
92+
user_role_set = set(user.roles)
93+
log.debug(f"Superset roles for user [{user.username}]: {user_role_set}")
94+
95+
if user_role_set != resolved_opa_roles:
96+
log.info(
97+
f"Superset and OPA roles diverge. Updating OPA roles for user [{user.username}]"
8398
)
84-
user.roles = list(opa_roles)
85-
self.update_user(user)
99+
user.roles = list(resolved_opa_roles)
100+
if not self.update_user(user):
101+
log.error(f"Failed to update user roles for user {user.username}")
86102

87103
return user.roles
88104

89-
@cachedmethod(lambda self: self.opa_cache)
90105
def get_opa_user_roles(self, username: str) -> list[str]:
91106
"""
92107
Queries an Open Policy Agent instance for the roles of a given user.
93108
94-
:returns: A list of role names or an empty list if an exception during
95-
the connection to OPA is encountered or if OPA didn't return a list.
109+
:returns: A list of Role objects assigned to the user or an empty list.
96110
"""
97-
98-
opa_url = current_app.config.get(
99-
"AUTH_OPA_REQUEST_URL", self.AUTH_OPA_REQUEST_URL_DEFAULT
100-
)
101-
package = current_app.config.get(
102-
"AUTH_OPA_PACKAGE", self.AUTH_OPA_PACKAGE_DEFAULT
103-
)
104-
rule = current_app.config.get("AUTH_OPA_RULE", self.AUTH_OPA_RULE_DEFAULT)
105-
timeout = current_app.config.get(
106-
"AUTH_OPA_REQUEST_TIMEOUT", self.AUTH_OPA_REQUEST_TIMEOUT_DEFAULT
107-
)
108111
input = {"input": {"username": username}}
109112
try:
113+
req_url = f"{self.auth_opa_url}/v1/data/{self.auth_opa_package}/{self.auth_opa_rule}"
110114
response = self.call_opa(
111-
url=f"{opa_url}/v1/data/{package}/{rule}",
115+
url=req_url,
112116
json=input,
113-
timeout=timeout,
117+
timeout=self.auth_opa_request_timeout,
114118
)
115119

116-
if response.status_code is None or response.status_code != 200:
117-
logging.error("OPA returned status code %s", response.status_code)
120+
if response.status_code != 200:
121+
log.error(
122+
f"OPA request [{req_url}] for user [{username}] failed with response code [{response.status_code}]"
123+
)
118124
return []
119125

120-
roles = response.json().get("result")
121-
if roles is None or type(roles).__name__ != "list":
122-
logging.error("The OPA query didn't return a list: %s", response.json())
126+
roles: list[str] = response.json().get("result")
127+
if type(roles) is not list:
128+
log.error(f"Expected a list or role names from OPA but got [{roles}]")
123129
return []
124130

131+
if not roles:
132+
log.error(
133+
"Expected a list or role names from OPA but got an empty list"
134+
)
135+
return []
136+
137+
log.debug(f"Retrieved OPA role names for user [{username}]: [{roles}]")
125138
return roles
139+
126140
except Exception as e:
127-
logging.error("Request to OPA failed", exc_info=e)
141+
log.error("Failed to get OPA role names", exc_info=e)
128142
return []
129143

144+
@cachedmethod(lambda self: self.role_cache)
145+
def resolve_opa_roles(self, username: str) -> set[Role]:
146+
"""
147+
Queries an Open Policy Agent instance for the roles of a given user.
148+
149+
Then maps the OPA role names to Superset Role objects.
150+
151+
The result is cached.
152+
153+
:returns: A set of Role objects assigned to the user or an empty set.
154+
"""
155+
try:
156+
roles = self.get_opa_user_roles(username)
157+
resolved_opa_roles = self.resolve_role_list(roles)
158+
159+
log.debug(
160+
f"Resolved OPA Roles for user [{username}]: [{resolved_opa_roles}]"
161+
)
162+
163+
return resolved_opa_roles
164+
165+
except Exception as e:
166+
log.error("Failed to resolve OPA roles", exc_info=e)
167+
return set()
168+
130169
def call_opa(self, url: str, json: dict, timeout: int) -> requests.Response:
131170
return self.opa_session.post(
132171
url=url,
133172
json=json,
134173
timeout=timeout,
135174
)
136175

137-
def resolve_role(self, role_name: str) -> Role:
176+
def resolve_role(self, role_name: str) -> Optional[Role]:
138177
"""
139178
Finds a role by name creating it if it doesn't exist in Superset yet.
140179
141-
:returns: A role.
180+
:returns: A role or None.
142181
"""
143-
role = self.find_role(role_name)
144-
if role is None:
145-
logging.debug("Creating role %s as it doesn't already exist.", role_name)
146-
self.add_role(role_name)
147-
return self.find_role(role_name)
182+
role: Optional[Role] = self.find_role(role_name)
183+
if not role:
184+
log.info(
185+
f"Failed to resolve role {role_name}, attempting to create it with no permissions"
186+
)
187+
role = self.add_role(role_name)
188+
if not role:
189+
log.error(f"Failed to create new role {role_name}")
190+
else:
191+
log.info(f"Resolved role name [{role_name}] to new role.")
192+
193+
log.debug(f"Resolved role name [{role_name}] to existing role.")
194+
195+
return role
196+
197+
def resolve_role_list(self, roles: list[str]) -> set[Role]:
198+
result: set[Role] = set()
199+
for role_name in roles:
200+
if resolved_role := self.resolve_role(role_name):
201+
result.add(resolved_role)
202+
return result

0 commit comments

Comments
 (0)