Skip to content

Commit 4bee06b

Browse files
darkcoderrisesdarkcoderrises
committed
feat(core): Add mcp server
1 parent c5cdaf0 commit 4bee06b

3 files changed

Lines changed: 62 additions & 1 deletion

File tree

‎dgraph/cmd/alpha/run.go‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ import (
2525
"time"
2626

2727
"github.com/golang/glog"
28+
"github.com/mark3labs/mcp-go/mcp"
29+
"github.com/mark3labs/mcp-go/server"
2830
"github.com/pkg/errors"
2931
"github.com/spf13/cobra"
3032
"go.opencensus.io/plugin/ocgrpc"
@@ -472,9 +474,61 @@ func serveGRPC(l net.Listener, tlsCfg *tls.Config, closer *z.Closer) {
472474
s.Stop()
473475
}
474476

477+
func setupMcp(baseMux *http.ServeMux) {
478+
s := server.NewMCPServer(
479+
"Dgraph MCP Server",
480+
"1.0.0",
481+
server.WithResourceCapabilities(true, true),
482+
server.WithLogging(),
483+
server.WithRecovery(),
484+
)
485+
486+
schemaTool := mcp.NewTool("Get-Schema",
487+
mcp.WithDescription("Get schema from dgraph db"),
488+
)
489+
490+
queryTool := mcp.NewTool("Run-DQL",
491+
mcp.WithDescription("Run DQL mutation / query on dgraph db"),
492+
mcp.WithString("query",
493+
mcp.Required(),
494+
mcp.Description("The query to perform. For mutations the format should look like {set {}}"),
495+
),
496+
)
497+
498+
s.AddTool(queryTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
499+
op := request.Params.Arguments["query"].(string)
500+
req := &apiv25.RunDQLRequest{
501+
NsName: "root",
502+
DqlQuery: op,
503+
}
504+
resp, err := (&edgraph.ServerV25{}).RunDQL(ctx, req)
505+
if err != nil {
506+
return mcp.NewToolResultError(err.Error()), nil
507+
}
508+
return mcp.NewToolResultText(string(resp.QueryResult)), nil
509+
})
510+
511+
s.AddTool(schemaTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
512+
req := &apiv25.RunDQLRequest{
513+
NsName: "root",
514+
DqlQuery: "schema {}",
515+
}
516+
resp, err := (&edgraph.ServerV25{}).RunDQL(ctx, req)
517+
if err != nil {
518+
return mcp.NewToolResultError(err.Error()), nil
519+
}
520+
return mcp.NewToolResultText(string(resp.QueryResult)), nil
521+
})
522+
523+
sse := server.NewSSEServer(s,
524+
server.WithBasePath("/mcp"),
525+
)
526+
baseMux.HandleFunc("/mcp", sse.ServeHTTP)
527+
baseMux.HandleFunc("/mcp/", sse.ServeHTTP)
528+
}
529+
475530
func setupServer(closer *z.Closer) {
476531
go worker.RunServer(bindall) // For pb.communication.
477-
478532
laddr := "localhost"
479533
if bindall {
480534
laddr = "0.0.0.0"
@@ -509,6 +563,7 @@ func setupServer(closer *z.Closer) {
509563
baseMux.HandleFunc("/state", stateHandler)
510564
baseMux.HandleFunc("/debug/jemalloc", x.JemallocHandler)
511565
http.DefaultServeMux.Handle("/debug/z", zpages.NewTracezHandler(zpages.NewSpanProcessor()))
566+
setupMcp(baseMux)
512567

513568
// TODO: Figure out what this is for?
514569
http.HandleFunc("/debug/store", storeStatsHandler)

‎go.mod‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ require (
3333
github.com/graph-gophers/graphql-go v1.6.0
3434
github.com/hashicorp/vault/api v1.16.0
3535
github.com/klauspost/compress v1.18.0
36+
github.com/mark3labs/mcp-go v0.21.0
3637
github.com/minio/minio-go/v6 v6.0.57
3738
github.com/paulmach/go.geojson v1.5.0
3839
github.com/pkg/errors v0.9.1
@@ -151,6 +152,7 @@ require (
151152
github.com/subosito/gotenv v1.6.0 // indirect
152153
github.com/viterin/partial v1.1.0 // indirect
153154
github.com/xdg/stringprep v1.0.3 // indirect
155+
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
154156
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
155157
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
156158
go.opentelemetry.io/otel/metric v1.35.0 // indirect

‎go.sum‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,8 @@ github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+
394394
github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80/go.mod h1:imJHygn/1yfhB7XSJJKlFZKl/J+dCPAknuiaGOshXAs=
395395
github.com/logrusorgru/aurora v0.0.0-20200102142835-e9ef32dff381/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
396396
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
397+
github.com/mark3labs/mcp-go v0.21.0 h1:oyEtiXg8PnrVEFis9b1AwbiUWF2dTbyBP5yLo7SruXE=
398+
github.com/mark3labs/mcp-go v0.21.0/go.mod h1:KmJndYv7GIgcPVwEKJjNcbhVQ+hJGJhrCCB/9xITzpE=
397399
github.com/matryer/moq v0.0.0-20200106131100-75d0ddfc0007/go.mod h1:9ELz6aaclSIGnZBoaSLZ3NAl1VTufbOrXBPvtcy6WiQ=
398400
github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
399401
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
@@ -568,6 +570,8 @@ github.com/xdg/scram v1.0.5 h1:TuS0RFmt5Is5qm9Tm2SoD89OPqe4IRiFtyFY4iwWXsw=
568570
github.com/xdg/scram v1.0.5/go.mod h1:lB8K/P019DLNhemzwFU4jHLhdvlE6uDZjXFejJXr49I=
569571
github.com/xdg/stringprep v1.0.3 h1:cmL5Enob4W83ti/ZHuZLuKD/xqJfus4fVPwE+/BDm+4=
570572
github.com/xdg/stringprep v1.0.3/go.mod h1:Jhud4/sHMO4oL310DaZAKk9ZaJ08SJfe+sJh0HrGL1Y=
573+
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
574+
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
571575
github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
572576
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
573577
github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

0 commit comments

Comments
 (0)