Skip to content

Commit 957e9e4

Browse files
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]>
1 parent 883880d commit 957e9e4

File tree

3 files changed

+131
-19
lines changed

3 files changed

+131
-19
lines changed

dgraph/cmd/mcp/mcp_server.go

Lines changed: 42 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ package mcp
33
import (
44
"context"
55
_ "embed"
6+
"encoding/json"
67
"fmt"
8+
"strconv"
79
"sync"
810
"time"
911

@@ -71,11 +73,15 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
7173
)
7274

7375
queryTool := mcp.NewTool("run_query",
74-
mcp.WithDescription("Run Dgraph DQL Query on dgraph db"),
76+
mcp.WithDescription("Run DQL Query on Dgraph"),
7577
mcp.WithString("query",
7678
mcp.Required(),
7779
mcp.Description("The query to perform"),
7880
),
81+
mcp.WithString("variables",
82+
mcp.Required(),
83+
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}"),
84+
),
7985
mcp.WithToolAnnotation(mcp.ToolAnnotation{
8086
ReadOnlyHint: &True,
8187
DestructiveHint: &False,
@@ -86,10 +92,10 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
8692

8793
if !readOnly {
8894
alterSchemaTool := mcp.NewTool("alter_schema",
89-
mcp.WithDescription("Alter Dgraph DQL Schema in dgraph db"),
95+
mcp.WithDescription("Alter DQL Schema in Dgraph"),
9096
mcp.WithString("schema",
9197
mcp.Required(),
92-
mcp.Description("Updated schema to insert inside the db"),
98+
mcp.Description("DQL Schema to apply"),
9399
),
94100
mcp.WithToolAnnotation(mcp.ToolAnnotation{
95101
ReadOnlyHint: &False,
@@ -128,10 +134,10 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
128134
})
129135

130136
mutationTool := mcp.NewTool("run_mutation",
131-
mcp.WithDescription("Run DQL Mutation on dgraph db"),
137+
mcp.WithDescription("Run DQL Mutation on Dgraph"),
132138
mcp.WithString("mutation",
133139
mcp.Required(),
134-
mcp.Description("The mutation to perform in json format"),
140+
mcp.Description("The mutation to perform in JSON format"),
135141
),
136142
mcp.WithToolAnnotation(mcp.ToolAnnotation{
137143
ReadOnlyHint: &False,
@@ -183,7 +189,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
183189
if err != nil {
184190
return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
185191
}
186-
txn := conn.NewTxn()
192+
txn := conn.NewReadOnlyTxn()
187193
defer func() {
188194
err := txn.Discard(ctx)
189195
if err != nil {
@@ -202,7 +208,35 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
202208
if !ok {
203209
return mcp.NewToolResultError("Query must be a string"), nil
204210
}
205-
resp, err := txn.Query(ctx, op)
211+
vars := make(map[string]any)
212+
variablesArg, ok := args["variables"]
213+
if ok && variablesArg != nil {
214+
variables, ok := variablesArg.(string)
215+
if !ok {
216+
return mcp.NewToolResultError("Variables must be a JSON-formatted string"), nil
217+
}
218+
// create a map of variables from JSON string
219+
if err := json.Unmarshal([]byte(variables), &vars); err != nil {
220+
return mcp.NewToolResultErrorFromErr("Error parsing variables", err), nil
221+
}
222+
}
223+
// convert vars map[string]any to map[string]string as required by txn.QueryWithVars
224+
varsString := make(map[string]string)
225+
for k, v := range vars {
226+
switch val := v.(type) {
227+
case string:
228+
varsString[k] = val
229+
case float64:
230+
varsString[k] = strconv.FormatFloat(val, 'f', -1, 64)
231+
case bool:
232+
varsString[k] = strconv.FormatBool(val)
233+
case nil:
234+
varsString[k] = "null"
235+
default:
236+
return mcp.NewToolResultError(fmt.Sprintf("could not convert complex variable %q to string", k)), nil
237+
}
238+
}
239+
resp, err := txn.QueryWithVars(ctx, op, varsString)
206240
if err != nil {
207241
return mcp.NewToolResultErrorFromErr("Error running query", err), nil
208242
}
@@ -214,7 +248,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
214248
if err != nil {
215249
return mcp.NewToolResultErrorFromErr("Error opening connection with Dgraph Alpha", err), nil
216250
}
217-
txn := conn.NewTxn()
251+
txn := conn.NewReadOnlyTxn()
218252
defer func() {
219253
err := txn.Discard(ctx)
220254
if err != nil {

dgraph/cmd/mcp/mcp_server_sse_test.go

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ func TestMCPSSE(t *testing.T) {
176176

177177
t.Run("AlterSchema", func(t *testing.T) {
178178
args := map[string]interface{}{
179-
"schema": "n: string @index(term) .",
179+
"schema": "n: string @index(term) .\nm: int @index(int) .\np: float @index(float) .",
180180
}
181181
resultText, err := callTool("alter_schema", args)
182182

@@ -190,7 +190,9 @@ func TestMCPSSE(t *testing.T) {
190190
"set": [
191191
{
192192
"uid": "_:1",
193-
"n": "Foo"
193+
"n": "Foo",
194+
"m": 20,
195+
"p": 3.14
194196
}
195197
]
196198
}`,
@@ -201,6 +203,21 @@ func TestMCPSSE(t *testing.T) {
201203
require.Equal(t, "Mutation completed, 1 UIDs created", resultText, "Should receive run mutation result")
202204
})
203205

206+
checkUIDResult := func(t *testing.T, resultText string) {
207+
require.NotEmpty(t, resultText, "Should receive run query result")
208+
209+
var result map[string][]map[string]string
210+
err := json.Unmarshal([]byte(resultText), &result)
211+
require.NoError(t, err, "Should be able to parse JSON response")
212+
213+
require.Contains(t, result, "q", "Response should contain 'q' field")
214+
require.Len(t, result["q"], 1, "Should have exactly one result")
215+
require.Contains(t, result["q"][0], "uid", "Result should have 'uid' field")
216+
217+
uidValue := result["q"][0]["uid"]
218+
require.Regexp(t, `^0x[0-9a-f]+$`, uidValue, "UID should be in the format 0x followed by hexadecimal digits")
219+
}
220+
204221
t.Run("RunQuery", func(t *testing.T) {
205222
args := map[string]interface{}{
206223
"query": `{q(func: allofterms(n, "Foo")) { uid }}`,
@@ -230,6 +247,55 @@ func TestMCPSSE(t *testing.T) {
230247
require.Contains(t, err.Error(), "Query must be present")
231248
})
232249

250+
t.Run("RunQueryWithVars", func(t *testing.T) {
251+
t.Run("StringVars", func(t *testing.T) {
252+
args := map[string]interface{}{
253+
"query": `query me($name: string) {q(func: allofterms(n, $name)) { uid }}`,
254+
"variables": `{"$name": "Foo"}`,
255+
}
256+
resultText, err := callTool("run_query", args)
257+
258+
require.NoError(t, err, "RunQueryWithVars should not fail")
259+
require.NotEmpty(t, resultText, "Should receive run query result")
260+
checkUIDResult(t, resultText)
261+
})
262+
263+
t.Run("IntVars", func(t *testing.T) {
264+
args := map[string]interface{}{
265+
"query": `query me($age: int) {q(func: eq(m, $age)) { uid }}`,
266+
"variables": `{"$age": 20}`,
267+
}
268+
resultText, err := callTool("run_query", args)
269+
270+
require.NoError(t, err, "RunQueryWithVars should not fail")
271+
require.NotEmpty(t, resultText, "Should receive run query result")
272+
checkUIDResult(t, resultText)
273+
})
274+
275+
t.Run("FloatVars", func(t *testing.T) {
276+
args := map[string]interface{}{
277+
"query": `query me($v: float) {q(func: eq(p, $v)) { uid }}`,
278+
"variables": `{"$v": 3.14}`,
279+
}
280+
resultText, err := callTool("run_query", args)
281+
282+
require.NoError(t, err, "RunQueryWithVars should not fail")
283+
require.NotEmpty(t, resultText, "Should receive run query result")
284+
checkUIDResult(t, resultText)
285+
})
286+
287+
t.Run("InvalidVars", func(t *testing.T) {
288+
args := map[string]interface{}{
289+
"query": `query me($name: string) {q(func: allofterms(n, $name)) { uid }}`,
290+
"variables": `{"$name": {"bar": "Foo"}}`, // embedded map should fail
291+
}
292+
_, err := callTool("run_query", args)
293+
294+
require.Error(t, err, "RunQueryWithVars should fail")
295+
require.Contains(t, err.Error(), "could not convert complex variable")
296+
})
297+
})
298+
233299
t.Run("RunGetCommonQueries", func(t *testing.T) {
234300
resultText, err := callTool("get_common_queries", map[string]interface{}{})
235301

dgraph/cmd/mcp/prompt.txt

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ The assistant is designed to help users interact with Dgraph databases using DQL
44
This prompt enables the assistant to interpret user intents, generate appropriate queries or mutations, and maintain an informative and conversational tone.
55

66
The user provides a Dgraph connection string. This string may point to a local or remote instance; the assistant doesn't need to differentiate.
7-
Instead, it validates the connection using `get schema {}` to ensure connectivity.
7+
Instead, it validates the connection using `get_schema` tool to ensure connectivity.
88

99
The assistant must provide helpful feedback, preserve conversational context, and ensure safe operations (especially mutations).
1010

@@ -14,9 +14,9 @@ MCP TOOLS
1414
Tools Available:
1515
- "get_schema": Gets Graph schema
1616
- "alter_schema": Alters Graph Schema.
17-
- "query": Executes DQL statements queries on the provided Dgraph connection and returns results.
18-
- "mutation": Executes DQL statements mutations on the provided Dgraph connection and returns results.
19-
- "common_queries": Some common queries that can be done on the Graph
17+
- "run_query": Executes DQL statements queries on the provided Dgraph connection and returns results. Variables are optional.
18+
- "run_mutation": Executes DQL statements mutations on the provided Dgraph connection and returns results.
19+
- "get_common_queries": Some common queries that can be done on the Graph
2020
</mcp>
2121

2222
WORKFLOW
@@ -25,19 +25,19 @@ WORKFLOW
2525

2626
1. Connection Setup:
2727
- Begin by prompting the user for their Dgraph connection string.
28-
- Use `get schema {}` as a basic query to test if the connection is valid.
28+
- Use the `get_schema` tool as a basic test to validate the connection.
2929
- On success: Confirm the connection is active.
3030
- On failure: Display a clear error and ask for correction.
3131

3232
2. Schema Exploration:
33-
- Upon user request or when needed to understand available data, run `get schema {}`.
33+
- Upon user request or when needed to understand available data, run the `get_schema` tool.
3434
- Parse and extract all predicates and types.
3535
- Present them in a readable, structured format with examples (e.g., `<name>: string @index(term) .`).
3636

3737
3. Query Execution:
3838
- Interpret user queries about data retrieval or analysis.
3939
- Match user intent to schema structure and generate DQL.
40-
- Use the "query" tool to run the DQL and return results.
40+
- Use the `run_query` tool to run the DQL and return results.
4141
- Provide explanations for the query structure and results, especially for less technical users.
4242

4343
4. Mutations and Data Insertion:
@@ -95,7 +95,7 @@ CONVERSATION FLOW
9595
3. For analytical or search-related questions:
9696
- Confirm which predicates or types are relevant.
9797
- Fetch schema if necessary.
98-
- Generate the corresponding DQL.
98+
- Generate the corresponding DQL queries.
9999
- Execute and present results.
100100
- Visualize where helpful.
101101

@@ -111,7 +111,7 @@ ERROR HANDLING
111111
<error-handling>
112112

113113
- Connection Errors:
114-
- Run `get schema {}` to validate.
114+
- Run the `get_schema` tool to validate the connection.
115115
- If it fails, provide meaningful errors: network issues, invalid string, authentication, etc.
116116

117117
- Schema Errors:
@@ -136,6 +136,7 @@ Dgraph uses DQL (Dgraph Query Language), a superset of GraphQL designed for grap
136136

137137
Common Statements:
138138
- `query { ... }` → For querying nodes and edges
139+
- `query me($foo: string, $bar: int) { ... }` → For querying nodes and edges with variables
139140
- `mutation { set { ... } }` → For inserting/updating data
140141
- `mutation { delete { ... } }` → For deleting data
141142
- `schema {}` → For introspecting current schema
@@ -171,6 +172,17 @@ Simple Query:
171172
}
172173
```
173174

175+
Query with variables:
176+
```
177+
query me($foo: string) {
178+
me(func: eq(name, $foo)) {
179+
uid
180+
name
181+
age
182+
}
183+
}
184+
```
185+
174186
Mutation:
175187
```
176188
mutation {

0 commit comments

Comments
 (0)