diff --git a/README.md b/README.md index 9e85248..d27293f 100644 --- a/README.md +++ b/README.md @@ -111,15 +111,41 @@ Set the following environment variables: ``` 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= AIRFLOW_PASSWORD= -AIRFLOW_API_VERSION=v1 # Optional, defaults to v1 ``` +**JWT Token Authentication:** +``` +AIRFLOW_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": "", "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 diff --git a/pyproject.toml b/pyproject.toml index 3b8e97e..817cffc 100644 --- a/pyproject.toml +++ b/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 = "me@gmyang.dev" } @@ -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" diff --git a/src/airflow/airflow_client.py b/src/airflow/airflow_client.py index 902fbdf..7111431 100644 --- a/src/airflow/airflow_client.py +++ b/src/airflow/airflow_client.py @@ -5,6 +5,7 @@ from src.envs import ( AIRFLOW_API_VERSION, AIRFLOW_HOST, + AIRFLOW_JWT_TOKEN, AIRFLOW_PASSWORD, AIRFLOW_USERNAME, ) @@ -12,7 +13,14 @@ # 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) diff --git a/src/envs.py b/src/envs.py index 3870f95..a6407e6 100644 --- a/src/envs.py +++ b/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") diff --git a/test/test_airflow_client.py b/test/test_airflow_client.py new file mode 100644 index 0000000..79c357c --- /dev/null +++ b/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": ""} diff --git a/uv.lock b/uv.lock index 7543ce6..d588f59 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.13'", @@ -750,7 +750,7 @@ wheels = [ [[package]] name = "mcp-server-apache-airflow" -version = "0.2.7" +version = "0.2.9" source = { editable = "." } dependencies = [ { name = "apache-airflow-client" }, @@ -758,6 +758,7 @@ dependencies = [ { name = "fastmcp" }, { name = "httpx" }, { name = "mcp" }, + { name = "pyjwt" }, ] [package.optional-dependencies] @@ -782,6 +783,7 @@ requires-dist = [ { name = "fastmcp", specifier = ">=2.11.3" }, { name = "httpx", specifier = ">=0.24.1" }, { name = "mcp", specifier = ">=0.1.0" }, + { name = "pyjwt", specifier = ">=2.8.0" }, { name = "twine", marker = "extra == 'dev'", specifier = ">=6.1.0" }, ] provides-extras = ["dev"] @@ -1078,6 +1080,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293, upload-time = "2025-01-06T17:26:25.553Z" }, ] +[[package]] +name = "pyjwt" +version = "2.10.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/46/bd74733ff231675599650d3e47f361794b22ef3e3770998dda30d3b63726/pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953", size = 87785, upload-time = "2024-11-28T03:43:29.933Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997, upload-time = "2024-11-28T03:43:27.893Z" }, +] + [[package]] name = "pyperclip" version = "1.9.0"