@@ -60,7 +60,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
60
60
server .WithRecovery (),
61
61
)
62
62
63
- schemaTool := mcp .NewTool ("Get-Schema " ,
63
+ schemaTool := mcp .NewTool ("get_schema " ,
64
64
mcp .WithDescription ("Get Dgraph DQL Schema from dgraph db" ),
65
65
mcp .WithToolAnnotation (mcp.ToolAnnotation {
66
66
ReadOnlyHint : & True ,
@@ -70,8 +70,8 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
70
70
}),
71
71
)
72
72
73
- queryTool := mcp .NewTool ("Run-Query " ,
74
- mcp .WithDescription ("Run Dgraph Query on dgraph db" ),
73
+ queryTool := mcp .NewTool ("run_query " ,
74
+ mcp .WithDescription ("Run Dgraph DQL Query on dgraph db" ),
75
75
mcp .WithString ("query" ,
76
76
mcp .Required (),
77
77
mcp .Description ("The query to perform" ),
@@ -85,7 +85,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
85
85
)
86
86
87
87
if ! readOnly {
88
- alterSchemaTool := mcp .NewTool ("Alter-Schema " ,
88
+ alterSchemaTool := mcp .NewTool ("alter_schema " ,
89
89
mcp .WithDescription ("Alter Dgraph DQL Schema in dgraph db" ),
90
90
mcp .WithString ("schema" ,
91
91
mcp .Required (),
@@ -100,24 +100,34 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
100
100
)
101
101
102
102
s .AddTool (alterSchemaTool , func (ctx context.Context , request mcp.CallToolRequest ) (* mcp.CallToolResult , error ) {
103
- schema , ok := request .GetArguments ()["schema" ].(string )
103
+ args := request .GetArguments ()
104
+ if args == nil {
105
+ return mcp .NewToolResultError ("Schema must be present" ), nil
106
+ }
107
+
108
+ schemaArg , ok := args ["schema" ]
109
+ if ! ok || schemaArg == nil {
110
+ return mcp .NewToolResultError ("Schema must be present" ), nil
111
+ }
112
+
113
+ schema , ok := schemaArg .(string )
104
114
if ! ok {
105
- return nil , fmt . Errorf ( "schema must be present" )
115
+ return mcp . NewToolResultError ( "Schema must be a string" ), nil
106
116
}
107
117
108
118
// Execute alter operation
109
119
conn , err := getConn (connectionString )
110
120
if err != nil {
111
- return nil , fmt . Errorf ( "error opening connection with Dgraph Alpha: %v " , err )
121
+ return mcp . NewToolResultErrorFromErr ( "Error opening connection with Dgraph Alpha" , err ), nil
112
122
}
113
123
if err = conn .SetSchema (ctx , dgo .RootNamespace , schema ); err != nil {
114
- return nil , fmt . Errorf ( "schema alteration failed: %v " , err )
124
+ return mcp . NewToolResultErrorFromErr ( "Schema alteration failed" , err ), nil
115
125
}
116
126
117
127
return mcp .NewToolResultText ("Schema updated successfully" ), nil
118
128
})
119
129
120
- mutationTool := mcp .NewTool ("Run-Mutation " ,
130
+ mutationTool := mcp .NewTool ("run_mutation " ,
121
131
mcp .WithDescription ("Run DQL Mutation on dgraph db" ),
122
132
mcp .WithString ("mutation" ,
123
133
mcp .Required (),
@@ -134,7 +144,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
134
144
s .AddTool (mutationTool , func (ctx context.Context , request mcp.CallToolRequest ) (* mcp.CallToolResult , error ) {
135
145
conn , err := getConn (connectionString )
136
146
if err != nil {
137
- return nil , err
147
+ return mcp . NewToolResultErrorFromErr ( "Error opening connection with Dgraph Alpha" , err ), nil
138
148
}
139
149
txn := conn .NewTxn ()
140
150
defer func () {
@@ -143,25 +153,35 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
143
153
glog .Errorf ("failed to discard transaction: %v" , err )
144
154
}
145
155
}()
146
- mutation , ok := request .GetArguments ()["mutation" ].(string )
156
+ args := request .GetArguments ()
157
+ if args == nil {
158
+ return mcp .NewToolResultError ("Mutation must be present" ), nil
159
+ }
160
+
161
+ mutationArg , ok := args ["mutation" ]
162
+ if ! ok || mutationArg == nil {
163
+ return mcp .NewToolResultError ("Mutation must be present" ), nil
164
+ }
165
+
166
+ mutation , ok := mutationArg .(string )
147
167
if ! ok {
148
- return nil , fmt . Errorf ( "mutation must present" )
168
+ return mcp . NewToolResultError ( "Mutation must be a string" ), nil
149
169
}
150
170
resp , err := txn .Mutate (ctx , & api.Mutation {
151
171
SetJson : []byte (mutation ),
152
172
CommitNow : true ,
153
173
})
154
174
if err != nil {
155
- return mcp .NewToolResultError ( err . Error () ), nil
175
+ return mcp .NewToolResultErrorFromErr ( " Error running mutation" , err ), nil
156
176
}
157
- return mcp .NewToolResultText (string ( resp .GetJson () )), nil
177
+ return mcp .NewToolResultText (fmt . Sprintf ( "Mutation completed, %d UIDs created" , len ( resp .Uids ) / 2 )), nil
158
178
})
159
179
}
160
180
161
181
s .AddTool (queryTool , func (ctx context.Context , request mcp.CallToolRequest ) (* mcp.CallToolResult , error ) {
162
182
conn , err := getConn (connectionString )
163
183
if err != nil {
164
- return nil , err
184
+ return mcp . NewToolResultErrorFromErr ( "Error opening connection with Dgraph Alpha" , err ), nil
165
185
}
166
186
txn := conn .NewTxn ()
167
187
defer func () {
@@ -170,18 +190,29 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
170
190
glog .Errorf ("failed to discard transaction: %v" , err )
171
191
}
172
192
}()
173
- op := request .GetArguments ()["query" ].(string )
193
+ args := request .GetArguments ()
194
+ if args == nil {
195
+ return mcp .NewToolResultError ("Query must be present" ), nil
196
+ }
197
+ queryArg , ok := args ["query" ]
198
+ if ! ok || queryArg == nil {
199
+ return mcp .NewToolResultError ("Query must be present" ), nil
200
+ }
201
+ op , ok := queryArg .(string )
202
+ if ! ok {
203
+ return mcp .NewToolResultError ("Query must be a string" ), nil
204
+ }
174
205
resp , err := txn .Query (ctx , op )
175
206
if err != nil {
176
- return mcp .NewToolResultError ( err . Error () ), nil
207
+ return mcp .NewToolResultErrorFromErr ( " Error running query" , err ), nil
177
208
}
178
209
return mcp .NewToolResultText (string (resp .GetJson ())), nil
179
210
})
180
211
181
212
s .AddTool (schemaTool , func (ctx context.Context , request mcp.CallToolRequest ) (* mcp.CallToolResult , error ) {
182
213
conn , err := getConn (connectionString )
183
214
if err != nil {
184
- return nil , err
215
+ return mcp . NewToolResultErrorFromErr ( "Error opening connection with Dgraph Alpha" , err ), nil
185
216
}
186
217
txn := conn .NewTxn ()
187
218
defer func () {
@@ -192,27 +223,27 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
192
223
}()
193
224
resp , err := txn .Query (ctx , "schema {}" )
194
225
if err != nil {
195
- return mcp .NewToolResultError ( err . Error () ), nil
226
+ return mcp .NewToolResultErrorFromErr ( " Error running query" , err ), nil
196
227
}
197
228
return mcp .NewToolResultText (string (resp .GetJson ())), nil
198
229
})
199
230
200
231
schemaResource := mcp .NewResource (
201
232
"dgraph://schema" ,
202
- "Dgraph Schema " ,
203
- mcp .WithResourceDescription ("The current Dgraph schema" ),
233
+ "dgraph_schema " ,
234
+ mcp .WithResourceDescription ("The current Dgraph DQL schema" ),
204
235
mcp .WithMIMEType ("text/plain" ),
205
236
)
206
237
207
238
s .AddResource (schemaResource , func (ctx context.Context , request mcp.ReadResourceRequest ) ([]mcp.ResourceContents , error ) {
208
239
// Execute operation
209
240
conn , err := getConn (connectionString )
210
241
if err != nil {
211
- return nil , err
242
+ return nil , fmt . Errorf ( "error opening connection with Dgraph Alpha: %w" , err )
212
243
}
213
244
resp , err := conn .NewTxn ().Query (ctx , "schema {}" )
214
245
if err != nil {
215
- return nil , fmt .Errorf ("failed to get schema : %v " , err )
246
+ return nil , fmt .Errorf ("error running query : %w " , err )
216
247
}
217
248
218
249
return []mcp.ResourceContents {
@@ -224,7 +255,7 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
224
255
}, nil
225
256
})
226
257
227
- commonQueriesTool := mcp .NewTool ("Get-Common-Queries " ,
258
+ commonQueriesTool := mcp .NewTool ("get_common_queries " ,
228
259
mcp .WithDescription ("Get common queries that you can run on the db. If you are seeing issues with your queries, you can check this tool once." ),
229
260
mcp .WithToolAnnotation (mcp.ToolAnnotation {
230
261
ReadOnlyHint : & True ,
@@ -257,16 +288,16 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
257
288
})
258
289
259
290
commonQueries := mcp .NewResource (
260
- "dgraph://common-queries " ,
261
- "Dgraph common queries " ,
291
+ "dgraph://common_queries " ,
292
+ "dgraph_common_queries " ,
262
293
mcp .WithResourceDescription ("The current Dgraph common queries that you can use to fix your queries" ),
263
294
mcp .WithMIMEType ("text/plain" ),
264
295
)
265
296
266
297
s .AddResource (commonQueries , func (ctx context.Context , request mcp.ReadResourceRequest ) ([]mcp.ResourceContents , error ) {
267
298
return []mcp.ResourceContents {
268
299
mcp.TextResourceContents {
269
- URI : "dgraph://commmon-queries " ,
300
+ URI : "dgraph://commmon_queries " ,
270
301
MIMEType : "text/plain" ,
271
302
Text : `
272
303
{
@@ -322,11 +353,11 @@ func NewMCPServer(connectionString string, readOnly bool) (*server.MCPServer, er
322
353
323
354
func addPrompt (s * server.MCPServer ) {
324
355
prompt := string (promptBytes )
325
- s .AddPrompt (mcp .NewPrompt ("Quick start prompt " ,
356
+ s .AddPrompt (mcp .NewPrompt ("quick_start_prompt " ,
326
357
mcp .WithPromptDescription ("A quick Start prompt for new users and llms" ),
327
358
), func (ctx context.Context , request mcp.GetPromptRequest ) (* mcp.GetPromptResult , error ) {
328
359
return mcp .NewGetPromptResult (
329
- "A quick start prompt " ,
360
+ "quick_start_prompt " ,
330
361
[]mcp.PromptMessage {
331
362
mcp .NewPromptMessage (
332
363
mcp .RoleAssistant ,
0 commit comments