feat: add intent and query_parameters fields to match dbt schema

Author: atvaccaro

mcp_server_snowflake/query_manager/tools.py

@@ -91,14 +91,16 @@ def run_query_tool(
 specific agents, models, or sessions in Snowflake's query history.
 
 Common context keys:
-- model: The AI model name (e.g., "claude-sonnet-4", "gpt-4")
-- agent_name: Name of the agent or application making queries (e.g., "Claude Code")
+- 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
 
-You can also set custom key-value pairs that will be available as template variables.
-Context persists for the lifetime of the MCP server connection.""",
+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[
@@ -129,6 +131,20 @@ def set_query_context_tool(
                 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(
@@ -154,6 +170,10 @@ def set_query_context_tool(
             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:

mcp_server_snowflake/server.py

@@ -61,11 +61,13 @@
         "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}",
@@ -542,6 +544,9 @@ def build_query_comment(
             "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}",
         }
@@ -553,9 +558,20 @@ def build_query_comment(
         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():
-                    result = result.replace(f"{{{key}}}", str(sub_value))
+                    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()}

services/configuration.yaml

@@ -65,10 +65,12 @@ query_comment:
   #   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}"
@@ -82,6 +84,8 @@ query_comment:
   #   {agent_name} - Agent name (from set_query_context, defaults to "mcp-server-snowflake")
   #   {user_email} - User email (from set_query_context or SNOWFLAKE_MCP_USER_EMAIL env var)
   #   {user_name} - User name (from set_query_context or SNOWFLAKE_MCP_USER_NAME env var)
+  #   {intent} - Query intent object: {category, confidence, domains, question}
+  #   {query_parameters} - Query params object: {datasets, dimensions, time_range}
   #   {session_id} - Session ID (from set_query_context)
   #   {server_name} - MCP server name
   #   {server_version} - Server version
@@ -92,5 +96,7 @@ query_comment:
   #     model="claude-sonnet-4-5-20250929",
   #     agent_name="Claude Code",
   #     user_email="user@example.com",
-  #     user_name="username"
+  #     user_name="username",
+  #     intent={"category": "aggregation", "question": "Total users"},
+  #     query_parameters={"datasets": ["my_table"], "dimensions": ["date"]}
   #   )