Skip to content

Commit 5e2d02b

Browse files
chore(airflow): Reformat Python code
1 parent 0884f38 commit 5e2d02b

File tree

2 files changed

+400
-398
lines changed

2 files changed

+400
-398
lines changed

airflow/opa-auth-manager/opa_auth_manager/opa_fab_auth_manager.py

Lines changed: 124 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@
2222
import json
2323
import 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

2728
class 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+
4546
class 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+
5961
class 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

Comments
 (0)