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(
1299 profile_name: &str,
1300 password: Option<&str>,
1301 ) -> Result<(Option<String>, Option<String>), ToolError> {
1302 nexql_conn::route_password_to_keyring(profile_name, password)
1303 .map_err(|e| ToolError::Execution(e.to_string()))
1304 }
1305
1306 async fn setup_connection_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1307 let profile_name = args
1308 .get("name")
1309 .and_then(|v| v.as_str())
1310 .unwrap_or("default");
1311
1312 let candidates = crate::detect::ConnectionDetector::detect_all(None);
1313
1314 let url = args.get("url").and_then(|v| v.as_str());
1315 let host = args.get("host").and_then(|v| v.as_str());
1316 let port = args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16);
1317 let dbname = args.get("dbname").and_then(|v| v.as_str());
1318 let user = args.get("user").and_then(|v| v.as_str());
1319 let password = args.get("password").and_then(|v| v.as_str());
1320 let sslmode = args.get("sslmode").and_then(|v| v.as_str());
1321
1322 let best_cand = candidates
1323 .iter()
1324 .find(|c| c.is_complete)
1325 .or_else(|| candidates.first());
1326
1327 let res_host = host.or_else(|| best_cand.and_then(|c| c.host.as_deref()));
1328 let res_port = port.or_else(|| best_cand.and_then(|c| c.port));
1329 let res_dbname = dbname.or_else(|| best_cand.and_then(|c| c.dbname.as_deref()));
1330 let res_user = user.or_else(|| best_cand.and_then(|c| c.user.as_deref()));
1331 let res_password = password.or_else(|| best_cand.and_then(|c| c.password.as_deref()));
1332 let res_url = url.or_else(|| best_cand.and_then(|c| c.url.as_deref()));
1333 let res_sslmode = sslmode.or_else(|| best_cand.and_then(|c| c.sslmode.as_deref()));
1334
1335 if res_url.is_none() && (res_host.is_none() || res_dbname.is_none() || res_user.is_none()) {
1336 let missing: Vec<&str> = vec![
1337 if res_host.is_none() {
1338 Some("host")
1339 } else {
1340 None
1341 },
1342 if res_dbname.is_none() {
1343 Some("dbname")
1344 } else {
1345 None
1346 },
1347 if res_user.is_none() {
1348 Some("user")
1349 } else {
1350 None
1351 },
1352 ]
1353 .into_iter()
1354 .flatten()
1355 .collect();
1356
1357 return Ok(ToolOutcome::ok_json(json!({
1358 "status": "needs_input",
1359 "message": "Insufficient connection details. Please supply missing fields.",
1360 "detectedCandidates": candidates.iter().map(|c| c.redacted_json()).collect::<Vec<_>>(),
1361 "missingFields": missing
1362 })));
1363 }
1364
1365 let params = nexql_conn::ConnectionParams {
1366 url: res_url.map(String::from),
1367 host: res_host.map(String::from),
1368 port: res_port,
1369 dbname: res_dbname.map(String::from),
1370 user: res_user.map(String::from),
1371 password: res_password.map(String::from),
1372 sslmode: res_sslmode.map(String::from),
1373 ..Default::default()
1374 };
1375
1376 match nexql_conn::test_connection(¶ms).await {
1377 Ok(report) => {
1378 let (kr_password, kr_provider) =
1379 Self::route_password_to_keyring(profile_name, params.password.as_deref())?;
1380 let p_config = nexql_conn::ProfileConfig {
1381 url: params.url.clone(),
1382 host: params.host.clone(),
1383 port: params.port,
1384 dbname: params.dbname.clone(),
1385 user: params.user.clone(),
1386 password: kr_password,
1387 sslmode: params.sslmode.clone(),
1388 credential_provider: kr_provider,
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 (kr_password, kr_provider) =
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: kr_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: kr_provider,
1477 ..Default::default()
1478 };
1479
1480 let path = nexql_conn::ConfigFile::default_path()
1481 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1482
1483 let mut cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1484 .map(|(c, _)| c)
1485 .unwrap_or_default();
1486 cfg.upsert_profile(name, p_config.clone());
1487 let backup = cfg
1488 .save(&path)
1489 .map_err(|e| ToolError::Execution(e.to_string()))?;
1490 self.register_profile_in_session(name, &p_config)?;
1491
1492 Ok(ToolOutcome::ok_json(json!({
1493 "status": "saved",
1494 "profile": name,
1495 "configPath": path.to_string_lossy().to_string(),
1496 "backup": backup.map(|b| b.to_string_lossy().to_string()),
1497 "sessionReloaded": true,
1498 })))
1499 }
1500
1501 async fn test_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1502 let name = args.get("name").and_then(|v| v.as_str());
1503
1504 let params = if let Some(pname) = name {
1505 let conn = self
1506 .session
1507 .connections()
1508 .into_iter()
1509 .find(|c| c.id == pname)
1510 .ok_or_else(|| ToolError::InvalidArgs(format!("Profile '{pname}' not found")))?;
1511 conn.params.clone()
1512 } else {
1513 nexql_conn::ConnectionParams {
1514 url: args.get("url").and_then(|v| v.as_str()).map(String::from),
1515 host: args.get("host").and_then(|v| v.as_str()).map(String::from),
1516 port: args.get("port").and_then(|v| v.as_u64()).map(|n| n as u16),
1517 dbname: args
1518 .get("dbname")
1519 .and_then(|v| v.as_str())
1520 .map(String::from),
1521 user: args.get("user").and_then(|v| v.as_str()).map(String::from),
1522 password: args
1523 .get("password")
1524 .and_then(|v| v.as_str())
1525 .map(String::from),
1526 sslmode: args
1527 .get("sslmode")
1528 .and_then(|v| v.as_str())
1529 .map(String::from),
1530 ..Default::default()
1531 }
1532 };
1533
1534 match nexql_conn::test_connection(¶ms).await {
1535 Ok(report) => Ok(ToolOutcome::ok_json(json!({
1536 "success": true,
1537 "serverVersion": report.server_version,
1538 "isSuperuser": report.is_superuser,
1539 "latencyMs": report.latency.as_millis()
1540 }))),
1541 Err(e) => Ok(ToolOutcome::ok_json(json!({
1542 "success": false,
1543 "error": e.to_string()
1544 }))),
1545 }
1546 }
1547
1548 async fn export_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1549 let format = args
1550 .get("format")
1551 .and_then(|v| v.as_str())
1552 .unwrap_or("full");
1553 let path = nexql_conn::ConfigFile::default_path()
1554 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1555 let cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1556 .map(|(c, _)| c)
1557 .unwrap_or_default();
1558
1559 if format == "project" {
1560 let proj = cfg.export_shareable();
1561 let toml_str =
1562 toml::to_string_pretty(&proj).map_err(|e| ToolError::Execution(e.to_string()))?;
1563 Ok(ToolOutcome::ok_json(json!({
1564 "format": "project",
1565 "filename": ".nexql/config.toml",
1566 "description": "Project policy overlay (no credentials). Use format=full for shareable connection profiles.",
1567 "content": toml_str,
1568 })))
1569 } else {
1570 let sanitized = cfg.export_full_sanitized();
1571 let toml_str = sanitized
1572 .to_toml_string()
1573 .map_err(|e| ToolError::Execution(e.to_string()))?;
1574 Ok(ToolOutcome::ok_json(json!({
1575 "format": "full",
1576 "description": "Full user config with passwords and secrets stripped — suitable for team sharing.",
1577 "profileCount": sanitized.profiles.len(),
1578 "content": toml_str,
1579 })))
1580 }
1581 }
1582
1583 async fn import_profile_tool(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1584 let content = if let Some(c) = args.get("content").and_then(|v| v.as_str()) {
1585 c.to_string()
1586 } else if let Some(p) = args.get("path").and_then(|v| v.as_str()) {
1587 std::fs::read_to_string(p)
1588 .map_err(|e| ToolError::Execution(format!("failed to read file {p}: {e}")))?
1589 } else {
1590 return Err(ToolError::Execution(
1591 "either 'content' or 'path' must be specified".into(),
1592 ));
1593 };
1594
1595 let path = nexql_conn::ConfigFile::default_path()
1596 .ok_or_else(|| ToolError::Execution("Could not resolve config directory".into()))?;
1597 let mut cfg = nexql_conn::ConfigFile::load_path_migrated(&path)
1598 .map(|(c, _)| c)
1599 .unwrap_or_default();
1600
1601 let imported: nexql_conn::ConfigFile = toml::from_str(&content)
1602 .map_err(|e| ToolError::Execution(format!("failed to parse TOML content: {e}")))?;
1603
1604 let mut count = 0;
1605 let mut imported_names: Vec<String> = Vec::new();
1606 for (name, prof) in imported.profiles {
1607 let prepared = nexql_conn::prepare_profile_for_persist(&name, prof)
1608 .map_err(|e| ToolError::Execution(e.to_string()))?;
1609 cfg.upsert_profile(name.clone(), prepared.clone());
1610 self.register_profile_in_session(&name, &prepared)?;
1611 imported_names.push(name);
1612 count += 1;
1613 }
1614 if imported.default_profile.is_some() {
1615 cfg.default_profile = imported.default_profile;
1616 }
1617
1618 let backup = cfg
1619 .save(&path)
1620 .map_err(|e| ToolError::Execution(e.to_string()))?;
1621
1622 Ok(ToolOutcome::ok_json(json!({
1623 "status": "imported",
1624 "imported_profiles": count,
1625 "profiles": imported_names,
1626 "configPath": path.to_string_lossy().to_string(),
1627 "backup": backup.map(|b| b.to_string_lossy().to_string()),
1628 "sessionReloaded": true,
1629 })))
1630 }
1631
1632 async fn ensure_index_warm_for(&self, ctx: &ScopedContext) -> Result<(), ToolError> {
1634 let store = self
1635 .index_store()
1636 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1637 let base = store.base_dir(&ctx.connection_id, &ctx.database);
1638 if store.read_manifest(&base)?.is_some() {
1639 return Ok(());
1640 }
1641 let req = BuildRequest {
1642 connection_id: ctx.connection_id.clone(),
1643 database: ctx.database.clone(),
1644 scope: IndexScope {
1645 included_schemas: vec![],
1646 excluded_objects: vec![],
1647 pii_excluded_columns: vec![],
1648 },
1649 depth: BuildDepth::Structure,
1650 build_mode: BuildMode::Guided,
1651 environment: "development".into(),
1652 embeddings: self.use_semantic,
1653 };
1654 let (client, _) = self
1655 .session
1656 .checkout_for(CheckoutTarget::Scoped(ctx))
1657 .await
1658 .map_err(|_| {
1659 ToolError::Execution(format!(
1660 "No schema index for database \"{}\" — call the 'rebuild_index' tool to build an index.",
1661 ctx.database
1662 ))
1663 })?;
1664 let db = PgCatalogDb::new(&client);
1665 build_index(store, &db, &req, None, None, self.embedder.as_deref())
1666 .await
1667 .map_err(|e| ToolError::Execution(format!("Automatic index build failed: {e}")))?;
1668 self.session
1669 .clear_index_stale(&ctx.connection_id, &ctx.database);
1670 Ok(())
1671 }
1672
1673 async fn ensure_index_warm(&self) -> Result<(), ToolError> {
1674 let (connection_id, database) = self.session.active_context().await;
1675 self.ensure_index_warm_for(&ScopedContext {
1676 connection_id,
1677 database,
1678 })
1679 .await
1680 }
1681
1682 async fn index_service_for(
1683 &self,
1684 ctx: &ScopedContext,
1685 ) -> Result<(&IndexStore, String, String), ToolError> {
1686 let store = self
1687 .index_store()
1688 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1689 let base = store.base_dir(&ctx.connection_id, &ctx.database);
1690 if store.read_manifest(&base)?.is_none() {
1691 self.ensure_index_warm_for(ctx).await?;
1692 }
1693 Ok((store, ctx.connection_id.clone(), ctx.database.clone()))
1694 }
1695
1696 async fn index_service(&self) -> Result<(&IndexStore, String, String), ToolError> {
1697 let store = self
1698 .index_store()
1699 .ok_or_else(|| ToolError::Execution(NO_INDEX_HINT.into()))?;
1700 let (connection_id, database) = self.session.active_context().await;
1701 let base = store.base_dir(&connection_id, &database);
1702 if store.read_manifest(&base)?.is_none() {
1703 self.ensure_index_warm().await?;
1704 }
1705 Ok((store, connection_id, database))
1706 }
1707
1708 async fn collect_index_refs(&self) -> Vec<String> {
1709 if let Ok((store, connection_id, database)) = self.index_service().await {
1710 let base = store.base_dir(&connection_id, &database);
1711 if let Ok(Some(manifest)) = store.read_manifest(&base) {
1712 let mut refs = Vec::new();
1713 for shard in &manifest.shards {
1714 if let Ok(Some(entries)) = store.read_shard_entries(&base, &shard.file) {
1715 refs.extend(entries.keys().cloned());
1716 }
1717 }
1718 return refs;
1719 }
1720 }
1721 Vec::new()
1722 }
1723
1724 async fn enrich_query_error(&self, pg_err: &tokio_postgres::Error) -> String {
1725 let base = nexql_conn::format_postgres_error(pg_err);
1726 let refs = self.collect_index_refs().await;
1727 sql::enhance_sql_error(&base, &refs)
1728 }
1729
1730 fn ref_resolution_error(ref_: &str, resolution: &RefResolution) -> Option<ToolError> {
1734 match resolution {
1735 RefResolution::Resolved(_) => None,
1736 RefResolution::Ambiguous(candidates) => Some(ToolError::InvalidArgs(format!(
1737 "ambiguous relation \"{ref_}\": {}",
1738 candidates.join(", ")
1739 ))),
1740 RefResolution::Unknown {
1741 suggestion: Some(s),
1742 } => Some(ToolError::InvalidArgs(format!(
1743 "unknown relation \"{ref_}\" — did you mean \"{s}\"?"
1744 ))),
1745 RefResolution::Unknown { suggestion: None } => Some(ToolError::InvalidArgs(format!(
1746 "unknown relation \"{ref_}\" — call search_schema to find valid refs."
1747 ))),
1748 }
1749 }
1750
1751 fn resolve_indexed_ref_strict(
1755 svc: &IndexQueryService<'_>,
1756 ref_: &str,
1757 ) -> Result<String, ToolError> {
1758 let resolution = svc.resolve_ref(ref_)?;
1759 if let Some(err) = Self::ref_resolution_error(ref_, &resolution) {
1760 return Err(err);
1761 }
1762 match resolution {
1763 RefResolution::Resolved(r) => Ok(r),
1764 _ => unreachable!("ref_resolution_error covers every non-Resolved case"),
1765 }
1766 }
1767
1768 fn resolve_indexed_ref_soft(
1773 svc: &IndexQueryService<'_>,
1774 ref_: &str,
1775 ) -> Result<String, ToolError> {
1776 match svc.resolve_ref(ref_)? {
1777 RefResolution::Resolved(r) => Ok(r),
1778 RefResolution::Ambiguous(candidates) => Err(ToolError::InvalidArgs(format!(
1779 "ambiguous relation \"{ref_}\": {}",
1780 candidates.join(", ")
1781 ))),
1782 RefResolution::Unknown { .. } => Ok(ref_.to_owned()),
1783 }
1784 }
1785
1786 async fn resolve_ref_best_effort(&self, ref_: &str) -> Result<String, ToolError> {
1792 let Some(store) = self.index_store() else {
1793 return Ok(ref_.to_owned());
1794 };
1795 let (connection_id, database) = self.session.active_context().await;
1796 let base = store.base_dir(&connection_id, &database);
1797 if store.read_manifest(&base)?.is_none() {
1798 return Ok(ref_.to_owned());
1799 }
1800 let svc = IndexQueryService::new(store, &connection_id, &database);
1801 Self::resolve_indexed_ref_soft(&svc, ref_)
1802 }
1803
1804 async fn search_schema(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1805 let query = args
1806 .get("query")
1807 .and_then(|v| v.as_str())
1808 .unwrap_or("")
1809 .trim();
1810 if query.is_empty() {
1811 return Ok(ToolOutcome::ok_json(json!([])));
1812 }
1813 let (store, connection_id, database) = self.index_service().await?;
1814 let svc = IndexQueryService::new(store, &connection_id, &database);
1815 let filter = self.query_filter();
1816 let hits = svc.search_schema(
1817 query,
1818 SEARCH_SCHEMA_LIMIT,
1819 Some(&filter),
1820 SearchOptions {
1821 use_semantic: self.use_semantic,
1822 embedder: self.embedder.as_deref(),
1823 },
1824 )?;
1825 let rows: Vec<Value> = hits
1826 .into_iter()
1827 .map(|h| {
1828 json!({
1829 "ref": h.ref_,
1830 "score": h.score,
1831 "kind": h.kind,
1832 })
1833 })
1834 .collect();
1835 Ok(ToolOutcome::ok_json(json!(rows)))
1836 }
1837
1838 async fn inspect_or_search(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1839 let query = args
1840 .get("query")
1841 .and_then(|v| v.as_str())
1842 .unwrap_or("")
1843 .trim();
1844 if query.is_empty() {
1845 return Err(ToolError::InvalidArgs("query is required".into()));
1846 }
1847 let include_columns = args
1848 .get("include_columns")
1849 .and_then(|v| v.as_bool())
1850 .unwrap_or(true);
1851 let limit_objects = args
1852 .get("limit_objects")
1853 .and_then(|v| v.as_u64())
1854 .map(|n| n as usize)
1855 .unwrap_or(3)
1856 .clamp(1, 20);
1857
1858 let (store, connection_id, database) = self.index_service().await?;
1859 let svc = IndexQueryService::new(store, &connection_id, &database);
1860 let filter = self.query_filter();
1861 let hits = svc.search_schema(
1862 query,
1863 limit_objects,
1864 Some(&filter),
1865 SearchOptions {
1866 use_semantic: self.use_semantic,
1867 embedder: self.embedder.as_deref(),
1868 },
1869 )?;
1870
1871 let mut matches = Vec::with_capacity(hits.len());
1872 for hit in hits {
1873 let entry = svc.describe_object(&hit.ref_, Some(&filter))?;
1874 let mut obj = json!({
1875 "ref": hit.ref_,
1876 "score": hit.score,
1877 "kind": hit.kind,
1878 "row_estimate": entry.row_estimate.round() as i64,
1879 "primary_key": entry.primary_key,
1880 });
1881 if include_columns {
1882 let fk_cols: std::collections::HashSet<String> = entry
1883 .foreign_keys
1884 .as_ref()
1885 .map(|fks| {
1886 fks.iter()
1887 .flat_map(|fk| fk.columns.clone())
1888 .collect()
1889 })
1890 .unwrap_or_default();
1891 let columns: Vec<Value> = entry
1892 .columns
1893 .iter()
1894 .map(|c| {
1895 json!({
1896 "name": c.name,
1897 "type": c.type_name,
1898 "is_pk": c.is_pk.unwrap_or(false),
1899 "is_fk": fk_cols.contains(&c.name),
1900 "not_null": c.not_null,
1901 })
1902 })
1903 .collect();
1904 obj["columns"] = json!(columns);
1905 }
1906 matches.push(obj);
1907 }
1908
1909 Ok(ToolOutcome::ok_json(json!({
1910 "query": query,
1911 "matches": matches,
1912 })))
1913 }
1914
1915 async fn search_all_databases(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1916 let query = args
1917 .get("query")
1918 .and_then(|v| v.as_str())
1919 .unwrap_or("")
1920 .trim();
1921 if query.is_empty() {
1922 return Err(ToolError::InvalidArgs("query is required".into()));
1923 }
1924 let limit_per_db = args
1925 .get("limit_per_database")
1926 .and_then(|v| v.as_u64())
1927 .map(|n| n as usize)
1928 .unwrap_or(3)
1929 .clamp(1, 10);
1930 let limit_pairs = args
1931 .get("limit_connections")
1932 .and_then(|v| v.as_u64())
1933 .map(|n| n as usize)
1934 .unwrap_or(20)
1935 .clamp(1, 50);
1936
1937 let Some(store) = self.index_store() else {
1938 return Err(ToolError::Execution(NO_INDEX_HINT.into()));
1939 };
1940
1941 let indexed = store.list_indexed_databases().unwrap_or_default();
1942 let connections = self.session.connections();
1943 let filter = self.query_filter();
1944 let mut hits: Vec<Value> = Vec::new();
1945 let mut searched = 0usize;
1946
1947 for (connection_id, database) in indexed {
1948 if searched >= limit_pairs {
1949 break;
1950 }
1951 if !connections.iter().any(|c| c.id == connection_id) {
1952 continue;
1953 }
1954 searched += 1;
1955 let svc = IndexQueryService::new(store, &connection_id, &database);
1956 if let Ok(results) = svc.search_schema(
1957 query,
1958 limit_per_db,
1959 Some(&filter),
1960 SearchOptions {
1961 use_semantic: self.use_semantic,
1962 embedder: self.embedder.as_deref(),
1963 },
1964 ) {
1965 for hit in results {
1966 hits.push(json!({
1967 "connectionId": connection_id,
1968 "database": database,
1969 "ref": hit.ref_,
1970 "score": hit.score,
1971 "kind": hit.kind,
1972 }));
1973 }
1974 }
1975 }
1976
1977 hits.sort_by(|a, b| {
1978 let sa = a.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
1979 let sb = b.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
1980 sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)
1981 });
1982
1983 Ok(ToolOutcome::ok_json(json!({
1984 "query": query,
1985 "hits": hits,
1986 "searched_pairs": searched,
1987 })))
1988 }
1989
1990 async fn describe_object(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
1991 let ref_ = args
1992 .get("ref")
1993 .and_then(|v| v.as_str())
1994 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
1995 let scope = self.execution_scope_from_args(args).await?;
1996 let resolve_refs = args
1997 .get("resolve_refs")
1998 .and_then(|v| v.as_bool())
1999 .unwrap_or(false);
2000 let resolve_limit = args
2001 .get("resolve_refs_limit")
2002 .and_then(|v| v.as_u64())
2003 .map(|n| n as usize)
2004 .unwrap_or(DEFAULT_RESOLVE_REFS_LIMIT);
2005 let (store, connection_id, database) = self.index_service_for(&scope.ctx).await?;
2006 let svc = IndexQueryService::new(store, &connection_id, &database);
2007 let resolved = Self::resolve_indexed_ref_soft(&svc, ref_)?;
2008 let filter = self.query_filter_for(&scope.ctx.connection_id);
2009 let entry = svc.describe_object(&resolved, Some(&filter))?;
2010 let mut value =
2011 serde_json::to_value(&entry).map_err(|e| ToolError::Execution(e.to_string()))?;
2012 let client = if resolve_refs {
2013 Some(
2014 self.session
2015 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2016 .await?
2017 .0,
2018 )
2019 } else {
2020 None
2021 };
2022 value = resolve::enrich_describe_object_with_store(
2023 value,
2024 &entry,
2025 store,
2026 &svc,
2027 resolve_refs,
2028 resolve_limit,
2029 client.as_ref(),
2030 )
2031 .await?;
2032 Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(value)))
2033 }
2034
2035 async fn get_join_path(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2036 let a = args
2037 .get("a")
2038 .and_then(|v| v.as_str())
2039 .ok_or_else(|| ToolError::InvalidArgs("a is required".into()))?;
2040 let b = args
2041 .get("b")
2042 .and_then(|v| v.as_str())
2043 .ok_or_else(|| ToolError::InvalidArgs("b is required".into()))?;
2044 let (store, connection_id, database) = self.index_service().await?;
2045 let svc = IndexQueryService::new(store, &connection_id, &database);
2046 let resolved_a = Self::resolve_indexed_ref_strict(&svc, a)?;
2051 let resolved_b = Self::resolve_indexed_ref_strict(&svc, b)?;
2052 let path = svc.get_join_path(&resolved_a, &resolved_b)?;
2053 let path_value =
2054 serde_json::to_value(path).map_err(|e| ToolError::Execution(e.to_string()))?;
2055 Ok(ToolOutcome::ok_json(json!({
2056 "path": path_value,
2057 "resolved_a": resolved_a,
2058 "resolved_b": resolved_b,
2059 })))
2060 }
2061
2062 async fn sample_values(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2063 let ref_ = args
2064 .get("ref")
2065 .and_then(|v| v.as_str())
2066 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2067 let col = args
2068 .get("col")
2069 .and_then(|v| v.as_str())
2070 .ok_or_else(|| ToolError::InvalidArgs("col is required".into()))?;
2071 let (store, connection_id, database) = self.index_service().await?;
2072 let svc = IndexQueryService::new(store, &connection_id, &database);
2073 let resolved = Self::resolve_indexed_ref_soft(&svc, ref_)?;
2074 let ref_ = resolved.as_str();
2075 let filter = self.query_filter();
2076 let result = svc.sample_values(ref_, col, Some(&filter), None)?;
2077
2078 let mut values = result.values;
2079 let mut message = result.message;
2080
2081 if values.is_empty()
2082 && let Ok(client) = self.session.checkout().await
2083 {
2084 let parts: Vec<&str> = ref_.split('.').collect();
2085 let (schema, table) = match parts.as_slice() {
2086 [s, t] => (*s, *t),
2087 _ => ("public", ref_),
2088 };
2089 let safe_schema = schema.replace('"', "\"\"");
2090 let safe_table = table.replace('"', "\"\"");
2091 let safe_col = col.replace('"', "\"\"");
2092 let query = format!(
2093 "SELECT DISTINCT \"{safe_col}\"::text FROM \"{safe_schema}\".\"{safe_table}\" WHERE \"{safe_col}\" IS NOT NULL LIMIT 20"
2094 );
2095 if let Ok(rows) = client.query(&query, &[]).await {
2096 let sampled: Vec<String> = rows
2097 .iter()
2098 .filter_map(|r| r.get::<_, Option<String>>(0))
2099 .collect();
2100 if !sampled.is_empty() {
2101 values = sampled;
2102 message = None;
2103 }
2104 }
2105 }
2106
2107 let mut payload = json!({ "values": values });
2108 if let Some(msg) = message {
2109 payload["message"] = json!(msg);
2110 }
2111 Ok(ToolOutcome::ok_json(payload))
2112 }
2113
2114 fn list_connections(&self) -> ToolOutcome {
2115 let rows: Vec<Value> = self
2116 .session
2117 .connections()
2118 .iter()
2119 .map(|c| {
2120 json!({
2121 "id": c.id,
2122 "name": c.name,
2123 "host": c.host,
2124 "port": c.port,
2125 "database": c.database,
2126 })
2127 })
2128 .collect();
2129 ToolOutcome::ok_json(json!(rows))
2130 }
2131
2132 async fn list_databases(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2133 let connection_id = args
2134 .get("connectionId")
2135 .and_then(|v| v.as_str())
2136 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
2137 let conn = self
2138 .session
2139 .connections()
2140 .into_iter()
2141 .find(|c| c.id == connection_id)
2142 .ok_or_else(|| {
2143 ToolError::Execution(format!(
2144 "Connection not found for ID: {connection_id} — call list_connections"
2145 ))
2146 })?;
2147 let client = {
2149 if self.session.active_context().await.0 == connection_id {
2151 self.session.checkout().await?
2152 } else {
2153 let pool_opts = self.session.pool_opts();
2154 let pool = nexql_conn::create_pool(&conn.params, &pool_opts).await?;
2155 nexql_conn::checkout_guarded(&pool, &pool_opts).await?
2156 }
2157 };
2158 let rows = client
2159 .query(
2160 "SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname",
2161 &[],
2162 )
2163 .await?;
2164 let names: Vec<String> = rows.iter().map(|r| r.get(0)).collect();
2165 Ok(ToolOutcome::ok_json(json!(names)))
2166 }
2167
2168 async fn list_schemas(&self) -> Result<ToolOutcome, ToolError> {
2169 let client = self.session.checkout().await?;
2170 let rows = client
2171 .query(
2172 r#"
2173 SELECT nspname AS schema_name
2174 FROM pg_namespace
2175 WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
2176 AND nspname NOT LIKE 'pg_%'
2177 ORDER BY nspname
2178 "#,
2179 &[],
2180 )
2181 .await?;
2182 let out: Vec<Value> = rows
2183 .iter()
2184 .filter(|r| {
2185 let name: String = r.get(0);
2186 self.session.filter().allows_schema(&name)
2187 })
2188 .map(|r| json!({ "schema_name": r.get::<_, String>(0) }))
2189 .collect();
2190 Ok(ToolOutcome::ok_json(json!(out)))
2191 }
2192
2193 async fn list_objects(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2194 let schema = args
2195 .get("schema")
2196 .and_then(|v| v.as_str())
2197 .unwrap_or("public");
2198 if !schema
2199 .chars()
2200 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
2201 {
2202 return Err(ToolError::InvalidArgs(
2203 "Invalid or missing schema name format".into(),
2204 ));
2205 }
2206 let scope = self.execution_scope_from_args(args).await?;
2207 if !scope.filter.allows_schema(schema) {
2208 return Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!([]))));
2209 }
2210 let include_partitions = args
2211 .get("include_partitions")
2212 .and_then(|v| v.as_bool())
2213 .unwrap_or(false);
2214 let kind = args.get("kind").and_then(|v| v.as_str());
2215 let partition_filter = if include_partitions {
2216 String::new()
2217 } else {
2218 " AND NOT c.relispartition".to_string()
2219 };
2220 let mut queries = Vec::new();
2221 let push_rel = |queries: &mut Vec<String>, relkinds: &[&str], label: &str| {
2222 let kinds = relkinds
2223 .iter()
2224 .map(|k| format!("'{k}'"))
2225 .collect::<Vec<_>>()
2226 .join(",");
2227 let partition_count_expr = if label == "partitioned_table" {
2228 ", (SELECT COUNT(*)::int FROM pg_inherits i WHERE i.inhparent = c.oid) AS partition_count"
2229 } else {
2230 ", NULL::int AS partition_count"
2231 };
2232 queries.push(format!(
2233 r#"
2234 SELECT n.nspname AS schema, c.relname AS name, '{label}' AS kind,
2235 d.description AS comment{partition_count_expr}
2236 FROM pg_class c
2237 JOIN pg_namespace n ON n.oid = c.relnamespace
2238 LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0
2239 WHERE n.nspname = $1 AND c.relkind IN ({kinds}){partition_filter}
2240 "#
2241 ));
2242 };
2243 if kind.is_none() || kind == Some("table") {
2244 push_rel(&mut queries, &["r", "f"], "table");
2245 if !include_partitions {
2246 push_rel(&mut queries, &["p"], "partitioned_table");
2247 } else {
2248 push_rel(&mut queries, &["r", "f", "p"], "table");
2249 }
2250 }
2251 if kind.is_none() || kind == Some("view") {
2252 push_rel(&mut queries, &["v"], "view");
2253 }
2254 if kind.is_none() || kind == Some("matview") {
2255 push_rel(&mut queries, &["m"], "matview");
2256 }
2257 if queries.is_empty() {
2258 return Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!([]))));
2259 }
2260 let sql = queries.join("\nUNION ALL\n") + "\nORDER BY kind, name";
2261 let (client, _) = self
2262 .session
2263 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2264 .await?;
2265 let rows = client.query(&sql, &[&schema]).await?;
2266 let out: Vec<Value> = rows
2267 .iter()
2268 .filter(|r| {
2269 let s: String = r.get("schema");
2270 let name: String = r.get("name");
2271 scope.filter.allows_table(&s, &name)
2272 })
2273 .map(|r| {
2274 let mut obj = json!({
2275 "schema": r.get::<_, String>("schema"),
2276 "name": r.get::<_, String>("name"),
2277 "kind": r.get::<_, String>("kind"),
2278 "comment": r.get::<_, Option<String>>("comment"),
2279 });
2280 if let Some(count) = r.get::<_, Option<i32>>("partition_count")
2281 && let Some(obj_map) = obj.as_object_mut()
2282 {
2283 obj_map.insert("partition_count".into(), json!(count));
2284 }
2285 obj
2286 })
2287 .collect();
2288 Ok(Self::scope_tag(&scope, ToolOutcome::ok_json(json!(out))))
2289 }
2290
2291 async fn get_current_context(&self) -> Result<ToolOutcome, ToolError> {
2292 let (connection_id, database) = self.session.active_context().await;
2293 let conn = self
2294 .session
2295 .connections()
2296 .into_iter()
2297 .find(|c| c.id == connection_id);
2298 Ok(ToolOutcome::ok_json(json!({
2299 "connectionId": connection_id,
2300 "connectionName": conn.as_ref().map(|c| c.name.clone()).unwrap_or_else(|| "Unknown".into()),
2301 "database": database,
2302 "host": conn.as_ref().and_then(|c| c.host.clone()),
2303 "port": conn.as_ref().and_then(|c| c.port),
2304 "access_mode": match self.session.access_mode() {
2305 nexql_policy::AccessMode::Read => "read",
2306 nexql_policy::AccessMode::Write => "write",
2307 nexql_policy::AccessMode::Admin => "admin",
2308 },
2309 })))
2310 }
2311
2312 async fn switch_connection(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2313 let connection_id = args
2314 .get("connectionId")
2315 .and_then(|v| v.as_str())
2316 .ok_or_else(|| ToolError::InvalidArgs("connectionId is required".into()))?;
2317 let database = args
2318 .get("database")
2319 .and_then(|v| v.as_str())
2320 .map(str::to_owned);
2321 self.session.switch(connection_id, database).await?;
2322 let _ = self.ensure_index_warm().await;
2323 self.get_current_context().await
2324 }
2325
2326 async fn run_select(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2327 let sql = args
2328 .get("sql")
2329 .and_then(|v| v.as_str())
2330 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2331 let scope = self.execution_scope_from_args(args).await?;
2332 match validate_readonly_sql(sql)? {
2333 SqlDecision::Allow => {}
2334 SqlDecision::Reject => {
2335 return Err(ToolError::Execution(
2336 "Security Error: Only read-only SELECT, WITH, or EXPLAIN statements are permitted."
2337 .into(),
2338 ));
2339 }
2340 }
2341 enforce_read_table_policy(&scope.filter, sql)?;
2342 let trimmed = sql.trim().to_ascii_lowercase();
2343 let params = parse_sql_params(args);
2344 let limit = args
2345 .get("limit")
2346 .and_then(|v| v.as_u64())
2347 .map(|n| n as u32)
2348 .unwrap_or(RUN_SELECT_DEFAULT_LIMIT)
2349 .min(scope.caps.max_rows);
2350 let format = args
2351 .get("format")
2352 .and_then(|v| v.as_str())
2353 .map(RunSelectFormat::parse)
2354 .transpose()?
2355 .unwrap_or(RunSelectFormat::Compact);
2356 let timeout_ms = args
2357 .get("timeout_ms")
2358 .and_then(|v| v.as_u64())
2359 .map(|n| n as u32)
2360 .map(|n| n.min(scope.caps.statement_timeout_ms))
2361 .unwrap_or(scope.caps.statement_timeout_ms);
2362 let resolve_fks = args
2363 .get("resolve_fks")
2364 .and_then(|v| v.as_bool())
2365 .unwrap_or(false);
2366 let columnar = matches!(format, RunSelectFormat::Compact);
2367 let outcome = if trimmed.starts_with("explain") {
2368 self.run_select_internal(
2369 sql,
2370 None,
2371 columnar,
2372 ¶ms,
2373 &scope,
2374 format,
2375 timeout_ms,
2376 resolve_fks,
2377 )
2378 .await?
2379 } else {
2380 self.run_select_internal(
2381 sql,
2382 Some(limit),
2383 columnar,
2384 ¶ms,
2385 &scope,
2386 format,
2387 timeout_ms,
2388 resolve_fks,
2389 )
2390 .await?
2391 };
2392 Ok(self.attach_critique(sql, outcome).await)
2393 }
2394
2395 const FAN_OUT_MULTIPLIER: f64 = 2.0;
2404 const SEQ_SCAN_ROW_THRESHOLD: f64 = 100_000.0;
2408
2409 async fn attach_critique(&self, sql: &str, mut outcome: ToolOutcome) -> ToolOutcome {
2410 if outcome.is_error {
2411 return outcome;
2412 }
2413 let Some(structured) = outcome.structured.as_ref() else {
2414 return outcome;
2415 };
2416 if structured.get("truncated_chars").is_some() {
2417 return outcome;
2420 }
2421 let Some(row_count) = structured
2422 .get("rows")
2423 .and_then(|v| v.as_array())
2424 .map(Vec::len)
2425 else {
2426 return outcome;
2427 };
2428 let tables = select_table_refs(sql).unwrap_or_default();
2429
2430 let mut critique = Vec::new();
2431 if let Some(item) = critique::limit_without_order_by(sql) {
2432 critique.push(item.to_json());
2433 }
2434 if row_count == 0
2435 && let Some((col, val)) = critique::simple_equality_filter(sql)
2436 && let [table] = tables.as_slice()
2437 && let Some(values) = self.sample_values_best_effort(table, &col).await
2438 {
2439 let sample = values
2440 .iter()
2441 .take(5)
2442 .cloned()
2443 .collect::<Vec<_>>()
2444 .join(", ");
2445 critique.push(json!({
2446 "signal": "zero_rows",
2447 "message": format!(
2448 "no rows: {col} = '{val}'; observed values include: {sample}"
2449 ),
2450 }));
2451 }
2452 if let Some((func, col)) = critique::null_skipping_aggregate(sql)
2453 && let [table] = tables.as_slice()
2454 && let Some((null_frac, row_estimate)) =
2455 self.column_null_frac_best_effort(table, &col).await
2456 && null_frac > 0.0
2457 {
2458 let skipped = (null_frac * row_estimate).round() as i64;
2459 critique.push(json!({
2460 "signal": "null_skipping_aggregate",
2461 "message": format!(
2462 "{}({col}) skips an estimated {skipped} NULL row(s) (~{:.0}% of {}.{col}) — intended?",
2463 func.to_ascii_uppercase(), null_frac * 100.0, table.name
2464 ),
2465 }));
2466 }
2467 let max_table_rows = self.max_table_row_estimate_best_effort(&tables).await;
2468 if row_count > 0
2469 && let Some(max_rows) = max_table_rows
2470 && max_rows > 0.0
2471 && (row_count as f64) > Self::FAN_OUT_MULTIPLIER * max_rows
2472 {
2473 let ratio = row_count as f64 / max_rows;
2474 critique.push(json!({
2475 "signal": "join_fan_out",
2476 "message": format!(
2477 "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."
2478 ),
2479 }));
2480 }
2481 if let Some(max_rows) = max_table_rows
2482 && max_rows > Self::SEQ_SCAN_ROW_THRESHOLD
2483 {
2484 for (relation, plan_rows) in self.large_seq_scans_best_effort(sql).await {
2485 critique.push(json!({
2486 "signal": "seq_scan_large_table",
2487 "message": format!(
2488 "seq scan on {relation} (est. {plan_rows:.0} rows) — consider an index on the filtered/joined column(s)."
2489 ),
2490 }));
2491 }
2492 }
2493
2494 if critique.is_empty() {
2495 return outcome;
2496 }
2497 if let Some(obj) = outcome.structured.as_mut().and_then(|v| v.as_object_mut()) {
2498 obj.insert("critique".into(), json!(critique));
2499 }
2500 if let Some(structured) = &outcome.structured
2501 && let Ok(text) = serde_json::to_string(structured)
2502 {
2503 outcome.text = text;
2504 }
2505 outcome
2506 }
2507
2508 async fn sample_values_best_effort(
2511 &self,
2512 table: &nexql_policy::ObjectRef,
2513 col: &str,
2514 ) -> Option<Vec<String>> {
2515 let store = self.index_store()?;
2516 let (connection_id, database) = self.session.active_context().await;
2517 let base = store.base_dir(&connection_id, &database);
2518 store.read_manifest(&base).ok()??;
2519 let svc = IndexQueryService::new(store, &connection_id, &database);
2520 let ref_ = format!("{}.{}", table.schema, table.name);
2521 let filter = self.query_filter();
2522 let result = svc.sample_values(&ref_, col, Some(&filter), None).ok()?;
2523 if result.values.is_empty() {
2524 None
2525 } else {
2526 Some(result.values)
2527 }
2528 }
2529
2530 async fn column_null_frac_best_effort(
2533 &self,
2534 table: &nexql_policy::ObjectRef,
2535 col: &str,
2536 ) -> Option<(f64, f64)> {
2537 let store = self.index_store()?;
2538 let (connection_id, database) = self.session.active_context().await;
2539 let base = store.base_dir(&connection_id, &database);
2540 let manifest = store.read_manifest(&base).ok()??;
2541 let entry = store
2542 .get_object_entry(&base, &manifest, &table.schema, &table.name)
2543 .ok()??;
2544 let profile = entry
2545 .columns
2546 .iter()
2547 .find(|c| c.name == col)?
2548 .profile
2549 .as_ref()?;
2550 Some((profile.null_frac, entry.row_estimate))
2551 }
2552
2553 async fn max_table_row_estimate_best_effort(
2557 &self,
2558 tables: &[nexql_policy::ObjectRef],
2559 ) -> Option<f64> {
2560 let store = self.index_store()?;
2561 let (connection_id, database) = self.session.active_context().await;
2562 let base = store.base_dir(&connection_id, &database);
2563 let manifest = store.read_manifest(&base).ok()??;
2564 tables
2565 .iter()
2566 .filter_map(|t| {
2567 store
2568 .get_object_entry(&base, &manifest, &t.schema, &t.name)
2569 .ok()
2570 .flatten()
2571 .map(|e| e.row_estimate)
2572 })
2573 .fold(None, |max, v| Some(max.map_or(v, |m: f64| m.max(v))))
2574 }
2575
2576 async fn large_seq_scans_best_effort(&self, sql: &str) -> Vec<(String, f64)> {
2581 let explain = build_explain_sql(sql, false);
2582 let Ok(outcome) = self.run_explain_in_transaction(&explain).await else {
2583 return Vec::new();
2584 };
2585 let Some(structured) = outcome.structured else {
2586 return Vec::new();
2587 };
2588 let Some(plan) = structured
2589 .get("rows")
2590 .and_then(|v| v.as_array())
2591 .and_then(|a| a.first())
2592 .and_then(|r| r.get("QUERY PLAN"))
2593 else {
2594 return Vec::new();
2595 };
2596 critique::large_seq_scans(plan, Self::SEQ_SCAN_ROW_THRESHOLD)
2597 }
2598
2599 async fn explain_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2600 let sql = args
2601 .get("sql")
2602 .and_then(|v| v.as_str())
2603 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
2604 let scope = self.execution_scope_from_args(args).await?;
2605 match validate_readonly_sql(sql)? {
2606 SqlDecision::Allow => {}
2607 SqlDecision::Reject => {
2608 return Err(ToolError::Execution(
2609 "Security Error: Only SELECT, WITH, or EXPLAIN statements can be analyzed."
2610 .into(),
2611 ));
2612 }
2613 }
2614 enforce_read_table_policy(&scope.filter, sql)?;
2615 let clean = if sql.trim().to_ascii_lowercase().starts_with("explain") {
2616 sql.to_string()
2617 } else {
2618 format!("EXPLAIN {sql}")
2619 };
2620 if validate_readonly_sql(&clean)? == SqlDecision::Reject {
2621 return Err(ToolError::Execution(
2622 "Security Error: EXPLAIN target is not read-only.".into(),
2623 ));
2624 }
2625 let timeout_ms = scope.caps.statement_timeout_ms;
2626 self.run_select_internal(
2627 &clean,
2628 None,
2629 false,
2630 &[],
2631 &scope,
2632 RunSelectFormat::Json,
2633 timeout_ms,
2634 false,
2635 )
2636 .await
2637 }
2638
2639 async fn get_ddl(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2640 let ref_ = args
2641 .get("ref")
2642 .and_then(|v| v.as_str())
2643 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2644 let scope = self.execution_scope_from_args(args).await?;
2645 let resolved = self.resolve_ref_best_effort(ref_).await?;
2646 let (schema, name) = parse_ref(&resolved).map_err(ToolError::InvalidArgs)?;
2647 let kind = args.get("kind").and_then(|v| v.as_str()).unwrap_or("table");
2648 let reg = sql::regclass_literal(&schema, &name);
2649 let (client, _) = self
2650 .session
2651 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2652 .await?;
2653
2654 match kind {
2655 "view" | "matview" => {
2656 let sql = format!("SELECT pg_get_viewdef({reg}, true) AS definition");
2657 let rows = client.query(&sql, &[]).await?;
2658 Ok(Self::scope_tag(
2659 &scope,
2660 ToolOutcome::ok_json(rows_to_json(&rows)),
2661 ))
2662 }
2663 "function" => {
2664 let sql = format!(
2665 r#"SELECT p.proname AS name, pg_get_functiondef(p.oid) AS definition
2666 FROM pg_proc p
2667 JOIN pg_namespace n ON n.oid = p.pronamespace
2668 WHERE n.nspname = '{schema}' AND p.proname = '{name}'"#
2669 );
2670 let rows = client.query(&sql, &[]).await?;
2671 Ok(Self::scope_tag(
2672 &scope,
2673 ToolOutcome::ok_json(rows_to_json(&rows)),
2674 ))
2675 }
2676 "index" => {
2677 let sql = format!("SELECT pg_get_indexdef({reg}) AS definition");
2678 let rows = client.query(&sql, &[]).await?;
2679 Ok(Self::scope_tag(
2680 &scope,
2681 ToolOutcome::ok_json(rows_to_json(&rows)),
2682 ))
2683 }
2684 "table" => {
2685 let columns = client
2686 .query(&sql::column_details(&schema, &name), &[])
2687 .await?;
2688 let constraints = client
2689 .query(
2690 &format!(
2691 r#"SELECT conname AS name, pg_get_constraintdef(oid) AS definition
2692 FROM pg_constraint WHERE conrelid = {reg} ORDER BY conname"#
2693 ),
2694 &[],
2695 )
2696 .await?;
2697 let indexes = client
2698 .query(
2699 &format!(
2700 r#"SELECT indexname AS name, indexdef AS definition
2701 FROM pg_indexes
2702 WHERE schemaname = '{schema}' AND tablename = '{name}'
2703 ORDER BY indexname"#
2704 ),
2705 &[],
2706 )
2707 .await?;
2708 Ok(Self::scope_tag(
2709 &scope,
2710 ToolOutcome::ok_json(json!({
2711 "table": format!("{schema}.{name}"),
2712 "columns": rows_to_json(&columns),
2713 "constraints": rows_to_json(&constraints),
2714 "indexes": rows_to_json(&indexes),
2715 })),
2716 ))
2717 }
2718 other => Err(ToolError::InvalidArgs(format!(
2719 "Unsupported DDL kind \"{other}\". Use table, view, matview, function, or index."
2720 ))),
2721 }
2722 }
2723
2724 async fn table_stats(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2725 let ref_ = args
2726 .get("ref")
2727 .and_then(|v| v.as_str())
2728 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2729 let scope = self.execution_scope_from_args(args).await?;
2730 let resolved = self.resolve_ref_best_effort(ref_).await?;
2731 let (schema, name) = parse_ref(&resolved).map_err(ToolError::InvalidArgs)?;
2732 let (client, _) = self
2733 .session
2734 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
2735 .await?;
2736 let stats = client.query(&sql::table_stats(&schema, &name), &[]).await?;
2737 let activity = client
2738 .query(&sql::table_activity(&schema, &name), &[])
2739 .await?;
2740 let columns = client
2741 .query(&sql::column_stats(&schema, &name), &[])
2742 .await?;
2743 let size = rows_to_json(&stats)
2744 .as_array()
2745 .and_then(|a| a.first())
2746 .cloned()
2747 .unwrap_or(Value::Null);
2748 let activity = rows_to_json(&activity)
2749 .as_array()
2750 .and_then(|a| a.first())
2751 .cloned()
2752 .unwrap_or(Value::Null);
2753 Ok(Self::scope_tag(
2754 &scope,
2755 ToolOutcome::ok_json(json!({
2756 "size": size,
2757 "activity": activity,
2758 "columns": rows_to_json(&columns),
2759 })),
2760 ))
2761 }
2762
2763 async fn index_usage(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2764 let ref_ = args
2765 .get("ref")
2766 .and_then(|v| v.as_str())
2767 .ok_or_else(|| ToolError::InvalidArgs("ref is required".into()))?;
2768 let (schema, name) = parse_ref(ref_).map_err(ToolError::InvalidArgs)?;
2769 let client = self.session.checkout().await?;
2770 let rows = client.query(&sql::index_usage(&schema, &name), &[]).await?;
2771 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2772 }
2773
2774 async fn list_running_queries(&self) -> Result<ToolOutcome, ToolError> {
2775 let client = self.session.checkout().await?;
2776 let rows = client.query(sql::running_queries(), &[]).await?;
2777 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2778 }
2779
2780 async fn find_blocking_locks(&self) -> Result<ToolOutcome, ToolError> {
2781 let client = self.session.checkout().await?;
2782 let rows = client.query(sql::blocking_locks(), &[]).await?;
2783 let values = rows_to_json(&rows);
2784 if values.as_array().map(|a| a.is_empty()).unwrap_or(true) {
2785 return Ok(ToolOutcome::ok_json(json!({
2786 "message": "No blocking locks found.",
2787 "locks": [],
2788 })));
2789 }
2790 Ok(ToolOutcome::ok_json(values))
2791 }
2792
2793 async fn slow_queries(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2794 let limit = args
2795 .get("limit")
2796 .and_then(|v| v.as_u64())
2797 .map(|n| n as u32)
2798 .unwrap_or(SLOW_QUERIES_DEFAULT);
2799 let client = self.session.checkout().await?;
2800 match client.query(&sql::slow_queries(limit), &[]).await {
2801 Ok(rows) => Ok(ToolOutcome::ok_json(rows_to_json(&rows))),
2802 Err(e) => {
2803 if let Some(message) = sql::map_stat_statements_error(&e) {
2804 Ok(ToolOutcome::ok_json(json!({
2805 "error": message,
2806 "hint": message,
2807 })))
2808 } else {
2809 Err(ToolError::Postgres(e))
2810 }
2811 }
2812 }
2813 }
2814
2815 async fn db_health_check(&self) -> Result<ToolOutcome, ToolError> {
2816 let client = self.session.checkout().await?;
2817 let sections: &[(&str, &str)] = &[
2818 ("overview", sql::database_stats()),
2819 ("cache", sql::cache_hit_ratio()),
2820 ("dead_tuples", sql::database_maintenance_stats()),
2821 ("connection_states", sql::connection_states()),
2822 ("blocking_locks", sql::blocking_locks()),
2823 ];
2824 let mut report = serde_json::Map::new();
2825 for (key, q) in sections {
2826 match client.query(*q, &[]).await {
2827 Ok(rows) => {
2828 report.insert((*key).into(), rows_to_json(&rows));
2829 }
2830 Err(e) => {
2831 report.insert((*key).into(), json!({ "error": e.to_string() }));
2832 }
2833 }
2834 }
2835 let lock_count = report
2836 .get("blocking_locks")
2837 .and_then(|v| v.as_array())
2838 .map(|a| a.len() as u64);
2839 report.insert("blocking_lock_count".into(), json!(lock_count));
2840 Ok(ToolOutcome::ok_json(Value::Object(report)))
2841 }
2842
2843 async fn run_explain_in_transaction(
2845 &self,
2846 explain_sql: &str,
2847 ) -> Result<ToolOutcome, ToolError> {
2848 let client = self.session.checkout().await?;
2849 client
2850 .batch_execute("SET statement_timeout = '30s'")
2851 .await?;
2852 client.batch_execute("BEGIN").await?;
2853 let result = async {
2854 client.batch_execute("SET TRANSACTION READ ONLY").await?;
2855 let rows = client.query(explain_sql, &[]).await?;
2856 Ok::<_, ToolError>(rows_to_json(&rows))
2857 }
2858 .await;
2859 let _ = client.batch_execute("ROLLBACK").await;
2861 match result {
2862 Ok(values) => Ok(ToolOutcome::ok_json(values)),
2863 Err(e) => Err(e),
2864 }
2865 }
2866
2867 async fn get_index_status(&self) -> Result<ToolOutcome, ToolError> {
2875 let (connection_id, database) = self.session.active_context().await;
2876 let Some(store) = self.index_store() else {
2877 return Ok(ToolOutcome::ok_json(json!({
2878 "status": "missing",
2879 "connectionId": connection_id,
2880 "database": database,
2881 "remediation": "rebuild_index",
2882 })));
2883 };
2884 let base = store.base_dir(&connection_id, &database);
2885 let Some(manifest) = store.read_manifest(&base)? else {
2886 return Ok(ToolOutcome::ok_json(json!({
2887 "status": "missing",
2888 "connectionId": connection_id,
2889 "database": database,
2890 "remediation": "rebuild_index",
2891 })));
2892 };
2893
2894 let mut live_fingerprint: Option<String> = None;
2895 let mut drift: Option<bool> = None;
2896 if let Ok(client) = self.session.checkout().await {
2897 let db = PgCatalogDb::new(&client);
2898 if let Ok(fp) = db.schema_fingerprint().await {
2899 drift = Some(fp != manifest.schema_fingerprint);
2900 live_fingerprint = Some(fp);
2901 }
2902 }
2903
2904 Ok(ToolOutcome::ok_json(json!({
2905 "status": "ok",
2906 "connectionId": manifest.connection_id,
2907 "database": manifest.database,
2908 "indexedAt": manifest.indexed_at,
2909 "fingerprint": manifest.schema_fingerprint,
2910 "liveFingerprint": live_fingerprint,
2911 "drift": drift,
2912 "pgVersion": manifest.pg_version,
2913 "counts": {
2914 "tables": manifest.counts.tables,
2915 "views": manifest.counts.views,
2916 "functions": manifest.counts.functions,
2917 "enums": manifest.counts.enums,
2918 },
2919 "buildMs": manifest.stats.build_ms,
2920 "warnings": manifest.stats.warnings,
2921 })))
2922 }
2923
2924 async fn list_extensions(&self) -> Result<ToolOutcome, ToolError> {
2925 let client = self.session.checkout().await?;
2926 let rows = client.query(sql::list_extensions(), &[]).await?;
2927 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2928 }
2929
2930 async fn server_settings(&self) -> Result<ToolOutcome, ToolError> {
2931 let client = self.session.checkout().await?;
2932 let rows = client.query(sql::server_settings(), &[]).await?;
2933 Ok(ToolOutcome::ok_json(rows_to_json(&rows)))
2934 }
2935
2936 async fn suggest_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
2937 let limit = args
2938 .get("limit")
2939 .and_then(|v| v.as_u64())
2940 .map(|n| n as u32)
2941 .unwrap_or(REPORT_LIMIT_DEFAULT);
2942 let client = self.session.checkout().await?;
2943 let mut query_errors = serde_json::Map::new();
2944
2945 let high_seq_json = match client.query(&sql::high_seq_scan_tables(limit), &[]).await {
2946 Ok(rows) => rows_to_json(&rows),
2947 Err(e) => {
2948 query_errors.insert(
2949 "high_seq_scan_tables".into(),
2950 json!(nexql_conn::format_postgres_error(&e)),
2951 );
2952 Value::Null
2953 }
2954 };
2955
2956 let unindexed_json = match client.query(&sql::unindexed_fk_columns(limit), &[]).await {
2957 Ok(rows) => rows_to_json(&rows),
2958 Err(e) => {
2959 query_errors.insert(
2960 "unindexed_fk_columns".into(),
2961 json!(nexql_conn::format_postgres_error(&e)),
2962 );
2963 Value::Null
2964 }
2965 };
2966
2967 let mut pg_stat_available = false;
2968 let mut slow_queries = Value::Null;
2969 let mut pg_stat_note: Option<String> = None;
2970 match client.query(&sql::slow_queries(limit.min(10)), &[]).await {
2971 Ok(rows) => {
2972 pg_stat_available = true;
2973 slow_queries = rows_to_json(&rows);
2974 }
2975 Err(e) => {
2976 if let Some(message) = sql::map_stat_statements_error(&e) {
2977 pg_stat_note = Some(message);
2978 } else {
2979 query_errors.insert(
2980 "slow_queries".into(),
2981 json!(nexql_conn::format_postgres_error(&e)),
2982 );
2983 }
2984 }
2985 }
2986
2987 let mut plan_heuristics = Value::Null;
2988 if let Some(sql_text) = args.get("sql").and_then(|v| v.as_str()) {
2989 require_select_or_with(&self.session.filter(), sql_text)?;
2990 let explain = build_explain_sql(sql_text, false);
2991 match self.run_explain_in_transaction(&explain).await {
2992 Ok(outcome) => {
2993 let rows = outcome.structured.unwrap_or(Value::Null);
2994 let plan = rows
2995 .as_array()
2996 .and_then(|a| a.first())
2997 .and_then(|r| r.get("QUERY PLAN"))
2998 .cloned()
2999 .unwrap_or(Value::Null);
3000 let metrics =
3001 extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
3002 plan_heuristics = json!({
3003 "metrics": metrics,
3004 "hint": "Use deep_plan_analysis with analyze=true for actual timings before creating indexes.",
3005 });
3006 }
3007 Err(e) => {
3008 query_errors.insert("plan_heuristics".into(), json!(e.to_string()));
3009 }
3010 }
3011 }
3012
3013 let has_candidates = high_seq_json
3014 .as_array()
3015 .map(|a| !a.is_empty())
3016 .unwrap_or(false)
3017 || unindexed_json
3018 .as_array()
3019 .map(|a| !a.is_empty())
3020 .unwrap_or(false)
3021 || plan_heuristics != Value::Null;
3022
3023 let mut payload = if !has_candidates && !pg_stat_available {
3024 json!({
3025 "suggestions": [],
3026 "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.",
3027 "hint": pg_stat_note,
3028 })
3029 } else if !has_candidates {
3030 json!({
3031 "high_seq_scan_tables": high_seq_json,
3032 "unindexed_fk_columns": unindexed_json,
3033 "slow_queries": slow_queries,
3034 "plan_heuristics": plan_heuristics,
3035 "message": "No strong index candidates from sequential-scan or unindexed-FK heuristics. Review slow_queries / pass sql for plan-level advice.",
3036 "hint": "CREATE INDEX CONCURRENTLY after validating with EXPLAIN (ANALYZE, BUFFERS).",
3037 })
3038 } else {
3039 json!({
3040 "high_seq_scan_tables": high_seq_json,
3041 "unindexed_fk_columns": unindexed_json,
3042 "slow_queries": slow_queries,
3043 "plan_heuristics": plan_heuristics,
3044 "pg_stat_statements": pg_stat_available,
3045 "hint": pg_stat_note.unwrap_or_else(|| {
3046 "Validate candidates with deep_plan_analysis / EXPLAIN before CREATE INDEX CONCURRENTLY.".into()
3047 }),
3048 })
3049 };
3050
3051 if !query_errors.is_empty() {
3052 payload["query_errors"] = Value::Object(query_errors);
3053 }
3054
3055 Ok(ToolOutcome::ok_json(payload))
3056 }
3057
3058 async fn find_unused_indexes(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3059 let limit = args
3060 .get("limit")
3061 .and_then(|v| v.as_u64())
3062 .map(|n| n as u32)
3063 .unwrap_or(REPORT_LIMIT_DEFAULT);
3064 let client = self.session.checkout().await?;
3065 let rows = client.query(&sql::find_unused_indexes(limit), &[]).await?;
3066 let indexes = rows_to_json(&rows);
3067 if indexes.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3068 return Ok(ToolOutcome::ok_json(json!({
3069 "indexes": [],
3070 "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.",
3071 })));
3072 }
3073 Ok(ToolOutcome::ok_json(json!({
3074 "indexes": indexes,
3075 "hint": "Prefer DROP INDEX CONCURRENTLY after confirming the workload (and that stats are mature).",
3076 })))
3077 }
3078
3079 async fn bloat_report(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3080 let limit = args
3081 .get("limit")
3082 .and_then(|v| v.as_u64())
3083 .map(|n| n as u32)
3084 .unwrap_or(REPORT_LIMIT_DEFAULT);
3085 let client = self.session.checkout().await?;
3086 let rows = client.query(&sql::bloat_report(limit), &[]).await?;
3087 let tables = rows_to_json(&rows);
3088 if tables.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3089 return Ok(ToolOutcome::ok_json(json!({
3090 "tables": [],
3091 "method": "dead_tuple_ratio",
3092 "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.",
3093 })));
3094 }
3095 Ok(ToolOutcome::ok_json(json!({
3096 "tables": tables,
3097 "method": "dead_tuple_ratio",
3098 "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.",
3099 "hint": "VACUUM ANALYZE on high bloat_pct tables; investigate autovacuum settings if last_autovacuum is stale.",
3100 })))
3101 }
3102
3103 async fn find_missing_fks(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3104 let limit = args
3105 .get("limit")
3106 .and_then(|v| v.as_u64())
3107 .map(|n| n as u32)
3108 .unwrap_or(REPORT_LIMIT_DEFAULT);
3109 let capped = limit.clamp(1, sql::REPORT_LIMIT_MAX) as usize;
3110
3111 if let Ok((store, connection_id, database)) = self.index_service().await {
3113 let base = store.base_dir(&connection_id, &database);
3114 if let Ok(Some(manifest)) = store.read_manifest(&base)
3115 && let Ok(Some(graph)) = store.read_join_graph(&base, &manifest)
3116 {
3117 let candidates: Vec<Value> = graph
3118 .edges
3119 .into_iter()
3120 .filter(|e| e.inferred == Some(true) && e.disabled != Some(true))
3121 .take(capped)
3122 .map(|e| {
3123 let cols: Vec<Value> = e
3124 .cols
3125 .iter()
3126 .map(|(a, b)| json!({ "from": a, "to": b }))
3127 .collect();
3128 json!({
3129 "from_table": e.from,
3130 "to_table": e.to,
3131 "via": e.via,
3132 "columns": cols,
3133 "detection": "join_graph_inferred",
3134 })
3135 })
3136 .collect();
3137 if !candidates.is_empty() {
3138 return Ok(ToolOutcome::ok_json(json!({
3139 "candidates": candidates,
3140 "source": "join_graph",
3141 "hint": "These edges were inferred by naming convention and have no declared FK. Review before ALTER TABLE … ADD FOREIGN KEY.",
3142 })));
3143 }
3144 }
3145 }
3146
3147 let client = self.session.checkout().await?;
3148 let rows = client
3149 .query(&sql::find_missing_fks_catalog(limit), &[])
3150 .await?;
3151 let candidates = rows_to_json(&rows);
3152 if candidates.as_array().map(|a| a.is_empty()).unwrap_or(true) {
3153 return Ok(ToolOutcome::ok_json(json!({
3154 "candidates": [],
3155 "source": "catalog",
3156 "message": "No missing FK candidates found via join-graph inferred edges or *_id naming against single-column PKs.",
3157 })));
3158 }
3159 Ok(ToolOutcome::ok_json(json!({
3160 "candidates": candidates,
3161 "source": "catalog",
3162 "hint": "Naming-inferred only — verify referential integrity and nullability before adding constraints. Run `nexql-mcp index build` for join-graph inferred edges.",
3163 })))
3164 }
3165
3166 async fn list_roles(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3167 let client = self.session.checkout().await?;
3168 let role = args
3169 .get("role")
3170 .and_then(|v| v.as_str())
3171 .map(str::trim)
3172 .filter(|s| !s.is_empty());
3173
3174 let Some(role_name) = role else {
3175 let rows = client.query(sql::list_roles(), &[]).await?;
3176 return Ok(ToolOutcome::ok_json(rows_to_json(&rows)));
3177 };
3178
3179 let details = client.query(sql::role_details(), &[&role_name]).await?;
3180 if details.is_empty() {
3181 return Err(ToolError::Execution(format!(
3182 "Role \"{role_name}\" not found"
3183 )));
3184 }
3185 let member_of = client.query(sql::role_member_of(), &[&role_name]).await?;
3186 let has_members = client.query(sql::role_has_members(), &[&role_name]).await?;
3187 let privileges = client
3188 .query(sql::role_table_privileges(), &[&role_name])
3189 .await?;
3190
3191 Ok(ToolOutcome::ok_json(json!({
3192 "role": rows_to_json(&details).as_array().and_then(|a| a.first().cloned()).unwrap_or(Value::Null),
3193 "member_of": rows_to_json(&member_of),
3194 "has_members": rows_to_json(&has_members),
3195 "table_privileges": rows_to_json(&privileges),
3196 })))
3197 }
3198
3199 async fn export_query(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3200 let sql = args
3201 .get("sql")
3202 .and_then(|v| v.as_str())
3203 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
3204 require_select_or_with(&self.session.filter(), sql)?;
3205
3206 let format = args
3207 .get("format")
3208 .and_then(|v| v.as_str())
3209 .map(|s| {
3210 ExportFormat::parse(s).ok_or_else(|| {
3211 ToolError::InvalidArgs(format!(
3212 "Unsupported format \"{s}\". Use csv, json, or sqlinsert."
3213 ))
3214 })
3215 })
3216 .transpose()?
3217 .unwrap_or(ExportFormat::Csv);
3218
3219 let table_target = match args.get("table").and_then(|v| v.as_str()) {
3220 Some(t) if !t.trim().is_empty() => Some(parse_ref(t).map_err(ToolError::InvalidArgs)?),
3221 _ => None,
3222 };
3223
3224 if format == ExportFormat::SqlInsert && table_target.is_none() {
3225 return Err(ToolError::InvalidArgs(
3226 "table (schema.name) is required when format=sqlinsert".into(),
3227 ));
3228 }
3229
3230 let scope = self.execution_scope_from_args(args).await?;
3231 let max_rows = scope.caps.max_rows;
3232 let outcome = self
3233 .run_select_internal(
3234 sql,
3235 Some(max_rows),
3236 false,
3237 &[],
3238 &scope,
3239 RunSelectFormat::Json,
3240 scope.caps.statement_timeout_ms,
3241 false,
3242 )
3243 .await?;
3244 if outcome.is_error {
3245 return Ok(outcome);
3246 }
3247
3248 let structured = outcome.structured.unwrap_or(Value::Null);
3249 let rows_val = structured
3250 .get("rows")
3251 .cloned()
3252 .or_else(|| structured.get("data").and_then(|d| d.get("rows").cloned()))
3253 .unwrap_or(Value::Array(vec![]));
3254 let rows = rows_val.as_array().cloned().unwrap_or_default();
3255 let columns = columns_from_rows(&rows);
3256 let truncated = structured
3257 .get("truncated")
3258 .and_then(|v| v.as_bool())
3259 .unwrap_or(false);
3260
3261 let payload = match format {
3262 ExportFormat::Json => {
3263 let grid: Vec<Value> = rows
3268 .iter()
3269 .map(|row| {
3270 json!(
3271 columns
3272 .iter()
3273 .map(|c| row.get(c).cloned().unwrap_or(Value::Null))
3274 .collect::<Vec<_>>()
3275 )
3276 })
3277 .collect();
3278 json!({
3279 "format": format.as_str(),
3280 "rowCount": rows.len(),
3281 "truncated": truncated,
3282 "columns": columns,
3283 "rows": grid,
3284 })
3285 }
3286 ExportFormat::Csv => {
3287 let content = rows_to_csv(&rows, &columns);
3288 let caps = self.session.caps();
3289 let (char_trunc, content) = caps.truncate_chars(&content);
3290 json!({
3291 "format": format.as_str(),
3292 "rowCount": rows.len(),
3293 "truncated": truncated || char_trunc,
3294 "columns": columns,
3295 "content": content,
3296 })
3297 }
3298 ExportFormat::SqlInsert => {
3299 let (schema, table) = table_target.expect("checked above");
3300 let content = rows_to_sql_insert(&rows, &columns, &schema, &table);
3301 let caps = self.session.caps();
3302 let (char_trunc, content) = caps.truncate_chars(&content);
3303 json!({
3304 "format": format.as_str(),
3305 "rowCount": rows.len(),
3306 "truncated": truncated || char_trunc,
3307 "table": format!("{schema}.{table}"),
3308 "columns": columns,
3309 "content": content,
3310 })
3311 }
3312 };
3313
3314 Ok(ToolOutcome::ok_json(payload))
3315 }
3316
3317 async fn db_dashboard(&self) -> Result<ToolOutcome, ToolError> {
3318 let client = self.session.checkout().await?;
3319 let sections: &[(&str, &str)] = &[
3320 ("db_info", sql::dashboard_db_info()),
3321 ("connection_states", sql::connection_states()),
3322 ("top_tables", sql::dashboard_top_tables()),
3323 ("object_counts", sql::dashboard_object_counts()),
3324 ("active_queries", sql::dashboard_active_queries()),
3325 ("blocking_locks", sql::blocking_locks()),
3326 ("max_connections", sql::dashboard_max_connections()),
3327 ("extension_count", sql::dashboard_extension_count()),
3328 ("cache", sql::cache_hit_ratio()),
3329 ];
3330 let mut report = serde_json::Map::new();
3331 for (key, q) in sections {
3332 match client.query(*q, &[]).await {
3333 Ok(rows) => {
3334 report.insert((*key).into(), rows_to_json(&rows));
3335 }
3336 Err(e) => {
3337 report.insert((*key).into(), json!({ "error": e.to_string() }));
3338 }
3339 }
3340 }
3341
3342 for key in ["db_info", "object_counts", "extension_count", "cache"] {
3344 if let Some(Value::Array(arr)) = report.get(key).cloned()
3345 && arr.len() == 1
3346 {
3347 report.insert(key.into(), arr.into_iter().next().unwrap());
3348 }
3349 }
3350 if let Some(Value::Array(arr)) = report.get("max_connections").cloned()
3351 && let Some(row) = arr.first()
3352 {
3353 report.insert(
3354 "max_connections".into(),
3355 row.get("max_connections")
3356 .cloned()
3357 .unwrap_or_else(|| row.clone()),
3358 );
3359 }
3360
3361 Ok(ToolOutcome::ok_json(Value::Object(report)))
3362 }
3363
3364 async fn deep_plan_analysis(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3365 let sql = args
3366 .get("sql")
3367 .and_then(|v| v.as_str())
3368 .ok_or_else(|| ToolError::InvalidArgs("sql is required".into()))?;
3369 require_select_or_with(&self.session.filter(), sql)?;
3370 let analyze = args
3371 .get("analyze")
3372 .and_then(|v| v.as_bool())
3373 .unwrap_or(true);
3374 let explain = build_explain_sql(sql, analyze);
3375 let outcome = self.run_explain_in_transaction(&explain).await?;
3376 let rows = outcome.structured.unwrap_or(Value::Null);
3377 let row_array = rows
3378 .get("rows")
3379 .and_then(|v| v.as_array())
3380 .or_else(|| rows.as_array());
3381 let plan = row_array
3382 .and_then(|a| a.first())
3383 .and_then(|r| r.get("QUERY PLAN"))
3384 .cloned()
3385 .unwrap_or(Value::Null);
3386 let deep = analyze_deep_plan(&plan, sql)
3387 .or_else(|| analyze_deep_plan(&rows, sql))
3388 .ok_or_else(|| {
3389 ToolError::Execution("Could not parse EXPLAIN JSON plan for deep analysis".into())
3390 })?;
3391 let metrics = extract_plan_metrics(&plan).or_else(|| extract_plan_metrics(&rows));
3392 Ok(ToolOutcome::ok_json(json!({
3393 "deep": deep,
3394 "metrics": metrics,
3395 "plan": plan,
3396 "analyzed": analyze,
3397 })))
3398 }
3399
3400 async fn schema_diff(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3401 let source_schema = args
3402 .get("sourceSchema")
3403 .and_then(|v| v.as_str())
3404 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
3405 let target_schema = args
3406 .get("targetSchema")
3407 .and_then(|v| v.as_str())
3408 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
3409 crate::schema_diff::require_safe_schema(source_schema)?;
3410 crate::schema_diff::require_safe_schema(target_schema)?;
3411
3412 let client = self.session.checkout().await?;
3413 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
3414 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
3415 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
3416 let changed = diffs
3417 .iter()
3418 .filter(|d| d.status != crate::schema_diff::DiffStatus::Unchanged)
3419 .count();
3420 Ok(ToolOutcome::ok_json(json!({
3421 "sourceSchema": source_schema,
3422 "targetSchema": target_schema,
3423 "tableCount": diffs.len(),
3424 "changedCount": changed,
3425 "diffs": crate::schema_diff::diffs_to_json(&diffs),
3426 })))
3427 }
3428
3429 async fn generate_migration(&self, args: &Value) -> Result<ToolOutcome, ToolError> {
3430 let source_schema = args
3431 .get("sourceSchema")
3432 .and_then(|v| v.as_str())
3433 .ok_or_else(|| ToolError::InvalidArgs("sourceSchema is required".into()))?;
3434 let target_schema = args
3435 .get("targetSchema")
3436 .and_then(|v| v.as_str())
3437 .ok_or_else(|| ToolError::InvalidArgs("targetSchema is required".into()))?;
3438 crate::schema_diff::require_safe_schema(source_schema)?;
3439 crate::schema_diff::require_safe_schema(target_schema)?;
3440
3441 let client = self.session.checkout().await?;
3442 let source = crate::schema_diff::load_schema_snapshot(&client, source_schema).await?;
3443 let target = crate::schema_diff::load_schema_snapshot(&client, target_schema).await?;
3444 let diffs = crate::schema_diff::compute_schema_diff(&source, &target);
3445 let statements =
3446 crate::schema_diff::build_migration_statements(source_schema, target_schema, &diffs);
3447 let sql = if statements.is_empty() {
3448 format!("-- No differences between {source_schema} and {target_schema}")
3449 } else {
3450 statements.join("\n\n")
3451 };
3452 Ok(ToolOutcome::ok_json(json!({
3453 "sourceSchema": source_schema,
3454 "targetSchema": target_schema,
3455 "statementCount": statements.len(),
3456 "sql": sql,
3457 "hint": "Read-only: review and run via execute_sql / apply_ddl only with --access-mode write|admin. Destructive drops are commented out.",
3458 })))
3459 }
3460
3461 #[allow(clippy::too_many_arguments)]
3462 async fn finalize_run_select_payload(
3463 &self,
3464 sql: &str,
3465 mut payload: Value,
3466 scope: &ExecutionScope,
3467 format: RunSelectFormat,
3468 columnar: bool,
3469 resolve_fks: bool,
3470 client: &deadpool_postgres::Object,
3471 ) -> Result<ToolOutcome, ToolError> {
3472 if resolve_fks
3473 && let Some(store) = self.index_store()
3474 && let Ok(tables) = select_table_refs(sql)
3475 {
3476 let filter = self.query_filter_for(&scope.ctx.connection_id);
3477 let _ = resolve::resolve_fks_on_payload(
3478 &mut payload,
3479 store,
3480 &scope.ctx.connection_id,
3481 &scope.ctx.database,
3482 &tables,
3483 &filter,
3484 client,
3485 )
3486 .await;
3487 }
3488
3489 if format == RunSelectFormat::Csv {
3490 let row_objs = payload_to_row_objects(&payload);
3491 let cols = columns_from_rows(&row_objs);
3492 let csv = rows_to_csv(&row_objs, &cols);
3493 payload = json!({
3494 "format": "csv",
3495 "csv": csv,
3496 });
3497 } else if format == RunSelectFormat::Markdown {
3498 let (cols, row_vecs) = payload_to_columnar(&payload);
3499 let md = rows_to_markdown(&cols, &row_vecs);
3500 payload = json!({
3501 "format": "markdown",
3502 "markdown": md,
3503 "columns": cols,
3504 "rows": row_vecs,
3505 });
3506 } else if columnar {
3507 payload = columnarize_read_payload(payload);
3508 }
3509
3510 let text = if matches!(format, RunSelectFormat::Compact | RunSelectFormat::Csv) {
3511 serde_json::to_string(&payload)
3512 } else {
3513 serde_json::to_string_pretty(&payload)
3514 }
3515 .map_err(|e| ToolError::Execution(e.to_string()))?;
3516 let (trunc, text) = scope.caps.truncate_chars(&text);
3517 let structured = if trunc {
3518 json!({ "truncated_chars": true, "data": payload })
3519 } else {
3520 payload
3521 };
3522 Ok(Self::scope_tag(
3523 scope,
3524 ToolOutcome {
3525 text: text.to_string(),
3526 structured: Some(structured),
3527 is_error: false,
3528 },
3529 ))
3530 }
3531
3532 #[allow(clippy::too_many_arguments)]
3541 async fn run_select_internal(
3542 &self,
3543 sql: &str,
3544 max_rows: Option<u32>,
3545 columnar: bool,
3546 params: &[Value],
3547 scope: &ExecutionScope,
3548 format: RunSelectFormat,
3549 timeout_ms: u32,
3550 resolve_fks: bool,
3551 ) -> Result<ToolOutcome, ToolError> {
3552 let (client, _) = self
3553 .session
3554 .checkout_for(CheckoutTarget::Scoped(&scope.ctx))
3555 .await?;
3556 ToolSession::set_statement_timeout(&client, timeout_ms).await?;
3557 let pg_params = sql_param_boxes(params)?;
3558 let pg_param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = pg_params
3559 .iter()
3560 .map(|p| p.as_ref() as &(dyn tokio_postgres::types::ToSql + Sync))
3561 .collect();
3562
3563 let Some(max_rows) = max_rows else {
3564 let rows = self
3565 .query_with_hints(&client, sql, &pg_param_refs, timeout_ms)
3566 .await?;
3567 let values = rows_to_json(&rows);
3568 let payload =
3569 self.apply_pii_redaction_for(sql, ensure_structured_object(values), &scope.filter);
3570 return self
3571 .finalize_run_select_payload(
3572 sql,
3573 payload,
3574 scope,
3575 format,
3576 columnar,
3577 resolve_fks,
3578 &client,
3579 )
3580 .await;
3581 };
3582
3583 let cleaned = sql.trim().trim_end_matches(';').trim();
3584 let wrapped = format!(
3585 "SELECT sub.*, COUNT(*) OVER() AS {NEXQL_TOTAL_COUNT_COL} FROM ({cleaned}) AS sub LIMIT {}",
3586 max_rows + 1
3587 );
3588 let rows = self
3589 .query_with_hints(&client, &wrapped, &pg_param_refs, timeout_ms)
3590 .await?;
3591 let truncated = rows.len() as u32 > max_rows;
3592 let keep_len = if truncated {
3593 max_rows as usize
3594 } else {
3595 rows.len()
3596 };
3597 let (total_count, values) =
3598 rows_to_json_array_with_total(&rows[..keep_len], NEXQL_TOTAL_COUNT_COL);
3599 let returned = keep_len;
3600 let mut payload =
3601 self.apply_pii_redaction_for(sql, ensure_structured_object(values), &scope.filter);
3602 if let Some(obj) = payload.as_object_mut() {
3603 obj.insert("limit".into(), json!(max_rows));
3604 if let Some(total) = total_count {
3605 obj.insert("total_count".into(), json!(total));
3606 obj.insert("has_more".into(), json!(total > max_rows as i64));
3607 } else if truncated {
3608 obj.insert("has_more".into(), json!(true));
3609 } else {
3610 obj.insert("has_more".into(), json!(false));
3611 obj.insert("total_count".into(), json!(returned));
3612 }
3613 if truncated {
3614 obj.insert("truncated".into(), json!(true));
3615 }
3616 }
3617 self.finalize_run_select_payload(
3618 sql,
3619 payload,
3620 scope,
3621 format,
3622 columnar,
3623 resolve_fks,
3624 &client,
3625 )
3626 .await
3627 }
3628
3629 async fn query_with_hints(
3630 &self,
3631 client: &deadpool_postgres::Object,
3632 sql: &str,
3633 params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
3634 timeout_ms: u32,
3635 ) -> Result<Vec<tokio_postgres::Row>, ToolError> {
3636 match client.query(sql, params).await {
3637 Ok(rows) => Ok(rows),
3638 Err(e) if is_statement_timeout(&e) => Err(ToolError::Execution(
3639 serde_json::to_string(&json!({
3640 "error": "statement_timeout",
3641 "timeout_ms": timeout_ms,
3642 "sql_preview": sql.chars().take(200).collect::<String>(),
3643 "hint": "Narrow the query, add indexes, pass a higher timeout_ms, or use terminate_query in admin mode.",
3644 }))
3645 .unwrap_or_else(|_| format!("statement timeout after {timeout_ms}ms")),
3646 )),
3647 Err(e) => Err(ToolError::Execution(self.enrich_query_error(&e).await)),
3648 }
3649 }
3650
3651 fn apply_pii_redaction_for(
3652 &self,
3653 sql: &str,
3654 mut payload: Value,
3655 filter: &PolicyFilter,
3656 ) -> Value {
3657 if filter.pii_columns.is_empty() {
3658 return payload;
3659 }
3660 let Ok(tables) = select_table_refs(sql) else {
3661 return payload;
3662 };
3663 let (redacted, cols) = redact_pii_in_payload(payload, &filter.pii_columns, &tables);
3664 payload = redacted;
3665 if !cols.is_empty()
3666 && let Some(obj) = payload.as_object_mut()
3667 {
3668 obj.insert("piiRedactedColumns".into(), json!(cols));
3669 }
3670 payload
3671 }
3672}
3673
3674fn is_statement_timeout(err: &tokio_postgres::Error) -> bool {
3675 err.code()
3676 .map(|code| code.code() == "57014")
3677 .unwrap_or(false)
3678}
3679
3680fn payload_to_row_objects(payload: &Value) -> Vec<Value> {
3681 if let Some(rows) = payload.get("rows").and_then(|v| v.as_array()) {
3682 if rows
3683 .first()
3684 .and_then(|r| r.as_object())
3685 .is_some()
3686 {
3687 return rows.clone();
3688 }
3689 if let Some(cols) = payload.get("columns").and_then(|v| v.as_array()) {
3690 let col_names: Vec<String> = cols
3691 .iter()
3692 .filter_map(|c| c.as_str().map(str::to_owned))
3693 .collect();
3694 return rows
3695 .iter()
3696 .filter_map(|row| row.as_array())
3697 .map(|cells| {
3698 let mut obj = serde_json::Map::new();
3699 for (idx, name) in col_names.iter().enumerate() {
3700 obj.insert(
3701 name.clone(),
3702 cells.get(idx).cloned().unwrap_or(Value::Null),
3703 );
3704 }
3705 Value::Object(obj)
3706 })
3707 .collect();
3708 }
3709 }
3710 Vec::new()
3711}
3712
3713fn payload_to_columnar(payload: &Value) -> (Vec<String>, Vec<Vec<Value>>) {
3714 if let (Some(cols), Some(rows)) = (
3715 payload.get("columns").and_then(|v| v.as_array()),
3716 payload.get("rows").and_then(|v| v.as_array()),
3717 ) {
3718 let col_names: Vec<String> = cols
3719 .iter()
3720 .filter_map(|c| c.as_str().map(str::to_owned))
3721 .collect();
3722 let row_vecs: Vec<Vec<Value>> = rows
3723 .iter()
3724 .filter_map(|r| r.as_array().cloned())
3725 .collect();
3726 return (col_names, row_vecs);
3727 }
3728 let row_objs = payload_to_row_objects(payload);
3729 let cols = columns_from_rows(&row_objs);
3730 let row_vecs: Vec<Vec<Value>> = row_objs
3731 .iter()
3732 .map(|row| {
3733 cols.iter()
3734 .map(|c| row.get(c).cloned().unwrap_or(Value::Null))
3735 .collect()
3736 })
3737 .collect();
3738 (cols, row_vecs)
3739}
3740
3741fn normalize_for_match(s: &str) -> String {
3743 let mut out = String::new();
3744 let mut last_was_sep = true;
3745 for ch in s.to_lowercase().chars() {
3746 if ch.is_ascii_alphanumeric() {
3747 out.push(ch);
3748 last_was_sep = false;
3749 } else if !last_was_sep {
3750 out.push(' ');
3751 last_was_sep = true;
3752 }
3753 }
3754 out.trim().to_string()
3755}
3756
3757fn fuzzy_score(hint: &str, candidate: &str) -> f64 {
3759 let h = normalize_for_match(hint);
3760 let c = normalize_for_match(candidate);
3761 if h.is_empty() || c.is_empty() {
3762 return 0.0;
3763 }
3764 if h == c {
3765 return 100.0;
3766 }
3767 if c.contains(&h) || h.contains(&c) {
3768 return 75.0;
3769 }
3770 let h_tokens: std::collections::HashSet<&str> =
3771 h.split(' ').filter(|s| !s.is_empty()).collect();
3772 let c_tokens: std::collections::HashSet<&str> =
3773 c.split(' ').filter(|s| !s.is_empty()).collect();
3774 let overlap = h_tokens.intersection(&c_tokens).count();
3775 if overlap == 0 {
3776 return 0.0;
3777 }
3778 (overlap as f64 / h_tokens.len().max(c_tokens.len()) as f64) * 60.0
3779}
3780
3781fn policy_to_query_filter(filter: &PolicyFilter) -> QueryPolicyFilter {
3782 QueryPolicyFilter {
3783 allow_schemas: filter.allow_schemas.clone(),
3784 deny_schemas: filter.deny_schemas.clone(),
3785 deny_tables: filter.deny_tables.clone(),
3786 pii_columns: filter.pii_columns.clone(),
3787 }
3788}
3789
3790fn require_select_or_with(filter: &PolicyFilter, sql: &str) -> Result<(), ToolError> {
3791 match validate_readonly_sql(sql)? {
3792 SqlDecision::Allow => {}
3793 SqlDecision::Reject => {
3794 return Err(ToolError::Execution(
3795 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
3796 ));
3797 }
3798 }
3799 enforce_read_table_policy(filter, sql)?;
3800 let trimmed = sql.trim().to_ascii_lowercase();
3801 if !(trimmed.starts_with("select") || trimmed.starts_with("with")) {
3802 return Err(ToolError::Execution(
3803 "Security Error: Only SELECT or WITH statements can be analyzed.".into(),
3804 ));
3805 }
3806 Ok(())
3807}
3808
3809fn rows_to_json(rows: &[tokio_postgres::Row]) -> Value {
3810 rows_to_json_array(rows)
3811}
3812
3813fn scores_equal(a: f64, b: f64) -> bool {
3814 (a - b).abs() <= f64::EPSILON * a.abs().max(b.abs()).max(1.0)
3815}
3816
3817fn read_recent_log_errors() -> Vec<String> {
3818 let path = std::env::var("NEXQL_MCP_LOG")
3819 .map(std::path::PathBuf::from)
3820 .ok()
3821 .or_else(|| {
3822 std::env::var_os("HOME").map(|h| {
3823 std::path::PathBuf::from(h)
3824 .join(".config")
3825 .join("nexql-mcp")
3826 .join("logs")
3827 .join("nexql-mcp.log")
3828 })
3829 });
3830
3831 let Some(log_path) = path else {
3832 return Vec::new();
3833 };
3834
3835 let Ok(content) = std::fs::read_to_string(&log_path) else {
3836 return Vec::new();
3837 };
3838
3839 content
3840 .lines()
3841 .rev()
3842 .take(50)
3843 .filter(|line| {
3844 line.contains("ERROR")
3845 || line.contains("WARN")
3846 || line.contains("failed")
3847 || line.contains("Error")
3848 })
3849 .map(String::from)
3850 .collect()
3851}
3852
3853fn parse_sql_params(args: &Value) -> Vec<Value> {
3854 args.get("params")
3855 .and_then(|v| v.as_array())
3856 .map(|a| a.to_vec())
3857 .unwrap_or_default()
3858}
3859
3860fn sql_param_boxes(params: &[Value]) -> Result<Vec<Box<dyn tokio_postgres::types::ToSql + Sync + Send>>, ToolError> {
3861 params
3862 .iter()
3863 .map(json_to_sql_param)
3864 .collect::<Result<Vec<_>, _>>()
3865}
3866
3867fn json_to_sql_param(
3868 v: &Value,
3869) -> Result<Box<dyn tokio_postgres::types::ToSql + Sync + Send>, ToolError> {
3870 match v {
3871 Value::Null => Ok(Box::new(None::<String>)),
3872 Value::Bool(b) => Ok(Box::new(*b)),
3873 Value::Number(n) => {
3874 if let Some(i) = n.as_i64() {
3875 Ok(Box::new(i))
3876 } else if let Some(f) = n.as_f64() {
3877 Ok(Box::new(f))
3878 } else {
3879 Err(ToolError::InvalidArgs("invalid numeric param".into()))
3880 }
3881 }
3882 Value::String(s) => Ok(Box::new(s.clone())),
3883 _ => Err(ToolError::InvalidArgs(
3884 "params must be string, number, boolean, or null".into(),
3885 )),
3886 }
3887}
3888
3889#[cfg(test)]
3890mod tests {
3891 use super::*;
3892 use crate::plan::build_explain_sql;
3893 use nexql_policy::PolicyFilter;
3894 use serde_json::json;
3895
3896 use crate::session::{ConnectionInfo, ConnectionPolicy, ToolSession};
3897 use nexql_policy::{AccessMode, PolicyCaps};
3898
3899 static CONFIG_ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
3905
3906 fn test_conn() -> ConnectionInfo {
3907 ConnectionInfo {
3908 id: "conn-1".into(),
3909 name: "conn-1".into(),
3910 host: Some("127.0.0.1".into()),
3911 port: Some(5432),
3912 database: Some("appdb".into()),
3913 params: Default::default(),
3914 policy: ConnectionPolicy {
3915 access_mode: AccessMode::Read,
3916 caps: PolicyCaps::default(),
3917 filter: PolicyFilter::default(),
3918 environment: None,
3919 },
3920 }
3921 }
3922
3923 #[test]
3924 fn scores_equal_treats_near_duplicates_as_tied() {
3925 let s = 3.295836866004329_f64;
3926 assert!(super::scores_equal(s, s));
3927 assert!(super::scores_equal(s, s + f64::EPSILON));
3928 }
3929
3930 #[test]
3931 fn policy_maps_one_to_one() {
3932 let f = PolicyFilter {
3933 allow_schemas: vec!["public".into()],
3934 deny_schemas: vec!["pgboss".into()],
3935 deny_tables: vec!["auth.*".into()],
3936 pii_columns: vec!["public.users.ssn".into()],
3937 };
3938 let q = policy_to_query_filter(&f);
3939 assert_eq!(q.allow_schemas, f.allow_schemas);
3940 assert_eq!(q.deny_schemas, f.deny_schemas);
3941 assert_eq!(q.deny_tables, f.deny_tables);
3942 assert_eq!(q.pii_columns, f.pii_columns);
3943 }
3944
3945 #[test]
3946 fn ok_json_wraps_arrays_for_cursor_structured_content() {
3947 let out = ToolOutcome::ok_json(json!([{ "id": 1 }, { "id": 2 }]));
3948 assert!(!out.is_error);
3949 let s = out.structured.as_ref().unwrap();
3950 assert!(s.is_object(), "structuredContent must be object, got {s}");
3951 assert_eq!(s["rows"].as_array().unwrap().len(), 2);
3952 assert!(out.text.contains("\"rows\""));
3953 }
3954
3955 #[test]
3956 fn ok_json_leaves_objects_unchanged() {
3957 let out = ToolOutcome::ok_json(json!({ "kind": "table", "name": "orders" }));
3958 let s = out.structured.as_ref().unwrap();
3959 assert_eq!(s["kind"], "table");
3960 assert!(s.get("rows").is_none());
3961 }
3962
3963 #[test]
3964 fn router_specs_include_phase4_and_phase9() {
3965 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
3966 let router = ToolRouter::with_index_store(session, None);
3967 assert_eq!(router.specs().len(), ToolName::ACTIVE.len());
3968 let names: Vec<_> = router.specs().iter().map(|s| s.name.as_str()).collect();
3969 assert!(names.contains(&"search_schema"));
3970 assert!(names.contains(&"get_ddl"));
3971 assert!(names.contains(&"deep_plan_analysis"));
3972 assert!(names.contains(&"get_index_status"));
3973 assert!(names.contains(&"list_extensions"));
3974 assert!(names.contains(&"server_settings"));
3975 assert!(names.contains(&"suggest_indexes"));
3976 assert!(names.contains(&"find_unused_indexes"));
3977 assert!(names.contains(&"bloat_report"));
3978 assert!(names.contains(&"find_missing_fks"));
3979 assert!(names.contains(&"export_query"));
3980 assert!(names.contains(&"list_roles"));
3981 assert!(names.contains(&"db_dashboard"));
3982 assert!(names.contains(&"deep_plan_analysis"));
3983 assert!(names.contains(&"execute_sql"));
3984 assert!(names.contains(&"edit_row"));
3985 assert!(names.contains(&"import_data"));
3986 assert!(names.contains(&"apply_ddl"));
3987 assert!(names.contains(&"create_index_concurrently"));
3988 assert!(names.contains(&"run_maintenance"));
3989 assert!(names.contains(&"terminate_query"));
3990 }
3991
3992 #[tokio::test]
3993 async fn write_tools_refuse_read_mode() {
3994 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
3995 let router = ToolRouter::with_index_store(session, None);
3996 for tool in [
3997 "execute_sql",
3998 "edit_row",
3999 "import_data",
4000 "apply_ddl",
4001 "create_index_concurrently",
4002 "run_maintenance",
4003 "terminate_query",
4004 ] {
4005 let out = router
4006 .call(tool, json!({ "sql": "SELECT 1", "table": "public.t", "rows": [], "action": "insert", "values": {}, "pid": 1 }))
4007 .await;
4008 assert!(out.is_error, "{tool}: {}", out.text);
4009 assert!(
4010 out.text.contains("write") || out.text.contains("admin"),
4011 "{tool}: {}",
4012 out.text
4013 );
4014 }
4015 }
4016
4017 #[tokio::test]
4018 async fn table_stats_rejects_injection_ref() {
4019 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4020 let router = ToolRouter::with_index_store(session, None);
4021 let out = router
4022 .call("table_stats", json!({ "ref": "public.users; DROP" }))
4023 .await;
4024 assert!(out.is_error, "{}", out.text);
4025 assert!(
4026 out.text.contains("Invalid object reference") || out.text.contains("invalid arguments"),
4027 "expected ref validation error, got: {}",
4028 out.text
4029 );
4030 }
4031
4032 #[test]
4033 fn explain_transaction_path_builds_readonly_sequence() {
4034 let explain = build_explain_sql("SELECT 1", true);
4036 assert!(explain.starts_with("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON)"));
4037 assert!(!explain.to_ascii_lowercase().contains("commit"));
4038 let steps = ["BEGIN", "SET TRANSACTION READ ONLY", &explain, "ROLLBACK"];
4039 assert_eq!(steps.len(), 4);
4040 assert_eq!(steps[0], "BEGIN");
4041 assert_eq!(steps[1], "SET TRANSACTION READ ONLY");
4042 assert_eq!(steps[3], "ROLLBACK");
4043 }
4044
4045 #[tokio::test]
4046 async fn missing_index_returns_actionable_error() {
4047 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4048 let router = ToolRouter::with_index_store(session, None);
4049 let out = router
4050 .call("search_schema", json!({ "query": "users" }))
4051 .await;
4052 assert!(out.is_error, "{}", out.text);
4053 assert!(
4054 out.text.contains("rebuild_index"),
4055 "expected actionable hint, got: {}",
4056 out.text
4057 );
4058 }
4059
4060 fn write_join_path_fixture(store: &IndexStore) {
4064 use nexql_index::{
4065 BuildDepth, BuildMode, ColumnEntry, DbObjectKind, IndexCounts, IndexDerived,
4066 IndexManifest, IndexScope, IndexStats, JOIN_GRAPH_FILE, JoinEdge, JoinGraph,
4067 ObjectEntry, ObjectShard, TOKENS_FILE,
4068 };
4069 use std::collections::HashMap;
4070
4071 let base = store.base_dir("conn-1", "appdb");
4072 let manifest = IndexManifest {
4073 format_version: 1,
4074 connection_id: "conn-1".into(),
4075 database: "appdb".into(),
4076 indexed_at: "2026-08-08T00:00:00.000Z".into(),
4077 build_mode: BuildMode::Auto,
4078 build_depth: BuildDepth::Structure,
4079 schema_fingerprint: "fp".into(),
4080 pg_version: "18.4".into(),
4081 environment: "development".into(),
4082 scope: IndexScope {
4083 included_schemas: vec!["public".into()],
4084 excluded_objects: vec![],
4085 pii_excluded_columns: vec![],
4086 },
4087 counts: IndexCounts {
4088 tables: 3,
4089 views: 0,
4090 functions: 0,
4091 enums: 0,
4092 },
4093 shards: vec![ObjectShard {
4094 file: "objects-public-0.json".into(),
4095 schema: "public".into(),
4096 objects: 3,
4097 bytes: 512,
4098 hash: "abc".into(),
4099 }],
4100 derived: IndexDerived {
4101 tokens: TOKENS_FILE.into(),
4102 join_graph: JOIN_GRAPH_FILE.into(),
4103 values: None,
4104 embeddings: None,
4105 embeddings_meta: None,
4106 },
4107 stats: IndexStats {
4108 build_ms: 1,
4109 queries_run: 1,
4110 warnings: vec![],
4111 },
4112 };
4113 store.write_manifest(&base, &manifest).unwrap();
4114
4115 fn entry(oid: u32) -> ObjectEntry {
4116 ObjectEntry {
4117 kind: DbObjectKind::Table,
4118 oid,
4119 object_hash: format!("hash{oid}"),
4120 comment: None,
4121 row_estimate: 10.0,
4122 size_bytes: 8192,
4123 columns: vec![ColumnEntry {
4124 name: "id".into(),
4125 type_name: "integer".into(),
4126 not_null: true,
4127 default_value: None,
4128 comment: None,
4129 ordinal: 1,
4130 is_pk: Some(true),
4131 profile: None,
4132 pii: None,
4133 }],
4134 primary_key: Some(vec!["id".into()]),
4135 foreign_keys: None,
4136 indexes: None,
4137 checks: None,
4138 excluded: None,
4139 definition: None,
4140 signature: None,
4141 language: None,
4142 volatility: None,
4143 body: None,
4144 values: None,
4145 base_type: None,
4146 constraint: None,
4147 }
4148 }
4149 let mut orders = entry(1);
4152 orders.columns.push(nexql_index::ColumnEntry {
4153 name: "status".into(),
4154 type_name: "text".into(),
4155 not_null: true,
4156 default_value: None,
4157 comment: None,
4158 ordinal: 2,
4159 is_pk: None,
4160 profile: Some(nexql_index::ColumnProfile {
4161 n_distinct: 2.0,
4162 null_frac: 0.0,
4163 common_values: Some(vec!["pending".into(), "paid".into()]),
4164 min: None,
4165 max: None,
4166 }),
4167 pii: None,
4168 });
4169 orders.columns.push(nexql_index::ColumnEntry {
4172 name: "amount".into(),
4173 type_name: "numeric".into(),
4174 not_null: false,
4175 default_value: None,
4176 comment: None,
4177 ordinal: 3,
4178 is_pk: None,
4179 profile: Some(nexql_index::ColumnProfile {
4180 n_distinct: 8.0,
4181 null_frac: 0.25,
4182 common_values: None,
4183 min: None,
4184 max: None,
4185 }),
4186 pii: None,
4187 });
4188
4189 let mut shard = HashMap::new();
4190 shard.insert("public.orders".into(), orders);
4191 shard.insert("public.customers".into(), entry(2));
4192 shard.insert("public.order_items".into(), entry(3));
4193 store
4194 .write_shard_entries(&base, "objects-public-0.json", &shard)
4195 .unwrap();
4196
4197 let graph = JoinGraph {
4198 edges: vec![
4199 JoinEdge {
4200 from: "public.order_items".into(),
4201 to: "public.orders".into(),
4202 via: "order_items_order_id_fkey".into(),
4203 cols: vec![("order_id".into(), "id".into())],
4204 inferred: None,
4205 disabled: None,
4206 },
4207 JoinEdge {
4208 from: "public.orders".into(),
4209 to: "public.customers".into(),
4210 via: "orders_customer_id_fkey".into(),
4211 cols: vec![("customer_id".into(), "id".into())],
4212 inferred: None,
4213 disabled: None,
4214 },
4215 ],
4216 };
4217 store.write_join_graph(&base, &graph).unwrap();
4218 }
4219
4220 fn join_path_router() -> ToolRouter {
4221 let tmp = tempfile::TempDir::new().unwrap();
4222 let store = IndexStore::new(tmp.path());
4223 write_join_path_fixture(&store);
4224 std::mem::forget(tmp); let session = ToolSession::for_tests(
4226 vec![test_conn()],
4227 PolicyFilter::default(),
4228 Some(IndexStore::new(store.root())),
4229 );
4230 ToolRouter::with_index_store(session, Some(store))
4231 }
4232
4233 #[test]
4236 fn columnarize_outcome_reshapes_flat_rows_for_generic_tool() {
4237 let outcome = ToolOutcome::ok_json(json!({
4238 "rows": [{ "name": "pgcrypto" }, { "name": "pg_stat_statements" }]
4239 }));
4240 let out = ToolRouter::columnarize_outcome("list_extensions", outcome);
4241 let structured = out.structured.unwrap();
4242 assert_eq!(structured["columns"], json!(["name"]));
4243 assert_eq!(
4244 structured["rows"],
4245 json!([["pgcrypto"], ["pg_stat_statements"]])
4246 );
4247 assert!(out.text.contains("pgcrypto"));
4249 assert!(
4250 !out.text.contains('\n'),
4251 "wire text should be compact, not pretty-printed"
4252 );
4253 }
4254
4255 #[test]
4262 fn build_tuning_summary_reads_row_object_shaped_suggestions() {
4263 let suggestions = json!({
4264 "high_seq_scan_tables": [{ "table_name": "orders" }],
4265 "unindexed_fk_columns": [{ "column_name": "customer_id" }],
4266 });
4267 let summary = ToolRouter::build_tuning_summary(&None, &suggestions);
4268 assert!(summary.contains("2 index recommendation"), "{summary}");
4269 }
4270
4271 #[test]
4274 fn columnarize_outcome_skips_orient_and_get_join_path() {
4275 for tool in ["orient", "get_join_path"] {
4276 let outcome = ToolOutcome::ok_json(json!({
4277 "tables": [{ "ref": "public.orders" }, { "ref": "public.customers" }]
4278 }));
4279 let out = ToolRouter::columnarize_outcome(tool, outcome);
4280 let structured = out.structured.unwrap();
4281 assert_eq!(
4282 structured["tables"],
4283 json!([{ "ref": "public.orders" }, { "ref": "public.customers" }]),
4284 "{tool} should be excluded from columnarization"
4285 );
4286 }
4287 }
4288
4289 #[tokio::test]
4292 async fn attach_critique_flags_limit_without_order_by() {
4293 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4294 let router = ToolRouter::with_index_store(session, None);
4295 let outcome = ToolOutcome::ok_json(json!({ "columns": ["n"], "rows": [[1]] }));
4296 let out = router
4297 .attach_critique("SELECT * FROM orders LIMIT 10", outcome)
4298 .await;
4299 let structured = out.structured.unwrap();
4300 let critique = structured["critique"].as_array().unwrap();
4301 assert!(
4302 critique
4303 .iter()
4304 .any(|c| c["signal"] == "limit_without_order_by")
4305 );
4306 }
4307
4308 #[tokio::test]
4309 async fn attach_critique_silent_when_nothing_fires() {
4310 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4311 let router = ToolRouter::with_index_store(session, None);
4312 let outcome = ToolOutcome::ok_json(json!({ "columns": ["n"], "rows": [[1]] }));
4313 let out = router
4314 .attach_critique("SELECT * FROM orders ORDER BY id LIMIT 10", outcome)
4315 .await;
4316 let structured = out.structured.unwrap();
4317 assert!(structured.get("critique").is_none());
4318 }
4319
4320 #[tokio::test]
4324 async fn attach_critique_zero_rows_suggests_observed_values() {
4325 let router = join_path_router();
4326 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": [] }));
4327 let out = router
4328 .attach_critique(
4329 "SELECT * FROM public.orders WHERE status = 'complete'",
4330 outcome,
4331 )
4332 .await;
4333 let structured = out.structured.unwrap();
4334 let critique = structured["critique"].as_array().unwrap();
4335 let zero_rows = critique
4336 .iter()
4337 .find(|c| c["signal"] == "zero_rows")
4338 .expect("zero_rows critique");
4339 let msg = zero_rows["message"].as_str().unwrap();
4340 assert!(msg.contains("status = 'complete'"), "{msg}");
4341 assert!(msg.contains("pending") && msg.contains("paid"), "{msg}");
4342 }
4343
4344 #[tokio::test]
4345 async fn attach_critique_zero_rows_silent_without_index() {
4346 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4347 let router = ToolRouter::with_index_store(session, None);
4348 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": [] }));
4349 let out = router
4350 .attach_critique(
4351 "SELECT * FROM public.orders WHERE status = 'complete'",
4352 outcome,
4353 )
4354 .await;
4355 let structured = out.structured.unwrap();
4356 assert!(structured.get("critique").is_none());
4357 }
4358
4359 #[tokio::test]
4364 async fn attach_critique_flags_join_fan_out() {
4365 let router = join_path_router();
4366 let rows: Vec<Value> = (0..25).map(|i| json!([i])).collect();
4367 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": rows }));
4368 let out = router
4369 .attach_critique("SELECT id FROM public.orders", outcome)
4370 .await;
4371 let structured = out.structured.unwrap();
4372 let critique = structured["critique"].as_array().unwrap();
4373 let fan_out = critique
4374 .iter()
4375 .find(|c| c["signal"] == "join_fan_out")
4376 .expect("join_fan_out critique");
4377 assert!(
4378 fan_out["message"].as_str().unwrap().contains("25"),
4379 "{}",
4380 fan_out["message"]
4381 );
4382 }
4383
4384 #[tokio::test]
4385 async fn attach_critique_silent_when_row_count_within_estimate() {
4386 let router = join_path_router();
4387 let rows: Vec<Value> = (0..5).map(|i| json!([i])).collect();
4388 let outcome = ToolOutcome::ok_json(json!({ "columns": ["id"], "rows": rows }));
4389 let out = router
4390 .attach_critique("SELECT id FROM public.orders", outcome)
4391 .await;
4392 let structured = out.structured.unwrap();
4393 assert!(
4394 structured
4395 .get("critique")
4396 .and_then(|c| c.as_array())
4397 .map(|a| !a.iter().any(|c| c["signal"] == "join_fan_out"))
4398 .unwrap_or(true)
4399 );
4400 }
4401
4402 #[tokio::test]
4405 async fn attach_critique_flags_null_skipping_aggregate() {
4406 let router = join_path_router();
4407 let outcome = ToolOutcome::ok_json(json!({ "columns": ["avg"], "rows": [[42]] }));
4408 let out = router
4409 .attach_critique("SELECT AVG(amount) FROM public.orders", outcome)
4410 .await;
4411 let structured = out.structured.unwrap();
4412 let critique = structured["critique"].as_array().unwrap();
4413 let signal = critique
4414 .iter()
4415 .find(|c| c["signal"] == "null_skipping_aggregate")
4416 .expect("null_skipping_aggregate critique");
4417 assert!(
4418 signal["message"].as_str().unwrap().contains("AVG(amount)"),
4419 "{}",
4420 signal["message"]
4421 );
4422 }
4423
4424 #[tokio::test]
4425 async fn attach_critique_silent_for_aggregate_without_nulls() {
4426 let router = join_path_router();
4427 let outcome = ToolOutcome::ok_json(json!({ "columns": ["c"], "rows": [[2]] }));
4429 let out = router
4430 .attach_critique("SELECT COUNT(status) FROM public.orders", outcome)
4431 .await;
4432 let structured = out.structured.unwrap();
4433 assert!(structured.get("critique").is_none());
4434 }
4435
4436 #[tokio::test]
4440 async fn attach_dml_critique_zero_rows_affected_suggests_observed_values() {
4441 let router = join_path_router();
4442 let outcome = ToolOutcome::ok_json(json!({
4443 "dry_run": false,
4444 "rolled_back": false,
4445 "rows_affected": 0,
4446 "rows": [],
4447 }));
4448 let out = router
4449 .attach_dml_critique(
4450 "UPDATE public.orders SET status = 'paid' WHERE status = 'complete'",
4451 outcome,
4452 )
4453 .await;
4454 let structured = out.structured.unwrap();
4455 let critique = structured["critique"].as_array().unwrap();
4456 let zero_rows = critique
4457 .iter()
4458 .find(|c| c["signal"] == "zero_rows")
4459 .expect("zero_rows critique");
4460 let msg = zero_rows["message"].as_str().unwrap();
4461 assert!(msg.contains("status = 'complete'"), "{msg}");
4462 assert!(msg.contains("pending") && msg.contains("paid"), "{msg}");
4463 }
4464
4465 #[tokio::test]
4466 async fn attach_dml_critique_silent_when_rows_affected() {
4467 let router = join_path_router();
4468 let outcome = ToolOutcome::ok_json(json!({
4469 "dry_run": false,
4470 "rolled_back": false,
4471 "rows_affected": 3,
4472 "rows": [],
4473 }));
4474 let out = router
4475 .attach_dml_critique(
4476 "UPDATE public.orders SET status = 'paid' WHERE status = 'complete'",
4477 outcome,
4478 )
4479 .await;
4480 let structured = out.structured.unwrap();
4481 assert!(structured.get("critique").is_none());
4482 }
4483
4484 #[tokio::test]
4485 async fn attach_dml_critique_silent_for_delete_without_filter() {
4486 let router = join_path_router();
4487 let outcome = ToolOutcome::ok_json(json!({
4488 "dry_run": false,
4489 "rolled_back": false,
4490 "rows_affected": 0,
4491 "rows": [],
4492 }));
4493 let out = router
4494 .attach_dml_critique("DELETE FROM public.orders", outcome)
4495 .await;
4496 let structured = out.structured.unwrap();
4497 assert!(structured.get("critique").is_none());
4498 }
4499
4500 #[tokio::test]
4504 async fn get_join_path_resolves_unqualified_names() {
4505 let router = join_path_router();
4506 let out = router
4507 .call(
4508 "get_join_path",
4509 json!({ "a": "order_items", "b": "orders" }),
4510 )
4511 .await;
4512 assert!(!out.is_error, "{}", out.text);
4513 let structured = out.structured.unwrap();
4514 assert_eq!(structured["resolved_a"], "public.order_items");
4515 assert_eq!(structured["resolved_b"], "public.orders");
4516 assert_eq!(structured["path"][0]["via"], "order_items_order_id_fkey");
4517 }
4518
4519 #[tokio::test]
4520 async fn get_join_path_qualified_still_works() {
4521 let router = join_path_router();
4522 let out = router
4523 .call(
4524 "get_join_path",
4525 json!({ "a": "public.order_items", "b": "public.customers" }),
4526 )
4527 .await;
4528 assert!(!out.is_error, "{}", out.text);
4529 let structured = out.structured.unwrap();
4530 assert_eq!(structured["path"].as_array().unwrap().len(), 2);
4531 }
4532
4533 #[tokio::test]
4534 async fn get_join_path_unknown_name_errors_with_suggestion() {
4535 let router = join_path_router();
4536 let out = router
4537 .call("get_join_path", json!({ "a": "ordrs", "b": "customers" }))
4538 .await;
4539 assert!(out.is_error, "{}", out.text);
4540 assert!(
4541 out.text.contains("did you mean") && out.text.contains("public.orders"),
4542 "expected a did-you-mean suggestion, got: {}",
4543 out.text
4544 );
4545 }
4546
4547 #[tokio::test]
4551 async fn get_index_status_reports_missing_without_erroring_no_store() {
4552 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4553 let router = ToolRouter::with_index_store(session, None);
4554 let out = router.call("get_index_status", json!({})).await;
4555 assert!(!out.is_error, "{}", out.text);
4556 let structured = out.structured.unwrap();
4557 assert_eq!(structured["status"], "missing");
4558 assert_eq!(structured["remediation"], "rebuild_index");
4559 assert_eq!(structured["database"], "appdb");
4560 }
4561
4562 #[tokio::test]
4563 async fn get_index_status_reports_missing_without_erroring_empty_manifest() {
4564 let tmp = tempfile::TempDir::new().unwrap();
4565 let store = IndexStore::new(tmp.path());
4566 let session = ToolSession::for_tests(
4567 vec![test_conn()],
4568 PolicyFilter::default(),
4569 Some(IndexStore::new(tmp.path())),
4570 );
4571 let router = ToolRouter::with_index_store(session, Some(store));
4572 let out = router.call("get_index_status", json!({})).await;
4573 assert!(!out.is_error, "{}", out.text);
4574 let structured = out.structured.unwrap();
4575 assert_eq!(structured["status"], "missing");
4576 assert_eq!(structured["remediation"], "rebuild_index");
4577 }
4578
4579 #[tokio::test]
4580 async fn get_index_status_reports_ok_when_indexed() {
4581 let router = join_path_router();
4582 let out = router.call("get_index_status", json!({})).await;
4583 assert!(!out.is_error, "{}", out.text);
4584 let structured = out.structured.unwrap();
4585 assert_eq!(structured["status"], "ok");
4586 assert_eq!(structured["database"], "appdb");
4587 }
4588
4589 #[tokio::test]
4593 async fn orient_lists_tables_and_declared_joins() {
4594 let router = join_path_router();
4595 let out = router.call("orient", json!({})).await;
4596 assert!(!out.is_error, "{}", out.text);
4597 let structured = out.structured.unwrap();
4598 assert_eq!(structured["database"], "appdb");
4599
4600 let tables = structured["tables"].as_array().unwrap();
4601 let refs: Vec<&str> = tables.iter().map(|t| t["ref"].as_str().unwrap()).collect();
4602 assert_eq!(
4603 refs,
4604 vec!["public.customers", "public.order_items", "public.orders"]
4605 );
4606 let orders = tables.iter().find(|t| t["ref"] == "public.orders").unwrap();
4607 assert_eq!(orders["pk"], "id");
4608 assert!(orders["columns"].as_str().unwrap().contains("id:integer!"));
4609
4610 let joins = structured["joins"].as_array().unwrap();
4611 assert_eq!(joins.len(), 2);
4612 assert!(joins.iter().any(|j| j["edge"]
4613 == "public.order_items.order_id -> public.orders.id"
4614 && j["declared"] == true));
4615 }
4616
4617 #[tokio::test]
4618 async fn orient_focus_narrows_tables_and_joins() {
4619 let router = join_path_router();
4620 let out = router.call("orient", json!({ "focus": "customers" })).await;
4621 assert!(!out.is_error, "{}", out.text);
4622 let structured = out.structured.unwrap();
4623 let tables = structured["tables"].as_array().unwrap();
4624 assert_eq!(tables.len(), 1);
4625 assert_eq!(tables[0]["ref"], "public.customers");
4626 let joins = structured["joins"].as_array().unwrap();
4628 assert_eq!(joins.len(), 1);
4629 assert!(
4630 joins[0]["edge"]
4631 .as_str()
4632 .unwrap()
4633 .contains("public.customers")
4634 );
4635 }
4636
4637 #[tokio::test]
4638 async fn orient_no_index_returns_notes_not_error() {
4639 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4640 let router = ToolRouter::with_index_store(session, None);
4641 let out = router.call("orient", json!({})).await;
4642 assert!(!out.is_error, "{}", out.text);
4643 let structured = out.structured.unwrap();
4644 assert_eq!(structured["tables"].as_array().unwrap().len(), 0);
4645 assert!(!structured["notes"].as_array().unwrap().is_empty());
4646 }
4647
4648 #[tokio::test]
4649 async fn empty_index_dir_returns_build_hint() {
4650 let tmp = tempfile::TempDir::new().unwrap();
4651 let store = IndexStore::new(tmp.path());
4652 let session = ToolSession::for_tests(
4653 vec![test_conn()],
4654 PolicyFilter::default(),
4655 Some(IndexStore::new(tmp.path())),
4656 );
4657 let router = ToolRouter::with_index_store(session, Some(store));
4658 let out = router
4659 .call("describe_object", json!({ "ref": "public.users" }))
4660 .await;
4661 assert!(out.is_error, "{}", out.text);
4662 assert!(
4663 out.text.contains("rebuild_index"),
4664 "expected build hint, got: {}",
4665 out.text
4666 );
4667 }
4668
4669 #[tokio::test]
4670 async fn outcome_tagged_with_connection_id_and_database() {
4671 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4672 let router = ToolRouter::new(session);
4673 let out = router.call("list_connections", json!({})).await;
4674 let structured = out.structured.expect("structured outcome");
4675 assert_eq!(
4676 structured.get("connectionId").and_then(|v| v.as_str()),
4677 Some("conn-1")
4678 );
4679 assert_eq!(
4680 structured.get("database").and_then(|v| v.as_str()),
4681 Some("appdb")
4682 );
4683 }
4684
4685 #[tokio::test]
4686 async fn setup_connection_returns_needs_input_when_incomplete() {
4687 unsafe {
4688 std::env::remove_var("DATABASE_URL");
4689 std::env::remove_var("POSTGRES_URL");
4690 std::env::remove_var("PGHOST");
4691 }
4692 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4693 let router = ToolRouter::new(session);
4694 let out = router.call("setup_connection", json!({})).await;
4695 let structured = out.structured.expect("structured outcome");
4696 assert!(structured.get("status").is_some());
4697 }
4698
4699 #[tokio::test]
4700 async fn save_profile_persists_config() {
4701 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4702 let router = ToolRouter::new(session);
4703 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4704 let temp_dir = tempfile::tempdir().unwrap();
4705 let cfg_path = temp_dir.path().join("config.toml");
4706 unsafe {
4707 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4708 }
4709
4710 let out = router
4711 .call(
4712 "save_profile",
4713 json!({
4714 "name": "staging",
4715 "host": "127.0.0.1",
4716 "port": 5432,
4717 "dbname": "stage_db",
4718 "user": "stage_user"
4719 }),
4720 )
4721 .await;
4722
4723 let structured = out.structured.expect("structured outcome");
4724 assert_eq!(
4725 structured.get("status").and_then(|v| v.as_str()),
4726 Some("saved")
4727 );
4728 assert_eq!(
4729 structured.get("profile").and_then(|v| v.as_str()),
4730 Some("staging")
4731 );
4732 }
4733
4734 #[tokio::test]
4742 async fn save_profile_never_persists_password_in_plaintext() {
4743 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4744 let router = ToolRouter::new(session);
4745 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4746 let temp_dir = tempfile::tempdir().unwrap();
4747 let cfg_path = temp_dir.path().join("config.toml");
4748 unsafe {
4749 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4750 }
4751 const SECRET: &str = "correct-horse-battery-staple";
4752
4753 let out = router
4754 .call(
4755 "save_profile",
4756 json!({
4757 "name": "prod-db",
4758 "host": "127.0.0.1",
4759 "password": SECRET,
4760 }),
4761 )
4762 .await;
4763
4764 if out.is_error {
4765 assert!(
4768 out.text.contains("keyring") || out.text.contains("password_command"),
4769 "{}",
4770 out.text
4771 );
4772 } else {
4773 let structured = out.structured.expect("structured outcome");
4774 assert_eq!(
4775 structured.get("status").and_then(|v| v.as_str()),
4776 Some("saved")
4777 );
4778 }
4779
4780 if cfg_path.exists() {
4781 let raw = std::fs::read_to_string(&cfg_path).unwrap();
4782 assert!(
4783 !raw.contains(SECRET),
4784 "password must never appear in the persisted config file: {raw}"
4785 );
4786 if !out.is_error {
4787 assert!(
4788 raw.contains("keyring"),
4789 "successful save must record credential_provider = \"keyring\": {raw}"
4790 );
4791 }
4792 }
4793 }
4794
4795 #[tokio::test]
4801 async fn save_profile_rejects_elevated_access_without_confirmation() {
4802 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4803 let router = ToolRouter::new(session);
4804 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4805 let temp_dir = tempfile::tempdir().unwrap();
4806 let cfg_path = temp_dir.path().join("config.toml");
4807 unsafe {
4808 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4809 }
4810
4811 let out = router
4812 .call(
4813 "save_profile",
4814 json!({
4815 "name": "prod",
4816 "host": "127.0.0.1",
4817 "access_mode": "admin",
4818 }),
4819 )
4820 .await;
4821 assert!(out.is_error, "{}", out.text);
4822 assert!(out.text.contains("confirm_elevated_access"), "{}", out.text);
4823 assert!(
4824 !cfg_path.exists(),
4825 "rejected escalation must not touch the config file"
4826 );
4827 }
4828
4829 #[tokio::test]
4830 async fn save_profile_allows_elevated_access_with_confirmation() {
4831 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4832 let router = ToolRouter::new(session);
4833 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4834 let temp_dir = tempfile::tempdir().unwrap();
4835 let cfg_path = temp_dir.path().join("config.toml");
4836 unsafe {
4837 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4838 }
4839
4840 let out = router
4841 .call(
4842 "save_profile",
4843 json!({
4844 "name": "prod",
4845 "host": "127.0.0.1",
4846 "access_mode": "admin",
4847 "confirm_elevated_access": true,
4848 }),
4849 )
4850 .await;
4851 assert!(!out.is_error, "{}", out.text);
4852 }
4853
4854 #[tokio::test]
4855 async fn save_profile_read_access_mode_needs_no_confirmation() {
4856 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4857 let router = ToolRouter::new(session);
4858 let _env_guard = CONFIG_ENV_LOCK.lock().await;
4859 let temp_dir = tempfile::tempdir().unwrap();
4860 let cfg_path = temp_dir.path().join("config.toml");
4861 unsafe {
4862 std::env::set_var("NEXQL_MCP_CONFIG", &cfg_path);
4863 }
4864
4865 let out = router
4866 .call(
4867 "save_profile",
4868 json!({
4869 "name": "readonly",
4870 "host": "127.0.0.1",
4871 "access_mode": "read",
4872 }),
4873 )
4874 .await;
4875 assert!(!out.is_error, "{}", out.text);
4876 }
4877
4878 #[tokio::test]
4879 async fn check_ddl_safety_tool_dispatches_ast_report() {
4880 let session = ToolSession::for_tests(vec![test_conn()], PolicyFilter::default(), None);
4881 let router = ToolRouter::new(session);
4882 let out = router
4883 .call(
4884 "check_ddl_safety",
4885 json!({ "ddl": "CREATE INDEX idx_col ON users(col);" }),
4886 )
4887 .await;
4888 let structured = out.structured.expect("structured outcome");
4889 assert_eq!(
4890 structured.get("overall_risk").and_then(|v| v.as_str()),
4891 Some("CRITICAL")
4892 );
4893 }
4894}