Skip to content

Commit 35d2ec0

Browse files
general-kroll-4-lifegeneral-kroll-4-life
committed
- Working locally again.
1 parent d7d43a8 commit 35d2ec0

File tree

4 files changed

+51
-239
lines changed

4 files changed

+51
-239
lines changed

‎AGENTS.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ Refer to this section whenever you encounter issues with resource keys containin
6464

6565
## Typed Tools & Resources
6666
- Preferred tools: `run_query(QueryInput)` and `run_query_json(QueryJSONInput)` with validated inputs (via Pydantic) and `row_limit` safeguards.
67-
- Legacy tools `query`/`query_json` remain for backward compatibility.
67+
- Legacy tools `query_v2`/`query_json` remain for backward compatibility. These return a json object with a property for rows.
68+
- Note the `query_v2` requires input of the form `{ "tool": "query", "input": { "sql": "SELECT 1;", "row_limit": 1 } }`
6869
- Table resources: `table://{schema}/{table}` (best-effort registration), with fallback tools `list_table_resources` and `read_table_resource`.
6970
- Prompts available as MCP prompts and tools: `write_safe_select`, `explain_plan_tips`.
7071

‎pkg/mcp_server/backend.go‎

Lines changed: 0 additions & 189 deletions
Original file line numberDiff line numberDiff line change
@@ -5,195 +5,6 @@ import (
55
"database/sql/driver"
66
)
77

8-
// Backend defines the interface for executing queries from MCP clients.
9-
// This abstraction allows for different backend implementations (in-memory, TCP, etc.)
10-
// while maintaining compatibility with the MCP protocol.
11-
12-
/*
13-
The Backend interface should include all of the tools from the below python snippet:
14-
15-
```python
16-
17-
@mcp.tool()
18-
def server_info() -> Dict[str, Any]:
19-
20-
"""Return server and environment info useful for clients."""
21-
return _BACKEND.server_info()
22-
23-
@mcp.tool()
24-
def db_identity() -> Dict[str, Any]:
25-
26-
"""Return current DB identity details: db, user, host, port, search_path, server version, cluster name."""
27-
return _BACKEND.db_identity()
28-
29-
@mcp.tool()
30-
def query(
31-
32-
sql: str,
33-
parameters: Optional[List[Any]] = None,
34-
row_limit: int = 500,
35-
format: str = "markdown",
36-
37-
) -> str:
38-
39-
"""Execute a SQL query (legacy signature). Prefer run_query with typed input."""
40-
return _BACKEND.query(sql, parameters, row_limit, format)
41-
42-
@mcp.tool()
43-
def query_json(sql: str, parameters: Optional[List[Any]] = None, row_limit: int = 500) -> List[Dict[str, Any]]:
44-
45-
"""Execute a SQL query and return JSON-serializable rows (legacy signature). Prefer run_query_json with typed input."""
46-
return _BACKEND.query_json(sql, parameters, row_limit)
47-
48-
@mcp.tool()
49-
def run_query(input: QueryInput) -> str:
50-
51-
"""Execute a SQL query with typed input (preferred)."""
52-
return _BACKEND.run_query(input)
53-
54-
@mcp.tool()
55-
def run_query_json(input: QueryJSONInput) -> List[Dict[str, Any]]:
56-
57-
"""Execute a SQL query and return JSON rows with typed input (preferred)."""
58-
return _BACKEND.run_query_json(input)
59-
60-
@mcp.tool()
61-
def list_table_resources(schema: str = 'public') -> List[str]:
62-
63-
"""List resource URIs for tables in a schema (fallback for clients without resource support)."""
64-
return _BACKEND.list_table_resources(schema)
65-
66-
@mcp.tool()
67-
def read_table_resource(schema: str, table: str, row_limit: int = 100) -> List[Dict[str, Any]]:
68-
69-
"""Read rows from a table resource (fallback)."""
70-
return _BACKEND.read_table_resource(schema, table, row_limit)
71-
72-
# Try to register proper MCP resources if available in FastMCP
73-
74-
try:
75-
76-
resource_decorator = getattr(mcp, "resource")
77-
if callable(resource_decorator):
78-
@resource_decorator("table://{schema}/{table}") # type: ignore
79-
def table_resource(schema: str, table: str, row_limit: int = 100):
80-
"""Resource reader for table rows."""
81-
rows = _BACKEND.read_table_resource(schema, table, row_limit=row_limit)
82-
# Return as JSON string to be universally consumable
83-
return json.dumps(rows, default=str)
84-
85-
except Exception as e:
86-
87-
logger.debug(f"Resource registration skipped: {e}")
88-
89-
try:
90-
91-
prompt_decorator = getattr(mcp, "prompt")
92-
if callable(prompt_decorator):
93-
@prompt_decorator("write_safe_select") # type: ignore
94-
def prompt_write_safe_select():
95-
return _BACKEND.prompt_write_safe_select_tool()
96-
97-
@prompt_decorator("explain_plan_tips") # type: ignore
98-
def prompt_explain_plan_tips():
99-
return _BACKEND.prompt_explain_plan_tips_tool()
100-
101-
except Exception as e:
102-
103-
logger.debug(f"Prompt registration skipped: {e}")
104-
105-
#
106-
107-
@mcp.tool()
108-
def prompt_write_safe_select_tool() -> str:
109-
110-
"""Prompt: guidelines for writing safe SELECT queries."""
111-
return _BACKEND.prompt_write_safe_select_tool()
112-
113-
@mcp.tool()
114-
def prompt_explain_plan_tips_tool() -> str:
115-
116-
"""Prompt: tips for reading EXPLAIN ANALYZE output."""
117-
return _BACKEND.prompt_explain_plan_tips_tool()
118-
119-
@mcp.tool()
120-
def list_schemas_json(input: ListSchemasInput) -> List[Dict[str, Any]]:
121-
122-
"""List schemas with filters and return JSON rows."""
123-
return _BACKEND.list_schemas_json(input)
124-
125-
@mcp.tool()
126-
def list_schemas_json_page(input: ListSchemasPageInput) -> Dict[str, Any]:
127-
128-
"""List schemas with pagination and filters. Returns { items: [...], next_cursor: str|null }"""
129-
return _BACKEND.list_schemas_json_page(input)
130-
131-
@mcp.tool()
132-
def list_tables_json(input: ListTablesInput) -> List[Dict[str, Any]]:
133-
134-
"""List tables in a schema with optional filters and return JSON rows."""
135-
return _BACKEND.list_tables_json(input)
136-
137-
@mcp.tool()
138-
def list_tables_json_page(input: ListTablesPageInput) -> Dict[str, Any]:
139-
140-
"""List tables with pagination and filters. Returns { items, next_cursor }."""
141-
return _BACKEND.list_tables_json_page(input)
142-
143-
@mcp.tool()
144-
def list_schemas() -> str:
145-
146-
"""List all schemas in the database."""
147-
return _BACKEND.list_schemas()
148-
149-
@mcp.tool()
150-
def list_tables(db_schema: Optional[str] = None) -> str:
151-
152-
"""List all tables in a specific schema.
153-
154-
Args:
155-
db_schema: The schema name to list tables from (defaults to 'public')
156-
"""
157-
return _BACKEND.list_tables(db_schema)
158-
159-
@mcp.tool()
160-
def describe_table(table_name: str, db_schema: Optional[str] = None) -> str:
161-
162-
"""Get detailed information about a table.
163-
When dealing with a stackql backend (ie: when the server is initialised to consume stackql using the 'dbapp' parameter), the required query input and returned schema can differ even across the one "resource" (table) object.
164-
This is because stackql has required where parameters for some access methods, where this can vary be SQL verb.
165-
In line with this, stackql responses will contain information about required where parameters, if applicable.
166-
167-
Args:
168-
table_name: The name of the table to describ
169-
db_schema: The schema name (defaults to 'public')
170-
"""
171-
return _BACKEND.describe_table(table_name, db_schema=db_schema)
172-
173-
@mcp.tool()
174-
def get_foreign_keys(table_name: str, db_schema: Optional[str] = None) -> str:
175-
176-
"""Get foreign key information for a table.
177-
178-
Args:
179-
table_name: The name of the table to get foreign keys from
180-
db_schema: The schema name (defaults to 'public')
181-
"""
182-
return _BACKEND.get_foreign_keys(table_name, db_schema)
183-
184-
@mcp.tool()
185-
def find_relationships(table_name: str, db_schema: Optional[str] = None) -> str:
186-
187-
"""Find both explicit and implied relationships for a table.
188-
189-
Args:
190-
table_name: The name of the table to analyze relationships for
191-
db_schema: The schema name (defaults to 'public')
192-
"""
193-
return _BACKEND.find_relationships(table_name, db_schema)
194-
195-
```
196-
*/
1978
type Backend interface {
1989

19910
// Ping verifies the backend connection is active.

‎pkg/mcp_server/dto.go‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -94,16 +94,14 @@ type serverInfoOutput struct {
9494
}
9595

9696
type queryInput struct {
97-
SQL string `json:"sql" yaml:"sql"`
98-
Parameters []string `json:"parameters,omitempty" yaml:"parameters,omitempty"`
99-
RowLimit int `json:"row_limit" yaml:"row_limit"`
100-
Format string `json:"format" yaml:"format"`
97+
SQL string `json:"sql" yaml:"sql"`
98+
RowLimit int `json:"row_limit" yaml:"row_limit"`
99+
Format string `json:"format" yaml:"format"`
101100
}
102101

103102
type queryJSONInput struct {
104-
SQL string `json:"sql" yaml:"sql"`
105-
Parameters []string `json:"parameters,omitempty" yaml:"parameters,omitempty"`
106-
RowLimit int `json:"row_limit" yaml:"row_limit"`
103+
SQL string `json:"sql" yaml:"sql"`
104+
RowLimit int `json:"row_limit" yaml:"row_limit"`
107105
}
108106

109107
type listSchemasInput struct {

‎pkg/mcp_server/server.go‎

Lines changed: 44 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package mcp_server //nolint:revive // fine for now
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"io"
78
"net/http"
@@ -146,48 +147,49 @@ func NewMCPServer(config *Config, backend Backend, logger *logrus.Logger) (MCPSe
146147
return nil, rv, nil
147148
},
148149
)
149-
// mcp.AddTool(
150-
// server,
151-
// &mcp.Tool{
152-
// Name: "query",
153-
// Description: "execute a SQL query",
154-
// // Input and output schemas can be defined here if needed.
155-
// },
156-
// func(ctx context.Context, req *mcp.CallToolRequest, args queryInput) (*mcp.CallToolResult, any, error) {
157-
// rv, rvErr := backend.RunQuery(ctx, args)
158-
// if rvErr != nil {
159-
// return nil, nil, rvErr
160-
// }
161-
// return &mcp.CallToolResult{
162-
// Content: []mcp.Content{
163-
// &mcp.TextContent{Text: rv},
164-
// },
165-
// }, nil, nil
166-
// },
167-
// )
168-
// mcp.AddTool(
169-
// server,
170-
// &mcp.Tool{
171-
// Name: "query_json",
172-
// Description: "execute a SQL query and return a JSON array of rows",
173-
// // Input and output schemas can be defined here if needed.
174-
// },
175-
// func(ctx context.Context, req *mcp.CallToolRequest, args queryJSONInput) (*mcp.CallToolResult, any, error) {
176-
// arr, err := backend.RunQueryJSON(ctx, args)
177-
// if err != nil {
178-
// return nil, nil, err
179-
// }
180-
// bytesArr, marshalErr := json.Marshal(arr)
181-
// if marshalErr != nil {
182-
// return nil, nil, fmt.Errorf("failed to marshal query result to JSON: %w", marshalErr)
183-
// }
184-
// return &mcp.CallToolResult{
185-
// Content: []mcp.Content{
186-
// &mcp.TextContent{Text: string(bytesArr)},
187-
// },
188-
// }, nil, nil
189-
// },
190-
// )
150+
mcp.AddTool(
151+
server,
152+
&mcp.Tool{
153+
Name: "query_v2",
154+
Description: "Execute a SQL query. Please adhere to the expected parameters. Returns a textual response",
155+
// Input and output schemas can be defined here if needed.
156+
},
157+
func(ctx context.Context, req *mcp.CallToolRequest, arg queryInput) (*mcp.CallToolResult, any, error) {
158+
logger.Warnf("Received query: %s", arg.SQL)
159+
rv, rvErr := backend.RunQuery(ctx, arg)
160+
if rvErr != nil {
161+
return nil, nil, rvErr
162+
}
163+
return &mcp.CallToolResult{
164+
Content: []mcp.Content{
165+
&mcp.TextContent{Text: rv},
166+
},
167+
}, nil, nil
168+
},
169+
)
170+
mcp.AddTool(
171+
server,
172+
&mcp.Tool{
173+
Name: "query_json_v2",
174+
Description: "Execute a SQL query and return a JSON array of rows, as text.",
175+
// Input and output schemas can be defined here if needed.
176+
},
177+
func(ctx context.Context, req *mcp.CallToolRequest, args queryJSONInput) (*mcp.CallToolResult, any, error) {
178+
arr, err := backend.RunQueryJSON(ctx, args)
179+
if err != nil {
180+
return nil, nil, err
181+
}
182+
bytesArr, marshalErr := json.Marshal(arr)
183+
if marshalErr != nil {
184+
return nil, nil, fmt.Errorf("failed to marshal query result to JSON: %w", marshalErr)
185+
}
186+
return &mcp.CallToolResult{
187+
Content: []mcp.Content{
188+
&mcp.TextContent{Text: string(bytesArr)},
189+
},
190+
}, nil, nil
191+
},
192+
)
191193

192194
return &simpleMCPServer{
193195
config: config,

0 commit comments

Comments
 (0)