Multi-team. Verify a task uses a pool it has access to when scheduling (#61227)

Author: vincbeck

airflow-core/src/airflow/models/pool.py

@@ -254,6 +254,7 @@ def to_json(self) -> dict[str, Any]:
             "slots": self.slots,
             "description": self.description,
             "include_deferred": self.include_deferred,
+            "team": self.team_name,
         }
 
     @provide_session

airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py

@@ -21,6 +21,7 @@
 
 from sqlalchemy import select
 
+from airflow.configuration import conf
 from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
 from airflow.utils.session import provide_session
 
@@ -41,6 +42,7 @@ def _get_dep_statuses(self, ti, session, dep_context=None):
         :param dep_context: the context for which this dependency should be evaluated for
         :return: True if there are available slots in the pool.
         """
+        from airflow.models.dag import DagModel
         from airflow.models.pool import Pool  # To avoid a circular dependency
 
         pool_name = ti.pool
@@ -53,6 +55,21 @@ def _get_dep_statuses(self, ti, session, dep_context=None):
             )
             return
 
+        # Check team compatibility
+        if pool.team_name:
+            if not conf.getboolean("core", "multi_team"):
+                raise ValueError(
+                    "Multi-team mode is not configured in the Airflow environment but the pool the task is using belongs to a team"
+                )
+
+            dag_team_name = DagModel.get_team_name(ti.dag_id, session=session)
+            if dag_team_name != pool.team_name:
+                yield self._failing_status(
+                    reason=f"Pool '{pool_name}' is assigned to team '{pool.team_name}' "
+                    f"but task's DAG belongs to team '{dag_team_name or 'None'}'"
+                )
+                return
+
         open_slots = pool.open_slots(session=session)
         if ti.state in pool.get_occupied_states():
             open_slots += ti.pool_slots

airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py

@@ -28,13 +28,15 @@
 from airflow.utils.state import TaskInstanceState
 
 from tests_common.test_utils import db
+from tests_common.test_utils.config import conf_vars
 
 pytestmark = pytest.mark.db_test
 
 
 class TestPoolSlotsAvailableDep:
     def setup_method(self):
         db.clear_db_pools()
+        db.clear_db_teams()
         with create_session() as session:
             test_pool = Pool(pool="test_pool", include_deferred=False)
             test_includes_deferred_pool = Pool(pool="test_includes_deferred_pool", include_deferred=True)
@@ -43,6 +45,7 @@ def setup_method(self):
 
     def teardown_method(self):
         db.clear_db_pools()
+        db.clear_db_teams()
 
     @patch("airflow.models.Pool.open_slots", return_value=0)
     def test_pooled_task_reached_concurrency(self, mock_open_slots):
@@ -70,3 +73,68 @@ def test_deferred_pooled_task_pass(self, mock_open_slots):
     def test_task_with_nonexistent_pool(self):
         ti = Mock(pool="nonexistent_pool", pool_slots=1)
         assert not PoolSlotsAvailableDep().is_met(ti=ti)
+
+    @conf_vars({("core", "multi_team"): "True"})
+    def test_pool_team_mismatch(self):
+        """Test that a task from one team cannot use a pool assigned to another team."""
+        from airflow.models.dag import DagModel
+        from airflow.models.team import Team
+
+        with create_session() as session:
+            # Create teams
+            team_a = Team(name="teamA")
+            team_b = Team(name="teamB")
+            session.add_all([team_a, team_b])
+            session.commit()
+
+            # Create a pool assigned to teamA
+            pool_team_a = Pool(pool="pool_team_a", slots=10, include_deferred=False, team_name="teamA")
+            session.add(pool_team_a)
+            session.commit()
+
+        # Mock a task instance from a DAG belonging to teamB
+        ti = Mock(pool="pool_team_a", pool_slots=1, dag_id="test_dag")
+
+        with patch.object(DagModel, "get_team_name", return_value="teamB"):
+            assert not PoolSlotsAvailableDep().is_met(ti=ti)
+
+    @conf_vars({("core", "multi_team"): "True"})
+    def test_pool_team_match(self):
+        """Test that a task from a team can use a pool assigned to the same team."""
+        from airflow.models.dag import DagModel
+        from airflow.models.team import Team
+
+        with create_session() as session:
+            # Create team
+            team_a = Team(name="teamA")
+            session.add(team_a)
+            session.commit()
+
+            # Create a pool assigned to teamA
+            pool_team_a = Pool(pool="pool_team_a", slots=10, include_deferred=False, team_name="teamA")
+            session.add(pool_team_a)
+            session.commit()
+
+        # Mock a task instance from a DAG belonging to teamA
+        ti = Mock(pool="pool_team_a", pool_slots=1, dag_id="test_dag", state=None)
+
+        with patch.object(DagModel, "get_team_name", return_value="teamA"):
+            with patch("airflow.models.Pool.open_slots", return_value=5):
+                assert PoolSlotsAvailableDep().is_met(ti=ti)
+
+    def test_pool_no_team_assignment(self):
+        """Test that a pool without team assignment can be used by any DAG."""
+        from airflow.models.dag import DagModel
+
+        with create_session() as session:
+            # Create a pool without team assignment
+            pool_no_team = Pool(pool="pool_no_team", slots=10, include_deferred=False, team_name=None)
+            session.add(pool_no_team)
+            session.commit()
+
+        # Mock a task instance from any DAG
+        ti = Mock(pool="pool_no_team", pool_slots=1, dag_id="test_dag", state=None)
+
+        with patch.object(DagModel, "get_team_name", return_value="teamA"):
+            with patch("airflow.models.Pool.open_slots", return_value=5):
+                assert PoolSlotsAvailableDep().is_met(ti=ti)