Skip to content

Commit 26946d6

Browse files
general-kroll-4-lifegeneral-kroll-4-life
authored
mcp-server-improve (#586)
Summary: - MCP `v0.1.1.` - Better MCP server docs. - `json`-centric MCP API. - Minimal `EXPLAIN` supported, for `SELECT` queries only, ie: `EXPLAIN SELECT...`. - Added robot test `Explain Select Repeatably Generates Messages`. - Added robot test `MCP HTTPS Server Validate Canonical`. - Added robot test `MCP HTTPS Server Query Canonical`. - Added robot test `MCP HTTPS List Providers Canonical`. - Added robot test `MCP HTTPS List Services Canonical`. - Added robot test `MCP HTTPS List Resources Canonical`. - Added robot test `MCP HTTPS List Methods Canonical`. - Added robot test `MCP HTTPS Describe Table Canonical`. - Added robot test `MCP HTTPS Server Exec Query Canonical`.
1 parent 9cc0cfe commit 26946d6

File tree

17 files changed

+683
-168
lines changed

17 files changed

+683
-168
lines changed

‎.vscode/launch.json‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -185,9 +185,21 @@
185185
"replace /*+ AWAIT */ google.compute.firewalls set data__disabled = 'true' where project = 'mutable-project' and firewall = 'replacable-firewall' returning *;",
186186
"select role_name from pgi.information_schema.applicable_roles order by role_name desc;",
187187
"SELECT * FROM datadog.organization.users;",
188+
"explain select 1 as foo;",
189+
"explain select name from google.compute.networks where project = 'stackql-demo';",
190+
"explain select * from aws.ec2.instances where region = 'ap-southeast-2';",
188191
],
189192
"default": "show providers;"
190193
},
194+
{
195+
"type": "pickString",
196+
"id": "mcpSrvCfgStr",
197+
"description": "MCP Server Configuration",
198+
"options": [
199+
"{\"server\": {\"transport\": \"http\", \"address\": \"127.0.0.1:9992\"} }"
200+
],
201+
"default": "{\"server\": {\"transport\": \"http\", \"address\": \"127.0.0.1:9992\"} }"
202+
},
191203
{
192204
"type": "pickString",
193205
"id": "registryString",
@@ -385,7 +397,7 @@
385397
"type": "pickString",
386398
"id": "mcpServerType",
387399
"description": "MCP server type",
388-
"default": "stdio",
400+
"default": "http",
389401
"options": [
390402
"http",
391403
"stdio"
@@ -464,6 +476,8 @@
464476
"program": "${workspaceFolder}/stackql",
465477
"args": [
466478
"mcp",
479+
"--mcp.server.type=${input:mcpServerType}",
480+
"--mcp.config=${input:mcpSrvCfgStr}",
467481
"--pgsrv.port=6555",
468482
"--tls.allowInsecure",
469483
"--auth=${input:authString}",
@@ -475,7 +489,6 @@
475489
"--dbInternal=${input:dbInternalString}",
476490
"--export.alias=${input:exportAliasString}",
477491
"--pgsrv.debug.enable=${input:serverDebugPublish}",
478-
"--mcp.server.type=${input:mcpServerType}",
479492
],
480493
},
481494
{

‎docs/mcp.md‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
2+
## Running the MCP server
3+
4+
If necessary, rebuild stackql with:
5+
6+
```bash
7+
python cicd/python/build.py --build
8+
```
9+
10+
**Note**: before starting an MCP server, remember to export all appropriate auth env vars.
11+
12+
We have a nice debug config for running an MCP server with `vscode`, please see [the `vscode` debug launch config](/.vscode/launch.json) for that. Otherwise, you can run with stackql (assuming locally built into `./build/stackql`):
13+
14+
15+
```bash
16+
17+
./build/stackql mcp --mcp.server.type=http --mcp.config '{"server": {"transport": "http", "address": "127.0.0.1:9992"} }'
18+
19+
20+
```
21+
22+
23+
## Using the MCP Client
24+
25+
This is very much a development tool, not currently recommended for production.
26+
27+
Build:
28+
29+
```bash
30+
python cicd/python/build.py --build-mcp-client
31+
```
32+
33+
Then, assuming you have a `stackql` MCP server serving streamable HTTP on port `9992`, you ca access any edpoint. The below examples are somewhat illustrative of a canonical agent pattern for agent behaviour.
34+
35+
36+
```bash
37+
38+
## List available providers.
39+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action list_providers
40+
41+
## List available services.
42+
## **must** supply <provider>
43+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action list_services --exec.args '{"provider": "google"}'
44+
45+
## List available resources.
46+
## **must** supply <provider>, <service>
47+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action list_resources --exec.args '{"provider": "google", "service": "compute"}'
48+
49+
## List access methods.
50+
## **must** supply <provider>, <service>, <resource>
51+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action list_methods --exec.args '{"provider": "google", "service": "compute", "resource": "networks"}'
52+
53+
## Describe published relation
54+
## **must** supply <provider>, <service>, <resource>
55+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action meta_describe_table --exec.args '{"provider": "google", "service": "compute", "resource": "networks"}'
56+
57+
## Validate query AOT. Only works for SELECT at this stage.
58+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action validate_query_json_v2 --exec.args '{"sql": "select name from google.compute.networks where project = '"'"'stackql-demo'"'"';"}'
59+
60+
## Run query
61+
./build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action query_json_v2 --exec.args '{"sql": "select name from google.compute.networks where project = '"'"'stackql-demo'"'"';"}'
62+
63+
## Exec query pattern; for non-read operations
64+
## Tread carefully!!!
65+
## These are almost always mutations
66+
# /build/stackql_mcp_client exec --client-type=http --url=http://127.0.0.1:9992 --exec.action exec_query_json_v2 --exec.args '{"sql": "delete from google.compute.networks where project = '"'"'<my-bucket-name>'"'"';"}'
67+
68+
```

‎internal/stackql/astanalysis/earlyanalysis/first_passes.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,12 @@ func (sp *standardInitialPasses) initialPasses(
219219
return err
220220
}
221221
ast := result.AST
222+
//nolint:gocritic // prefer switch
223+
switch node := ast.(type) {
224+
case *sqlparser.Explain:
225+
ast = node.Statement
226+
}
227+
222228
annotatedAST, err := annotatedast.NewAnnotatedAst(sp.parentAnnotatedAST, ast)
223229
if err != nil {
224230
return err
@@ -235,6 +241,8 @@ func (sp *standardInitialPasses) initialPasses(
235241
switch node := subjectAST.(type) {
236242
case *sqlparser.DDL:
237243
subjectAST = node.SelectStatement
244+
// case *sqlparser.Explain:
245+
// subjectAST = node.Statement
238246
}
239247
pbi, pbiErr := planbuilderinput.NewPlanBuilderInput(
240248
annotatedAST,
@@ -326,7 +334,7 @@ func (sp *standardInitialPasses) initialPasses(
326334
pbi, err := planbuilderinput.NewPlanBuilderInput(
327335
annotatedAST,
328336
handlerCtx,
329-
ast,
337+
result.AST,
330338
threeToFiveAgg.GetTables(),
331339
threeToFiveAgg.GetAliasedColumns(),
332340
threeToFiveAgg.GetAliasMap(),

‎internal/stackql/mcpbackend/mcp_reverse_proxy_backend_service.go‎

Lines changed: 51 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"database/sql"
66
"encoding/json"
77
"fmt"
8+
"time"
89

910
"github.com/sirupsen/logrus"
1011
"github.com/stackql/stackql/internal/stackql/handler"
@@ -79,6 +80,38 @@ func (b *stackqlMCPReverseProxyService) Greet(ctx context.Context, args dto.Gree
7980
return "Hi " + args.Name, nil
8081
}
8182

83+
func (b *stackqlMCPReverseProxyService) ExecQuery(ctx context.Context, query string) (map[string]any, error) {
84+
return b.execQuery(ctx, query)
85+
}
86+
87+
func (b *stackqlMCPReverseProxyService) ValidateQuery(ctx context.Context, query string) ([]map[string]any, error) {
88+
explainQuery := fmt.Sprintf("EXPLAIN %s", query)
89+
rows, err := b.query(ctx, explainQuery, unlimitedRowLimit)
90+
if err != nil {
91+
return nil, err
92+
}
93+
return rows, nil
94+
}
95+
96+
func (b *stackqlMCPReverseProxyService) execQuery(ctx context.Context, query string) (map[string]any, error) {
97+
r, sqlErr := b.db.Exec(query)
98+
if sqlErr != nil {
99+
return nil, sqlErr
100+
}
101+
rowsAffected, rowsAffectedErr := r.RowsAffected()
102+
lastInsertId, lastInsertIdErr := r.LastInsertId()
103+
rv := map[string]any{}
104+
if rowsAffectedErr == nil {
105+
rv["rows_affected"] = rowsAffected
106+
}
107+
if lastInsertIdErr == nil {
108+
rv["last_insert_id"] = lastInsertId
109+
}
110+
oneLinerOutput := time.Now().Format("2006-01-02T15:04:05-07:00 MST")
111+
rv["timestamp"] = oneLinerOutput
112+
return rv, nil
113+
}
114+
82115
//nolint:gocognit,funlen // acceptable
83116
func (b *stackqlMCPReverseProxyService) query(ctx context.Context, query string, rowLimit int) ([]map[string]any, error) {
84117
r, sqlErr := b.db.Query(query)
@@ -230,44 +263,44 @@ func (b *stackqlMCPReverseProxyService) PromptWriteSafeSelectTool(ctx context.Co
230263
// return "stub", nil
231264
// }
232265

233-
func (b *stackqlMCPReverseProxyService) DescribeTable(ctx context.Context, hI dto.HierarchyInput) (string, error) {
266+
func (b *stackqlMCPReverseProxyService) DescribeTable(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
234267
q, qErr := b.interrogator.GetDescribeTable(hI)
235268
if qErr != nil {
236-
return "", qErr
269+
return nil, qErr
237270
}
238-
return b.renderQueryResults(q, hI.Format, hI.RowLimit)
271+
return b.query(ctx, q, hI.RowLimit)
239272
}
240273

241-
func (b *stackqlMCPReverseProxyService) GetForeignKeys(ctx context.Context, hI dto.HierarchyInput) (string, error) {
242-
return b.interrogator.GetForeignKeys(hI)
274+
func (b *stackqlMCPReverseProxyService) GetForeignKeys(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
275+
return nil, fmt.Errorf("GetForeignKeys not implemented")
243276
}
244277

245278
func (b *stackqlMCPReverseProxyService) FindRelationships(ctx context.Context, hI dto.HierarchyInput) (string, error) {
246279
return b.interrogator.FindRelationships(hI)
247280
}
248281

249-
func (b *stackqlMCPReverseProxyService) ListProviders(ctx context.Context) (string, error) {
282+
func (b *stackqlMCPReverseProxyService) ListProviders(ctx context.Context) ([]map[string]interface{}, error) {
250283
q, qErr := b.interrogator.GetShowProviders(dto.HierarchyInput{}, "")
251284
if qErr != nil {
252-
return "", qErr
285+
return nil, qErr
253286
}
254-
return b.renderQueryResults(q, "", unlimitedRowLimit)
287+
return b.query(ctx, q, unlimitedRowLimit)
255288
}
256289

257-
func (b *stackqlMCPReverseProxyService) ListServices(ctx context.Context, hI dto.HierarchyInput) (string, error) {
290+
func (b *stackqlMCPReverseProxyService) ListServices(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
258291
q, qErr := b.interrogator.GetShowServices(hI, "")
259292
if qErr != nil {
260-
return "", qErr
293+
return nil, qErr
261294
}
262-
return b.renderQueryResults(q, hI.Format, hI.RowLimit)
295+
return b.query(ctx, q, hI.RowLimit)
263296
}
264297

265-
func (b *stackqlMCPReverseProxyService) ListResources(ctx context.Context, hI dto.HierarchyInput) (string, error) {
298+
func (b *stackqlMCPReverseProxyService) ListResources(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
266299
q, qErr := b.interrogator.GetShowResources(hI, "")
267300
if qErr != nil {
268-
return "", qErr
301+
return nil, qErr
269302
}
270-
return b.renderQueryResults(q, hI.Format, hI.RowLimit)
303+
return b.query(ctx, q, hI.RowLimit)
271304
}
272305

273306
func (b *stackqlMCPReverseProxyService) ListTablesJSON(ctx context.Context, input dto.ListTablesInput) ([]map[string]interface{}, error) {
@@ -290,14 +323,14 @@ func (b *stackqlMCPReverseProxyService) ListTablesJSONPage(ctx context.Context,
290323
return map[string]interface{}{}, nil
291324
}
292325

293-
func (b *stackqlMCPReverseProxyService) ListTables(ctx context.Context, hI dto.HierarchyInput) (string, error) {
326+
func (b *stackqlMCPReverseProxyService) ListTables(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
294327
return b.ListResources(ctx, hI)
295328
}
296329

297-
func (b *stackqlMCPReverseProxyService) ListMethods(ctx context.Context, hI dto.HierarchyInput) (string, error) {
330+
func (b *stackqlMCPReverseProxyService) ListMethods(ctx context.Context, hI dto.HierarchyInput) ([]map[string]interface{}, error) {
298331
q, qErr := b.interrogator.GetShowMethods(hI)
299332
if qErr != nil {
300-
return "", qErr
333+
return nil, qErr
301334
}
302-
return b.renderQueryResults(q, hI.Format, hI.RowLimit)
335+
return b.query(ctx, q, hI.RowLimit)
303336
}

0 commit comments

Comments
 (0)