abhishekbhakat/airflow-mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved Debug logging

Browse Source
This commit is contained in:
abhishekbhakat
2025-02-14 10:13:44 +00:00
parent 289a20650d
commit 509cdedbc5
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
airflow-mcp-server/src/airflow_mcp_server/tools/airflow_tool.py
View File

@@ -56,6 +56,7 @@ class AirflowTool(BaseTools):
"""Execute the operation with provided parameters.""" """Execute the operation with provided parameters."""
try: try:
mapping = self.operation.input_model.model_config["parameter_mapping"] mapping = self.operation.input_model.model_config["parameter_mapping"]
body = body or {}
path_params = {k: body[k] for k in mapping.get("path", []) if k in body} path_params = {k: body[k] for k in mapping.get("path", []) if k in body}
query_params = {k: body[k] for k in mapping.get("query", []) if k in body} query_params = {k: body[k] for k in mapping.get("query", []) if k in body}
body_params = {k: body[k] for k in mapping.get("body", []) if k in body} body_params = {k: body[k] for k in mapping.get("body", []) if k in body}
@@ -68,14 +69,22 @@ class AirflowTool(BaseTools):
body=body_params, body=body_params,
) )
logger.debug("Raw response: %s", response)
# Validate response if model exists # Validate response if model exists
if self.operation.response_model and isinstance(response, dict): if self.operation.response_model and isinstance(response, dict):
try: try:
logger.debug("Response model schema: %s", self.operation.response_model.model_json_schema())
model: type[BaseModel] = self.operation.response_model model: type[BaseModel] = self.operation.response_model
logger.debug("Attempting to validate response with model: %s", model.__name__)
validated_response = model(**response) validated_response = model(**response)
return validated_response.model_dump() logger.debug("Response validation successful")
result = validated_response.model_dump()
logger.debug("Final response after model_dump: %s", result)
return result
except ValidationError as e: except ValidationError as e:
logger.error("Response validation failed: %s", e) logger.error("Response validation failed: %s", e)
logger.error("Validation error details: %s", e.errors())
raise RuntimeError(f"Invalid response format: {e}") raise RuntimeError(f"Invalid response format: {e}")
return response return response