Skip to content

Commit 488bc90

Browse files
darkcoderrisesdarkcoderrises
committed
fix comments
1 parent 54c98f1 commit 488bc90

2 files changed

Lines changed: 80 additions & 45 deletions

File tree

‎dgraph/cmd/alpha/run.go‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ they form a Raft group and provide synchronous replication.
113113
flag.String("custom_tokenizers", "",
114114
"Comma separated list of tokenizer plugins for custom indices.")
115115

116+
flag.Bool("mcp", true, "Enable MCP server.")
117+
116118
// By default Go GRPC traces all requests.
117119
grpc.EnableTracing = false
118120

@@ -493,7 +495,7 @@ func buildConnectionString(addr string, port int) string {
493495
return fmt.Sprintf("dgraph://%s:%d", addr, port)
494496
}
495497

496-
func setupServer(closer *z.Closer) {
498+
func setupServer(closer *z.Closer, enableMcp bool) {
497499
go worker.RunServer(bindall) // For pb.communication.
498500
laddr := "localhost"
499501
if bindall {
@@ -597,9 +599,12 @@ func setupServer(closer *z.Closer) {
597599
x.ServerCloser.AddRunning(3)
598600
go serveGRPC(grpcListener, tlsCfg, x.ServerCloser)
599601

600-
if err := setupMcp(baseMux, buildConnectionString(laddr, grpcPort()), "true"); err != nil {
601-
log.Fatal(err)
602+
if enableMcp {
603+
if err := setupMcp(baseMux, buildConnectionString(laddr, grpcPort()), "true"); err != nil {
604+
log.Fatal(err)
605+
}
602606
}
607+
603608
go x.StartListenHttpAndHttps(httpListener, tlsCfg, x.ServerCloser)
604609

605610
go func() {
@@ -665,6 +670,8 @@ func run() {
665670
x.Config.Limit = z.NewSuperFlag(Alpha.Conf.GetString("limit")).MergeAndCheckDefault(
666671
worker.LimitDefaults)
667672

673+
enableMcp := Alpha.Conf.GetBool("mcp")
674+
668675
opts := worker.Options{
669676
PostingDir: Alpha.Conf.GetString("postings"),
670677
WALDir: Alpha.Conf.GetString("wal"),
@@ -845,7 +852,7 @@ func run() {
845852
// close alpha. This closer is for closing and waiting that subscription.
846853
adminCloser := z.NewCloser(1)
847854

848-
setupServer(adminCloser)
855+
setupServer(adminCloser, enableMcp)
849856
glog.Infoln("GRPC and HTTP stopped.")
850857

851858
// This might not close until group is given the signal to close. So, only signal here,

‎mcp/mcp_server.go‎

Lines changed: 69 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,7 @@ import (
1313
"github.com/mark3labs/mcp-go/server"
1414
)
1515

16-
// NewMCPServer initializes and returns a new MCPServer instance.
17-
func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error) {
18-
s := server.NewMCPServer(
19-
"Dgraph MCP Server",
20-
"1.0.0",
21-
server.WithResourceCapabilities(true, true),
22-
server.WithLogging(),
23-
server.WithRecovery(),
24-
)
25-
16+
func getConn(connectionString string) (*dgo.Dgraph, error) {
2617
conn, err := dgo.Open(connectionString)
2718
if err != nil {
2819
for i := range 10 {
@@ -36,9 +27,21 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
3627
return nil, fmt.Errorf("error opening connection: %v", err)
3728
}
3829
}
30+
return conn, nil
31+
}
32+
33+
// NewMCPServer initializes and returns a new MCPServer instance.
34+
func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error) {
35+
s := server.NewMCPServer(
36+
"Dgraph MCP Server",
37+
"1.0.0",
38+
server.WithResourceCapabilities(true, true),
39+
server.WithLogging(),
40+
server.WithRecovery(),
41+
)
3942

4043
schemaTool := mcp.NewTool("Get-Schema",
41-
mcp.WithDescription("Get schema from dgraph db"),
44+
mcp.WithDescription("Get Dgraph DQL Schema from dgraph db"),
4245
mcp.WithToolAnnotation(mcp.ToolAnnotation{
4346
ReadOnlyHint: true,
4447
DestructiveHint: false,
@@ -47,22 +50,8 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
4750
}),
4851
)
4952

50-
alterSchemaTool := mcp.NewTool("Alter-Schema",
51-
mcp.WithDescription("Alter schema in dgraph db"),
52-
mcp.WithString("schema",
53-
mcp.Required(),
54-
mcp.Description("Updated schema to insert inside the db"),
55-
),
56-
mcp.WithToolAnnotation(mcp.ToolAnnotation{
57-
ReadOnlyHint: false,
58-
DestructiveHint: true,
59-
IdempotentHint: false,
60-
OpenWorldHint: false,
61-
}),
62-
)
63-
6453
queryTool := mcp.NewTool("Run-Query",
65-
mcp.WithDescription("Run Query on dgraph db"),
54+
mcp.WithDescription("Run Dgraph Query on dgraph db"),
6655
mcp.WithString("query",
6756
mcp.Required(),
6857
mcp.Description("The query to perform"),
@@ -75,37 +64,60 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
7564
}),
7665
)
7766

78-
mutationTool := mcp.NewTool("Run-Mutation",
79-
mcp.WithDescription("Run Mutation on dgraph db"),
80-
mcp.WithString("mutation",
81-
mcp.Required(),
82-
mcp.Description("The mutation to perform in json format"),
83-
),
84-
mcp.WithToolAnnotation(mcp.ToolAnnotation{
85-
ReadOnlyHint: false,
86-
DestructiveHint: true,
87-
IdempotentHint: false,
88-
OpenWorldHint: false,
89-
}),
90-
)
91-
9267
if readOnly != "true" {
68+
69+
alterSchemaTool := mcp.NewTool("Alter-Schema",
70+
mcp.WithDescription("Alter Dgraph DQL Schema in dgraph db"),
71+
mcp.WithString("schema",
72+
mcp.Required(),
73+
mcp.Description("Updated schema to insert inside the db"),
74+
),
75+
mcp.WithToolAnnotation(mcp.ToolAnnotation{
76+
ReadOnlyHint: false,
77+
DestructiveHint: true,
78+
IdempotentHint: false,
79+
OpenWorldHint: false,
80+
}),
81+
)
82+
9383
s.AddTool(alterSchemaTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
9484
schema, ok := request.Params.Arguments["schema"].(string)
9585
if !ok {
9686
return nil, fmt.Errorf("schema must present")
9787
}
9888

9989
// Execute alter operation
100-
err := conn.SetSchema(ctx, "root", schema)
90+
conn, err := getConn(connectionString)
91+
if err != nil {
92+
return nil, fmt.Errorf("error opening connection: %v", err)
93+
}
94+
err = conn.SetSchema(ctx, "root", schema)
10195
if err != nil {
10296
return nil, fmt.Errorf("schema alteration failed: %v", err)
10397
}
10498

10599
return mcp.NewToolResultText("Schema updated successfully"), nil
106100
})
107101

102+
mutationTool := mcp.NewTool("Run-Mutation",
103+
mcp.WithDescription("Run DQL Mutation on dgraph db"),
104+
mcp.WithString("mutation",
105+
mcp.Required(),
106+
mcp.Description("The mutation to perform in json format"),
107+
),
108+
mcp.WithToolAnnotation(mcp.ToolAnnotation{
109+
ReadOnlyHint: false,
110+
DestructiveHint: true,
111+
IdempotentHint: false,
112+
OpenWorldHint: false,
113+
}),
114+
)
115+
108116
s.AddTool(mutationTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
117+
conn, err := getConn(connectionString)
118+
if err != nil {
119+
return nil, fmt.Errorf("error opening connection: %v", err)
120+
}
109121
txn := conn.NewTxn()
110122
defer func() {
111123
err := txn.Discard(ctx)
@@ -130,6 +142,10 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
130142
}
131143

132144
s.AddTool(queryTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
145+
conn, err := getConn(connectionString)
146+
if err != nil {
147+
return nil, fmt.Errorf("error opening connection: %v", err)
148+
}
133149
txn := conn.NewTxn()
134150
defer func() {
135151
err := txn.Discard(ctx)
@@ -146,6 +162,10 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
146162
})
147163

148164
s.AddTool(schemaTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
165+
conn, err := getConn(connectionString)
166+
if err != nil {
167+
return nil, fmt.Errorf("error opening connection: %v", err)
168+
}
149169
txn := conn.NewTxn()
150170
defer func() {
151171
err := txn.Discard(ctx)
@@ -169,6 +189,10 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
169189

170190
s.AddResource(schemaResource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
171191
// Execute operation
192+
conn, err := getConn(connectionString)
193+
if err != nil {
194+
return nil, fmt.Errorf("error opening connection: %v", err)
195+
}
172196
resp, err := conn.NewTxn().Query(ctx, "schema {}")
173197
if err != nil {
174198
return nil, fmt.Errorf("failed to get schema: %v", err)
@@ -256,6 +280,10 @@ func NewMCPServer(connectionString, readOnly string) (*server.MCPServer, error)
256280
// Add resource with its handler
257281
s.AddResource(schemaResource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
258282
// Execute operation
283+
conn, err := getConn(connectionString)
284+
if err != nil {
285+
return nil, fmt.Errorf("error opening connection: %v", err)
286+
}
259287
resp, err := conn.NewTxn().Query(ctx, "schema {}")
260288
if err != nil {
261289
return nil, fmt.Errorf("failed to get schema: %v", err)

0 commit comments

Comments
 (0)