Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions airflow-core/docs/core-concepts/multi-team.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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_<team_name>``. 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
"""""""""""""""""""""""""""""""""""""""""""

Expand Down
1 change: 1 addition & 0 deletions airflow-core/newsfragments/69768.feature.rst
Original file line number Diff line number Diff line change
@@ -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.
17 changes: 17 additions & 0 deletions airflow-core/src/airflow/cli/commands/team_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}")
Expand Down
26 changes: 26 additions & 0 deletions airflow-core/src/airflow/dag_processing/dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Comment on lines +184 to +186

@justinpakzad justinpakzad Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The addition to the team create command would only create the default pools for new teams going forward. That would mean for already existing teams (without explicitly assigned pools), their tasks would be assigned a pool that doesn't actually exist. Users would need to manually create the team default pools. I know multi-team is still "experimental" and relatively new, so this might not be an issue, but worth flagging.

Another potential edge case is if a user explicitly assigns the global default pool to a task, this would actually overwrite it with the default team pool. Probably not very common, but worth mentioning as this would effectively mean that the default global pool would be unreachable for team Dags.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we check for if task.pool is None or task.pool == NOTSET instead?



class DagBag(LoggingMixin):
"""
A dagbag is a collection of dags, parsed out of a folder tree and has high level configuration settings.
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 4 additions & 0 deletions airflow-core/src/airflow/models/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
19 changes: 19 additions & 0 deletions airflow-core/tests/unit/cli/commands/test_team_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down
68 changes: 68 additions & 0 deletions airflow-core/tests/unit/dag_processing/test_dagbag.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions airflow-core/tests/unit/models/test_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down