diff --git a/airflow/api/__init__.py b/airflow/api/__init__.py index b37b29248703b..ac1050946ac55 100644 --- a/airflow/api/__init__.py +++ b/airflow/api/__init__.py @@ -25,15 +25,6 @@ log = logging.getLogger(__name__) -class ApiAuth: # pylint: disable=too-few-public-methods - """Class to keep module of Authentication API """ - def __init__(self): - self.api_auth = None - - -API_AUTH = ApiAuth() - - def load_auth(): """Loads authentication backend""" auth_backend = 'airflow.api.auth.backend.default' @@ -43,7 +34,7 @@ def load_auth(): pass try: - API_AUTH.api_auth = import_module(auth_backend) + return import_module(auth_backend) except ImportError as err: log.critical( "Cannot import %s for API authentication due to: %s", diff --git a/airflow/api/auth/backend/default.py b/airflow/api/auth/backend/default.py index d444c80579049..5f04ca7f7c499 100644 --- a/airflow/api/auth/backend/default.py +++ b/airflow/api/auth/backend/default.py @@ -17,7 +17,7 @@ # under the License. """Default authentication backend - everything is allowed""" from functools import wraps -from typing import Callable, Optional, TypeVar, cast +from typing import Callable, TypeVar, cast from airflow.typing_compat import Protocol @@ -33,9 +33,6 @@ def handle_response(self, _): ... -CLIENT_AUTH = None # type: Optional[ClientAuthProtocol] - - def init_app(_): """Initializes authentication backend""" diff --git a/airflow/api/client/__init__.py b/airflow/api/client/__init__.py index aca4df06cb0fb..ffb8b6a9413bd 100644 --- a/airflow/api/client/__init__.py +++ b/airflow/api/client/__init__.py @@ -33,6 +33,6 @@ def get_current_api_client() -> Client: api_module = import_module(conf.get('cli', 'api_client')) # type: Any api_client = api_module.Client( api_base_url=conf.get('cli', 'endpoint_url'), - auth=api.API_AUTH.api_auth.CLIENT_AUTH + auth=api.load_auth() ) return api_client diff --git a/airflow/models/pool.py b/airflow/models/pool.py index 3819d6d135fda..f351fd593eacf 100644 --- a/airflow/models/pool.py +++ b/airflow/models/pool.py @@ -16,11 +16,12 @@ # specific language governing permissions and limitations # under the License. -from typing import Dict, Iterable, Tuple +from typing import Dict, Iterable, Optional, Tuple from sqlalchemy import Column, Integer, String, Text, func from sqlalchemy.orm.session import Session +from airflow.exceptions import AirflowException from airflow.models.base import Base from airflow.ti_deps.dependencies_states import EXECUTION_STATES from airflow.typing_compat import TypedDict @@ -33,6 +34,7 @@ class PoolStats(TypedDict): total: int running: int queued: int + open: int class Pool(Base): @@ -79,7 +81,7 @@ def get_default_pool(session: Session = None): @provide_session def slots_stats(session: Session = None) -> Dict[str, PoolStats]: """ - Get Pool stats (Number of Running, Queued & Total tasks) + Get Pool stats (Number of Running, Queued, Open & Total tasks) :param session: SQLAlchemy ORM Session """ @@ -89,7 +91,7 @@ def slots_stats(session: Session = None) -> Dict[str, PoolStats]: pool_rows: Iterable[Tuple[str, int]] = session.query(Pool.pool, Pool.slots).all() for (pool_name, total_slots) in pool_rows: - pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0) + pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0) state_count_by_pool = ( session.query(TaskInstance.pool, TaskInstance.state, func.count()) @@ -98,19 +100,28 @@ def slots_stats(session: Session = None) -> Dict[str, PoolStats]: ).all() # calculate queued and running metrics + count: int for (pool_name, state, count) in state_count_by_pool: - pool = pools.get(pool_name) - if not pool: + stats_dict: Optional[PoolStats] = pools.get(pool_name) + if not stats_dict: continue - pool[state] = count + # TypedDict key must be a string literal, so we use if-statements to set value + if state == "running": + stats_dict["running"] = count + elif state == "queued": + stats_dict["queued"] = count + else: + raise AirflowException( + f"Unexpected state. Expected values: {EXECUTION_STATES}." + ) # calculate open metric - for pool_name, stats in pools.items(): - if stats["total"] == -1: + for pool_name, stats_dict in pools.items(): + if stats_dict["total"] == -1: # -1 means infinite - stats["open"] = -1 + stats_dict["open"] = -1 else: - stats["open"] = stats["total"] - stats[State.RUNNING] - stats[State.QUEUED] + stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"] return pools diff --git a/airflow/settings.py b/airflow/settings.py index cb4ac4549ae32..5aad336d993ec 100644 --- a/airflow/settings.py +++ b/airflow/settings.py @@ -29,8 +29,6 @@ from sqlalchemy.orm.session import Session as SASession from sqlalchemy.pool import NullPool -# noinspection PyUnresolvedReferences -from airflow import api # pylint: disable=unused-import from airflow.configuration import AIRFLOW_HOME, WEBSERVER_CONFIG, conf # NOQA F401 from airflow.logging_config import configure_logging @@ -330,7 +328,6 @@ def initialize(): # The webservers import this file from models.py with the default settings. configure_orm() configure_action_logging() - api.load_auth() # Ensure we close DB connections at scheduler and gunicon worker terminations atexit.register(dispose_orm) diff --git a/airflow/www/api/experimental/endpoints.py b/airflow/www/api/experimental/endpoints.py index 6a74f5ce181fd..94ca2519486d6 100644 --- a/airflow/www/api/experimental/endpoints.py +++ b/airflow/www/api/experimental/endpoints.py @@ -16,10 +16,11 @@ # specific language governing permissions and limitations # under the License. import logging +from functools import wraps +from typing import Callable, TypeVar, cast -from flask import Blueprint, g, jsonify, request, url_for +from flask import Blueprint, current_app, g, jsonify, request, url_for -import airflow.api from airflow import models from airflow.api.common.experimental import delete_dag as delete, pool as pool_api, trigger_dag as trigger from airflow.api.common.experimental.get_code import get_code @@ -35,7 +36,17 @@ log = logging.getLogger(__name__) -requires_authentication = airflow.api.API_AUTH.api_auth.requires_authentication +T = TypeVar("T", bound=Callable) # pylint: disable=invalid-name + + +def requires_authentication(function: T): + """Decorator for functions that require authentication""" + @wraps(function) + def decorated(*args, **kwargs): + return current_app.api_auth.requires_authentication(function)(*args, **kwargs) + + return cast(T, decorated) + api_experimental = Blueprint('api_experimental', __name__) diff --git a/airflow/www/extensions/init_security.py b/airflow/www/extensions/init_security.py index 484543a94808c..e2cc7bf6f8c78 100644 --- a/airflow/www/extensions/init_security.py +++ b/airflow/www/extensions/init_security.py @@ -14,8 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import logging +from importlib import import_module from airflow.configuration import conf +from airflow.exceptions import AirflowConfigException, AirflowException + +log = logging.getLogger(__name__) def init_xframe_protection(app): @@ -37,8 +42,19 @@ def apply_caching(response): def init_api_experimental_auth(app): - """Initialize authorization in Experimental API""" - from airflow import api - - api.load_auth() - api.API_AUTH.api_auth.init_app(app) + """Loads authentication backend""" + auth_backend = 'airflow.api.auth.backend.default' + try: + auth_backend = conf.get("api", "auth_backend") + except AirflowConfigException: + pass + + try: + app.api_auth = import_module(auth_backend) + app.api_auth.init_app(app) + except ImportError as err: + log.critical( + "Cannot import %s for API authentication due to: %s", + auth_backend, err + ) + raise AirflowException(err) diff --git a/airflow/www/extensions/init_views.py b/airflow/www/extensions/init_views.py index 7cfda964644f8..ac1e7a4661ea8 100644 --- a/airflow/www/extensions/init_views.py +++ b/airflow/www/extensions/init_views.py @@ -112,12 +112,5 @@ def init_api_experimental(app): """Initialize Experimental API""" from airflow.www.api.experimental import endpoints - # required for testing purposes otherwise the module retains - # a link to the default_auth - if app.config['TESTING']: - import importlib - - importlib.reload(endpoints) - app.register_blueprint(endpoints.api_experimental, url_prefix='/api/experimental') app.extensions['csrf'].exempt(endpoints.api_experimental)