Skip to main content

nexql_tools/
schema.rs

1//! Tool descriptors for the active MCP surface (Phase 2–4).
2
3use 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
14/// Tools filtered by the requested `ToolProfile`.
15pub 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
23/// Generate a formatted Mermaid ERD snippet for an object's column & key structure.
24pub 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
53/// Generate a formatted Mermaid ERD diagram snippet for a FK join path.
54pub 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
82/// Phase 2 catalog tools (live Postgres; no index required).
83pub 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
139/// Phase 3 index tools (require `nexql-mcp index build`).
140pub 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
170/// Phase 4 monitoring / DDL tools (descriptions from ToolSpec.ts where available).
171pub 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
256/// Phase 4b read-only breadth (export / role introspection).
257pub 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
312/// Phase 9 write/admin tools (always listed; access-gated at dispatch).
313pub 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
366/// Full tools/list surface for the current phase (catalog + index + Phase 4 + 4b + 9).
367pub 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}