Override get_authorized_connections, get_authorized_pools and get_authorized_variables in Fab auth manager

airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -423,133 +423,133 @@ def batch_is_authorized_variable(
         )
 
     @provide_session
-    def get_authorized_dag_ids(
+    def get_authorized_connections(
         self,
         *,
         user: T,
         method: ResourceMethod = "GET",
         session: Session = NEW_SESSION,
     ) -> set[str]:
         """
-        Get DAGs the user has access to.
+        Get connection ids (``conn_id``) the user has access to.
 
         :param user: the user
         :param method: the method to filter on
         :param session: the session
         """
-        stmt = (
-            select(DagModel.dag_id, Team.name)
-            .join(DagBundleModel, DagModel.bundle_name == DagBundleModel.name)
-            .join(
-                dag_bundle_team_association_table,
-                DagBundleModel.name == dag_bundle_team_association_table.c.dag_bundle_name,
-                isouter=True,
-            )
-            .join(Team, Team.id == dag_bundle_team_association_table.c.team_id, isouter=True)
-        )
+        stmt = select(Connection.conn_id, Team.name).join(Team, Connection.team_id == Team.id, isouter=True)
         rows = session.execute(stmt).all()
-        dags_by_team: dict[str | None, set[str]] = defaultdict(set)
-        for dag_id, team_name in rows:
-            dags_by_team[team_name].add(dag_id)
+        connections_by_team: dict[str | None, set[str]] = defaultdict(set)
+        for conn_id, team_name in rows:
+            connections_by_team[team_name].add(conn_id)
 
-        dag_ids: set[str] = set()
-        for team_name, team_dag_ids in dags_by_team.items():
-            dag_ids.update(
-                self.filter_authorized_dag_ids(
-                    dag_ids=team_dag_ids, user=user, method=method, team_name=team_name
+        conn_ids: set[str] = set()
+        for team_name, team_conn_ids in connections_by_team.items():
+            conn_ids.update(
+                self.filter_authorized_connections(
+                    conn_ids=team_conn_ids, user=user, method=method, team_name=team_name
                 )
             )
 
-        return dag_ids
+        return conn_ids
 
-    def filter_authorized_dag_ids(
+    def filter_authorized_connections(
         self,
         *,
-        dag_ids: set[str],
+        conn_ids: set[str],
         user: T,
         method: ResourceMethod = "GET",
         team_name: str | None = None,
     ) -> set[str]:
         """
-        Filter DAGs the user has access to.
+        Filter connections the user has access to.
 
-        By default, check individually if the user has permissions to access the DAG.
+        By default, check individually if the user has permissions to access the connection.
         Can lead to some poor performance. It is recommended to override this method in the auth manager
         implementation to provide a more efficient implementation.
 
-        :param dag_ids: the set of DAG ids
+        :param conn_ids: the set of connection ids (``conn_id``)
         :param user: the user
         :param method: the method to filter on
-        :param team_name: the name of the team associated to the Dags if Airflow environment runs in
+        :param team_name: the name of the team associated to the connections if Airflow environment runs in
             multi-team mode
         """
 
-        def _is_authorized_dag_id(dag_id: str):
-            return self.is_authorized_dag(
-                method=method, details=DagDetails(id=dag_id, team_name=team_name), user=user
+        def _is_authorized_connection(conn_id: str):
+            return self.is_authorized_connection(
+                method=method, details=ConnectionDetails(conn_id=conn_id, team_name=team_name), user=user
             )
 
-        return {dag_id for dag_id in dag_ids if _is_authorized_dag_id(dag_id)}
+        return {conn_id for conn_id in conn_ids if _is_authorized_connection(conn_id)}
 
     @provide_session
-    def get_authorized_connections(
+    def get_authorized_dag_ids(
         self,
         *,
         user: T,
         method: ResourceMethod = "GET",
         session: Session = NEW_SESSION,
     ) -> set[str]:
         """
-        Get connection ids (``conn_id``) the user has access to.
+        Get DAGs the user has access to.
 
         :param user: the user
         :param method: the method to filter on
         :param session: the session
         """
-        stmt = select(Connection.conn_id, Team.name).join(Team, Connection.team_id == Team.id, isouter=True)
+        stmt = (
+            select(DagModel.dag_id, Team.name)
+            .join(DagBundleModel, DagModel.bundle_name == DagBundleModel.name)
+            .join(
+                dag_bundle_team_association_table,
+                DagBundleModel.name == dag_bundle_team_association_table.c.dag_bundle_name,
+                isouter=True,
+            )
+            .join(Team, Team.id == dag_bundle_team_association_table.c.team_id, isouter=True)
+        )
         rows = session.execute(stmt).all()
-        connections_by_team: dict[str | None, set[str]] = defaultdict(set)
-        for conn_id, team_name in rows:
-            connections_by_team[team_name].add(conn_id)
+        dags_by_team: dict[str | None, set[str]] = defaultdict(set)
+        for dag_id, team_name in rows:
+            dags_by_team[team_name].add(dag_id)
 
-        conn_ids: set[str] = set()
-        for team_name, team_conn_ids in connections_by_team.items():
-            conn_ids.update(
-                self.filter_authorized_connections(
-                    conn_ids=team_conn_ids, user=user, method=method, team_name=team_name
+        dag_ids: set[str] = set()
+        for team_name, team_dag_ids in dags_by_team.items():
+            dag_ids.update(
+                self.filter_authorized_dag_ids(
+                    dag_ids=team_dag_ids, user=user, method=method, team_name=team_name
                 )
             )
 
-        return conn_ids
+        return dag_ids
 
-    def filter_authorized_connections(
+    def filter_authorized_dag_ids(
         self,
         *,
-        conn_ids: set[str],
+        dag_ids: set[str],
         user: T,
         method: ResourceMethod = "GET",
         team_name: str | None = None,
     ) -> set[str]:
         """
-        Filter connections the user has access to.
+        Filter DAGs the user has access to.
 
-        By default, check individually if the user has permissions to access the connection.
+        By default, check individually if the user has permissions to access the DAG.
         Can lead to some poor performance. It is recommended to override this method in the auth manager
         implementation to provide a more efficient implementation.
 
-        :param conn_ids: the set of connection ids (``conn_id``)
+        :param dag_ids: the set of DAG ids
         :param user: the user
         :param method: the method to filter on
-        :param team_name: the name of the team associated to the connections if Airflow environment runs in
+        :param team_name: the name of the team associated to the Dags if Airflow environment runs in
             multi-team mode
         """
 
-        def _is_authorized_connection(conn_id: str):
-            return self.is_authorized_connection(
-                method=method, details=ConnectionDetails(conn_id=conn_id, team_name=team_name), user=user
+        def _is_authorized_dag_id(dag_id: str):
+            return self.is_authorized_dag(
+                method=method, details=DagDetails(id=dag_id, team_name=team_name), user=user
             )
 
-        return {conn_id for conn_id in conn_ids if _is_authorized_connection(conn_id)}
+        return {dag_id for dag_id in dag_ids if _is_authorized_dag_id(dag_id)}
 
     @provide_session
     def get_authorized_variables(

providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -57,7 +57,7 @@
 )
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException, AirflowException
-from airflow.models import DagModel
+from airflow.models import Connection, DagModel, Pool, Variable
 from airflow.providers.fab.auth_manager.cli_commands.definition import (
     DB_COMMANDS,
     PERMISSIONS_CLEANUP_COMMAND,
@@ -437,6 +437,26 @@ def filter_authorized_menu_items(self, menu_items: list[MenuItem], user: User) -
             )
         ]
 
+    @provide_session
+    def get_authorized_connections(
+        self,
+        *,
+        user: User,
+        method: ResourceMethod = "GET",
+        session: Session = NEW_SESSION,
+    ) -> set[str]:
+        """
+        Get connection ids (``conn_id``) the user has access to.
+
+        Fab auth manager does not allow fine-grained access with connections. Thus, return all the connection ids.
+
+        :param user: the user
+        :param method: the method to filter on
+        :param session: the session
+        """
+        rows = session.execute(select(Connection.conn_id)).scalars().all()
+        return set(rows)
+
     @provide_session
     def get_authorized_dag_ids(
         self,
@@ -477,6 +497,46 @@ def get_authorized_dag_ids(
                         resources.add(resource[len(permissions.RESOURCE_DAG_PREFIX) :])
         return set(session.scalars(select(DagModel.dag_id).where(DagModel.dag_id.in_(resources))))
 
+    @provide_session
+    def get_authorized_pools(
+        self,
+        *,
+        user: User,
+        method: ResourceMethod = "GET",
+        session: Session = NEW_SESSION,
+    ) -> set[str]:
+        """
+        Get pools the user has access to.
+
+        Fab auth manager does not allow fine-grained access with pools. Thus, return all the pool names.
+
+        :param user: the user
+        :param method: the method to filter on
+        :param session: the session
+        """
+        rows = session.execute(select(Pool.pool)).scalars().all()
+        return set(rows)
+
+    @provide_session
+    def get_authorized_variables(
+        self,
+        *,
+        user: User,
+        method: ResourceMethod = "GET",
+        session: Session = NEW_SESSION,
+    ) -> set[str]:
+        """
+        Get variable keys the user has access to.
+
+        Fab auth manager does not allow fine-grained access with variables. Thus, return all the variable keys.
+
+        :param user: the user
+        :param method: the method to filter on
+        :param session: the session
+        """
+        rows = session.execute(select(Variable.key)).scalars().all()
+        return set(rows)
+
     @cached_property
     def security_manager(self) -> FabAirflowSecurityManagerOverride:
         """Return the security manager specific to FAB."""

providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py

@@ -686,6 +686,12 @@ def test_filter_authorized_menu_items(
         result = auth_manager.filter_authorized_menu_items(menu_items, user=user)
         assert result == expected_result
 
+    def test_get_authorized_connections(self, auth_manager):
+        session = Mock()
+        session.execute.return_value.scalars.return_value.all.return_value = ["conn1", "conn2"]
+        result = auth_manager.get_authorized_connections(user=Mock(), method="GET", session=session)
+        assert result == {"conn1", "conn2"}
+
     @pytest.mark.parametrize(
         "method, user_permissions, expected_results",
         [
@@ -778,6 +784,18 @@ def test_get_authorized_dag_ids(
 
         delete_user(flask_app, "username")
 
+    def test_get_authorized_pools(self, auth_manager):
+        session = Mock()
+        session.execute.return_value.scalars.return_value.all.return_value = ["pool1", "pool2"]
+        result = auth_manager.get_authorized_pools(user=Mock(), method="GET", session=session)
+        assert result == {"pool1", "pool2"}
+
+    def test_get_authorized_variables(self, auth_manager):
+        session = Mock()
+        session.execute.return_value.scalars.return_value.all.return_value = ["var1", "var2"]
+        result = auth_manager.get_authorized_variables(user=Mock(), method="GET", session=session)
+        assert result == {"var1", "var2"}
+
     def test_security_manager_return_fab_security_manager_override(self, auth_manager_with_appbuilder):
         assert isinstance(auth_manager_with_appbuilder.security_manager, FabAirflowSecurityManagerOverride)