feat: align query comment format with dbt query tag schema

Author: atvaccaro

mcp_server_snowflake/query_manager/tools.py

@@ -92,9 +92,10 @@ def run_query_tool(
 
 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")
+- user_email: Email of the user running the agent
+- user_name: Name of the user running the agent
 - session_id: A unique session identifier for grouping related queries
-- agent_name: Name of the agent or application making queries
-- user_context: Additional user-provided context
 
 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.""",
@@ -104,21 +105,35 @@ def set_query_context_tool(
             str,
             Field(
                 default=None,
-                description="AI model name (e.g., 'claude-sonnet-4')",
+                description="AI model name (e.g., 'claude-sonnet-4-5-20250929')",
             ),
         ] = None,
-        session_id: Annotated[
+        agent_name: Annotated[
             str,
             Field(
                 default=None,
-                description="Unique session identifier for grouping queries",
+                description="Name of the agent or application (e.g., 'Claude Code')",
             ),
         ] = None,
-        agent_name: Annotated[
+        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,
+        session_id: Annotated[
             str,
             Field(
                 default=None,
-                description="Name of the agent or application",
+                description="Unique session identifier for grouping queries",
             ),
         ] = None,
         custom_context: Annotated[
@@ -133,10 +148,14 @@ def set_query_context_tool(
         context = {}
         if model is not None:
             context["model"] = model
-        if session_id is not None:
-            context["session_id"] = session_id
         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 session_id is not None:
+            context["session_id"] = session_id
         if custom_context is not None:
             context.update(custom_context)
 

mcp_server_snowflake/server.py

@@ -54,14 +54,22 @@
 tag_minor_version = 3
 query_tag = {"origin": "sf_sit", "name": "mcp_server"}
 
-# Default query comment template - provides observability metadata for queries
+# Default query comment template - matches dbt query tag format for observability
 DEFAULT_QUERY_COMMENT_TEMPLATE = {
-    "source": "mcp-server-snowflake",
-    "request_id": "{request_id}",
-    "timestamp": "{timestamp}",
-    "tool": "{tool_name}",
-    "statement_type": "{statement_type}",
-    "model": "{model}",
+    "agent": "{agent_name}",
+    "context": {
+        "request_id": "{request_id}",
+        "timestamp": "{timestamp}",
+    },
+    "model_version": "{model}",
+    "query": {
+        "tool": "{tool_name}",
+        "statement_type": "{statement_type}",
+    },
+    "user": {
+        "email": "{user_email}",
+        "name": "{user_name}",
+    },
 }
 
 logger = get_logger(server_name)
@@ -525,8 +533,15 @@ def build_query_comment(
             ),
             # Session ID: from runtime context or default
             "session_id": self.query_context.get("session_id", "unknown"),
-            # Agent name: from runtime context or default
-            "agent_name": self.query_context.get("agent_name", "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")
+            ),
             "server_name": server_name,
             "server_version": f"{tag_major_version}.{tag_minor_version}",
         }

mcp_server_snowflake/tests/test_query_comment.py

@@ -177,16 +177,20 @@ def test_substitutes_template_variables(self, tmp_path):
         assert result is not None
         comment = json.loads(result)
 
-        # Check substituted values
-        assert comment["tool"] == "run_snowflake_query"
-        assert comment["statement_type"] == "Select"
-        assert comment["model"] == "claude-sonnet-4"
-        assert comment["source"] == "mcp-server-snowflake"
+        # Check substituted values (new nested format matching dbt query tag)
+        assert comment["query"]["tool"] == "run_snowflake_query"
+        assert comment["query"]["statement_type"] == "Select"
+        assert comment["model_version"] == "claude-sonnet-4"
+        assert comment["agent"] == "mcp-server-snowflake"  # Default agent name
 
         # Check that UUID and timestamp were generated (not literal template strings)
-        assert "{request_id}" not in comment["request_id"]
-        assert "{timestamp}" not in comment["timestamp"]
-        assert len(comment["request_id"]) == 36  # UUID format
+        assert "{request_id}" not in comment["context"]["request_id"]
+        assert "{timestamp}" not in comment["context"]["timestamp"]
+        assert len(comment["context"]["request_id"]) == 36  # UUID format
+
+        # Check user fields default to unknown
+        assert comment["user"]["email"] == "unknown"
+        assert comment["user"]["name"] == "unknown"
 
     def test_nested_template_substitution(self, tmp_path):
         """Test that nested template variables are substituted."""
@@ -257,7 +261,7 @@ def test_model_defaults_to_unknown(self, tmp_path):
             )
 
         comment = json.loads(result)
-        assert comment["model"] == "unknown"
+        assert comment["model_version"] == "unknown"
 
 
 class TestSetQueryContext:
@@ -325,7 +329,7 @@ def test_runtime_context_overrides_env_var(self, tmp_path):
 
         comment = json.loads(result)
         # Runtime context should override env var
-        assert comment["model"] == "runtime-model"
+        assert comment["model_version"] == "runtime-model"
 
     def test_get_query_context_returns_copy(self, tmp_path):
         """Test that get_query_context returns a copy of the context."""

services/configuration.yaml

@@ -54,32 +54,43 @@ sql_statement_permissions: # List SQL statements to explicitly allow (True) or d
   - Use: True
 
 # Query comment configuration - adds metadata to queries for observability
+# Format matches dbt query tag schema for consistency across tools
 # Context can be set via:
-#   1. Environment variable: SNOWFLAKE_MCP_MODEL
+#   1. Environment variables: SNOWFLAKE_MCP_MODEL, SNOWFLAKE_MCP_USER_EMAIL, SNOWFLAKE_MCP_USER_NAME
 #   2. Runtime: Call the set_query_context tool at session start
 query_comment:
   enabled: True # Enabled by default for observability
-  # template: # Optional custom template (uses default if not specified)
-  #   source: "my-mcp-agent"
-  #   request_id: "{request_id}"
-  #   timestamp: "{timestamp}"
-  #   tool: "{tool_name}"
-  #   statement_type: "{statement_type}"
-  #   model: "{model}"
-  #   session_id: "{session_id}"
-  #   agent_name: "{agent_name}"
+  # template: # Optional custom template (default matches dbt query tag format)
+  #   agent: "{agent_name}"
+  #   context:
+  #     request_id: "{request_id}"
+  #     timestamp: "{timestamp}"
+  #   model_version: "{model}"
+  #   query:
+  #     tool: "{tool_name}"
+  #     statement_type: "{statement_type}"
+  #   user:
+  #     email: "{user_email}"
+  #     name: "{user_name}"
   #
   # Supported template variables:
   #   {request_id} - Unique UUID for this request
   #   {timestamp} - ISO 8601 timestamp (UTC)
   #   {tool_name} - Name of the MCP tool executing the query
   #   {statement_type} - SQL statement type (Select, Insert, etc.)
-  #   {model} - AI model name (from set_query_context or SNOWFLAKE_MCP_MODEL env var)
+  #   {model} - AI model version (from set_query_context or SNOWFLAKE_MCP_MODEL env var)
+  #   {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)
   #   {session_id} - Session ID (from set_query_context)
-  #   {agent_name} - Agent name (from set_query_context)
   #   {server_name} - MCP server name
   #   {server_version} - Server version
   #   {custom_key} - Any custom key set via set_query_context tool
   #
   # To set context at runtime, call the set_query_context tool:
-  #   set_query_context(model="claude-sonnet-4", session_id="sess-123", agent_name="my-agent")
+  #   set_query_context(
+  #     model="claude-sonnet-4-5-20250929",
+  #     agent_name="Claude Code",
+  #     user_email="user@example.com",
+  #     user_name="username"
+  #   )