Merge pull request #18 from redhat-ai-tools/fix-remote-exec

Fix remote setup

Author: natifridman

src/config.py

@@ -16,6 +16,14 @@
 SNOWFLAKE_DATABASE = os.environ.get("SNOWFLAKE_DATABASE")
 SNOWFLAKE_SCHEMA = os.environ.get("SNOWFLAKE_SCHEMA")
 
+# Snowflake token handling - for stdio transport, get from environment
+# For other transports, it will be retrieved from request context in tools layer
+if MCP_TRANSPORT == "stdio":
+    SNOWFLAKE_TOKEN = os.environ.get("SNOWFLAKE_TOKEN")
+else:
+    # For non-stdio transports, token will be passed from tools layer
+    SNOWFLAKE_TOKEN = None
+
 # Prometheus metrics configuration
 ENABLE_METRICS = os.environ.get("ENABLE_METRICS", "false").lower() == "true"
 METRICS_PORT = int(os.environ.get("METRICS_PORT", "8000"))

src/database.py

@@ -3,15 +3,16 @@
 import re
 import time
 import logging
-from typing import Any, List, Dict
+from typing import Any, List, Dict, Optional
 
 import httpx
 
 from config import (
     MCP_TRANSPORT, 
     SNOWFLAKE_BASE_URL, 
     SNOWFLAKE_DATABASE, 
-    SNOWFLAKE_SCHEMA
+    SNOWFLAKE_SCHEMA,
+    SNOWFLAKE_TOKEN
 )
 from metrics import track_snowflake_query
 
@@ -28,22 +29,19 @@ def sanitize_sql_value(value: str) -> str:
 async def make_snowflake_request(
     endpoint: str, 
     method: str = "POST", 
-    data: dict[str, Any] = None
+    data: dict[str, Any] = None,
+    snowflake_token: Optional[str] = None
 ) -> dict[str, Any] | None:
     """Make a request to Snowflake API"""
-    # Get token based on transport type
-    if MCP_TRANSPORT == "stdio":
-        snowflake_token = os.environ.get("SNOWFLAKE_TOKEN")
-    else:
-        # This would need to be passed in or handled differently in non-stdio mode
-        snowflake_token = os.environ.get("SNOWFLAKE_TOKEN")
+    # Use provided token or fall back to config
+    token = snowflake_token or SNOWFLAKE_TOKEN
 
-    if not snowflake_token:
+    if not token:
         logger.error("SNOWFLAKE_TOKEN environment variable is required but not set")
         return None
 
     headers = {
-        "Authorization": f"Bearer {snowflake_token}",
+        "Authorization": f"Bearer {token}",
         "Accept": "application/json",
         "Content-Type": "application/json"
     }
@@ -75,7 +73,7 @@ async def make_snowflake_request(
         logger.error(f"Unexpected error in Snowflake API request: {str(e)}")
         return None
 
-async def execute_snowflake_query(sql: str) -> List[Dict[str, Any]]:
+async def execute_snowflake_query(sql: str, snowflake_token: Optional[str] = None) -> List[Dict[str, Any]]:
     """Execute a SQL query against Snowflake and return results"""
     start_time = time.time()
     success = False
@@ -92,7 +90,7 @@ async def execute_snowflake_query(sql: str) -> List[Dict[str, Any]]:
 
         logger.info(f"Executing Snowflake query: {sql[:100]}...")  # Log first 100 chars of query
 
-        response = await make_snowflake_request(endpoint, "POST", payload)
+        response = await make_snowflake_request(endpoint, "POST", payload, snowflake_token)
 
         # Check if response is None (indicating an error in API request or JSON parsing)
         if response is None:
@@ -130,7 +128,7 @@ def format_snowflake_row(row_data: List[Any], columns: List[str]) -> Dict[str, A
 
     return {columns[i]: row_data[i] for i in range(len(columns))}
 
-async def get_issue_labels(issue_ids: List[str]) -> Dict[str, List[str]]:
+async def get_issue_labels(issue_ids: List[str], snowflake_token: Optional[str] = None) -> Dict[str, List[str]]:
     """Get labels for given issue IDs from Snowflake"""
     if not issue_ids:
         return {}
@@ -157,7 +155,7 @@ async def get_issue_labels(issue_ids: List[str]) -> Dict[str, List[str]]:
         WHERE ISSUE IN ({ids_str}) AND LABEL IS NOT NULL
         """
 
-        rows = await execute_snowflake_query(sql)
+        rows = await execute_snowflake_query(sql, snowflake_token)
         columns = ["ISSUE", "LABEL"]
 
         for row in rows:

src/tools.py

@@ -3,6 +3,7 @@
 
 from mcp.server.fastmcp import FastMCP
 
+from config import MCP_TRANSPORT, SNOWFLAKE_TOKEN
 from database import (
     execute_snowflake_query,
     format_snowflake_row,
@@ -13,12 +14,27 @@
 
 logger = logging.getLogger(__name__)
 
+def get_snowflake_token(mcp: FastMCP) -> Optional[str]:
+    """Get Snowflake token from either config (stdio) or request headers (non-stdio)"""
+    if MCP_TRANSPORT == "stdio":
+        return SNOWFLAKE_TOKEN
+    else:
+        try:
+            # Get token from request headers for non-stdio transports
+            context = mcp.get_context()
+            if context and hasattr(context, 'request_context') and context.request_context:
+                headers = context.request_context.request.headers
+                return headers.get("X-Snowflake-Token")
+        except Exception as e:
+            logger.error(f"Error getting token from request context: {e}")
+        return None
+
 def register_tools(mcp: FastMCP) -> None:
     """Register all MCP tools"""
 
     @mcp.tool()
-    @track_tool_usage("list_issues")
-    async def list_issues(
+    @track_tool_usage("list_jira_issues")
+    async def list_jira_issues(
         project: Optional[str] = None,
         issue_type: Optional[str] = None,
         status: Optional[str] = None,
@@ -41,6 +57,11 @@ async def list_issues(
             Dictionary containing issues list and metadata
         """
         try:
+            # Get the Snowflake token
+            snowflake_token = get_snowflake_token(mcp)
+            if not snowflake_token:
+                return {"error": "Snowflake token not available", "issues": []}
+            
             # Build SQL query with filters
             sql_conditions = []
 
@@ -77,7 +98,7 @@ async def list_issues(
             LIMIT {limit}
             """
 
-            rows = await execute_snowflake_query(sql)
+            rows = await execute_snowflake_query(sql, snowflake_token)
 
             issues = []
             issue_ids = []
@@ -121,7 +142,7 @@ async def list_issues(
                     issue_ids.append(str(row_dict.get("ID")))
 
             # Get labels for enrichment
-            labels_data = await get_issue_labels(issue_ids)
+            labels_data = await get_issue_labels(issue_ids, snowflake_token)
 
             # Enrich issues with labels
             for issue in issues:
@@ -145,8 +166,8 @@ async def list_issues(
             return {"error": f"Error reading issues from Snowflake: {str(e)}", "issues": []}
 
     @mcp.tool()
-    @track_tool_usage("get_issue_details")
-    async def get_issue_details(issue_key: str) -> Dict[str, Any]:
+    @track_tool_usage("get_jira_issue_details")
+    async def get_jira_issue_details(issue_key: str) -> Dict[str, Any]:
         """
         Get detailed information for a specific JIRA issue by its key from Snowflake.
         
@@ -157,6 +178,11 @@ async def get_issue_details(issue_key: str) -> Dict[str, Any]:
             Dictionary containing detailed issue information
         """
         try:
+            # Get the Snowflake token
+            snowflake_token = get_snowflake_token(mcp)
+            if not snowflake_token:
+                return {"error": "Snowflake token not available"}
+            
             sql = f"""
             SELECT 
                 ID, ISSUE_KEY, PROJECT, ISSUENUM, ISSUETYPE, SUMMARY, DESCRIPTION,
@@ -169,7 +195,7 @@ async def get_issue_details(issue_key: str) -> Dict[str, Any]:
             LIMIT 1
             """
 
-            rows = await execute_snowflake_query(sql)
+            rows = await execute_snowflake_query(sql, snowflake_token)
 
             if not rows:
                 return {"error": f"Issue with key '{issue_key}' not found"}
@@ -215,24 +241,29 @@ async def get_issue_details(issue_key: str) -> Dict[str, Any]:
             }
 
             # Get labels for this issue
-            labels_data = await get_issue_labels([str(issue['id'])])
+            labels_data = await get_issue_labels([str(issue['id'])], snowflake_token)
             issue['labels'] = labels_data.get(str(issue['id']), [])
 
-            return {"issue": issue}
+            return issue
 
         except Exception as e:
-            return {"error": f"Error retrieving issue details from Snowflake: {str(e)}"}
+            return {"error": f"Error reading issue details from Snowflake: {str(e)}"}
 
     @mcp.tool()
-    @track_tool_usage("get_project_summary")
-    async def get_project_summary() -> Dict[str, Any]:
+    @track_tool_usage("get_jira_project_summary")
+    async def get_jira_project_summary() -> Dict[str, Any]:
         """
         Get a summary of all projects in the JIRA data from Snowflake.
         
         Returns:
             Dictionary containing project statistics
         """
         try:
+            # Get the Snowflake token
+            snowflake_token = get_snowflake_token(mcp)
+            if not snowflake_token:
+                return {"error": "Snowflake token not available"}
+            
             sql = """
             SELECT 
                 PROJECT,
@@ -244,7 +275,7 @@ async def get_project_summary() -> Dict[str, Any]:
             ORDER BY PROJECT, ISSUESTATUS, PRIORITY
             """
 
-            rows = await execute_snowflake_query(sql)
+            rows = await execute_snowflake_query(sql, snowflake_token)
             columns = ["PROJECT", "ISSUESTATUS", "PRIORITY", "COUNT"]
 
             project_stats = {}
@@ -256,7 +287,7 @@ async def get_project_summary() -> Dict[str, Any]:
                 project = row_dict.get("PROJECT", "Unknown")
                 status = row_dict.get("ISSUESTATUS", "Unknown")
                 priority = row_dict.get("PRIORITY", "Unknown")
-                count = row_dict.get("COUNT", 0)
+                count = int(row_dict.get("COUNT", 0)) if row_dict.get("COUNT") is not None else 0
 
                 if project not in project_stats:
                     project_stats[project] = {