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)