1use std::sync::Arc;
7
8use nexql_index::{
9 BuildDepth, BuildMode, BuildRequest, CatalogDb, Embedder, IndexQueryService, IndexScope,
10 IndexStore, ObjectEntry, PgCatalogDb, QueryPolicyFilter, RefResolution, SearchOptions,
11 build_index,
12};
13use nexql_policy::{
14 PolicyCaps, PolicyFilter, SqlDecision, enforce_read_table_policy, select_table_refs,
15 validate_readonly_sql,
16};
17use serde_json::{Value, json};
18
19use crate::cell_json::{
20 columnarize_read_payload, columnarize_row_arrays, redact_pii_in_payload,
21 rows_to_json_array, rows_to_json_array_with_total,
22};
23use crate::critique;
24use crate::error::ToolError;
25use crate::export::{ExportFormat, columns_from_rows, rows_to_csv, rows_to_sql_insert};
26use crate::format::rows_to_markdown;
27use crate::plan::{analyze_deep_plan, build_explain_sql, extract_plan_metrics};
28use crate::registry::ToolName;
29use crate::resolve::{self, DEFAULT_RESOLVE_REFS_LIMIT};
30use crate::schema::{ToolSpec, active_tools};
31use crate::session::{CheckoutTarget, ScopedContext, ToolSession};
32use crate::sql::{self, REPORT_LIMIT_DEFAULT, SLOW_QUERIES_DEFAULT, parse_ref};
33use crate::write::{
34 apply_ddl, create_index_concurrently, edit_row, execute_sql, import_data, run_maintenance,
35 terminate_query,
36};
37
38const SEARCH_SCHEMA_LIMIT: usize = 10;
40
41const RUN_SELECT_DEFAULT_LIMIT: u32 = 50;
43
44const NEXQL_TOTAL_COUNT_COL: &str = "nexql_total_count";
46
47const NO_INDEX_HINT: &str =
48 "No schema index configured — call the 'rebuild_index' tool to build an index.";
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51enum RunSelectFormat {
52 Compact,
53 Json,
54 Markdown,
55 Csv,
56}
57
58impl RunSelectFormat {
59 fn parse(s: &str) -> Result<Self, ToolError> {
60 match s.to_ascii_lowercase().as_str() {
61 "compact" => Ok(Self::Compact),
62 "json" => Ok(Self::Json),
63 "markdown" => Ok(Self::Markdown),
64 "csv" => Ok(Self::Csv),
65 other => Err(ToolError::InvalidArgs(format!(
66 "Unsupported format \"{other}\". Use compact, json, markdown, or csv."
67 ))),
68 }
69 }
70}
71
72struct ExecutionScope {
73 ctx: ScopedContext,
74 filter: PolicyFilter,
75 caps: PolicyCaps,
76}
77
78#[derive(Debug, Clone)]
79pub struct ToolOutcome {
80 pub text: String,
81 pub structured: Option<Value>,
82 pub is_error: bool,
83}
84
85impl ToolOutcome {
86 pub fn ok_json(value: Value) -> Self {
92 let value = ensure_structured_object(value);
93 let text = serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string());
94 Self {
95 text,
96 structured: Some(value),
97 is_error: false,
98 }
99 }
100
101 pub fn err(msg: impl Into<String>) -> Self {
102 let message = msg.into();
103 Self {
104 text: message.clone(),
105 structured: Some(json!({ "error": message })),
106 is_error: true,
107 }
108 }
109}
110
111fn ensure_structured_object(value: Value) -> Value {
113 match value {
114 Value::Array(rows) => json!({ "rows": rows }),
115 other => other,
116 }
117}
118
119pub struct ToolRouter {
120 session: Arc<ToolSession>,
121 index_override: Option<Option<IndexStore>>,
123 use_semantic: bool,
125 embedder: Option<Arc<dyn Embedder>>,
126 specs: Vec<ToolSpec>,
127 managed_extension: bool,
128}
129
130impl ToolRouter {
131 pub fn new(session: Arc<ToolSession>) -> Self {
132 Self {
133 session,
134 index_override: None,
135 use_semantic: false,
136 embedder: None,
137 specs: active_tools(),
138 managed_extension: false,
139 }
140 }
141
142 pub fn with_index_store(session: Arc<ToolSession>, store: Option<IndexStore>) -> Self {
144 Self {
145 session,
146 index_override: Some(store),
147 use_semantic: false,
148 embedder: None,
149 specs: active_tools(),
150 managed_extension: false,
151 }
152 }
153
154 pub fn with_semantic(
156 mut self,
157 use_semantic: bool,
158 embedder: Option<Arc<dyn Embedder>>,
159 ) -> Self {
160 self.use_semantic = use_semantic;
161 self.embedder = embedder;
162 self
163 }
164
165 pub fn with_profile(mut self, profile: crate::registry::ToolProfile) -> Self {
167 self.specs = crate::schema::tools_for_profile(profile);
168 self
169 }
170
171 pub fn with_managed_extension(mut self, enabled: bool) -> Self {
173 self.managed_extension = enabled;
174 if enabled {
175 const BLOCKED: &[ToolName] = &[
176 ToolName::SetupConnection,
177 ToolName::SaveProfile,
178 ToolName::TestProfile,
179 ToolName::ExportProfile,
180 ToolName::ImportProfile,
181 ];
182 self.specs.retain(|s| !BLOCKED.contains(&s.name));
183 }
184 self
185 }
186
187 pub fn specs(&self) -> &[ToolSpec] {
188 &self.specs
189 }
190
191 fn index_store(&self) -> Option<&IndexStore> {
192 match &self.index_override {
193 Some(inner) => inner.as_ref(),
194 None => self.session.index_store.as_ref(),
195 }
196 }
197
198 fn query_filter(&self) -> QueryPolicyFilter {
199 policy_to_query_filter(&self.session.filter())
200 }
201
202 fn query_filter_for(&self, connection_id: &str) -> QueryPolicyFilter {
203 policy_to_query_filter(&self.session.filter_for(connection_id))
204 }
205
206 async fn execution_scope_from_args(&self, args: &Value) -> Result<ExecutionScope, ToolError> {
207 let ctx = self
208 .session
209 .resolve_scoped_context(
210 args.get("connectionId").and_then(|v| v.as_str()),
211 args.get("database").and_then(|v| v.as_str()),
212 )
213 .await?;
214 Ok(ExecutionScope {
215 filter: self.session.filter_for(&ctx.connection_id),
216 caps: self.session.caps_for(&ctx.connection_id),
217 ctx,
218 })
219 }
220
221 fn scope_tag(scope: &ExecutionScope, mut outcome: ToolOutcome) -> ToolOutcome {
222 if let Some(obj) = outcome.structured.as_mut().and_then(|v| v.as_object_mut()) {
223 obj.insert("connectionId".into(), json!(scope.ctx.connection_id));
224 obj.insert("database".into(), json!(scope.ctx.database));
225 }
226 outcome
227 }
228
229 pub async fn call(&self, name: &str, args: Value) -> ToolOutcome {
230 let outcome = match self.call_inner(name, args).await {
231 Ok(out) => out,
232 Err(e) => ToolOutcome::err(e.to_string()),
233 };
234 let outcome = Self::columnarize_outcome(name, outcome);
235 self.tag_outcome_with_context(outcome).await
236 }
237
238 fn columnarize_outcome(name: &str, mut outcome: ToolOutcome) -> ToolOutcome {
257 if outcome.is_error {
258 return outcome;
259 }
260 if matches!(
261 ToolName::parse(name),
262 Some(ToolName::Orient) | Some(ToolName::GetJoinPath)
263 ) {
264 return outcome;
265 }
266 let Some(structured) = outcome.structured.take() else {
267 return outcome;
268 };
269 let transformed = columnarize_row_arrays(structured);
270 if let Ok(text) = serde_json::to_string(&transformed) {
271 outcome.text = text;
272 }
273 outcome.structured = Some(transformed);
274 outcome
275 }
276
277 async fn tag_outcome_with_context(&self, mut outcome: ToolOutcome) -> ToolOutcome {
278 let (connection_id, database) = if let Some(structured) = &outcome.structured {
279 let cid = structured.get("connectionId").and_then(|v| v.as_str());
280 let db = structured.get("database").and_then(|v| v.as_str());
281 if let (Some(c), Some(d)) = (cid, db) {
282 (c.to_string(), d.to_string())
283 } else {
284 self.session.active_context().await
285 }
286 } else {
287 self.session.active_context().await
288 };
289 let access_mode = match self.session.access_mode() {
290 nexql_policy::AccessMode::Read => "read",
291 nexql_policy::AccessMode::Write => "write",
292 nexql_policy::AccessMode::Admin => "admin",
293 };
294 let mut freshness: Option<serde_json::Value> = None;
295 if let Some(store) = self.session.index_store.as_ref() {
296 let base = store.base_dir(&connection_id, &database);
297 if let Ok(Some(manifest)) = store.read_manifest(&base) {
298 let stale = self.session.is_index_stale(&connection_id, &database);
299 let mut freshness_obj = json!({
300 "indexedAt": manifest.indexed_at,
301 "schemaFingerprint": manifest.schema_fingerprint,
302 "stale": stale,
303 });
304 if stale {
305 freshness_obj["reason"] = json!("schema_changed");
306 }
307 freshness = Some(freshness_obj);
308 } else {
309 freshness = Some(json!({ "stale": true, "reason": "no_index" }));
310 }
311 }
312 if let Some(ref mut structured) = outcome.structured
313 && let Some(obj) = structured.as_object_mut()
314 {
315 if !obj.contains_key("connectionId") {
316 obj.insert("connectionId".into(), json!(connection_id));
317 }
318 if !obj.contains_key("database") {
319 obj.insert("database".into(), json!(database));
320 }
321 if !obj.contains_key("accessMode") {
322 obj.insert("accessMode".into(), json!(access_mode));
323 }
324 if let Some(ref f) = freshness {
325 obj.insert("freshness".into(), f.clone());
326 }
327 }
328 let header = format!(
329 "[context connectionId={connection_id} database={database} accessMode={access_mode}]\n"
330 );
331 if !outcome.text.starts_with("[context ") {
332 outcome.text = format!("{header}{}", outcome.text);
333 }
334 outcome
335 }
336
337 async fn call_inner(&self, name: &str, args: Value) -> Result<ToolOutcome, ToolError> {
338 let tool = ToolName::parse(name).ok_or_else(|| ToolError::Unknown(name.to_string()))?;
339 match tool {
340 ToolName::ListConnections => Ok(self.list_connections()),
341 ToolName::ListDatabases => self.list_databases(&args).await,
342 ToolName::ListSchemas => self.list_schemas().await,
343 ToolName::ListObjects => self.list_objects(&args).await,
344 ToolName::GetCurrentContext => self.get_current_context().await,
345 ToolName::SwitchConnection => self.switch_connection(&args).await,
346 ToolName::RunSelect => self.run_select(&args).await,
347 ToolName::ExplainQuery => self.explain_query(&args).await,
348 ToolName::DiscoverTools => self.discover_tools(&args).await,
349 ToolName::RunDoctor => self.run_doctor_tool().await,
350 ToolName::SetupConnection => self.setup_connection_tool(&args).await,
351 ToolName::SaveProfile => self.save_profile_tool(&args).await,
352 ToolName::TestProfile => self.test_profile_tool(&args).await,
353 ToolName::ExportProfile => self.export_profile_tool(&args).await,
354 ToolName::ImportProfile => self.import_profile_tool(&args).await,
355 ToolName::ResolveTarget => self.resolve_target(&args).await,
356 ToolName::Orient => self.orient(&args).await,
357 ToolName::InspectOrSearch => self.inspect_or_search(&args).await,
358 ToolName::SearchAllDatabases => self.search_all_databases(&args).await,
359 ToolName::SearchSchema => self.search_schema(&args).await,
360 ToolName::DescribeObject => self.describe_object(&args).await,
361 ToolName::GetJoinPath => self.get_join_path(&args).await,
362 ToolName::SampleValues => self.sample_values(&args).await,
363 ToolName::GetDdl => self.get_ddl(&args).await,
364 ToolName::TableStats => self.table_stats(&args).await,
365 ToolName::IndexUsage => self.index_usage(&args).await,
366 ToolName::ListRunningQueries => self.list_running_queries().await,
367 ToolName::FindBlockingLocks => self.find_blocking_locks().await,
368 ToolName::SlowQueries => self.slow_queries(&args).await,
369 ToolName::DbHealthCheck => self.db_health_check().await,
370 ToolName::GetIndexStatus => self.get_index_status().await,
371 ToolName::ListExtensions => self.list_extensions().await,
372 ToolName::ServerSettings => self.server_settings().await,
373 ToolName::SuggestIndexes => self.suggest_indexes(&args).await,
374 ToolName::FindUnusedIndexes => self.find_unused_indexes(&args).await,
375 ToolName::BloatReport => self.bloat_report(&args).await,
376 ToolName::FindMissingFks => self.find_missing_fks(&args).await,
377 ToolName::ExportQuery => self.export_query(&args).await,
378 ToolName::ListRoles => self.list_roles(&args).await,
379 ToolName::DbDashboard => self.db_dashboard().await,
380 ToolName::DeepPlanAnalysis => self.deep_plan_analysis(&args).await,
381 ToolName::SchemaDiff => self.schema_diff(&args).await,
382 ToolName::GenerateMigration => self.generate_migration(&args).await,
383 ToolName::ExecuteSql => self.execute_sql_tool(&args).await,
384 ToolName::EditRow => self.edit_row_tool(&args).await,
385 ToolName::ImportData => self.import_data_tool(&args).await,
386 ToolName::ApplyDdl => self.apply_ddl_tool(&args).await,
387 ToolName::CreateIndexConcurrently => self.create_index_concurrently_tool(&args).await,
388 ToolName::RunMaintenance => self.run_maintenance_tool(&args).await,
389 ToolName::TerminateQuery => self.terminate_query_tool(&args).await,
390 ToolName::AutoTuneQuery => self.auto_tune_query(&args).await,
391 ToolName::CheckDdlSafety => self.check_ddl_safety_tool(&args).await,
392 ToolName::RebuildIndex => self.rebuild_index_tool(&args).await,
393 ToolName::RefreshIndex => self.refresh_index_tool(&args).await,
394 }
395 }
396
397 fn require_write(&self) -> Result<(), ToolError> {
398 if !self.session.access_mode().allows_writes() {
399 return Err(ToolError::Execution(
400 "write tools require --access-mode write or admin (current session: read)".into(),
401 ));
402 }
403 Ok(())
404 }
405
406 fn require_admin(&self) -> Result<(), ToolError> {
407 if !self.session.access_mode().allows_admin() {
408 return Err(ToolError::Execution(
409 "admin tools require --access-mode admin".into(),
410 ));
411 }
412 Ok(())
413 }
414
415 async fn execute_sql_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
416 self.require_write()?;
417 let sql = args
418 .get("sql")
419 .and_then(|v| v.as_str())
420 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
421 let dry_run = args
422 .get("dry_run")
423 .and_then(|v| v.as_bool())
424 .unwrap_or(false);
425 let include_diff = args
426 .get("include_diff")
427 .and_then(|v| v.as_bool())
428 .unwrap_or(dry_run);
429 let outcome = execute_sql(&self.session, sql, dry_run, include_diff).await?;
430 Ok(self.attach_dml_critique(sql, outcome).await)
431 }
432
433 async fn attach_dml_critique(&self, sql: &str, mut outcome: ToolOutcome) -> ToolOutcome {
439 if outcome.is_error {
440 return outcome;
441 }
442 let Some(structured) = outcome.structured.as_ref() else {
443 return outcome;
444 };
445 if structured.get("rows_affected").and_then(Value::as_u64) != Some(0) {
446 return outcome;
447 }
448 let Some((table, col, val)) = critique::dml_equality_filter(sql) else {
449 return outcome;
450 };
451 let Some((schema, name)) = table.split_once('.') else {
452 return outcome;
453 };
454 let table_ref = nexql_policy::ObjectRef::new(schema, name);
455 let Some(values) = self.sample_values_best_effort(&table_ref, &col).await else {
456 return outcome;
457 };
458 let sample = values
459 .iter()
460 .take(5)
461 .cloned()
462 .collect::<Vec<_>>()
463 .join(", ");
464 if let Some(obj) = outcome.structured.as_mut().and_then(|v| v.as_object_mut()) {
465 obj.insert(
466 "critique".into(),
467 json!([{
468 "signal": "zero_rows",
469 "message": format!(
470 "0 rows affected: {col} = '{val}'; observed values include: {sample}"
471 ),
472 }]),
473 );
474 }
475 if let Some(structured) = &outcome.structured
476 && let Ok(text) = serde_json::to_string(structured)
477 {
478 outcome.text = text;
479 }
480 outcome
481 }
482
483 async fn edit_row_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
484 self.require_write()?;
485 edit_row(&self.session, args).await
486 }
487
488 async fn import_data_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
489 self.require_write()?;
490 import_data(&self.session, args).await
491 }
492
493 async fn apply_ddl_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
494 self.require_admin()?;
495 let sql = args
496 .get("sql")
497 .and_then(|v| v.as_str())
498 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
499 let dry_run = args
500 .get("dry_run")
501 .and_then(|v| v.as_bool())
502 .unwrap_or(false);
503 apply_ddl(&self.session, sql, dry_run).await
504 }
505
506 async fn create_index_concurrently_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
507 self.require_admin()?;
508 let sql = args
509 .get("sql")
510 .and_then(|v| v.as_str())
511 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
512 create_index_concurrently(&self.session, sql).await
513 }
514
515 async fn run_maintenance_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
516 self.require_admin()?;
517 run_maintenance(&self.session, args).await
518 }
519
520 async fn terminate_query_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
521 self.require_admin()?;
522 terminate_query(&self.session, args).await
523 }
524
525 async fn live_databases_by_connection(
528 &self,
529 ) -> std::collections::HashMap<String, std::collections::HashSet<String>> {
530 use std::collections::HashMap;
531 let mut map = HashMap::new();
532 for conn in self.session.connections() {
533 if let Ok(names) = self.list_database_names_for(&conn).await {
534 map.insert(conn.id.clone(), names.into_iter().collect());
535 }
536 }
537 map
538 }
539
540 async fn list_database_names_for(
541 &self,
542 conn: &crate::session::ConnectionInfo,
543 ) -> Result<Vec<String>, ToolError> {
544 let client = if self.session.active_context().await.0 == conn.id {
545 self.session.checkout().await?
546 } else {
547 let pool_opts = self.session.pool_opts();
548 let pool = nexql_conn::create_pool(&conn.params, &pool_opts).await?;
549 nexql_conn::checkout_guarded(&pool, &pool_opts).await?
550 };
551 let rows = client
552 .query(
553 "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname",
554 &[],
555 )
556 .await?;
557 Ok(rows.iter().map(|r| r.get(0)).collect())
558 }
559
560 async fn resolve_target(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
561 let hint = args
562 .get("hint")
563 .and_then(|v| v.as_str())
564 .map(str::trim)
565 .filter(|s| !s.is_empty());
566 let object_hint = args
567 .get("objectHint")
568 .and_then(|v| v.as_str())
569 .map(str::trim)
570 .filter(|s| !s.is_empty());
571 if hint.is_none() && object_hint.is_none() {
572 return Err(ToolError::InvalidArgs(
573 "At least one of \"hint\" or \"objectHint\" is required.".into(),
574 ));
575 }
576
577 let connections = self.session.connections();
578 if connections.is_empty() {
579 return Ok(ToolOutcome::err("No connections configured."));
580 }
581
582 #[derive(Clone)]
583 struct Candidate {
584 connection_id: String,
585 database: String,
586 }
587 fn key_of(c: &Candidate) -> String {
588 format!("{}\u{0}{}", c.connection_id, c.database)
589 }
590
591 let indexed: Vec<(String, String)> = self
592 .index_store()
593 .map(|store| store.list_indexed_databases().unwrap_or_default())
594 .unwrap_or_default();
595
596 let live_dbs = self.live_databases_by_connection().await;
597
598 let mut seen = std::collections::HashSet::new();
599 let mut candidates: Vec<Candidate> = Vec::new();
600 let mut add_candidate = |connection_id: &str, database: &str| {
601 if !connections.iter().any(|c| c.id == connection_id) {
602 return;
603 }
604 let key = format!("{connection_id}\u{0}{database}");
605 if !seen.insert(key) {
606 return;
607 }
608 candidates.push(Candidate {
609 connection_id: connection_id.to_string(),
610 database: database.to_string(),
611 });
612 };
613 for (cid, db) in &indexed {
614 if let Some(set) = live_dbs.get(cid) {
615 if set.contains(db) {
616 add_candidate(cid, db);
617 } else if let Some(store) = self.index_store() {
618 let _ = store.clear_index(cid, db);
619 }
620 }
621 }
622 for c in &connections {
623 if let Some(dbs) = live_dbs.get(&c.id) {
624 for db in dbs {
625 add_candidate(&c.id, db);
626 }
627 } else {
628 let db = c.database.clone().unwrap_or_else(|| "postgres".into());
629 add_candidate(&c.id, &db);
630 }
631 }
632
633 let mut scored: std::collections::HashMap<String, (Candidate, f64, Vec<String>)> =
634 std::collections::HashMap::new();
635
636 if let Some(hint) = hint {
637 for c in &candidates {
638 let Some(conn) = connections.iter().find(|x| x.id == c.connection_id) else {
639 continue;
640 };
641 let fields: [(&str, &str); 3] = [
642 ("connection name", conn.name.as_str()),
643 ("host", conn.host.as_deref().unwrap_or("")),
644 ("database", c.database.as_str()),
645 ];
646 let mut best = 0.0f64;
647 let mut best_field = "";
648 for (label, value) in fields {
649 let s = fuzzy_score(hint, value);
650 if s > best {
651 best = s;
652 best_field = label;
653 }
654 }
655 if best > 0.0 {
656 let entry = scored
657 .entry(key_of(c))
658 .or_insert_with(|| (c.clone(), 0.0, Vec::new()));
659 entry.1 += best;
660 entry
661 .2
662 .push(format!("{best_field} matched hint \"{hint}\" ({best:.0})"));
663 }
664 }
665 }
666
667 if let Some(object_hint) = object_hint
668 && let Some(store) = self.index_store()
669 {
670 let filter = self.query_filter();
671 for (cid, db) in &indexed {
672 let svc = IndexQueryService::new(store, cid.clone(), db.clone());
673 if let Ok(hits) = svc.search_schema(
674 object_hint,
675 3,
676 Some(&filter),
677 SearchOptions {
678 use_semantic: self.use_semantic,
679 embedder: self.embedder.as_deref(),
680 },
681 ) && let Some(top) = hits.first()
682 {
683 let c = Candidate {
684 connection_id: cid.clone(),
685 database: db.clone(),
686 };
687 let entry = scored
688 .entry(key_of(&c))
689 .or_insert_with(|| (c.clone(), 0.0, Vec::new()));
690 entry.1 += top.score * 10.0;
691 entry.2.push(format!(
692 "schema search for \"{object_hint}\" found {} (score {:.2})",
693 top.ref_, top.score
694 ));
695 }
696 }
697 }
698
699 let mut ranked: Vec<(Candidate, f64, Vec<String>)> = scored.into_values().collect();
700 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
701
702 if ranked.is_empty() {
703 let candidates_json: Vec<Value> = connections
704 .iter()
705 .map(|c| {
706 json!({
707 "connectionId": c.id,
708 "connectionName": c.name,
709 "database": c.database.clone().unwrap_or_else(|| "postgres".into()),
710 })
711 })
712 .collect();
713 return Ok(ToolOutcome::ok_json(json!({
714 "ambiguous": true,
715 "message": format!(
716 "No connection/database matched \"{}\". Choose from the configured connections.",
717 hint.or(object_hint).unwrap_or_default()
718 ),
719 "candidates": candidates_json
720 })));
721 }
722
723 let winner = &ranked[0];
724 let is_tied = ranked
725 .get(1)
726 .is_some_and(|runner_up| runner_up.1 >= winner.1 * 0.85);
727
728 if is_tied {
729 let threshold = winner.1 * 0.85;
730 let tied: Vec<&(Candidate, f64, Vec<String>)> =
731 ranked.iter().filter(|r| r.1 >= threshold).take(5).collect();
732 let candidates_json: Vec<Value> = tied
733 .iter()
734 .filter_map(|(c, score, evidence)| {
735 connections
736 .iter()
737 .find(|x| x.id == c.connection_id)
738 .map(|conn| {
739 json!({
740 "connectionId": c.connection_id,
741 "connectionName": conn.name,
742 "database": c.database,
743 "score": score,
744 "evidence": evidence,
745 })
746 })
747 })
748 .collect();
749 return Ok(ToolOutcome::ok_json(json!({
750 "ambiguous": true,
751 "message": format!("{} equally-plausible candidates matched.", tied.len()),
752 "candidates": candidates_json
753 })));
754 }
755
756 let (winner_candidate, winner_score, winner_evidence) = winner;
757
758 if let Some(object_hint) = object_hint
759 && let Some(store) = self.index_store()
760 {
761 let filter = self.query_filter();
762 let svc = IndexQueryService::new(
763 store,
764 &winner_candidate.connection_id,
765 &winner_candidate.database,
766 );
767 if let Ok(hits) = svc.search_schema(
768 object_hint,
769 5,
770 Some(&filter),
771 SearchOptions {
772 use_semantic: self.use_semantic,
773 embedder: self.embedder.as_deref(),
774 },
775 ) && hits.len() >= 2
776 {
777 let top_score = hits[0].score;
778 let tied: Vec<&nexql_index::RankedHit> = hits
779 .iter()
780 .filter(|h| scores_equal(h.score, top_score))
781 .collect();
782 if tied.len() > 1 {
783 let candidates_json: Vec<Value> = tied
784 .iter()
785 .map(|h| {
786 json!({
787 "ref": h.ref_,
788 "score": h.score,
789 "kind": h.kind,
790 "connectionId": winner_candidate.connection_id,
791 "database": winner_candidate.database,
792 })
793 })
794 .collect();
795 return Ok(ToolOutcome::ok_json(json!({
796 "ambiguous": true,
797 "message": format!(
798 "{} objects matched \"{object_hint}\" with equal scores — choose explicitly.",
799 tied.len()
800 ),
801 "candidates": candidates_json,
802 })));
803 }
804 }
805 }
806
807 self.session
808 .switch(
809 &winner_candidate.connection_id,
810 Some(winner_candidate.database.clone()),
811 )
812 .await?;
813 let conn = connections
814 .iter()
815 .find(|x| x.id == winner_candidate.connection_id)
816 .ok_or_else(|| ToolError::Execution("resolved connection vanished".into()))?;
817
818 Ok(ToolOutcome::ok_json(json!({
819 "resolved": true,
820 "connectionId": winner_candidate.connection_id,
821 "connectionName": conn.name,
822 "database": winner_candidate.database,
823 "confidence": winner_score,
824 "evidence": winner_evidence,
825 })))
826 }
827
828 async fn orient(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
836 const TABLE_LIMIT: usize = 40;
837 const ENUM_MAX_DISTINCT: f64 = 20.0;
842
843 let focus = args
844 .get("focus")
845 .and_then(|v| v.as_str())
846 .map(str::to_ascii_lowercase);
847 let (connection_id, database) = self.session.active_context().await;
848
849 let Some(store) = self.index_store() else {
850 return Ok(ToolOutcome::ok_json(json!({
851 "database": database,
852 "tables": [],
853 "joins": [],
854 "enums": {},
855 "notes": [NO_INDEX_HINT],
856 })));
857 };
858 let base = store.base_dir(&connection_id, &database);
859 if store.read_manifest(&base)?.is_none()
860 && let Err(e) = self.ensure_index_warm().await
861 {
862 return Ok(ToolOutcome::ok_json(json!({
863 "database": database,
864 "tables": [],
865 "joins": [],
866 "enums": {},
867 "notes": [e.to_string()],
868 })));
869 }
870 let Some(manifest) = store.read_manifest(&base)? else {
871 return Ok(ToolOutcome::ok_json(json!({
872 "database": database,
873 "tables": [],
874 "joins": [],
875 "enums": {},
876 "notes": [format!(
877 "No schema index for database \"{database}\" — call the 'rebuild_index' tool to build an index."
878 )],
879 })));
880 };
881
882 let mut entries: Vec<(String, ObjectEntry)> = Vec::new();
883 for shard in &manifest.shards {
884 if let Some(shard_entries) = store.read_shard_entries(&base, &shard.file)? {
885 entries.extend(shard_entries);
886 }
887 }
888 entries.sort_by(|a, b| a.0.cmp(&b.0));
889 if let Some(f) = &focus {
890 entries.retain(|(ref_, _)| ref_.to_ascii_lowercase().contains(f.as_str()));
891 }
892
893 let mut notes: Vec<String> = manifest.stats.warnings.clone();
894 let total = entries.len();
895 if total > TABLE_LIMIT {
896 notes.push(format!(
897 "Showing {TABLE_LIMIT} of {total} objects — pass `focus` to narrow the digest."
898 ));
899 entries.truncate(TABLE_LIMIT);
900 }
901 let shown_refs: std::collections::HashSet<&str> =
902 entries.iter().map(|(r, _)| r.as_str()).collect();
903
904 let mut tables = Vec::with_capacity(entries.len());
905 let mut enums = serde_json::Map::new();
906 for (ref_, entry) in &entries {
907 if entry.excluded == Some(true) {
908 continue;
909 }
910 let pk = entry
911 .primary_key
912 .as_ref()
913 .map(|cols| cols.join(","))
914 .filter(|s| !s.is_empty());
915 let columns = entry
916 .columns
917 .iter()
918 .map(|c| {
919 let bang = if c.not_null { "!" } else { "" };
920 format!("{}:{}{bang}", c.name, c.type_name)
921 })
922 .collect::<Vec<_>>()
923 .join(", ");
924 tables.push(json!({
925 "ref": ref_,
926 "kind": entry.kind.as_str(),
927 "rows": format!("~{}", entry.row_estimate.round() as i64),
928 "pk": pk,
929 "columns": columns,
930 }));
931
932 for col in &entry.columns {
933 if col.pii == Some(true) {
934 continue;
935 }
936 let is_textish = col.type_name.contains("char") || col.type_name.contains("text");
937 let Some(profile) = &col.profile else {
938 continue;
939 };
940 if !is_textish
941 || profile.n_distinct <= 0.0
942 || profile.n_distinct > ENUM_MAX_DISTINCT
943 {
944 continue;
945 }
946 if let Some(vals) = &profile.common_values
947 && !vals.is_empty()
948 {
949 enums.insert(format!("{ref_}.{}", col.name), json!(vals));
950 }
951 }
952 }
953
954 let mut joins = Vec::new();
955 if let Some(graph) = store.read_join_graph(&base, &manifest)? {
956 for edge in &graph.edges {
957 if focus.is_some()
958 && !(shown_refs.contains(edge.from.as_str())
959 || shown_refs.contains(edge.to.as_str()))
960 {
961 continue;
962 }
963 let edge_str = edge
964 .cols
965 .iter()
966 .map(|(from_col, to_col)| {
967 format!("{}.{from_col} -> {}.{to_col}", edge.from, edge.to)
968 })
969 .collect::<Vec<_>>()
970 .join("; ");
971 let inferred = edge.inferred == Some(true);
972 joins.push(json!({
973 "edge": edge_str,
974 "declared": !inferred,
975 "detection": if inferred { "join_graph_inferred" } else { "declared_fk" },
976 "via": edge.via,
977 "disabled": edge.disabled == Some(true),
978 }));
979 }
980 }
981
982 Ok(ToolOutcome::ok_json(json!({
983 "database": database,
984 "tables": tables,
985 "joins": joins,
986 "enums": Value::Object(enums),
987 "notes": notes,
988 })))
989 }
990
991 async fn discover_tools(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
992 let query = args
993 .get("query")
994 .and_then(|v| v.as_str())
995 .map(str::to_lowercase);
996 let category = args
997 .get("category")
998 .and_then(|v| v.as_str())
999 .map(str::to_lowercase);
1000
1001 let all_specs = active_tools();
1003 let filtered: Vec<Value> = all_specs
1004 .into_iter()
1005 .filter(|spec| {
1006 if spec.name == ToolName::DiscoverTools {
1007 return false;
1008 }
1009 if let Some(ref cat) = category {
1010 match cat.as_str() {
1011 "query" if !ToolName::QUERY_PROFILE.contains(&spec.name) => return false,
1012 "dba" if !ToolName::DBA_PROFILE.contains(&spec.name) => return false,
1013 "write" if !ToolName::PHASE9.contains(&spec.name) => return false,
1014 _ => {}
1015 }
1016 }
1017 if let Some(ref q) = query {
1018 let name_match = spec.name.as_str().contains(q.as_str());
1019 let desc_match = spec.description.to_lowercase().contains(q.as_str());
1020 if !name_match && !desc_match {
1021 return false;
1022 }
1023 }
1024 true
1025 })
1026 .map(|spec| {
1027 json!({
1028 "name": spec.name.as_str(),
1029 "description": spec.description,
1030 "input_schema": spec.input_schema,
1031 })
1032 })
1033 .collect();
1034
1035 Ok(ToolOutcome::ok_json(json!({
1036 "query": args.get("query"),
1037 "category": args.get("category"),
1038 "count": filtered.len(),
1039 "tools": filtered,
1040 })))
1041 }
1042
1043 fn build_tuning_summary(plan_structured: &Option<Value>, suggestions: &Value) -> String {
1044 let mut parts = Vec::new();
1045 if let Some(structured) = plan_structured
1046 && let Some(metrics) = structured.get("metrics")
1047 {
1048 if let Some(exec_time) = metrics.get("executionTime").and_then(|v| v.as_f64()) {
1049 parts.push(format!("Query executed in {:.2}ms.", exec_time));
1050 }
1051 if let Some(seq_scans) = metrics.get("sequentialScans").and_then(|v| v.as_u64())
1052 && seq_scans > 0
1053 {
1054 parts.push(format!("Found {seq_scans} sequential scan(s)."));
1055 }
1056 }
1057
1058 let candidate_count = suggestions
1059 .get("high_seq_scan_tables")
1060 .and_then(|v| v.as_array())
1061 .map(|a| a.len())
1062 .unwrap_or(0)
1063 + suggestions
1064 .get("unindexed_fk_columns")
1065 .and_then(|v| v.as_array())
1066 .map(|a| a.len())
1067 .unwrap_or(0);
1068
1069 if candidate_count > 0 {
1070 parts.push(format!(
1071 "{candidate_count} index recommendation(s) identified."
1072 ));
1073 } else {
1074 parts.push("No explicit index candidate recommendations generated.".into());
1075 }
1076
1077 if parts.is_empty() {
1078 "Auto-tune evaluation complete. Inspect execution plan and index recommendations."
1079 .into()
1080 } else {
1081 parts.join(" ")
1082 }
1083 }
1084
1085 async fn auto_tune_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1086 let sql = args
1087 .get("sql")
1088 .and_then(|v| v.as_str())
1089 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
1090
1091 let deep_plan = self
1092 .deep_plan_analysis(&json!({ "sql": sql, "analyze": true }))
1093 .await?;
1094
1095 let suggestions_res = self.suggest_indexes(&json!({ "sql": sql })).await;
1096 let (suggestions_data, suggestions_error) = match suggestions_res {
1097 Ok(outcome) => (outcome.structured.unwrap_or(json!([])), None),
1098 Err(e) => (json!([]), Some(e.to_string())),
1099 };
1100
1101 let summary_text = Self::build_tuning_summary(&deep_plan.structured, &suggestions_data);
1102
1103 let mut payload = json!({
1104 "target_query": sql,
1105 "deep_plan_analysis": deep_plan.structured,
1106 "index_suggestions": suggestions_data,
1107 "tuning_summary": summary_text,
1108 });
1109
1110 if let Some(err) = suggestions_error {
1111 payload["suggestions_error"] = json!(err);
1112 }
1113
1114 Ok(ToolOutcome::ok_json(payload))
1115 }
1116
1117 async fn check_ddl_safety_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1118 let ddl = args
1119 .get("ddl")
1120 .and_then(|v| v.as_str())
1121 .ok_or_else(|| ToolError::InvalidArgs("ddl is required".into()))?;
1122
1123 let report = crate::dba_guard::analyze_ddl_safety(ddl);
1124 Ok(ToolOutcome::ok_json(report))
1125 }
1126
1127 async fn rebuild_index_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1128 let store = self
1129 .index_store()
1130 .ok_or_else(|| ToolError::Execution("Index store unavailable".into()))?;
1131 let (connection_id, database) = self.session.active_context().await;
1132 let depth_str = args
1133 .get("depth")
1134 .and_then(|v| v.as_str())
1135 .unwrap_or("structure");
1136 let depth: BuildDepth = match depth_str.to_lowercase().as_str() {
1137 "profiles" | "full" => BuildDepth::Profiles,
1138 _ => BuildDepth::Structure,
1139 };
1140
1141 let req = BuildRequest {
1142 connection_id: connection_id.clone(),
1143 database: database.clone(),
1144 scope: IndexScope {
1145 included_schemas: vec![],
1146 excluded_objects: vec![],
1147 pii_excluded_columns: vec![],
1148 },
1149 depth,
1150 build_mode: BuildMode::Guided,
1151 environment: "development".into(),
1152 embeddings: self.use_semantic,
1153 };
1154
1155 let client = self.session.checkout().await?;
1156 let db = PgCatalogDb::new(&client);
1157 let manifest = build_index(store, &db, &req, None, None, self.embedder.as_deref())
1158 .await
1159 .map_err(|e| ToolError::Execution(format!("Index build failed: {e}")))?;
1160 self.session.clear_index_stale(&connection_id, &database);
1161
1162 Ok(ToolOutcome::ok_json(json!({
1163 "status": "completed",
1164 "connection_id": connection_id,
1165 "database": database,
1166 "schema_fingerprint": manifest.schema_fingerprint,
1167 "counts": manifest.counts,
1168 "build_ms": manifest.stats.build_ms,
1169 })))
1170 }
1171
1172 async fn refresh_index_tool(&self, _args: &Value) -> Result<ToolOutcome, ToolError> {
1173 let store = self
1174 .index_store()
1175 .ok_or_else(|| ToolError::Execution("Index store unavailable".into()))?;
1176 let (connection_id, database) = self.session.active_context().await;
1177 let base = store.base_dir(&connection_id, &database);
1178 let manifest = store.read_manifest(&base)?.ok_or_else(|| {
1179 ToolError::Execution(
1180 "No existing index manifest to refresh — call 'rebuild_index'.".into(),
1181 )
1182 })?;
1183
1184 let req = BuildRequest {
1185 connection_id: connection_id.clone(),
1186 database: database.clone(),
1187 scope: manifest.scope,
1188 depth: manifest.build_depth,
1189 build_mode: manifest.build_mode,
1190 environment: manifest.environment,
1191 embeddings: self.use_semantic,
1192 };
1193
1194 let client = self.session.checkout().await?;
1195 let db = PgCatalogDb::new(&client);
1196 let new_manifest = build_index(store, &db, &req, None, None, self.embedder.as_deref())
1197 .await
1198 .map_err(|e| ToolError::Execution(format!("Index refresh failed: {e}")))?;
1199 self.session.clear_index_stale(&connection_id, &database);
1200
1201 Ok(ToolOutcome::ok_json(json!({
1202 "status": "refreshed",
1203 "connection_id": connection_id,
1204 "database": database,
1205 "schema_fingerprint": new_manifest.schema_fingerprint,
1206 "counts": new_manifest.counts,
1207 "build_ms": new_manifest.stats.build_ms,
1208 })))
1209 }
1210
1211 async fn run_doctor_tool(&self) -> Result<ToolOutcome, ToolError> {
1212 let (connection_id, database) = self.session.active_context().await;
1213 let client = self.session.checkout().await?;
1214
1215 let version: String = client
1216 .query_one("SELECT version()", &[])
1217 .await
1218 .map_err(|e| ToolError::Execution(e.to_string()))?
1219 .get(0);
1220
1221 let is_super: String = client
1222 .query_one("SELECT current_setting('is_superuser')", &[])
1223 .await
1224 .map_err(|e| ToolError::Execution(e.to_string()))?
1225 .get(0);
1226 let is_superuser = is_super.eq_ignore_ascii_case("on");
1227
1228 let ro: String = client
1229 .query_one("SHOW default_transaction_read_only", &[])
1230 .await
1231 .map_err(|e| ToolError::Execution(e.to_string()))?
1232 .get(0);
1233
1234 let timeout: String = client
1235 .query_one("SHOW statement_timeout", &[])
1236 .await
1237 .map_err(|e| ToolError::Execution(e.to_string()))?
1238 .get(0);
1239
1240 let pgs_present: bool = match client
1241 .query_one(
1242 "SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements')",
1243 &[],
1244 )
1245 .await
1246 {
1247 Ok(row) => row.get(0),
1248 Err(_) => false,
1249 };
1250
1251 let index_status = if let Some(store) = self.index_store() {
1252 let base = store.base_dir(&connection_id, &database);
1253 match store.read_manifest(&base) {
1254 Ok(Some(m)) => json!({
1255 "present": true,
1256 "indexed_at": m.indexed_at,
1257 "fingerprint": m.schema_fingerprint,
1258 "tables": m.counts.tables,
1259 }),
1260 _ => json!({ "present": false }),
1261 }
1262 } else {
1263 json!({ "present": false, "reason": "no_index_store" })
1264 };
1265
1266 let recent_errors = read_recent_log_errors();
1267
1268 Ok(ToolOutcome::ok_json(json!({
1269 "status": "ok",
1270 "connection_id": connection_id,
1271 "database": database,
1272 "version": version.split(',').next().unwrap_or(&version),
1273 "access_mode": format!("{:?}", self.session.access_mode()),
1274 "superuser": is_superuser,
1275 "read_only": ro,
1276 "statement_timeout": timeout,
1277 "pg_stat_statements": pgs_present,
1278 "index": index_status,
1279 "recent_errors": recent_errors,
1280 })))
1281 }
1282
1283 fn register_profile_in_session(
1284 &self,
1285 name: &str,
1286 profile: &nexql_conn::ProfileConfig,
1287 ) -> Result<(), ToolError> {
1288 self.session.register_profile(
1289 name,
1290 profile,
1291 self.session.access_mode(),
1292 self.session.caps(),
1293 )
1294 }
1295
1296 fn route_password_to_keyring(
1298 profile_name: &str,
1299 password: Option<&str>,
1300 ) -> Result<nexql_conn::RoutedCredential, ToolError> {
1301 nexql_conn::route_password_to_keyring(profile_name, password)
1302 .map_err(|e| ToolError::Execution(e.to_string()))
1303 }
1304
1305 async fn setup_connection_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1306 let profile_name = args
1307 .get("name")
1308 .and_then(|v| v.as_str())
1309 .unwrap_or("default");
1310
1311 let candidates = crate::detect::ConnectionDetector::detect_all(None);
1312
1313 let url = args.get("url").and_then(|v| v.as_str());
1314 let host = args.get("host").and_then(|v| v.as_str());
1315 let port = args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16);
1316 let dbname = args.get("dbname").and_then(|v| v.as_str());
1317 let user = args.get("user").and_then(|v| v.as_str());
1318 let password = args.get("password").and_then(|v| v.as_str());
1319 let sslmode = args.get("sslmode").and_then(|v| v.as_str());
1320
1321 let best_cand = candidates
1322 .iter()
1323 .find(|c| c.is_complete)
1324 .or_else(|| candidates.first());
1325
1326 let res_host = host.or_else(|| best_cand.and_then(|c| c.host.as_deref()));
1327 let res_port = port.or_else(|| best_cand.and_then(|c| c.port));
1328 let res_dbname = dbname.or_else(|| best_cand.and_then(|c| c.dbname.as_deref()));
1329 let res_user = user.or_else(|| best_cand.and_then(|c| c.user.as_deref()));
1330 let res_password = password.or_else(|| best_cand.and_then(|c| c.password.as_deref()));
1331 let res_url = url.or_else(|| best_cand.and_then(|c| c.url.as_deref()));
1332 let res_sslmode = sslmode.or_else(|| best_cand.and_then(|c| c.sslmode.as_deref()));
1333
1334 if res_url.is_none() && (res_host.is_none() || res_dbname.is_none() || res_user.is_none()) {
1335 let missing: Vec<&str> = vec![
1336 if res_host.is_none() {
1337 Some("host")
1338 } else {
1339 None
1340 },
1341 if res_dbname.is_none() {
1342 Some("dbname")
1343 } else {
1344 None
1345 },
1346 if res_user.is_none() {
1347 Some("user")
1348 } else {
1349 None
1350 },
1351 ]
1352 .into_iter()
1353 .flatten()
1354 .collect();
1355
1356 return Ok(ToolOutcome::ok_json(json!({
1357 "status": "needs_input",
1358 "message": "Insufficient connection details. Please supply missing fields.",
1359 "detectedCandidates": candidates.iter().map(|c| c.redacted_json()).collect::<Vec<_>>(),
1360 "missingFields": missing
1361 })));
1362 }
1363
1364 let params = nexql_conn::ConnectionParams {
1365 url: res_url.map(String::from),
1366 host: res_host.map(String::from),
1367 port: res_port,
1368 dbname: res_dbname.map(String::from),
1369 user: res_user.map(String::from),
1370 password: res_password.map(String::from),
1371 sslmode: res_sslmode.map(String::from),
1372 ..Default::default()
1373 };
1374
1375 match nexql_conn::test_connection(¶ms).await {
1376 Ok(report) => {
1377 let routed =
1378 Self::route_password_to_keyring(profile_name, params.password.as_deref())?;
1379 let p_config = nexql_conn::ProfileConfig {
1380 url: params.url.clone(),
1381 host: params.host.clone(),
1382 port: params.port,
1383 dbname: params.dbname.clone(),
1384 user: params.user.clone(),
1385 password: routed.password,
1386 sslmode: params.sslmode.clone(),
1387 credential_provider: routed.credential_provider,
1388 password_file: routed.password_file,
1389 ..Default::default()
1390 };
1391
1392 let path = nexql_conn::ConfigFile::default_path().ok_or_else(|| {
1393 ToolError::Execution("Could not resolve config directory".into())
1394 })?;
1395 let mut cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1396 .map(|(c, _)| c)
1397 .unwrap_or_default();
1398 cfg.upsert_profile(profile_name, p_config.clone());
1399 let backup = cfg
1400 .save(&path)
1401 .map_err(|e| ToolError::Execution(e.to_string()))?;
1402 self.register_profile_in_session(profile_name, &p_config)?;
1403
1404 Ok(ToolOutcome::ok_json(json!({
1405 "status": "configured",
1406 "profileName": profile_name,
1407 "serverVersion": report.server_version,
1408 "isSuperuser": report.is_superuser,
1409 "latencyMs": report.latency.as_millis(),
1410 "configPath": path.to_string_lossy().to_string(),
1411 "backup": backup.map(|b| b.to_string_lossy().to_string()),
1412 "sessionReloaded": true,
1413 })))
1414 }
1415 Err(e) => Ok(ToolOutcome::ok_json(json!({
1416 "status": "failed",
1417 "error": e.to_string(),
1418 "detectedCandidates": candidates.iter().map(|c| c.redacted_json()).collect::<Vec<_>>()
1419 }))),
1420 }
1421 }
1422
1423 async fn save_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1424 let name = args
1425 .get("name")
1426 .and_then(|v| v.as_str())
1427 .ok_or_else(|| ToolError::InvalidArgs("name parameter is required".into()))?;
1428
1429 if let Some(mode_str) = args.get("access_mode").and_then(|v| v.as_str()) {
1435 let mode: nexql_policy::AccessMode = mode_str.parse().map_err(|_| {
1436 ToolError::InvalidArgs(format!(
1437 "invalid access_mode \"{mode_str}\" — expected read, write, or admin"
1438 ))
1439 })?;
1440 let confirmed = args
1441 .get("confirm_elevated_access")
1442 .and_then(|v| v.as_bool())
1443 .unwrap_or(false);
1444 if mode.allows_writes() && !confirmed {
1445 return Err(ToolError::InvalidArgs(format!(
1446 "refusing to save profile \"{name}\" with access_mode \"{mode_str}\" — pass confirm_elevated_access: true to override"
1447 )));
1448 }
1449 }
1450
1451 let routed =
1452 Self::route_password_to_keyring(name, args.get("password").and_then(|v| v.as_str()))?;
1453
1454 let p_config = nexql_conn::ProfileConfig {
1455 url: args.get("url").and_then(|v| v.as_str()).map(String::from),
1456 host: args.get("host").and_then(|v| v.as_str()).map(String::from),
1457 port: args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16),
1458 dbname: args
1459 .get("dbname")
1460 .and_then(|v| v.as_str())
1461 .map(String::from),
1462 user: args.get("user").and_then(|v| v.as_str()).map(String::from),
1463 password: routed.password,
1464 sslmode: args
1465 .get("sslmode")
1466 .and_then(|v| v.as_str())
1467 .map(String::from),
1468 access_mode: args
1469 .get("access_mode")
1470 .and_then(|v| v.as_str())
1471 .map(String::from),
1472 max_rows: args
1473 .get("max_rows")
1474 .and_then(|v| v.as_u64())
1475 .map(|n| n as u32),
1476 credential_provider: routed.credential_provider,
1477 password_file: routed.password_file,
1478 ..Default::default()
1479 };
1480
1481 let path = nexql_conn::ConfigFile::default_path()
1482 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1483
1484 let mut cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1485 .map(|(c, _)| c)
1486 .unwrap_or_default();
1487 cfg.upsert_profile(name, p_config.clone());
1488 let backup = cfg
1489 .save(&path)
1490 .map_err(|e| ToolError::Execution(e.to_string()))?;
1491 self.register_profile_in_session(name, &p_config)?;
1492
1493 Ok(ToolOutcome::ok_json(json!({
1494 "status": "saved",
1495 "profile": name,
1496 "configPath": path.to_string_lossy().to_string(),
1497 "backup": backup.map(|b| b.to_string_lossy().to_string()),
1498 "sessionReloaded": true,
1499 })))
1500 }
1501
1502 async fn test_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1503 let name = args.get("name").and_then(|v| v.as_str());
1504
1505 let params = if let Some(pname) = name {
1506 let conn = self
1507 .session
1508 .connections()
1509 .into_iter()
1510 .find(|c| c.id == pname)
1511 .ok_or_else(|| ToolError::InvalidArgs(format!("Profile '{pname}' not found")))?;
1512 conn.params.clone()
1513 } else {
1514 nexql_conn::ConnectionParams {
1515 url: args.get("url").and_then(|v| v.as_str()).map(String::from),
1516 host: args.get("host").and_then(|v| v.as_str()).map(String::from),
1517 port: args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16),
1518 dbname: args
1519 .get("dbname")
1520 .and_then(|v| v.as_str())
1521 .map(String::from),
1522 user: args.get("user").and_then(|v| v.as_str()).map(String::from),
1523 password: args
1524 .get("password")
1525 .and_then(|v| v.as_str())
1526 .map(String::from),
1527 sslmode: args
1528 .get("sslmode")
1529 .and_then(|v| v.as_str())
1530 .map(String::from),
1531 ..Default::default()
1532 }
1533 };
1534
1535 match nexql_conn::test_connection(¶ms).await {
1536 Ok(report) => Ok(ToolOutcome::ok_json(json!({
1537 "success": true,
1538 "serverVersion": report.server_version,
1539 "isSuperuser": report.is_superuser,
1540 "latencyMs": report.latency.as_millis()
1541 }))),
1542 Err(e) => Ok(ToolOutcome::ok_json(json!({
1543 "success": false,
1544 "error": e.to_string()
1545 }))),
1546 }
1547 }
1548
1549 async fn export_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1550 let format = args
1551 .get("format")
1552 .and_then(|v| v.as_str())
1553 .unwrap_or("full");
1554 let path = nexql_conn::ConfigFile::default_path()
1555 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1556 let cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1557 .map(|(c, _)| c)
1558 .unwrap_or_default();
1559
1560 if format == "project" {
1561 let proj = cfg.export_shareable();
1562 let toml_str =
1563 toml::to_string_pretty(&proj).map_err(|e| ToolError::Execution(e.to_string()))?;
1564 Ok(ToolOutcome::ok_json(json!({
1565 "format": "project",
1566 "filename": ".nexql/config.toml",
1567 "description": "Project policy overlay (no credentials). Use format=full for shareable connection profiles.",
1568 "content": toml_str,
1569 })))
1570 } else {
1571 let sanitized = cfg.export_full_sanitized();
1572 let toml_str = sanitized
1573 .to_toml_string()
1574 .map_err(|e| ToolError::Execution(e.to_string()))?;
1575 Ok(ToolOutcome::ok_json(json!({
1576 "format": "full",
1577 "description": "Full user config with passwords and secrets stripped — suitable for team sharing.",
1578 "profileCount": sanitized.profiles.len(),
1579 "content": toml_str,
1580 })))
1581 }
1582 }
1583
1584 async fn import_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1585 let content = if let Some(c) = args.get("content").and_then(|v| v.as_str()) {
1586 c.to_string()
1587 } else if let Some(p) = args.get("path").and_then(|v| v.as_str()) {
1588 std::fs::read_to_string(p)
1589 .map_err(|e| ToolError::Execution(format!("failed to read file {p}: {e}")))?
1590 } else {
1591 return Err(ToolError::Execution(
1592 "either 'content' or 'path' must be specified".into(),
1593 ));
1594 };
1595
1596 let path = nexql_conn::ConfigFile::default_path()
1597 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1598 let mut cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1599 .map(|(c, _)| c)
1600 .unwrap_or_default();
1601
1602 let imported: nexql_conn::ConfigFile = toml::from_str(&content)
1603 .map_err(|e| ToolError::Execution(format!("failed to parse TOML content: {e}")))?;
1604
1605 let mut count = 0;
1606 let mut imported_names: Vec<String> = Vec::new();
1607 for (name, prof) in imported.profiles {
1608 let prepared = nexql_conn::prepare_profile_for_persist(&name, prof)
1609 .map_err(|e| ToolError::Execution(e.to_string()))?;
1610 cfg.upsert_profile(name.clone(), prepared.clone());
1611 self.register_profile_in_session(&name, &prepared)?;
1612 imported_names.push(name);
1613 count += 1;
1614 }
1615 if imported.default_profile.is_some() {
1616 cfg.default_profile = imported.default_profile;
1617 }
1618
1619 let backup = cfg
1620 .save(&path)
1621 .map_err(|e| ToolError::Execution(e.to_string()))?;
1622
1623 Ok(ToolOutcome::ok_json(json!({
1624 "status": "imported",
1625 "imported_profiles": count,
1626 "profiles": imported_names,
1627 "configPath": path.to_string_lossy().to_string(),
1628 "backup": backup.map(|b| b.to_string_lossy().to_string()),
1629 "sessionReloaded": true,
1630 })))
1631 }
1632
1633 async fn ensure_index_warm_for(&self, ctx: &ScopedContext) -> Result<(), ToolError> {
1635 let store = self
1636 .index_store()
1637 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1638 let base = store.base_dir(&ctx.connection_id, &ctx.database);
1639 if store.read_manifest(&base)?.is_some() {
1640 return Ok(());
1641 }
1642 let req = BuildRequest {
1643 connection_id: ctx.connection_id.clone(),
1644 database: ctx.database.clone(),
1645 scope: IndexScope {
1646 included_schemas: vec![],
1647 excluded_objects: vec![],
1648 pii_excluded_columns: vec![],
1649 },
1650 depth: BuildDepth::Structure,
1651 build_mode: BuildMode::Guided,
1652 environment: "development".into(),
1653 embeddings: self.use_semantic,
1654 };
1655 let (client, _) = self
1656 .session
1657 .checkout_for(CheckoutTarget::Scoped(ctx))
1658 .await
1659 .map_err(|_| {
1660 ToolError::Execution(format!(
1661 "No schema index for database \"{}\" — call the 'rebuild_index' tool to build an index.",
1662 ctx.database
1663 ))
1664 })?;
1665 let db = PgCatalogDb::new(&client);
1666 build_index(store, &db, &req, None, None, self.embedder.as_deref())
1667 .await
1668 .map_err(|e| ToolError::Execution(format!("Automatic index build failed: {e}")))?;
1669 self.session
1670 .clear_index_stale(&ctx.connection_id, &ctx.database);
1671 Ok(())
1672 }
1673
1674 async fn ensure_index_warm(&self) -> Result<(), ToolError> {
1675 let (connection_id, database) = self.session.active_context().await;
1676 self.ensure_index_warm_for(&ScopedContext {
1677 connection_id,
1678 database,
1679 })
1680 .await
1681 }
1682
1683 async fn index_service_for(
1684 &self,
1685 ctx: &ScopedContext,
1686 ) -> Result<(&IndexStore, String, String), ToolError> {
1687 let store = self
1688 .index_store()
1689 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1690 let base = store.base_dir(&ctx.connection_id, &ctx.database);
1691 if store.read_manifest(&base)?.is_none() {
1692 self.ensure_index_warm_for(ctx).await?;
1693 }
1694 Ok((store, ctx.connection_id.clone(), ctx.database.clone()))
1695 }
1696
1697 async fn index_service(&self) -> Result<(&IndexStore, String, String), ToolError> {
1698 let store = self
1699 .index_store()
1700 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1701 let (connection_id, database) = self.session.active_context().await;
1702 let base = store.base_dir(&connection_id, &database);
1703 if store.read_manifest(&base)?.is_none() {
1704 self.ensure_index_warm().await?;
1705 }
1706 Ok((store, connection_id, database))
1707 }
1708
1709 async fn collect_index_refs(&self) -> Vec<String> {
1710 if let Ok((store, connection_id, database)) = self.index_service().await {
1711 let base = store.base_dir(&connection_id, &database);
1712 if let Ok(Some(manifest)) = store.read_manifest(&base) {
1713 let mut refs = Vec::new();
1714 for shard in &manifest.shards {
1715 if let Ok(Some(entries)) = store.read_shard_entries(&base, &shard.file) {
1716 refs.extend(entries.keys().cloned());
1717 }
1718 }
1719 return refs;
1720 }
1721 }
1722 Vec::new()
1723 }
1724
1725 async fn enrich_query_error(&self, pg_err: &tokio_postgres::Error) -> String {
1726 let base = nexql_conn::format_postgres_error(pg_err);
1727 let refs = self.collect_index_refs().await;
1728 sql::enhance_sql_error(&base, &refs)
1729 }
1730
1731 fn ref_resolution_error(ref_: &str, resolution: &RefResolution) -> Option<ToolError> {
1735 match resolution {
1736 RefResolution::Resolved(_) => None,
1737 RefResolution::Ambiguous(candidates) => Some(ToolError::InvalidArgs(format!(
1738 "ambiguous relation \"{ref_}\": {}",
1739 candidates.join(", ")
1740 ))),
1741 RefResolution::Unknown {
1742 suggestion: Some(s),
1743 } => Some(ToolError::InvalidArgs(format!(
1744 "unknown relation \"{ref_}\" — did you mean \"{s}\"?"
1745 ))),
1746 RefResolution::Unknown { suggestion: None } => Some(ToolError::InvalidArgs(format!(
1747 "unknown relation \"{ref_}\" — call search_schema to find valid refs."
1748 ))),
1749 }
1750 }
1751
1752 fn resolve_indexed_ref_strict(
1756 svc: &IndexQueryService<'_>,
1757 ref_: &str,
1758 ) -> Result<String, ToolError> {
1759 let resolution = svc.resolve_ref(ref_)?;
1760 if let Some(err) = Self::ref_resolution_error(ref_, &resolution) {
1761 return Err(err);
1762 }
1763 match resolution {
1764 RefResolution::Resolved(r) => Ok(r),
1765 _ => unreachable!("ref_resolution_error covers every non-Resolved case"),
1766 }
1767 }
1768
1769 fn resolve_indexed_ref_soft(
1774 svc: &IndexQueryService<'_>,
1775 ref_: &str,
1776 ) -> Result<String, ToolError> {
1777 match svc.resolve_ref(ref_)? {
1778 RefResolution::Resolved(r) => Ok(r),
1779 RefResolution::Ambiguous(candidates) => Err(ToolError::InvalidArgs(format!(
1780 "ambiguous relation \"{ref_}\": {}",
1781 candidates.join(", ")
1782 ))),
1783 RefResolution::Unknown { .. } => Ok(ref_.to_owned()),
1784 }
1785 }
1786
1787 async fn resolve_ref_best_effort(&self, ref_: &str) -> Result<String, ToolError> {
1793 let Some(store) = self.index_store() else {
1794 return Ok(ref_.to_owned());
1795 };
1796 let (connection_id, database) = self.session.active_context().await;
1797 let base = store.base_dir(&connection_id, &database);
1798 if store.read_manifest(&base)?.is_none() {
1799 return Ok(ref_.to_owned());
1800 }
1801 let svc = IndexQueryService::new(store, &connection_id, &database);
1802 Self::resolve_indexed_ref_soft(&svc, ref_)
1803 }
1804
1805 async fn search_schema(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1806 let query = args
1807 .get("query")
1808 .and_then(|v| v.as_str())
1809 .unwrap_or("")
1810 .trim();
1811 if query.is_empty() {
1812 return Ok(ToolOutcome::ok_json(json!([])));
1813 }
1814 let (store, connection_id, database) = self.index_service().await?;
1815 let svc = IndexQueryService::new(store, &connection_id, &database);
1816 let filter = self.query_filter();
1817 let hits = svc.search_schema(
1818 query,
1819 SEARCH_SCHEMA_LIMIT,
1820 Some(&filter),
1821 SearchOptions {
1822 use_semantic: self.use_semantic,
1823 embedder: self.embedder.as_deref(),
1824 },
1825 )?;
1826 let rows: Vec<Value> = hits
1827 .into_iter()
1828 .map(|h| {
1829 json!({
1830 "ref": h.ref_,
1831 "score": h.score,
1832 "kind": h.kind,
1833 })
1834 })
1835 .collect();
1836 Ok(ToolOutcome::ok_json(json!(rows)))
1837 }
1838
1839 async fn inspect_or_search(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1840 let query = args
1841 .get("query")
1842 .and_then(|v| v.as_str())
1843 .unwrap_or("")
1844 .trim();
1845 if query.is_empty() {
1846 return Err(ToolError::InvalidArgs("query is required".into()));
1847 }
1848 let include_columns = args
1849 .get("include_columns")
1850 .and_then(|v| v.as_bool())
1851 .unwrap_or(true);
1852 let limit_objects = args
1853 .get("limit_objects")
1854 .and_then(|v| v.as_u64())
1855 .map(|n| n as usize)
1856 .unwrap_or(3)
1857 .clamp(1, 20);
1858
1859 let (store, connection_id, database) = self.index_service().await?;
1860 let svc = IndexQueryService::new(store, &connection_id, &database);
1861 let filter = self.query_filter();
1862 let hits = svc.search_schema(
1863 query,
1864 limit_objects,
1865 Some(&filter),
1866 SearchOptions {
1867 use_semantic: self.use_semantic,
1868 embedder: self.embedder.as_deref(),
1869 },
1870 )?;
1871
1872 let mut matches = Vec::with_capacity(hits.len());
1873 for hit in hits {
1874 let entry = svc.describe_object(&hit.ref_, Some(&filter))?;
1875 let mut obj = json!({
1876 "ref": hit.ref_,
1877 "score": hit.score,
1878 "kind": hit.kind,
1879 "row_estimate": entry.row_estimate.round() as i64,
1880 "primary_key": entry.primary_key,
1881 });
1882 if include_columns {
1883 let fk_cols: std::collections::HashSet<String> = entry
1884 .foreign_keys
1885 .as_ref()
1886 .map(|fks| {
1887 fks.iter()
1888 .flat_map(|fk| fk.columns.clone())
1889 .collect()
1890 })
1891 .unwrap_or_default();
1892 let columns: Vec<Value> = entry
1893 .columns
1894 .iter()
1895 .map(|c| {
1896 json!({
1897 "name": c.name,
1898 "type": c.type_name,
1899 "is_pk": c.is_pk.unwrap_or(false),
1900 "is_fk": fk_cols.contains(&c.name),
1901 "not_null": c.not_null,
1902 })
1903 })
1904 .collect();
1905 obj["columns"] = json!(columns);
1906 }
1907 matches.push(obj);
1908 }
1909
1910 Ok(ToolOutcome::ok_json(json!({
1911 "query": query,
1912 "matches": matches,
1913 })))
1914 }
1915
1916 async fn search_all_databases(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1917 let query = args
1918 .get("query")
1919 .and_then(|v| v.as_str())
1920 .unwrap_or("")
1921 .trim();
1922 if query.is_empty() {
1923 return Err(ToolError::InvalidArgs("query is required".into()));
1924 }
1925 let limit_per_db = args
1926 .get("limit_per_database")
1927 .and_then(|v| v.as_u64())
1928 .map(|n| n as usize)
1929 .unwrap_or(3)
1930 .clamp(1, 10);
1931 let limit_pairs = args
1932 .get("limit_connections")
1933 .and_then(|v| v.as_u64())
1934 .map(|n| n as usize)
1935 .unwrap_or(20)
1936 .clamp(1, 50);
1937
1938 let Some(store) = self.index_store() else {
1939 return Err(ToolError::Execution(NO_INDEX_HINT.into()));
1940 };
1941
1942 let indexed = store.list_indexed_databases().unwrap_or_default();
1943 let connections = self.session.connections();
1944 let filter = self.query_filter();
1945 let mut hits: Vec<Value> = Vec::new();
1946 let mut searched = 0usize;
1947
1948 for (connection_id, database) in indexed {
1949 if searched >= limit_pairs {
1950 break;
1951 }
1952 if !connections.iter().any(|c| c.id == connection_id) {
1953 continue;
1954 }
1955 searched += 1;
1956 let svc = IndexQueryService::new(store, &connection_id, &database);
1957 if let Ok(results) = svc.search_schema(
1958 query,
1959 limit_per_db,
1960 Some(&filter),
1961 SearchOptions {
1962 use_semantic: self.use_semantic,
1963 embedder: self.embedder.as_deref(),
1964 },
1965 ) {
1966 for hit in results {
1967 hits.push(json!({
1968 "connectionId": connection_id,
1969 "database": database,
1970 "ref": hit.ref_,
1971 "score": hit.score,
1972 "kind": hit.kind,
1973 }));
1974 }
1975 }
1976 }
1977
1978 hits.sort_by(|a, b| {
1979 let sa = a.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
1980 let sb = b.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
1981 sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)
1982 });
1983
1984 Ok(ToolOutcome::ok_json(json!({
1985 "query": query,
1986 "hits": hits,
1987 "searched_pairs": searched,
1988 })))
1989 }
1990
1991 async fn describe_object(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1992 let ref_ = args
1993 .get("ref")
1994 .and_then(|v| v.as_str())
1995 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1996 let scope = self.execution_scope_from_args(args).await?;
1997 let resolve_refs = args
1998 .get("resolve_refs")
1999 .and_then(|v| v.as_bool())
2000 .unwrap_or(false);
2001 let resolve_limit = args
2002 .get("resolve_refs_limit")
2003 .and_then(|v| v.as_u64())
2004 .map(|n| n as usize)
2005 .unwrap_or(DEFAULT_RESOLVE_REFS_LIMIT);
2006 let (store, connection_id, database) = self.index_service_for(&scope.ctx).await?;
2007 let svc = IndexQueryService::new(store, &connection_id, &database);
2008 let resolved = Self::resolve_indexed_ref_soft(&svc, ref_)?;
2009 let filter = self.query_filter_for(&scope.ctx.connection_id);
2010 let entry = svc.describe_object(&resolved, Some(&filter))?;
2011 let mut value =
2012 serde_json::to_value(&entry).map_err(|e| ToolError::Execution(e.to_string()))?;
2013 let client = if resolve_refs {
2014 Some(
2015 self.session
2016 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2017 .await?
2018 .0,
2019 )
2020 } else {
2021 None
2022 };
2023 value = resolve::enrich_describe_object_with_store(
2024 value,
2025 &entry,
2026 store,
2027 &svc,
2028 resolve_refs,
2029 resolve_limit,
2030 client.as_ref(),
2031 )
2032 .await?;
2033 Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(value)))
2034 }
2035
2036 async fn get_join_path(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2037 let a = args
2038 .get("a")
2039 .and_then(|v| v.as_str())
2040 .ok_or_else(|| ToolError::InvalidArgs("a is required".into()))?;
2041 let b = args
2042 .get("b")
2043 .and_then(|v| v.as_str())
2044 .ok_or_else(|| ToolError::InvalidArgs("b is required".into()))?;
2045 let (store, connection_id, database) = self.index_service().await?;
2046 let svc = IndexQueryService::new(store, &connection_id, &database);
2047 let resolved_a = Self::resolve_indexed_ref_strict(&svc, a)?;
2052 let resolved_b = Self::resolve_indexed_ref_strict(&svc, b)?;
2053 let path = svc.get_join_path(&resolved_a, &resolved_b)?;
2054 let path_value =
2055 serde_json::to_value(path).map_err(|e| ToolError::Execution(e.to_string()))?;
2056 Ok(ToolOutcome::ok_json(json!({
2057 "path": path_value,
2058 "resolved_a": resolved_a,
2059 "resolved_b": resolved_b,
2060 })))
2061 }
2062
2063 async fn sample_values(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2064 let ref_ = args
2065 .get("ref")
2066 .and_then(|v| v.as_str())
2067 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2068 let col = args
2069 .get("col")
2070 .and_then(|v| v.as_str())
2071 .ok_or_else(|| ToolError::InvalidArgs("col is required".into()))?;
2072 let (store, connection_id, database) = self.index_service().await?;
2073 let svc = IndexQueryService::new(store, &connection_id, &database);
2074 let resolved = Self::resolve_indexed_ref_soft(&svc, ref_)?;
2075 let ref_ = resolved.as_str();
2076 let filter = self.query_filter();
2077 let result = svc.sample_values(ref_, col, Some(&filter), None)?;
2078
2079 let mut values = result.values;
2080 let mut message = result.message;
2081
2082 if values.is_empty()
2083 && let Ok(client) = self.session.checkout().await
2084 {
2085 let parts: Vec<&str> = ref_.split('.').collect();
2086 let (schema, table) = match parts.as_slice() {
2087 [s, t] => (*s, *t),
2088 _ => ("public", ref_),
2089 };
2090 let safe_schema = schema.replace('"', "\"\"");
2091 let safe_table = table.replace('"', "\"\"");
2092 let safe_col = col.replace('"', "\"\"");
2093 let query = format!(
2094 "SELECT DISTINCT \"{safe_col}\"::text FROM \"{safe_schema}\".\"{safe_table}\" WHERE \"{safe_col}\" IS NOT NULL LIMIT 20"
2095 );
2096 if let Ok(rows) = client.query(&query, &[]).await {
2097 let sampled: Vec<String> = rows
2098 .iter()
2099 .filter_map(|r| r.get::<_, Option<String>>(0))
2100 .collect();
2101 if !sampled.is_empty() {
2102 values = sampled;
2103 message = None;
2104 }
2105 }
2106 }
2107
2108 let mut payload = json!({ "values": values });
2109 if let Some(msg) = message {
2110 payload["message"] = json!(msg);
2111 }
2112 Ok(ToolOutcome::ok_json(payload))
2113 }
2114
2115 fn list_connections(&self) -> ToolOutcome {
2116 let rows: Vec<Value> = self
2117 .session
2118 .connections()
2119 .iter()
2120 .map(|c| {
2121 json!({
2122 "id": c.id,
2123 "name": c.name,
2124 "host": c.host,
2125 "port": c.port,
2126 "database": c.database,
2127 })
2128 })
2129 .collect();
2130 ToolOutcome::ok_json(json!(rows))
2131 }
2132
2133 async fn list_databases(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2134 let connection_id = args
2135 .get("connectionId")
2136 .and_then(|v| v.as_str())
2137 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
2138 let conn = self
2139 .session
2140 .connections()
2141 .into_iter()
2142 .find(|c| c.id == connection_id)
2143 .ok_or_else(|| {
2144 ToolError::Execution(format!(
2145 "Connection not found for ID: {connection_id} — call list_connections"
2146 ))
2147 })?;
2148 let client = {
2150 if self.session.active_context().await.0 == connection_id {
2152 self.session.checkout().await?
2153 } else {
2154 let pool_opts = self.session.pool_opts();
2155 let pool = nexql_conn::create_pool(&conn.params, &pool_opts).await?;
2156 nexql_conn::checkout_guarded(&pool, &pool_opts).await?
2157 }
2158 };
2159 let rows = client
2160 .query(
2161 "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname",
2162 &[],
2163 )
2164 .await?;
2165 let names: Vec<String> = rows.iter().map(|r| r.get(0)).collect();
2166 Ok(ToolOutcome::ok_json(json!(names)))
2167 }
2168
2169 async fn list_schemas(&self) -> Result<ToolOutcome, ToolError> {
2170 let client = self.session.checkout().await?;
2171 let rows = client
2172 .query(
2173 r#"
2174 SELECT nspname AS schema_name
2175 FROM pg_namespace
2176 WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
2177 AND nspname NOT LIKE 'pg_%'
2178 ORDER BY nspname
2179 "#,
2180 &[],
2181 )
2182 .await?;
2183 let out: Vec<Value> = rows
2184 .iter()
2185 .filter(|r| {
2186 let name: String = r.get(0);
2187 self.session.filter().allows_schema(&name)
2188 })
2189 .map(|r| json!({ "schema_name": r.get::<_, String>(0) }))
2190 .collect();
2191 Ok(ToolOutcome::ok_json(json!(out)))
2192 }
2193
2194 async fn list_objects(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2195 let schema = args
2196 .get("schema")
2197 .and_then(|v| v.as_str())
2198 .unwrap_or("public");
2199 if !schema
2200 .chars()
2201 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
2202 {
2203 return Err(ToolError::InvalidArgs(
2204 "Invalid or missing schema name format".into(),
2205 ));
2206 }
2207 let scope = self.execution_scope_from_args(args).await?;
2208 if !scope.filter.allows_schema(schema) {
2209 return Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!([]))));
2210 }
2211 let include_partitions = args
2212 .get("include_partitions")
2213 .and_then(|v| v.as_bool())
2214 .unwrap_or(false);
2215 let kind = args.get("kind").and_then(|v| v.as_str());
2216 let partition_filter = if include_partitions {
2217 String::new()
2218 } else {
2219 " AND NOT c.relispartition".to_string()
2220 };
2221 let mut queries = Vec::new();
2222 let push_rel = |queries: &mut Vec<String>, relkinds: &[&str], label: &str| {
2223 let kinds = relkinds
2224 .iter()
2225 .map(|k| format!("'{k}'"))
2226 .collect::<Vec<_>>()
2227 .join(",");
2228 let partition_count_expr = if label == "partitioned_table" {
2229 ", (SELECT COUNT(*)::int FROM pg_inherits i WHERE i.inhparent = c.oid) AS partition_count"
2230 } else {
2231 ", NULL::int AS partition_count"
2232 };
2233 queries.push(format!(
2234 r#"
2235 SELECT n.nspname AS schema, c.relname AS name, '{label}' AS kind,
2236 d.description AS comment{partition_count_expr}
2237 FROM pg_class c
2238 JOIN pg_namespace n ON n.oid = c.relnamespace
2239 LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0
2240 WHERE n.nspname = $1 AND c.relkind IN ({kinds}){partition_filter}
2241 "#
2242 ));
2243 };
2244 if kind.is_none() || kind == Some("table") {
2245 push_rel(&mut queries, &["r", "f"], "table");
2246 if !include_partitions {
2247 push_rel(&mut queries, &["p"], "partitioned_table");
2248 } else {
2249 push_rel(&mut queries, &["r", "f", "p"], "table");
2250 }
2251 }
2252 if kind.is_none() || kind == Some("view") {
2253 push_rel(&mut queries, &["v"], "view");
2254 }
2255 if kind.is_none() || kind == Some("matview") {
2256 push_rel(&mut queries, &["m"], "matview");
2257 }
2258 if queries.is_empty() {
2259 return Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!([]))));
2260 }
2261 let sql = queries.join("\nUNION ALL\n") + "\nORDER BY kind, name";
2262 let (client, _) = self
2263 .session
2264 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2265 .await?;
2266 let rows = client.query(&sql, &[&schema]).await?;
2267 let out: Vec<Value> = rows
2268 .iter()
2269 .filter(|r| {
2270 let s: String = r.get("schema");
2271 let name: String = r.get("name");
2272 scope.filter.allows_table(&s, &name)
2273 })
2274 .map(|r| {
2275 let mut obj = json!({
2276 "schema": r.get::<_, String>("schema"),
2277 "name": r.get::<_, String>("name"),
2278 "kind": r.get::<_, String>("kind"),
2279 "comment": r.get::<_, Option<String>>("comment"),
2280 });
2281 if let Some(count) = r.get::<_, Option<i32>>("partition_count")
2282 && let Some(obj_map) = obj.as_object_mut()
2283 {
2284 obj_map.insert("partition_count".into(), json!(count));
2285 }
2286 obj
2287 })
2288 .collect();
2289 Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!(out))))
2290 }
2291
2292 async fn get_current_context(&self) -> Result<ToolOutcome, ToolError> {
2293 let (connection_id, database) = self.session.active_context().await;
2294 let conn = self
2295 .session
2296 .connections()
2297 .into_iter()
2298 .find(|c| c.id == connection_id);
2299 Ok(ToolOutcome::ok_json(json!({
2300 "connectionId": connection_id,
2301 "connectionName": conn.as_ref().map(|c| c.name.clone()).unwrap_or_else(|| "Unknown".into()),
2302 "database": database,
2303 "host": conn.as_ref().and_then(|c| c.host.clone()),
2304 "port": conn.as_ref().and_then(|c| c.port),
2305 "access_mode": match self.session.access_mode() {
2306 nexql_policy::AccessMode::Read => "read",
2307 nexql_policy::AccessMode::Write => "write",
2308 nexql_policy::AccessMode::Admin => "admin",
2309 },
2310 })))
2311 }
2312
2313 async fn switch_connection(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2314 let connection_id = args
2315 .get("connectionId")
2316 .and_then(|v| v.as_str())
2317 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
2318 let database = args
2319 .get("database")
2320 .and_then(|v| v.as_str())
2321 .map(str::to_owned);
2322 self.session.switch(connection_id, database).await?;
2323 let _ = self.ensure_index_warm().await;
2324 self.get_current_context().await
2325 }
2326
2327 async fn run_select(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2328 let sql = args
2329 .get("sql")
2330 .and_then(|v| v.as_str())
2331 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2332 let scope = self.execution_scope_from_args(args).await?;
2333 match validate_readonly_sql(sql)? {
2334 SqlDecision::Allow => {}
2335 SqlDecision::Reject => {
2336 return Err(ToolError::Execution(
2337 "Security Error: Only read-only SELECT, WITH, or EXPLAIN statements are permitted."
2338 .into(),
2339 ));
2340 }
2341 }
2342 enforce_read_table_policy(&scope.filter, sql)?;
2343 let trimmed = sql.trim().to_ascii_lowercase();
2344 let params = parse_sql_params(args);
2345 let limit = args
2346 .get("limit")
2347 .and_then(|v| v.as_u64())
2348 .map(|n| n as u32)
2349 .unwrap_or(RUN_SELECT_DEFAULT_LIMIT)
2350 .min(scope.caps.max_rows);
2351 let format = args
2352 .get("format")
2353 .and_then(|v| v.as_str())
2354 .map(RunSelectFormat::parse)
2355 .transpose()?
2356 .unwrap_or(RunSelectFormat::Compact);
2357 let timeout_ms = args
2358 .get("timeout_ms")
2359 .and_then(|v| v.as_u64())
2360 .map(|n| n as u32)
2361 .map(|n| n.min(scope.caps.statement_timeout_ms))
2362 .unwrap_or(scope.caps.statement_timeout_ms);
2363 let resolve_fks = args
2364 .get("resolve_fks")
2365 .and_then(|v| v.as_bool())
2366 .unwrap_or(false);
2367 let columnar = matches!(format, RunSelectFormat::Compact);
2368 let outcome = if trimmed.starts_with("explain") {
2369 self.run_select_internal(
2370 sql,
2371 None,
2372 columnar,
2373 ¶ms,
2374 &scope,
2375 format,
2376 timeout_ms,
2377 resolve_fks,
2378 )
2379 .await?
2380 } else {
2381 self.run_select_internal(
2382 sql,
2383 Some(limit),
2384 columnar,
2385 ¶ms,
2386 &scope,
2387 format,
2388 timeout_ms,
2389 resolve_fks,
2390 )
2391 .await?
2392 };
2393 Ok(self.attach_critique(sql, outcome).await)
2394 }
2395
2396 const FAN_OUT_MULTIPLIER: f64 = 2.0;
2405 const SEQ_SCAN_ROW_THRESHOLD: f64 = 100_000.0;
2409
2410 async fn attach_critique(&self, sql: &str, mut outcome: ToolOutcome) -> ToolOutcome {
2411 if outcome.is_error {
2412 return outcome;
2413 }
2414 let Some(structured) = outcome.structured.as_ref() else {
2415 return outcome;
2416 };
2417 if structured.get("truncated_chars").is_some() {
2418 return outcome;
2421 }
2422 let Some(row_count) = structured
2423 .get("rows")
2424 .and_then(|v| v.as_array())
2425 .map(Vec::len)
2426 else {
2427 return outcome;
2428 };
2429 let tables = select_table_refs(sql).unwrap_or_default();
2430
2431 let mut critique = Vec::new();
2432 if let Some(item) = critique::limit_without_order_by(sql) {
2433 critique.push(item.to_json());
2434 }
2435 if row_count == 0
2436 && let Some((col, val)) = critique::simple_equality_filter(sql)
2437 && let [table] = tables.as_slice()
2438 && let Some(values) = self.sample_values_best_effort(table, &col).await
2439 {
2440 let sample = values
2441 .iter()
2442 .take(5)
2443 .cloned()
2444 .collect::<Vec<_>>()
2445 .join(", ");
2446 critique.push(json!({
2447 "signal": "zero_rows",
2448 "message": format!(
2449 "no rows: {col} = '{val}'; observed values include: {sample}"
2450 ),
2451 }));
2452 }
2453 if let Some((func, col)) = critique::null_skipping_aggregate(sql)
2454 && let [table] = tables.as_slice()
2455 && let Some((null_frac, row_estimate)) =
2456 self.column_null_frac_best_effort(table, &col).await
2457 && null_frac > 0.0
2458 {
2459 let skipped = (null_frac * row_estimate).round() as i64;
2460 critique.push(json!({
2461 "signal": "null_skipping_aggregate",
2462 "message": format!(
2463 "{}({col}) skips an estimated {skipped} NULL row(s) (~{:.0}% of {}.{col}) — intended?",
2464 func.to_ascii_uppercase(), null_frac * 100.0, table.name
2465 ),
2466 }));
2467 }
2468 let max_table_rows = self.max_table_row_estimate_best_effort(&tables).await;
2469 if row_count > 0
2470 && let Some(max_rows) = max_table_rows
2471 && max_rows > 0.0
2472 && (row_count as f64) > Self::FAN_OUT_MULTIPLIER * max_rows
2473 {
2474 let ratio = row_count as f64 / max_rows;
2475 critique.push(json!({
2476 "signal": "join_fan_out",
2477 "message": format!(
2478 "output rows ({row_count}) are {ratio:.1}x the largest referenced table's estimated size ({max_rows:.0}) — possible missing join key or unintended many-to-many join."
2479 ),
2480 }));
2481 }
2482 if let Some(max_rows) = max_table_rows
2483 && max_rows > Self::SEQ_SCAN_ROW_THRESHOLD
2484 {
2485 for (relation, plan_rows) in self.large_seq_scans_best_effort(sql).await {
2486 critique.push(json!({
2487 "signal": "seq_scan_large_table",
2488 "message": format!(
2489 "seq scan on {relation} (est. {plan_rows:.0} rows) — consider an index on the filtered/joined column(s)."
2490 ),
2491 }));
2492 }
2493 }
2494
2495 if critique.is_empty() {
2496 return outcome;
2497 }
2498 if let Some(obj) = outcome.structured.as_mut().and_then(|v| v.as_object_mut()) {
2499 obj.insert("critique".into(), json!(critique));
2500 }
2501 if let Some(structured) = &outcome.structured
2502 && let Ok(text) = serde_json::to_string(structured)
2503 {
2504 outcome.text = text;
2505 }
2506 outcome
2507 }
2508
2509 async fn sample_values_best_effort(
2512 &self,
2513 table: &nexql_policy::ObjectRef,
2514 col: &str,
2515 ) -> Option<Vec<String>> {
2516 let store = self.index_store()?;
2517 let (connection_id, database) = self.session.active_context().await;
2518 let base = store.base_dir(&connection_id, &database);
2519 store.read_manifest(&base).ok()??;
2520 let svc = IndexQueryService::new(store, &connection_id, &database);
2521 let ref_ = format!("{}.{}", table.schema, table.name);
2522 let filter = self.query_filter();
2523 let result = svc.sample_values(&ref_, col, Some(&filter), None).ok()?;
2524 if result.values.is_empty() {
2525 None
2526 } else {
2527 Some(result.values)
2528 }
2529 }
2530
2531 async fn column_null_frac_best_effort(
2534 &self,
2535 table: &nexql_policy::ObjectRef,
2536 col: &str,
2537 ) -> Option<(f64, f64)> {
2538 let store = self.index_store()?;
2539 let (connection_id, database) = self.session.active_context().await;
2540 let base = store.base_dir(&connection_id, &database);
2541 let manifest = store.read_manifest(&base).ok()??;
2542 let entry = store
2543 .get_object_entry(&base, &manifest, &table.schema, &table.name)
2544 .ok()??;
2545 let profile = entry
2546 .columns
2547 .iter()
2548 .find(|c| c.name == col)?
2549 .profile
2550 .as_ref()?;
2551 Some((profile.null_frac, entry.row_estimate))
2552 }
2553
2554 async fn max_table_row_estimate_best_effort(
2558 &self,
2559 tables: &[nexql_policy::ObjectRef],
2560 ) -> Option<f64> {
2561 let store = self.index_store()?;
2562 let (connection_id, database) = self.session.active_context().await;
2563 let base = store.base_dir(&connection_id, &database);
2564 let manifest = store.read_manifest(&base).ok()??;
2565 tables
2566 .iter()
2567 .filter_map(|t| {
2568 store
2569 .get_object_entry(&base, &manifest, &t.schema, &t.name)
2570 .ok()
2571 .flatten()
2572 .map(|e| e.row_estimate)
2573 })
2574 .fold(None, |max, v| Some(max.map_or(v, |m: f64| m.max(v))))
2575 }
2576
2577 async fn large_seq_scans_best_effort(&self, sql: &str) -> Vec<(String, f64)> {
2582 let explain = build_explain_sql(sql, false);
2583 let Ok(outcome) = self.run_explain_in_transaction(&explain).await else {
2584 return Vec::new();
2585 };
2586 let Some(structured) = outcome.structured else {
2587 return Vec::new();
2588 };
2589 let Some(plan) = structured
2590 .get("rows")
2591 .and_then(|v| v.as_array())
2592 .and_then(|a| a.first())
2593 .and_then(|r| r.get("QUERY PLAN"))
2594 else {
2595 return Vec::new();
2596 };
2597 critique::large_seq_scans(plan, Self::SEQ_SCAN_ROW_THRESHOLD)
2598 }
2599
2600 async fn explain_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2601 let sql = args
2602 .get("sql")
2603 .and_then(|v| v.as_str())
2604 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2605 let scope = self.execution_scope_from_args(args).await?;
2606 match validate_readonly_sql(sql)? {
2607 SqlDecision::Allow => {}
2608 SqlDecision::Reject => {
2609 return Err(ToolError::Execution(
2610 "Security Error: Only SELECT, WITH, or EXPLAIN statements can be analyzed."
2611 .into(),
2612 ));
2613 }
2614 }
2615 enforce_read_table_policy(&scope.filter, sql)?;
2616 let clean = if sql.trim().to_ascii_lowercase().starts_with("explain") {
2617 sql.to_string()
2618 } else {
2619 format!("EXPLAIN {sql}")
2620 };
2621 if validate_readonly_sql(&clean)? == SqlDecision::Reject {
2622 return Err(ToolError::Execution(
2623 "Security Error: EXPLAIN target is not read-only.".into(),
2624 ));
2625 }
2626 let timeout_ms = scope.caps.statement_timeout_ms;
2627 self.run_select_internal(
2628 &clean,
2629 None,
2630 false,
2631 &[],
2632 &scope,
2633 RunSelectFormat::Json,
2634 timeout_ms,
2635 false,
2636 )
2637 .await
2638 }
2639
2640 async fn get_ddl(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2641 let ref_ = args
2642 .get("ref")
2643 .and_then(|v| v.as_str())
2644 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2645 let scope = self.execution_scope_from_args(args).await?;
2646 let resolved = self.resolve_ref_best_effort(ref_).await?;
2647 let (schema, name) = parse_ref(&resolved).map_err(ToolError::InvalidArgs)?;
2648 let kind = args.get("kind").and_then(|v| v.as_str()).unwrap_or("table");
2649 let reg = sql::regclass_literal(&schema, &name);
2650 let (client, _) = self
2651 .session
2652 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2653 .await?;
2654
2655 match kind {
2656 "view" | "matview" => {
2657 let sql = format!("SELECT pg_get_viewdef({reg}, true) AS definition");
2658 let rows = client.query(&sql, &[]).await?;
2659 Ok(Self::scope_tag(
2660 &scope,
2661 ToolOutcome::ok_json(rows_to_json(&rows)),
2662 ))
2663 }
2664 "function" => {
2665 let sql = format!(
2666 r#"SELECT p.proname AS name, pg_get_functiondef(p.oid) AS definition
2667 FROM pg_proc p
2668 JOIN pg_namespace n ON n.oid = p.pronamespace
2669 WHERE n.nspname = '{schema}' AND p.proname = '{name}'"#
2670 );
2671 let rows = client.query(&sql, &[]).await?;
2672 Ok(Self::scope_tag(
2673 &scope,
2674 ToolOutcome::ok_json(rows_to_json(&rows)),
2675 ))
2676 }
2677 "index" => {
2678 let sql = format!("SELECT pg_get_indexdef({reg}) AS definition");
2679 let rows = client.query(&sql, &[]).await?;
2680 Ok(Self::scope_tag(
2681 &scope,
2682 ToolOutcome::ok_json(rows_to_json(&rows)),
2683 ))
2684 }
2685 "table" => {
2686 let columns = client
2687 .query(&sql::column_details(&schema, &name), &[])
2688 .await?;
2689 let constraints = client
2690 .query(
2691 &format!(
2692 r#"SELECT conname AS name, pg_get_constraintdef(oid) AS definition
2693 FROM pg_constraint WHERE conrelid = {reg} ORDER BY conname"#
2694 ),
2695 &[],
2696 )
2697 .await?;
2698 let indexes = client
2699 .query(
2700 &format!(
2701 r#"SELECT indexname AS name, indexdef AS definition
2702 FROM pg_indexes
2703 WHERE schemaname = '{schema}' AND tablename = '{name}'
2704 ORDER BY indexname"#
2705 ),
2706 &[],
2707 )
2708 .await?;
2709 Ok(Self::scope_tag(
2710 &scope,
2711 ToolOutcome::ok_json(json!({
2712 "table": format!("{schema}.{name}"),
2713 "columns": rows_to_json(&columns),
2714 "constraints": rows_to_json(&constraints),
2715 "indexes": rows_to_json(&indexes),
2716 })),
2717 ))
2718 }
2719 other => Err(ToolError::InvalidArgs(format!(
2720 "Unsupported DDL kind \"{other}\". Use table, view, matview, function, or index."
2721 ))),
2722 }
2723 }
2724
2725 async fn table_stats(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2726 let ref_ = args
2727 .get("ref")
2728 .and_then(|v| v.as_str())
2729 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2730 let scope = self.execution_scope_from_args(args).await?;
2731 let resolved = self.resolve_ref_best_effort(ref_).await?;
2732 let (schema, name) = parse_ref(&resolved).map_err(ToolError::InvalidArgs)?;
2733 let (client, _) = self
2734 .session
2735 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2736 .await?;
2737 let stats = client.query(&sql::table_stats(&schema, &name), &[]).await?;
2738 let activity = client
2739 .query(&sql::table_activity(&schema, &name), &[])
2740 .await?;
2741 let columns = client
2742 .query(&sql::column_stats(&schema, &name), &[])
2743 .await?;
2744 let size = rows_to_json(&stats)
2745 .as_array()
2746 .and_then(|a| a.first())
2747 .cloned()
2748 .unwrap_or(Value::Null);
2749 let activity = rows_to_json(&activity)
2750 .as_array()
2751 .and_then(|a| a.first())
2752 .cloned()
2753 .unwrap_or(Value::Null);
2754 Ok(Self::scope_tag(
2755 &scope,
2756 ToolOutcome::ok_json(json!({
2757 "size": size,
2758 "activity": activity,
2759 "columns": rows_to_json(&columns),
2760 })),
2761 ))
2762 }
2763
2764 async fn index_usage(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2765 let ref_ = args
2766 .get("ref")
2767 .and_then(|v| v.as_str())
2768 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2769 let (schema, name) = parse_ref(ref_).map_err(ToolError::InvalidArgs)?;
2770 let client = self.session.checkout().await?;
2771 let rows = client.query(&sql::index_usage(&schema, &name), &[]).await?;
2772 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2773 }
2774
2775 async fn list_running_queries(&self) -> Result<ToolOutcome, ToolError> {
2776 let client = self.session.checkout().await?;
2777 let rows = client.query(sql::running_queries(), &[]).await?;
2778 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2779 }
2780
2781 async fn find_blocking_locks(&self) -> Result<ToolOutcome, ToolError> {
2782 let client = self.session.checkout().await?;
2783 let rows = client.query(sql::blocking_locks(), &[]).await?;
2784 let values = rows_to_json(&rows);
2785 if values.as_array().map(|a| a.is_empty()).unwrap_or(true) {
2786 return Ok(ToolOutcome::ok_json(json!({
2787 "message": "No blocking locks found.",
2788 "locks": [],
2789 })));
2790 }
2791 Ok(ToolOutcome::ok_json(values))
2792 }
2793
2794 async fn slow_queries(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2795 let limit = args
2796 .get("limit")
2797 .and_then(|v| v.as_u64())
2798 .map(|n| n as u32)
2799 .unwrap_or(SLOW_QUERIES_DEFAULT);
2800 let client = self.session.checkout().await?;
2801 match client.query(&sql::slow_queries(limit), &[]).await {
2802 Ok(rows) => Ok(ToolOutcome::ok_json(rows_to_json(&rows))),
2803 Err(e) => {
2804 if let Some(message) = sql::map_stat_statements_error(&e) {
2805 Ok(ToolOutcome::ok_json(json!({
2806 "error": message,
2807 "hint": message,
2808 })))
2809 } else {
2810 Err(ToolError::Postgres(e))
2811 }
2812 }
2813 }
2814 }
2815
2816 async fn db_health_check(&self) -> Result<ToolOutcome, ToolError> {
2817 let client = self.session.checkout().await?;
2818 let sections: &[(&str, &str)] = &[
2819 ("overview", sql::database_stats()),
2820 ("cache", sql::cache_hit_ratio()),
2821 ("dead_tuples", sql::database_maintenance_stats()),
2822 ("connection_states", sql::connection_states()),
2823 ("blocking_locks", sql::blocking_locks()),
2824 ];
2825 let mut report = serde_json::Map::new();
2826 for (key, q) in sections {
2827 match client.query(*q, &[]).await {
2828 Ok(rows) => {
2829 report.insert((*key).into(), rows_to_json(&rows));
2830 }
2831 Err(e) => {
2832 report.insert((*key).into(), json!({ "error": e.to_string() }));
2833 }
2834 }
2835 }
2836 let lock_count = report
2837 .get("blocking_locks")
2838 .and_then(|v| v.as_array())
2839 .map(|a| a.len() as u64);
2840 report.insert("blocking_lock_count".into(), json!(lock_count));
2841 Ok(ToolOutcome::ok_json(Value::Object(report)))
2842 }
2843
2844 async fn run_explain_in_transaction(
2846 &self,
2847 explain_sql: &str,
2848 ) -> Result<ToolOutcome, ToolError> {
2849 let client = self.session.checkout().await?;
2850 client
2851 .batch_execute("SET statement_timeout = '30s'")
2852 .await?;
2853 client.batch_execute("BEGIN").await?;
2854 let result = async {
2855 client.batch_execute("SET TRANSACTION READ ONLY").await?;
2856 let rows = client.query(explain_sql, &[]).await?;
2857 Ok::<_, ToolError>(rows_to_json(&rows))
2858 }
2859 .await;
2860 let _ = client.batch_execute("ROLLBACK").await;
2862 match result {
2863 Ok(values) => Ok(ToolOutcome::ok_json(values)),
2864 Err(e) => Err(e),
2865 }
2866 }
2867
2868 async fn get_index_status(&self) -> Result<ToolOutcome, ToolError> {
2876 let (connection_id, database) = self.session.active_context().await;
2877 let Some(store) = self.index_store() else {
2878 return Ok(ToolOutcome::ok_json(json!({
2879 "status": "missing",
2880 "connectionId": connection_id,
2881 "database": database,
2882 "remediation": "rebuild_index",
2883 })));
2884 };
2885 let base = store.base_dir(&connection_id, &database);
2886 let Some(manifest) = store.read_manifest(&base)? else {
2887 return Ok(ToolOutcome::ok_json(json!({
2888 "status": "missing",
2889 "connectionId": connection_id,
2890 "database": database,
2891 "remediation": "rebuild_index",
2892 })));
2893 };
2894
2895 let mut live_fingerprint: Option<String> = None;
2896 let mut drift: Option<bool> = None;
2897 if let Ok(client) = self.session.checkout().await {
2898 let db = PgCatalogDb::new(&client);
2899 if let Ok(fp) = db.schema_fingerprint().await {
2900 drift = Some(fp != manifest.schema_fingerprint);
2901 live_fingerprint = Some(fp);
2902 }
2903 }
2904
2905 Ok(ToolOutcome::ok_json(json!({
2906 "status": "ok",
2907 "connectionId": manifest.connection_id,
2908 "database": manifest.database,
2909 "indexedAt": manifest.indexed_at,
2910 "fingerprint": manifest.schema_fingerprint,
2911 "liveFingerprint": live_fingerprint,
2912 "drift": drift,
2913 "pgVersion": manifest.pg_version,
2914 "counts": {
2915 "tables": manifest.counts.tables,
2916 "views": manifest.counts.views,
2917 "functions": manifest.counts.functions,
2918 "enums": manifest.counts.enums,
2919 },
2920 "buildMs": manifest.stats.build_ms,
2921 "warnings": manifest.stats.warnings,
2922 })))
2923 }
2924
2925 async fn list_extensions(&self) -> Result<ToolOutcome, ToolError> {
2926 let client = self.session.checkout().await?;
2927 let rows = client.query(sql::list_extensions(), &[]).await?;
2928 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2929 }
2930
2931 async fn server_settings(&self) -> Result<ToolOutcome, ToolError> {
2932 let client = self.session.checkout().await?;
2933 let rows = client.query(sql::server_settings(), &[]).await?;
2934 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2935 }
2936
2937 async fn suggest_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2938 let limit = args
2939 .get("limit")
2940 .and_then(|v| v.as_u64())
2941 .map(|n| n as u32)
2942 .unwrap_or(REPORT_LIMIT_DEFAULT);
2943 let client = self.session.checkout().await?;
2944 let mut query_errors = serde_json::Map::new();
2945
2946 let high_seq_json = match client.query(&sql::high_seq_scan_tables(limit), &[]).await {
2947 Ok(rows) => rows_to_json(&rows),
2948 Err(e) => {
2949 query_errors.insert(
2950 "high_seq_scan_tables".into(),
2951 json!(nexql_conn::format_postgres_error(&e)),
2952 );
2953 Value::Null
2954 }
2955 };
2956
2957 let unindexed_json = match client.query(&sql::unindexed_fk_columns(limit), &[]).await {
2958 Ok(rows) => rows_to_json(&rows),
2959 Err(e) => {
2960 query_errors.insert(
2961 "unindexed_fk_columns".into(),
2962 json!(nexql_conn::format_postgres_error(&e)),
2963 );
2964 Value::Null
2965 }
2966 };
2967
2968 let mut pg_stat_available = false;
2969 let mut slow_queries = Value::Null;
2970 let mut pg_stat_note: Option<String> = None;
2971 match client.query(&sql::slow_queries(limit.min(10)), &[]).await {
2972 Ok(rows) => {
2973 pg_stat_available = true;
2974 slow_queries = rows_to_json(&rows);
2975 }
2976 Err(e) => {
2977 if let Some(message) = sql::map_stat_statements_error(&e) {
2978 pg_stat_note = Some(message);
2979 } else {
2980 query_errors.insert(
2981 "slow_queries".into(),
2982 json!(nexql_conn::format_postgres_error(&e)),
2983 );
2984 }
2985 }
2986 }
2987
2988 let mut plan_heuristics = Value::Null;
2989 if let Some(sql_text) = args.get("sql").and_then(|v| v.as_str()) {
2990 require_select_or_with(&self.session.filter(), sql_text)?;
2991 let explain = build_explain_sql(sql_text, false);
2992 match self.run_explain_in_transaction(&explain).await {
2993 Ok(outcome) => {
2994 let rows = outcome.structured.unwrap_or(Value::Null);
2995 let plan = rows
2996 .as_array()
2997 .and_then(|a| a.first())
2998 .and_then(|r| r.get("QUERY PLAN"))
2999 .cloned()
3000 .unwrap_or(Value::Null);
3001 let metrics =
3002 extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
3003 plan_heuristics = json!({
3004 "metrics": metrics,
3005 "hint": "Use deep_plan_analysis with analyze=true for actual timings before creating indexes.",
3006 });
3007 }
3008 Err(e) => {
3009 query_errors.insert("plan_heuristics".into(), json!(e.to_string()));
3010 }
3011 }
3012 }
3013
3014 let has_candidates = high_seq_json
3015 .as_array()
3016 .map(|a| !a.is_empty())
3017 .unwrap_or(false)
3018 || unindexed_json
3019 .as_array()
3020 .map(|a| !a.is_empty())
3021 .unwrap_or(false)
3022 || plan_heuristics != Value::Null;
3023
3024 let mut payload = if !has_candidates && !pg_stat_available {
3025 json!({
3026 "suggestions": [],
3027 "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.",
3028 "hint": pg_stat_note,
3029 })
3030 } else if !has_candidates {
3031 json!({
3032 "high_seq_scan_tables": high_seq_json,
3033 "unindexed_fk_columns": unindexed_json,
3034 "slow_queries": slow_queries,
3035 "plan_heuristics": plan_heuristics,
3036 "message": "No strong index candidates from sequential-scan or unindexed-FK heuristics. Review slow_queries / pass sql for plan-level advice.",
3037 "hint": "CREATE INDEX CONCURRENTLY after validating with EXPLAIN (ANALYZE, BUFFERS).",
3038 })
3039 } else {
3040 json!({
3041 "high_seq_scan_tables": high_seq_json,
3042 "unindexed_fk_columns": unindexed_json,
3043 "slow_queries": slow_queries,
3044 "plan_heuristics": plan_heuristics,
3045 "pg_stat_statements": pg_stat_available,
3046 "hint": pg_stat_note.unwrap_or_else(|| {
3047 "Validate candidates with deep_plan_analysis / EXPLAIN before CREATE INDEX CONCURRENTLY.".into()
3048 }),
3049 })
3050 };
3051
3052 if !query_errors.is_empty() {
3053 payload["query_errors"] = Value::Object(query_errors);
3054 }
3055
3056 Ok(ToolOutcome::ok_json(payload))
3057 }
3058
3059 async fn find_unused_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3060 let limit = args
3061 .get("limit")
3062 .and_then(|v| v.as_u64())
3063 .map(|n| n as u32)
3064 .unwrap_or(REPORT_LIMIT_DEFAULT);
3065 let client = self.session.checkout().await?;
3066 let rows = client.query(&sql::find_unused_indexes(limit), &[]).await?;
3067 let indexes = rows_to_json(&rows);
3068 if indexes.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3069 return Ok(ToolOutcome::ok_json(json!({
3070 "indexes": [],
3071 "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.",
3072 })));
3073 }
3074 Ok(ToolOutcome::ok_json(json!({
3075 "indexes": indexes,
3076 "hint": "Prefer DROP INDEX CONCURRENTLY after confirming the workload (and that stats are mature).",
3077 })))
3078 }
3079
3080 async fn bloat_report(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3081 let limit = args
3082 .get("limit")
3083 .and_then(|v| v.as_u64())
3084 .map(|n| n as u32)
3085 .unwrap_or(REPORT_LIMIT_DEFAULT);
3086 let client = self.session.checkout().await?;
3087 let rows = client.query(&sql::bloat_report(limit), &[]).await?;
3088 let tables = rows_to_json(&rows);
3089 if tables.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3090 return Ok(ToolOutcome::ok_json(json!({
3091 "tables": [],
3092 "method": "dead_tuple_ratio",
3093 "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.",
3094 })));
3095 }
3096 Ok(ToolOutcome::ok_json(json!({
3097 "tables": tables,
3098 "method": "dead_tuple_ratio",
3099 "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.",
3100 "hint": "VACUUM ANALYZE on high bloat_pct tables; investigate autovacuum settings if last_autovacuum is stale.",
3101 })))
3102 }
3103
3104 async fn find_missing_fks(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3105 let limit = args
3106 .get("limit")
3107 .and_then(|v| v.as_u64())
3108 .map(|n| n as u32)
3109 .unwrap_or(REPORT_LIMIT_DEFAULT);
3110 let capped = limit.clamp(1, sql::REPORT_LIMIT_MAX) as usize;
3111
3112 if let Ok((store, connection_id, database)) = self.index_service().await {
3114 let base = store.base_dir(&connection_id, &database);
3115 if let Ok(Some(manifest)) = store.read_manifest(&base)
3116 && let Ok(Some(graph)) = store.read_join_graph(&base, &manifest)
3117 {
3118 let candidates: Vec<Value> = graph
3119 .edges
3120 .into_iter()
3121 .filter(|e| e.inferred == Some(true) && e.disabled != Some(true))
3122 .take(capped)
3123 .map(|e| {
3124 let cols: Vec<Value> = e
3125 .cols
3126 .iter()
3127 .map(|(a, b)| json!({ "from": a, "to": b }))
3128 .collect();
3129 json!({
3130 "from_table": e.from,
3131 "to_table": e.to,
3132 "via": e.via,
3133 "columns": cols,
3134 "detection": "join_graph_inferred",
3135 })
3136 })
3137 .collect();
3138 if !candidates.is_empty() {
3139 return Ok(ToolOutcome::ok_json(json!({
3140 "candidates": candidates,
3141 "source": "join_graph",
3142 "hint": "These edges were inferred by naming convention and have no declared FK. Review before ALTER TABLE … ADD FOREIGN KEY.",
3143 })));
3144 }
3145 }
3146 }
3147
3148 let client = self.session.checkout().await?;
3149 let rows = client
3150 .query(&sql::find_missing_fks_catalog(limit), &[])
3151 .await?;
3152 let candidates = rows_to_json(&rows);
3153 if candidates.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3154 return Ok(ToolOutcome::ok_json(json!({
3155 "candidates": [],
3156 "source": "catalog",
3157 "message": "No missing FK candidates found via join-graph inferred edges or *_id naming against single-column PKs.",
3158 })));
3159 }
3160 Ok(ToolOutcome::ok_json(json!({
3161 "candidates": candidates,
3162 "source": "catalog",
3163 "hint": "Naming-inferred only — verify referential integrity and nullability before adding constraints. Run `nexql-mcp index build` for join-graph inferred edges.",
3164 })))
3165 }
3166
3167 async fn list_roles(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3168 let client = self.session.checkout().await?;
3169 let role = args
3170 .get("role")
3171 .and_then(|v| v.as_str())
3172 .map(str::trim)
3173 .filter(|s| !s.is_empty());
3174
3175 let Some(role_name) = role else {
3176 let rows = client.query(sql::list_roles(), &[]).await?;
3177 return Ok(ToolOutcome::ok_json(rows_to_json(&rows)));
3178 };
3179
3180 let details = client.query(sql::role_details(), &[&role_name]).await?;
3181 if details.is_empty() {
3182 return Err(ToolError::Execution(format!(
3183 "Role \"{role_name}\" not found"
3184 )));
3185 }
3186 let member_of = client.query(sql::role_member_of(), &[&role_name]).await?;
3187 let has_members = client.query(sql::role_has_members(), &[&role_name]).await?;
3188 let privileges = client
3189 .query(sql::role_table_privileges(), &[&role_name])
3190 .await?;
3191
3192 Ok(ToolOutcome::ok_json(json!({
3193 "role": rows_to_json(&details).as_array().and_then(|a| a.first().cloned()).unwrap_or(Value::Null),
3194 "member_of": rows_to_json(&member_of),
3195 "has_members": rows_to_json(&has_members),
3196 "table_privileges": rows_to_json(&privileges),
3197 })))
3198 }
3199
3200 async fn export_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3201 let sql = args
3202 .get("sql")
3203 .and_then(|v| v.as_str())
3204 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
3205 require_select_or_with(&self.session.filter(), sql)?;
3206
3207 let format = args
3208 .get("format")
3209 .and_then(|v| v.as_str())
3210 .map(|s| {
3211 ExportFormat::parse(s).ok_or_else(|| {
3212 ToolError::InvalidArgs(format!(
3213 "Unsupported format \"{s}\". Use csv, json, or sqlinsert."
3214 ))
3215 })
3216 })
3217 .transpose()?
3218 .unwrap_or(ExportFormat::Csv);
3219
3220 let table_target = match args.get("table").and_then(|v| v.as_str()) {
3221 Some(t) if !t.trim().is_empty() => Some(parse_ref(t).map_err(ToolError::InvalidArgs)?),
3222 _ => None,
3223 };
3224
3225 if format == ExportFormat::SqlInsert && table_target.is_none() {
3226 return Err(ToolError::InvalidArgs(
3227 "table (schema.name) is required when format=sqlinsert".into(),
3228 ));
3229 }
3230
3231 let scope = self.execution_scope_from_args(args).await?;
3232 let max_rows = scope.caps.max_rows;
3233 let outcome = self
3234 .run_select_internal(
3235 sql,
3236 Some(max_rows),
3237 false,
3238 &[],
3239 &scope,
3240 RunSelectFormat::Json,
3241 scope.caps.statement_timeout_ms,
3242 false,
3243 )
3244 .await?;
3245 if outcome.is_error {
3246 return Ok(outcome);
3247 }
3248
3249 let structured = outcome.structured.unwrap_or(Value::Null);
3250 let rows_val = structured
3251 .get("rows")
3252 .cloned()
3253 .or_else(|| structured.get("data").and_then(|d| d.get("rows").cloned()))
3254 .unwrap_or(Value::Array(vec![]));
3255 let rows = rows_val.as_array().cloned().unwrap_or_default();
3256 let columns = columns_from_rows(&rows);
3257 let truncated = structured
3258 .get("truncated")
3259 .and_then(|v| v.as_bool())
3260 .unwrap_or(false);
3261
3262 let payload = match format {
3263 ExportFormat::Json => {
3264 let grid: Vec<Value> = rows
3269 .iter()
3270 .map(|row| {
3271 json!(
3272 columns
3273 .iter()
3274 .map(|c| row.get(c).cloned().unwrap_or(Value::Null))
3275 .collect::<Vec<_>>()
3276 )
3277 })
3278 .collect();
3279 json!({
3280 "format": format.as_str(),
3281 "rowCount": rows.len(),
3282 "truncated": truncated,
3283 "columns": columns,
3284 "rows": grid,
3285 })
3286 }
3287 ExportFormat::Csv => {
3288 let content = rows_to_csv(&rows, &columns);
3289 let caps = self.session.caps();
3290 let (char_trunc, content) = caps.truncate_chars(&content);
3291 json!({
3292 "format": format.as_str(),
3293 "rowCount": rows.len(),
3294 "truncated": truncated || char_trunc,
3295 "columns": columns,
3296 "content": content,
3297 })
3298 }
3299 ExportFormat::SqlInsert => {
3300 let (schema, table) = table_target.expect("checked above");
3301 let content = rows_to_sql_insert(&rows, &columns, &schema, &table);
3302 let caps = self.session.caps();
3303 let (char_trunc, content) = caps.truncate_chars(&content);
3304 json!({
3305 "format": format.as_str(),
3306 "rowCount": rows.len(),
3307 "truncated": truncated || char_trunc,
3308 "table": format!("{schema}.{table}"),
3309 "columns": columns,
3310 "content": content,
3311 })
3312 }
3313 };
3314
3315 Ok(ToolOutcome::ok_json(payload))
3316 }
3317
3318 async fn db_dashboard(&self) -> Result<ToolOutcome, ToolError> {
3319 let client = self.session.checkout().await?;
3320 let sections: &[(&str, &str)] = &[
3321 ("db_info", sql::dashboard_db_info()),
3322 ("connection_states", sql::connection_states()),
3323 ("top_tables", sql::dashboard_top_tables()),
3324 ("object_counts", sql::dashboard_object_counts()),
3325 ("active_queries", sql::dashboard_active_queries()),
3326 ("blocking_locks", sql::blocking_locks()),
3327 ("max_connections", sql::dashboard_max_connections()),
3328 ("extension_count", sql::dashboard_extension_count()),
3329 ("cache", sql::cache_hit_ratio()),
3330 ];
3331 let mut report = serde_json::Map::new();
3332 for (key, q) in sections {
3333 match client.query(*q, &[]).await {
3334 Ok(rows) => {
3335 report.insert((*key).into(), rows_to_json(&rows));
3336 }
3337 Err(e) => {
3338 report.insert((*key).into(), json!({ "error": e.to_string() }));
3339 }
3340 }
3341 }
3342
3343 for key in ["db_info", "object_counts", "extension_count", "cache"] {
3345 if let Some(Value::Array(arr)) = report.get(key).cloned()
3346 && arr.len() == 1
3347 {
3348 report.insert(key.into(), arr.into_iter().next().unwrap());
3349 }
3350 }
3351 if let Some(Value::Array(arr)) = report.get("max_connections").cloned()
3352 && let Some(row) = arr.first()
3353 {
3354 report.insert(
3355 "max_connections".into(),
3356 row.get("max_connections")
3357 .cloned()
3358 .unwrap_or_else(|| row.clone()),
3359 );
3360 }
3361
3362 Ok(ToolOutcome::ok_json(Value::Object(report)))
3363 }
3364
3365 async fn deep_plan_analysis(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3366 let sql = args
3367 .get("sql")
3368 .and_then(|v| v.as_str())
3369 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
3370 require_select_or_with(&self.session.filter(), sql)?;
3371 let analyze = args
3372 .get("analyze")
3373 .and_then(|v| v.as_bool())
3374 .unwrap_or(true);
3375 let explain = build_explain_sql(sql, analyze);
3376 let outcome = self.run_explain_in_transaction(&explain).await?;
3377 let rows = outcome.structured.unwrap_or(Value::Null);
3378 let row_array = rows
3379 .get("rows")
3380 .and_then(|v| v.as_array())
3381 .or_else(|| rows.as_array());
3382 let plan = row_array
3383 .and_then(|a| a.first())
3384 .and_then(|r| r.get("QUERY PLAN"))
3385 .cloned()
3386 .unwrap_or(Value::Null);
3387 let deep = analyze_deep_plan(&plan, sql)
3388 .or_else(|| analyze_deep_plan(&rows, sql))
3389 .ok_or_else(|| {
3390 ToolError::Execution("Could not parse EXPLAIN JSON plan for deep analysis".into())
3391 })?;
3392 let metrics = extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
3393 Ok(ToolOutcome::ok_json(json!({
3394 "deep": deep,
3395 "metrics": metrics,
3396 "plan": plan,
3397 "analyzed": analyze,
3398 })))
3399 }
3400
3401 async fn schema_diff(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3402 let source_schema = args
3403 .get("sourceSchema")
3404 .and_then(|v| v.as_str())
3405 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
3406 let target_schema = args
3407 .get("targetSchema")
3408 .and_then(|v| v.as_str())
3409 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
3410 crate::schema_diff::require_safe_schema(source_schema)?;
3411 crate::schema_diff::require_safe_schema(target_schema)?;
3412
3413 let client = self.session.checkout().await?;
3414 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
3415 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
3416 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
3417 let changed = diffs
3418 .iter()
3419 .filter(|d| d.status != crate::schema_diff::DiffStatus::Unchanged)
3420 .count();
3421 Ok(ToolOutcome::ok_json(json!({
3422 "sourceSchema": source_schema,
3423 "targetSchema": target_schema,
3424 "tableCount": diffs.len(),
3425 "changedCount": changed,
3426 "diffs": crate::schema_diff::diffs_to_json(&diffs),
3427 })))
3428 }
3429
3430 async fn generate_migration(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3431 let source_schema = args
3432 .get("sourceSchema")
3433 .and_then(|v| v.as_str())
3434 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
3435 let target_schema = args
3436 .get("targetSchema")
3437 .and_then(|v| v.as_str())
3438 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
3439 crate::schema_diff::require_safe_schema(source_schema)?;
3440 crate::schema_diff::require_safe_schema(target_schema)?;
3441
3442 let client = self.session.checkout().await?;
3443 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
3444 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
3445 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
3446 let statements =
3447 crate::schema_diff::build_migration_statements(source_schema, target_schema, &diffs);
3448 let sql = if statements.is_empty() {
3449 format!("-- No differences between {source_schema} and {target_schema}")
3450 } else {
3451 statements.join("\n\n")
3452 };
3453 Ok(ToolOutcome::ok_json(json!({
3454 "sourceSchema": source_schema,
3455 "targetSchema": target_schema,
3456 "statementCount": statements.len(),
3457 "sql": sql,
3458 "hint": "Read-only: review and run via execute_sql / apply_ddl only with --access-mode write|admin. Destructive drops are commented out.",
3459 })))
3460 }
3461
3462 #[allow(clippy::too_many_arguments)]
3463 async fn finalize_run_select_payload(
3464 &self,
3465 sql: &str,
3466 mut payload: Value,
3467 scope: &ExecutionScope,
3468 format: RunSelectFormat,
3469 columnar: bool,
3470 resolve_fks: bool,
3471 client: &deadpool_postgres::Object,
3472 ) -> Result<ToolOutcome, ToolError> {
3473 if resolve_fks
3474 && let Some(store) = self.index_store()
3475 && let Ok(tables) = select_table_refs(sql)
3476 {
3477 let filter = self.query_filter_for(&scope.ctx.connection_id);
3478 let _ = resolve::resolve_fks_on_payload(
3479 &mut payload,
3480 store,
3481 &scope.ctx.connection_id,
3482 &scope.ctx.database,
3483 &tables,
3484 &filter,
3485 client,
3486 )
3487 .await;
3488 }
3489
3490 if format == RunSelectFormat::Csv {
3491 let row_objs = payload_to_row_objects(&payload);
3492 let cols = columns_from_rows(&row_objs);
3493 let csv = rows_to_csv(&row_objs, &cols);
3494 payload = json!({
3495 "format": "csv",
3496 "csv": csv,
3497 });
3498 } else if format == RunSelectFormat::Markdown {
3499 let (cols, row_vecs) = payload_to_columnar(&payload);
3500 let md = rows_to_markdown(&cols, &row_vecs);
3501 payload = json!({
3502 "format": "markdown",
3503 "markdown": md,
3504 "columns": cols,
3505 "rows": row_vecs,
3506 });
3507 } else if columnar {
3508 payload = columnarize_read_payload(payload);
3509 }
3510
3511 let text = if matches!(format, RunSelectFormat::Compact | RunSelectFormat::Csv) {
3512 serde_json::to_string(&payload)
3513 } else {
3514 serde_json::to_string_pretty(&payload)
3515 }
3516 .map_err(|e| ToolError::Execution(e.to_string()))?;
3517 let (trunc, text) = scope.caps.truncate_chars(&text);
3518 let structured = if trunc {
3519 json!({ "truncated_chars": true, "data": payload })
3520 } else {
3521 payload
3522 };
3523 Ok(Self::scope_tag(
3524 scope,
3525 ToolOutcome {
3526 text: text.to_string(),
3527 structured: Some(structured),
3528 is_error: false,
3529 },
3530 ))
3531 }
3532
3533 #[allow(clippy::too_many_arguments)]
3542 async fn run_select_internal(
3543 &self,
3544 sql: &str,
3545 max_rows: Option<u32>,
3546 columnar: bool,
3547 params: &[Value],
3548 scope: &ExecutionScope,
3549 format: RunSelectFormat,
3550 timeout_ms: u32,
3551 resolve_fks: bool,
3552 ) -> Result<ToolOutcome, ToolError> {
3553 let (client, _) = self
3554 .session
3555 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
3556 .await?;
3557 ToolSession::set_statement_timeout(&client, timeout_ms).await?;
3558 let pg_params = sql_param_boxes(params)?;
3559 let pg_param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = pg_params
3560 .iter()
3561 .map(|p| p.as_ref() as &(dyn tokio_postgres::types::ToSql + Sync))
3562 .collect();
3563
3564 let Some(max_rows) = max_rows else {
3565 let rows = self
3566 .query_with_hints(&client, sql, &pg_param_refs, timeout_ms)
3567 .await?;
3568 let values = rows_to_json(&rows);
3569 let payload =
3570 self.apply_pii_redaction_for(sql, ensure_structured_object(values), &scope.filter);
3571 return self
3572 .finalize_run_select_payload(
3573 sql,
3574 payload,
3575 scope,
3576 format,
3577 columnar,
3578 resolve_fks,
3579 &client,
3580 )
3581 .await;
3582 };
3583
3584 let cleaned = sql.trim().trim_end_matches(';').trim();
3585 let wrapped = format!(
3586 "SELECT sub.*, COUNT(*) OVER() AS {NEXQL_TOTAL_COUNT_COL} FROM ({cleaned}) AS sub LIMIT {}",
3587 max_rows + 1
3588 );
3589 let rows = self
3590 .query_with_hints(&client, &wrapped, &pg_param_refs, timeout_ms)
3591 .await?;
3592 let truncated = rows.len() as u32 > max_rows;
3593 let keep_len = if truncated {
3594 max_rows as usize
3595 } else {
3596 rows.len()
3597 };
3598 let (total_count, values) =
3599 rows_to_json_array_with_total(&rows[..keep_len], NEXQL_TOTAL_COUNT_COL);
3600 let returned = keep_len;
3601 let mut payload =
3602 self.apply_pii_redaction_for(sql, ensure_structured_object(values), &scope.filter);
3603 if let Some(obj) = payload.as_object_mut() {
3604 obj.insert("limit".into(), json!(max_rows));
3605 if let Some(total) = total_count {
3606 obj.insert("total_count".into(), json!(total));
3607 obj.insert("has_more".into(), json!(total > max_rows as i64));
3608 } else if truncated {
3609 obj.insert("has_more".into(), json!(true));
3610 } else {
3611 obj.insert("has_more".into(), json!(false));
3612 obj.insert("total_count".into(), json!(returned));
3613 }
3614 if truncated {
3615 obj.insert("truncated".into(), json!(true));
3616 }
3617 }
3618 self.finalize_run_select_payload(
3619 sql,
3620 payload,
3621 scope,
3622 format,
3623 columnar,
3624 resolve_fks,
3625 &client,
3626 )
3627 .await
3628 }
3629
3630 async fn query_with_hints(
3631 &self,
3632 client: &deadpool_postgres::Object,
3633 sql: &str,
3634 params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
3635 timeout_ms: u32,
3636 ) -> Result<Vec<tokio_postgres::Row>, ToolError> {
3637 match client.query(sql, params).await {
3638 Ok(rows) => Ok(rows),
3639 Err(e) if is_statement_timeout(&e) => Err(ToolError::Execution(
3640 serde_json::to_string(&json!({
3641 "error": "statement_timeout",
3642 "timeout_ms": timeout_ms,
3643 "sql_preview": sql.chars().take(200).collect::<String>(),
3644 "hint": "Narrow the query, add indexes, pass a higher timeout_ms, or use terminate_query in admin mode.",
3645 }))
3646 .unwrap_or_else(|_| format!("statement timeout after {timeout_ms}ms")),
3647 )),
3648 Err(e) => Err(ToolError::Execution(self.enrich_query_error(&e).await)),
3649 }
3650 }
3651
3652 fn apply_pii_redaction_for(
3653 &self,
3654 sql: &str,
3655 mut payload: Value,
3656 filter: &PolicyFilter,
3657 ) -> Value {
3658 if filter.pii_columns.is_empty() {
3659 return payload;
3660 }
3661 let Ok(tables) = select_table_refs(sql) else {
3662 return payload;
3663 };
3664 let (redacted, cols) = redact_pii_in_payload(payload, &filter.pii_columns, &tables);
3665 payload = redacted;
3666 if !cols.is_empty()
3667 && let Some(obj) = payload.as_object_mut()
3668 {
3669 obj.insert("piiRedactedColumns".into(), json!(cols));
3670 }
3671 payload
3672 }
3673}
3674
3675fn is_statement_timeout(err: &tokio_postgres::Error) -> bool {
3676 err.code()
3677 .map(|code| code.code() == "57014")
3678 .unwrap_or(false)
3679}
3680
3681fn payload_to_row_objects(payload: &Value) -> Vec<Value> {
3682 if let Some(rows) = payload.get("rows").and_then(|v| v.as_array()) {
3683 if rows
3684 .first()
3685 .and_then(|r| r.as_object())
3686 .is_some()
3687 {
3688 return rows.clone();
3689 }
3690 if let Some(cols) = payload.get("columns").and_then(|v| v.as_array()) {
3691 let col_names: Vec<String> = cols
3692 .iter()
3693 .filter_map(|c| c.as_str().map(str::to_owned))
3694 .collect();
3695 return rows
3696 .iter()
3697 .filter_map(|row| row.as_array())
3698 .map(|cells| {
3699 let mut obj = serde_json::Map::new();
3700 for (idx, name) in col_names.iter().enumerate() {
3701 obj.insert(
3702 name.clone(),
3703 cells.get(idx).cloned().unwrap_or(Value::Null),
3704 );
3705 }
3706 Value::Object(obj)
3707 })
3708 .collect();
3709 }
3710 }
3711 Vec::new()
3712}
3713
3714fn payload_to_columnar(payload: &Value) -> (Vec<String>, Vec<Vec<Value>>) {
3715 if let (Some(cols), Some(rows)) = (
3716 payload.get("columns").and_then(|v| v.as_array()),
3717 payload.get("rows").and_then(|v| v.as_array()),
3718 ) {
3719 let col_names: Vec<String> = cols
3720 .iter()
3721 .filter_map(|c| c.as_str().map(str::to_owned))
3722 .collect();
3723 let row_vecs: Vec<Vec<Value>> = rows
3724 .iter()
3725 .filter_map(|r| r.as_array().cloned())
3726 .collect();
3727 return (col_names, row_vecs);
3728 }
3729 let row_objs = payload_to_row_objects(payload);
3730 let cols = columns_from_rows(&row_objs);
3731 let row_vecs: Vec<Vec<Value>> = row_objs
3732 .iter()
3733 .map(|row| {
3734 cols.iter()
3735 .map(|c| row.get(c).cloned().unwrap_or(Value::Null))
3736 .collect()
3737 })
3738 .collect();
3739 (cols, row_vecs)
3740}
3741
3742fn normalize_for_match(s: &str) -> String {
3744 let mut out = String::new();
3745 let mut last_was_sep = true;
3746 for ch in s.to_lowercase().chars() {
3747 if ch.is_ascii_alphanumeric() {
3748 out.push(ch);
3749 last_was_sep = false;
3750 } else if !last_was_sep {
3751 out.push(' ');
3752 last_was_sep = true;
3753 }
3754 }
3755 out.trim().to_string()
3756}
3757
3758fn fuzzy_score(hint: &str, candidate: &str) -> f64 {
3760 let h = normalize_for_match(hint);
3761 let c = normalize_for_match(candidate);
3762 if h.is_empty() || c.is_empty() {
3763 return 0.0;
3764 }
3765 if h == c {
3766 return 100.0;
3767 }
3768 if c.contains(&h) || h.contains(&c) {
3769 return 75.0;
3770 }
3771 let h_tokens: std::collections::HashSet<&str> =
3772 h.split(' ').filter(|s| !s.is_empty()).collect();
3773 let c_tokens: std::collections::HashSet<&str> =
3774 c.split(' ').filter(|s| !s.is_empty()).collect();
3775 let overlap = h_tokens.intersection(&c_tokens).count();
3776 if overlap == 0 {
3777 return 0.0;
3778 }
3779 (overlap as f64 / h_tokens.len().max(c_tokens.len()) as f64) * 60.0
3780}
3781
3782fn policy_to_query_filter(filter: &PolicyFilter) -> QueryPolicyFilter {
3783 QueryPolicyFilter {
3784 allow_schemas: filter.allow_schemas.clone(),
3785 deny_schemas: filter.deny_schemas.clone(),
3786 deny_tables: filter.deny_tables.clone(),
3787 pii_columns: filter.pii_columns.clone(),
3788 }
3789}
3790
3791fn require_select_or_with(filter: &PolicyFilter, sql: &str) -> Result<(), ToolError> {
3792 match validate_readonly_sql(sql)? {
3793 SqlDecision::Allow => {}
3794 SqlDecision::Reject => {
3795 return Err(ToolError::Execution(
3796 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
3797 ));
3798 }
3799 }
3800 enforce_read_table_policy(filter, sql)?;
3801 let trimmed = sql.trim().to_ascii_lowercase();
3802 if !(trimmed.starts_with("select") || trimmed.starts_with("with")) {
3803 return Err(ToolError::Execution(
3804 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
3805 ));
3806 }
3807 Ok(())
3808}
3809
3810fn rows_to_json(rows: &[tokio_postgres::Row]) -> Value {
3811 rows_to_json_array(rows)
3812}
3813
3814fn scores_equal(a: f64, b: f64) -> bool {
3815 (a - b).abs() <= f64::EPSILON * a.abs().max(b.abs()).max(1.0)
3816}
3817
3818fn read_recent_log_errors() -> Vec<String> {
3819 let path = std::env::var("NEXQL_MCP_LOG")
3820 .map(std::path::PathBuf::from)
3821 .ok()
3822 .or_else(|| {
3823 std::env::var_os("HOME").map(|h| {
3824 std::path::PathBuf::from(h)
3825 .join(".config")
3826 .join("nexql-mcp")
3827 .join("logs")
3828 .join("nexql-mcp.log")
3829 })
3830 });
3831
3832 let Some(log_path) = path else {
3833 return Vec::new();
3834 };
3835
3836 let Ok(content) = std::fs::read_to_string(&log_path) else {
3837 return Vec::new();
3838 };
3839
3840 content
3841 .lines()
3842 .rev()
3843 .take(50)
3844 .filter(|line| {
3845 line.contains("ERROR")
3846 || line.contains("WARN")
3847 || line.contains("failed")
3848 || line.contains("Error")
3849 })
3850 .map(String::from)
3851 .collect()
3852}
3853
3854fn parse_sql_params(args: &Value) -> Vec<Value> {
3855 args.get("params")
3856 .and_then(|v| v.as_array())
3857 .map(|a| a.to_vec())
3858 .unwrap_or_default()
3859}
3860
3861fn sql_param_boxes(params: &[Value]) -> Result<Vec<Box<dyn tokio_postgres::types::ToSql + Sync + Send>>, ToolError> {
3862 params
3863 .iter()
3864 .map(json_to_sql_param)
3865 .collect::<Result<Vec<_>, _>>()
3866}
3867
3868fn json_to_sql_param(
3869 v: &Value,
3870) -> Result<Box<dyn tokio_postgres::types::ToSql + Sync + Send>, ToolError> {
3871 match v {
3872 Value::Null => Ok(Box::new(None::<String>)),
3873 Value::Bool(b) => Ok(Box::new(*b)),
3874 Value::Number(n) => {
3875 if let Some(i) = n.as_i64() {
3876 Ok(Box::new(i))
3877 } else if let Some(f) = n.as_f64() {
3878 Ok(Box::new(f))
3879 } else {
3880 Err(ToolError::InvalidArgs("invalid numeric param".into()))
3881 }
3882 }
3883 Value::String(s) => Ok(Box::new(s.clone())),
3884 _ => Err(ToolError::InvalidArgs(
3885 "params must be string, number, boolean, or null".into(),
3886 )),
3887 }
3888}
3889
3890#[cfg(test)]
3891mod tests {
3892 use super::*;
3893 use crate::plan::build_explain_sql;
3894 use nexql_policy::PolicyFilter;
3895 use serde_json::json;
3896
3897 use crate::session::{ConnectionInfo, ConnectionPolicy, ToolSession};
3898 use nexql_policy::{AccessMode, PolicyCaps};
3899
3900 static CONFIG_ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
3906
3907 fn test_conn() -> ConnectionInfo {
3908 ConnectionInfo {
3909 id: "conn-1".into(),
3910 name: "conn-1".into(),
3911 host: Some("127.0.0.1".into()),
3912 port: Some(5432),
3913 database: Some("appdb".into()),
3914 params: Default::default(),
3915 policy: ConnectionPolicy {
3916 access_mode: AccessMode::Read,
3917 caps: PolicyCaps::default(),
3918 filter: PolicyFilter::default(),
3919 environment: None,
3920 },
3921 }
3922 }
3923
3924 #[test]
3925 fn scores_equal_treats_near_duplicates_as_tied() {
3926 let s = 3.295836866004329_f64;
3927 assert!(super::scores_equal(s, s));
3928 assert!(super::scores_equal(s, s + f64::EPSILON));
3929 }
3930
3931 #[test]
3932 fn policy_maps_one_to_one() {
3933 let f = PolicyFilter {
3934 allow_schemas: vec!["public".into()],
3935 deny_schemas: vec!["pgboss".into()],
3936 deny_tables: vec!["auth.*".into()],
3937 pii_columns: vec!["public.users.ssn".into()],
3938 };
3939 let q = policy_to_query_filter(&f);
3940 assert_eq!(q.allow_schemas, f.allow_schemas);
3941 assert_eq!(q.deny_schemas, f.deny_schemas);
3942 assert_eq!(q.deny_tables, f.deny_tables);
3943 assert_eq!(q.pii_columns, f.pii_columns);
3944 }
3945
3946 #[test]
3947 fn ok_json_wraps_arrays_for_cursor_structured_content() {
3948 let out = ToolOutcome::ok_json(json!([{ "id": 1 }, { "id": 2 }]));
3949 assert!(!out.is_error);
3950 let s = out.structured.as_ref().unwrap();
3951 assert!(s.is_object(), "structuredContent must be object, got {s}");
3952 assert_eq!(s["rows"].as_array().unwrap().len(), 2);
3953 assert!(out.text.contains("\"rows\""));
3954 }
3955
3956 #[test]
3957 fn ok_json_leaves_objects_unchanged() {
3958 let out = ToolOutcome::ok_json(json!({ "kind": "table", "name": "orders" }));
3959 let s = out.structured.as_ref().unwrap();
3960 assert_eq!(s["kind"], "table");
3961 assert!(s.get("rows").is_none());
3962 }
3963
3964 #[test]
3965 fn router_specs_include_phase4_and_phase9() {
3966 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
3967 let router = ToolRouter::with_index_store(session, None);
3968 assert_eq!(router.specs().len(), ToolName::ACTIVE.len());
3969 let names: Vec<_> = router.specs().iter().map(|s| s.name.as_str()).collect();
3970 assert!(names.contains(&"search_schema"));
3971 assert!(names.contains(&"get_ddl"));
3972 assert!(names.contains(&"deep_plan_analysis"));
3973 assert!(names.contains(&"get_index_status"));
3974 assert!(names.contains(&"list_extensions"));
3975 assert!(names.contains(&"server_settings"));
3976 assert!(names.contains(&"suggest_indexes"));
3977 assert!(names.contains(&"find_unused_indexes"));
3978 assert!(names.contains(&"bloat_report"));
3979 assert!(names.contains(&"find_missing_fks"));
3980 assert!(names.contains(&"export_query"));
3981 assert!(names.contains(&"list_roles"));
3982 assert!(names.contains(&"db_dashboard"));
3983 assert!(names.contains(&"deep_plan_analysis"));
3984 assert!(names.contains(&"execute_sql"));
3985 assert!(names.contains(&"edit_row"));
3986 assert!(names.contains(&"import_data"));
3987 assert!(names.contains(&"apply_ddl"));
3988 assert!(names.contains(&"create_index_concurrently"));
3989 assert!(names.contains(&"run_maintenance"));
3990 assert!(names.contains(&"terminate_query"));
3991 }
3992
3993 #[tokio::test]
3994 async fn write_tools_refuse_read_mode() {
3995 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
3996 let router = ToolRouter::with_index_store(session, None);
3997 for tool in [
3998 "execute_sql",
3999 "edit_row",
4000 "import_data",
4001 "apply_ddl",
4002 "create_index_concurrently",
4003 "run_maintenance",
4004 "terminate_query",
4005 ] {
4006 let out = router
4007 .call(tool, json!({ "sql": "SELECT 1", "table": "public.t", "rows": [], "action": "insert", "values": {}, "pid": 1 }))
4008 .await;
4009 assert!(out.is_error, "{tool}: {}", out.text);
4010 assert!(
4011 out.text.contains("write") || out.text.contains("admin"),
4012 "{tool}: {}",
4013 out.text
4014 );
4015 }
4016 }
4017
4018 #[tokio::test]
4019 async fn table_stats_rejects_injection_ref() {
4020 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4021 let router = ToolRouter::with_index_store(session, None);
4022 let out = router
4023 .call("table_stats", json!({ "ref": "public.users; DROP" }))
4024 .await;
4025 assert!(out.is_error, "{}", out.text);
4026 assert!(
4027 out.text.contains("Invalid object reference") || out.text.contains("invalid arguments"),
4028 "expected ref validation error, got: {}",
4029 out.text
4030 );
4031 }
4032
4033 #[test]
4034 fn explain_transaction_path_builds_readonly_sequence() {
4035 let explain = build_explain_sql("SELECT 1", true);
4037 assert!(explain.starts_with("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)"));
4038 assert!(!explain.to_ascii_lowercase().contains("commit"));
4039 let steps = ["BEGIN", "SET TRANSACTION READ ONLY", &explain, "ROLLBACK"];
4040 assert_eq!(steps.len(), 4);
4041 assert_eq!(steps[0], "BEGIN");
4042 assert_eq!(steps[1], "SET TRANSACTION READ ONLY");
4043 assert_eq!(steps[3], "ROLLBACK");
4044 }
4045
4046 #[tokio::test]
4047 async fn missing_index_returns_actionable_error() {
4048 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4049 let router = ToolRouter::with_index_store(session, None);
4050 let out = router
4051 .call("search_schema", json!({ "query": "users" }))
4052 .await;
4053 assert!(out.is_error, "{}", out.text);
4054 assert!(
4055 out.text.contains("rebuild_index"),
4056 "expected actionable hint, got: {}",
4057 out.text
4058 );
4059 }
4060
4061 fn write_join_path_fixture(store: &IndexStore) {
4065 use nexql_index::{
4066 BuildDepth, BuildMode, ColumnEntry, DbObjectKind, IndexCounts, IndexDerived,
4067 IndexManifest, IndexScope, IndexStats, JOIN_GRAPH_FILE, JoinEdge, JoinGraph,
4068 ObjectEntry, ObjectShard, TOKENS_FILE,
4069 };
4070 use std::collections::HashMap;
4071
4072 let base = store.base_dir("conn-1", "appdb");
4073 let manifest = IndexManifest {
4074 format_version: 1,
4075 connection_id: "conn-1".into(),
4076 database: "appdb".into(),
4077 indexed_at: "2026-08-08T00:00:00.000Z".into(),
4078 build_mode: BuildMode::Auto,
4079 build_depth: BuildDepth::Structure,
4080 schema_fingerprint: "fp".into(),
4081 pg_version: "18.4".into(),
4082 environment: "development".into(),
4083 scope: IndexScope {
4084 included_schemas: vec!["public".into()],
4085 excluded_objects: vec![],
4086 pii_excluded_columns: vec![],
4087 },
4088 counts: IndexCounts {
4089 tables: 3,
4090 views: 0,
4091 functions: 0,
4092 enums: 0,
4093 },
4094 shards: vec![ObjectShard {
4095 file: "objects-public-0.json".into(),
4096 schema: "public".into(),
4097 objects: 3,
4098 bytes: 512,
4099 hash: "abc".into(),
4100 }],
4101 derived: IndexDerived {
4102 tokens: TOKENS_FILE.into(),
4103 join_graph: JOIN_GRAPH_FILE.into(),
4104 values: None,
4105 embeddings: None,
4106 embeddings_meta: None,
4107 },
4108 stats: IndexStats {
4109 build_ms: 1,
4110 queries_run: 1,
4111 warnings: vec![],
4112 },
4113 };
4114 store.write_manifest(&base, &manifest).unwrap();
4115
4116 fn entry(oid: u32) -> ObjectEntry {
4117 ObjectEntry {
4118 kind: DbObjectKind::Table,
4119 oid,
4120 object_hash: format!("hash{oid}"),
4121 comment: None,
4122 row_estimate: 10.0,
4123 size_bytes: 8192,
4124 columns: vec![ColumnEntry {
4125 name: "id".into(),
4126 type_name: "integer".into(),
4127 not_null: true,
4128 default_value: None,
4129 comment: None,
4130 ordinal: 1,
4131 is_pk: Some(true),
4132 profile: None,
4133 pii: None,
4134 }],
4135 primary_key: Some(vec!["id".into()]),
4136 foreign_keys: None,
4137 indexes: None,
4138 checks: None,
4139 excluded: None,
4140 definition: None,
4141 signature: None,
4142 language: None,
4143 volatility: None,
4144 body: None,
4145 values: None,
4146 base_type: None,
4147 constraint: None,
4148 }
4149 }
4150 let mut orders = entry(1);
4153 orders.columns.push(nexql_index::ColumnEntry {
4154 name: "status".into(),
4155 type_name: "text".into(),
4156 not_null: true,
4157 default_value: None,
4158 comment: None,
4159 ordinal: 2,
4160 is_pk: None,
4161 profile: Some(nexql_index::ColumnProfile {
4162 n_distinct: 2.0,
4163 null_frac: 0.0,
4164 common_values: Some(vec!["pending".into(), "paid".into()]),
4165 min: None,
4166 max: None,
4167 }),
4168 pii: None,
4169 });
4170 orders.columns.push(nexql_index::ColumnEntry {
4173 name: "amount".into(),
4174 type_name: "numeric".into(),
4175 not_null: false,
4176 default_value: None,
4177 comment: None,
4178 ordinal: 3,
4179 is_pk: None,
4180 profile: Some(nexql_index::ColumnProfile {
4181 n_distinct: 8.0,
4182 null_frac: 0.25,
4183 common_values: None,
4184 min: None,
4185 max: None,
4186 }),
4187 pii: None,
4188 });
4189
4190 let mut shard = HashMap::new();
4191 shard.insert("public.orders".into(), orders);
4192 shard.insert("public.customers".into(), entry(2));
4193 shard.insert("public.order_items".into(), entry(3));
4194 store
4195 .write_shard_entries(&base, "objects-public-0.json", &shard)
4196 .unwrap();
4197
4198 let graph = JoinGraph {
4199 edges: vec![
4200 JoinEdge {
4201 from: "public.order_items".into(),
4202 to: "public.orders".into(),
4203 via: "order_items_order_id_fkey".into(),
4204 cols: vec![("order_id".into(), "id".into())],
4205 inferred: None,
4206 disabled: None,
4207 },
4208 JoinEdge {
4209 from: "public.orders".into(),
4210 to: "public.customers".into(),
4211 via: "orders_customer_id_fkey".into(),
4212 cols: vec![("customer_id".into(), "id".into())],
4213 inferred: None,
4214 disabled: None,
4215 },
4216 ],
4217 };
4218 store.write_join_graph(&base, &graph).unwrap();
4219 }
4220
4221 fn join_path_router() -> ToolRouter {
4222 let tmp = tempfile::TempDir::new().unwrap();
4223 let store = IndexStore::new(tmp.path());
4224 write_join_path_fixture(&store);
4225 std::mem::forget(tmp); let session = ToolSession::for_tests(
4227 vec![test_conn()],
4228 PolicyFilter::default(),
4229 Some(IndexStore::new(store.root())),
4230 );
4231 ToolRouter::with_index_store(session, Some(store))
4232 }
4233
4234 #[test]
4237 fn columnarize_outcome_reshapes_flat_rows_for_generic_tool() {
4238 let outcome = ToolOutcome::ok_json(json!({
4239 "rows": [{ "name": "pgcrypto" }, { "name": "pg_stat_statements" }]
4240 }));
4241 let out = ToolRouter::columnarize_outcome("list_extensions", outcome);
4242 let structured = out.structured.unwrap();
4243 assert_eq!(structured["columns"], json!(["name"]));
4244 assert_eq!(
4245 structured["rows"],
4246 json!([["pgcrypto"], ["pg_stat_statements"]])
4247 );
4248 assert!(out.text.contains("pgcrypto"));
4250 assert!(
4251 !out.text.contains('\n'),
4252 "wire text should be compact, not pretty-printed"
4253 );
4254 }
4255
4256 #[test]
4263 fn build_tuning_summary_reads_row_object_shaped_suggestions() {
4264 let suggestions = json!({
4265 "high_seq_scan_tables": [{ "table_name": "orders" }],
4266 "unindexed_fk_columns": [{ "column_name": "customer_id" }],
4267 });
4268 let summary = ToolRouter::build_tuning_summary(&None, &suggestions);
4269 assert!(summary.contains("2 index recommendation"), "{summary}");
4270 }
4271
4272 #[test]
4275 fn columnarize_outcome_skips_orient_and_get_join_path() {
4276 for tool in ["orient", "get_join_path"] {
4277 let outcome = ToolOutcome::ok_json(json!({
4278 "tables": [{ "ref": "public.orders" }, { "ref": "public.customers" }]
4279 }));
4280 let out = ToolRouter::columnarize_outcome(tool, outcome);
4281 let structured = out.structured.unwrap();
4282 assert_eq!(
4283 structured["tables"],
4284 json!([{ "ref": "public.orders" }, { "ref": "public.customers" }]),
4285 "{tool} should be excluded from columnarization"
4286 );
4287 }
4288 }
4289
4290 #[tokio::test]
4293 async fn attach_critique_flags_limit_without_order_by() {
4294 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4295 let router = ToolRouter::with_index_store(session, None);
4296 let outcome = ToolOutcome::ok_json(json!({ "columns": ["n"], "rows": [[1]] }));
4297 let out = router
4298 .attach_critique("SELECT * FROM orders LIMIT 10", outcome)
4299 .await;
4300 let structured = out.structured.unwrap();
4301 let critique = structured["critique"].as_array().unwrap();
4302 assert!(
4303 critique
4304 .iter()
4305 .any(|c| c["signal"] == "limit_without_order_by")
4306 );
4307 }
4308
4309 #[tokio::test]
4310 async fn attach_critique_silent_when_nothing_fires() {
4311 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4312 let router = ToolRouter::with_index_store(session, None);
4313 let outcome = ToolOutcome::ok_json(json!({ "columns": ["n"], "rows": [[1]] }));
4314 let out = router
4315 .attach_critique("SELECT * FROM orders ORDER BY id LIMIT 10", outcome)
4316 .await;
4317 let structured = out.structured.unwrap();
4318 assert!(structured.get("critique").is_none());
4319 }
4320
4321 #[tokio::test]
4325 async fn attach_critique_zero_rows_suggests_observed_values() {
4326 let router = join_path_router();
4327 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": [] }));
4328 let out = router
4329 .attach_critique(
4330 "SELECT * FROM public.orders WHERE status = 'complete'",
4331 outcome,
4332 )
4333 .await;
4334 let structured = out.structured.unwrap();
4335 let critique = structured["critique"].as_array().unwrap();
4336 let zero_rows = critique
4337 .iter()
4338 .find(|c| c["signal"] == "zero_rows")
4339 .expect("zero_rows critique");
4340 let msg = zero_rows["message"].as_str().unwrap();
4341 assert!(msg.contains("status = 'complete'"), "{msg}");
4342 assert!(msg.contains("pending") && msg.contains("paid"), "{msg}");
4343 }
4344
4345 #[tokio::test]
4346 async fn attach_critique_zero_rows_silent_without_index() {
4347 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4348 let router = ToolRouter::with_index_store(session, None);
4349 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": [] }));
4350 let out = router
4351 .attach_critique(
4352 "SELECT * FROM public.orders WHERE status = 'complete'",
4353 outcome,
4354 )
4355 .await;
4356 let structured = out.structured.unwrap();
4357 assert!(structured.get("critique").is_none());
4358 }
4359
4360 #[tokio::test]
4365 async fn attach_critique_flags_join_fan_out() {
4366 let router = join_path_router();
4367 let rows: Vec<Value> = (0..25).map(|i| json!([i])).collect();
4368 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": rows }));
4369 let out = router
4370 .attach_critique("SELECT id FROM public.orders", outcome)
4371 .await;
4372 let structured = out.structured.unwrap();
4373 let critique = structured["critique"].as_array().unwrap();
4374 let fan_out = critique
4375 .iter()
4376 .find(|c| c["signal"] == "join_fan_out")
4377 .expect("join_fan_out critique");
4378 assert!(
4379 fan_out["message"].as_str().unwrap().contains("25"),
4380 "{}",
4381 fan_out["message"]
4382 );
4383 }
4384
4385 #[tokio::test]
4386 async fn attach_critique_silent_when_row_count_within_estimate() {
4387 let router = join_path_router();
4388 let rows: Vec<Value> = (0..5).map(|i| json!([i])).collect();
4389 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": rows }));
4390 let out = router
4391 .attach_critique("SELECT id FROM public.orders", outcome)
4392 .await;
4393 let structured = out.structured.unwrap();
4394 assert!(
4395 structured
4396 .get("critique")
4397 .and_then(|c| c.as_array())
4398 .map(|a| !a.iter().any(|c| c["signal"] == "join_fan_out"))
4399 .unwrap_or(true)
4400 );
4401 }
4402
4403 #[tokio::test]
4406 async fn attach_critique_flags_null_skipping_aggregate() {
4407 let router = join_path_router();
4408 let outcome = ToolOutcome::ok_json(json!({ "columns": ["avg"], "rows": [[42]] }));
4409 let out = router
4410 .attach_critique("SELECT AVG(amount) FROM public.orders", outcome)
4411 .await;
4412 let structured = out.structured.unwrap();
4413 let critique = structured["critique"].as_array().unwrap();
4414 let signal = critique
4415 .iter()
4416 .find(|c| c["signal"] == "null_skipping_aggregate")
4417 .expect("null_skipping_aggregate critique");
4418 assert!(
4419 signal["message"].as_str().unwrap().contains("AVG(amount)"),
4420 "{}",
4421 signal["message"]
4422 );
4423 }
4424
4425 #[tokio::test]
4426 async fn attach_critique_silent_for_aggregate_without_nulls() {
4427 let router = join_path_router();
4428 let outcome = ToolOutcome::ok_json(json!({ "columns": ["c"], "rows": [[2]] }));
4430 let out = router
4431 .attach_critique("SELECT COUNT(status) FROM public.orders", outcome)
4432 .await;
4433 let structured = out.structured.unwrap();
4434 assert!(structured.get("critique").is_none());
4435 }
4436
4437 #[tokio::test]
4441 async fn attach_dml_critique_zero_rows_affected_suggests_observed_values() {
4442 let router = join_path_router();
4443 let outcome = ToolOutcome::ok_json(json!({
4444 "dry_run": false,
4445 "rolled_back": false,
4446 "rows_affected": 0,
4447 "rows": [],
4448 }));
4449 let out = router
4450 .attach_dml_critique(
4451 "UPDATE public.orders SET status = 'paid' WHERE status = 'complete'",
4452 outcome,
4453 )
4454 .await;
4455 let structured = out.structured.unwrap();
4456 let critique = structured["critique"].as_array().unwrap();
4457 let zero_rows = critique
4458 .iter()
4459 .find(|c| c["signal"] == "zero_rows")
4460 .expect("zero_rows critique");
4461 let msg = zero_rows["message"].as_str().unwrap();
4462 assert!(msg.contains("status = 'complete'"), "{msg}");
4463 assert!(msg.contains("pending") && msg.contains("paid"), "{msg}");
4464 }
4465
4466 #[tokio::test]
4467 async fn attach_dml_critique_silent_when_rows_affected() {
4468 let router = join_path_router();
4469 let outcome = ToolOutcome::ok_json(json!({
4470 "dry_run": false,
4471 "rolled_back": false,
4472 "rows_affected": 3,
4473 "rows": [],
4474 }));
4475 let out = router
4476 .attach_dml_critique(
4477 "UPDATE public.orders SET status = 'paid' WHERE status = 'complete'",
4478 outcome,
4479 )
4480 .await;
4481 let structured = out.structured.unwrap();
4482 assert!(structured.get("critique").is_none());
4483 }
4484
4485 #[tokio::test]
4486 async fn attach_dml_critique_silent_for_delete_without_filter() {
4487 let router = join_path_router();
4488 let outcome = ToolOutcome::ok_json(json!({
4489 "dry_run": false,
4490 "rolled_back": false,
4491 "rows_affected": 0,
4492 "rows": [],
4493 }));
4494 let out = router
4495 .attach_dml_critique("DELETE FROM public.orders", outcome)
4496 .await;
4497 let structured = out.structured.unwrap();
4498 assert!(structured.get("critique").is_none());
4499 }
4500
4501 #[tokio::test]
4505 async fn get_join_path_resolves_unqualified_names() {
4506 let router = join_path_router();
4507 let out = router
4508 .call(
4509 "get_join_path",
4510 json!({ "a": "order_items", "b": "orders" }),
4511 )
4512 .await;
4513 assert!(!out.is_error, "{}", out.text);
4514 let structured = out.structured.unwrap();
4515 assert_eq!(structured["resolved_a"], "public.order_items");
4516 assert_eq!(structured["resolved_b"], "public.orders");
4517 assert_eq!(structured["path"][0]["via"], "order_items_order_id_fkey");
4518 }
4519
4520 #[tokio::test]
4521 async fn get_join_path_qualified_still_works() {
4522 let router = join_path_router();
4523 let out = router
4524 .call(
4525 "get_join_path",
4526 json!({ "a": "public.order_items", "b": "public.customers" }),
4527 )
4528 .await;
4529 assert!(!out.is_error, "{}", out.text);
4530 let structured = out.structured.unwrap();
4531 assert_eq!(structured["path"].as_array().unwrap().len(), 2);
4532 }
4533
4534 #[tokio::test]
4535 async fn get_join_path_unknown_name_errors_with_suggestion() {
4536 let router = join_path_router();
4537 let out = router
4538 .call("get_join_path", json!({ "a": "ordrs", "b": "customers" }))
4539 .await;
4540 assert!(out.is_error, "{}", out.text);
4541 assert!(
4542 out.text.contains("did you mean") && out.text.contains("public.orders"),
4543 "expected a did-you-mean suggestion, got: {}",
4544 out.text
4545 );
4546 }
4547
4548 #[tokio::test]
4552 async fn get_index_status_reports_missing_without_erroring_no_store() {
4553 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4554 let router = ToolRouter::with_index_store(session, None);
4555 let out = router.call("get_index_status", json!({})).await;
4556 assert!(!out.is_error, "{}", out.text);
4557 let structured = out.structured.unwrap();
4558 assert_eq!(structured["status"], "missing");
4559 assert_eq!(structured["remediation"], "rebuild_index");
4560 assert_eq!(structured["database"], "appdb");
4561 }
4562
4563 #[tokio::test]
4564 async fn get_index_status_reports_missing_without_erroring_empty_manifest() {
4565 let tmp = tempfile::TempDir::new().unwrap();
4566 let store = IndexStore::new(tmp.path());
4567 let session = ToolSession::for_tests(
4568 vec![test_conn()],
4569 PolicyFilter::default(),
4570 Some(IndexStore::new(tmp.path())),
4571 );
4572 let router = ToolRouter::with_index_store(session, Some(store));
4573 let out = router.call("get_index_status", json!({})).await;
4574 assert!(!out.is_error, "{}", out.text);
4575 let structured = out.structured.unwrap();
4576 assert_eq!(structured["status"], "missing");
4577 assert_eq!(structured["remediation"], "rebuild_index");
4578 }
4579
4580 #[tokio::test]
4581 async fn get_index_status_reports_ok_when_indexed() {
4582 let router = join_path_router();
4583 let out = router.call("get_index_status", json!({})).await;
4584 assert!(!out.is_error, "{}", out.text);
4585 let structured = out.structured.unwrap();
4586 assert_eq!(structured["status"], "ok");
4587 assert_eq!(structured["database"], "appdb");
4588 }
4589
4590 #[tokio::test]
4594 async fn orient_lists_tables_and_declared_joins() {
4595 let router = join_path_router();
4596 let out = router.call("orient", json!({})).await;
4597 assert!(!out.is_error, "{}", out.text);
4598 let structured = out.structured.unwrap();
4599 assert_eq!(structured["database"], "appdb");
4600
4601 let tables = structured["tables"].as_array().unwrap();
4602 let refs: Vec<&str> = tables.iter().map(|t| t["ref"].as_str().unwrap()).collect();
4603 assert_eq!(
4604 refs,
4605 vec!["public.customers", "public.order_items", "public.orders"]
4606 );
4607 let orders = tables.iter().find(|t| t["ref"] == "public.orders").unwrap();
4608 assert_eq!(orders["pk"], "id");
4609 assert!(orders["columns"].as_str().unwrap().contains("id:integer!"));
4610
4611 let joins = structured["joins"].as_array().unwrap();
4612 assert_eq!(joins.len(), 2);
4613 assert!(joins.iter().any(|j| j["edge"]
4614 == "public.order_items.order_id -> public.orders.id"
4615 && j["declared"] == true));
4616 }
4617
4618 #[tokio::test]
4619 async fn orient_focus_narrows_tables_and_joins() {
4620 let router = join_path_router();
4621 let out = router.call("orient", json!({ "focus": "customers" })).await;
4622 assert!(!out.is_error, "{}", out.text);
4623 let structured = out.structured.unwrap();
4624 let tables = structured["tables"].as_array().unwrap();
4625 assert_eq!(tables.len(), 1);
4626 assert_eq!(tables[0]["ref"], "public.customers");
4627 let joins = structured["joins"].as_array().unwrap();
4629 assert_eq!(joins.len(), 1);
4630 assert!(
4631 joins[0]["edge"]
4632 .as_str()
4633 .unwrap()
4634 .contains("public.customers")
4635 );
4636 }
4637
4638 #[tokio::test]
4639 async fn orient_no_index_returns_notes_not_error() {
4640 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4641 let router = ToolRouter::with_index_store(session, None);
4642 let out = router.call("orient", json!({})).await;
4643 assert!(!out.is_error, "{}", out.text);
4644 let structured = out.structured.unwrap();
4645 assert_eq!(structured["tables"].as_array().unwrap().len(), 0);
4646 assert!(!structured["notes"].as_array().unwrap().is_empty());
4647 }
4648
4649 #[tokio::test]
4650 async fn empty_index_dir_returns_build_hint() {
4651 let tmp = tempfile::TempDir::new().unwrap();
4652 let store = IndexStore::new(tmp.path());
4653 let session = ToolSession::for_tests(
4654 vec![test_conn()],
4655 PolicyFilter::default(),
4656 Some(IndexStore::new(tmp.path())),
4657 );
4658 let router = ToolRouter::with_index_store(session, Some(store));
4659 let out = router
4660 .call("describe_object", json!({ "ref": "public.users" }))
4661 .await;
4662 assert!(out.is_error, "{}", out.text);
4663 assert!(
4664 out.text.contains("rebuild_index"),
4665 "expected build hint, got: {}",
4666 out.text
4667 );
4668 }
4669
4670 #[tokio::test]
4671 async fn outcome_tagged_with_connection_id_and_database() {
4672 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4673 let router = ToolRouter::new(session);
4674 let out = router.call("list_connections", json!({})).await;
4675 let structured = out.structured.expect("structured outcome");
4676 assert_eq!(
4677 structured.get("connectionId").and_then(|v| v.as_str()),
4678 Some("conn-1")
4679 );
4680 assert_eq!(
4681 structured.get("database").and_then(|v| v.as_str()),
4682 Some("appdb")
4683 );
4684 }
4685
4686 #[tokio::test]
4687 async fn setup_connection_returns_needs_input_when_incomplete() {
4688 unsafe {
4689 std::env::remove_var("DATABASE_URL");
4690 std::env::remove_var("POSTGRES_URL");
4691 std::env::remove_var("PGHOST");
4692 }
4693 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4694 let router = ToolRouter::new(session);
4695 let out = router.call("setup_connection", json!({})).await;
4696 let structured = out.structured.expect("structured outcome");
4697 assert!(structured.get("status").is_some());
4698 }
4699
4700 #[tokio::test]
4701 async fn save_profile_persists_config() {
4702 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4703 let router = ToolRouter::new(session);
4704 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4705 let temp_dir = tempfile::tempdir().unwrap();
4706 let cfg_path = temp_dir.path().join("config.toml");
4707 unsafe {
4708 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4709 }
4710
4711 let out = router
4712 .call(
4713 "save_profile",
4714 json!({
4715 "name": "staging",
4716 "host": "127.0.0.1",
4717 "port": 5432,
4718 "dbname": "stage_db",
4719 "user": "stage_user"
4720 }),
4721 )
4722 .await;
4723
4724 let structured = out.structured.expect("structured outcome");
4725 assert_eq!(
4726 structured.get("status").and_then(|v| v.as_str()),
4727 Some("saved")
4728 );
4729 assert_eq!(
4730 structured.get("profile").and_then(|v| v.as_str()),
4731 Some("staging")
4732 );
4733 }
4734
4735 #[tokio::test]
4743 async fn save_profile_never_persists_password_in_plaintext() {
4744 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4745 let router = ToolRouter::new(session);
4746 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4747 let temp_dir = tempfile::tempdir().unwrap();
4748 let cfg_path = temp_dir.path().join("config.toml");
4749 unsafe {
4750 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4751 }
4752 const SECRET: &str = "correct-horse-battery-staple";
4753
4754 let out = router
4755 .call(
4756 "save_profile",
4757 json!({
4758 "name": "prod-db",
4759 "host": "127.0.0.1",
4760 "password": SECRET,
4761 }),
4762 )
4763 .await;
4764
4765 if out.is_error {
4766 assert!(
4769 out.text.contains("keyring") || out.text.contains("password_command"),
4770 "{}",
4771 out.text
4772 );
4773 } else {
4774 let structured = out.structured.expect("structured outcome");
4775 assert_eq!(
4776 structured.get("status").and_then(|v| v.as_str()),
4777 Some("saved")
4778 );
4779 }
4780
4781 if cfg_path.exists() {
4782 let raw = std::fs::read_to_string(&cfg_path).unwrap();
4783 assert!(
4784 !raw.contains(SECRET),
4785 "password must never appear in the persisted config file: {raw}"
4786 );
4787 if !out.is_error {
4788 assert!(
4793 raw.contains("keyring") || raw.contains(nexql_conn::ENCRYPTED_FILE_PROVIDER),
4794 "successful save must record credential_provider = \"keyring\" or \"encrypted_file\": {raw}"
4795 );
4796 }
4797 }
4798 }
4799
4800 #[tokio::test]
4806 async fn save_profile_rejects_elevated_access_without_confirmation() {
4807 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4808 let router = ToolRouter::new(session);
4809 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4810 let temp_dir = tempfile::tempdir().unwrap();
4811 let cfg_path = temp_dir.path().join("config.toml");
4812 unsafe {
4813 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4814 }
4815
4816 let out = router
4817 .call(
4818 "save_profile",
4819 json!({
4820 "name": "prod",
4821 "host": "127.0.0.1",
4822 "access_mode": "admin",
4823 }),
4824 )
4825 .await;
4826 assert!(out.is_error, "{}", out.text);
4827 assert!(out.text.contains("confirm_elevated_access"), "{}", out.text);
4828 assert!(
4829 !cfg_path.exists(),
4830 "rejected escalation must not touch the config file"
4831 );
4832 }
4833
4834 #[tokio::test]
4835 async fn save_profile_allows_elevated_access_with_confirmation() {
4836 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4837 let router = ToolRouter::new(session);
4838 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4839 let temp_dir = tempfile::tempdir().unwrap();
4840 let cfg_path = temp_dir.path().join("config.toml");
4841 unsafe {
4842 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4843 }
4844
4845 let out = router
4846 .call(
4847 "save_profile",
4848 json!({
4849 "name": "prod",
4850 "host": "127.0.0.1",
4851 "access_mode": "admin",
4852 "confirm_elevated_access": true,
4853 }),
4854 )
4855 .await;
4856 assert!(!out.is_error, "{}", out.text);
4857 }
4858
4859 #[tokio::test]
4860 async fn save_profile_read_access_mode_needs_no_confirmation() {
4861 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4862 let router = ToolRouter::new(session);
4863 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4864 let temp_dir = tempfile::tempdir().unwrap();
4865 let cfg_path = temp_dir.path().join("config.toml");
4866 unsafe {
4867 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4868 }
4869
4870 let out = router
4871 .call(
4872 "save_profile",
4873 json!({
4874 "name": "readonly",
4875 "host": "127.0.0.1",
4876 "access_mode": "read",
4877 }),
4878 )
4879 .await;
4880 assert!(!out.is_error, "{}", out.text);
4881 }
4882
4883 #[tokio::test]
4884 async fn check_ddl_safety_tool_dispatches_ast_report() {
4885 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4886 let router = ToolRouter::new(session);
4887 let out = router
4888 .call(
4889 "check_ddl_safety",
4890 json!({ "ddl": "CREATE INDEX idx_col ON users(col);" }),
4891 )
4892 .await;
4893 let structured = out.structured.expect("structured outcome");
4894 assert_eq!(
4895 structured.get("overall_risk").and_then(|v| v.as_str()),
4896 Some("CRITICAL")
4897 );
4898 }
4899}