1use serde_json::{Value, json};
7
8use crate::registry::{ToolName, ToolProfile};
9
10#[derive(Debug, Clone)]
11pub struct ToolSpec {
12 pub name: ToolName,
13 pub description: &'static str,
14 pub input_schema: Value,
15}
16
17pub fn tools_for_profile(profile: ToolProfile) -> Vec<ToolSpec> {
19 let names = ToolName::for_profile(profile);
20 active_tools()
21 .into_iter()
22 .filter(|spec| names.contains(&spec.name))
23 .collect()
24}
25
26pub fn generate_mermaid_erd_for_object(obj: &serde_json::Map<String, Value>) -> Option<String> {
28 let ref_name = obj.get("ref").and_then(|v| v.as_str()).unwrap_or("table");
29 let safe_table_name = ref_name.replace(['.', '-'], "_");
30 let mut diagram = String::from("erDiagram\n");
31 diagram.push_str(&format!(" {safe_table_name} {{\n"));
32 if let Some(columns) = obj.get("columns").and_then(|v| v.as_array()) {
33 for col in columns {
34 let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("col");
35 let data_type = col.get("type").and_then(|v| v.as_str()).unwrap_or("string");
36 let pk = col.get("is_pk").and_then(|v| v.as_bool()).unwrap_or(false);
37 let fk = col.get("is_fk").and_then(|v| v.as_bool()).unwrap_or(false);
38 let key_str = match (pk, fk) {
39 (true, true) => " PK,FK",
40 (true, false) => " PK",
41 (false, true) => " FK",
42 _ => "",
43 };
44 diagram.push_str(&format!(
45 " {} {}{}\n",
46 data_type.replace(' ', "_"),
47 name,
48 key_str
49 ));
50 }
51 }
52 diagram.push_str(" }\n");
53 Some(diagram)
54}
55
56pub fn generate_mermaid_diagram_for_path(path_val: &Value) -> Option<String> {
58 let edges = path_val
59 .as_array()
60 .or_else(|| path_val.get("path").and_then(|v| v.as_array()))?;
61 if edges.is_empty() {
62 return None;
63 }
64 let mut diagram = String::from("erDiagram\n");
65 for edge in edges {
66 let from = edge
67 .get("from")
68 .and_then(|v| v.as_str())
69 .unwrap_or("A")
70 .replace(['.', '-'], "_");
71 let to = edge
72 .get("to")
73 .and_then(|v| v.as_str())
74 .unwrap_or("B")
75 .replace(['.', '-'], "_");
76 let from_col = edge.get("from_col").and_then(|v| v.as_str()).unwrap_or("");
77 let to_col = edge.get("to_col").and_then(|v| v.as_str()).unwrap_or("");
78 diagram.push_str(&format!(
79 " {from} }}|--|| {to} : \"{from_col} -> {to_col}\"\n"
80 ));
81 }
82 Some(diagram)
83}
84
85pub fn phase2_catalog_tools() -> Vec<ToolSpec> {
87 vec![
88 ToolSpec {
89 name: ToolName::ListConnections,
90 description: "List configured connection profiles (never includes passwords).",
91 input_schema: object_schema(&[]),
92 },
93 ToolSpec {
94 name: ToolName::ListDatabases,
95 description: "List databases available for a connection profile.",
96 input_schema: object_schema(&[("connectionId", "string", true)]),
97 },
98 ToolSpec {
99 name: ToolName::ListSchemas,
100 description: "List non-system schemas in the currently selected database.",
101 input_schema: object_schema(&[]),
102 },
103 ToolSpec {
104 name: ToolName::ListObjects,
105 description: "List objects (tables, views, …) in a schema.",
106 input_schema: object_schema(&[("schema", "string", false), ("kind", "string", false)]),
107 },
108 ToolSpec {
109 name: ToolName::GetCurrentContext,
110 description: "Return the active profile, database, and access mode.",
111 input_schema: object_schema(&[]),
112 },
113 ToolSpec {
114 name: ToolName::SwitchConnection,
115 description: "Switch the session to another connection profile / database.",
116 input_schema: object_schema(&[
117 ("connectionId", "string", true),
118 ("database", "string", false),
119 ]),
120 },
121 ToolSpec {
122 name: ToolName::RunSelect,
123 description: "Run a read-only SELECT or WITH query. DML/DDL are rejected. Only reference tables/columns confirmed via list_schemas / list_objects.",
124 input_schema: object_schema(&[("sql", "string", true)]),
125 },
126 ToolSpec {
127 name: ToolName::ExplainQuery,
128 description: "Run EXPLAIN (no ANALYZE execute) for a SELECT/WITH query.",
129 input_schema: object_schema(&[("sql", "string", true)]),
130 },
131 ToolSpec {
132 name: ToolName::DiscoverTools,
133 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.",
134 input_schema: object_schema(&[
135 ("query", "string", false),
136 ("category", "string", false),
137 ]),
138 },
139 ToolSpec {
140 name: ToolName::RunDoctor,
141 description: "Run diagnostic health checks on active database connection, permissions, session guards, and index status.",
142 input_schema: object_schema(&[]),
143 },
144 ToolSpec {
145 name: ToolName::SetupConnection,
146 description: "Automatically detect or configure a database connection. Scans environment variables, workspace files, and local settings, eliciting missing credentials when supported.",
147 input_schema: object_schema(&[
148 ("name", "string", false),
149 ("url", "string", false),
150 ("host", "string", false),
151 ("port", "number", false),
152 ("dbname", "string", false),
153 ("user", "string", false),
154 ("password", "string", false),
155 ("sslmode", "string", false),
156 ("interactive", "boolean", false),
157 ]),
158 },
159 ToolSpec {
160 name: ToolName::SaveProfile,
161 description: "Save or update a database connection profile in user configuration with atomic backup and dynamic session reload.",
162 input_schema: object_schema(&[
163 ("name", "string", true),
164 ("url", "string", false),
165 ("host", "string", false),
166 ("port", "number", false),
167 ("dbname", "string", false),
168 ("user", "string", false),
169 ("password", "string", false),
170 ("sslmode", "string", false),
171 ("access_mode", "string", false),
172 ("max_rows", "number", false),
173 ]),
174 },
175 ToolSpec {
176 name: ToolName::TestProfile,
177 description: "Test a database connection profile or inline parameters and return server version, superuser status, and round-trip latency.",
178 input_schema: object_schema(&[
179 ("name", "string", false),
180 ("url", "string", false),
181 ("host", "string", false),
182 ("port", "number", false),
183 ("dbname", "string", false),
184 ("user", "string", false),
185 ("password", "string", false),
186 ("sslmode", "string", false),
187 ]),
188 },
189 ToolSpec {
190 name: ToolName::ExportProfile,
191 description: "Export a secret-sanitized TOML configuration for team sharing (.nexql/config.toml) with all passwords and credentials stripped.",
192 input_schema: object_schema(&[("format", "string", false)]),
193 },
194 ToolSpec {
195 name: ToolName::ImportProfile,
196 description: "Import a team configuration file (.nexql/config.toml) or TOML content into local user configuration.",
197 input_schema: object_schema(&[("content", "string", false), ("path", "string", false)]),
198 },
199 ]
200}
201
202pub fn phase3_index_tools() -> Vec<ToolSpec> {
204 vec![
205 ToolSpec {
206 name: ToolName::ResolveTarget,
207 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.",
208 input_schema: object_schema(&[
209 ("hint", "string", false),
210 ("objectHint", "string", false),
211 ]),
212 },
213 ToolSpec {
214 name: ToolName::SearchSchema,
215 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.",
216 input_schema: object_schema(&[("query", "string", true)]),
217 },
218 ToolSpec {
219 name: ToolName::DescribeObject,
220 description: "Get structural details of a specific database object (table, view, or materialized view) including columns, data types, constraints, and indexes.",
221 input_schema: object_schema(&[("ref", "string", true)]),
222 },
223 ToolSpec {
224 name: ToolName::GetJoinPath,
225 description: "Find the shortest path of join relationships and foreign keys between two database tables.",
226 input_schema: object_schema(&[("a", "string", true), ("b", "string", true)]),
227 },
228 ToolSpec {
229 name: ToolName::SampleValues,
230 description: "Retrieve a list of sample values from a specific table column to inspect its contents. Only works on read-only SELECT queries.",
231 input_schema: object_schema(&[("ref", "string", true), ("col", "string", true)]),
232 },
233 ToolSpec {
234 name: ToolName::RebuildIndex,
235 description: "Rebuild the schema index for the active database connection.",
236 input_schema: object_schema(&[("depth", "string", false)]),
237 },
238 ToolSpec {
239 name: ToolName::RefreshIndex,
240 description: "Refresh the schema index for the active database connection using previous build scope.",
241 input_schema: object_schema(&[]),
242 },
243 ]
244}
245
246pub fn phase4_tools() -> Vec<ToolSpec> {
248 vec![
249 ToolSpec {
250 name: ToolName::GetDdl,
251 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).",
252 input_schema: object_schema(&[("ref", "string", true), ("kind", "string", false)]),
253 },
254 ToolSpec {
255 name: ToolName::TableStats,
256 description: "Get size, row-count, activity (scans, inserts/updates/deletes, dead tuples, vacuum/analyze times) and per-column statistics for a specific table.",
257 input_schema: object_schema(&[("ref", "string", true)]),
258 },
259 ToolSpec {
260 name: ToolName::IndexUsage,
261 description: "Get index usage statistics (scan counts, size, definition, type) for a specific table's indexes. Useful for finding unused or missing indexes.",
262 input_schema: object_schema(&[("ref", "string", true)]),
263 },
264 ToolSpec {
265 name: ToolName::ListRunningQueries,
266 description: "List currently executing (non-idle) queries in the connected database with pid, user, state, wait events, and duration.",
267 input_schema: object_schema(&[]),
268 },
269 ToolSpec {
270 name: ToolName::FindBlockingLocks,
271 description: "Find lock contention: which queries are blocked waiting on locks and which pids/queries are blocking them.",
272 input_schema: object_schema(&[]),
273 },
274 ToolSpec {
275 name: ToolName::SlowQueries,
276 description: "List the slowest statements by mean execution time from pg_stat_statements (requires the extension; returns a hint if not installed).",
277 input_schema: object_schema(&[("limit", "number", false)]),
278 },
279 ToolSpec {
280 name: ToolName::DbHealthCheck,
281 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.",
282 input_schema: object_schema(&[]),
283 },
284 ToolSpec {
285 name: ToolName::ExplainAnalyze,
286 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.",
287 input_schema: object_schema(&[("sql", "string", true)]),
288 },
289 ToolSpec {
290 name: ToolName::AnalyzeQueryPlan,
291 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.",
292 input_schema: object_schema(&[("sql", "string", true), ("analyze", "boolean", false)]),
293 },
294 ToolSpec {
295 name: ToolName::GetIndexStatus,
296 description: "Return schema-index status for the active connection/database: indexed_at, fingerprint, object counts, and optional live fingerprint drift.",
297 input_schema: object_schema(&[]),
298 },
299 ToolSpec {
300 name: ToolName::ListExtensions,
301 description: "List installed PostgreSQL extensions (name, version, schema).",
302 input_schema: object_schema(&[]),
303 },
304 ToolSpec {
305 name: ToolName::ServerSettings,
306 description: "Return key PostgreSQL server settings from pg_settings (memory, connections, timeouts, autovacuum, version).",
307 input_schema: object_schema(&[]),
308 },
309 ToolSpec {
310 name: ToolName::SuggestIndexes,
311 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.",
312 input_schema: object_schema(&[("limit", "number", false), ("sql", "string", false)]),
313 },
314 ToolSpec {
315 name: ToolName::FindUnusedIndexes,
316 description: "List indexes with idx_scan = 0 (never used since stats reset), excluding primary keys, unique indexes, and constraint-backed indexes.",
317 input_schema: object_schema(&[("limit", "number", false)]),
318 },
319 ToolSpec {
320 name: ToolName::BloatReport,
321 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 %.",
322 input_schema: object_schema(&[("limit", "number", false)]),
323 },
324 ToolSpec {
325 name: ToolName::FindMissingFks,
326 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).",
327 input_schema: object_schema(&[("limit", "number", false)]),
328 },
329 ]
330}
331
332pub fn phase4b_tools() -> Vec<ToolSpec> {
334 vec![
335 ToolSpec {
336 name: ToolName::ExportQuery,
337 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.",
338 input_schema: object_schema(&[
339 ("sql", "string", true),
340 ("format", "string", false),
341 ("table", "string", false),
342 ]),
343 },
344 ToolSpec {
345 name: ToolName::ListRoles,
346 description: "List PostgreSQL roles (attributes). Pass role to get memberships and table privileges for one role.",
347 input_schema: object_schema(&[("role", "string", false)]),
348 },
349 ToolSpec {
350 name: ToolName::DbDashboard,
351 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.",
352 input_schema: object_schema(&[]),
353 },
354 ToolSpec {
355 name: ToolName::DeepPlanAnalysis,
356 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).",
357 input_schema: object_schema(&[("sql", "string", true), ("analyze", "boolean", false)]),
358 },
359 ToolSpec {
360 name: ToolName::SchemaDiff,
361 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.",
362 input_schema: object_schema(&[
363 ("sourceSchema", "string", true),
364 ("targetSchema", "string", true),
365 ]),
366 },
367 ToolSpec {
368 name: ToolName::GenerateMigration,
369 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.",
370 input_schema: object_schema(&[
371 ("sourceSchema", "string", true),
372 ("targetSchema", "string", true),
373 ]),
374 },
375 ToolSpec {
376 name: ToolName::AutoTuneQuery,
377 description: "Autonomous query tuner: executes EXPLAIN ANALYZE, checks table statistics, evaluates missing indexes, and outputs step-by-step performance tuning recommendations.",
378 input_schema: object_schema(&[("sql", "string", true)]),
379 },
380 ToolSpec {
381 name: ToolName::CheckDdlSafety,
382 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.",
383 input_schema: object_schema(&[("ddl", "string", true)]),
384 },
385 ]
386}
387
388pub fn phase9_write_tools() -> Vec<ToolSpec> {
390 vec![
391 ToolSpec {
392 name: ToolName::ExecuteSql,
393 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.",
394 input_schema: object_schema(&[("sql", "string", true), ("dry_run", "boolean", false)]),
395 },
396 ToolSpec {
397 name: ToolName::EditRow,
398 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.",
399 input_schema: object_schema(&[
400 ("table", "string", true),
401 ("action", "string", true),
402 ("values", "object", false),
403 ("pk", "object", false),
404 ]),
405 },
406 ToolSpec {
407 name: ToolName::ImportData,
408 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.",
409 input_schema: object_schema(&[
410 ("table", "string", true),
411 ("rows", "array", true),
412 ("columns", "array", false),
413 ]),
414 },
415 ToolSpec {
416 name: ToolName::ApplyDdl,
417 description: "Apply a DDL statement (CREATE, ALTER, DROP, TRUNCATE, …) in admin mode inside a transaction. Set dry_run=true to roll back.",
418 input_schema: object_schema(&[("sql", "string", true), ("dry_run", "boolean", false)]),
419 },
420 ToolSpec {
421 name: ToolName::CreateIndexConcurrently,
422 description: "Run CREATE INDEX CONCURRENTLY outside a transaction (non-blocking index build). Admin mode only.",
423 input_schema: object_schema(&[("sql", "string", true)]),
424 },
425 ToolSpec {
426 name: ToolName::RunMaintenance,
427 description: "Run VACUUM, ANALYZE, or REINDEX outside a transaction. Admin mode only. Optional table (schema.name); vacuum supports full=true.",
428 input_schema: object_schema(&[
429 ("action", "string", true),
430 ("table", "string", false),
431 ("full", "boolean", false),
432 ]),
433 },
434 ToolSpec {
435 name: ToolName::TerminateQuery,
436 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.",
437 input_schema: object_schema(&[("pid", "number", true), ("force", "boolean", false)]),
438 },
439 ]
440}
441
442pub fn active_tools() -> Vec<ToolSpec> {
444 let mut specs = phase2_catalog_tools();
445 specs.extend(phase3_index_tools());
446 specs.extend(phase4_tools());
447 specs.extend(phase4b_tools());
448 specs.extend(phase9_write_tools());
449 specs
450}
451
452fn object_schema(props: &[(&str, &str, bool)]) -> Value {
453 let mut properties = serde_json::Map::new();
454 let mut required = Vec::new();
455 for (name, ty, req) in props {
456 let prop_val = match *ty {
457 "array" => match *name {
458 "columns" => json!({ "type": "array", "items": { "type": "string" } }),
459 "rows" => json!({ "type": "array", "items": { "type": "object" } }),
460 _ => json!({ "type": "array", "items": {} }),
461 },
462 _ => json!({ "type": *ty }),
463 };
464 properties.insert((*name).into(), prop_val);
465 if *req {
466 required.push(json!(*name));
467 }
468 }
469 json!({
470 "type": "object",
471 "properties": properties,
472 "required": required
473 })
474}
475
476#[cfg(test)]
477mod tests {
478 use super::*;
479 use crate::registry::ToolName;
480
481 #[test]
482 fn active_tools_lists_fifty_three() {
483 let specs = active_tools();
484 assert_eq!(specs.len(), 53);
485 assert_eq!(specs.len(), ToolName::ACTIVE.len());
486 for (spec, name) in specs.iter().zip(ToolName::ACTIVE.iter()) {
487 assert_eq!(spec.name, *name);
488 }
489 }
490
491 #[test]
492 fn phase9_write_tools_count() {
493 assert_eq!(phase9_write_tools().len(), ToolName::PHASE9.len());
494 }
495
496 #[test]
497 fn array_properties_have_items() {
498 for tool in active_tools() {
499 if let Some(props) = tool
500 .input_schema
501 .get("properties")
502 .and_then(|p| p.as_object())
503 {
504 for (prop_name, prop_val) in props {
505 if prop_val.get("type").and_then(|t| t.as_str()) == Some("array") {
506 assert!(
507 prop_val.get("items").is_some(),
508 "Tool '{}' parameter '{}' is array type but missing 'items'",
509 tool.name.as_str(),
510 prop_name
511 );
512 }
513 }
514 }
515 }
516 }
517
518 #[test]
519 fn profile_tools_filtering() {
520 let query_specs = tools_for_profile(ToolProfile::Query);
521 assert_eq!(query_specs.len(), 18);
522
523 let dba_specs = tools_for_profile(ToolProfile::Dba);
524 assert_eq!(dba_specs.len(), 28);
525
526 let meta_specs = tools_for_profile(ToolProfile::Meta);
527 assert_eq!(meta_specs.len(), 10);
528
529 let full_specs = tools_for_profile(ToolProfile::Full);
530 assert_eq!(full_specs.len(), 53);
531 }
532
533 #[test]
534 fn generate_mermaid_erd_test() {
535 let obj = json!({
536 "ref": "public.users",
537 "columns": [
538 { "name": "id", "type": "uuid", "is_pk": true, "is_fk": false },
539 { "name": "email", "type": "varchar", "is_pk": false, "is_fk": false },
540 { "name": "org_id", "type": "uuid", "is_pk": false, "is_fk": true }
541 ]
542 });
543 let diagram = generate_mermaid_erd_for_object(obj.as_object().unwrap()).unwrap();
544 assert!(diagram.contains("erDiagram"));
545 assert!(diagram.contains("public_users"));
546 assert!(diagram.contains("uuid id PK"));
547 assert!(diagram.contains("uuid org_id FK"));
548 }
549
550 #[test]
551 fn generate_mermaid_diagram_for_path_test() {
552 let path = json!([
553 { "from": "public.orders", "to": "public.users", "from_col": "user_id", "to_col": "id" }
554 ]);
555 let diagram = generate_mermaid_diagram_for_path(&path).unwrap();
556 assert!(diagram.contains("erDiagram"));
557 assert!(diagram.contains("public_orders }|--|| public_users : \"user_id -> id\""));
558 }
559}