From 96918421e05aad55d102c1ad34f3a734387fe356 Mon Sep 17 00:00:00 2001 From: Sameer Mesiah Date: Sat, 11 Jul 2026 22:29:14 +0100 Subject: [PATCH] Create a default pool for each team when it is created and automatically assign tasks without an explicit pool to the team's default pool during DAG parsing. Add unit tests covering pool assignment and team creation, and update the multi-team documentation to describe the new behaviour. --- .../docs/core-concepts/multi-team.rst | 5 ++ airflow-core/newsfragments/69768.feature.rst | 1 + .../src/airflow/cli/commands/team_command.py | 17 +++++ .../src/airflow/dag_processing/dagbag.py | 26 +++++++ airflow-core/src/airflow/models/pool.py | 4 ++ .../unit/cli/commands/test_team_command.py | 19 ++++++ .../tests/unit/dag_processing/test_dagbag.py | 68 +++++++++++++++++++ airflow-core/tests/unit/models/test_pool.py | 3 + 8 files changed, 143 insertions(+) create mode 100644 airflow-core/newsfragments/69768.feature.rst diff --git a/airflow-core/docs/core-concepts/multi-team.rst b/airflow-core/docs/core-concepts/multi-team.rst index 6a9d8a2ec63e7..a20c1b8be7193 100644 --- a/airflow-core/docs/core-concepts/multi-team.rst +++ b/airflow-core/docs/core-concepts/multi-team.rst @@ -269,6 +269,11 @@ Use the ``--team-name`` option with ``airflow pools set`` to assign a pool to a The ``--team-name`` option is rejected when ``core.multi_team`` is disabled. The specified team must exist in the database (create it first with ``airflow teams create``). + When ``core.multi_team`` is enabled, ``airflow teams create`` automatically + creates a default pool named ``default_pool_``. Tasks in DAG + bundles associated with that team that do not explicitly specify a pool will + use the team's default pool automatically. + Creating Team-scoped Pools via the REST API """"""""""""""""""""""""""""""""""""""""""" diff --git a/airflow-core/newsfragments/69768.feature.rst b/airflow-core/newsfragments/69768.feature.rst new file mode 100644 index 0000000000000..b8ced1f7296b2 --- /dev/null +++ b/airflow-core/newsfragments/69768.feature.rst @@ -0,0 +1 @@ +Add automatic creation of team default pools in multi-team deployments and assign tasks without an explicitly configured pool to their team's default pool during DAG parsing. diff --git a/airflow-core/src/airflow/cli/commands/team_command.py b/airflow-core/src/airflow/cli/commands/team_command.py index 931b5507be325..04a7a445fb83e 100644 --- a/airflow-core/src/airflow/cli/commands/team_command.py +++ b/airflow-core/src/airflow/cli/commands/team_command.py @@ -25,6 +25,7 @@ from sqlalchemy.exc import IntegrityError from airflow.cli.simple_table import AirflowConsole +from airflow.configuration import conf from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.models.connection import Connection from airflow.models.pool import Pool @@ -74,8 +75,24 @@ def team_create(args, *, session=NEW_SESSION): try: session.add(new_team) + session.flush() + + if conf.getboolean("core", "multi_team"): + Pool.create_or_update_pool( + name=Pool.get_default_team_pool_name(team_name), + slots=conf.getint( + "core", + "default_pool_task_slot_count", + ), + description=f"Default pool for team '{team_name}'", + include_deferred=False, + team_name=team_name, + session=session, + ) + session.commit() print(f"Team '{team_name}' created successfully.") + except IntegrityError as e: session.rollback() raise SystemExit(f"Failed to create team '{team_name}': {e}") diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 3bf82804bffe1..bf088c04c1b2a 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -43,6 +43,7 @@ ) from airflow.executors.executor_loader import ExecutorLoader from airflow.listeners.listener import get_listener_manager +from airflow.models.pool import Pool from airflow.serialization.definitions.notset import NOTSET, ArgNotSet, is_arg_set from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.file import correct_maybe_zipped @@ -161,6 +162,30 @@ def _validate_executor_fields(dag: DAG, bundle_name: str | None = None) -> None: ) +def _assign_default_team_pools( + dag: DAG, + bundle_name: str | None = None, +) -> None: + """Assign the default team pool to tasks that do not explicitly specify a pool.""" + dag_team_name = None + + if conf.getboolean("core", "multi_team"): + if bundle_name: + from airflow.dag_processing.bundles.manager import DagBundlesManager + + bundle_manager = DagBundlesManager() + bundle_config = bundle_manager._bundle_config[bundle_name] + + dag_team_name = bundle_config.team_name + + if not dag_team_name: + return + + for task in dag.tasks: + if task.pool == Pool.DEFAULT_POOL_NAME: + task.pool = Pool.get_default_team_pool_name(dag_team_name) + + class DagBag(LoggingMixin): """ A dagbag is a collection of dags, parsed out of a folder tree and has high level configuration settings. @@ -344,6 +369,7 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True): # Validate before adding to bag (matches original _process_modules behavior) dag.validate() _validate_executor_fields(dag, self.bundle_name) + _assign_default_team_pools(dag, self.bundle_name) self.bag_dag(dag=dag) bagged_dags.append(dag) except AirflowClusterPolicySkipDag: diff --git a/airflow-core/src/airflow/models/pool.py b/airflow-core/src/airflow/models/pool.py index d6a4915ea2bb2..339abaddf5dc8 100644 --- a/airflow-core/src/airflow/models/pool.py +++ b/airflow-core/src/airflow/models/pool.py @@ -121,6 +121,10 @@ def get_default_pool(*, session: Session = NEW_SESSION) -> Pool | None: """ return Pool.get_pool(Pool.DEFAULT_POOL_NAME, session=session) + @staticmethod + def get_default_team_pool_name(team_name: str) -> str: + return f"default_pool_{team_name}" + @staticmethod @provide_session def create_or_update_pool( diff --git a/airflow-core/tests/unit/cli/commands/test_team_command.py b/airflow-core/tests/unit/cli/commands/test_team_command.py index 1df56d9c3cc3f..12720854f1557 100644 --- a/airflow-core/tests/unit/cli/commands/test_team_command.py +++ b/airflow-core/tests/unit/cli/commands/test_team_command.py @@ -79,6 +79,25 @@ def test_team_create_success(self, stdout_capture): assert "Team 'test-team' created successfully" in output assert str(team.name) in output + def test_team_create_creates_default_pool(self, stdout_capture): + """Test that creating a team also creates its default pool.""" + with conf_vars( + { + ("core", "multi_team"): "True", + ("multi_team", "default_team_pool_slots"): "128", + } + ): + with stdout_capture: + team_command.team_create(self.parser.parse_args(["teams", "create", "team-a"])) + + pool = self.session.scalar(select(Pool).where(Pool.pool == Pool.get_default_team_pool_name("team-a"))) + + assert pool is not None + assert pool.team_name == "team-a" + assert pool.slots == 128 + assert pool.include_deferred is False + assert pool.description == "Default pool for team 'team-a'" + def test_team_create_empty_name(self): """Test team creation with empty name.""" with pytest.raises(SystemExit, match="Team name cannot be empty"): diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 0ea0a29fc145b..ac3e06fb177c9 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -46,6 +46,7 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.models.dag import DagModel from airflow.models.dagwarning import DagWarning, DagWarningType +from airflow.models.pool import Pool from airflow.models.serialized_dag import SerializedDagModel from airflow.sdk import DAG, BaseOperator @@ -383,6 +384,73 @@ def test_dag_with_bundle_name(self, tmp_path): for dag in dagbag2.dags.values(): assert dag.bundle_name is None + @pytest.mark.parametrize( + ("team_name", "operator_args", "expected_pool"), + [ + pytest.param( + "team_a", + "", + Pool.get_default_team_pool_name("team_a"), + id="default-pool-replaced", + ), + pytest.param( + "team_a", + ', pool="custom_pool"', + "custom_pool", + id="custom-pool-preserved", + ), + pytest.param( + None, + "", + Pool.DEFAULT_POOL_NAME, + id="no-team", + ), + ], + ) + @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") + def test_default_pool_replaced_with_team_pool( + self, + mock_manager, + tmp_path, + team_name, + operator_args, + expected_pool, + ): + mock_bundle = mock.MagicMock() + mock_bundle.team_name = team_name + mock_manager.return_value._bundle_config = { + "test_bundle": mock_bundle, + } + + dag_file = tmp_path / "test_dag.py" + dag_file.write_text( + textwrap.dedent( + f""" + from airflow.sdk import dag + + from airflow.providers.standard.operators.empty import EmptyOperator + + @dag(schedule=None) + def my_dag(): + EmptyOperator(task_id="task1"{operator_args}) + + my_dag() + """ + ) + ) + + with conf_vars({("core", "multi_team"): "True"}): + dagbag = DagBag( + dag_folder=os.fspath(tmp_path), + bundle_name="test_bundle", + ) + + dag = dagbag.get_dag("my_dag") + + assert dag is not None + assert not dagbag.import_errors + assert dag.task_dict["task1"].pool == expected_pool + def test_get_existing_dag(self, tmp_path, standard_example_dags_folder): """ Test that we're able to parse some example DAGs and retrieve them diff --git a/airflow-core/tests/unit/models/test_pool.py b/airflow-core/tests/unit/models/test_pool.py index 2c7db9d9800e4..d6bccc7ccf99c 100644 --- a/airflow-core/tests/unit/models/test_pool.py +++ b/airflow-core/tests/unit/models/test_pool.py @@ -286,6 +286,9 @@ def test_get_pools(self): assert pools[0].pool == self.pools[0].pool assert pools[1].pool == self.pools[1].pool + def test_default_team_pool_name(self): + assert Pool.get_default_team_pool_name("team_a") == "default_pool_team_a" + def test_create_pool(self, session): self.add_pools() pool = Pool.create_or_update_pool(name="foo", slots=5, description="", include_deferred=True)