1use std::sync::Arc;
7
8use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
9use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
10use nexql_index::{
11 BuildDepth, BuildMode, BuildRequest, CatalogDb, Embedder, IndexQueryService, IndexScope,
12 IndexStore, PgCatalogDb, QueryPolicyFilter, SearchOptions, build_index,
13};
14use nexql_policy::{PolicyFilter, SqlDecision, validate_readonly_sql};
15use rust_decimal::Decimal;
16use serde_json::{Value, json};
17use tokio_postgres::types::{FromSql, Kind, Type};
18use uuid::Uuid;
19
20use crate::error::ToolError;
21use crate::export::{ExportFormat, columns_from_rows, rows_to_csv, rows_to_sql_insert};
22use crate::plan::{analyze_deep_plan, build_explain_sql, extract_plan_metrics};
23use crate::registry::ToolName;
24use crate::schema::{ToolSpec, active_tools};
25use crate::session::ToolSession;
26use crate::sql::{self, REPORT_LIMIT_DEFAULT, SLOW_QUERIES_DEFAULT, parse_ref};
27use crate::write::{
28 apply_ddl, create_index_concurrently, edit_row, execute_sql, import_data, run_maintenance,
29 terminate_query,
30};
31
32const SEARCH_SCHEMA_LIMIT: usize = 10;
34
35const NO_INDEX_HINT: &str =
36 "No schema index configured — call the 'rebuild_index' tool to build an index.";
37
38#[derive(Debug, Clone)]
39pub struct ToolOutcome {
40 pub text: String,
41 pub structured: Option<Value>,
42 pub is_error: bool,
43}
44
45impl ToolOutcome {
46 pub fn ok_json(value: Value) -> Self {
52 let value = ensure_structured_object(value);
53 let text = serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string());
54 Self {
55 text,
56 structured: Some(value),
57 is_error: false,
58 }
59 }
60
61 pub fn err(msg: impl Into<String>) -> Self {
62 let message = msg.into();
63 Self {
64 text: message.clone(),
65 structured: Some(json!({ "error": message })),
66 is_error: true,
67 }
68 }
69}
70
71fn ensure_structured_object(value: Value) -> Value {
73 match value {
74 Value::Array(rows) => json!({ "rows": rows }),
75 other => other,
76 }
77}
78
79pub struct ToolRouter {
80 session: Arc<ToolSession>,
81 index_override: Option<Option<IndexStore>>,
83 use_semantic: bool,
85 embedder: Option<Arc<dyn Embedder>>,
86 specs: Vec<ToolSpec>,
87 managed_extension: bool,
88}
89
90impl ToolRouter {
91 pub fn new(session: Arc<ToolSession>) -> Self {
92 Self {
93 session,
94 index_override: None,
95 use_semantic: false,
96 embedder: None,
97 specs: active_tools(),
98 managed_extension: false,
99 }
100 }
101
102 pub fn with_index_store(session: Arc<ToolSession>, store: Option<IndexStore>) -> Self {
104 Self {
105 session,
106 index_override: Some(store),
107 use_semantic: false,
108 embedder: None,
109 specs: active_tools(),
110 managed_extension: false,
111 }
112 }
113
114 pub fn with_semantic(
116 mut self,
117 use_semantic: bool,
118 embedder: Option<Arc<dyn Embedder>>,
119 ) -> Self {
120 self.use_semantic = use_semantic;
121 self.embedder = embedder;
122 self
123 }
124
125 pub fn with_profile(mut self, profile: crate::registry::ToolProfile) -> Self {
127 self.specs = crate::schema::tools_for_profile(profile);
128 self
129 }
130
131 pub fn with_managed_extension(mut self, enabled: bool) -> Self {
133 self.managed_extension = enabled;
134 if enabled {
135 const BLOCKED: &[ToolName] = &[
136 ToolName::SetupConnection,
137 ToolName::SaveProfile,
138 ToolName::TestProfile,
139 ToolName::ExportProfile,
140 ToolName::ImportProfile,
141 ];
142 self.specs.retain(|s| !BLOCKED.contains(&s.name));
143 }
144 self
145 }
146
147 pub fn specs(&self) -> &[ToolSpec] {
148 &self.specs
149 }
150
151 fn index_store(&self) -> Option<&IndexStore> {
152 match &self.index_override {
153 Some(inner) => inner.as_ref(),
154 None => self.session.index_store.as_ref(),
155 }
156 }
157
158 fn query_filter(&self) -> QueryPolicyFilter {
159 policy_to_query_filter(&self.session.filter())
160 }
161
162 pub async fn call(&self, name: &str, args: Value) -> ToolOutcome {
163 let outcome = match self.call_inner(name, args).await {
164 Ok(out) => out,
165 Err(e) => ToolOutcome::err(e.to_string()),
166 };
167 self.tag_outcome_with_context(outcome).await
168 }
169
170 async fn tag_outcome_with_context(&self, mut outcome: ToolOutcome) -> ToolOutcome {
171 let (connection_id, database) = self.session.active_context().await;
172 let access_mode = match self.session.access_mode() {
173 nexql_policy::AccessMode::Read => "read",
174 nexql_policy::AccessMode::Write => "write",
175 nexql_policy::AccessMode::Admin => "admin",
176 };
177 let mut freshness: Option<serde_json::Value> = None;
178 if let Some(store) = self.session.index_store.as_ref() {
179 let base = store.base_dir(&connection_id, &database);
180 if let Ok(Some(manifest)) = store.read_manifest(&base) {
181 freshness = Some(json!({
182 "indexedAt": manifest.indexed_at,
183 "schemaFingerprint": manifest.schema_fingerprint,
184 "stale": false,
185 }));
186 } else {
187 freshness = Some(json!({ "stale": true, "reason": "no_index" }));
188 }
189 }
190 if let Some(ref mut structured) = outcome.structured {
191 if let Some(obj) = structured.as_object_mut() {
192 if !obj.contains_key("connectionId") {
193 obj.insert("connectionId".into(), json!(connection_id));
194 }
195 if !obj.contains_key("database") {
196 obj.insert("database".into(), json!(database));
197 }
198 if !obj.contains_key("accessMode") {
199 obj.insert("accessMode".into(), json!(access_mode));
200 }
201 if let Some(ref f) = freshness {
202 obj.insert("freshness".into(), f.clone());
203 }
204 }
205 }
206 let header = format!(
207 "[context connectionId={connection_id} database={database} accessMode={access_mode}]\n"
208 );
209 if !outcome.text.starts_with("[context ") {
210 outcome.text = format!("{header}{}", outcome.text);
211 }
212 outcome
213 }
214
215 async fn call_inner(&self, name: &str, args: Value) -> Result<ToolOutcome, ToolError> {
216 let tool = ToolName::parse(name).ok_or_else(|| ToolError::Unknown(name.to_string()))?;
217 match tool {
218 ToolName::ListConnections => Ok(self.list_connections()),
219 ToolName::ListDatabases => self.list_databases(&args).await,
220 ToolName::ListSchemas => self.list_schemas().await,
221 ToolName::ListObjects => self.list_objects(&args).await,
222 ToolName::GetCurrentContext => self.get_current_context().await,
223 ToolName::SwitchConnection => self.switch_connection(&args).await,
224 ToolName::RunSelect => self.run_select(&args).await,
225 ToolName::ExplainQuery => self.explain_query(&args).await,
226 ToolName::SearchSchema => self.search_schema(&args).await,
227 ToolName::DescribeObject => self.describe_object(&args).await,
228 ToolName::GetJoinPath => self.get_join_path(&args).await,
229 ToolName::SampleValues => self.sample_values(&args).await,
230 ToolName::GetDdl => self.get_ddl(&args).await,
231 ToolName::TableStats => self.table_stats(&args).await,
232 ToolName::IndexUsage => self.index_usage(&args).await,
233 ToolName::ListRunningQueries => self.list_running_queries().await,
234 ToolName::FindBlockingLocks => self.find_blocking_locks().await,
235 ToolName::SlowQueries => self.slow_queries(&args).await,
236 ToolName::DbHealthCheck => self.db_health_check().await,
237 ToolName::ExplainAnalyze => self.explain_analyze(&args).await,
238 ToolName::AnalyzeQueryPlan => self.analyze_query_plan(&args).await,
239 ToolName::GetIndexStatus => self.get_index_status().await,
240 ToolName::ListExtensions => self.list_extensions().await,
241 ToolName::ServerSettings => self.server_settings().await,
242 ToolName::SuggestIndexes => self.suggest_indexes(&args).await,
243 ToolName::FindUnusedIndexes => self.find_unused_indexes(&args).await,
244 ToolName::BloatReport => self.bloat_report(&args).await,
245 ToolName::FindMissingFks => self.find_missing_fks(&args).await,
246 ToolName::ExportQuery => self.export_query(&args).await,
247 ToolName::ListRoles => self.list_roles(&args).await,
248 ToolName::DbDashboard => self.db_dashboard().await,
249 ToolName::DeepPlanAnalysis => self.deep_plan_analysis(&args).await,
250 ToolName::SchemaDiff => self.schema_diff(&args).await,
251 ToolName::GenerateMigration => self.generate_migration(&args).await,
252 ToolName::ExecuteSql => self.execute_sql_tool(&args).await,
253 ToolName::EditRow => self.edit_row_tool(&args).await,
254 ToolName::ImportData => self.import_data_tool(&args).await,
255 ToolName::ApplyDdl => self.apply_ddl_tool(&args).await,
256 ToolName::CreateIndexConcurrently => self.create_index_concurrently_tool(&args).await,
257 ToolName::RunMaintenance => self.run_maintenance_tool(&args).await,
258 ToolName::TerminateQuery => self.terminate_query_tool(&args).await,
259 ToolName::ResolveTarget => self.resolve_target(&args).await,
260 ToolName::DiscoverTools => self.discover_tools(&args).await,
261 ToolName::AutoTuneQuery => self.auto_tune_query(&args).await,
262 ToolName::CheckDdlSafety => self.check_ddl_safety_tool(&args).await,
263 ToolName::RebuildIndex => self.rebuild_index_tool(&args).await,
264 ToolName::RefreshIndex => self.refresh_index_tool(&args).await,
265 ToolName::RunDoctor => self.run_doctor_tool().await,
266 ToolName::SetupConnection => self.setup_connection_tool(&args).await,
267 ToolName::SaveProfile => self.save_profile_tool(&args).await,
268 ToolName::TestProfile => self.test_profile_tool(&args).await,
269 ToolName::ExportProfile => self.export_profile_tool(&args).await,
270 ToolName::ImportProfile => self.import_profile_tool(&args).await,
271 }
272 }
273
274 fn require_write(&self) -> Result<(), ToolError> {
275 if !self.session.access_mode().allows_writes() {
276 return Err(ToolError::Execution(
277 "write tools require --access-mode write or admin (current session: read)".into(),
278 ));
279 }
280 Ok(())
281 }
282
283 fn require_admin(&self) -> Result<(), ToolError> {
284 if !self.session.access_mode().allows_admin() {
285 return Err(ToolError::Execution(
286 "admin tools require --access-mode admin".into(),
287 ));
288 }
289 Ok(())
290 }
291
292 async fn execute_sql_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
293 self.require_write()?;
294 let sql = args
295 .get("sql")
296 .and_then(|v| v.as_str())
297 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
298 let dry_run = args
299 .get("dry_run")
300 .and_then(|v| v.as_bool())
301 .unwrap_or(false);
302 execute_sql(&self.session, sql, dry_run).await
303 }
304
305 async fn edit_row_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
306 self.require_write()?;
307 edit_row(&self.session, args).await
308 }
309
310 async fn import_data_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
311 self.require_write()?;
312 import_data(&self.session, args).await
313 }
314
315 async fn apply_ddl_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
316 self.require_admin()?;
317 let sql = args
318 .get("sql")
319 .and_then(|v| v.as_str())
320 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
321 let dry_run = args
322 .get("dry_run")
323 .and_then(|v| v.as_bool())
324 .unwrap_or(false);
325 apply_ddl(&self.session, sql, dry_run).await
326 }
327
328 async fn create_index_concurrently_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
329 self.require_admin()?;
330 let sql = args
331 .get("sql")
332 .and_then(|v| v.as_str())
333 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
334 create_index_concurrently(&self.session, sql).await
335 }
336
337 async fn run_maintenance_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
338 self.require_admin()?;
339 run_maintenance(&self.session, args).await
340 }
341
342 async fn terminate_query_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
343 self.require_admin()?;
344 terminate_query(&self.session, args).await
345 }
346
347 async fn resolve_target(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
350 let hint = args
351 .get("hint")
352 .and_then(|v| v.as_str())
353 .map(str::trim)
354 .filter(|s| !s.is_empty());
355 let object_hint = args
356 .get("objectHint")
357 .and_then(|v| v.as_str())
358 .map(str::trim)
359 .filter(|s| !s.is_empty());
360 if hint.is_none() && object_hint.is_none() {
361 return Err(ToolError::InvalidArgs(
362 "At least one of \"hint\" or \"objectHint\" is required.".into(),
363 ));
364 }
365
366 let connections = &self.session.connections;
367 if connections.is_empty() {
368 return Ok(ToolOutcome::err("No connections configured."));
369 }
370
371 #[derive(Clone)]
372 struct Candidate {
373 connection_id: String,
374 database: String,
375 }
376 fn key_of(c: &Candidate) -> String {
377 format!("{}\u{0}{}", c.connection_id, c.database)
378 }
379
380 let indexed: Vec<(String, String)> = self
381 .index_store()
382 .map(|store| store.list_indexed_databases().unwrap_or_default())
383 .unwrap_or_default();
384
385 let mut seen = std::collections::HashSet::new();
386 let mut candidates: Vec<Candidate> = Vec::new();
387 let mut add_candidate = |connection_id: &str, database: &str| {
388 if !connections.iter().any(|c| c.id == connection_id) {
389 return;
390 }
391 let key = format!("{connection_id}\u{0}{database}");
392 if !seen.insert(key) {
393 return;
394 }
395 candidates.push(Candidate {
396 connection_id: connection_id.to_string(),
397 database: database.to_string(),
398 });
399 };
400 for (cid, db) in &indexed {
401 add_candidate(cid, db);
402 }
403 for c in connections {
404 let db = c.database.clone().unwrap_or_else(|| "postgres".into());
405 add_candidate(&c.id, &db);
406 }
407
408 let mut scored: std::collections::HashMap<String, (Candidate, f64, Vec<String>)> =
409 std::collections::HashMap::new();
410
411 if let Some(hint) = hint {
412 for c in &candidates {
413 let Some(conn) = connections.iter().find(|x| x.id == c.connection_id) else {
414 continue;
415 };
416 let fields: [(&str, &str); 3] = [
417 ("connection name", conn.name.as_str()),
418 ("host", conn.host.as_deref().unwrap_or("")),
419 ("database", c.database.as_str()),
420 ];
421 let mut best = 0.0f64;
422 let mut best_field = "";
423 for (label, value) in fields {
424 let s = fuzzy_score(hint, value);
425 if s > best {
426 best = s;
427 best_field = label;
428 }
429 }
430 if best > 0.0 {
431 let entry = scored
432 .entry(key_of(c))
433 .or_insert_with(|| (c.clone(), 0.0, Vec::new()));
434 entry.1 += best;
435 entry
436 .2
437 .push(format!("{best_field} matched hint \"{hint}\" ({best:.0})"));
438 }
439 }
440 }
441
442 if let Some(object_hint) = object_hint
443 && let Some(store) = self.index_store()
444 {
445 let filter = self.query_filter();
446 for (cid, db) in &indexed {
447 let svc = IndexQueryService::new(store, cid.clone(), db.clone());
448 if let Ok(hits) = svc.search_schema(
449 object_hint,
450 3,
451 Some(&filter),
452 SearchOptions {
453 use_semantic: self.use_semantic,
454 embedder: self.embedder.as_deref(),
455 },
456 ) && let Some(top) = hits.first()
457 {
458 let c = Candidate {
459 connection_id: cid.clone(),
460 database: db.clone(),
461 };
462 let entry = scored
463 .entry(key_of(&c))
464 .or_insert_with(|| (c.clone(), 0.0, Vec::new()));
465 entry.1 += top.score * 10.0;
466 entry.2.push(format!(
467 "schema search for \"{object_hint}\" found {} (score {:.2})",
468 top.ref_, top.score
469 ));
470 }
471 }
472 }
473
474 let mut ranked: Vec<(Candidate, f64, Vec<String>)> = scored.into_values().collect();
475 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
476
477 if ranked.is_empty() {
478 let candidates_json: Vec<Value> = connections
479 .iter()
480 .map(|c| {
481 json!({
482 "connectionId": c.id,
483 "connectionName": c.name,
484 "database": c.database.clone().unwrap_or_else(|| "postgres".into()),
485 })
486 })
487 .collect();
488 return Ok(ToolOutcome::ok_json(json!({
489 "ambiguous": true,
490 "message": format!(
491 "No connection/database matched \"{}\". Choose from the configured connections.",
492 hint.or(object_hint).unwrap_or_default()
493 ),
494 "candidates": candidates_json
495 })));
496 }
497
498 let winner = &ranked[0];
499 let is_tied = ranked
500 .get(1)
501 .is_some_and(|runner_up| runner_up.1 >= winner.1 * 0.85);
502
503 if is_tied {
504 let threshold = winner.1 * 0.85;
505 let tied: Vec<&(Candidate, f64, Vec<String>)> =
506 ranked.iter().filter(|r| r.1 >= threshold).take(5).collect();
507 let candidates_json: Vec<Value> = tied
508 .iter()
509 .filter_map(|(c, score, evidence)| {
510 connections
511 .iter()
512 .find(|x| x.id == c.connection_id)
513 .map(|conn| {
514 json!({
515 "connectionId": c.connection_id,
516 "connectionName": conn.name,
517 "database": c.database,
518 "score": score,
519 "evidence": evidence,
520 })
521 })
522 })
523 .collect();
524 return Ok(ToolOutcome::ok_json(json!({
525 "ambiguous": true,
526 "message": format!("{} equally-plausible candidates matched.", tied.len()),
527 "candidates": candidates_json
528 })));
529 }
530
531 let (winner_candidate, winner_score, winner_evidence) = winner;
532 self.session
533 .switch(
534 &winner_candidate.connection_id,
535 Some(winner_candidate.database.clone()),
536 )
537 .await?;
538 let conn = connections
539 .iter()
540 .find(|x| x.id == winner_candidate.connection_id)
541 .ok_or_else(|| ToolError::Execution("resolved connection vanished".into()))?;
542
543 Ok(ToolOutcome::ok_json(json!({
544 "resolved": true,
545 "connectionId": winner_candidate.connection_id,
546 "connectionName": conn.name,
547 "database": winner_candidate.database,
548 "confidence": winner_score,
549 "evidence": winner_evidence,
550 })))
551 }
552
553 async fn discover_tools(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
554 let query = args
555 .get("query")
556 .and_then(|v| v.as_str())
557 .map(str::to_lowercase);
558 let category = args
559 .get("category")
560 .and_then(|v| v.as_str())
561 .map(str::to_lowercase);
562
563 let all_specs = active_tools();
565 let filtered: Vec<Value> = all_specs
566 .into_iter()
567 .filter(|spec| {
568 if spec.name == ToolName::DiscoverTools {
569 return false;
570 }
571 if let Some(ref cat) = category {
572 match cat.as_str() {
573 "query" if !ToolName::QUERY_PROFILE.contains(&spec.name) => return false,
574 "dba" if !ToolName::DBA_PROFILE.contains(&spec.name) => return false,
575 "write" if !ToolName::PHASE9.contains(&spec.name) => return false,
576 _ => {}
577 }
578 }
579 if let Some(ref q) = query {
580 let name_match = spec.name.as_str().contains(q.as_str());
581 let desc_match = spec.description.to_lowercase().contains(q.as_str());
582 if !name_match && !desc_match {
583 return false;
584 }
585 }
586 true
587 })
588 .map(|spec| {
589 json!({
590 "name": spec.name.as_str(),
591 "description": spec.description,
592 "input_schema": spec.input_schema,
593 })
594 })
595 .collect();
596
597 Ok(ToolOutcome::ok_json(json!({
598 "query": args.get("query"),
599 "category": args.get("category"),
600 "count": filtered.len(),
601 "tools": filtered,
602 })))
603 }
604
605 fn build_tuning_summary(plan_structured: &Option<Value>, suggestions: &Value) -> String {
606 let mut parts = Vec::new();
607 if let Some(structured) = plan_structured {
608 if let Some(metrics) = structured.get("metrics") {
609 if let Some(exec_time) = metrics.get("executionTime").and_then(|v| v.as_f64()) {
610 parts.push(format!("Query executed in {:.2}ms.", exec_time));
611 }
612 if let Some(seq_scans) = metrics.get("sequentialScans").and_then(|v| v.as_u64()) {
613 if seq_scans > 0 {
614 parts.push(format!("Found {seq_scans} sequential scan(s)."));
615 }
616 }
617 }
618 }
619
620 let candidate_count = suggestions
621 .get("high_seq_scan_tables")
622 .and_then(|v| v.as_array())
623 .map(|a| a.len())
624 .unwrap_or(0)
625 + suggestions
626 .get("unindexed_fk_columns")
627 .and_then(|v| v.as_array())
628 .map(|a| a.len())
629 .unwrap_or(0);
630
631 if candidate_count > 0 {
632 parts.push(format!(
633 "{candidate_count} index recommendation(s) identified."
634 ));
635 } else {
636 parts.push("No explicit index candidate recommendations generated.".into());
637 }
638
639 if parts.is_empty() {
640 "Auto-tune evaluation complete. Inspect execution plan and index recommendations."
641 .into()
642 } else {
643 parts.join(" ")
644 }
645 }
646
647 async fn auto_tune_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
648 let sql = args
649 .get("sql")
650 .and_then(|v| v.as_str())
651 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
652
653 let deep_plan = self
654 .deep_plan_analysis(&json!({ "sql": sql, "analyze": true }))
655 .await?;
656
657 let suggestions_res = self.suggest_indexes(&json!({ "sql": sql })).await;
658 let (suggestions_data, suggestions_error) = match suggestions_res {
659 Ok(outcome) => (outcome.structured.unwrap_or(json!([])), None),
660 Err(e) => (json!([]), Some(e.to_string())),
661 };
662
663 let summary_text = Self::build_tuning_summary(&deep_plan.structured, &suggestions_data);
664
665 let mut payload = json!({
666 "target_query": sql,
667 "deep_plan_analysis": deep_plan.structured,
668 "index_suggestions": suggestions_data,
669 "tuning_summary": summary_text,
670 });
671
672 if let Some(err) = suggestions_error {
673 payload["suggestions_error"] = json!(err);
674 }
675
676 Ok(ToolOutcome::ok_json(payload))
677 }
678
679 async fn check_ddl_safety_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
680 let ddl = args
681 .get("ddl")
682 .and_then(|v| v.as_str())
683 .ok_or_else(|| ToolError::InvalidArgs("ddl is required".into()))?;
684
685 let report = crate::dba_guard::analyze_ddl_safety(ddl);
686 Ok(ToolOutcome::ok_json(report))
687 }
688
689 async fn rebuild_index_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
690 let store = self
691 .index_store()
692 .ok_or_else(|| ToolError::Execution("Index store unavailable".into()))?;
693 let (connection_id, database) = self.session.active_context().await;
694 let depth_str = args
695 .get("depth")
696 .and_then(|v| v.as_str())
697 .unwrap_or("structure");
698 let depth: BuildDepth = match depth_str.to_lowercase().as_str() {
699 "profiles" | "full" => BuildDepth::Profiles,
700 _ => BuildDepth::Structure,
701 };
702
703 let req = BuildRequest {
704 connection_id: connection_id.clone(),
705 database: database.clone(),
706 scope: IndexScope {
707 included_schemas: vec![],
708 excluded_objects: vec![],
709 pii_excluded_columns: vec![],
710 },
711 depth,
712 build_mode: BuildMode::Guided,
713 environment: "development".into(),
714 embeddings: self.use_semantic,
715 };
716
717 let client = self.session.checkout().await?;
718 let db = PgCatalogDb::new(&client);
719 let manifest = build_index(store, &db, &req, None, None, self.embedder.as_deref())
720 .await
721 .map_err(|e| ToolError::Execution(format!("Index build failed: {e}")))?;
722
723 Ok(ToolOutcome::ok_json(json!({
724 "status": "completed",
725 "connection_id": connection_id,
726 "database": database,
727 "schema_fingerprint": manifest.schema_fingerprint,
728 "counts": manifest.counts,
729 "build_ms": manifest.stats.build_ms,
730 })))
731 }
732
733 async fn refresh_index_tool(&self, _args: &Value) -> Result<ToolOutcome, ToolError> {
734 let store = self
735 .index_store()
736 .ok_or_else(|| ToolError::Execution("Index store unavailable".into()))?;
737 let (connection_id, database) = self.session.active_context().await;
738 let base = store.base_dir(&connection_id, &database);
739 let manifest = store.read_manifest(&base)?.ok_or_else(|| {
740 ToolError::Execution(
741 "No existing index manifest to refresh — call 'rebuild_index'.".into(),
742 )
743 })?;
744
745 let req = BuildRequest {
746 connection_id: connection_id.clone(),
747 database: database.clone(),
748 scope: manifest.scope,
749 depth: manifest.build_depth,
750 build_mode: manifest.build_mode,
751 environment: manifest.environment,
752 embeddings: self.use_semantic,
753 };
754
755 let client = self.session.checkout().await?;
756 let db = PgCatalogDb::new(&client);
757 let new_manifest = build_index(store, &db, &req, None, None, self.embedder.as_deref())
758 .await
759 .map_err(|e| ToolError::Execution(format!("Index refresh failed: {e}")))?;
760
761 Ok(ToolOutcome::ok_json(json!({
762 "status": "refreshed",
763 "connection_id": connection_id,
764 "database": database,
765 "schema_fingerprint": new_manifest.schema_fingerprint,
766 "counts": new_manifest.counts,
767 "build_ms": new_manifest.stats.build_ms,
768 })))
769 }
770
771 async fn run_doctor_tool(&self) -> Result<ToolOutcome, ToolError> {
772 let (connection_id, database) = self.session.active_context().await;
773 let client = self.session.checkout().await?;
774
775 let version: String = client
776 .query_one("SELECT version()", &[])
777 .await
778 .map_err(|e| ToolError::Execution(e.to_string()))?
779 .get(0);
780
781 let is_super: String = client
782 .query_one("SELECT current_setting('is_superuser')", &[])
783 .await
784 .map_err(|e| ToolError::Execution(e.to_string()))?
785 .get(0);
786 let is_superuser = is_super.eq_ignore_ascii_case("on");
787
788 let ro: String = client
789 .query_one("SHOW default_transaction_read_only", &[])
790 .await
791 .map_err(|e| ToolError::Execution(e.to_string()))?
792 .get(0);
793
794 let timeout: String = client
795 .query_one("SHOW statement_timeout", &[])
796 .await
797 .map_err(|e| ToolError::Execution(e.to_string()))?
798 .get(0);
799
800 let pgs_present: bool = match client
801 .query_one(
802 "SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements')",
803 &[],
804 )
805 .await
806 {
807 Ok(row) => row.get(0),
808 Err(_) => false,
809 };
810
811 let index_status = if let Some(store) = self.index_store() {
812 let base = store.base_dir(&connection_id, &database);
813 match store.read_manifest(&base) {
814 Ok(Some(m)) => json!({
815 "present": true,
816 "indexed_at": m.indexed_at,
817 "fingerprint": m.schema_fingerprint,
818 "tables": m.counts.tables,
819 }),
820 _ => json!({ "present": false }),
821 }
822 } else {
823 json!({ "present": false, "reason": "no_index_store" })
824 };
825
826 let recent_errors = read_recent_log_errors();
827
828 Ok(ToolOutcome::ok_json(json!({
829 "status": "ok",
830 "connection_id": connection_id,
831 "database": database,
832 "version": version.split(',').next().unwrap_or(&version),
833 "access_mode": format!("{:?}", self.session.access_mode()),
834 "superuser": is_superuser,
835 "read_only": ro,
836 "statement_timeout": timeout,
837 "pg_stat_statements": pgs_present,
838 "index": index_status,
839 "recent_errors": recent_errors,
840 })))
841 }
842
843 async fn setup_connection_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
844 let profile_name = args
845 .get("name")
846 .and_then(|v| v.as_str())
847 .unwrap_or("default");
848
849 let candidates = crate::detect::ConnectionDetector::detect_all(None);
850
851 let url = args.get("url").and_then(|v| v.as_str());
852 let host = args.get("host").and_then(|v| v.as_str());
853 let port = args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16);
854 let dbname = args.get("dbname").and_then(|v| v.as_str());
855 let user = args.get("user").and_then(|v| v.as_str());
856 let password = args.get("password").and_then(|v| v.as_str());
857 let sslmode = args.get("sslmode").and_then(|v| v.as_str());
858
859 let best_cand = candidates
860 .iter()
861 .find(|c| c.is_complete)
862 .or_else(|| candidates.first());
863
864 let res_host = host.or_else(|| best_cand.and_then(|c| c.host.as_deref()));
865 let res_port = port.or_else(|| best_cand.and_then(|c| c.port));
866 let res_dbname = dbname.or_else(|| best_cand.and_then(|c| c.dbname.as_deref()));
867 let res_user = user.or_else(|| best_cand.and_then(|c| c.user.as_deref()));
868 let res_password = password.or_else(|| best_cand.and_then(|c| c.password.as_deref()));
869 let res_url = url.or_else(|| best_cand.and_then(|c| c.url.as_deref()));
870 let res_sslmode = sslmode.or_else(|| best_cand.and_then(|c| c.sslmode.as_deref()));
871
872 if res_url.is_none() && (res_host.is_none() || res_dbname.is_none() || res_user.is_none()) {
873 let missing: Vec<&str> = vec![
874 if res_host.is_none() {
875 Some("host")
876 } else {
877 None
878 },
879 if res_dbname.is_none() {
880 Some("dbname")
881 } else {
882 None
883 },
884 if res_user.is_none() {
885 Some("user")
886 } else {
887 None
888 },
889 ]
890 .into_iter()
891 .flatten()
892 .collect();
893
894 return Ok(ToolOutcome::ok_json(json!({
895 "status": "needs_input",
896 "message": "Insufficient connection details. Please supply missing fields.",
897 "detectedCandidates": candidates.iter().map(|c| c.redacted_json()).collect::<Vec<_>>(),
898 "missingFields": missing
899 })));
900 }
901
902 let params = nexql_conn::ConnectionParams {
903 url: res_url.map(String::from),
904 host: res_host.map(String::from),
905 port: res_port,
906 dbname: res_dbname.map(String::from),
907 user: res_user.map(String::from),
908 password: res_password.map(String::from),
909 sslmode: res_sslmode.map(String::from),
910 ..Default::default()
911 };
912
913 match nexql_conn::test_connection(¶ms).await {
914 Ok(report) => {
915 let p_config = nexql_conn::ProfileConfig {
916 url: params.url.clone(),
917 host: params.host.clone(),
918 port: params.port,
919 dbname: params.dbname.clone(),
920 user: params.user.clone(),
921 password: params.password.clone(),
922 sslmode: params.sslmode.clone(),
923 ..Default::default()
924 };
925
926 let path = nexql_conn::ConfigFile::default_path().ok_or_else(|| {
927 ToolError::Execution("Could not resolve config directory".into())
928 })?;
929 let mut cfg = nexql_conn::ConfigFile::load_path(&path).unwrap_or_default();
930 cfg.upsert_profile(profile_name, p_config);
931 let backup = cfg
932 .save(&path)
933 .map_err(|e| ToolError::Execution(e.to_string()))?;
934
935 Ok(ToolOutcome::ok_json(json!({
936 "status": "configured",
937 "profileName": profile_name,
938 "serverVersion": report.server_version,
939 "isSuperuser": report.is_superuser,
940 "latencyMs": report.latency.as_millis(),
941 "configPath": path.to_string_lossy().to_string(),
942 "backup": backup.map(|b| b.to_string_lossy().to_string())
943 })))
944 }
945 Err(e) => Ok(ToolOutcome::ok_json(json!({
946 "status": "failed",
947 "error": e.to_string(),
948 "detectedCandidates": candidates.iter().map(|c| c.redacted_json()).collect::<Vec<_>>()
949 }))),
950 }
951 }
952
953 async fn save_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
954 let name = args
955 .get("name")
956 .and_then(|v| v.as_str())
957 .ok_or_else(|| ToolError::InvalidArgs("name parameter is required".into()))?;
958
959 let p_config = nexql_conn::ProfileConfig {
960 url: args.get("url").and_then(|v| v.as_str()).map(String::from),
961 host: args.get("host").and_then(|v| v.as_str()).map(String::from),
962 port: args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16),
963 dbname: args
964 .get("dbname")
965 .and_then(|v| v.as_str())
966 .map(String::from),
967 user: args.get("user").and_then(|v| v.as_str()).map(String::from),
968 password: args
969 .get("password")
970 .and_then(|v| v.as_str())
971 .map(String::from),
972 sslmode: args
973 .get("sslmode")
974 .and_then(|v| v.as_str())
975 .map(String::from),
976 access_mode: args
977 .get("access_mode")
978 .and_then(|v| v.as_str())
979 .map(String::from),
980 max_rows: args
981 .get("max_rows")
982 .and_then(|v| v.as_u64())
983 .map(|n| n as u32),
984 ..Default::default()
985 };
986
987 let path = nexql_conn::ConfigFile::default_path()
988 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
989
990 let mut cfg = nexql_conn::ConfigFile::load_path(&path).unwrap_or_default();
991 cfg.upsert_profile(name, p_config);
992 let backup = cfg
993 .save(&path)
994 .map_err(|e| ToolError::Execution(e.to_string()))?;
995
996 Ok(ToolOutcome::ok_json(json!({
997 "status": "saved",
998 "profile": name,
999 "configPath": path.to_string_lossy().to_string(),
1000 "backup": backup.map(|b| b.to_string_lossy().to_string())
1001 })))
1002 }
1003
1004 async fn test_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1005 let name = args.get("name").and_then(|v| v.as_str());
1006
1007 let params = if let Some(pname) = name {
1008 let conn = self
1009 .session
1010 .connections
1011 .iter()
1012 .find(|c| c.id == pname)
1013 .ok_or_else(|| ToolError::InvalidArgs(format!("Profile '{pname}' not found")))?;
1014 conn.params.clone()
1015 } else {
1016 nexql_conn::ConnectionParams {
1017 url: args.get("url").and_then(|v| v.as_str()).map(String::from),
1018 host: args.get("host").and_then(|v| v.as_str()).map(String::from),
1019 port: args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16),
1020 dbname: args
1021 .get("dbname")
1022 .and_then(|v| v.as_str())
1023 .map(String::from),
1024 user: args.get("user").and_then(|v| v.as_str()).map(String::from),
1025 password: args
1026 .get("password")
1027 .and_then(|v| v.as_str())
1028 .map(String::from),
1029 sslmode: args
1030 .get("sslmode")
1031 .and_then(|v| v.as_str())
1032 .map(String::from),
1033 ..Default::default()
1034 }
1035 };
1036
1037 match nexql_conn::test_connection(¶ms).await {
1038 Ok(report) => Ok(ToolOutcome::ok_json(json!({
1039 "success": true,
1040 "serverVersion": report.server_version,
1041 "isSuperuser": report.is_superuser,
1042 "latencyMs": report.latency.as_millis()
1043 }))),
1044 Err(e) => Ok(ToolOutcome::ok_json(json!({
1045 "success": false,
1046 "error": e.to_string()
1047 }))),
1048 }
1049 }
1050
1051 async fn export_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1052 let format = args
1053 .get("format")
1054 .and_then(|v| v.as_str())
1055 .unwrap_or("project");
1056 let path = nexql_conn::ConfigFile::default_path()
1057 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1058 let cfg = nexql_conn::ConfigFile::load_path(&path).unwrap_or_default();
1059
1060 if format == "full" {
1061 let sanitized = cfg.export_full_sanitized();
1062 let toml_str = sanitized
1063 .to_toml_string()
1064 .map_err(|e| ToolError::Execution(e.to_string()))?;
1065 Ok(ToolOutcome::ok_json(json!({
1066 "format": "full",
1067 "content": toml_str,
1068 })))
1069 } else {
1070 let proj = cfg.export_shareable();
1071 let toml_str =
1072 toml::to_string_pretty(&proj).map_err(|e| ToolError::Execution(e.to_string()))?;
1073 Ok(ToolOutcome::ok_json(json!({
1074 "format": "project",
1075 "filename": ".nexql/config.toml",
1076 "content": toml_str,
1077 })))
1078 }
1079 }
1080
1081 async fn import_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1082 let content = if let Some(c) = args.get("content").and_then(|v| v.as_str()) {
1083 c.to_string()
1084 } else if let Some(p) = args.get("path").and_then(|v| v.as_str()) {
1085 std::fs::read_to_string(p)
1086 .map_err(|e| ToolError::Execution(format!("failed to read file {p}: {e}")))?
1087 } else {
1088 return Err(ToolError::Execution(
1089 "either 'content' or 'path' must be specified".into(),
1090 ));
1091 };
1092
1093 let path = nexql_conn::ConfigFile::default_path()
1094 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1095 let mut cfg = nexql_conn::ConfigFile::load_path(&path).unwrap_or_default();
1096
1097 let imported: nexql_conn::ConfigFile = toml::from_str(&content)
1098 .map_err(|e| ToolError::Execution(format!("failed to parse TOML content: {e}")))?;
1099
1100 let mut count = 0;
1101 for (name, prof) in imported.profiles {
1102 cfg.upsert_profile(name, prof);
1103 count += 1;
1104 }
1105 if imported.default_profile.is_some() {
1106 cfg.default_profile = imported.default_profile;
1107 }
1108
1109 let backup = cfg
1110 .save(&path)
1111 .map_err(|e| ToolError::Execution(e.to_string()))?;
1112
1113 Ok(ToolOutcome::ok_json(json!({
1114 "status": "imported",
1115 "imported_profiles": count,
1116 "configPath": path.to_string_lossy().to_string(),
1117 "backup": backup.map(|b| b.to_string_lossy().to_string())
1118 })))
1119 }
1120
1121 async fn index_service(&self) -> Result<(&IndexStore, String, String), ToolError> {
1122 let store = self
1123 .index_store()
1124 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1125 let (connection_id, database) = self.session.active_context().await;
1126 let base = store.base_dir(&connection_id, &database);
1127 if store.read_manifest(&base)?.is_none() {
1128 return Err(ToolError::Execution(format!(
1129 "No schema index for database \"{database}\" — call the 'rebuild_index' tool to build an index."
1130 )));
1131 }
1132 Ok((store, connection_id, database))
1133 }
1134
1135 async fn search_schema(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1136 let query = args
1137 .get("query")
1138 .and_then(|v| v.as_str())
1139 .unwrap_or("")
1140 .trim();
1141 if query.is_empty() {
1142 return Ok(ToolOutcome::ok_json(json!([])));
1143 }
1144 let (store, connection_id, database) = self.index_service().await?;
1145 let svc = IndexQueryService::new(store, &connection_id, &database);
1146 let filter = self.query_filter();
1147 let hits = svc.search_schema(
1148 query,
1149 SEARCH_SCHEMA_LIMIT,
1150 Some(&filter),
1151 SearchOptions {
1152 use_semantic: self.use_semantic,
1153 embedder: self.embedder.as_deref(),
1154 },
1155 )?;
1156 let rows: Vec<Value> = hits
1157 .into_iter()
1158 .map(|h| {
1159 json!({
1160 "ref": h.ref_,
1161 "score": h.score,
1162 "kind": h.kind,
1163 })
1164 })
1165 .collect();
1166 Ok(ToolOutcome::ok_json(json!(rows)))
1167 }
1168
1169 async fn describe_object(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1170 let ref_ = args
1171 .get("ref")
1172 .and_then(|v| v.as_str())
1173 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1174 let (store, connection_id, database) = self.index_service().await?;
1175 let svc = IndexQueryService::new(store, &connection_id, &database);
1176 let filter = self.query_filter();
1177 let entry = svc.describe_object(ref_, Some(&filter))?;
1178 let value = serde_json::to_value(entry).map_err(|e| ToolError::Execution(e.to_string()))?;
1179 Ok(ToolOutcome::ok_json(value))
1180 }
1181
1182 async fn get_join_path(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1183 let a = args
1184 .get("a")
1185 .and_then(|v| v.as_str())
1186 .ok_or_else(|| ToolError::InvalidArgs("a is required".into()))?;
1187 let b = args
1188 .get("b")
1189 .and_then(|v| v.as_str())
1190 .ok_or_else(|| ToolError::InvalidArgs("b is required".into()))?;
1191 let (store, connection_id, database) = self.index_service().await?;
1192 let svc = IndexQueryService::new(store, &connection_id, &database);
1193 let path = svc.get_join_path(a, b)?;
1194 let value = serde_json::to_value(path).map_err(|e| ToolError::Execution(e.to_string()))?;
1195 Ok(ToolOutcome::ok_json(value))
1196 }
1197
1198 async fn sample_values(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1199 let ref_ = args
1200 .get("ref")
1201 .and_then(|v| v.as_str())
1202 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1203 let col = args
1204 .get("col")
1205 .and_then(|v| v.as_str())
1206 .ok_or_else(|| ToolError::InvalidArgs("col is required".into()))?;
1207 let (store, connection_id, database) = self.index_service().await?;
1208 let svc = IndexQueryService::new(store, &connection_id, &database);
1209 let filter = self.query_filter();
1210 let result = svc.sample_values(ref_, col, Some(&filter), None)?;
1211
1212 let mut values = result.values;
1213 let mut message = result.message;
1214
1215 if values.is_empty() {
1216 if let Ok(client) = self.session.checkout().await {
1217 let parts: Vec<&str> = ref_.split('.').collect();
1218 let (schema, table) = match parts.as_slice() {
1219 [s, t] => (*s, *t),
1220 _ => ("public", ref_),
1221 };
1222 let safe_schema = schema.replace('"', "\"\"");
1223 let safe_table = table.replace('"', "\"\"");
1224 let safe_col = col.replace('"', "\"\"");
1225 let query = format!(
1226 "SELECT DISTINCT \"{safe_col}\"::text FROM \"{safe_schema}\".\"{safe_table}\" WHERE \"{safe_col}\" IS NOT NULL LIMIT 20"
1227 );
1228 if let Ok(rows) = client.query(&query, &[]).await {
1229 let sampled: Vec<String> = rows
1230 .iter()
1231 .filter_map(|r| r.get::<_, Option<String>>(0))
1232 .collect();
1233 if !sampled.is_empty() {
1234 values = sampled;
1235 message = None;
1236 }
1237 }
1238 }
1239 }
1240
1241 let mut payload = json!({ "values": values });
1242 if let Some(msg) = message {
1243 payload["message"] = json!(msg);
1244 }
1245 Ok(ToolOutcome::ok_json(payload))
1246 }
1247
1248 fn list_connections(&self) -> ToolOutcome {
1249 let rows: Vec<Value> = self
1250 .session
1251 .connections
1252 .iter()
1253 .map(|c| {
1254 json!({
1255 "id": c.id,
1256 "name": c.name,
1257 "host": c.host,
1258 "port": c.port,
1259 "database": c.database,
1260 })
1261 })
1262 .collect();
1263 ToolOutcome::ok_json(json!(rows))
1264 }
1265
1266 async fn list_databases(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1267 let connection_id = args
1268 .get("connectionId")
1269 .and_then(|v| v.as_str())
1270 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
1271 let conn = self
1272 .session
1273 .connections
1274 .iter()
1275 .find(|c| c.id == connection_id)
1276 .ok_or_else(|| {
1277 ToolError::Execution(format!(
1278 "Connection not found for ID: {connection_id} — call list_connections"
1279 ))
1280 })?;
1281 let client = {
1283 if self.session.active_context().await.0 == connection_id {
1285 self.session.checkout().await?
1286 } else {
1287 let pool_opts = self.session.pool_opts();
1288 let pool = nexql_conn::create_pool(&conn.params, &pool_opts).await?;
1289 nexql_conn::checkout_guarded(&pool, &pool_opts).await?
1290 }
1291 };
1292 let rows = client
1293 .query(
1294 "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname",
1295 &[],
1296 )
1297 .await?;
1298 let names: Vec<String> = rows.iter().map(|r| r.get(0)).collect();
1299 Ok(ToolOutcome::ok_json(json!(names)))
1300 }
1301
1302 async fn list_schemas(&self) -> Result<ToolOutcome, ToolError> {
1303 let client = self.session.checkout().await?;
1304 let rows = client
1305 .query(
1306 r#"
1307 SELECT nspname AS schema_name
1308 FROM pg_namespace
1309 WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
1310 AND nspname NOT LIKE 'pg_%'
1311 ORDER BY nspname
1312 "#,
1313 &[],
1314 )
1315 .await?;
1316 let out: Vec<Value> = rows
1317 .iter()
1318 .filter(|r| {
1319 let name: String = r.get(0);
1320 self.session.filter().allows_schema(&name)
1321 })
1322 .map(|r| json!({ "schema_name": r.get::<_, String>(0) }))
1323 .collect();
1324 Ok(ToolOutcome::ok_json(json!(out)))
1325 }
1326
1327 async fn list_objects(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1328 let schema = args
1329 .get("schema")
1330 .and_then(|v| v.as_str())
1331 .unwrap_or("public");
1332 if !schema
1333 .chars()
1334 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
1335 {
1336 return Err(ToolError::InvalidArgs(
1337 "Invalid or missing schema name format".into(),
1338 ));
1339 }
1340 if !self.session.filter().allows_schema(schema) {
1341 return Ok(ToolOutcome::ok_json(json!([])));
1342 }
1343 let kind = args.get("kind").and_then(|v| v.as_str());
1344 let mut queries = Vec::new();
1345 let push_rel = |queries: &mut Vec<String>, relkinds: &[&str], label: &str| {
1346 let kinds = relkinds
1347 .iter()
1348 .map(|k| format!("'{k}'"))
1349 .collect::<Vec<_>>()
1350 .join(",");
1351 queries.push(format!(
1352 r#"
1353 SELECT n.nspname AS schema, c.relname AS name, '{label}' AS kind,
1354 d.description AS comment
1355 FROM pg_class c
1356 JOIN pg_namespace n ON n.oid = c.relnamespace
1357 LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0
1358 WHERE n.nspname = $1 AND c.relkind IN ({kinds})
1359 "#
1360 ));
1361 };
1362 if kind.is_none() || kind == Some("table") {
1363 push_rel(&mut queries, &["r", "f", "p"], "table");
1364 }
1365 if kind.is_none() || kind == Some("view") {
1366 push_rel(&mut queries, &["v"], "view");
1367 }
1368 if kind.is_none() || kind == Some("matview") {
1369 push_rel(&mut queries, &["m"], "matview");
1370 }
1371 if queries.is_empty() {
1372 return Ok(ToolOutcome::ok_json(json!([])));
1373 }
1374 let sql = queries.join("\nUNION ALL\n") + "\nORDER BY kind, name";
1375 let client = self.session.checkout().await?;
1376 let rows = client.query(&sql, &[&schema]).await?;
1377 let out: Vec<Value> = rows
1378 .iter()
1379 .filter(|r| {
1380 let s: String = r.get("schema");
1381 let name: String = r.get("name");
1382 self.session.filter().allows_table(&s, &name)
1383 })
1384 .map(|r| {
1385 json!({
1386 "schema": r.get::<_, String>("schema"),
1387 "name": r.get::<_, String>("name"),
1388 "kind": r.get::<_, String>("kind"),
1389 "comment": r.get::<_, Option<String>>("comment"),
1390 })
1391 })
1392 .collect();
1393 Ok(ToolOutcome::ok_json(json!(out)))
1394 }
1395
1396 async fn get_current_context(&self) -> Result<ToolOutcome, ToolError> {
1397 let (connection_id, database) = self.session.active_context().await;
1398 let conn = self
1399 .session
1400 .connections
1401 .iter()
1402 .find(|c| c.id == connection_id);
1403 Ok(ToolOutcome::ok_json(json!({
1404 "connectionId": connection_id,
1405 "connectionName": conn.map(|c| c.name.as_str()).unwrap_or("Unknown"),
1406 "database": database,
1407 "host": conn.and_then(|c| c.host.clone()),
1408 "port": conn.and_then(|c| c.port),
1409 "access_mode": match self.session.access_mode() {
1410 nexql_policy::AccessMode::Read => "read",
1411 nexql_policy::AccessMode::Write => "write",
1412 nexql_policy::AccessMode::Admin => "admin",
1413 },
1414 })))
1415 }
1416
1417 async fn switch_connection(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1418 let connection_id = args
1419 .get("connectionId")
1420 .and_then(|v| v.as_str())
1421 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
1422 let database = args
1423 .get("database")
1424 .and_then(|v| v.as_str())
1425 .map(str::to_owned);
1426 self.session.switch(connection_id, database).await?;
1427 self.get_current_context().await
1428 }
1429
1430 async fn run_select(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1431 let sql = args
1432 .get("sql")
1433 .and_then(|v| v.as_str())
1434 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
1435 match validate_readonly_sql(sql)? {
1436 SqlDecision::Allow => {}
1437 SqlDecision::Reject => {
1438 return Err(ToolError::Execution(
1439 "Security Error: Only read-only SELECT, WITH, or EXPLAIN statements are permitted."
1440 .into(),
1441 ));
1442 }
1443 }
1444 let trimmed = sql.trim().to_ascii_lowercase();
1445 if trimmed.starts_with("explain") {
1446 return self.run_select_internal(sql, None).await;
1447 }
1448 let max_rows = self.session.caps().max_rows;
1449 self.run_select_internal(sql, Some(max_rows)).await
1450 }
1451
1452 async fn explain_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1453 let sql = args
1454 .get("sql")
1455 .and_then(|v| v.as_str())
1456 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
1457 match validate_readonly_sql(sql)? {
1458 SqlDecision::Allow => {}
1459 SqlDecision::Reject => {
1460 return Err(ToolError::Execution(
1461 "Security Error: Only SELECT, WITH, or EXPLAIN statements can be analyzed."
1462 .into(),
1463 ));
1464 }
1465 }
1466 let clean = if sql.trim().to_ascii_lowercase().starts_with("explain") {
1467 sql.to_string()
1468 } else {
1469 format!("EXPLAIN {sql}")
1470 };
1471 if validate_readonly_sql(&clean)? == SqlDecision::Reject {
1473 return Err(ToolError::Execution(
1474 "Security Error: EXPLAIN target is not read-only.".into(),
1475 ));
1476 }
1477 self.run_select_internal(&clean, None).await
1478 }
1479
1480 async fn get_ddl(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1481 let ref_ = args
1482 .get("ref")
1483 .and_then(|v| v.as_str())
1484 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1485 let (schema, name) = parse_ref(ref_).map_err(ToolError::InvalidArgs)?;
1486 let kind = args.get("kind").and_then(|v| v.as_str()).unwrap_or("table");
1487 let reg = sql::regclass_literal(&schema, &name);
1488 let client = self.session.checkout().await?;
1489
1490 match kind {
1491 "view" | "matview" => {
1492 let sql = format!("SELECT pg_get_viewdef({reg}, true) AS definition");
1493 let rows = client.query(&sql, &[]).await?;
1494 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1495 }
1496 "function" => {
1497 let sql = format!(
1498 r#"SELECT p.proname AS name, pg_get_functiondef(p.oid) AS definition
1499 FROM pg_proc p
1500 JOIN pg_namespace n ON n.oid = p.pronamespace
1501 WHERE n.nspname = '{schema}' AND p.proname = '{name}'"#
1502 );
1503 let rows = client.query(&sql, &[]).await?;
1504 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1505 }
1506 "index" => {
1507 let sql = format!("SELECT pg_get_indexdef({reg}) AS definition");
1508 let rows = client.query(&sql, &[]).await?;
1509 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1510 }
1511 "table" => {
1512 let columns = client
1513 .query(&sql::column_details(&schema, &name), &[])
1514 .await?;
1515 let constraints = client
1516 .query(
1517 &format!(
1518 r#"SELECT conname AS name, pg_get_constraintdef(oid) AS definition
1519 FROM pg_constraint WHERE conrelid = {reg} ORDER BY conname"#
1520 ),
1521 &[],
1522 )
1523 .await?;
1524 let indexes = client
1525 .query(
1526 &format!(
1527 r#"SELECT indexname AS name, indexdef AS definition
1528 FROM pg_indexes
1529 WHERE schemaname = '{schema}' AND tablename = '{name}'
1530 ORDER BY indexname"#
1531 ),
1532 &[],
1533 )
1534 .await?;
1535 Ok(ToolOutcome::ok_json(json!({
1536 "table": format!("{schema}.{name}"),
1537 "columns": rows_to_json(&columns),
1538 "constraints": rows_to_json(&constraints),
1539 "indexes": rows_to_json(&indexes),
1540 })))
1541 }
1542 other => Err(ToolError::InvalidArgs(format!(
1543 "Unsupported DDL kind \"{other}\". Use table, view, matview, function, or index."
1544 ))),
1545 }
1546 }
1547
1548 async fn table_stats(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1549 let ref_ = args
1550 .get("ref")
1551 .and_then(|v| v.as_str())
1552 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1553 let (schema, name) = parse_ref(ref_).map_err(ToolError::InvalidArgs)?;
1554 let client = self.session.checkout().await?;
1555 let stats = client.query(&sql::table_stats(&schema, &name), &[]).await?;
1556 let activity = client
1557 .query(&sql::table_activity(&schema, &name), &[])
1558 .await?;
1559 let columns = client
1560 .query(&sql::column_stats(&schema, &name), &[])
1561 .await?;
1562 let size = rows_to_json(&stats)
1563 .as_array()
1564 .and_then(|a| a.first())
1565 .cloned()
1566 .unwrap_or(Value::Null);
1567 let activity = rows_to_json(&activity)
1568 .as_array()
1569 .and_then(|a| a.first())
1570 .cloned()
1571 .unwrap_or(Value::Null);
1572 Ok(ToolOutcome::ok_json(json!({
1573 "size": size,
1574 "activity": activity,
1575 "columns": rows_to_json(&columns),
1576 })))
1577 }
1578
1579 async fn index_usage(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1580 let ref_ = args
1581 .get("ref")
1582 .and_then(|v| v.as_str())
1583 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1584 let (schema, name) = parse_ref(ref_).map_err(ToolError::InvalidArgs)?;
1585 let client = self.session.checkout().await?;
1586 let rows = client.query(&sql::index_usage(&schema, &name), &[]).await?;
1587 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1588 }
1589
1590 async fn list_running_queries(&self) -> Result<ToolOutcome, ToolError> {
1591 let client = self.session.checkout().await?;
1592 let rows = client.query(sql::running_queries(), &[]).await?;
1593 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1594 }
1595
1596 async fn find_blocking_locks(&self) -> Result<ToolOutcome, ToolError> {
1597 let client = self.session.checkout().await?;
1598 let rows = client.query(sql::blocking_locks(), &[]).await?;
1599 let values = rows_to_json(&rows);
1600 if values.as_array().map(|a| a.is_empty()).unwrap_or(true) {
1601 return Ok(ToolOutcome::ok_json(json!({
1602 "message": "No blocking locks found.",
1603 "locks": [],
1604 })));
1605 }
1606 Ok(ToolOutcome::ok_json(values))
1607 }
1608
1609 async fn slow_queries(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1610 let limit = args
1611 .get("limit")
1612 .and_then(|v| v.as_u64())
1613 .map(|n| n as u32)
1614 .unwrap_or(SLOW_QUERIES_DEFAULT);
1615 let client = self.session.checkout().await?;
1616 match client.query(&sql::slow_queries(limit), &[]).await {
1617 Ok(rows) => Ok(ToolOutcome::ok_json(rows_to_json(&rows))),
1618 Err(e) => {
1619 if let Some(message) = sql::map_stat_statements_error(&e) {
1620 Ok(ToolOutcome::ok_json(json!({
1621 "error": message,
1622 "hint": message,
1623 })))
1624 } else {
1625 Err(ToolError::Postgres(e))
1626 }
1627 }
1628 }
1629 }
1630
1631 async fn db_health_check(&self) -> Result<ToolOutcome, ToolError> {
1632 let client = self.session.checkout().await?;
1633 let sections: &[(&str, &str)] = &[
1634 ("overview", sql::database_stats()),
1635 ("cache", sql::cache_hit_ratio()),
1636 ("dead_tuples", sql::database_maintenance_stats()),
1637 ("connection_states", sql::connection_states()),
1638 ("blocking_locks", sql::blocking_locks()),
1639 ];
1640 let mut report = serde_json::Map::new();
1641 for (key, q) in sections {
1642 match client.query(*q, &[]).await {
1643 Ok(rows) => {
1644 report.insert((*key).into(), rows_to_json(&rows));
1645 }
1646 Err(e) => {
1647 report.insert((*key).into(), json!({ "error": e.to_string() }));
1648 }
1649 }
1650 }
1651 let lock_count = report
1652 .get("blocking_locks")
1653 .and_then(|v| v.as_array())
1654 .map(|a| a.len() as u64);
1655 report.insert("blocking_lock_count".into(), json!(lock_count));
1656 Ok(ToolOutcome::ok_json(Value::Object(report)))
1657 }
1658
1659 async fn explain_analyze(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1660 let sql = args
1661 .get("sql")
1662 .and_then(|v| v.as_str())
1663 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
1664 require_select_or_with(sql)?;
1665 let explain = build_explain_sql(sql, true);
1666 self.run_explain_in_transaction(&explain).await
1667 }
1668
1669 async fn analyze_query_plan(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1670 let sql = args
1671 .get("sql")
1672 .and_then(|v| v.as_str())
1673 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
1674 require_select_or_with(sql)?;
1675 let analyze = args
1676 .get("analyze")
1677 .and_then(|v| v.as_bool())
1678 .unwrap_or(false);
1679 let explain = build_explain_sql(sql, analyze);
1680 let outcome = self.run_explain_in_transaction(&explain).await?;
1681 let rows = outcome.structured.unwrap_or(Value::Null);
1682 let row_array = rows
1683 .get("rows")
1684 .and_then(|v| v.as_array())
1685 .or_else(|| rows.as_array());
1686 let plan = row_array
1687 .and_then(|a| a.first())
1688 .and_then(|r| r.get("QUERY PLAN"))
1689 .cloned()
1690 .unwrap_or(Value::Null);
1691 let metrics = extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
1692 let recommendations = metrics
1693 .as_ref()
1694 .and_then(|m| m.get("recommendations"))
1695 .cloned()
1696 .unwrap_or_else(|| json!([]));
1697 Ok(ToolOutcome::ok_json(json!({
1698 "metrics": metrics,
1699 "recommendations": recommendations,
1700 "plan": plan,
1701 })))
1702 }
1703
1704 async fn run_explain_in_transaction(
1706 &self,
1707 explain_sql: &str,
1708 ) -> Result<ToolOutcome, ToolError> {
1709 let client = self.session.checkout().await?;
1710 client
1711 .batch_execute("SET statement_timeout = '30s'")
1712 .await?;
1713 client.batch_execute("BEGIN").await?;
1714 let result = async {
1715 client.batch_execute("SET TRANSACTION READ ONLY").await?;
1716 let rows = client.query(explain_sql, &[]).await?;
1717 Ok::<_, ToolError>(rows_to_json(&rows))
1718 }
1719 .await;
1720 let _ = client.batch_execute("ROLLBACK").await;
1722 match result {
1723 Ok(values) => Ok(ToolOutcome::ok_json(values)),
1724 Err(e) => Err(e),
1725 }
1726 }
1727
1728 async fn get_index_status(&self) -> Result<ToolOutcome, ToolError> {
1729 let (store, connection_id, database) = self.index_service().await?;
1730 let base = store.base_dir(&connection_id, &database);
1731 let Some(manifest) = store.read_manifest(&base)? else {
1732 return Err(ToolError::Execution(format!(
1733 "No schema index for database \"{database}\" — run `nexql-mcp index build`."
1734 )));
1735 };
1736
1737 let mut live_fingerprint: Option<String> = None;
1738 let mut drift: Option<bool> = None;
1739 if let Ok(client) = self.session.checkout().await {
1740 let db = PgCatalogDb::new(&client);
1741 if let Ok(fp) = db.schema_fingerprint().await {
1742 drift = Some(fp != manifest.schema_fingerprint);
1743 live_fingerprint = Some(fp);
1744 }
1745 }
1746
1747 Ok(ToolOutcome::ok_json(json!({
1748 "connectionId": manifest.connection_id,
1749 "database": manifest.database,
1750 "indexedAt": manifest.indexed_at,
1751 "fingerprint": manifest.schema_fingerprint,
1752 "liveFingerprint": live_fingerprint,
1753 "drift": drift,
1754 "pgVersion": manifest.pg_version,
1755 "counts": {
1756 "tables": manifest.counts.tables,
1757 "views": manifest.counts.views,
1758 "functions": manifest.counts.functions,
1759 "enums": manifest.counts.enums,
1760 },
1761 "buildMs": manifest.stats.build_ms,
1762 "warnings": manifest.stats.warnings,
1763 })))
1764 }
1765
1766 async fn list_extensions(&self) -> Result<ToolOutcome, ToolError> {
1767 let client = self.session.checkout().await?;
1768 let rows = client.query(sql::list_extensions(), &[]).await?;
1769 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1770 }
1771
1772 async fn server_settings(&self) -> Result<ToolOutcome, ToolError> {
1773 let client = self.session.checkout().await?;
1774 let rows = client.query(sql::server_settings(), &[]).await?;
1775 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
1776 }
1777
1778 async fn suggest_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1779 let limit = args
1780 .get("limit")
1781 .and_then(|v| v.as_u64())
1782 .map(|n| n as u32)
1783 .unwrap_or(REPORT_LIMIT_DEFAULT);
1784 let client = self.session.checkout().await?;
1785
1786 let high_seq = client.query(&sql::high_seq_scan_tables(limit), &[]).await?;
1787 let unindexed_fks = client.query(&sql::unindexed_fk_columns(limit), &[]).await?;
1788
1789 let mut pg_stat_available = false;
1790 let mut slow_queries = Value::Null;
1791 let mut pg_stat_note: Option<String> = None;
1792 match client.query(&sql::slow_queries(limit.min(10)), &[]).await {
1793 Ok(rows) => {
1794 pg_stat_available = true;
1795 slow_queries = rows_to_json(&rows);
1796 }
1797 Err(e) => {
1798 if let Some(message) = sql::map_stat_statements_error(&e) {
1799 pg_stat_note = Some(message);
1800 } else {
1801 return Err(ToolError::Postgres(e));
1802 }
1803 }
1804 }
1805
1806 let mut plan_heuristics = Value::Null;
1807 if let Some(sql_text) = args.get("sql").and_then(|v| v.as_str()) {
1808 require_select_or_with(sql_text)?;
1809 let explain = build_explain_sql(sql_text, false);
1810 let outcome = self.run_explain_in_transaction(&explain).await?;
1811 let rows = outcome.structured.unwrap_or(Value::Null);
1812 let plan = rows
1813 .as_array()
1814 .and_then(|a| a.first())
1815 .and_then(|r| r.get("QUERY PLAN"))
1816 .cloned()
1817 .unwrap_or(Value::Null);
1818 let metrics = extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
1819 plan_heuristics = json!({
1820 "metrics": metrics,
1821 "hint": "Use analyze_query_plan with analyze=true for actual timings before creating indexes.",
1822 });
1823 }
1824
1825 let high_seq_json = rows_to_json(&high_seq);
1826 let unindexed_json = rows_to_json(&unindexed_fks);
1827 let has_candidates = high_seq_json
1828 .as_array()
1829 .map(|a| !a.is_empty())
1830 .unwrap_or(false)
1831 || unindexed_json
1832 .as_array()
1833 .map(|a| !a.is_empty())
1834 .unwrap_or(false)
1835 || plan_heuristics != Value::Null;
1836
1837 if !has_candidates && !pg_stat_available {
1838 return Ok(ToolOutcome::ok_json(json!({
1839 "suggestions": [],
1840 "message": "No index suggestions yet. Either table stats show healthy index use, or there is not enough scan history. Enable pg_stat_statements and/or pass a sql argument for EXPLAIN plan heuristics.",
1841 "hint": pg_stat_note,
1842 })));
1843 }
1844
1845 if !has_candidates {
1846 return Ok(ToolOutcome::ok_json(json!({
1847 "high_seq_scan_tables": high_seq_json,
1848 "unindexed_fk_columns": unindexed_json,
1849 "slow_queries": slow_queries,
1850 "plan_heuristics": plan_heuristics,
1851 "message": "No strong index candidates from sequential-scan or unindexed-FK heuristics. Review slow_queries / pass sql for plan-level advice.",
1852 "hint": "CREATE INDEX CONCURRENTLY after validating with EXPLAIN (ANALYZE, BUFFERS).",
1853 })));
1854 }
1855
1856 Ok(ToolOutcome::ok_json(json!({
1857 "high_seq_scan_tables": high_seq_json,
1858 "unindexed_fk_columns": unindexed_json,
1859 "slow_queries": slow_queries,
1860 "plan_heuristics": plan_heuristics,
1861 "pg_stat_statements": pg_stat_available,
1862 "hint": pg_stat_note.unwrap_or_else(|| {
1863 "Validate candidates with analyze_query_plan / EXPLAIN before CREATE INDEX CONCURRENTLY.".into()
1864 }),
1865 })))
1866 }
1867
1868 async fn find_unused_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1869 let limit = args
1870 .get("limit")
1871 .and_then(|v| v.as_u64())
1872 .map(|n| n as u32)
1873 .unwrap_or(REPORT_LIMIT_DEFAULT);
1874 let client = self.session.checkout().await?;
1875 let rows = client.query(&sql::find_unused_indexes(limit), &[]).await?;
1876 let indexes = rows_to_json(&rows);
1877 if indexes.as_array().map(|a| a.is_empty()).unwrap_or(true) {
1878 return Ok(ToolOutcome::ok_json(json!({
1879 "indexes": [],
1880 "message": "No unused non-constraint indexes found (idx_scan = 0). Note: pg_stat_reset / server restart clears scan counts — treat never-scanned indexes cautiously on fresh stats.",
1881 })));
1882 }
1883 Ok(ToolOutcome::ok_json(json!({
1884 "indexes": indexes,
1885 "hint": "Prefer DROP INDEX CONCURRENTLY after confirming the workload (and that stats are mature).",
1886 })))
1887 }
1888
1889 async fn bloat_report(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1890 let limit = args
1891 .get("limit")
1892 .and_then(|v| v.as_u64())
1893 .map(|n| n as u32)
1894 .unwrap_or(REPORT_LIMIT_DEFAULT);
1895 let client = self.session.checkout().await?;
1896 let rows = client.query(&sql::bloat_report(limit), &[]).await?;
1897 let tables = rows_to_json(&rows);
1898 if tables.as_array().map(|a| a.is_empty()).unwrap_or(true) {
1899 return Ok(ToolOutcome::ok_json(json!({
1900 "tables": [],
1901 "method": "dead_tuple_ratio",
1902 "message": "No tables with significant dead-tuple pressure (>1000 dead tuples). This is a simplified estimate from pg_stat_user_tables, not physical page bloat.",
1903 })));
1904 }
1905 Ok(ToolOutcome::ok_json(json!({
1906 "tables": tables,
1907 "method": "dead_tuple_ratio",
1908 "note": "Approximate bloat via n_dead_tup / (n_live_tup + n_dead_tup). Not a physical page-bloat estimate (pgstattuple / check_postgres). Consider VACUUM / VACUUM FULL only after confirming impact.",
1909 "hint": "VACUUM ANALYZE on high bloat_pct tables; investigate autovacuum settings if last_autovacuum is stale.",
1910 })))
1911 }
1912
1913 async fn find_missing_fks(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1914 let limit = args
1915 .get("limit")
1916 .and_then(|v| v.as_u64())
1917 .map(|n| n as u32)
1918 .unwrap_or(REPORT_LIMIT_DEFAULT);
1919 let capped = limit.clamp(1, sql::REPORT_LIMIT_MAX) as usize;
1920
1921 if let Ok((store, connection_id, database)) = self.index_service().await {
1923 let base = store.base_dir(&connection_id, &database);
1924 if let Ok(Some(manifest)) = store.read_manifest(&base) {
1925 if let Ok(Some(graph)) = store.read_join_graph(&base, &manifest) {
1926 let candidates: Vec<Value> = graph
1927 .edges
1928 .into_iter()
1929 .filter(|e| e.inferred == Some(true) && e.disabled != Some(true))
1930 .take(capped)
1931 .map(|e| {
1932 let cols: Vec<Value> = e
1933 .cols
1934 .iter()
1935 .map(|(a, b)| json!({ "from": a, "to": b }))
1936 .collect();
1937 json!({
1938 "from_table": e.from,
1939 "to_table": e.to,
1940 "via": e.via,
1941 "columns": cols,
1942 "detection": "join_graph_inferred",
1943 })
1944 })
1945 .collect();
1946 if !candidates.is_empty() {
1947 return Ok(ToolOutcome::ok_json(json!({
1948 "candidates": candidates,
1949 "source": "join_graph",
1950 "hint": "These edges were inferred by naming convention and have no declared FK. Review before ALTER TABLE … ADD FOREIGN KEY.",
1951 })));
1952 }
1953 }
1954 }
1955 }
1956
1957 let client = self.session.checkout().await?;
1958 let rows = client
1959 .query(&sql::find_missing_fks_catalog(limit), &[])
1960 .await?;
1961 let candidates = rows_to_json(&rows);
1962 if candidates.as_array().map(|a| a.is_empty()).unwrap_or(true) {
1963 return Ok(ToolOutcome::ok_json(json!({
1964 "candidates": [],
1965 "source": "catalog",
1966 "message": "No missing FK candidates found via join-graph inferred edges or *_id naming against single-column PKs.",
1967 })));
1968 }
1969 Ok(ToolOutcome::ok_json(json!({
1970 "candidates": candidates,
1971 "source": "catalog",
1972 "hint": "Naming-inferred only — verify referential integrity and nullability before adding constraints. Run `nexql-mcp index build` for join-graph inferred edges.",
1973 })))
1974 }
1975
1976 async fn list_roles(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1977 let client = self.session.checkout().await?;
1978 let role = args
1979 .get("role")
1980 .and_then(|v| v.as_str())
1981 .map(str::trim)
1982 .filter(|s| !s.is_empty());
1983
1984 let Some(role_name) = role else {
1985 let rows = client.query(sql::list_roles(), &[]).await?;
1986 return Ok(ToolOutcome::ok_json(rows_to_json(&rows)));
1987 };
1988
1989 let details = client.query(sql::role_details(), &[&role_name]).await?;
1990 if details.is_empty() {
1991 return Err(ToolError::Execution(format!(
1992 "Role \"{role_name}\" not found"
1993 )));
1994 }
1995 let member_of = client.query(sql::role_member_of(), &[&role_name]).await?;
1996 let has_members = client.query(sql::role_has_members(), &[&role_name]).await?;
1997 let privileges = client
1998 .query(sql::role_table_privileges(), &[&role_name])
1999 .await?;
2000
2001 Ok(ToolOutcome::ok_json(json!({
2002 "role": rows_to_json(&details).as_array().and_then(|a| a.first().cloned()).unwrap_or(Value::Null),
2003 "member_of": rows_to_json(&member_of),
2004 "has_members": rows_to_json(&has_members),
2005 "table_privileges": rows_to_json(&privileges),
2006 })))
2007 }
2008
2009 async fn export_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2010 let sql = args
2011 .get("sql")
2012 .and_then(|v| v.as_str())
2013 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2014 require_select_or_with(sql)?;
2015
2016 let format = args
2017 .get("format")
2018 .and_then(|v| v.as_str())
2019 .map(|s| {
2020 ExportFormat::parse(s).ok_or_else(|| {
2021 ToolError::InvalidArgs(format!(
2022 "Unsupported format \"{s}\". Use csv, json, or sqlinsert."
2023 ))
2024 })
2025 })
2026 .transpose()?
2027 .unwrap_or(ExportFormat::Csv);
2028
2029 let table_target = match args.get("table").and_then(|v| v.as_str()) {
2030 Some(t) if !t.trim().is_empty() => Some(parse_ref(t).map_err(ToolError::InvalidArgs)?),
2031 _ => None,
2032 };
2033
2034 if format == ExportFormat::SqlInsert && table_target.is_none() {
2035 return Err(ToolError::InvalidArgs(
2036 "table (schema.name) is required when format=sqlinsert".into(),
2037 ));
2038 }
2039
2040 let max_rows = self.session.caps().max_rows;
2041 let outcome = self.run_select_internal(sql, Some(max_rows)).await?;
2042 if outcome.is_error {
2043 return Ok(outcome);
2044 }
2045
2046 let structured = outcome.structured.unwrap_or(Value::Null);
2047 let rows_val = structured
2048 .get("rows")
2049 .cloned()
2050 .or_else(|| structured.get("data").and_then(|d| d.get("rows").cloned()))
2051 .unwrap_or(Value::Array(vec![]));
2052 let rows = rows_val.as_array().cloned().unwrap_or_default();
2053 let columns = columns_from_rows(&rows);
2054 let truncated = structured
2055 .get("truncated")
2056 .and_then(|v| v.as_bool())
2057 .unwrap_or(false);
2058
2059 let payload = match format {
2060 ExportFormat::Json => json!({
2061 "format": format.as_str(),
2062 "rowCount": rows.len(),
2063 "truncated": truncated,
2064 "columns": columns,
2065 "rows": rows,
2066 }),
2067 ExportFormat::Csv => {
2068 let content = rows_to_csv(&rows, &columns);
2069 let caps = self.session.caps();
2070 let (char_trunc, content) = caps.truncate_chars(&content);
2071 json!({
2072 "format": format.as_str(),
2073 "rowCount": rows.len(),
2074 "truncated": truncated || char_trunc,
2075 "columns": columns,
2076 "content": content,
2077 })
2078 }
2079 ExportFormat::SqlInsert => {
2080 let (schema, table) = table_target.expect("checked above");
2081 let content = rows_to_sql_insert(&rows, &columns, &schema, &table);
2082 let caps = self.session.caps();
2083 let (char_trunc, content) = caps.truncate_chars(&content);
2084 json!({
2085 "format": format.as_str(),
2086 "rowCount": rows.len(),
2087 "truncated": truncated || char_trunc,
2088 "table": format!("{schema}.{table}"),
2089 "columns": columns,
2090 "content": content,
2091 })
2092 }
2093 };
2094
2095 Ok(ToolOutcome::ok_json(payload))
2096 }
2097
2098 async fn db_dashboard(&self) -> Result<ToolOutcome, ToolError> {
2099 let client = self.session.checkout().await?;
2100 let sections: &[(&str, &str)] = &[
2101 ("db_info", sql::dashboard_db_info()),
2102 ("connection_states", sql::connection_states()),
2103 ("top_tables", sql::dashboard_top_tables()),
2104 ("object_counts", sql::dashboard_object_counts()),
2105 ("active_queries", sql::dashboard_active_queries()),
2106 ("blocking_locks", sql::blocking_locks()),
2107 ("max_connections", sql::dashboard_max_connections()),
2108 ("extension_count", sql::dashboard_extension_count()),
2109 ("cache", sql::cache_hit_ratio()),
2110 ];
2111 let mut report = serde_json::Map::new();
2112 for (key, q) in sections {
2113 match client.query(*q, &[]).await {
2114 Ok(rows) => {
2115 report.insert((*key).into(), rows_to_json(&rows));
2116 }
2117 Err(e) => {
2118 report.insert((*key).into(), json!({ "error": e.to_string() }));
2119 }
2120 }
2121 }
2122
2123 for key in ["db_info", "object_counts", "extension_count", "cache"] {
2125 if let Some(Value::Array(arr)) = report.get(key).cloned() {
2126 if arr.len() == 1 {
2127 report.insert(key.into(), arr.into_iter().next().unwrap());
2128 }
2129 }
2130 }
2131 if let Some(Value::Array(arr)) = report.get("max_connections").cloned() {
2132 if let Some(row) = arr.first() {
2133 report.insert(
2134 "max_connections".into(),
2135 row.get("max_connections")
2136 .cloned()
2137 .unwrap_or_else(|| row.clone()),
2138 );
2139 }
2140 }
2141
2142 Ok(ToolOutcome::ok_json(Value::Object(report)))
2143 }
2144
2145 async fn deep_plan_analysis(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2146 let sql = args
2147 .get("sql")
2148 .and_then(|v| v.as_str())
2149 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2150 require_select_or_with(sql)?;
2151 let analyze = args
2152 .get("analyze")
2153 .and_then(|v| v.as_bool())
2154 .unwrap_or(true);
2155 let explain = build_explain_sql(sql, analyze);
2156 let outcome = self.run_explain_in_transaction(&explain).await?;
2157 let rows = outcome.structured.unwrap_or(Value::Null);
2158 let row_array = rows
2159 .get("rows")
2160 .and_then(|v| v.as_array())
2161 .or_else(|| rows.as_array());
2162 let plan = row_array
2163 .and_then(|a| a.first())
2164 .and_then(|r| r.get("QUERY PLAN"))
2165 .cloned()
2166 .unwrap_or(Value::Null);
2167 let deep = analyze_deep_plan(&plan, sql)
2168 .or_else(|| analyze_deep_plan(&rows, sql))
2169 .ok_or_else(|| {
2170 ToolError::Execution("Could not parse EXPLAIN JSON plan for deep analysis".into())
2171 })?;
2172 let metrics = extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
2173 Ok(ToolOutcome::ok_json(json!({
2174 "deep": deep,
2175 "metrics": metrics,
2176 "plan": plan,
2177 "analyzed": analyze,
2178 })))
2179 }
2180
2181 async fn schema_diff(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2182 let source_schema = args
2183 .get("sourceSchema")
2184 .and_then(|v| v.as_str())
2185 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
2186 let target_schema = args
2187 .get("targetSchema")
2188 .and_then(|v| v.as_str())
2189 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
2190 crate::schema_diff::require_safe_schema(source_schema)?;
2191 crate::schema_diff::require_safe_schema(target_schema)?;
2192
2193 let client = self.session.checkout().await?;
2194 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
2195 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
2196 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
2197 let changed = diffs
2198 .iter()
2199 .filter(|d| d.status != crate::schema_diff::DiffStatus::Unchanged)
2200 .count();
2201 Ok(ToolOutcome::ok_json(json!({
2202 "sourceSchema": source_schema,
2203 "targetSchema": target_schema,
2204 "tableCount": diffs.len(),
2205 "changedCount": changed,
2206 "diffs": crate::schema_diff::diffs_to_json(&diffs),
2207 })))
2208 }
2209
2210 async fn generate_migration(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2211 let source_schema = args
2212 .get("sourceSchema")
2213 .and_then(|v| v.as_str())
2214 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
2215 let target_schema = args
2216 .get("targetSchema")
2217 .and_then(|v| v.as_str())
2218 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
2219 crate::schema_diff::require_safe_schema(source_schema)?;
2220 crate::schema_diff::require_safe_schema(target_schema)?;
2221
2222 let client = self.session.checkout().await?;
2223 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
2224 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
2225 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
2226 let statements =
2227 crate::schema_diff::build_migration_statements(source_schema, target_schema, &diffs);
2228 let sql = if statements.is_empty() {
2229 format!("-- No differences between {source_schema} and {target_schema}")
2230 } else {
2231 statements.join("\n\n")
2232 };
2233 Ok(ToolOutcome::ok_json(json!({
2234 "sourceSchema": source_schema,
2235 "targetSchema": target_schema,
2236 "statementCount": statements.len(),
2237 "sql": sql,
2238 "hint": "Read-only: review and run via execute_sql / apply_ddl only with --access-mode write|admin. Destructive drops are commented out.",
2239 })))
2240 }
2241
2242 async fn run_select_internal(
2243 &self,
2244 sql: &str,
2245 max_rows: Option<u32>,
2246 ) -> Result<ToolOutcome, ToolError> {
2247 let client = self.session.checkout().await?;
2248 let Some(max_rows) = max_rows else {
2249 let rows = client.query(sql, &[]).await?;
2250 let values = rows_to_json(&rows);
2251 let payload = ensure_structured_object(values);
2253 let text = serde_json::to_string_pretty(&payload)
2254 .map_err(|e| ToolError::Execution(e.to_string()))?;
2255 let caps = self.session.caps();
2256 let (trunc, text) = caps.truncate_chars(&text);
2257 let structured = if trunc {
2258 json!({ "truncated_chars": true, "data": payload })
2259 } else {
2260 payload
2261 };
2262 return Ok(ToolOutcome {
2263 text: text.to_string(),
2264 structured: Some(structured),
2265 is_error: false,
2266 });
2267 };
2268
2269 let cleaned = sql.trim().trim_end_matches(';').trim();
2270 let wrapped = format!(
2271 "SELECT * FROM ({cleaned}) AS nexql_limited LIMIT {}",
2272 max_rows + 1
2273 );
2274 let rows = client.query(&wrapped, &[]).await.map_err(|e| {
2275 ToolError::Execution(format!(
2276 "Failed to execute row-limited query (refusing unbounded fallback): {e}"
2277 ))
2278 })?;
2279 let truncated = rows.len() as u32 > max_rows;
2280 let keep = if truncated {
2281 &rows[..max_rows as usize]
2282 } else {
2283 &rows[..]
2284 };
2285 let values = rows_to_json(keep);
2286 let mut payload = ensure_structured_object(values);
2288 if truncated {
2289 if let Some(obj) = payload.as_object_mut() {
2290 obj.insert("truncated".into(), json!(true));
2291 obj.insert("maxRows".into(), json!(max_rows));
2292 }
2293 }
2294 let text = serde_json::to_string_pretty(&payload)
2295 .map_err(|e| ToolError::Execution(e.to_string()))?;
2296 let caps = self.session.caps();
2297 let (char_trunc, text) = caps.truncate_chars(&text);
2298 let structured = if char_trunc {
2299 json!({ "truncated_chars": true, "data": payload })
2300 } else {
2301 payload
2302 };
2303 Ok(ToolOutcome {
2304 text: text.to_string(),
2305 structured: Some(structured),
2306 is_error: false,
2307 })
2308 }
2309}
2310
2311fn normalize_for_match(s: &str) -> String {
2313 let mut out = String::new();
2314 let mut last_was_sep = true;
2315 for ch in s.to_lowercase().chars() {
2316 if ch.is_ascii_alphanumeric() {
2317 out.push(ch);
2318 last_was_sep = false;
2319 } else if !last_was_sep {
2320 out.push(' ');
2321 last_was_sep = true;
2322 }
2323 }
2324 out.trim().to_string()
2325}
2326
2327fn fuzzy_score(hint: &str, candidate: &str) -> f64 {
2329 let h = normalize_for_match(hint);
2330 let c = normalize_for_match(candidate);
2331 if h.is_empty() || c.is_empty() {
2332 return 0.0;
2333 }
2334 if h == c {
2335 return 100.0;
2336 }
2337 if c.contains(&h) || h.contains(&c) {
2338 return 75.0;
2339 }
2340 let h_tokens: std::collections::HashSet<&str> =
2341 h.split(' ').filter(|s| !s.is_empty()).collect();
2342 let c_tokens: std::collections::HashSet<&str> =
2343 c.split(' ').filter(|s| !s.is_empty()).collect();
2344 let overlap = h_tokens.intersection(&c_tokens).count();
2345 if overlap == 0 {
2346 return 0.0;
2347 }
2348 (overlap as f64 / h_tokens.len().max(c_tokens.len()) as f64) * 60.0
2349}
2350
2351fn policy_to_query_filter(filter: &PolicyFilter) -> QueryPolicyFilter {
2352 QueryPolicyFilter {
2353 allow_schemas: filter.allow_schemas.clone(),
2354 deny_schemas: filter.deny_schemas.clone(),
2355 deny_tables: filter.deny_tables.clone(),
2356 pii_columns: filter.pii_columns.clone(),
2357 }
2358}
2359
2360fn require_select_or_with(sql: &str) -> Result<(), ToolError> {
2361 match validate_readonly_sql(sql)? {
2362 SqlDecision::Allow => {}
2363 SqlDecision::Reject => {
2364 return Err(ToolError::Execution(
2365 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
2366 ));
2367 }
2368 }
2369 let trimmed = sql.trim().to_ascii_lowercase();
2370 if !(trimmed.starts_with("select") || trimmed.starts_with("with")) {
2371 return Err(ToolError::Execution(
2372 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
2373 ));
2374 }
2375 Ok(())
2376}
2377
2378fn rows_to_json(rows: &[tokio_postgres::Row]) -> Value {
2379 let arr: Vec<Value> = rows
2380 .iter()
2381 .map(|row| {
2382 let mut map = serde_json::Map::new();
2383 for (i, col) in row.columns().iter().enumerate() {
2384 map.insert(col.name().to_string(), cell_to_json(row, i));
2385 }
2386 Value::Object(map)
2387 })
2388 .collect();
2389 Value::Array(arr)
2390}
2391
2392enum SqlNullness {
2394 Null,
2395 Value,
2396}
2397
2398impl<'a> FromSql<'a> for SqlNullness {
2399 fn from_sql(_: &Type, _: &'a [u8]) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
2400 Ok(SqlNullness::Value)
2401 }
2402
2403 fn from_sql_null(_: &Type) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
2404 Ok(SqlNullness::Null)
2405 }
2406
2407 fn accepts(_: &Type) -> bool {
2408 true
2409 }
2410}
2411
2412fn try_cell<T, F>(row: &tokio_postgres::Row, idx: usize, map: F) -> Option<Value>
2413where
2414 T: for<'a> FromSql<'a>,
2415 F: FnOnce(T) -> Value,
2416{
2417 match row.try_get::<_, Option<T>>(idx) {
2418 Ok(Some(v)) => Some(map(v)),
2419 Ok(None) => Some(Value::Null),
2420 Err(_) => None,
2421 }
2422}
2423
2424fn cell_to_json(row: &tokio_postgres::Row, idx: usize) -> Value {
2425 let col_type = row.columns()[idx].type_();
2426 if matches!(row.try_get::<_, SqlNullness>(idx), Ok(SqlNullness::Null)) {
2427 return Value::Null;
2428 }
2429
2430 if let Kind::Array(elem) = col_type.kind() {
2431 return array_cell_to_json(row, idx, elem);
2432 }
2433
2434 if let Some(v) = match *col_type {
2435 Type::BOOL => try_cell::<bool, _>(row, idx, |b| json!(b)),
2436 Type::INT2 => try_cell::<i16, _>(row, idx, |n| json!(n)),
2437 Type::INT4 | Type::OID => try_cell::<i32, _>(row, idx, |n| json!(n)),
2438 Type::INT8 => try_cell::<i64, _>(row, idx, |n| json!(n)),
2439 Type::FLOAT4 => try_cell::<f32, _>(row, idx, |n| json!(n)),
2440 Type::FLOAT8 => try_cell::<f64, _>(row, idx, |n| json!(n)),
2441 Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::NAME => {
2442 try_cell::<String, _>(row, idx, Value::String)
2443 }
2444 Type::TIMESTAMP => try_cell::<NaiveDateTime, _>(row, idx, |t| {
2445 json!(t.format("%Y-%m-%dT%H:%M:%S%.f").to_string())
2446 }),
2447 Type::TIMESTAMPTZ => {
2448 try_cell::<DateTime<FixedOffset>, _>(row, idx, |t| json!(t.to_rfc3339()))
2449 }
2450 Type::DATE => {
2451 try_cell::<NaiveDate, _>(row, idx, |d| json!(d.format("%Y-%m-%d").to_string()))
2452 }
2453 Type::TIME => {
2454 try_cell::<NaiveTime, _>(row, idx, |t| json!(t.format("%H:%M:%S%.f").to_string()))
2455 }
2456 Type::UUID => try_cell::<Uuid, _>(row, idx, |u| json!(u.to_string())),
2457 Type::JSON | Type::JSONB => try_cell::<Value, _>(row, idx, |j| j),
2458 Type::NUMERIC => try_cell::<Decimal, _>(row, idx, |d| json!(d.to_string())),
2459 Type::MONEY => try_cell::<i64, _>(row, idx, |v| json!(money_to_string(v))),
2460 Type::BYTEA => try_cell::<Vec<u8>, _>(row, idx, |b| json!(BASE64.encode(b))),
2461 _ => None,
2462 } {
2463 return v;
2464 }
2465
2466 cell_to_json_untyped(row, idx, col_type)
2467}
2468
2469fn array_cell_to_json(row: &tokio_postgres::Row, idx: usize, elem: &Type) -> Value {
2470 let try_array = |result: Result<Option<Vec<Value>>, tokio_postgres::Error>| -> Option<Value> {
2471 match result {
2472 Ok(Some(items)) => Some(Value::Array(items)),
2473 Ok(None) => Some(Value::Null),
2474 Err(_) => None,
2475 }
2476 };
2477
2478 match *elem {
2479 Type::BOOL => {
2480 if let Some(v) = try_array(
2481 row.try_get::<_, Option<Vec<bool>>>(idx)
2482 .map(|v| v.map(|a| a.into_iter().map(|x| json!(x)).collect())),
2483 ) {
2484 return v;
2485 }
2486 }
2487 Type::INT2 => {
2488 if let Some(v) = try_array(
2489 row.try_get::<_, Option<Vec<i16>>>(idx)
2490 .map(|v| v.map(|a| a.into_iter().map(|x| json!(x)).collect())),
2491 ) {
2492 return v;
2493 }
2494 }
2495 Type::INT4 | Type::OID => {
2496 if let Some(v) = try_array(
2497 row.try_get::<_, Option<Vec<i32>>>(idx)
2498 .map(|v| v.map(|a| a.into_iter().map(|x| json!(x)).collect())),
2499 ) {
2500 return v;
2501 }
2502 }
2503 Type::INT8 => {
2504 if let Some(v) = try_array(
2505 row.try_get::<_, Option<Vec<i64>>>(idx)
2506 .map(|v| v.map(|a| a.into_iter().map(|x| json!(x)).collect())),
2507 ) {
2508 return v;
2509 }
2510 }
2511 Type::FLOAT4 => {
2512 if let Some(v) = try_array(
2513 row.try_get::<_, Option<Vec<f32>>>(idx)
2514 .map(|v| v.map(|a| a.into_iter().map(|n| json!(n)).collect())),
2515 ) {
2516 return v;
2517 }
2518 }
2519 Type::FLOAT8 => {
2520 if let Some(v) = try_array(
2521 row.try_get::<_, Option<Vec<f64>>>(idx)
2522 .map(|v| v.map(|a| a.into_iter().map(|n| json!(n)).collect())),
2523 ) {
2524 return v;
2525 }
2526 }
2527 Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::NAME => {
2528 if let Some(v) = try_array(
2529 row.try_get::<_, Option<Vec<String>>>(idx)
2530 .map(|v| v.map(|a| a.into_iter().map(Value::String).collect())),
2531 ) {
2532 return v;
2533 }
2534 }
2535 Type::UUID => {
2536 if let Some(v) = try_array(
2537 row.try_get::<_, Option<Vec<Uuid>>>(idx)
2538 .map(|v| v.map(|a| a.into_iter().map(|u| json!(u.to_string())).collect())),
2539 ) {
2540 return v;
2541 }
2542 }
2543 Type::TIMESTAMP => {
2544 if let Some(v) = try_array(row.try_get::<_, Option<Vec<NaiveDateTime>>>(idx).map(|v| {
2545 v.map(|a| {
2546 a.into_iter()
2547 .map(|t| json!(t.format("%Y-%m-%dT%H:%M:%S%.f").to_string()))
2548 .collect()
2549 })
2550 })) {
2551 return v;
2552 }
2553 }
2554 Type::TIMESTAMPTZ => {
2555 if let Some(v) = try_array(
2556 row.try_get::<_, Option<Vec<DateTime<FixedOffset>>>>(idx)
2557 .map(|v| v.map(|a| a.into_iter().map(|t| json!(t.to_rfc3339())).collect())),
2558 ) {
2559 return v;
2560 }
2561 }
2562 Type::DATE => {
2563 if let Some(v) = try_array(row.try_get::<_, Option<Vec<NaiveDate>>>(idx).map(|v| {
2564 v.map(|a| {
2565 a.into_iter()
2566 .map(|d| json!(d.format("%Y-%m-%d").to_string()))
2567 .collect()
2568 })
2569 })) {
2570 return v;
2571 }
2572 }
2573 Type::JSON | Type::JSONB => {
2574 if let Some(v) = try_array(row.try_get::<_, Option<Vec<Value>>>(idx)) {
2575 return v;
2576 }
2577 }
2578 Type::NUMERIC => {
2579 if let Some(v) = try_array(
2580 row.try_get::<_, Option<Vec<Decimal>>>(idx)
2581 .map(|v| v.map(|a| a.into_iter().map(|d| json!(d.to_string())).collect())),
2582 ) {
2583 return v;
2584 }
2585 }
2586 Type::MONEY => {
2587 if let Some(v) = try_array(
2588 row.try_get::<_, Option<Vec<i64>>>(idx)
2589 .map(|v| v.map(|a| a.into_iter().map(|m| json!(money_to_string(m))).collect())),
2590 ) {
2591 return v;
2592 }
2593 }
2594 Type::BYTEA => {
2595 if let Some(v) = try_array(
2596 row.try_get::<_, Option<Vec<Vec<u8>>>>(idx)
2597 .map(|v| v.map(|a| a.into_iter().map(|b| json!(BASE64.encode(b))).collect())),
2598 ) {
2599 return v;
2600 }
2601 }
2602 _ => {}
2603 }
2604
2605 cell_to_json_untyped(row, idx, row.columns()[idx].type_())
2606}
2607
2608fn money_to_string(v: i64) -> String {
2610 let sign = if v < 0 { "-" } else { "" };
2611 let abs = v.unsigned_abs();
2612 format!("{}{}.{:04}", sign, abs / 10_000, abs % 10_000)
2613}
2614
2615fn cell_to_json_untyped(row: &tokio_postgres::Row, idx: usize, pg_type: &Type) -> Value {
2617 if let Ok(Some(s)) = row.try_get::<_, Option<String>>(idx) {
2618 return Value::String(s);
2619 }
2620 json!({
2621 "__untyped": true,
2622 "type": pg_type.name()
2623 })
2624}
2625
2626fn read_recent_log_errors() -> Vec<String> {
2627 let path = std::env::var("NEXQL_MCP_LOG")
2628 .map(std::path::PathBuf::from)
2629 .ok()
2630 .or_else(|| {
2631 std::env::var_os("HOME").map(|h| {
2632 std::path::PathBuf::from(h)
2633 .join(".config")
2634 .join("nexql-mcp")
2635 .join("logs")
2636 .join("nexql-mcp.log")
2637 })
2638 });
2639
2640 let Some(log_path) = path else {
2641 return Vec::new();
2642 };
2643
2644 let Ok(content) = std::fs::read_to_string(&log_path) else {
2645 return Vec::new();
2646 };
2647
2648 content
2649 .lines()
2650 .rev()
2651 .take(50)
2652 .filter(|line| {
2653 line.contains("ERROR")
2654 || line.contains("WARN")
2655 || line.contains("failed")
2656 || line.contains("Error")
2657 })
2658 .map(String::from)
2659 .collect()
2660}
2661
2662#[cfg(test)]
2663mod tests {
2664 use super::*;
2665 use crate::plan::build_explain_sql;
2666 use nexql_policy::PolicyFilter;
2667 use serde_json::json;
2668
2669 use crate::session::{ConnectionInfo, ConnectionPolicy, ToolSession};
2670 use nexql_policy::{AccessMode, PolicyCaps};
2671
2672 fn test_conn() -> ConnectionInfo {
2673 ConnectionInfo {
2674 id: "conn-1".into(),
2675 name: "conn-1".into(),
2676 host: Some("127.0.0.1".into()),
2677 port: Some(5432),
2678 database: Some("appdb".into()),
2679 params: Default::default(),
2680 policy: ConnectionPolicy {
2681 access_mode: AccessMode::Read,
2682 caps: PolicyCaps::default(),
2683 filter: PolicyFilter::default(),
2684 environment: None,
2685 },
2686 }
2687 }
2688
2689 #[test]
2690 fn policy_maps_one_to_one() {
2691 let f = PolicyFilter {
2692 allow_schemas: vec!["public".into()],
2693 deny_schemas: vec!["pgboss".into()],
2694 deny_tables: vec!["auth.*".into()],
2695 pii_columns: vec!["public.users.ssn".into()],
2696 };
2697 let q = policy_to_query_filter(&f);
2698 assert_eq!(q.allow_schemas, f.allow_schemas);
2699 assert_eq!(q.deny_schemas, f.deny_schemas);
2700 assert_eq!(q.deny_tables, f.deny_tables);
2701 assert_eq!(q.pii_columns, f.pii_columns);
2702 }
2703
2704 #[test]
2705 fn ok_json_wraps_arrays_for_cursor_structured_content() {
2706 let out = ToolOutcome::ok_json(json!([{ "id": 1 }, { "id": 2 }]));
2707 assert!(!out.is_error);
2708 let s = out.structured.as_ref().unwrap();
2709 assert!(s.is_object(), "structuredContent must be object, got {s}");
2710 assert_eq!(s["rows"].as_array().unwrap().len(), 2);
2711 assert!(out.text.contains("\"rows\""));
2712 }
2713
2714 #[test]
2715 fn ok_json_leaves_objects_unchanged() {
2716 let out = ToolOutcome::ok_json(json!({ "kind": "table", "name": "orders" }));
2717 let s = out.structured.as_ref().unwrap();
2718 assert_eq!(s["kind"], "table");
2719 assert!(s.get("rows").is_none());
2720 }
2721
2722 #[test]
2723 fn router_specs_include_phase4_and_phase9() {
2724 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2725 let router = ToolRouter::with_index_store(session, None);
2726 assert_eq!(router.specs().len(), ToolName::ACTIVE.len());
2727 let names: Vec<_> = router.specs().iter().map(|s| s.name.as_str()).collect();
2728 assert!(names.contains(&"search_schema"));
2729 assert!(names.contains(&"get_ddl"));
2730 assert!(names.contains(&"explain_analyze"));
2731 assert!(names.contains(&"get_index_status"));
2732 assert!(names.contains(&"list_extensions"));
2733 assert!(names.contains(&"server_settings"));
2734 assert!(names.contains(&"suggest_indexes"));
2735 assert!(names.contains(&"find_unused_indexes"));
2736 assert!(names.contains(&"bloat_report"));
2737 assert!(names.contains(&"find_missing_fks"));
2738 assert!(names.contains(&"export_query"));
2739 assert!(names.contains(&"list_roles"));
2740 assert!(names.contains(&"db_dashboard"));
2741 assert!(names.contains(&"deep_plan_analysis"));
2742 assert!(names.contains(&"execute_sql"));
2743 assert!(names.contains(&"edit_row"));
2744 assert!(names.contains(&"import_data"));
2745 assert!(names.contains(&"apply_ddl"));
2746 assert!(names.contains(&"create_index_concurrently"));
2747 assert!(names.contains(&"run_maintenance"));
2748 assert!(names.contains(&"terminate_query"));
2749 }
2750
2751 #[tokio::test]
2752 async fn write_tools_refuse_read_mode() {
2753 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2754 let router = ToolRouter::with_index_store(session, None);
2755 for tool in [
2756 "execute_sql",
2757 "edit_row",
2758 "import_data",
2759 "apply_ddl",
2760 "create_index_concurrently",
2761 "run_maintenance",
2762 "terminate_query",
2763 ] {
2764 let out = router
2765 .call(tool, json!({ "sql": "SELECT 1", "table": "public.t", "rows": [], "action": "insert", "values": {}, "pid": 1 }))
2766 .await;
2767 assert!(out.is_error, "{tool}: {}", out.text);
2768 assert!(
2769 out.text.contains("write") || out.text.contains("admin"),
2770 "{tool}: {}",
2771 out.text
2772 );
2773 }
2774 }
2775
2776 #[tokio::test]
2777 async fn table_stats_rejects_injection_ref() {
2778 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2779 let router = ToolRouter::with_index_store(session, None);
2780 let out = router
2781 .call("table_stats", json!({ "ref": "public.users; DROP" }))
2782 .await;
2783 assert!(out.is_error, "{}", out.text);
2784 assert!(
2785 out.text.contains("Invalid object reference") || out.text.contains("invalid arguments"),
2786 "expected ref validation error, got: {}",
2787 out.text
2788 );
2789 }
2790
2791 #[test]
2792 fn explain_transaction_path_builds_readonly_sequence() {
2793 let explain = build_explain_sql("SELECT 1", true);
2795 assert!(explain.starts_with("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)"));
2796 assert!(!explain.to_ascii_lowercase().contains("commit"));
2797 let steps = ["BEGIN", "SET TRANSACTION READ ONLY", &explain, "ROLLBACK"];
2798 assert_eq!(steps.len(), 4);
2799 assert_eq!(steps[0], "BEGIN");
2800 assert_eq!(steps[1], "SET TRANSACTION READ ONLY");
2801 assert_eq!(steps[3], "ROLLBACK");
2802 }
2803
2804 #[tokio::test]
2805 async fn missing_index_returns_actionable_error() {
2806 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2807 let router = ToolRouter::with_index_store(session, None);
2808 let out = router
2809 .call("search_schema", json!({ "query": "users" }))
2810 .await;
2811 assert!(out.is_error, "{}", out.text);
2812 assert!(
2813 out.text.contains("rebuild_index"),
2814 "expected actionable hint, got: {}",
2815 out.text
2816 );
2817 }
2818
2819 #[tokio::test]
2820 async fn empty_index_dir_returns_build_hint() {
2821 let tmp = tempfile::TempDir::new().unwrap();
2822 let store = IndexStore::new(tmp.path());
2823 let session = ToolSession::for_tests(
2824 vec![test_conn()],
2825 PolicyFilter::default(),
2826 Some(IndexStore::new(tmp.path())),
2827 );
2828 let router = ToolRouter::with_index_store(session, Some(store));
2829 let out = router
2830 .call("describe_object", json!({ "ref": "public.users" }))
2831 .await;
2832 assert!(out.is_error, "{}", out.text);
2833 assert!(
2834 out.text.contains("rebuild_index"),
2835 "expected build hint, got: {}",
2836 out.text
2837 );
2838 }
2839
2840 #[tokio::test]
2841 async fn outcome_tagged_with_connection_id_and_database() {
2842 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2843 let router = ToolRouter::new(session);
2844 let out = router.call("list_connections", json!({})).await;
2845 let structured = out.structured.expect("structured outcome");
2846 assert_eq!(
2847 structured.get("connectionId").and_then(|v| v.as_str()),
2848 Some("conn-1")
2849 );
2850 assert_eq!(
2851 structured.get("database").and_then(|v| v.as_str()),
2852 Some("appdb")
2853 );
2854 }
2855
2856 #[tokio::test]
2857 async fn setup_connection_returns_needs_input_when_incomplete() {
2858 unsafe {
2859 std::env::remove_var("DATABASE_URL");
2860 std::env::remove_var("POSTGRES_URL");
2861 std::env::remove_var("PGHOST");
2862 }
2863 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2864 let router = ToolRouter::new(session);
2865 let out = router.call("setup_connection", json!({})).await;
2866 let structured = out.structured.expect("structured outcome");
2867 assert!(structured.get("status").is_some());
2868 }
2869
2870 #[tokio::test]
2871 async fn save_profile_persists_config() {
2872 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2873 let router = ToolRouter::new(session);
2874 let temp_dir = tempfile::tempdir().unwrap();
2875 let cfg_path = temp_dir.path().join("config.toml");
2876 unsafe {
2877 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
2878 }
2879
2880 let out = router
2881 .call(
2882 "save_profile",
2883 json!({
2884 "name": "staging",
2885 "host": "127.0.0.1",
2886 "port": 5432,
2887 "dbname": "stage_db",
2888 "user": "stage_user"
2889 }),
2890 )
2891 .await;
2892
2893 let structured = out.structured.expect("structured outcome");
2894 assert_eq!(
2895 structured.get("status").and_then(|v| v.as_str()),
2896 Some("saved")
2897 );
2898 assert_eq!(
2899 structured.get("profile").and_then(|v| v.as_str()),
2900 Some("staging")
2901 );
2902 }
2903
2904 #[tokio::test]
2905 async fn check_ddl_safety_tool_dispatches_ast_report() {
2906 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
2907 let router = ToolRouter::new(session);
2908 let out = router
2909 .call(
2910 "check_ddl_safety",
2911 json!({ "ddl": "CREATE INDEX idx_col ON users(col);" }),
2912 )
2913 .await;
2914 let structured = out.structured.expect("structured outcome");
2915 assert_eq!(
2916 structured.get("overall_risk").and_then(|v| v.as_str()),
2917 Some("CRITICAL")
2918 );
2919 }
2920}