Skip to content

Commit 3006ab4

Browse files
general-kroll-4-lifegeneral-kroll-4-life
committed
- Backend changes.
1 parent a5d6a68 commit 3006ab4

File tree

4 files changed

+437
-31
lines changed

4 files changed

+437
-31
lines changed

‎pkg/mcp_server/backend.go‎

Lines changed: 242 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,255 @@ import (
88
// Backend defines the interface for executing queries from MCP clients.
99
// This abstraction allows for different backend implementations (in-memory, TCP, etc.)
1010
// while maintaining compatibility with the MCP protocol.
11-
type Backend interface {
12-
// Execute runs a query and returns the results.
13-
// The query string and parameters are provided by the MCP client.
14-
Execute(ctx context.Context, query string, params map[string]interface{}) (QueryResult, error)
1511

16-
// GetSchema returns metadata about available resources and their structure.
17-
// This is used by MCP clients to understand what data is available.
18-
GetSchema(ctx context.Context) (SchemaProvider, error)
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+
*/
197+
type Backend interface {
19198

20199
// Ping verifies the backend connection is active.
21200
Ping(ctx context.Context) error
22201

23202
// Close gracefully shuts down the backend connection.
24203
Close() error
204+
// Server and environment info
205+
ServerInfo(ctx context.Context) (map[string]interface{}, error)
206+
207+
// Current DB identity details
208+
DBIdentity(ctx context.Context) (map[string]interface{}, error)
209+
210+
// Execute a SQL query (legacy signature)
211+
Query(ctx context.Context, sql string, parameters []interface{}, rowLimit int, format string) (string, error)
212+
213+
// Execute a SQL query and return JSON-serializable rows (legacy signature)
214+
QueryJSON(ctx context.Context, sql string, parameters []interface{}, rowLimit int) ([]map[string]interface{}, error)
215+
216+
// Execute a SQL query with typed input (preferred)
217+
RunQuery(ctx context.Context, input QueryInput) (string, error)
218+
219+
// Execute a SQL query and return JSON rows with typed input (preferred)
220+
RunQueryJSON(ctx context.Context, input QueryJSONInput) ([]map[string]interface{}, error)
221+
222+
// List resource URIs for tables in a schema
223+
ListTableResources(ctx context.Context, schema string) ([]string, error)
224+
225+
// Read rows from a table resource
226+
ReadTableResource(ctx context.Context, schema string, table string, rowLimit int) ([]map[string]interface{}, error)
227+
228+
// Prompt: guidelines for writing safe SELECT queries
229+
PromptWriteSafeSelectTool(ctx context.Context) (string, error)
230+
231+
// Prompt: tips for reading EXPLAIN ANALYZE output
232+
PromptExplainPlanTipsTool(ctx context.Context) (string, error)
233+
234+
// List schemas with filters and return JSON rows
235+
ListSchemasJSON(ctx context.Context, input ListSchemasInput) ([]map[string]interface{}, error)
236+
237+
// List schemas with pagination and filters
238+
ListSchemasJSONPage(ctx context.Context, input ListSchemasPageInput) (map[string]interface{}, error)
239+
240+
// List tables in a schema with optional filters and return JSON rows
241+
ListTablesJSON(ctx context.Context, input ListTablesInput) ([]map[string]interface{}, error)
242+
243+
// List tables with pagination and filters
244+
ListTablesJSONPage(ctx context.Context, input ListTablesPageInput) (map[string]interface{}, error)
245+
246+
// List all schemas in the database
247+
ListSchemas(ctx context.Context) (string, error)
248+
249+
// List all tables in a specific schema
250+
ListTables(ctx context.Context, dbSchema string) (string, error)
251+
252+
// Get detailed information about a table
253+
DescribeTable(ctx context.Context, tableName string, dbSchema string) (string, error)
254+
255+
// Get foreign key information for a table
256+
GetForeignKeys(ctx context.Context, tableName string, dbSchema string) (string, error)
257+
258+
// Find both explicit and implied relationships for a table
259+
FindRelationships(ctx context.Context, tableName string, dbSchema string) (string, error)
25260
}
26261

27262
// QueryResult represents the result of a query execution.

‎pkg/mcp_server/dto.go‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,120 @@ type GreetingInput struct {
77
type GreetingOutput struct {
88
Greeting string `json:"greeting" jsonschema:"the greeting to tell to the user"`
99
}
10+
11+
/*
12+
13+
Comment AA
14+
15+
Please turn the below python classes into golang structures of the same name with json and yaml attributes exposed
16+
17+
```python
18+
class QueryInput(BaseModel):
19+
sql: str = Field(description="SQL statement to execute")
20+
parameters: Optional[List[Any]] = Field(default=None, description="Positional parameters for the SQL")
21+
row_limit: int = Field(default=500, ge=1, le=10000, description="Max rows to return for SELECT queries")
22+
format: Literal["markdown", "json"] = Field(default="markdown", description="Output format for results")
23+
24+
25+
class QueryJSONInput(BaseModel):
26+
sql: str
27+
parameters: Optional[List[Any]] = None
28+
row_limit: int = 500
29+
30+
class ListSchemasInput(BaseModel):
31+
include_system: bool = Field(default=False, description="Include pg_* and information_schema")
32+
include_temp: bool = Field(default=False, description="Include temporary schemas (pg_temp_*)")
33+
require_usage: bool = Field(default=True, description="Only list schemas with USAGE privilege")
34+
row_limit: int = Field(default=10000, ge=1, le=100000, description="Maximum number of schemas to return")
35+
name_like: Optional[str] = Field(default=None, description="Filter schema names by LIKE pattern (use % and _). '*' and '?' will be translated.")
36+
case_sensitive: bool = Field(default=False, description="When true, use LIKE instead of ILIKE for name_like")
37+
38+
class ListSchemasPageInput(BaseModel):
39+
include_system: bool = False
40+
include_temp: bool = False
41+
require_usage: bool = True
42+
page_size: int = Field(default=500, ge=1, le=10000)
43+
cursor: Optional[str] = None
44+
name_like: Optional[str] = None
45+
case_sensitive: bool = False
46+
47+
48+
class ListTablesInput(BaseModel):
49+
db_schema: Optional[str] = Field(default=None, description="Schema to list tables from; defaults to current_schema()")
50+
name_like: Optional[str] = Field(default=None, description="Filter table_name by pattern; '*' and '?' translate to SQL wildcards")
51+
case_sensitive: bool = Field(default=False, description="Use LIKE (true) or ILIKE (false) for name_like")
52+
table_types: Optional[List[str]] = Field(
53+
default=None,
54+
description="Limit to specific information_schema table_type values (e.g., 'BASE TABLE','VIEW')",
55+
)
56+
row_limit: int = Field(default=10000, ge=1, le=100000)
57+
58+
59+
class ListTablesPageInput(BaseModel):
60+
db_schema: Optional[str] = None
61+
name_like: Optional[str] = None
62+
case_sensitive: bool = False
63+
table_types: Optional[List[str]] = None
64+
page_size: int = Field(default=500, ge=1, le=10000)
65+
cursor: Optional[str] = None
66+
67+
def to_list_tables_input(self) -> ListTablesInput:
68+
return ListTablesInput(
69+
db_schema=self.db_schema,
70+
name_like=self.name_like,
71+
case_sensitive=self.case_sensitive,
72+
table_types=self.table_types,
73+
row_limit=self.page_size
74+
)
75+
```
76+
77+
*/
78+
79+
type QueryInput struct {
80+
SQL string `json:"sql" yaml:"sql"`
81+
Parameters []interface{} `json:"parameters,omitempty" yaml:"parameters,omitempty"`
82+
RowLimit int `json:"row_limit" yaml:"row_limit"`
83+
Format string `json:"format" yaml:"format"`
84+
}
85+
86+
type QueryJSONInput struct {
87+
SQL string `json:"sql" yaml:"sql"`
88+
Parameters []interface{} `json:"parameters,omitempty" yaml:"parameters,omitempty"`
89+
RowLimit int `json:"row_limit" yaml:"row_limit"`
90+
}
91+
92+
type ListSchemasInput struct {
93+
IncludeSystem bool `json:"include_system" yaml:"include_system"`
94+
IncludeTemp bool `json:"include_temp" yaml:"include_temp"`
95+
RequireUsage bool `json:"require_usage" yaml:"require_usage"`
96+
RowLimit int `json:"row_limit" yaml:"row_limit"`
97+
NameLike *string `json:"name_like,omitempty" yaml:"name_like,omitempty"`
98+
CaseSensitive bool `json:"case_sensitive" yaml:"case_sensitive"`
99+
}
100+
101+
type ListSchemasPageInput struct {
102+
IncludeSystem bool `json:"include_system" yaml:"include_system"`
103+
IncludeTemp bool `json:"include_temp" yaml:"include_temp"`
104+
RequireUsage bool `json:"require_usage" yaml:"require_usage"`
105+
PageSize int `json:"page_size" yaml:"page_size"`
106+
Cursor *string `json:"cursor,omitempty" yaml:"cursor,omitempty"`
107+
NameLike *string `json:"name_like,omitempty" yaml:"name_like,omitempty"`
108+
CaseSensitive bool `json:"case_sensitive" yaml:"case_sensitive"`
109+
}
110+
111+
type ListTablesInput struct {
112+
DBSchema *string `json:"db_schema,omitempty" yaml:"db_schema,omitempty"`
113+
NameLike *string `json:"name_like,omitempty" yaml:"name_like,omitempty"`
114+
CaseSensitive bool `json:"case_sensitive" yaml:"case_sensitive"`
115+
TableTypes []string `json:"table_types,omitempty" yaml:"table_types,omitempty"`
116+
RowLimit int `json:"row_limit" yaml:"row_limit"`
117+
}
118+
119+
type ListTablesPageInput struct {
120+
DBSchema *string `json:"db_schema,omitempty" yaml:"db_schema,omitempty"`
121+
NameLike *string `json:"name_like,omitempty" yaml:"name_like,omitempty"`
122+
CaseSensitive bool `json:"case_sensitive" yaml:"case_sensitive"`
123+
TableTypes []string `json:"table_types,omitempty" yaml:"table_types,omitempty"`
124+
PageSize int `json:"page_size" yaml:"page_size"`
125+
Cursor *string `json:"cursor,omitempty" yaml:"cursor,omitempty"`
126+
}

0 commit comments

Comments
 (0)