Merge pull request #5 from makenotion/atvaccaro/query-comment-template

feat: add set_query_context tool for runtime observability

Author: atvaccaro

mcp_server_snowflake/query_manager/tools.py

@@ -8,7 +8,9 @@
 from mcp_server_snowflake.utils import SnowflakeException
 
 
-def run_query(statement: str, snowflake_service):
+def run_query(
+    statement: str, snowflake_service, tool_name: str = "run_snowflake_query"
+):
     """
     Execute SQL statement and fetch all results using Snowflake connector.
 
@@ -21,6 +23,8 @@ def run_query(statement: str, snowflake_service):
         SQL statement to execute
     snowflake_service : SnowflakeService
         The Snowflake service instance to use for connection
+    tool_name : str
+        Name of the tool executing the query (for query comments)
 
     Returns
     -------
@@ -33,14 +37,29 @@ def run_query(statement: str, snowflake_service):
         If connection fails or SQL execution encounters an error
     """
     try:
+        # Get statement type for query comment
+        statement_type = get_statement_type(statement)
+
+        # Build query comment if enabled
+        query_comment = snowflake_service.build_query_comment(
+            tool_name=tool_name,
+            statement_type=statement_type,
+        )
+
+        # Prepend comment to statement if enabled
+        if query_comment:
+            statement_with_comment = f"/* {query_comment} */\n{statement}"
+        else:
+            statement_with_comment = statement
+
         with snowflake_service.get_connection(
             use_dict_cursor=True,
             session_parameters=snowflake_service.get_query_tag_param(),
         ) as (
             con,
             cur,
         ):
-            cur.execute(statement)
+            cur.execute(statement_with_comment)
             return cur.fetchall()
     except Exception as e:
         raise SnowflakeException(
@@ -63,6 +82,121 @@ def run_query_tool(
     ):
         return run_query(statement, snowflake_service)
 
+    @server.tool(
+        name="set_query_context",
+        description="""Set runtime context for query comments and observability.
+
+Call this tool at the start of a session to register context information that will be
+included in all subsequent SQL query comments. This enables tracking queries back to
+specific agents, models, or sessions in Snowflake's query history.
+
+Common context keys:
+- model: The AI model name (e.g., "claude-sonnet-4-5-20250929")
+- agent_name: Name of the agent or application (e.g., "Claude Code")
+- user_email: Email of the user running the agent
+- user_name: Name of the user running the agent
+- intent: Object describing query intent (category, confidence, domains, question)
+- query_parameters: Object describing query details (datasets, dimensions, time_range)
+- session_id: A unique session identifier for grouping related queries
+
+Context persists for the lifetime of the MCP server connection. Call this tool again
+to update intent/query_parameters for different queries.""",
+    )
+    def set_query_context_tool(
+        model: Annotated[
+            str,
+            Field(
+                default=None,
+                description="AI model name (e.g., 'claude-sonnet-4-5-20250929')",
+            ),
+        ] = None,
+        agent_name: Annotated[
+            str,
+            Field(
+                default=None,
+                description="Name of the agent or application (e.g., 'Claude Code')",
+            ),
+        ] = None,
+        user_email: Annotated[
+            str,
+            Field(
+                default=None,
+                description="Email of the user running the agent",
+            ),
+        ] = None,
+        user_name: Annotated[
+            str,
+            Field(
+                default=None,
+                description="Name of the user running the agent",
+            ),
+        ] = None,
+        intent: Annotated[
+            dict,
+            Field(
+                default=None,
+                description="Query intent: {category, confidence, domains, question}",
+            ),
+        ] = None,
+        query_parameters: Annotated[
+            dict,
+            Field(
+                default=None,
+                description="Query parameters: {datasets, dimensions, time_range}",
+            ),
+        ] = None,
+        session_id: Annotated[
+            str,
+            Field(
+                default=None,
+                description="Unique session identifier for grouping queries",
+            ),
+        ] = None,
+        custom_context: Annotated[
+            dict,
+            Field(
+                default=None,
+                description="Additional custom key-value pairs for context",
+            ),
+        ] = None,
+    ):
+        """Set query context for observability."""
+        context = {}
+        if model is not None:
+            context["model"] = model
+        if agent_name is not None:
+            context["agent_name"] = agent_name
+        if user_email is not None:
+            context["user_email"] = user_email
+        if user_name is not None:
+            context["user_name"] = user_name
+        if intent is not None:
+            context["intent"] = intent
+        if query_parameters is not None:
+            context["query_parameters"] = query_parameters
+        if session_id is not None:
+            context["session_id"] = session_id
+        if custom_context is not None:
+            context.update(custom_context)
+
+        updated_context = snowflake_service.set_query_context(**context)
+        return {
+            "status": "success",
+            "message": "Query context updated successfully",
+            "context": updated_context,
+        }
+
+    @server.tool(
+        name="get_query_context",
+        description="Get the current query context that will be included in query comments.",
+    )
+    def get_query_context_tool():
+        """Get current query context."""
+        return {
+            "context": snowflake_service.get_query_context(),
+            "query_comments_enabled": snowflake_service.query_comment_enabled,
+        }
+
 
 def get_statement_type(sql_string):
     """

mcp_server_snowflake/server.py

@@ -12,8 +12,10 @@
 import argparse
 import json
 import os
+import uuid
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Dict, Generator, Literal, Optional, Tuple, cast
 
@@ -52,6 +54,26 @@
 tag_minor_version = 3
 query_tag = {"origin": "sf_sit", "name": "mcp_server"}
 
+# Default query comment template - matches dbt query tag format for observability
+DEFAULT_QUERY_COMMENT_TEMPLATE = {
+    "agent": "{agent_name}",
+    "context": {
+        "request_id": "{request_id}",
+        "timestamp": "{timestamp}",
+    },
+    "intent": "{intent}",
+    "model_version": "{model}",
+    "query": {
+        "tool": "{tool_name}",
+        "statement_type": "{statement_type}",
+    },
+    "query_parameters": "{query_parameters}",
+    "user": {
+        "email": "{user_email}",
+        "name": "{user_name}",
+    },
+}
+
 logger = get_logger(server_name)
 
 
@@ -130,6 +152,10 @@ def __init__(
         self.semantic_manager = False
         self.default_session_parameters: Dict[str, Any] = {}
         self.query_tag = query_tag if query_tag is not None else None
+        self.query_comment_template: Optional[Dict[str, Any]] = None
+        self.query_comment_enabled = False
+        # Runtime query context set by agents via set_query_context tool
+        self.query_context: Dict[str, str] = {}
         self.tag_major_version = (
             tag_major_version if tag_major_version is not None else None
         )
@@ -184,6 +210,16 @@ def unpack_service_specs(self) -> None:
                 self.query_manager = other_services.get("query_manager", False)
                 self.semantic_manager = other_services.get("semantic_manager", False)
 
+            # Parse query comment configuration
+            query_comment_config = service_config.get("query_comment", {})
+            if query_comment_config:
+                self.query_comment_enabled = query_comment_config.get("enabled", False)
+                custom_template = query_comment_config.get("template")
+                if custom_template:
+                    self.query_comment_template = custom_template
+                elif self.query_comment_enabled:
+                    self.query_comment_template = DEFAULT_QUERY_COMMENT_TEMPLATE.copy()
+
         except Exception as e:
             logger.error(f"Error extracting service specifications: {e}")
             raise
@@ -414,6 +450,139 @@ def get_query_tag_param(
         else:
             return None
 
+    def set_query_context(self, **kwargs: str) -> Dict[str, str]:
+        """
+        Set runtime query context values for query comments.
+
+        This method allows agents to set context information (like model name,
+        session ID, etc.) that will be included in subsequent query comments.
+        Context values override environment variables.
+
+        Parameters
+        ----------
+        **kwargs : str
+            Key-value pairs to set in the query context.
+            Common keys: model, session_id, agent_name, user_context
+
+        Returns
+        -------
+        Dict[str, str]
+            The updated query context dictionary
+        """
+        self.query_context.update(kwargs)
+        return self.query_context.copy()
+
+    def get_query_context(self) -> Dict[str, str]:
+        """
+        Get the current runtime query context.
+
+        Returns
+        -------
+        Dict[str, str]
+            Current query context dictionary
+        """
+        return self.query_context.copy()
+
+    def clear_query_context(self) -> None:
+        """Clear all runtime query context values."""
+        self.query_context.clear()
+
+    def build_query_comment(
+        self,
+        tool_name: str = "unknown",
+        statement_type: str = "unknown",
+    ) -> Optional[str]:
+        """
+        Build a query comment string with template variable substitution.
+
+        Substitutes template variables in the query comment template with actual values.
+        Supported variables:
+        - {request_id}: Unique UUID for this request
+        - {timestamp}: ISO 8601 timestamp
+        - {tool_name}: Name of the MCP tool being used
+        - {statement_type}: Type of SQL statement (Select, Insert, etc.)
+        - {model}: AI model name (from query_context, env var, or 'unknown')
+        - {session_id}: Session ID (from query_context or 'unknown')
+        - {agent_name}: Agent name (from query_context or 'unknown')
+        - {server_name}: MCP server name
+        - {server_version}: Server version string
+        - Any custom keys set via set_query_context
+
+        Parameters
+        ----------
+        tool_name : str
+            Name of the MCP tool making the query
+        statement_type : str
+            Type of SQL statement being executed
+
+        Returns
+        -------
+        str or None
+            JSON string of the query comment, or None if disabled
+        """
+        if not self.query_comment_enabled or self.query_comment_template is None:
+            return None
+
+        # Build substitution values - runtime context takes precedence over env vars
+        substitutions = {
+            "request_id": str(uuid.uuid4()),
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "tool_name": tool_name,
+            "statement_type": statement_type,
+            # Model: check runtime context first, then env var, then default
+            "model": self.query_context.get(
+                "model", os.environ.get("SNOWFLAKE_MCP_MODEL", "unknown")
+            ),
+            # Session ID: from runtime context or default
+            "session_id": self.query_context.get("session_id", "unknown"),
+            # Agent name: from runtime context or default to server name
+            "agent_name": self.query_context.get("agent_name", server_name),
+            # User info: from runtime context
+            "user_email": self.query_context.get(
+                "user_email", os.environ.get("SNOWFLAKE_MCP_USER_EMAIL", "unknown")
+            ),
+            "user_name": self.query_context.get(
+                "user_name", os.environ.get("SNOWFLAKE_MCP_USER_NAME", "unknown")
+            ),
+            # Intent and query_parameters: complex objects from agent (default to null)
+            "intent": self.query_context.get("intent"),
+            "query_parameters": self.query_context.get("query_parameters"),
+            "server_name": server_name,
+            "server_version": f"{tag_major_version}.{tag_minor_version}",
+        }
+        # Add any additional custom context values
+        for key, value in self.query_context.items():
+            if key not in substitutions:
+                substitutions[key] = value
+
+        def substitute_value(value: Any) -> Any:
+            """Recursively substitute template variables in values."""
+            if isinstance(value, str):
+                # Check if the entire string is a single placeholder like "{intent}"
+                # If so, return the actual value (could be dict, None, etc.)
+                import re
+
+                match = re.fullmatch(r"\{(\w+)\}", value)
+                if match:
+                    key = match.group(1)
+                    if key in substitutions:
+                        return substitutions[key]
+                # Otherwise do string replacement
+                result = value
+                for key, sub_value in substitutions.items():
+                    if sub_value is not None:
+                        result = result.replace(f"{{{key}}}", str(sub_value))
+                return result
+            elif isinstance(value, dict):
+                return {k: substitute_value(v) for k, v in value.items()}
+            elif isinstance(value, list):
+                return [substitute_value(item) for item in value]
+            else:
+                return value
+
+        comment = substitute_value(self.query_comment_template)
+        return json.dumps(comment)
+
 
 def get_var(var_name: str, env_var_name: str, args) -> Optional[str]:
     """