2222import json
2323import requests
2424
25- METRIC_NAME_OPA_CACHE_LIMIT_REACHED = "opa_cache_limit_reached"
25+ METRIC_NAME_OPA_CACHE_LIMIT_REACHED = "opa_cache_limit_reached"
26+
2627
2728class OpaInput :
2829 """
@@ -33,15 +34,15 @@ def __init__(self, input: dict) -> None:
3334 self .input = input
3435
3536 def __eq__ (self , other : object ) -> bool :
36- return isinstance (other , OpaInput ) \
37- and self .input == other .input
37+ return isinstance (other , OpaInput ) and self .input == other .input
3838
3939 def __hash__ (self ) -> int :
4040 return hash (json .dumps (self .input , sort_keys = True ))
4141
4242 def to_dict (self ) -> dict :
4343 return self .input
4444
45+
4546class Cache (TTLCache ):
4647 """
4748 LRU Cache implementation with per-item time-to-live (TTL) value.
@@ -56,16 +57,17 @@ def popitem(self):
5657 Stats .incr (METRIC_NAME_OPA_CACHE_LIMIT_REACHED )
5758 return super ().popitem ()
5859
60+
5961class OpaFabAuthManager (FabAuthManager , LoggingMixin ):
6062 """
6163 Auth manager based on the FabAuthManager which delegates the authorization to an Open Policy
6264 Agent
6365 """
6466
65- AUTH_OPA_CACHE_MAXSIZE_DEFAULT = 1000
66- AUTH_OPA_CACHE_TTL_IN_SEC_DEFAULT = 30
67- AUTH_OPA_REQUEST_URL_DEFAULT = ' http://opa:8081/v1/data/airflow'
68- AUTH_OPA_REQUEST_TIMEOUT_DEFAULT = 10
67+ AUTH_OPA_CACHE_MAXSIZE_DEFAULT = 1000
68+ AUTH_OPA_CACHE_TTL_IN_SEC_DEFAULT = 30
69+ AUTH_OPA_REQUEST_URL_DEFAULT = " http://opa:8081/v1/data/airflow"
70+ AUTH_OPA_REQUEST_TIMEOUT_DEFAULT = 10
6971
7072 def init (self ) -> None :
7173 """
@@ -79,12 +81,10 @@ def init(self) -> None:
7981 config = self .appbuilder .get_app .config
8082 self .opa_cache = Cache (
8183 maxsize = config .get (
82- 'AUTH_OPA_CACHE_MAXSIZE' ,
83- self .AUTH_OPA_CACHE_MAXSIZE_DEFAULT
84+ "AUTH_OPA_CACHE_MAXSIZE" , self .AUTH_OPA_CACHE_MAXSIZE_DEFAULT
8485 ),
8586 ttl = config .get (
86- 'AUTH_OPA_CACHE_TTL_IN_SEC' ,
87- self .AUTH_OPA_CACHE_TTL_IN_SEC_DEFAULT
87+ "AUTH_OPA_CACHE_TTL_IN_SEC" , self .AUTH_OPA_CACHE_TTL_IN_SEC_DEFAULT
8888 ),
8989 )
9090 self .opa_session = requests .Session ()
@@ -113,20 +113,16 @@ def _is_authorized_in_opa(self, endpoint: str, input: OpaInput) -> bool:
113113 self .log .debug ("Forward authorization request to OPA" )
114114
115115 config = self .appbuilder .get_app .config
116- opa_url = config .get (
117- 'AUTH_OPA_REQUEST_URL' ,
118- self .AUTH_OPA_REQUEST_URL_DEFAULT
119- )
116+ opa_url = config .get ("AUTH_OPA_REQUEST_URL" , self .AUTH_OPA_REQUEST_URL_DEFAULT )
120117 try :
121118 response = self .call_opa (
122- f' { opa_url } /{ endpoint } ' ,
119+ f" { opa_url } /{ endpoint } " ,
123120 json = input .to_dict (),
124121 timeout = config .get (
125- 'AUTH_OPA_REQUEST_TIMEOUT' ,
126- self .AUTH_OPA_REQUEST_TIMEOUT_DEFAULT
127- )
122+ "AUTH_OPA_REQUEST_TIMEOUT" , self .AUTH_OPA_REQUEST_TIMEOUT_DEFAULT
123+ ),
128124 )
129- return response .json ().get (' result' )
125+ return response .json ().get (" result" )
130126 except Exception as e :
131127 self .log .error ("Request to OPA failed" , exc_info = e )
132128 return False
@@ -160,19 +156,21 @@ def is_authorized_configuration(
160156 section = details .section
161157
162158 return self ._is_authorized_in_opa (
163- 'is_authorized_configuration' ,
164- OpaInput ({
165- 'input' : {
166- 'method' : method ,
167- 'details' : {
168- 'section' : section ,
169- },
170- 'user' : {
171- 'id' : user .get_id (),
172- 'name' : user .get_name (),
173- },
159+ "is_authorized_configuration" ,
160+ OpaInput (
161+ {
162+ "input" : {
163+ "method" : method ,
164+ "details" : {
165+ "section" : section ,
166+ },
167+ "user" : {
168+ "id" : user .get_id (),
169+ "name" : user .get_name (),
170+ },
171+ }
174172 }
175- })
173+ ),
176174 )
177175
178176 @override
@@ -203,19 +201,21 @@ def is_authorized_connection(
203201 conn_id = details .conn_id
204202
205203 return self ._is_authorized_in_opa (
206- 'is_authorized_connection' ,
207- OpaInput ({
208- 'input' : {
209- 'method' : method ,
210- 'details' : {
211- 'conn_id' : conn_id ,
212- },
213- 'user' : {
214- 'id' : user .get_id (),
215- 'name' : user .get_name (),
216- },
204+ "is_authorized_connection" ,
205+ OpaInput (
206+ {
207+ "input" : {
208+ "method" : method ,
209+ "details" : {
210+ "conn_id" : conn_id ,
211+ },
212+ "user" : {
213+ "id" : user .get_id (),
214+ "name" : user .get_name (),
215+ },
216+ }
217217 }
218- })
218+ ),
219219 )
220220
221221 @override
@@ -254,20 +254,22 @@ def is_authorized_dag(
254254 dag_id = details .id
255255
256256 return self ._is_authorized_in_opa (
257- 'is_authorized_dag' ,
258- OpaInput ({
259- 'input' : {
260- 'method' : method ,
261- 'access_entity' : entity ,
262- 'details' : {
263- 'id' : dag_id ,
264- },
265- 'user' : {
266- 'id' : user .get_id (),
267- 'name' : user .get_name (),
268- },
257+ "is_authorized_dag" ,
258+ OpaInput (
259+ {
260+ "input" : {
261+ "method" : method ,
262+ "access_entity" : entity ,
263+ "details" : {
264+ "id" : dag_id ,
265+ },
266+ "user" : {
267+ "id" : user .get_id (),
268+ "name" : user .get_name (),
269+ },
270+ }
269271 }
270- })
272+ ),
271273 )
272274
273275 @override
@@ -298,19 +300,21 @@ def is_authorized_dataset(
298300 uri = details .uri
299301
300302 return self ._is_authorized_in_opa (
301- 'is_authorized_dataset' ,
302- OpaInput ({
303- 'input' : {
304- 'method' : method ,
305- 'details' : {
306- 'uri' : uri ,
307- },
308- 'user' : {
309- 'id' : user .get_id (),
310- 'name' : user .get_name (),
311- },
303+ "is_authorized_dataset" ,
304+ OpaInput (
305+ {
306+ "input" : {
307+ "method" : method ,
308+ "details" : {
309+ "uri" : uri ,
310+ },
311+ "user" : {
312+ "id" : user .get_id (),
313+ "name" : user .get_name (),
314+ },
315+ }
312316 }
313- })
317+ ),
314318 )
315319
316320 @override
@@ -341,19 +345,21 @@ def is_authorized_pool(
341345 name = details .name
342346
343347 return self ._is_authorized_in_opa (
344- 'is_authorized_pool' ,
345- OpaInput ({
346- 'input' : {
347- 'method' : method ,
348- 'details' : {
349- 'name' : name ,
350- },
351- 'user' : {
352- 'id' : user .get_id (),
353- 'name' : user .get_name (),
354- },
348+ "is_authorized_pool" ,
349+ OpaInput (
350+ {
351+ "input" : {
352+ "method" : method ,
353+ "details" : {
354+ "name" : name ,
355+ },
356+ "user" : {
357+ "id" : user .get_id (),
358+ "name" : user .get_name (),
359+ },
360+ }
355361 }
356- })
362+ ),
357363 )
358364
359365 @override
@@ -384,19 +390,21 @@ def is_authorized_variable(
384390 key = details .key
385391
386392 return self ._is_authorized_in_opa (
387- 'is_authorized_variable' ,
388- OpaInput ({
389- 'input' : {
390- 'method' : method ,
391- 'details' : {
392- 'key' : key ,
393- },
394- 'user' : {
395- 'id' : user .get_id (),
396- 'name' : user .get_name (),
397- },
393+ "is_authorized_variable" ,
394+ OpaInput (
395+ {
396+ "input" : {
397+ "method" : method ,
398+ "details" : {
399+ "key" : key ,
400+ },
401+ "user" : {
402+ "id" : user .get_id (),
403+ "name" : user .get_name (),
404+ },
405+ }
398406 }
399- })
407+ ),
400408 )
401409
402410 @override
@@ -420,16 +428,18 @@ def is_authorized_view(
420428 user = self .get_user ()
421429
422430 return self ._is_authorized_in_opa (
423- 'is_authorized_view' ,
424- OpaInput ({
425- 'input' : {
426- 'access_view' : access_view .name ,
427- 'user' : {
428- 'id' : user .get_id (),
429- 'name' : user .get_name (),
430- },
431+ "is_authorized_view" ,
432+ OpaInput (
433+ {
434+ "input" : {
435+ "access_view" : access_view .name ,
436+ "user" : {
437+ "id" : user .get_id (),
438+ "name" : user .get_name (),
439+ },
440+ }
431441 }
432- })
442+ ),
433443 )
434444
435445 @override
@@ -462,15 +472,17 @@ def is_authorized_custom_view(
462472 user = self .get_user ()
463473
464474 return self ._is_authorized_in_opa (
465- 'is_authorized_custom_view' ,
466- OpaInput ({
467- 'input' : {
468- 'method' : method ,
469- 'resource_name' : resource_name ,
470- 'user' : {
471- 'id' : user .get_id (),
472- 'name' : user .get_name (),
473- },
475+ "is_authorized_custom_view" ,
476+ OpaInput (
477+ {
478+ "input" : {
479+ "method" : method ,
480+ "resource_name" : resource_name ,
481+ "user" : {
482+ "id" : user .get_id (),
483+ "name" : user .get_name (),
484+ },
485+ }
474486 }
475- })
487+ ),
476488 )
0 commit comments