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