Skip to content

Commit e2f1950

Browse files
potiukpotiuk
authored
Upgrade FAB to 4.1.1 (#24399)
* Upgrade FAB to 4.1.1 The Flask Application Builder have been updated recently to support a number of newer dependencies. This PR is the attempt to migrate FAB to newer version. This includes: * update setup.py and setup.cfg upper and lower bounds to account for proper version of dependencies that FAB < 4.0.0 was blocking from upgrade * added typed Flask application retrieval with a custom application fields available for MyPy typing checks. * fix typing to account for typing hints added in multiple upgraded libraries optional values and content of request returned as Mapping * switch to PyJWT 2.* by using non-deprecated "required" claim as list rather than separate fields * add possibiliyt to install providers without constraints so that we could avoid errors on conflicting constraints when upgrade-to-newer-dependencies is used * add pre-commit to check that 2.4+ only get_airflow_app is not used in providers * avoid Bad Request in case the request sent to Flask 2.0 is not JSon content type * switch imports of internal classes to direct packages where classes are available rather than from "airflow.models" to satisfy MyPY * synchronize changes of FAB Security Manager 4.1.1 with our copy of the Security Manager. * add error handling for a few "None" cases detected by MyPY * corrected test cases that were broken by immutability of Flask 2 objects and better escaping done by Flask 2 * updated test cases to account for redirection to "path" rather than full URL by Flask2 Fixes: #22397 * fixup! Upgrade FAB to 4.1.1
1 parent c3c1f7e commit e2f1950

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

59 files changed

+762
-561
lines changed

‎.github/workflows/ci.yml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -805,6 +805,8 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
805805
run: >
806806
breeze verify-provider-packages --use-airflow-version wheel --use-packages-from-dist
807807
--package-format wheel
808+
env:
809+
SKIP_CONSTRAINTS: "${{ needs.build-info.outputs.upgradeToNewerDependencies }}"
808810
- name: "Remove airflow package and replace providers with 2.2-compliant versions"
809811
run: |
810812
rm -vf dist/apache_airflow-*.whl \
@@ -882,6 +884,8 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
882884
run: >
883885
breeze verify-provider-packages --use-airflow-version sdist --use-packages-from-dist
884886
--package-format sdist
887+
env:
888+
SKIP_CONSTRAINTS: "${{ needs.build-info.outputs.upgradeToNewerDependencies }}"
885889
- name: "Fix ownership"
886890
run: breeze fix-ownership
887891
if: always()

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -347,7 +347,7 @@ repos:
347347
language: python
348348
files: ^BREEZE\.rst$|^dev/breeze/.*$
349349
pass_filenames: false
350-
additional_dependencies: ['rich>=12.4.4', 'rich-click']
350+
additional_dependencies: ['rich>=12.4.4', 'rich-click>=1.5']
351351
- id: update-local-yml-file
352352
name: Update mounts in the local yml file
353353
entry: ./scripts/ci/pre_commit/pre_commit_local_yml_mounts.py

‎Dockerfile.ci‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -686,29 +686,47 @@ if [[ ${SKIP_ENVIRONMENT_INITIALIZATION=} != "true" ]]; then
686686
echo "${COLOR_BLUE}Uninstalling airflow and providers"
687687
echo
688688
uninstall_airflow_and_providers
689-
echo "${COLOR_BLUE}Install airflow from wheel package with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
690-
echo
691-
install_airflow_from_wheel "${AIRFLOW_EXTRAS}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
689+
if [[ ${SKIP_CONSTRAINTS,,=} == "true" ]]; then
690+
echo "${COLOR_BLUE}Install airflow from wheel package with extras: '${AIRFLOW_EXTRAS}' with no constraints.${COLOR_RESET}"
691+
echo
692+
install_airflow_from_wheel "${AIRFLOW_EXTRAS}" "none"
693+
else
694+
echo "${COLOR_BLUE}Install airflow from wheel package with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
695+
echo
696+
install_airflow_from_wheel "${AIRFLOW_EXTRAS}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
697+
fi
692698
uninstall_providers
693699
elif [[ ${USE_AIRFLOW_VERSION} == "sdist" ]]; then
694700
echo
695701
echo "${COLOR_BLUE}Uninstalling airflow and providers"
696702
echo
697703
uninstall_airflow_and_providers
698704
echo
699-
echo "${COLOR_BLUE}Install airflow from sdist package with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
700-
echo
701-
install_airflow_from_sdist "${AIRFLOW_EXTRAS}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
705+
if [[ ${SKIP_CONSTRAINTS,,=} == "true" ]]; then
706+
echo "${COLOR_BLUE}Install airflow from sdist package with extras: '${AIRFLOW_EXTRAS}' with no constraints.${COLOR_RESET}"
707+
echo
708+
install_airflow_from_sdist "${AIRFLOW_EXTRAS}" "none"
709+
else
710+
echo "${COLOR_BLUE}Install airflow from sdist package with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
711+
echo
712+
install_airflow_from_sdist "${AIRFLOW_EXTRAS}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
713+
fi
702714
uninstall_providers
703715
else
704716
echo
705717
echo "${COLOR_BLUE}Uninstalling airflow and providers"
706718
echo
707719
uninstall_airflow_and_providers
708720
echo
709-
echo "${COLOR_BLUE}Install released airflow from PyPI with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
710-
echo
711-
install_released_airflow_version "${USE_AIRFLOW_VERSION}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
721+
if [[ ${SKIP_CONSTRAINTS,,=} == "true" ]]; then
722+
echo "${COLOR_BLUE}Install released airflow from PyPI with extras: '${AIRFLOW_EXTRAS}' with no constraints.${COLOR_RESET}"
723+
echo
724+
install_released_airflow_version "${USE_AIRFLOW_VERSION}" "none"
725+
else
726+
echo "${COLOR_BLUE}Install released airflow from PyPI with extras: '${AIRFLOW_EXTRAS}' and constraints reference ${AIRFLOW_CONSTRAINTS_REFERENCE}.${COLOR_RESET}"
727+
echo
728+
install_released_airflow_version "${USE_AIRFLOW_VERSION}" "${AIRFLOW_CONSTRAINTS_REFERENCE}"
729+
fi
712730
fi
713731
if [[ ${USE_PACKAGES_FROM_DIST=} == "true" ]]; then
714732
echo

‎airflow/api/auth/backend/basic_auth.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,11 @@
1818
from functools import wraps
1919
from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast
2020

21-
from flask import Response, current_app, request
21+
from flask import Response, request
2222
from flask_appbuilder.const import AUTH_LDAP
2323
from flask_login import login_user
2424

25+
from airflow.utils.airflow_flask_app import get_airflow_app
2526
from airflow.www.fab_security.sqla.models import User
2627

2728
CLIENT_AUTH: Optional[Union[Tuple[str, str], Any]] = None
@@ -40,7 +41,7 @@ def auth_current_user() -> Optional[User]:
4041
if auth is None or not auth.username or not auth.password:
4142
return None
4243

43-
ab_security_manager = current_app.appbuilder.sm
44+
ab_security_manager = get_airflow_app().appbuilder.sm
4445
user = None
4546
if ab_security_manager.auth_type == AUTH_LDAP:
4647
user = ab_security_manager.auth_user_ldap(auth.username, auth.password)

‎airflow/api_connexion/endpoints/dag_endpoint.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from typing import Collection, Optional
2020

2121
from connexion import NoContent
22-
from flask import current_app, g, request
22+
from flask import g, request
2323
from marshmallow import ValidationError
2424
from sqlalchemy.orm import Session
2525
from sqlalchemy.sql.expression import or_
@@ -38,6 +38,7 @@
3838
from airflow.exceptions import AirflowException, DagNotFound
3939
from airflow.models.dag import DagModel, DagTag
4040
from airflow.security import permissions
41+
from airflow.utils.airflow_flask_app import get_airflow_app
4142
from airflow.utils.session import NEW_SESSION, provide_session
4243

4344

@@ -56,7 +57,7 @@ def get_dag(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
5657
@security.requires_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG)])
5758
def get_dag_details(*, dag_id: str) -> APIResponse:
5859
"""Get details of DAG."""
59-
dag: DAG = current_app.dag_bag.get_dag(dag_id)
60+
dag: DAG = get_airflow_app().dag_bag.get_dag(dag_id)
6061
if not dag:
6162
raise NotFound("DAG not found", detail=f"The DAG with dag_id: {dag_id} was not found")
6263
return dag_detail_schema.dump(dag)
@@ -83,7 +84,7 @@ def get_dags(
8384
if dag_id_pattern:
8485
dags_query = dags_query.filter(DagModel.dag_id.ilike(f'%{dag_id_pattern}%'))
8586

86-
readable_dags = current_app.appbuilder.sm.get_accessible_dag_ids(g.user)
87+
readable_dags = get_airflow_app().appbuilder.sm.get_accessible_dag_ids(g.user)
8788

8889
dags_query = dags_query.filter(DagModel.dag_id.in_(readable_dags))
8990
if tags:
@@ -143,7 +144,7 @@ def patch_dags(limit, session, offset=0, only_active=True, tags=None, dag_id_pat
143144
if dag_id_pattern == '~':
144145
dag_id_pattern = '%'
145146
dags_query = dags_query.filter(DagModel.dag_id.ilike(f'%{dag_id_pattern}%'))
146-
editable_dags = current_app.appbuilder.sm.get_editable_dag_ids(g.user)
147+
editable_dags = get_airflow_app().appbuilder.sm.get_editable_dag_ids(g.user)
147148

148149
dags_query = dags_query.filter(DagModel.dag_id.in_(editable_dags))
149150
if tags:

‎airflow/api_connexion/endpoints/dag_run_endpoint.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
import pendulum
2121
from connexion import NoContent
22-
from flask import current_app, g, request
22+
from flask import g
2323
from marshmallow import ValidationError
2424
from sqlalchemy import or_
2525
from sqlalchemy.orm import Query, Session
@@ -30,6 +30,7 @@
3030
set_dag_run_state_to_success,
3131
)
3232
from airflow.api_connexion import security
33+
from airflow.api_connexion.endpoints.request_dict import get_json_request_dict
3334
from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound
3435
from airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters
3536
from airflow.api_connexion.schemas.dag_run_schema import (
@@ -47,6 +48,7 @@
4748
from airflow.api_connexion.types import APIResponse
4849
from airflow.models import DagModel, DagRun
4950
from airflow.security import permissions
51+
from airflow.utils.airflow_flask_app import get_airflow_app
5052
from airflow.utils.session import NEW_SESSION, provide_session
5153
from airflow.utils.state import DagRunState
5254
from airflow.utils.types import DagRunType
@@ -167,7 +169,7 @@ def get_dag_runs(
167169

168170
# This endpoint allows specifying ~ as the dag_id to retrieve DAG Runs for all DAGs.
169171
if dag_id == "~":
170-
appbuilder = current_app.appbuilder
172+
appbuilder = get_airflow_app().appbuilder
171173
query = query.filter(DagRun.dag_id.in_(appbuilder.sm.get_readable_dag_ids(g.user)))
172174
else:
173175
query = query.filter(DagRun.dag_id == dag_id)
@@ -199,13 +201,13 @@ def get_dag_runs(
199201
@provide_session
200202
def get_dag_runs_batch(*, session: Session = NEW_SESSION) -> APIResponse:
201203
"""Get list of DAG Runs"""
202-
body = request.get_json()
204+
body = get_json_request_dict()
203205
try:
204206
data = dagruns_batch_form_schema.load(body)
205207
except ValidationError as err:
206208
raise BadRequest(detail=str(err.messages))
207209

208-
appbuilder = current_app.appbuilder
210+
appbuilder = get_airflow_app().appbuilder
209211
readable_dag_ids = appbuilder.sm.get_readable_dag_ids(g.user)
210212
query = session.query(DagRun)
211213
if data.get("dag_ids"):
@@ -252,7 +254,7 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
252254
detail=f"DAG with dag_id: '{dag_id}' has import errors",
253255
)
254256
try:
255-
post_body = dagrun_schema.load(request.json, session=session)
257+
post_body = dagrun_schema.load(get_json_request_dict(), session=session)
256258
except ValidationError as err:
257259
raise BadRequest(detail=str(err))
258260

@@ -268,7 +270,7 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
268270
)
269271
if not dagrun_instance:
270272
try:
271-
dag = current_app.dag_bag.get_dag(dag_id)
273+
dag = get_airflow_app().dag_bag.get_dag(dag_id)
272274
dag_run = dag.create_dagrun(
273275
run_type=DagRunType.MANUAL,
274276
run_id=run_id,
@@ -277,7 +279,7 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
277279
state=DagRunState.QUEUED,
278280
conf=post_body.get("conf"),
279281
external_trigger=True,
280-
dag_hash=current_app.dag_bag.dags_hash.get(dag_id),
282+
dag_hash=get_airflow_app().dag_bag.dags_hash.get(dag_id),
281283
)
282284
return dagrun_schema.dump(dag_run)
283285
except ValueError as ve:
@@ -310,12 +312,12 @@ def update_dag_run_state(*, dag_id: str, dag_run_id: str, session: Session = NEW
310312
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
311313
raise NotFound(error_message)
312314
try:
313-
post_body = set_dagrun_state_form_schema.load(request.json)
315+
post_body = set_dagrun_state_form_schema.load(get_json_request_dict())
314316
except ValidationError as err:
315317
raise BadRequest(detail=str(err))
316318

317319
state = post_body['state']
318-
dag = current_app.dag_bag.get_dag(dag_id)
320+
dag = get_airflow_app().dag_bag.get_dag(dag_id)
319321
if state == DagRunState.SUCCESS:
320322
set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True)
321323
elif state == DagRunState.QUEUED:
@@ -339,15 +341,15 @@ def clear_dag_run(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSIO
339341
session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none()
340342
)
341343
if dag_run is None:
342-
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
344+
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
343345
raise NotFound(error_message)
344346
try:
345-
post_body = clear_dagrun_form_schema.load(request.json)
347+
post_body = clear_dagrun_form_schema.load(get_json_request_dict())
346348
except ValidationError as err:
347349
raise BadRequest(detail=str(err))
348350

349351
dry_run = post_body.get('dry_run', False)
350-
dag = current_app.dag_bag.get_dag(dag_id)
352+
dag = get_airflow_app().dag_bag.get_dag(dag_id)
351353
start_date = dag_run.logical_date
352354
end_date = dag_run.logical_date
353355

‎airflow/api_connexion/endpoints/extra_link_endpoint.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from flask import current_app
1918
from sqlalchemy.orm.session import Session
2019

2120
from airflow import DAG
@@ -25,6 +24,7 @@
2524
from airflow.exceptions import TaskNotFound
2625
from airflow.models.dagbag import DagBag
2726
from airflow.security import permissions
27+
from airflow.utils.airflow_flask_app import get_airflow_app
2828
from airflow.utils.session import NEW_SESSION, provide_session
2929

3030

@@ -46,7 +46,7 @@ def get_extra_links(
4646
"""Get extra links for task instance"""
4747
from airflow.models.taskinstance import TaskInstance
4848

49-
dagbag: DagBag = current_app.dag_bag
49+
dagbag: DagBag = get_airflow_app().dag_bag
5050
dag: DAG = dagbag.get_dag(dag_id)
5151
if not dag:
5252
raise NotFound("DAG not found", detail=f'DAG with ID = "{dag_id}" not found')

‎airflow/api_connexion/endpoints/log_endpoint.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,9 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17-
1817
from typing import Any, Optional
1918

20-
from flask import Response, current_app, request
19+
from flask import Response, request
2120
from itsdangerous.exc import BadSignature
2221
from itsdangerous.url_safe import URLSafeSerializer
2322
from sqlalchemy.orm.session import Session
@@ -29,6 +28,7 @@
2928
from airflow.exceptions import TaskNotFound
3029
from airflow.models import TaskInstance
3130
from airflow.security import permissions
31+
from airflow.utils.airflow_flask_app import get_airflow_app
3232
from airflow.utils.log.log_reader import TaskLogReader
3333
from airflow.utils.session import NEW_SESSION, provide_session
3434

@@ -52,7 +52,7 @@ def get_log(
5252
session: Session = NEW_SESSION,
5353
) -> APIResponse:
5454
"""Get logs for specific task instance"""
55-
key = current_app.config["SECRET_KEY"]
55+
key = get_airflow_app().config["SECRET_KEY"]
5656
if not token:
5757
metadata = {}
5858
else:
@@ -87,7 +87,7 @@ def get_log(
8787
metadata['end_of_log'] = True
8888
raise NotFound(title="TaskInstance not found")
8989

90-
dag = current_app.dag_bag.get_dag(dag_id)
90+
dag = get_airflow_app().dag_bag.get_dag(dag_id)
9191
if dag:
9292
try:
9393
ti.task = dag.get_task(ti.task_id)
@@ -101,7 +101,8 @@ def get_log(
101101
if return_type == 'application/json' or return_type is None: # default
102102
logs, metadata = task_log_reader.read_log_chunks(ti, task_try_number, metadata)
103103
logs = logs[0] if task_try_number is not None else logs
104-
token = URLSafeSerializer(key).dumps(metadata)
104+
# we must have token here, so we can safely ignore it
105+
token = URLSafeSerializer(key).dumps(metadata) # type: ignore[assignment]
105106
return logs_schema.dump(LogResponseObject(continuation_token=token, content=logs))
106107
# text/plain. Stream
107108
logs = task_log_reader.read_log_stream(ti, task_try_number, metadata)

‎airflow/api_connexion/endpoints/pool_endpoint.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,14 @@
1818
from http import HTTPStatus
1919
from typing import Optional
2020

21-
from flask import Response, request
21+
from flask import Response
2222
from marshmallow import ValidationError
2323
from sqlalchemy import func
2424
from sqlalchemy.exc import IntegrityError
2525
from sqlalchemy.orm import Session
2626

2727
from airflow.api_connexion import security
28+
from airflow.api_connexion.endpoints.request_dict import get_json_request_dict
2829
from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound
2930
from airflow.api_connexion.parameters import apply_sorting, check_limit, format_parameters
3031
from airflow.api_connexion.schemas.pool_schema import PoolCollection, pool_collection_schema, pool_schema
@@ -85,9 +86,10 @@ def patch_pool(
8586
session: Session = NEW_SESSION,
8687
) -> APIResponse:
8788
"""Update a pool"""
89+
request_dict = get_json_request_dict()
8890
# Only slots can be modified in 'default_pool'
8991
try:
90-
if pool_name == Pool.DEFAULT_POOL_NAME and request.json["name"] != Pool.DEFAULT_POOL_NAME:
92+
if pool_name == Pool.DEFAULT_POOL_NAME and request_dict["name"] != Pool.DEFAULT_POOL_NAME:
9193
if update_mask and len(update_mask) == 1 and update_mask[0].strip() == "slots":
9294
pass
9395
else:
@@ -100,7 +102,7 @@ def patch_pool(
100102
raise NotFound(detail=f"Pool with name:'{pool_name}' not found")
101103

102104
try:
103-
patch_body = pool_schema.load(request.json)
105+
patch_body = pool_schema.load(request_dict)
104106
except ValidationError as err:
105107
raise BadRequest(detail=str(err.messages))
106108

@@ -121,7 +123,7 @@ def patch_pool(
121123

122124
else:
123125
required_fields = {"name", "slots"}
124-
fields_diff = required_fields - set(request.json.keys())
126+
fields_diff = required_fields - set(get_json_request_dict().keys())
125127
if fields_diff:
126128
raise BadRequest(detail=f"Missing required property(ies): {sorted(fields_diff)}")
127129

@@ -136,12 +138,12 @@ def patch_pool(
136138
def post_pool(*, session: Session = NEW_SESSION) -> APIResponse:
137139
"""Create a pool"""
138140
required_fields = {"name", "slots"} # Pool would require both fields in the post request
139-
fields_diff = required_fields - set(request.json.keys())
141+
fields_diff = required_fields - set(get_json_request_dict().keys())
140142
if fields_diff:
141143
raise BadRequest(detail=f"Missing required property(ies): {sorted(fields_diff)}")
142144

143145
try:
144-
post_body = pool_schema.load(request.json, session=session)
146+
post_body = pool_schema.load(get_json_request_dict(), session=session)
145147
except ValidationError as err:
146148
raise BadRequest(detail=str(err.messages))
147149

0 commit comments

Comments
 (0)