Integrate KeycloakAuthManager with airflowctl

providers/keycloak/src/airflow/providers/keycloak/auth_manager/openapi/v2-keycloak-auth-manager-generated.yaml

@@ -85,6 +85,43 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/HTTPValidationError'
+  /auth/token/cli:
+    post:
+      tags:
+      - KeycloakAuthManagerToken
+      summary: Create Token Cli
+      operationId: create_token_cli
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/TokenBody'
+        required: true
+      responses:
+        '201':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/TokenResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '401':
+          description: Unauthorized
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     HTTPExceptionResponse:

providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/token.py

@@ -19,16 +19,13 @@
 
 import logging
 
-from fastapi import HTTPException
-from keycloak import KeycloakAuthenticationError
 from starlette import status
 
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.configuration import conf
 from airflow.providers.keycloak.auth_manager.datamodels.token import TokenBody, TokenResponse
-from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager
-from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser
+from airflow.providers.keycloak.auth_manager.services.token import create_token_for
 
 log = logging.getLogger(__name__)
 token_router = AirflowRouter(tags=["KeycloakAuthManagerToken"])
@@ -40,23 +37,19 @@
     responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]),
 )
 def create_token(body: TokenBody) -> TokenResponse:
-    client = KeycloakAuthManager.get_keycloak_client()
-
-    try:
-        tokens = client.token(body.username, body.password)
-    except KeycloakAuthenticationError:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid credentials",
-        )
-
-    userinfo = client.userinfo(tokens["access_token"])
-    user = KeycloakAuthManagerUser(
-        user_id=userinfo["sub"],
-        name=userinfo["preferred_username"],
-        access_token=tokens["access_token"],
-        refresh_token=tokens["refresh_token"],
-    )
-    token = get_auth_manager().generate_jwt(user)
+    token = create_token_for(body.username, body.password)
+    return TokenResponse(access_token=token)
 
+
+@token_router.post(
+    "/token/cli",
+    status_code=status.HTTP_201_CREATED,
+    responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_401_UNAUTHORIZED]),
+)
+def create_token_cli(body: TokenBody) -> TokenResponse:
+    token = create_token_for(
+        body.username,
+        body.password,
+        expiration_time_in_seconds=int(conf.getint("api_auth", "jwt_cli_expiration_time")),
+    )
     return TokenResponse(access_token=token)

providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

providers/keycloak/src/airflow/providers/keycloak/auth_manager/services/token.py

@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from fastapi import HTTPException
+from keycloak import KeycloakAuthenticationError
+from starlette import status
+
+from airflow.api_fastapi.app import get_auth_manager
+from airflow.configuration import conf
+from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager
+from airflow.providers.keycloak.auth_manager.user import KeycloakAuthManagerUser
+
+
+def create_token_for(
+    username: str,
+    password: str,
+    expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time"),
+) -> str:
+    client = KeycloakAuthManager.get_keycloak_client()
+
+    try:
+        tokens = client.token(username, password)
+    except KeycloakAuthenticationError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid credentials",
+        )
+
+    userinfo = client.userinfo(tokens["access_token"])
+    user = KeycloakAuthManagerUser(
+        user_id=userinfo["sub"],
+        name=userinfo["preferred_username"],
+        access_token=tokens["access_token"],
+        refresh_token=tokens["refresh_token"],
+    )
+
+    return get_auth_manager().generate_jwt(user, expiration_time_in_seconds=expiration_time_in_seconds)

providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_token.py

@@ -16,41 +16,46 @@
 # under the License.
 from __future__ import annotations
 
-from unittest.mock import Mock, patch
-
-from keycloak import KeycloakAuthenticationError
+from unittest.mock import patch
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 
+from tests_common.test_utils.config import conf_vars
+
 
 class TestTokenRouter:
-    @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client")
-    def test_create_token(self, mock_get_keycloak_client, client):
-        mock_keycloak_client = Mock()
-        mock_keycloak_client.token.return_value = {
-            "access_token": "access_token",
-            "refresh_token": "refresh_token",
+    token = "token"
+    token_body_dict = {"username": "username", "password": "password"}
+
+    @conf_vars(
+        {
+            ("api_auth", "jwt_expiration_time"): "10",
         }
-        mock_keycloak_client.userinfo.return_value = {"sub": "sub", "preferred_username": "username"}
-        mock_get_keycloak_client.return_value = mock_keycloak_client
+    )
+    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for")
+    def test_create_token(self, mock_create_token_for, client):
+        mock_create_token_for.return_value = self.token
         response = client.post(
             AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token",
-            json={"username": "username", "password": "password"},
+            json=self.token_body_dict,
         )
 
         assert response.status_code == 201
-        mock_keycloak_client.token.assert_called_once_with("username", "password")
-        mock_keycloak_client.userinfo.assert_called_once_with("access_token")
-
-    @patch("airflow.providers.keycloak.auth_manager.routes.login.KeycloakAuthManager.get_keycloak_client")
-    def test_create_token_with_invalid_creds(self, mock_get_keycloak_client, client):
-        mock_keycloak_client = Mock()
-        mock_keycloak_client.token.side_effect = KeycloakAuthenticationError()
-        mock_get_keycloak_client.return_value = mock_keycloak_client
+        assert response.json() == {"access_token": self.token}
+
+    @conf_vars(
+        {
+            ("api_auth", "jwt_cli_expiration_time"): "10",
+            ("api_auth", "jwt_expiration_time"): "10",
+        }
+    )
+    @patch("airflow.providers.keycloak.auth_manager.routes.token.create_token_for")
+    def test_create_token_cli(self, mock_create_token_for, client):
+        mock_create_token_for.return_value = self.token
         response = client.post(
-            AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token",
-            json={"username": "username", "password": "password"},
+            AUTH_MANAGER_FASTAPI_APP_PREFIX + "/token/cli",
+            json=self.token_body_dict,
         )
 
-        assert response.status_code == 401
-        mock_keycloak_client.token.assert_called_once_with("username", "password")
+        assert response.status_code == 201
+        assert response.json() == {"access_token": self.token}

providers/keycloak/tests/unit/keycloak/auth_manager/services/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

providers/keycloak/tests/unit/keycloak/auth_manager/services/test_token.py

@@ -0,0 +1,78 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import fastapi
+import pytest
+from keycloak import KeycloakAuthenticationError
+
+from airflow.configuration import conf
+from airflow.providers.keycloak.auth_manager.services.token import create_token_for
+
+from tests_common.test_utils.config import conf_vars
+
+
+class TestTokenService:
+    token = "token"
+    test_username = "test_user"
+    test_password = "test_pass"
+    test_access_token = "access_token"
+    test_refresh_token = "refresh_token"
+
+    @conf_vars(
+        {
+            ("api_auth", "jwt_expiration_time"): "10",
+        }
+    )
+    @patch("airflow.providers.keycloak.auth_manager.services.token.get_auth_manager")
+    @patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+    def test_create_token(self, mock_get_keycloak_client, mock_get_auth_manager):
+        mock_keycloak_client = Mock()
+        mock_keycloak_client.token.return_value = {
+            "access_token": self.test_access_token,
+            "refresh_token": self.test_refresh_token,
+        }
+        mock_keycloak_client.userinfo.return_value = {"sub": "sub", "preferred_username": "username"}
+        mock_get_keycloak_client.return_value = mock_keycloak_client
+        mock_auth_manager = Mock()
+        mock_get_auth_manager.return_value = mock_auth_manager
+        mock_auth_manager.generate_jwt.return_value = self.token
+
+        assert create_token_for(username=self.test_username, password=self.test_password) == self.token
+        mock_keycloak_client.token.assert_called_once_with(self.test_username, self.test_password)
+        mock_keycloak_client.userinfo.assert_called_once_with(self.test_access_token)
+
+    @conf_vars(
+        {
+            ("api_auth", "jwt_cli_expiration_time"): "10",
+            ("api_auth", "jwt_expiration_time"): "10",
+        }
+    )
+    @patch("airflow.providers.keycloak.auth_manager.services.token.KeycloakAuthManager.get_keycloak_client")
+    def test_create_token_with_invalid_creds(self, mock_get_keycloak_client):
+        mock_keycloak_client = Mock()
+        mock_keycloak_client.token.side_effect = KeycloakAuthenticationError()
+        mock_get_keycloak_client.return_value = mock_keycloak_client
+
+        with pytest.raises(fastapi.exceptions.HTTPException):
+            create_token_for(
+                username=self.test_username,
+                password=self.test_password,
+                expiration_time_in_seconds=conf.getint("api_auth", "jwt_cli_expiration_time"),
+            )