From 48bee64810e5d54fa19deb3753774e1dbccc7b48 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 27 Aug 2023 22:46:48 +0200 Subject: [PATCH] Improve importing the modules in Airflow www package --- airflow/www/api/experimental/endpoints.py | 7 +++++-- airflow/www/extensions/init_appbuilder.py | 10 +++++++--- airflow/www/extensions/init_views.py | 6 +++++- .../www/extensions/init_wsgi_middlewares.py | 2 +- airflow/www/fab_security/manager.py | 12 ++++++++---- airflow/www/security.py | 4 +++- airflow/www/utils.py | 9 +++++---- airflow/www/views.py | 19 +++++++++++-------- 8 files changed, 45 insertions(+), 24 deletions(-) diff --git a/airflow/www/api/experimental/endpoints.py b/airflow/www/api/experimental/endpoints.py index ccf024543f328..0b7d9a91348f6 100644 --- a/airflow/www/api/experimental/endpoints.py +++ b/airflow/www/api/experimental/endpoints.py @@ -19,9 +19,9 @@ import logging from functools import wraps -from typing import Callable, TypeVar, cast +from typing import TYPE_CHECKING, Callable, TypeVar, cast -from flask import Blueprint, Response, current_app, g, jsonify, request, url_for +from flask import Blueprint, current_app, g, jsonify, request, url_for from airflow import models from airflow.api.common.experimental import delete_dag as delete, pool as pool_api, trigger_dag as trigger @@ -37,6 +37,9 @@ from airflow.utils.strings import to_boolean from airflow.version import version +if TYPE_CHECKING: + from flask import Response + log = logging.getLogger(__name__) T = TypeVar("T", bound=Callable) diff --git a/airflow/www/extensions/init_appbuilder.py b/airflow/www/extensions/init_appbuilder.py index bbc6343fcf05a..5ad9f0eabcf4d 100644 --- a/airflow/www/extensions/init_appbuilder.py +++ b/airflow/www/extensions/init_appbuilder.py @@ -20,9 +20,10 @@ import logging from functools import reduce +from typing import TYPE_CHECKING from flask import Blueprint, current_app, url_for -from flask_appbuilder import BaseView, __version__ +from flask_appbuilder import __version__ from flask_appbuilder.babel.manager import BabelManager from flask_appbuilder.const import ( LOGMSG_ERR_FAB_ADD_PERMISSION_MENU, @@ -35,14 +36,17 @@ ) from flask_appbuilder.filters import TemplateFilters from flask_appbuilder.menu import Menu -from flask_appbuilder.security.manager import BaseSecurityManager from flask_appbuilder.views import IndexView, UtilView -from sqlalchemy.orm import Session from airflow import settings from airflow.configuration import conf from airflow.www.extensions.init_auth_manager import get_auth_manager +if TYPE_CHECKING: + from flask_appbuilder import BaseView + from flask_appbuilder.security.manager import BaseSecurityManager + from sqlalchemy.orm import Session + # This product contains a modified portion of 'Flask App Builder' developed by Daniel Vaz Gaspar. # (https://github.com/dpgaspar/Flask-AppBuilder). # Copyright 2013, Daniel Vaz Gaspar diff --git a/airflow/www/extensions/init_views.py b/airflow/www/extensions/init_views.py index 492f325bb9c99..16813031e3946 100644 --- a/airflow/www/extensions/init_views.py +++ b/airflow/www/extensions/init_views.py @@ -20,11 +20,12 @@ import warnings from functools import cached_property from os import path +from typing import TYPE_CHECKING from connexion import FlaskApi, ProblemException, Resolver from connexion.decorators.validation import RequestBodyValidator from connexion.exceptions import BadRequestProblem -from flask import Flask, request +from flask import request from airflow.api_connexion.exceptions import common_error_handler from airflow.configuration import conf @@ -32,6 +33,9 @@ from airflow.security import permissions from airflow.utils.yaml import safe_load +if TYPE_CHECKING: + from flask import Flask + log = logging.getLogger(__name__) # airflow/www/extensions/init_views.py => airflow/ diff --git a/airflow/www/extensions/init_wsgi_middlewares.py b/airflow/www/extensions/init_wsgi_middlewares.py index 37d4a074b4d18..27368e71eed9b 100644 --- a/airflow/www/extensions/init_wsgi_middlewares.py +++ b/airflow/www/extensions/init_wsgi_middlewares.py @@ -20,7 +20,6 @@ from typing import TYPE_CHECKING, Iterable from urllib.parse import urlsplit -from flask import Flask from werkzeug.middleware.dispatcher import DispatcherMiddleware from werkzeug.middleware.proxy_fix import ProxyFix @@ -29,6 +28,7 @@ if TYPE_CHECKING: from _typeshed.wsgi import StartResponse, WSGIEnvironment + from flask import Flask def _root_app(env: WSGIEnvironment, resp: StartResponse) -> Iterable[bytes]: diff --git a/airflow/www/fab_security/manager.py b/airflow/www/fab_security/manager.py index 06ec74bbc556f..0e6be93b7bd54 100644 --- a/airflow/www/fab_security/manager.py +++ b/airflow/www/fab_security/manager.py @@ -20,12 +20,11 @@ import datetime import logging -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 import re2 -from flask import Flask, g, session, url_for -from flask_appbuilder import AppBuilder +from flask import g, session, url_for from flask_appbuilder.const import ( AUTH_DB, AUTH_LDAP, @@ -63,10 +62,15 @@ from flask_limiter.util import get_remote_address from werkzeug.security import check_password_hash -from airflow.auth.managers.fab.models import Action, Permission, RegisterUser, Resource, Role, User from airflow.configuration import conf from airflow.www.extensions.init_auth_manager import get_auth_manager +if TYPE_CHECKING: + from flask import Flask + from flask_appbuilder import AppBuilder + + from airflow.auth.managers.fab.models import Action, Permission, RegisterUser, Resource, Role, User + # This product contains a modified portion of 'Flask App Builder' developed by Daniel Vaz Gaspar. # (https://github.com/dpgaspar/Flask-AppBuilder). # Copyright 2013, Daniel Vaz Gaspar diff --git a/airflow/www/security.py b/airflow/www/security.py index 3e61008501fca..a558d92b15621 100644 --- a/airflow/www/security.py +++ b/airflow/www/security.py @@ -21,7 +21,7 @@ from flask import g from sqlalchemy import or_, select -from sqlalchemy.orm import Session, joinedload +from sqlalchemy.orm import joinedload from airflow.auth.managers.fab.models import Permission, Resource, Role, User from airflow.auth.managers.fab.views.permissions import ( @@ -61,6 +61,8 @@ } if TYPE_CHECKING: + from sqlalchemy.orm import Session + SecurityManagerOverride: type = object else: # Fetch the security manager override from the auth manager diff --git a/airflow/www/utils.py b/airflow/www/utils.py index a8638d328c26f..5c1de3d05b649 100644 --- a/airflow/www/utils.py +++ b/airflow/www/utils.py @@ -27,19 +27,16 @@ from flask.helpers import flash from flask_appbuilder.forms import FieldConverter from flask_appbuilder.models.filters import BaseFilter -from flask_appbuilder.models.sqla import Model, filters as fab_sqlafilters +from flask_appbuilder.models.sqla import filters as fab_sqlafilters from flask_appbuilder.models.sqla.filters import get_field_setup_query, set_value_to_type from flask_appbuilder.models.sqla.interface import SQLAInterface from flask_babel import lazy_gettext from markdown_it import MarkdownIt from markupsafe import Markup -from pendulum.datetime import DateTime from pygments import highlight, lexers from pygments.formatters import HtmlFormatter -from pygments.lexer import Lexer from sqlalchemy import delete, func, select, types from sqlalchemy.ext.associationproxy import AssociationProxy -from sqlalchemy.sql import Select from airflow.exceptions import RemovedInAirflow3Warning from airflow.models import errors @@ -56,7 +53,11 @@ from airflow.www.widgets import AirflowDateTimePickerWidget if TYPE_CHECKING: + from flask_appbuilder.models.sqla import Model + from pendulum.datetime import DateTime + from pygments.lexer import Lexer from sqlalchemy.orm.session import Session + from sqlalchemy.sql import Select from sqlalchemy.sql.operators import ColumnOperators from airflow.www.fab_security.sqla.manager import SecurityManager diff --git a/airflow/www/views.py b/airflow/www/views.py index 8d39a4bcfaed6..2e70f514f0350 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -31,7 +31,7 @@ from collections import defaultdict from functools import cached_property, wraps from json import JSONDecodeError -from typing import Any, Callable, Collection, Iterator, Mapping, MutableMapping, Sequence +from typing import TYPE_CHECKING, Any, Callable, Collection, Iterator, Mapping, MutableMapping, Sequence from urllib.parse import unquote, urljoin, urlsplit import configupdater @@ -69,7 +69,7 @@ from pendulum.parsing.exceptions import ParserError from sqlalchemy import Date, and_, case, desc, func, inspect, select, union_all from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, joinedload +from sqlalchemy.orm import joinedload from wtforms import BooleanField, validators import airflow @@ -96,12 +96,10 @@ from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import Connection, DagModel, DagTag, Log, SlaMiss, TaskFail, Trigger, XCom, errors -from airflow.models.abstractoperator import AbstractOperator -from airflow.models.dag import DAG, get_dataset_triggered_next_run_info +from airflow.models.dag import get_dataset_triggered_next_run_info from airflow.models.dagrun import RUN_ID_REGEX, DagRun, DagRunType from airflow.models.dataset import DagScheduleDatasetReference, DatasetDagRunQueue, DatasetEvent, DatasetModel from airflow.models.mappedoperator import MappedOperator -from airflow.models.operator import Operator from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance, TaskInstanceNote from airflow.providers_manager import ProvidersManager @@ -138,6 +136,13 @@ ) from airflow.www.widgets import AirflowModelListWidget, AirflowVariableShowWidget +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.abstractoperator import AbstractOperator + from airflow.models.dag import DAG + from airflow.models.operator import Operator + PAGE_SIZE = conf.getint("webserver", "page_size") FILTER_TAGS_COOKIE = "tags_filter" FILTER_STATUS_COOKIE = "dag_status_filter" @@ -835,9 +840,7 @@ def index(self): is_paused_count = dict( session.execute( - all_dags.with_only_columns([DagModel.is_paused, func.count()]).group_by( - DagModel.is_paused - ) + select(DagModel.is_paused, func.count(DagModel.dag_id)).group_by(DagModel.is_paused) ).all() )