feat(mcp): Add variable support to MCP run_query tool (#9464)

Add support for variables to the MCP query tool.
Accept 'variables' argument.
variables must be a JSON map of parameters used in the query string when
using a parametrized query.

**Description**

Add variables argument to query tool
use QueryWithVars instead of Query and pass var map to QueryWithVars

**Checklist**

- [x] Code compiles correctly and linting passes locally
- [ ] For all _code_ changes, an entry added to the `CHANGELOG.md` file
describing and linking to
      this PR
- [x] Tests added for new functionality, or regression tests for bug
fixes added as applicable
- [ ] For public APIs, new features, etc., PR on
[docs repo](https://github.com/hypermodeinc/docs) staged and linked here

---------

Co-authored-by: mattthew <[email protected]>

Author: rderbier

dgraph/cmd/mcp/mcp_server.go

@@ -3,7 +3,9 @@ package mcp
 import (
 	"context"
 	_ "embed"
+	"encoding/json"
 	"fmt"
+	"strconv"
 	"sync"
 	"time"
 
@@ -71,11 +73,15 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 	)
 
 	queryTool := mcp.NewTool("run_query",
-		mcp.WithDescription("Run Dgraph DQL Query on dgraph db"),
+		mcp.WithDescription("Run DQL Query on Dgraph"),
 		mcp.WithString("query",
 			mcp.Required(),
 			mcp.Description("The query to perform"),
 		),
+		mcp.WithString("variables",
+			mcp.Required(),
+			mcp.Description("The parameters to pass to the query in JSON format. The JSON should be a map of string keys to string, number or boolean values. Example: {\"$param1\": \"value1\", \"$param2\": 123, \"$param3\": true}"),
+		),
 		mcp.WithToolAnnotation(mcp.ToolAnnotation{
 			ReadOnlyHint:    &True,
 			DestructiveHint: &False,
@@ -86,10 +92,10 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 
 	if !readOnly {
 		alterSchemaTool := mcp.NewTool("alter_schema",
-			mcp.WithDescription("Alter Dgraph DQL Schema in dgraph db"),
+			mcp.WithDescription("Alter DQL Schema in Dgraph"),
 			mcp.WithString("schema",
 				mcp.Required(),
-				mcp.Description("Updated schema to insert inside the db"),
+				mcp.Description("DQL Schema to apply"),
 			),
 			mcp.WithToolAnnotation(mcp.ToolAnnotation{
 				ReadOnlyHint:    &False,
@@ -128,10 +134,10 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 		})
 
 		mutationTool := mcp.NewTool("run_mutation",
-			mcp.WithDescription("Run DQL Mutation on dgraph db"),
+			mcp.WithDescription("Run DQL Mutation on Dgraph"),
 			mcp.WithString("mutation",
 				mcp.Required(),
-				mcp.Description("The mutation to perform in json format"),
+				mcp.Description("The mutation to perform in JSON format"),
 			),
 			mcp.WithToolAnnotation(mcp.ToolAnnotation{
 				ReadOnlyHint:    &False,
@@ -183,7 +189,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 		if err != nil {
 			return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
 		}
-		txn := conn.NewTxn()
+		txn := conn.NewReadOnlyTxn()
 		defer func() {
 			err := txn.Discard(ctx)
 			if err != nil {
@@ -202,7 +208,35 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 		if !ok {
 			return mcp.NewToolResultError("Query must be a string"), nil
 		}
-		resp, err := txn.Query(ctx, op)
+		vars := make(map[string]any)
+		variablesArg, ok := args["variables"]
+		if ok && variablesArg != nil {
+			variables, ok := variablesArg.(string)
+			if !ok {
+				return mcp.NewToolResultError("Variables must be a JSON-formatted string"), nil
+			}
+			// create a map of variables from JSON string
+			if err := json.Unmarshal([]byte(variables), &vars); err != nil {
+				return mcp.NewToolResultErrorFromErr("Error parsing variables", err), nil
+			}
+		}
+		// convert vars map[string]any to map[string]string as required by txn.QueryWithVars
+		varsString := make(map[string]string)
+		for k, v := range vars {
+			switch val := v.(type) {
+			case string:
+				varsString[k] = val
+			case float64:
+				varsString[k] = strconv.FormatFloat(val, 'f', -1, 64)
+			case bool:
+				varsString[k] = strconv.FormatBool(val)
+			case nil:
+				varsString[k] = "null"
+			default:
+				return mcp.NewToolResultError(fmt.Sprintf("could not convert complex variable %q to string", k)), nil
+			}
+		}
+		resp, err := txn.QueryWithVars(ctx, op, varsString)
 		if err != nil {
 			return mcp.NewToolResultErrorFromErr("Error running query", err), nil
 		}
@@ -214,7 +248,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
 		if err != nil {
 			return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
 		}
-		txn := conn.NewTxn()
+		txn := conn.NewReadOnlyTxn()
 		defer func() {
 			err := txn.Discard(ctx)
 			if err != nil {

dgraph/cmd/mcp/mcp_server_sse_test.go

@@ -176,7 +176,7 @@ func TestMCPSSE(t *testing.T) {
 
 	t.Run("AlterSchema", func(t *testing.T) {
 		args := map[string]interface{}{
-			"schema": "n: string @index(term) .",
+			"schema": "n: string @index(term) .\nm: int @index(int) .\np: float @index(float) .",
 		}
 		resultText, err := callTool("alter_schema", args)
 
@@ -190,7 +190,9 @@ func TestMCPSSE(t *testing.T) {
   "set": [
     {
       "uid": "_:1",
-      "n": "Foo"
+      "n": "Foo",
+      "m": 20,
+      "p": 3.14
     }
   ]
 }`,
@@ -201,6 +203,21 @@ func TestMCPSSE(t *testing.T) {
 		require.Equal(t, "Mutation completed, 1 UIDs created", resultText, "Should receive run mutation result")
 	})
 
+	checkUIDResult := func(t *testing.T, resultText string) {
+		require.NotEmpty(t, resultText, "Should receive run query result")
+
+		var result map[string][]map[string]string
+		err := json.Unmarshal([]byte(resultText), &result)
+		require.NoError(t, err, "Should be able to parse JSON response")
+
+		require.Contains(t, result, "q", "Response should contain 'q' field")
+		require.Len(t, result["q"], 1, "Should have exactly one result")
+		require.Contains(t, result["q"][0], "uid", "Result should have 'uid' field")
+
+		uidValue := result["q"][0]["uid"]
+		require.Regexp(t, `^0x[0-9a-f]+$`, uidValue, "UID should be in the format 0x followed by hexadecimal digits")
+	}
+
 	t.Run("RunQuery", func(t *testing.T) {
 		args := map[string]interface{}{
 			"query": `{q(func: allofterms(n, "Foo")) { uid }}`,
@@ -230,6 +247,55 @@ func TestMCPSSE(t *testing.T) {
 		require.Contains(t, err.Error(), "Query must be present")
 	})
 
+	t.Run("RunQueryWithVars", func(t *testing.T) {
+		t.Run("StringVars", func(t *testing.T) {
+			args := map[string]interface{}{
+				"query":     `query me($name: string) {q(func: allofterms(n, $name)) { uid }}`,
+				"variables": `{"$name": "Foo"}`,
+			}
+			resultText, err := callTool("run_query", args)
+
+			require.NoError(t, err, "RunQueryWithVars should not fail")
+			require.NotEmpty(t, resultText, "Should receive run query result")
+			checkUIDResult(t, resultText)
+		})
+
+		t.Run("IntVars", func(t *testing.T) {
+			args := map[string]interface{}{
+				"query":     `query me($age: int) {q(func: eq(m, $age)) { uid }}`,
+				"variables": `{"$age": 20}`,
+			}
+			resultText, err := callTool("run_query", args)
+
+			require.NoError(t, err, "RunQueryWithVars should not fail")
+			require.NotEmpty(t, resultText, "Should receive run query result")
+			checkUIDResult(t, resultText)
+		})
+
+		t.Run("FloatVars", func(t *testing.T) {
+			args := map[string]interface{}{
+				"query":     `query me($v: float) {q(func: eq(p, $v)) { uid }}`,
+				"variables": `{"$v": 3.14}`,
+			}
+			resultText, err := callTool("run_query", args)
+
+			require.NoError(t, err, "RunQueryWithVars should not fail")
+			require.NotEmpty(t, resultText, "Should receive run query result")
+			checkUIDResult(t, resultText)
+		})
+
+		t.Run("InvalidVars", func(t *testing.T) {
+			args := map[string]interface{}{
+				"query":     `query me($name: string) {q(func: allofterms(n, $name)) { uid }}`,
+				"variables": `{"$name": {"bar": "Foo"}}`, // embedded map should fail
+			}
+			_, err := callTool("run_query", args)
+
+			require.Error(t, err, "RunQueryWithVars should fail")
+			require.Contains(t, err.Error(), "could not convert complex variable")
+		})
+	})
+
 	t.Run("RunGetCommonQueries", func(t *testing.T) {
 		resultText, err := callTool("get_common_queries", map[string]interface{}{})
 

dgraph/cmd/mcp/prompt.txt

@@ -4,7 +4,7 @@ The assistant is designed to help users interact with Dgraph databases using DQL
 This prompt enables the assistant to interpret user intents, generate appropriate queries or mutations, and maintain an informative and conversational tone.
 
 The user provides a Dgraph connection string. This string may point to a local or remote instance; the assistant doesn't need to differentiate.
-Instead, it validates the connection using `get schema {}` to ensure connectivity.
+Instead, it validates the connection using `get_schema` tool to ensure connectivity.
 
 The assistant must provide helpful feedback, preserve conversational context, and ensure safe operations (especially mutations).
 
@@ -14,9 +14,9 @@ MCP TOOLS
 Tools Available:
 - "get_schema": Gets Graph schema
 - "alter_schema": Alters Graph Schema.
-- "query": Executes DQL statements queries on the provided Dgraph connection and returns results.
-- "mutation": Executes DQL statements mutations on the provided Dgraph connection and returns results.
-- "common_queries": Some common queries that can be done on the Graph
+- "run_query": Executes DQL statements queries on the provided Dgraph connection and returns results. Variables are optional.
+- "run_mutation": Executes DQL statements mutations on the provided Dgraph connection and returns results.
+- "get_common_queries": Some common queries that can be done on the Graph
 </mcp>
 
 WORKFLOW
@@ -25,19 +25,19 @@ WORKFLOW
 
 1. Connection Setup:
    - Begin by prompting the user for their Dgraph connection string.
-   - Use `get schema {}` as a basic query to test if the connection is valid.
+   - Use the `get_schema` tool as a basic test to validate the connection.
    - On success: Confirm the connection is active.
    - On failure: Display a clear error and ask for correction.
 
 2. Schema Exploration:
-   - Upon user request or when needed to understand available data, run `get schema {}`.
+   - Upon user request or when needed to understand available data, run the `get_schema` tool.
    - Parse and extract all predicates and types.
    - Present them in a readable, structured format with examples (e.g., `<name>: string @index(term) .`).
 
 3. Query Execution:
    - Interpret user queries about data retrieval or analysis.
    - Match user intent to schema structure and generate DQL.
-   - Use the "query" tool to run the DQL and return results.
+   - Use the `run_query` tool to run the DQL and return results.
    - Provide explanations for the query structure and results, especially for less technical users.
 
 4. Mutations and Data Insertion:
@@ -95,7 +95,7 @@ CONVERSATION FLOW
 3. For analytical or search-related questions:
    - Confirm which predicates or types are relevant.
    - Fetch schema if necessary.
-   - Generate the corresponding DQL.
+   - Generate the corresponding DQL queries.
    - Execute and present results.
    - Visualize where helpful.
 
@@ -111,7 +111,7 @@ ERROR HANDLING
 <error-handling>
 
 - Connection Errors:
-  - Run `get schema {}` to validate.
+  - Run the `get_schema` tool to validate the connection.
   - If it fails, provide meaningful errors: network issues, invalid string, authentication, etc.
 
 - Schema Errors:
@@ -136,6 +136,7 @@ Dgraph uses DQL (Dgraph Query Language), a superset of GraphQL designed for grap
 
 Common Statements:
 - `query { ... }` → For querying nodes and edges
+- `query me($foo: string, $bar: int) { ... }` → For querying nodes and edges with variables
 - `mutation { set { ... } }` → For inserting/updating data
 - `mutation { delete { ... } }` → For deleting data
 - `schema {}` → For introspecting current schema
@@ -171,6 +172,17 @@ Simple Query:
 }
 ```
 
+Query with variables:
+```
+query me($foo: string) {
+  me(func: eq(name, $foo)) {
+    uid
+    name
+    age
+  }
+}
+```
+
 Mutation:
 ```
 mutation {