o-nikolas commented on code in PR #69768:
URL: https://github.com/apache/airflow/pull/69768#discussion_r3634954630
##########
airflow-core/docs/core-concepts/multi-team.rst:
##########
@@ -269,6 +269,14 @@ 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.
+
+ Existing multi-team deployments should run ``airflow teams sync`` after
upgrading
Review Comment:
Multi-team is in experimental phase right now. I'm not sure we want to a
very time based piece in our docs. Update the news fragment to call out this
guidance about syncing existing multi-team deployments.
##########
airflow-core/docs/core-concepts/multi-team.rst:
##########
@@ -269,6 +269,14 @@ 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
Review Comment:
```suggestion
bundles associated with that team which do not explicitly specify a pool
will
```
##########
airflow-core/docs/core-concepts/multi-team.rst:
##########
@@ -269,6 +269,14 @@ 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
Review Comment:
```suggestion
creates a default pool named ``default_pool_<team_name>``. Tasks in Dag
```
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -172,10 +199,27 @@ def team_sync(args, *, session=NEW_SESSION):
teams_added = 0
try:
- for team_name in dag_bundle_teams -
Team.get_all_team_names(session=session):
- team = Team(name=team_name)
- session.add(team)
- teams_added += 1
+ existing_teams = Team.get_all_team_names(session=session)
+ for team_name in dag_bundle_teams:
+ if team_name not in existing_teams:
+ session.add(Team(name=team_name))
+ teams_added += 1
+
+ session.flush()
+
+ if conf.getboolean("core", "multi_team"):
Review Comment:
If you're going to add a check if multi-team is enabled then maybe do it
higher? Do we even want to add the teams if multi-team is not enabled?
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -118,9 +135,19 @@ def team_delete(args, *, session=NEW_SESSION):
associations.append(f"{variable_count} variable(s)")
# Check pool associations
- if pool_count :=
session.scalar(select(func.count(Pool.id)).where(Pool.team_name == team.name)):
+ if pool_count := session.scalar(
+ select(func.count(Pool.id)).where(
+ Pool.team_name == team.name,
+ Pool.pool != Pool.get_default_team_pool_name(team.name),
+ )
+ ):
associations.append(f"{pool_count} pool(s)")
+ default_pool = session.scalar(select(Pool).where(Pool.pool ==
Pool.get_default_team_pool_name(team.name)))
+
+ if default_pool:
+ session.delete(default_pool)
Review Comment:
This delete is happening way too soon. There is a check below "are you sure
you want to delete..." you should not delete anything until after that.
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -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",
Review Comment:
This config already exists:
https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html#default-pool-task-slot-count
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -74,8 +75,24 @@ def team_create(args, *, session=NEW_SESSION):
try:
session.add(new_team)
+ session.flush()
Review Comment:
Do you need the flush here?
##########
airflow-core/src/airflow/dag_processing/dagbag.py:
##########
@@ -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
Review Comment:
Can any of this logic be extracted from the similar steps done in
_validate_executor_fields()?
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -172,10 +199,27 @@ def team_sync(args, *, session=NEW_SESSION):
teams_added = 0
try:
- for team_name in dag_bundle_teams -
Team.get_all_team_names(session=session):
- team = Team(name=team_name)
- session.add(team)
- teams_added += 1
+ existing_teams = Team.get_all_team_names(session=session)
+ for team_name in dag_bundle_teams:
+ if team_name not in existing_teams:
+ session.add(Team(name=team_name))
+ teams_added += 1
+
+ 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",
+ ),
Review Comment:
This clobbers any specific pool slots admins may have set if the pool
already exists. You really just want create if missing not create or update
##########
airflow-core/src/airflow/cli/commands/team_command.py:
##########
@@ -172,10 +199,27 @@ def team_sync(args, *, session=NEW_SESSION):
teams_added = 0
try:
- for team_name in dag_bundle_teams -
Team.get_all_team_names(session=session):
- team = Team(name=team_name)
- session.add(team)
- teams_added += 1
+ existing_teams = Team.get_all_team_names(session=session)
+ for team_name in dag_bundle_teams:
+ if team_name not in existing_teams:
+ session.add(Team(name=team_name))
+ teams_added += 1
+
+ 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,
Review Comment:
Wrap this in a helper or move the magic strings to be constants. It is a
duplicated of the one in create_team and they will surely drift from each other
over time.
##########
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_pool_task_slot_count"): "128",
Review Comment:
```suggestion
("core", "default_pool_task_slot_count"): "111",
```
This is a `core` config, and the default is already 128, setting it to
something else ensures that this is working in your test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
