1use std::collections::HashMap;
10
11use serde::{Deserialize, Serialize};
12use thiserror::Error;
13
14pub const ERR_INVALID_PARAMS: i32 = -32602;
15
16#[derive(Debug, Error)]
17pub enum PromptError {
18 #[error("{0}")]
19 InvalidParams(String),
20}
21
22impl PromptError {
23 pub fn code(&self) -> i32 {
24 ERR_INVALID_PARAMS
25 }
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct PromptArgument {
30 pub name: String,
31 pub description: String,
32 pub required: bool,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PromptInfo {
37 pub name: String,
38 pub description: String,
39 pub arguments: Vec<PromptArgument>,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct PromptMessageContent {
44 #[serde(rename = "type")]
45 pub type_: String,
46 pub text: String,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct PromptMessage {
51 pub role: String,
52 pub content: PromptMessageContent,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct PromptGetResult {
57 pub description: String,
58 pub messages: Vec<PromptMessage>,
59}
60
61struct PromptDef {
62 name: &'static str,
63 description: &'static str,
64 arguments: &'static [ArgDef],
65 build: fn(&HashMap<String, String>) -> String,
66}
67
68struct ArgDef {
69 name: &'static str,
70 description: &'static str,
71 required: bool,
72}
73
74const PROMPTS: &[PromptDef] = &[
75 PromptDef {
76 name: "health-check",
77 description: "Run a full database health assessment and summarize issues by severity.",
78 arguments: &[],
79 build: |_| {
80 [
81 "Assess the health of the connected PostgreSQL database:",
82 "1. Run the db_health_check tool for the overview (size, connections, cache hit ratio, dead tuples).",
83 "2. Run find_blocking_locks to check for lock contention.",
84 "3. Run list_running_queries to spot long-running or stuck queries.",
85 "Then produce a summary grouped by severity (critical / warning / ok):",
86 "- Flag cache hit ratio below 0.95, any blocking locks, queries running longer than 5 minutes, and tables with high dead-tuple counts.",
87 "- For each issue, state the evidence and a concrete remediation (e.g. VACUUM, index, terminate pid).",
88 ]
89 .join("\n")
90 },
91 },
92 PromptDef {
93 name: "analyze-slow-queries",
94 description: "Find the slowest queries and propose index or rewrite improvements.",
95 arguments: &[],
96 build: |_| {
97 [
98 "Identify and improve the slowest queries in the connected database:",
99 "1. Run the slow_queries tool to get the top statements by mean execution time.",
100 "2. For each of the top 3 offenders, run deep_plan_analysis on the query text to get plan metrics and bottlenecks.",
101 "3. Before proposing any index, verify the referenced tables and columns exist using describe_object.",
102 "Deliver: for each slow query — the bottleneck (seq scan, spill, misestimate), a proposed fix (CREATE INDEX CONCURRENTLY statement or query rewrite), and the expected impact.",
103 ]
104 .join("\n")
105 },
106 },
107 PromptDef {
108 name: "explore-schema",
109 description: "Explore and summarize the database schema around a topic.",
110 arguments: &[ArgDef {
111 name: "topic",
112 description: "What to explore, e.g. \"orders\", \"user accounts\", \"billing\".",
113 required: true,
114 }],
115 build: |args| {
116 let topic = args.get("topic").map(String::as_str).unwrap_or("");
117 [
118 format!("Explore the database schema related to: {topic}"),
119 "1. Run search_schema with the topic to find relevant tables, views, and functions.".into(),
120 "2. Run describe_object on each of the top hits to get columns, keys, and indexes.".into(),
121 "3. Run get_join_path between related tables to understand how they connect.".into(),
122 "Deliver a schema summary: the core tables with their purpose, key columns, relationships (as a join diagram in text), and any views or functions that operate on them.".into(),
123 ]
124 .join("\n")
125 },
126 },
127 PromptDef {
128 name: "debug-blocking",
129 description: "Diagnose lock contention and identify the root blocking session.",
130 arguments: &[],
131 build: |_| {
132 [
133 "Diagnose lock contention in the connected database:",
134 "1. Run find_blocking_locks to get blocked/blocking pid pairs with their queries.",
135 "2. Run list_running_queries to see the full activity picture (states, wait events, durations).",
136 "Then explain the lock chain: which pid is the root blocker, what query it is running, how long it has been running, and which sessions are waiting on it (directly or transitively).",
137 "Recommend an action: wait, or terminate the root blocker via terminate_query only if --access-mode admin and the user explicitly confirms.",
138 ]
139 .join("\n")
140 },
141 },
142 PromptDef {
143 name: "write-migration",
144 description: "Draft a safe PostgreSQL migration for a described schema change.",
145 arguments: &[ArgDef {
146 name: "change",
147 description: "The schema change to implement, e.g. \"add soft-delete to orders\".",
148 required: true,
149 }],
150 build: |args| {
151 let change = args.get("change").map(String::as_str).unwrap_or("");
152 [
153 format!("Draft a PostgreSQL migration for: {change}"),
154 "1. Use search_schema and describe_object to ground every table/column you will touch in the live index.".into(),
155 "2. If comparing two schemas, run schema_diff then generate_migration (read-only — emits SQL only).".into(),
156 "3. Prefer non-blocking patterns (CREATE INDEX CONCURRENTLY, ADD COLUMN nullable first, backfill, then constrain).".into(),
157 "4. Produce up and down SQL as separate scripts, with a short risk note (locks, rewrite, invalid indexes).".into(),
158 "Do not run write SQL unless the session is --access-mode write|admin and the user explicitly asked to apply.".into(),
159 ]
160 .join("\n")
161 },
162 },
163 PromptDef {
164 name: "diff-schemas",
165 description: "Compare two PostgreSQL schemas and summarize structural differences.",
166 arguments: &[
167 ArgDef {
168 name: "sourceSchema",
169 description: "Current / left schema name (e.g. public).",
170 required: true,
171 },
172 ArgDef {
173 name: "targetSchema",
174 description: "Desired / right schema name.",
175 required: true,
176 },
177 ],
178 build: |args| {
179 let source = args.get("sourceSchema").map(String::as_str).unwrap_or("");
180 let target = args.get("targetSchema").map(String::as_str).unwrap_or("");
181 [
182 format!("Compare schema \"{source}\" to \"{target}\":"),
183 "1. Run schema_diff with sourceSchema and targetSchema.".into(),
184 "2. Summarize added/removed/changed tables and the highest-risk column/constraint changes.".into(),
185 "3. Optionally run generate_migration for review-only SQL (do not execute).".into(),
186 ]
187 .join("\n")
188 },
189 },
190 PromptDef {
191 name: "plan-deep-dive",
192 description: "Deep-analyze a query plan with severity-graded findings.",
193 arguments: &[ArgDef {
194 name: "sql",
195 description: "The SELECT/WITH query to analyze.",
196 required: true,
197 }],
198 build: |args| {
199 let sql = args.get("sql").map(String::as_str).unwrap_or("");
200 [
201 "Deep-dive this query plan:".to_owned(),
202 format!("```sql\n{sql}\n```"),
203 "1. Ground referenced objects with search_schema / describe_object.".into(),
204 "2. Run deep_plan_analysis (analyze=true) for severity-graded skew / CTE / function / subquery findings and parsed plan metrics.".into(),
205 "3. Cross-check with suggest_indexes; propose indexes or rewrites with evidence.".into(),
206 ]
207 .join("\n")
208 },
209 },
210 PromptDef {
211 name: "optimize-table",
212 description: "Analyze a table's indexes, bloat signals, and access patterns; propose improvements.",
213 arguments: &[ArgDef {
214 name: "ref",
215 description: "Table ref as schema.name, e.g. \"public.orders\".",
216 required: true,
217 }],
218 build: |args| {
219 let ref_ = args.get("ref").map(String::as_str).unwrap_or("");
220 [
221 format!("Optimize table {ref_}:"),
222 "1. Run describe_object on the ref to get columns, keys, and indexes.".into(),
223 "2. Run table_stats and index_usage for the same ref.".into(),
224 "3. Cross-check with slow_queries / deep_plan_analysis for statements that hit this table.".into(),
225 "Deliver: unused or redundant indexes, missing indexes (with CREATE INDEX CONCURRENTLY), and VACUUM/ANALYZE advice with evidence.".into(),
226 ]
227 .join("\n")
228 },
229 },
230 PromptDef {
231 name: "explain-this-query",
232 description: "Explain a SQL query against the live schema and propose plan improvements.",
233 arguments: &[ArgDef {
234 name: "sql",
235 description: "The SQL SELECT (or other read query) to explain.",
236 required: true,
237 }],
238 build: |args| {
239 let sql = args.get("sql").map(String::as_str).unwrap_or("");
240 [
241 "Explain and improve this query:".to_owned(),
242 format!("```sql\n{sql}\n```"),
243 "1. Ground every referenced object with describe_object / search_schema before commenting on columns.".into(),
244 "2. Run explain_query (and deep_plan_analysis if available) on the SQL.".into(),
245 "3. Call out seq scans, misestimates, spills, and missing indexes; propose a rewritten query or index when justified.".into(),
246 ]
247 .join("\n")
248 },
249 },
250];
251
252pub struct PromptCatalog;
254
255impl PromptCatalog {
256 pub fn list() -> Vec<PromptInfo> {
257 PROMPTS
258 .iter()
259 .map(|p| PromptInfo {
260 name: p.name.into(),
261 description: p.description.into(),
262 arguments: p
263 .arguments
264 .iter()
265 .map(|a| PromptArgument {
266 name: a.name.into(),
267 description: a.description.into(),
268 required: a.required,
269 })
270 .collect(),
271 })
272 .collect()
273 }
274
275 pub fn get(name: &str, args: &HashMap<String, String>) -> Result<PromptGetResult, PromptError> {
276 let prompt = PROMPTS
277 .iter()
278 .find(|p| p.name == name)
279 .ok_or_else(|| PromptError::InvalidParams(format!("Unknown prompt: {name}")))?;
280
281 for arg in prompt.arguments {
282 if arg.required {
283 let missing = match args.get(arg.name) {
284 None => true,
285 Some(v) if v.is_empty() => true,
286 Some(_) => false,
287 };
288 if missing {
289 return Err(PromptError::InvalidParams(format!(
290 "Missing required argument \"{}\" for prompt \"{name}\"",
291 arg.name
292 )));
293 }
294 }
295 }
296
297 let text = (prompt.build)(args);
298 Ok(PromptGetResult {
299 description: prompt.description.into(),
300 messages: vec![PromptMessage {
301 role: "user".into(),
302 content: PromptMessageContent {
303 type_: "text".into(),
304 text,
305 },
306 }],
307 })
308 }
309
310 pub fn names() -> Vec<&'static str> {
311 PROMPTS.iter().map(|p| p.name).collect()
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318
319 #[test]
320 fn lists_nine_prompts_including_diff_and_deep_plan() {
321 let names = PromptCatalog::names();
322 assert_eq!(names.len(), 9);
323 for expected in [
324 "health-check",
325 "analyze-slow-queries",
326 "explore-schema",
327 "debug-blocking",
328 "write-migration",
329 "diff-schemas",
330 "plan-deep-dive",
331 "optimize-table",
332 "explain-this-query",
333 ] {
334 assert!(names.contains(&expected), "missing {expected}");
335 }
336 let explore = PromptCatalog::list()
337 .into_iter()
338 .find(|p| p.name == "explore-schema")
339 .unwrap();
340 assert_eq!(explore.arguments[0].name, "topic");
341 assert!(explore.arguments[0].required);
342 }
343
344 #[test]
345 fn get_rejects_missing_required_arg() {
346 let err = PromptCatalog::get("explore-schema", &HashMap::new()).unwrap_err();
347 assert_eq!(err.code(), ERR_INVALID_PARAMS);
348 assert!(err.to_string().contains("topic"));
349 }
350
351 #[test]
352 fn get_unknown_prompt() {
353 let err = PromptCatalog::get("nope", &HashMap::new()).unwrap_err();
354 assert!(err.to_string().contains("Unknown prompt"));
355 }
356
357 #[test]
358 fn get_debug_blocking_mentions_tool() {
359 let result = PromptCatalog::get("debug-blocking", &HashMap::new()).unwrap();
360 assert!(
361 result.messages[0]
362 .content
363 .text
364 .contains("find_blocking_locks")
365 );
366 }
367}