Skip to content

Commit 04aae8b

Browse files
general-kroll-4-lifegeneral-kroll-4-life
committed
- Local working ok.
1 parent 35d2ec0 commit 04aae8b

File tree

4 files changed

+242
-13
lines changed

4 files changed

+242
-13
lines changed

‎pkg/mcp_server/backend.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,10 @@ type Backend interface {
3333
RunQueryJSON(ctx context.Context, input queryJSONInput) ([]map[string]interface{}, error)
3434

3535
// List resource URIs for tables in a schema
36-
ListTableResources(ctx context.Context, schema string) ([]string, error)
36+
ListTableResources(ctx context.Context, hI hierarchyInput) ([]string, error)
3737

3838
// Read rows from a table resource
39-
ReadTableResource(ctx context.Context, schema string, table string, rowLimit int) ([]map[string]interface{}, error)
39+
ReadTableResource(ctx context.Context, hI hierarchyInput) ([]map[string]interface{}, error)
4040

4141
// Prompt: guidelines for writing safe SELECT queries
4242
PromptWriteSafeSelectTool(ctx context.Context) (string, error)

‎pkg/mcp_server/dto.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ type hierarchyInput struct {
8585
Service string `json:"service" yaml:"service"`
8686
Resource string `json:"resource" yaml:"resource"`
8787
Method string `json:"method" yaml:"method"`
88+
RowLimit int `json:"row_limit" yaml:"row_limit"`
8889
}
8990

9091
type serverInfoOutput struct {

‎pkg/mcp_server/example_backend.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,11 +55,11 @@ func (b *ExampleBackend) RunQueryJSON(ctx context.Context, input queryJSONInput)
5555
return []map[string]interface{}{}, nil
5656
}
5757

58-
func (b *ExampleBackend) ListTableResources(ctx context.Context, schema string) ([]string, error) {
58+
func (b *ExampleBackend) ListTableResources(ctx context.Context, hI hierarchyInput) ([]string, error) {
5959
return []string{}, nil
6060
}
6161

62-
func (b *ExampleBackend) ReadTableResource(ctx context.Context, schema string, table string, rowLimit int) ([]map[string]interface{}, error) {
62+
func (b *ExampleBackend) ReadTableResource(ctx context.Context, hI hierarchyInput) ([]map[string]interface{}, error) {
6363
return []map[string]interface{}{}, nil
6464
}
6565

‎pkg/mcp_server/server.go‎

Lines changed: 237 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -43,14 +43,6 @@ type simpleMCPServer struct {
4343
servers []io.Closer // Track all running servers for cleanup
4444
}
4545

46-
func sayHi(_ context.Context, _ *mcp.CallToolRequest, input GreetingInput) (
47-
*mcp.CallToolResult,
48-
GreetingOutput,
49-
error,
50-
) {
51-
return nil, GreetingOutput{Greeting: "Hi " + input.Name}, nil
52-
}
53-
5446
func (s *simpleMCPServer) runHTTPServer(server *mcp.Server, url string) error {
5547
// Create the streamable HTTP handler.
5648
handler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server {
@@ -98,7 +90,7 @@ func NewMCPServer(config *Config, backend Backend, logger *logrus.Logger) (MCPSe
9890
}
9991

10092
server := mcp.NewServer(
101-
&mcp.Implementation{Name: "greeter", Version: "v0.1.0"},
93+
&mcp.Implementation{Name: "stackql", Version: "v0.1.0"},
10294
nil,
10395
)
10496
mcp.AddTool(
@@ -191,6 +183,242 @@ func NewMCPServer(config *Config, backend Backend, logger *logrus.Logger) (MCPSe
191183
},
192184
)
193185

186+
mcp.AddTool(
187+
server,
188+
&mcp.Tool{
189+
Name: "list_table_resources",
190+
Description: "List resource URIs for tables in a schema.",
191+
},
192+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
193+
result, err := backend.ListTableResources(ctx, args)
194+
if err != nil {
195+
return nil, nil, err
196+
}
197+
return &mcp.CallToolResult{
198+
Content: []mcp.Content{
199+
&mcp.TextContent{Text: fmt.Sprintf("%v", result)},
200+
},
201+
}, result, nil
202+
},
203+
)
204+
205+
mcp.AddTool(
206+
server,
207+
&mcp.Tool{
208+
Name: "read_table_resource",
209+
Description: "Read rows from a table resource.",
210+
},
211+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
212+
result, err := backend.ReadTableResource(ctx, args)
213+
if err != nil {
214+
return nil, nil, err
215+
}
216+
return &mcp.CallToolResult{
217+
Content: []mcp.Content{
218+
&mcp.TextContent{Text: fmt.Sprintf("%v", result)},
219+
},
220+
}, result, nil
221+
},
222+
)
223+
224+
mcp.AddTool(
225+
server,
226+
&mcp.Tool{
227+
Name: "prompt_write_safe_select_tool",
228+
Description: "Prompt: guidelines for writing safe SELECT queries.",
229+
},
230+
func(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) {
231+
result, err := backend.PromptWriteSafeSelectTool(ctx)
232+
if err != nil {
233+
return nil, nil, err
234+
}
235+
return &mcp.CallToolResult{
236+
Content: []mcp.Content{
237+
&mcp.TextContent{Text: result},
238+
},
239+
}, result, nil
240+
},
241+
)
242+
243+
mcp.AddTool(
244+
server,
245+
&mcp.Tool{
246+
Name: "prompt_explain_plan_tips_tool",
247+
Description: "Prompt: tips for reading EXPLAIN ANALYZE output.",
248+
},
249+
func(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) {
250+
result, err := backend.PromptExplainPlanTipsTool(ctx)
251+
if err != nil {
252+
return nil, nil, err
253+
}
254+
return &mcp.CallToolResult{
255+
Content: []mcp.Content{
256+
&mcp.TextContent{Text: result},
257+
},
258+
}, result, nil
259+
},
260+
)
261+
262+
mcp.AddTool(
263+
server,
264+
&mcp.Tool{
265+
Name: "list_tables_json",
266+
Description: "List tables in a schema and return JSON rows.",
267+
},
268+
func(ctx context.Context, req *mcp.CallToolRequest, args listTablesInput) (*mcp.CallToolResult, any, error) {
269+
result, err := backend.ListTablesJSON(ctx, args)
270+
if err != nil {
271+
return nil, nil, err
272+
}
273+
bytesArr, marshalErr := json.Marshal(result)
274+
if marshalErr != nil {
275+
return nil, nil, fmt.Errorf("failed to marshal result to JSON: %w", marshalErr)
276+
}
277+
return &mcp.CallToolResult{
278+
Content: []mcp.Content{
279+
&mcp.TextContent{Text: string(bytesArr)},
280+
},
281+
}, result, nil
282+
},
283+
)
284+
285+
mcp.AddTool(
286+
server,
287+
&mcp.Tool{
288+
Name: "list_tables_json_page",
289+
Description: "List tables with pagination and filters, returns JSON.",
290+
},
291+
func(ctx context.Context, req *mcp.CallToolRequest, args listTablesPageInput) (*mcp.CallToolResult, any, error) {
292+
result, err := backend.ListTablesJSONPage(ctx, args)
293+
if err != nil {
294+
return nil, nil, err
295+
}
296+
bytesArr, marshalErr := json.Marshal(result)
297+
if marshalErr != nil {
298+
return nil, nil, fmt.Errorf("failed to marshal result to JSON: %w", marshalErr)
299+
}
300+
return &mcp.CallToolResult{
301+
Content: []mcp.Content{
302+
&mcp.TextContent{Text: string(bytesArr)},
303+
},
304+
}, result, nil
305+
},
306+
)
307+
308+
mcp.AddTool(
309+
server,
310+
&mcp.Tool{
311+
Name: "list_providers",
312+
Description: "List all schemas/providers in the database.",
313+
},
314+
func(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) {
315+
result, err := backend.ListProviders(ctx)
316+
if err != nil {
317+
return nil, nil, err
318+
}
319+
return &mcp.CallToolResult{
320+
Content: []mcp.Content{
321+
&mcp.TextContent{Text: result},
322+
},
323+
}, result, nil
324+
},
325+
)
326+
327+
mcp.AddTool(
328+
server,
329+
&mcp.Tool{
330+
Name: "list_services",
331+
Description: "List services for a provider.",
332+
},
333+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
334+
result, err := backend.ListServices(ctx, args)
335+
if err != nil {
336+
return nil, nil, err
337+
}
338+
return &mcp.CallToolResult{
339+
Content: []mcp.Content{
340+
&mcp.TextContent{Text: result},
341+
},
342+
}, result, nil
343+
},
344+
)
345+
346+
mcp.AddTool(
347+
server,
348+
&mcp.Tool{
349+
Name: "list_resources",
350+
Description: "List resources for a service.",
351+
},
352+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
353+
result, err := backend.ListResources(ctx, args)
354+
if err != nil {
355+
return nil, nil, err
356+
}
357+
return &mcp.CallToolResult{
358+
Content: []mcp.Content{
359+
&mcp.TextContent{Text: result},
360+
},
361+
}, result, nil
362+
},
363+
)
364+
365+
mcp.AddTool(
366+
server,
367+
&mcp.Tool{
368+
Name: "describe_table",
369+
Description: "Get detailed information about a table.",
370+
},
371+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
372+
result, err := backend.DescribeTable(ctx, args)
373+
if err != nil {
374+
return nil, nil, err
375+
}
376+
return &mcp.CallToolResult{
377+
Content: []mcp.Content{
378+
&mcp.TextContent{Text: result},
379+
},
380+
}, result, nil
381+
},
382+
)
383+
384+
mcp.AddTool(
385+
server,
386+
&mcp.Tool{
387+
Name: "get_foreign_keys",
388+
Description: "Get foreign key information for a table.",
389+
},
390+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
391+
result, err := backend.GetForeignKeys(ctx, args)
392+
if err != nil {
393+
return nil, nil, err
394+
}
395+
return &mcp.CallToolResult{
396+
Content: []mcp.Content{
397+
&mcp.TextContent{Text: result},
398+
},
399+
}, result, nil
400+
},
401+
)
402+
403+
mcp.AddTool(
404+
server,
405+
&mcp.Tool{
406+
Name: "find_relationships",
407+
Description: "Find explicit and implied relationships for a table.",
408+
},
409+
func(ctx context.Context, req *mcp.CallToolRequest, args hierarchyInput) (*mcp.CallToolResult, any, error) {
410+
result, err := backend.FindRelationships(ctx, args)
411+
if err != nil {
412+
return nil, nil, err
413+
}
414+
return &mcp.CallToolResult{
415+
Content: []mcp.Content{
416+
&mcp.TextContent{Text: result},
417+
},
418+
}, result, nil
419+
},
420+
)
421+
194422
return &simpleMCPServer{
195423
config: config,
196424
backend: backend,

0 commit comments

Comments
 (0)