1use serde_json::{Value, json};
4
5use crate::registry::{ToolName, ToolProfile};
6
7#[derive(Debug, Clone)]
8pub struct ToolSpec {
9 pub name: ToolName,
10 pub description: &'static str,
11 pub input_schema: Value,
12}
13
14pub fn tools_for_profile(profile: ToolProfile) -> Vec<ToolSpec> {
16 let names = ToolName::for_profile(profile);
17 active_tools()
18 .into_iter()
19 .filter(|spec| names.contains(&spec.name))
20 .collect()
21}
22
23pub fn generate_mermaid_erd_for_object(obj: &serde_json::Map<String, Value>) -> Option<String> {
25 let ref_name = obj.get("ref").and_then(|v| v.as_str()).unwrap_or("table");
26 let safe_table_name = ref_name.replace(['.', '-'], "_");
27 let mut diagram = String::from("erDiagram\n");
28 diagram.push_str(&format!(" {safe_table_name} {{\n"));
29 if let Some(columns) = obj.get("columns").and_then(|v| v.as_array()) {
30 for col in columns {
31 let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("col");
32 let data_type = col.get("type").and_then(|v| v.as_str()).unwrap_or("string");
33 let pk = col.get("is_pk").and_then(|v| v.as_bool()).unwrap_or(false);
34 let fk = col.get("is_fk").and_then(|v| v.as_bool()).unwrap_or(false);
35 let key_str = match (pk, fk) {
36 (true, true) => " PK,FK",
37 (true, false) => " PK",
38 (false, true) => " FK",
39 _ => "",
40 };
41 diagram.push_str(&format!(
42 " {} {}{}\n",
43 data_type.replace(' ', "_"),
44 name,
45 key_str
46 ));
47 }
48 }
49 diagram.push_str(" }\n");
50 Some(diagram)
51}
52
53pub fn generate_mermaid_diagram_for_path(path_val: &Value) -> Option<String> {
55 let edges = path_val
56 .as_array()
57 .or_else(|| path_val.get("path").and_then(|v| v.as_array()))?;
58 if edges.is_empty() {
59 return None;
60 }
61 let mut diagram = String::from("erDiagram\n");
62 for edge in edges {
63 let from = edge
64 .get("from")
65 .and_then(|v| v.as_str())
66 .unwrap_or("A")
67 .replace(['.', '-'], "_");
68 let to = edge
69 .get("to")
70 .and_then(|v| v.as_str())
71 .unwrap_or("B")
72 .replace(['.', '-'], "_");
73 let from_col = edge.get("from_col").and_then(|v| v.as_str()).unwrap_or("");
74 let to_col = edge.get("to_col").and_then(|v| v.as_str()).unwrap_or("");
75 diagram.push_str(&format!(
76 " {from} }}|--|| {to} : \"{from_col} -> {to_col}\"\n"
77 ));
78 }
79 Some(diagram)
80}
81
82pub fn phase2_catalog_tools() -> Vec<ToolSpec> {
84 vec![
85 ToolSpec {
86 name: ToolName::ListConnections,
87 description: "List configured connection profiles (never includes passwords).",
88 input_schema: object_schema(&[]),
89 },
90 ToolSpec {
91 name: ToolName::ListDatabases,
92 description: "List databases available for a connection profile.",
93 input_schema: object_schema(&[("connectionId", "string", true)]),
94 },
95 ToolSpec {
96 name: ToolName::ListSchemas,
97 description: "List non-system schemas in the currently selected database.",
98 input_schema: object_schema(&[]),
99 },
100 ToolSpec {
101 name: ToolName::ListObjects,
102 description: "List objects (tables, views, …) in a schema.",
103 input_schema: object_schema(&[("schema", "string", false), ("kind", "string", false)]),
104 },
105 ToolSpec {
106 name: ToolName::GetCurrentContext,
107 description: "Return the active profile, database, and access mode.",
108 input_schema: object_schema(&[]),
109 },
110 ToolSpec {
111 name: ToolName::SwitchConnection,
112 description: "Switch the session to another connection profile / database.",
113 input_schema: object_schema(&[
114 ("connectionId", "string", true),
115 ("database", "string", false),
116 ]),
117 },
118 ToolSpec {
119 name: ToolName::RunSelect,
120 description: "Run a read-only SELECT or WITH query. DML/DDL are rejected. Only reference tables/columns confirmed via list_schemas / list_objects.",
121 input_schema: object_schema(&[("sql", "string", true)]),
122 },
123 ToolSpec {
124 name: ToolName::ExplainQuery,
125 description: "Run EXPLAIN (no ANALYZE execute) for a SELECT/WITH query.",
126 input_schema: object_schema(&[("sql", "string", true)]),
127 },
128 ToolSpec {
129 name: ToolName::DiscoverTools,
130 description: "Dynamically discover and inspect specialized MCP database tools by keyword query (e.g., 'locks', 'bloat', 'index') or category ('query', 'dba', 'write'). Use this when you need specialized tools beyond the core surface.",
131 input_schema: object_schema(&[
132 ("query", "string", false),
133 ("category", "string", false),
134 ]),
135 },
136 ]
137}
138
139pub fn phase3_index_tools() -> Vec<ToolSpec> {
141 vec![
142 ToolSpec {
143 name: ToolName::ResolveTarget,
144 description: "Autonomously find which connection/database matches a user's hint (a database name, environment, host fragment) and/or an object hint (a table/view name), searching across ALL configured connections and their indexed schemas. Call this FIRST whenever the request references a database, environment, or object that is not the current session context — before search_schema, before list_connections. When the match is unambiguous it switches the session context automatically and returns the resolved connection/database; only returns `ambiguous: true` with a candidate list when multiple equally-plausible matches exist, in which case surface those candidates to the user rather than guessing.",
145 input_schema: object_schema(&[("hint", "string", false), ("objectHint", "string", false)]),
146 },
147 ToolSpec {
148 name: ToolName::SearchSchema,
149 description: "Search the live, auto-indexed database schema using natural language or keywords to find tables, views, materialized views, and functions matching the query. Call this FIRST before writing any SQL — do not assume a table exists without finding it here.",
150 input_schema: object_schema(&[("query", "string", true)]),
151 },
152 ToolSpec {
153 name: ToolName::DescribeObject,
154 description: "Get structural details of a specific database object (table, view, or materialized view) including columns, data types, constraints, and indexes.",
155 input_schema: object_schema(&[("ref", "string", true)]),
156 },
157 ToolSpec {
158 name: ToolName::GetJoinPath,
159 description: "Find the shortest path of join relationships and foreign keys between two database tables.",
160 input_schema: object_schema(&[("a", "string", true), ("b", "string", true)]),
161 },
162 ToolSpec {
163 name: ToolName::SampleValues,
164 description: "Retrieve a list of sample values from a specific table column to inspect its contents. Only works on read-only SELECT queries.",
165 input_schema: object_schema(&[("ref", "string", true), ("col", "string", true)]),
166 },
167 ]
168}
169
170pub fn phase4_tools() -> Vec<ToolSpec> {
172 vec![
173 ToolSpec {
174 name: ToolName::GetDdl,
175 description: "Get the DDL / definition of a database object. Views, materialized views, functions, and indexes return their CREATE statement; tables return structured DDL (columns, constraints, indexes).",
176 input_schema: object_schema(&[("ref", "string", true), ("kind", "string", false)]),
177 },
178 ToolSpec {
179 name: ToolName::TableStats,
180 description: "Get size, row-count, activity (scans, inserts/updates/deletes, dead tuples, vacuum/analyze times) and per-column statistics for a specific table.",
181 input_schema: object_schema(&[("ref", "string", true)]),
182 },
183 ToolSpec {
184 name: ToolName::IndexUsage,
185 description: "Get index usage statistics (scan counts, size, definition, type) for a specific table's indexes. Useful for finding unused or missing indexes.",
186 input_schema: object_schema(&[("ref", "string", true)]),
187 },
188 ToolSpec {
189 name: ToolName::ListRunningQueries,
190 description: "List currently executing (non-idle) queries in the connected database with pid, user, state, wait events, and duration.",
191 input_schema: object_schema(&[]),
192 },
193 ToolSpec {
194 name: ToolName::FindBlockingLocks,
195 description: "Find lock contention: which queries are blocked waiting on locks and which pids/queries are blocking them.",
196 input_schema: object_schema(&[]),
197 },
198 ToolSpec {
199 name: ToolName::SlowQueries,
200 description: "List the slowest statements by mean execution time from pg_stat_statements (requires the extension; returns a hint if not installed).",
201 input_schema: object_schema(&[("limit", "number", false)]),
202 },
203 ToolSpec {
204 name: ToolName::DbHealthCheck,
205 description: "Run a database health overview: size/connection stats, cache hit ratio, tables with dead tuples needing vacuum, active connections, and blocking-lock count. Sections that fail are reported individually; partial results are still returned.",
206 input_schema: object_schema(&[]),
207 },
208 ToolSpec {
209 name: ToolName::ExplainAnalyze,
210 description: "Run EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) on a SELECT/WITH query inside a read-only transaction that is always rolled back. WARNING: the query actually executes (volatile functions run), so expect real query runtime.",
211 input_schema: object_schema(&[("sql", "string", true)]),
212 },
213 ToolSpec {
214 name: ToolName::AnalyzeQueryPlan,
215 description: "Run EXPLAIN (FORMAT JSON) on a SELECT/WITH query and return parsed plan metrics (scan counts, bottlenecks, buffer stats) plus performance recommendations. Set analyze=true to also execute the query for actual timings.",
216 input_schema: object_schema(&[("sql", "string", true), ("analyze", "boolean", false)]),
217 },
218 ToolSpec {
219 name: ToolName::GetIndexStatus,
220 description: "Return schema-index status for the active connection/database: indexed_at, fingerprint, object counts, and optional live fingerprint drift.",
221 input_schema: object_schema(&[]),
222 },
223 ToolSpec {
224 name: ToolName::ListExtensions,
225 description: "List installed PostgreSQL extensions (name, version, schema).",
226 input_schema: object_schema(&[]),
227 },
228 ToolSpec {
229 name: ToolName::ServerSettings,
230 description: "Return key PostgreSQL server settings from pg_settings (memory, connections, timeouts, autovacuum, version).",
231 input_schema: object_schema(&[]),
232 },
233 ToolSpec {
234 name: ToolName::SuggestIndexes,
235 description: "Suggest indexes from high sequential-scan tables, unindexed FK columns, and optional pg_stat_statements / EXPLAIN plan heuristics. Pass sql to analyze a specific query plan.",
236 input_schema: object_schema(&[("limit", "number", false), ("sql", "string", false)]),
237 },
238 ToolSpec {
239 name: ToolName::FindUnusedIndexes,
240 description: "List indexes with idx_scan = 0 (never used since stats reset), excluding primary keys, unique indexes, and constraint-backed indexes.",
241 input_schema: object_schema(&[("limit", "number", false)]),
242 },
243 ToolSpec {
244 name: ToolName::BloatReport,
245 description: "Approximate table bloat via dead-tuple ratio from pg_stat_user_tables (simplified estimate — not physical page bloat). Tables with >1000 dead tuples, ordered by bloat %.",
246 input_schema: object_schema(&[("limit", "number", false)]),
247 },
248 ToolSpec {
249 name: ToolName::FindMissingFks,
250 description: "Find likely missing foreign keys: prefers schema-index join-graph inferred edges; falls back to catalog naming (*_id columns without an FK matching a PK).",
251 input_schema: object_schema(&[("limit", "number", false)]),
252 },
253 ]
254}
255
256pub fn phase4b_tools() -> Vec<ToolSpec> {
258 vec![
259 ToolSpec {
260 name: ToolName::ExportQuery,
261 description: "Run a read-only SELECT/WITH and format results as CSV, JSON, or SQL INSERT statements. Honors max-row / max-char caps. For sqlinsert, pass table as schema.name.",
262 input_schema: object_schema(&[
263 ("sql", "string", true),
264 ("format", "string", false),
265 ("table", "string", false),
266 ]),
267 },
268 ToolSpec {
269 name: ToolName::ListRoles,
270 description: "List PostgreSQL roles (attributes). Pass role to get memberships and table privileges for one role.",
271 input_schema: object_schema(&[("role", "string", false)]),
272 },
273 ToolSpec {
274 name: ToolName::DbDashboard,
275 description: "One-shot live metrics bundle: DB size/owner, connection-state breakdown, top tables by size, object counts, active queries, and blocking locks. Soft-fails per section.",
276 input_schema: object_schema(&[]),
277 },
278 ToolSpec {
279 name: ToolName::DeepPlanAnalysis,
280 description: "Run EXPLAIN (ANALYZE by default) and return severity-graded findings: estimate skew, expensive function/CTE/subquery nodes, and recommendations. Set analyze=false for plan-only (no execution).",
281 input_schema: object_schema(&[("sql", "string", true), ("analyze", "boolean", false)]),
282 },
283 ToolSpec {
284 name: ToolName::SchemaDiff,
285 description: "Compare two schemas in the current database (or sourceSchema vs targetSchema). Returns structured table/column/constraint/index diffs. Read-only — does not apply changes.",
286 input_schema: object_schema(&[
287 ("sourceSchema", "string", true),
288 ("targetSchema", "string", true),
289 ]),
290 },
291 ToolSpec {
292 name: ToolName::GenerateMigration,
293 description: "Emit migration SQL to evolve sourceSchema toward targetSchema (from a live schema_diff). Read-only — returns SQL text, never executes it. Destructive drops are commented out.",
294 input_schema: object_schema(&[
295 ("sourceSchema", "string", true),
296 ("targetSchema", "string", true),
297 ]),
298 },
299 ToolSpec {
300 name: ToolName::AutoTuneQuery,
301 description: "Autonomous query tuner: executes EXPLAIN ANALYZE, checks table statistics, evaluates missing indexes, and outputs step-by-step performance tuning recommendations.",
302 input_schema: object_schema(&[("sql", "string", true)]),
303 },
304 ToolSpec {
305 name: ToolName::CheckDdlSafety,
306 description: "Safety guard for migration DDL: inspects SQL for dangerous exclusive locks (e.g. non-concurrent index builds, column drops, table rewrites) and outputs risk scores and safe zero-downtime alternatives.",
307 input_schema: object_schema(&[("ddl", "string", true)]),
308 },
309 ]
310}
311
312pub fn phase9_write_tools() -> Vec<ToolSpec> {
314 vec![
315 ToolSpec {
316 name: ToolName::ExecuteSql,
317 description: "Execute DML (and DDL in admin mode) inside an explicit transaction. Set dry_run=true to roll back after execution. Errors always roll back.",
318 input_schema: object_schema(&[("sql", "string", true), ("dry_run", "boolean", false)]),
319 },
320 ToolSpec {
321 name: ToolName::EditRow,
322 description: "Structured insert, update, or delete by primary key. The server builds parameterized SQL — pass table (schema.name), action, values, and pk for update/delete.",
323 input_schema: object_schema(&[
324 ("table", "string", true),
325 ("action", "string", true),
326 ("values", "object", false),
327 ("pk", "object", false),
328 ]),
329 },
330 ToolSpec {
331 name: ToolName::ImportData,
332 description: "Batch INSERT rows from a JSON array of objects into a table. Optional columns array fixes column order; otherwise keys from the first row are used.",
333 input_schema: object_schema(&[
334 ("table", "string", true),
335 ("rows", "array", true),
336 ("columns", "array", false),
337 ]),
338 },
339 ToolSpec {
340 name: ToolName::ApplyDdl,
341 description: "Apply a DDL statement (CREATE, ALTER, DROP, TRUNCATE, …) in admin mode inside a transaction. Set dry_run=true to roll back.",
342 input_schema: object_schema(&[("sql", "string", true), ("dry_run", "boolean", false)]),
343 },
344 ToolSpec {
345 name: ToolName::CreateIndexConcurrently,
346 description: "Run CREATE INDEX CONCURRENTLY outside a transaction (non-blocking index build). Admin mode only.",
347 input_schema: object_schema(&[("sql", "string", true)]),
348 },
349 ToolSpec {
350 name: ToolName::RunMaintenance,
351 description: "Run VACUUM, ANALYZE, or REINDEX outside a transaction. Admin mode only. Optional table (schema.name); vacuum supports full=true.",
352 input_schema: object_schema(&[
353 ("action", "string", true),
354 ("table", "string", false),
355 ("full", "boolean", false),
356 ]),
357 },
358 ToolSpec {
359 name: ToolName::TerminateQuery,
360 description: "Cancel (pg_cancel_backend) or force-terminate (pg_terminate_backend) a backend by pid. Admin mode only. Refuses superuser targets and the current session.",
361 input_schema: object_schema(&[("pid", "number", true), ("force", "boolean", false)]),
362 },
363 ]
364}
365
366pub fn active_tools() -> Vec<ToolSpec> {
368 let mut specs = phase2_catalog_tools();
369 specs.extend(phase3_index_tools());
370 specs.extend(phase4_tools());
371 specs.extend(phase4b_tools());
372 specs.extend(phase9_write_tools());
373 specs
374}
375
376fn object_schema(props: &[(&str, &str, bool)]) -> Value {
377 let mut properties = serde_json::Map::new();
378 let mut required = Vec::new();
379 for (name, ty, req) in props {
380 let prop_val = match *ty {
381 "array" => match *name {
382 "columns" => json!({ "type": "array", "items": { "type": "string" } }),
383 "rows" => json!({ "type": "array", "items": { "type": "object" } }),
384 _ => json!({ "type": "array", "items": {} }),
385 },
386 _ => json!({ "type": *ty }),
387 };
388 properties.insert((*name).into(), prop_val);
389 if *req {
390 required.push(json!(*name));
391 }
392 }
393 json!({
394 "type": "object",
395 "properties": properties,
396 "required": required
397 })
398}
399
400#[cfg(test)]
401mod tests {
402 use super::*;
403 use crate::registry::ToolName;
404
405 #[test]
406 fn active_tools_lists_forty_five() {
407 let specs = active_tools();
408 assert_eq!(specs.len(), 45);
409 assert_eq!(specs.len(), ToolName::ACTIVE.len());
410 for (spec, name) in specs.iter().zip(ToolName::ACTIVE.iter()) {
411 assert_eq!(spec.name, *name);
412 }
413 }
414
415 #[test]
416 fn phase9_write_tools_count() {
417 assert_eq!(phase9_write_tools().len(), ToolName::PHASE9.len());
418 }
419
420 #[test]
421 fn array_properties_have_items() {
422 for tool in active_tools() {
423 if let Some(props) = tool
424 .input_schema
425 .get("properties")
426 .and_then(|p| p.as_object())
427 {
428 for (prop_name, prop_val) in props {
429 if prop_val.get("type").and_then(|t| t.as_str()) == Some("array") {
430 assert!(
431 prop_val.get("items").is_some(),
432 "Tool '{}' parameter '{}' is array type but missing 'items'",
433 tool.name.as_str(),
434 prop_name
435 );
436 }
437 }
438 }
439 }
440 }
441
442 #[test]
443 fn profile_tools_filtering() {
444 let query_specs = tools_for_profile(ToolProfile::Query);
445 assert_eq!(query_specs.len(), 15);
446
447 let dba_specs = tools_for_profile(ToolProfile::Dba);
448 assert_eq!(dba_specs.len(), 25);
449
450 let meta_specs = tools_for_profile(ToolProfile::Meta);
451 assert_eq!(meta_specs.len(), 6);
452
453 let full_specs = tools_for_profile(ToolProfile::Full);
454 assert_eq!(full_specs.len(), 45);
455 }
456
457 #[test]
458 fn generate_mermaid_erd_test() {
459 let obj = json!({
460 "ref": "public.users",
461 "columns": [
462 { "name": "id", "type": "uuid", "is_pk": true, "is_fk": false },
463 { "name": "email", "type": "varchar", "is_pk": false, "is_fk": false },
464 { "name": "org_id", "type": "uuid", "is_pk": false, "is_fk": true }
465 ]
466 });
467 let diagram = generate_mermaid_erd_for_object(obj.as_object().unwrap()).unwrap();
468 assert!(diagram.contains("erDiagram"));
469 assert!(diagram.contains("public_users"));
470 assert!(diagram.contains("uuid id PK"));
471 assert!(diagram.contains("uuid org_id FK"));
472 }
473
474 #[test]
475 fn generate_mermaid_diagram_for_path_test() {
476 let path = json!([
477 { "from": "public.orders", "to": "public.users", "from_col": "user_id", "to_col": "id" }
478 ]);
479 let diagram = generate_mermaid_diagram_for_path(&path).unwrap();
480 assert!(diagram.contains("erDiagram"));
481 assert!(diagram.contains("public_orders }|--|| public_users : \"user_id -> id\""));
482 }
483}