feat: Add jwt token authentication

README.md

@@ -111,15 +111,41 @@ Set the following environment variables:
 
 ```
 AIRFLOW_HOST=<your-airflow-host>        # Optional, defaults to http://localhost:8080
+AIRFLOW_API_VERSION=v1                  # Optional, defaults to v1
+```
+
+#### Authentication
+
+Choose one of the following authentication methods:
+
+**Basic Authentication (default):**
+```
 AIRFLOW_USERNAME=<your-airflow-username>
 AIRFLOW_PASSWORD=<your-airflow-password>
-AIRFLOW_API_VERSION=v1                  # Optional, defaults to v1
 ```
 
+**JWT Token Authentication:**
+```
+AIRFLOW_JWT_TOKEN=<your-jwt-token>
+```
+
+To obtain a JWT token, you can use Airflow's authentication endpoint:
+
+```bash
+ENDPOINT_URL="http://localhost:8080"  # Replace with your Airflow endpoint
+curl -X 'POST' \
+  "${ENDPOINT_URL}/auth/token" \
+  -H 'Content-Type: application/json' \
+  -d '{ "username": "<your-username>", "password": "<your-password>" }'
+```
+
+> **Note**: If both JWT token and basic authentication credentials are provided, JWT token takes precedence.
+
 ### Usage with Claude Desktop
 
 Add to your `claude_desktop_config.json`:
 
+**Basic Authentication:**
 ```json
 {
   "mcpServers": {
@@ -136,8 +162,25 @@ Add to your `claude_desktop_config.json`:
 }
 ```
 
+**JWT Token Authentication:**
+```json
+{
+  "mcpServers": {
+    "mcp-server-apache-airflow": {
+      "command": "uvx",
+      "args": ["mcp-server-apache-airflow"],
+      "env": {
+        "AIRFLOW_HOST": "https://your-airflow-host",
+        "AIRFLOW_JWT_TOKEN": "your-jwt-token"
+      }
+    }
+  }
+}
+```
+
 For read-only mode (recommended for safety):
 
+**Basic Authentication:**
 ```json
 {
   "mcpServers": {
@@ -154,8 +197,25 @@ For read-only mode (recommended for safety):
 }
 ```
 
+**JWT Token Authentication:**
+```json
+{
+  "mcpServers": {
+    "mcp-server-apache-airflow": {
+      "command": "uvx",
+      "args": ["mcp-server-apache-airflow", "--read-only"],
+      "env": {
+        "AIRFLOW_HOST": "https://your-airflow-host",
+        "AIRFLOW_JWT_TOKEN": "your-jwt-token"
+      }
+    }
+  }
+}
+```
+
 Alternative configuration using `uv`:
 
+**Basic Authentication:**
 ```json
 {
   "mcpServers": {
@@ -177,6 +237,27 @@ Alternative configuration using `uv`:
 }
 ```
 
+**JWT Token Authentication:**
+```json
+{
+  "mcpServers": {
+    "mcp-server-apache-airflow": {
+      "command": "uv",
+      "args": [
+        "--directory",
+        "/path/to/mcp-server-apache-airflow",
+        "run",
+        "mcp-server-apache-airflow"
+      ],
+      "env": {
+        "AIRFLOW_HOST": "https://your-airflow-host",
+        "AIRFLOW_JWT_TOKEN": "your-jwt-token"
+      }
+    }
+  }
+}
+```
+
 Replace `/path/to/mcp-server-apache-airflow` with the actual path where you've cloned the repository.
 
 ### Selecting the API groups

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp-server-apache-airflow"
-version = "0.2.8"
+version = "0.2.9"
 description = "Model Context Protocol (MCP) server for Apache Airflow"
 authors = [
     { name = "Gyeongmo Yang", email = "[email protected]" }
@@ -11,6 +11,7 @@ dependencies = [
     "mcp>=0.1.0",
     "apache-airflow-client>=2.7.0,<3",
     "fastmcp>=2.11.3",
+    "PyJWT>=2.8.0",
 ]
 requires-python = ">=3.10"
 readme = "README.md"

src/airflow/airflow_client.py

@@ -5,14 +5,22 @@
 from src.envs import (
     AIRFLOW_API_VERSION,
     AIRFLOW_HOST,
+    AIRFLOW_JWT_TOKEN,
     AIRFLOW_PASSWORD,
     AIRFLOW_USERNAME,
 )
 
 # Create a configuration and API client
 configuration = Configuration(
     host=urljoin(AIRFLOW_HOST, f"/api/{AIRFLOW_API_VERSION}"),
-    username=AIRFLOW_USERNAME,
-    password=AIRFLOW_PASSWORD,
 )
+
+# Set up authentication - prefer JWT token if available, fallback to basic auth
+if AIRFLOW_JWT_TOKEN:
+    configuration.api_key = {"Authorization": f"Bearer {AIRFLOW_JWT_TOKEN}"}
+    configuration.api_key_prefix = {"Authorization": ""}
+elif AIRFLOW_USERNAME and AIRFLOW_PASSWORD:
+    configuration.username = AIRFLOW_USERNAME
+    configuration.password = AIRFLOW_PASSWORD
+
 api_client = ApiClient(configuration)

src/envs.py

@@ -6,6 +6,8 @@
 _airflow_host_raw = os.getenv("AIRFLOW_HOST", "http://localhost:8080")
 AIRFLOW_HOST = urlparse(_airflow_host_raw)._replace(path="").geturl().rstrip("/")
 
+# Authentication - supports both basic auth and JWT token auth
 AIRFLOW_USERNAME = os.getenv("AIRFLOW_USERNAME")
 AIRFLOW_PASSWORD = os.getenv("AIRFLOW_PASSWORD")
+AIRFLOW_JWT_TOKEN = os.getenv("AIRFLOW_JWT_TOKEN")
 AIRFLOW_API_VERSION = os.getenv("AIRFLOW_API_VERSION", "v1")

test/test_airflow_client.py

@@ -0,0 +1,147 @@
+"""Tests for the airflow client authentication module."""
+
+import os
+import sys
+from unittest.mock import patch
+
+from airflow_client.client import ApiClient
+
+
+class TestAirflowClientAuthentication:
+    """Test cases for airflow client authentication configuration."""
+
+    def test_basic_auth_configuration(self):
+        """Test that basic authentication is configured correctly."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIRFLOW_HOST": "http://localhost:8080",
+                "AIRFLOW_USERNAME": "testuser",
+                "AIRFLOW_PASSWORD": "testpass",
+                "AIRFLOW_API_VERSION": "v1",
+            },
+            clear=True,
+        ):
+            # Clear any cached modules
+            modules_to_clear = ["src.envs", "src.airflow.airflow_client"]
+            for module in modules_to_clear:
+                if module in sys.modules:
+                    del sys.modules[module]
+
+            # Re-import after setting environment
+            from src.airflow.airflow_client import api_client, configuration
+
+            # Verify configuration
+            assert configuration.host == "http://localhost:8080/api/v1"
+            assert configuration.username == "testuser"
+            assert configuration.password == "testpass"
+            assert isinstance(api_client, ApiClient)
+
+    def test_jwt_token_auth_configuration(self):
+        """Test that JWT token authentication is configured correctly."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIRFLOW_HOST": "http://localhost:8080",
+                "AIRFLOW_JWT_TOKEN": "test.jwt.token",
+                "AIRFLOW_API_VERSION": "v1",
+            },
+            clear=True,
+        ):
+            # Clear any cached modules
+            modules_to_clear = ["src.envs", "src.airflow.airflow_client"]
+            for module in modules_to_clear:
+                if module in sys.modules:
+                    del sys.modules[module]
+
+            # Re-import after setting environment
+            from src.airflow.airflow_client import api_client, configuration
+
+            # Verify configuration
+            assert configuration.host == "http://localhost:8080/api/v1"
+            assert configuration.api_key == {"Authorization": "Bearer test.jwt.token"}
+            assert configuration.api_key_prefix == {"Authorization": ""}
+            assert isinstance(api_client, ApiClient)
+
+    def test_jwt_token_takes_precedence_over_basic_auth(self):
+        """Test that JWT token takes precedence when both auth methods are provided."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIRFLOW_HOST": "http://localhost:8080",
+                "AIRFLOW_USERNAME": "testuser",
+                "AIRFLOW_PASSWORD": "testpass",
+                "AIRFLOW_JWT_TOKEN": "test.jwt.token",
+                "AIRFLOW_API_VERSION": "v1",
+            },
+            clear=True,
+        ):
+            # Clear any cached modules
+            modules_to_clear = ["src.envs", "src.airflow.airflow_client"]
+            for module in modules_to_clear:
+                if module in sys.modules:
+                    del sys.modules[module]
+
+            # Re-import after setting environment
+            from src.airflow.airflow_client import api_client, configuration
+
+            # Verify JWT token is used (not basic auth)
+            assert configuration.host == "http://localhost:8080/api/v1"
+            assert configuration.api_key == {"Authorization": "Bearer test.jwt.token"}
+            assert configuration.api_key_prefix == {"Authorization": ""}
+            # Basic auth should not be set when JWT is present
+            assert not hasattr(configuration, "username") or configuration.username is None
+            assert not hasattr(configuration, "password") or configuration.password is None
+            assert isinstance(api_client, ApiClient)
+
+    def test_no_auth_configuration(self):
+        """Test that configuration works with no authentication (for testing/development)."""
+        with patch.dict(os.environ, {"AIRFLOW_HOST": "http://localhost:8080", "AIRFLOW_API_VERSION": "v1"}, clear=True):
+            # Clear any cached modules
+            modules_to_clear = ["src.envs", "src.airflow.airflow_client"]
+            for module in modules_to_clear:
+                if module in sys.modules:
+                    del sys.modules[module]
+
+            # Re-import after setting environment
+            from src.airflow.airflow_client import api_client, configuration
+
+            # Verify configuration
+            assert configuration.host == "http://localhost:8080/api/v1"
+            # No auth should be set
+            assert not hasattr(configuration, "username") or configuration.username is None
+            assert not hasattr(configuration, "password") or configuration.password is None
+            # api_key might be an empty dict by default, but should not have Authorization
+            assert "Authorization" not in getattr(configuration, "api_key", {})
+            assert isinstance(api_client, ApiClient)
+
+    def test_environment_variable_parsing(self):
+        """Test that environment variables are parsed correctly."""
+        with patch.dict(
+            os.environ,
+            {
+                "AIRFLOW_HOST": "https://airflow.example.com:8080/custom",
+                "AIRFLOW_JWT_TOKEN": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
+                "AIRFLOW_API_VERSION": "v2",
+            },
+            clear=True,
+        ):
+            # Clear any cached modules
+            modules_to_clear = ["src.envs", "src.airflow.airflow_client"]
+            for module in modules_to_clear:
+                if module in sys.modules:
+                    del sys.modules[module]
+
+            # Re-import after setting environment
+            from src.airflow.airflow_client import configuration
+            from src.envs import AIRFLOW_API_VERSION, AIRFLOW_HOST, AIRFLOW_JWT_TOKEN
+
+            # Verify environment variables are parsed correctly
+            assert AIRFLOW_HOST == "https://airflow.example.com:8080"
+            assert AIRFLOW_JWT_TOKEN == "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9..."
+            assert AIRFLOW_API_VERSION == "v2"
+
+            # Verify configuration uses parsed values
+            assert configuration.host == "https://airflow.example.com:8080/api/v2"
+            assert configuration.api_key == {"Authorization": "Bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9..."}
+            assert configuration.api_key_prefix == {"Authorization": ""}