precedence implementation

This commit is contained in:
2025-02-25 06:10:04 +00:00
parent 8b38a26e8a
commit 355fb55bdb
5 changed files with 42 additions and 10 deletions

View File

@@ -102,11 +102,11 @@ class AirflowClient:
self.base_url = base_url.rstrip("/") self.base_url = base_url.rstrip("/")
self.headers = {"Accept": "application/json"} self.headers = {"Accept": "application/json"}
# Set authentication header based on what was provided # Set authentication header based on precedence (cookie > auth_token)
if auth_token: if cookie:
self.headers["Authorization"] = f"Basic {auth_token}"
elif cookie:
self.headers["Cookie"] = cookie self.headers["Cookie"] = cookie
elif auth_token:
self.headers["Authorization"] = f"Basic {auth_token}"
except Exception as e: except Exception as e:
logger.error("Failed to initialize AirflowClient: %s", e) logger.error("Failed to initialize AirflowClient: %s", e)

View File

@@ -21,7 +21,17 @@ logger = logging.getLogger(__name__)
async def serve() -> None: async def serve() -> None:
"""Start MCP server.""" """Start MCP server.
Configuration precedence:
1. Environment variables (highest)
2. Command line arguments (if applicable)
3. Default values (lowest)
For authentication:
1. Cookie authentication (highest)
2. Auth token authentication (secondary)
"""
# Check for AIRFLOW_BASE_URL which is always required # Check for AIRFLOW_BASE_URL which is always required
if "AIRFLOW_BASE_URL" not in os.environ: if "AIRFLOW_BASE_URL" not in os.environ:
raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL") raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")

View File

@@ -12,7 +12,17 @@ logger = logging.getLogger(__name__)
async def serve() -> None: async def serve() -> None:
"""Start MCP server in safe mode (read-only operations).""" """Start MCP server in safe mode (read-only operations).
Configuration precedence:
1. Environment variables (highest)
2. Command line arguments (if applicable)
3. Default values (lowest)
For authentication:
1. Cookie authentication (highest)
2. Auth token authentication (secondary)
"""
# Check for AIRFLOW_BASE_URL which is always required # Check for AIRFLOW_BASE_URL which is always required
if "AIRFLOW_BASE_URL" not in os.environ: if "AIRFLOW_BASE_URL" not in os.environ:
raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL") raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")

View File

@@ -12,7 +12,17 @@ logger = logging.getLogger(__name__)
async def serve() -> None: async def serve() -> None:
"""Start MCP server in unsafe mode (all operations).""" """Start MCP server in unsafe mode (all operations).
Configuration precedence:
1. Environment variables (highest)
2. Command line arguments (if applicable)
3. Default values (lowest)
For authentication:
1. Cookie authentication (highest)
2. Auth token authentication (secondary)
"""
# Check for AIRFLOW_BASE_URL which is always required # Check for AIRFLOW_BASE_URL which is always required
if "AIRFLOW_BASE_URL" not in os.environ: if "AIRFLOW_BASE_URL" not in os.environ:
raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL") raise ValueError("Missing required environment variable: AIRFLOW_BASE_URL")

View File

@@ -46,10 +46,12 @@ def _initialize_client() -> AirflowClient:
# Initialize client with appropriate authentication method # Initialize client with appropriate authentication method
client_args = {"spec_path": spec_path, "base_url": os.environ["AIRFLOW_BASE_URL"]} client_args = {"spec_path": spec_path, "base_url": os.environ["AIRFLOW_BASE_URL"]}
if has_auth_token: # Apply cookie auth first if available (highest precedence)
client_args["auth_token"] = os.environ["AUTH_TOKEN"] if has_cookie:
elif has_cookie:
client_args["cookie"] = os.environ["COOKIE"] client_args["cookie"] = os.environ["COOKIE"]
# Otherwise use auth token if available
elif has_auth_token:
client_args["auth_token"] = os.environ["AUTH_TOKEN"]
return AirflowClient(**client_args) return AirflowClient(**client_args)