(airflow) branch main updated: feat(keycloak): add method to retrieve teams from Keycloak as resources (#62715)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new c6d8e167bd5 feat(keycloak): add method to retrieve teams from Keycloak as resources (#62715)
c6d8e167bd5 is described below

commit c6d8e167bd5693f61a7b8e19b706204b89b0ec67
Author: Mathieu Monet <60...@users.noreply.github.com>
AuthorDate: Mon Mar 2 21:33:40 2026 +0100

    feat(keycloak): add method to retrieve teams from Keycloak as resources (#62715)
---
 .../keycloak/auth_manager/keycloak_auth_manager.py | 18 +++++++++++++++
 .../auth_manager/test_keycloak_auth_manager.py     | 27 ++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
index 11f2c55b0f7..bc398020844 100644
--- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
+++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
@@ -471,6 +471,24 @@ class KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
             )
         raise AirflowException(f"Unexpected error: {resp.status_code} - {resp.text}")
 
+    def _get_teams(self) -> set[str]:
+        realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+        server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
+
+        pat = self.get_keycloak_client().token(grant_type="client_credentials")["access_token"]
+
+        prefix = f"{KeycloakResource.TEAM.value}:"
+        resource_url = f"{server_url.rstrip('/')}/realms/{realm}/authz/protection/resource_set"
+        resources_resp = self.http_session.get(
+            resource_url,
+            params={"name": prefix, "matchingUri": "false", "max": "-1", "deep": "true"},
+            headers={"Authorization": f"Bearer {pat}"},
+            timeout=5,
+        )
+        resources_resp.raise_for_status()
+
+        return {r["name"][len(prefix) :] for r in resources_resp.json() if r["name"].startswith(prefix)}
+
     @staticmethod
     def _get_token_url(server_url, realm):
         # Normalize server_url to avoid double slashes (required for Keycloak 26.4+ strict path validation).
diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
index 4610b45bf63..b1c28714f6b 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
@@ -933,3 +933,30 @@ class TestKeycloakAuthManager:
         """Test that _get_token_url normalizes server_url by stripping trailing slashes."""
         token_url = auth_manager._get_token_url(server_url, "myrealm")
         assert token_url == expected_url
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Team not available before Airflow 3.2.0")
+    @patch(
+        "airflow.providers.keycloak.auth_manager.keycloak_auth_manager.KeycloakAuthManager.get_keycloak_client"
+    )
+    def test_get_teams(self, mock_get_keycloak_client, auth_manager_multi_team):
+        """_get_teams fetches Team: resources from Keycloak and returns team names."""
+        mock_get_keycloak_client.return_value.token.return_value = {"access_token": "pat-token"}
+
+        mock_response = Mock()
+        mock_response.json.return_value = [
+            {"name": "Team:team-a", "type": "urn:airflow:resource"},
+            {"name": "Team:team-b", "type": "urn:airflow:resource"},
+            {"name": "Connection", "type": "urn:airflow:resource"},  # should be ignored
+        ]
+        mock_response.raise_for_status = Mock()
+        auth_manager_multi_team.http_session.get = Mock(return_value=mock_response)
+
+        result = auth_manager_multi_team._get_teams()
+
+        assert result == {"team-a", "team-b"}
+        auth_manager_multi_team.http_session.get.assert_called_once_with(
+            "server_url/realms/realm/authz/protection/resource_set",
+            params={"name": "Team:", "matchingUri": "false", "max": "-1", "deep": "true"},
+            headers={"Authorization": "Bearer pat-token"},
+            timeout=5,
+        )