feat: rescope to AI-first DB harness with multi-DB chat agent

Removes enterprise/DBA features and replaces the marginal AI bar with a
central chat agent that has progressive-discovery tools, cross-session
memory, saved-query reuse, and inline result actions. Adds ClickHouse
support alongside PostgreSQL/Greenplum.

Cleanup
- Drop ~10k LOC of advanced features: Docker, Snapshots, Validation,
  Index Advisor, Role/User Management, Data Generator, ERD, Lookup.
- Trim deps: drop @xyflow/react, dagre, @types/dagre; cut tokio features
  to rt-multi-thread/sync/time/net/macros.
- Remove unused TuskError variants and dead helpers (topological_sort,
  invalidate_schema_cache).

Multi-DB (PostgreSQL + ClickHouse)
- New src-tauri/src/db/ module: ChClient (HTTP-based, reuses reqwest),
  sql_guard (cross-flavor read-only whitelist with 8 tests).
- ConnectionConfig gains db_flavor and secure fields with serde defaults
  for backwards-compatible connections.json.
- All connection/query/schema/data commands dispatch by flavor; CH
  covers connect, execute_query, list_databases/schemas/tables/views/
  columns/completion_schema, paginated table fetch.
- Frontend: dbCapabilities matrix, ConnectionDialog engine selector
  with port auto-swap and HTTPS toggle, SqlEditor switches to
  StandardSQL dialect for CH, TableDataView surfaces CH connections as
  read-only.

AI-first chat agent
- New src/components/chat/ panel with composer, message rendering,
  collapsible tool-call/result blocks, top-level ErrorBoundary.
- Backend agent loop in commands/chat.rs with strict-JSON tool
  protocol. Nine tools: list_databases, list_tables, get_columns,
  switch_database, run_query, remember, save_query, find_queries, final.
  Forgiving parser accepts both flat and nested-input shapes.
- Compressed history: only the last 4 run_query results carry sample
  rows (≤10, cells truncated to 200 chars) into LLM context; older
  results marked omitted.
- System prompt uses lite OVERVIEW (DB list + active-DB tables only)
  instead of full DDL — schema details are loaded on demand via
  get_columns. CH OVERVIEW shows cross-DB tables since CH allows
  db.table queries.

Cross-session memory (F1)
- Per-connection markdown file at app_data_dir/memory/<connection_id>.md,
  16KB cap with oldest-block eviction. Agent appends via remember()
  tool; the file is injected into LEARNED NOTES section of every system
  prompt.
- New Memory sidebar tab with editable textarea, badge for note count,
  empty-state with template. Edits picked up on the next agent turn.

Saved-query reuse (F2)
- Tools save_query and find_queries scoped to current connection.
  save_query attaches a UUID + timestamp; find_queries returns top 10
  matches with SQL preview ≤500 chars.
- Storage shared with the sidebar Saved panel.

Inline result actions (F3)
- run_query result block in chat gets Open-full (90vw × 80vh modal with
  full ResultsTable, no row cap) and Export (reuses ExportDialog for
  CSV/JSON via existing exportCsv/exportJson commands).

Verification
- cargo check clean, zero warnings.
- cargo test --lib: 50 pass (20 chat parser + 4 memory + 8 sql_guard +
  6 clean_sql + 12 escape_ident).
- npx tsc --noEmit clean.
- npx vitest run: 20 pass.
This commit is contained in:
dev
2026-05-06 19:30:44 +03:00
parent 652937f7f5
commit 4f7afc17f4
89 changed files with 4151 additions and 10756 deletions
-254
View File
@@ -15,13 +15,10 @@
"@tauri-apps/api": "^2.10.1",
"@tauri-apps/plugin-dialog": "^2.6.0",
"@tauri-apps/plugin-shell": "^2.3.5",
"@types/dagre": "^0.7.54",
"@uiw/react-codemirror": "^4.25.4",
"@xyflow/react": "^12.10.2",
"class-variance-authority": "^0.7.1",
"clsx": "^2.1.1",
"cmdk": "^1.1.1",
"dagre": "^0.8.5",
"lucide-react": "^0.563.0",
"next-themes": "^0.4.6",
"radix-ui": "^1.4.3",
@@ -4810,61 +4807,6 @@
"assertion-error": "^2.0.1"
}
},
"node_modules/@types/d3-color": {
"version": "3.1.3",
"resolved": "https://registry.npmjs.org/@types/d3-color/-/d3-color-3.1.3.tgz",
"integrity": "sha512-iO90scth9WAbmgv7ogoq57O9YpKmFBbmoEoCHDB2xMBY0+/KVrqAaCDyCE16dUspeOvIxFFRI+0sEtqDqy2b4A==",
"license": "MIT"
},
"node_modules/@types/d3-drag": {
"version": "3.0.7",
"resolved": "https://registry.npmjs.org/@types/d3-drag/-/d3-drag-3.0.7.tgz",
"integrity": "sha512-HE3jVKlzU9AaMazNufooRJ5ZpWmLIoc90A37WU2JMmeq28w1FQqCZswHZ3xR+SuxYftzHq6WU6KJHvqxKzTxxQ==",
"license": "MIT",
"dependencies": {
"@types/d3-selection": "*"
}
},
"node_modules/@types/d3-interpolate": {
"version": "3.0.4",
"resolved": "https://registry.npmjs.org/@types/d3-interpolate/-/d3-interpolate-3.0.4.tgz",
"integrity": "sha512-mgLPETlrpVV1YRJIglr4Ez47g7Yxjl1lj7YKsiMCb27VJH9W8NVM6Bb9d8kkpG/uAQS5AmbA48q2IAolKKo1MA==",
"license": "MIT",
"dependencies": {
"@types/d3-color": "*"
}
},
"node_modules/@types/d3-selection": {
"version": "3.0.11",
"resolved": "https://registry.npmjs.org/@types/d3-selection/-/d3-selection-3.0.11.tgz",
"integrity": "sha512-bhAXu23DJWsrI45xafYpkQ4NtcKMwWnAC/vKrd2l+nxMFuvOT3XMYTIj2opv8vq8AO5Yh7Qac/nSeP/3zjTK0w==",
"license": "MIT"
},
"node_modules/@types/d3-transition": {
"version": "3.0.9",
"resolved": "https://registry.npmjs.org/@types/d3-transition/-/d3-transition-3.0.9.tgz",
"integrity": "sha512-uZS5shfxzO3rGlu0cC3bjmMFKsXv+SmZZcgp0KD22ts4uGXp5EVYGzu/0YdwZeKmddhcAccYtREJKkPfXkZuCg==",
"license": "MIT",
"dependencies": {
"@types/d3-selection": "*"
}
},
"node_modules/@types/d3-zoom": {
"version": "3.0.8",
"resolved": "https://registry.npmjs.org/@types/d3-zoom/-/d3-zoom-3.0.8.tgz",
"integrity": "sha512-iqMC4/YlFCSlO8+2Ii1GGGliCAY4XdeG748w5vQUbevlbDu0zSjH/+jojorQVBK/se0j6DUFNPBGSqD3YWYnDw==",
"license": "MIT",
"dependencies": {
"@types/d3-interpolate": "*",
"@types/d3-selection": "*"
}
},
"node_modules/@types/dagre": {
"version": "0.7.54",
"resolved": "https://registry.npmjs.org/@types/dagre/-/dagre-0.7.54.tgz",
"integrity": "sha512-QjcRY+adGbYvBFS7cwv5txhVIwX1XXIUswWl+kSQTbI6NjgZydrZkEKX/etzVd7i+bCsCb40Z/xlBY5eoFuvWQ==",
"license": "MIT"
},
"node_modules/@types/deep-eql": {
"version": "4.0.2",
"resolved": "https://registry.npmjs.org/@types/deep-eql/-/deep-eql-4.0.2.tgz",
@@ -5384,66 +5326,6 @@
"url": "https://opencollective.com/vitest"
}
},
"node_modules/@xyflow/react": {
"version": "12.10.2",
"resolved": "https://registry.npmjs.org/@xyflow/react/-/react-12.10.2.tgz",
"integrity": "sha512-CgIi6HwlcHXwlkTpr0fxLv/0sRVNZ8IdwKLzzeCscaYBwpvfcH1QFOCeaTCuEn1FQEs/B8CjnTSjhs8udgmBgQ==",
"license": "MIT",
"dependencies": {
"@xyflow/system": "0.0.76",
"classcat": "^5.0.3",
"zustand": "^4.4.0"
},
"peerDependencies": {
"react": ">=17",
"react-dom": ">=17"
}
},
"node_modules/@xyflow/react/node_modules/zustand": {
"version": "4.5.7",
"resolved": "https://registry.npmjs.org/zustand/-/zustand-4.5.7.tgz",
"integrity": "sha512-CHOUy7mu3lbD6o6LJLfllpjkzhHXSBlX8B9+qPddUsIfeF5S/UZ5q0kmCsnRqT1UHFQZchNFDDzMbQsuesHWlw==",
"license": "MIT",
"dependencies": {
"use-sync-external-store": "^1.2.2"
},
"engines": {
"node": ">=12.7.0"
},
"peerDependencies": {
"@types/react": ">=16.8",
"immer": ">=9.0.6",
"react": ">=16.8"
},
"peerDependenciesMeta": {
"@types/react": {
"optional": true
},
"immer": {
"optional": true
},
"react": {
"optional": true
}
}
},
"node_modules/@xyflow/system": {
"version": "0.0.76",
"resolved": "https://registry.npmjs.org/@xyflow/system/-/system-0.0.76.tgz",
"integrity": "sha512-hvwvnRS1B3REwVDlWexsq7YQaPZeG3/mKo1jv38UmnpWmxihp14bW6VtEOuHEwJX2FvzFw8k77LyKSk/wiZVNA==",
"license": "MIT",
"dependencies": {
"@types/d3-drag": "^3.0.7",
"@types/d3-interpolate": "^3.0.4",
"@types/d3-selection": "^3.0.10",
"@types/d3-transition": "^3.0.8",
"@types/d3-zoom": "^3.0.8",
"d3-drag": "^3.0.0",
"d3-interpolate": "^3.0.1",
"d3-selection": "^3.0.0",
"d3-zoom": "^3.0.0"
}
},
"node_modules/accepts": {
"version": "2.0.0",
"resolved": "https://registry.npmjs.org/accepts/-/accepts-2.0.0.tgz",
@@ -5877,12 +5759,6 @@
"url": "https://polar.sh/cva"
}
},
"node_modules/classcat": {
"version": "5.0.5",
"resolved": "https://registry.npmjs.org/classcat/-/classcat-5.0.5.tgz",
"integrity": "sha512-JhZUT7JFcQy/EzW605k/ktHtncoo9vnyW/2GspNYwFlN1C/WmjuV/xtS04e9SOkL2sTdw0VAZ2UGCcQ9lR6p6w==",
"license": "MIT"
},
"node_modules/cli-cursor": {
"version": "5.0.0",
"resolved": "https://registry.npmjs.org/cli-cursor/-/cli-cursor-5.0.0.tgz",
@@ -6268,121 +6144,6 @@
"devOptional": true,
"license": "MIT"
},
"node_modules/d3-color": {
"version": "3.1.0",
"resolved": "https://registry.npmjs.org/d3-color/-/d3-color-3.1.0.tgz",
"integrity": "sha512-zg/chbXyeBtMQ1LbD/WSoW2DpC3I0mpmPdW+ynRTj/x2DAWYrIY7qeZIHidozwV24m4iavr15lNwIwLxRmOxhA==",
"license": "ISC",
"engines": {
"node": ">=12"
}
},
"node_modules/d3-dispatch": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/d3-dispatch/-/d3-dispatch-3.0.1.tgz",
"integrity": "sha512-rzUyPU/S7rwUflMyLc1ETDeBj0NRuHKKAcvukozwhshr6g6c5d8zh4c2gQjY2bZ0dXeGLWc1PF174P2tVvKhfg==",
"license": "ISC",
"engines": {
"node": ">=12"
}
},
"node_modules/d3-drag": {
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/d3-drag/-/d3-drag-3.0.0.tgz",
"integrity": "sha512-pWbUJLdETVA8lQNJecMxoXfH6x+mO2UQo8rSmZ+QqxcbyA3hfeprFgIT//HW2nlHChWeIIMwS2Fq+gEARkhTkg==",
"license": "ISC",
"dependencies": {
"d3-dispatch": "1 - 3",
"d3-selection": "3"
},
"engines": {
"node": ">=12"
}
},
"node_modules/d3-ease": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/d3-ease/-/d3-ease-3.0.1.tgz",
"integrity": "sha512-wR/XK3D3XcLIZwpbvQwQ5fK+8Ykds1ip7A2Txe0yxncXSdq1L9skcG7blcedkOX+ZcgxGAmLX1FrRGbADwzi0w==",
"license": "BSD-3-Clause",
"engines": {
"node": ">=12"
}
},
"node_modules/d3-interpolate": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/d3-interpolate/-/d3-interpolate-3.0.1.tgz",
"integrity": "sha512-3bYs1rOD33uo8aqJfKP3JWPAibgw8Zm2+L9vBKEHJ2Rg+viTR7o5Mmv5mZcieN+FRYaAOWX5SJATX6k1PWz72g==",
"license": "ISC",
"dependencies": {
"d3-color": "1 - 3"
},
"engines": {
"node": ">=12"
}
},
"node_modules/d3-selection": {
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz",
"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
"license": "ISC",
"engines": {
"node": ">=12"
}
},
"node_modules/d3-timer": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/d3-timer/-/d3-timer-3.0.1.tgz",
"integrity": "sha512-ndfJ/JxxMd3nw31uyKoY2naivF+r29V+Lc0svZxe1JvvIRmi8hUsrMvdOwgS1o6uBHmiz91geQ0ylPP0aj1VUA==",
"license": "ISC",
"engines": {
"node": ">=12"
}
},
"node_modules/d3-transition": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/d3-transition/-/d3-transition-3.0.1.tgz",
"integrity": "sha512-ApKvfjsSR6tg06xrL434C0WydLr7JewBB3V+/39RMHsaXTOG0zmt/OAXeng5M5LBm0ojmxJrpomQVZ1aPvBL4w==",
"license": "ISC",
"dependencies": {
"d3-color": "1 - 3",
"d3-dispatch": "1 - 3",
"d3-ease": "1 - 3",
"d3-interpolate": "1 - 3",
"d3-timer": "1 - 3"
},
"engines": {
"node": ">=12"
},
"peerDependencies": {
"d3-selection": "2 - 3"
}
},
"node_modules/d3-zoom": {
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/d3-zoom/-/d3-zoom-3.0.0.tgz",
"integrity": "sha512-b8AmV3kfQaqWAuacbPuNbL6vahnOJflOhexLzMMNLga62+/nh0JzvJ0aO/5a5MVgUFGS7Hu1P9P03o3fJkDCyw==",
"license": "ISC",
"dependencies": {
"d3-dispatch": "1 - 3",
"d3-drag": "2 - 3",
"d3-interpolate": "1 - 3",
"d3-selection": "2 - 3",
"d3-transition": "2 - 3"
},
"engines": {
"node": ">=12"
}
},
"node_modules/dagre": {
"version": "0.8.5",
"resolved": "https://registry.npmjs.org/dagre/-/dagre-0.8.5.tgz",
"integrity": "sha512-/aTqmnRta7x7MCCpExk7HQL2O4owCT2h8NT//9I1OQ9vt29Pa0BzSAkR5lwFUcQ7491yVi/3CXU9jQ5o0Mn2Sw==",
"license": "MIT",
"dependencies": {
"graphlib": "^2.1.8",
"lodash": "^4.17.15"
}
},
"node_modules/data-uri-to-buffer": {
"version": "4.0.1",
"resolved": "https://registry.npmjs.org/data-uri-to-buffer/-/data-uri-to-buffer-4.0.1.tgz",
@@ -7610,15 +7371,6 @@
"dev": true,
"license": "ISC"
},
"node_modules/graphlib": {
"version": "2.1.8",
"resolved": "https://registry.npmjs.org/graphlib/-/graphlib-2.1.8.tgz",
"integrity": "sha512-jcLLfkpoVGmH7/InMC/1hIvOPSUh38oJtGhvrOFGzioE1DZ+0YW16RgmOJhHiuWTvGiJQ9Z1Ik43JvkRPRvE+A==",
"license": "MIT",
"dependencies": {
"lodash": "^4.17.15"
}
},
"node_modules/graphql": {
"version": "16.12.0",
"resolved": "https://registry.npmjs.org/graphql/-/graphql-16.12.0.tgz",
@@ -8555,12 +8307,6 @@
"url": "https://github.com/sponsors/sindresorhus"
}
},
"node_modules/lodash": {
"version": "4.17.23",
"resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.23.tgz",
"integrity": "sha512-LgVTMpQtIopCi79SJeDiP0TfWi5CNEc/L/aRdTh3yIvmZXTnheWpKjSZhnvMl8iXbC1tFg9gdHHDMLoV7CnG+w==",
"license": "MIT"
},
"node_modules/lodash.merge": {
"version": "4.6.2",
"resolved": "https://registry.npmjs.org/lodash.merge/-/lodash.merge-4.6.2.tgz",
-3
View File
@@ -20,13 +20,10 @@
"@tauri-apps/api": "^2.10.1",
"@tauri-apps/plugin-dialog": "^2.6.0",
"@tauri-apps/plugin-shell": "^2.3.5",
"@types/dagre": "^0.7.54",
"@uiw/react-codemirror": "^4.25.4",
"@xyflow/react": "^12.10.2",
"class-variance-authority": "^0.7.1",
"clsx": "^2.1.1",
"cmdk": "^1.1.1",
"dagre": "^0.8.5",
"lucide-react": "^0.563.0",
"next-themes": "^0.4.6",
"radix-ui": "^1.4.3",
-1
View File
@@ -4705,7 +4705,6 @@ dependencies = [
"libc",
"mio",
"pin-project-lite",
"signal-hook-registry",
"socket2",
"tokio-macros",
"windows-sys 0.61.2",
+1 -1
View File
@@ -21,7 +21,7 @@ tauri-plugin-shell = "2"
tauri-plugin-dialog = "2"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["rt-multi-thread", "sync", "time", "net", "macros", "process", "io-util"] }
tokio = { version = "1", features = ["rt-multi-thread", "sync", "time", "net", "macros"] }
sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "postgres", "json", "chrono", "uuid", "bigdecimal"] }
chrono = { version = "0.4", features = ["serde"] }
uuid = { version = "1", features = ["serde"] }
File diff suppressed because it is too large Load Diff
+869
View File
@@ -0,0 +1,869 @@
use crate::commands::ai::{build_overview_context, call_ollama_chat_messages};
use crate::commands::chat_tools::{
find_queries_tool, get_columns_tool, list_databases_tool, list_tables_tool, save_query_tool,
switch_database_tool,
};
use crate::commands::memory::{append_memory_core, read_memory_core};
use crate::commands::queries::execute_query_core;
use crate::error::{TuskError, TuskResult};
use crate::models::ai::OllamaChatMessage;
use crate::models::chat::ChatMessage;
use crate::models::query_result::QueryResult;
use crate::state::AppState;
use chrono::Utc;
use serde_json::Value;
use std::sync::Arc;
use tauri::{AppHandle, State};
const MAX_HOPS: usize = 8;
/// Number of MOST RECENT run_query tool_results that get full sample-rows in
/// LLM history. Older ones are reduced to a marker so very long threads stay
/// within model context budget.
const RECENT_TOOL_RESULTS_FULL: usize = 4;
/// Sample-row cap for compressed run_query results in LLM history.
const RUN_QUERY_SAMPLE_ROWS: usize = 10;
/// Per-cell character cap when stringifying sample rows.
const CELL_CHAR_CAP: usize = 200;
/// Per text-tool-result character cap (list_tables, get_columns, etc).
const TEXT_TOOL_CHAR_CAP: usize = 10_000;
// ---------------------------------------------------------------------------
// Action protocol
// ---------------------------------------------------------------------------
#[derive(Debug)]
enum AgentAction {
Final { text: String },
RunQuery { sql: String },
ListDatabases,
ListTables { database: Option<String> },
GetColumns { tables: Vec<String> },
SwitchDatabase { database: String },
Remember { note: String },
SaveQuery { name: String, sql: String },
FindQueries { text: String },
}
/// Parse the model's JSON response. Accepts both shapes the model tends to emit:
/// {"action":"X","field":"..."} — flat (matches our prompt)
/// {"action":"X","input":{"field":"..."}} — nested (common tool-use convention)
fn parse_agent_action(raw: &str) -> Result<AgentAction, String> {
let v: Value = serde_json::from_str(raw).map_err(|e| e.to_string())?;
let obj = v.as_object().ok_or_else(|| "expected JSON object".to_string())?;
let action = obj
.get("action")
.and_then(|a| a.as_str())
.ok_or_else(|| "missing field `action`".to_string())?;
let lookup = |key: &str| -> Option<&Value> {
obj.get(key)
.or_else(|| obj.get("input").and_then(|i| i.as_object()).and_then(|i| i.get(key)))
};
match action {
"final" => {
let text = lookup("text")
.and_then(|v| v.as_str())
.ok_or_else(|| "final action missing `text`".to_string())?
.to_string();
Ok(AgentAction::Final { text })
}
"run_query" => {
let sql = lookup("sql")
.and_then(|v| v.as_str())
.ok_or_else(|| "run_query action missing `sql`".to_string())?
.to_string();
Ok(AgentAction::RunQuery { sql })
}
"list_databases" => Ok(AgentAction::ListDatabases),
"list_tables" => {
let database = lookup("database")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
Ok(AgentAction::ListTables { database })
}
"get_columns" => {
let arr = lookup("tables")
.and_then(|v| v.as_array())
.ok_or_else(|| "get_columns action missing `tables`: [...]".to_string())?;
let tables: Vec<String> = arr
.iter()
.filter_map(|v| v.as_str().map(|s| s.to_string()))
.collect();
if tables.is_empty() {
return Err("get_columns `tables` array must not be empty".into());
}
Ok(AgentAction::GetColumns { tables })
}
"switch_database" => {
let database = lookup("database")
.and_then(|v| v.as_str())
.ok_or_else(|| "switch_database missing `database`".to_string())?
.to_string();
Ok(AgentAction::SwitchDatabase { database })
}
"remember" => {
let note = lookup("note")
.and_then(|v| v.as_str())
.ok_or_else(|| "remember action missing `note`".to_string())?
.trim()
.to_string();
if note.is_empty() {
return Err("remember `note` must not be empty".into());
}
Ok(AgentAction::Remember { note })
}
"save_query" => {
let name = lookup("name")
.and_then(|v| v.as_str())
.ok_or_else(|| "save_query missing `name`".to_string())?
.trim()
.to_string();
let sql = lookup("sql")
.and_then(|v| v.as_str())
.ok_or_else(|| "save_query missing `sql`".to_string())?
.trim()
.to_string();
if name.is_empty() {
return Err("save_query `name` must not be empty".into());
}
if sql.is_empty() {
return Err("save_query `sql` must not be empty".into());
}
Ok(AgentAction::SaveQuery { name, sql })
}
"find_queries" => {
let text = lookup("text")
.and_then(|v| v.as_str())
.ok_or_else(|| "find_queries missing `text`".to_string())?
.trim()
.to_string();
if text.is_empty() {
return Err("find_queries `text` must not be empty".into());
}
Ok(AgentAction::FindQueries { text })
}
// Legacy from earlier iterations — silently ignored at parse time so the
// model can recover with a different action.
"get_schema" => Err(
"get_schema is deprecated; use get_columns({\"tables\":[...]}) instead.".to_string(),
),
other => Err(format!("unknown action `{}`", other)),
}
}
// ---------------------------------------------------------------------------
// id / time helpers
// ---------------------------------------------------------------------------
fn now_ms() -> i64 {
Utc::now().timestamp_millis()
}
fn new_id(prefix: &str) -> String {
format!("{}-{}-{}", prefix, now_ms(), rand_suffix())
}
fn rand_suffix() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
format!("{:x}", nanos)
}
// ---------------------------------------------------------------------------
// System prompt
// ---------------------------------------------------------------------------
fn system_prompt(overview: &str, memory: &str) -> String {
let overview_block = if overview.is_empty() {
"(overview unavailable; respond with `final` asking the user to reconnect.)".to_string()
} else {
overview.to_string()
};
let memory_block = if memory.trim().is_empty() {
"(empty — call remember() when you discover non-obvious facts about this database)".to_string()
} else {
memory.to_string()
};
format!(
r#"ROLE: Tusk's data assistant. Reply in the user's language.
You operate as an agent in a single-tool-per-turn loop with hop limit {hops}. On every turn output STRICT JSON — exactly one of these shapes, with all fields at the root (no `input` wrapper, no markdown fences):
{{"action":"list_databases"}}
Refresh the database list when the OVERVIEW seems stale.
{{"action":"list_tables"}}
List tables in the active database.
{{"action":"list_tables","database":"<name>"}}
List tables in a specific database (PostgreSQL: requires switch_database before run_query).
{{"action":"get_columns","tables":["schema.table","schema.table2"]}}
Load full column info (types, PK, FK, comments, enums) for the listed tables. Use this BEFORE writing SQL when you don't already know the columns.
{{"action":"switch_database","database":"<name>"}}
Change the active database. Required for PostgreSQL when the user's question concerns data in another database. ClickHouse rarely needs this — `db.table` qualifiers are allowed without switching.
{{"action":"run_query","sql":"SELECT ..."}}
Execute read-only SQL (SELECT / WITH ... SELECT / EXPLAIN / SHOW / DESCRIBE). Mutating SQL is rejected by the read-only guard.
{{"action":"remember","note":"<short observation>"}}
Persist a non-obvious fact about THIS database for future sessions: column semantics, naming conventions, business-rule encodings, gotchas. Keep notes < 200 chars. The user sees and can edit your notes in the Memory sidebar tab.
{{"action":"find_queries","text":"<keywords>"}}
Search saved queries (your past work + user-saved). Use BEFORE writing complex SQL — a usable variant may already exist. Top 10 matches with SQL preview are returned.
{{"action":"save_query","name":"<short label>","sql":"<the SQL>"}}
Persist a non-trivial working SELECT for reuse later. Use AFTER a successful run_query when the query is likely to be re-run. Keep `name` short and descriptive (e.g. "GMV by carrier — last 30d"). The user sees these in sidebar → Saved.
{{"action":"final","text":"..."}}
End the turn with a plain-language answer for the user. Do NOT repeat the result table — the UI shows it. Mention caveats (LIMIT, NULL filters, sampling).
WORKFLOW
1. Read LEARNED NOTES below first — the user (or your past self) may have already documented relevant facts.
2. For non-trivial requests, run `find_queries({{text}})` once to check if a saved query already answers the question.
3. Pick candidate tables from the OVERVIEW (active DB) or call list_tables if you need other DBs.
4. If a candidate's columns are unknown, call get_columns FIRST. NEVER invent columns.
5. If the user's data lives in a different DB and engine is PostgreSQL, switch_database first.
6. Execute run_query.
7. If you discovered something non-obvious (semantics, gotcha, business rule that isn't visible from the schema alone), call `remember` BEFORE `final`. Future sessions will see your notes here.
8. If the query is likely to be re-run later (a real report-style request, not a one-off lookup), call `save_query` with a concise `name`.
9. Answer with `final`.
RULES
- Use ONLY identifiers visible to you (overview / list_tables / get_columns output). Don't pluralize, translate, or guess.
- LIMIT on ad-hoc SELECTs unless aggregating.
- On SQL error retry once with a fix; on the second failure respond with `final` explaining what's missing.
- `remember` is for durable facts, not transient observations. Don't memorise query results — only insights about the schema/data model that aren't already in the OVERVIEW.
═══════════════════════════════════════════════════════════════
LEARNED NOTES (per-connection memory; user can edit in sidebar → Memory)
═══════════════════════════════════════════════════════════════
{memory}
═══════════════════════════════════════════════════════════════
═══════════════════════════════════════════════════════════════
OVERVIEW (refreshed every turn)
═══════════════════════════════════════════════════════════════
{overview}
═══════════════════════════════════════════════════════════════
"#,
hops = MAX_HOPS,
memory = memory_block,
overview = overview_block,
)
}
// ---------------------------------------------------------------------------
// Compressed history projection
// ---------------------------------------------------------------------------
/// Compact view of a QueryResult for re-injection into the LLM history.
/// Keeps just enough for the model to reason about the next step (column
/// names, types, total row count, first N rows) without the full payload.
fn compact_query_result(result: &QueryResult) -> Value {
let total = result.rows.len();
let sample: Vec<Vec<Value>> = result
.rows
.iter()
.take(RUN_QUERY_SAMPLE_ROWS)
.map(|row| row.iter().map(truncate_cell).collect())
.collect();
serde_json::json!({
"columns": result.columns,
"types": result.types,
"row_count": total,
"execution_time_ms": result.execution_time_ms,
"sample_rows": sample,
"truncated": total > RUN_QUERY_SAMPLE_ROWS,
})
}
fn truncate_cell(v: &Value) -> Value {
match v {
Value::String(s) if s.chars().count() > CELL_CHAR_CAP => {
let truncated: String = s.chars().take(CELL_CHAR_CAP).collect();
Value::String(format!("{}…", truncated))
}
other => other.clone(),
}
}
fn truncate_text(text: &str) -> String {
if text.len() <= TEXT_TOOL_CHAR_CAP {
text.to_string()
} else {
let mut out = text[..TEXT_TOOL_CHAR_CAP].to_string();
out.push_str("\n…(truncated)");
out
}
}
fn build_history(
messages: &[ChatMessage],
overview_text: &str,
memory_text: &str,
) -> Vec<OllamaChatMessage> {
// Index of run_query tool_results in `messages`. Used to mark which ones
// get full sample rows vs the "(rows omitted)" placeholder.
let run_query_indices: Vec<usize> = messages
.iter()
.enumerate()
.filter_map(|(i, m)| match m {
ChatMessage::ToolResult { tool, .. } if tool == "run_query" => Some(i),
_ => None,
})
.collect();
let keep_full_after_index: usize = if run_query_indices.len() <= RECENT_TOOL_RESULTS_FULL {
0
} else {
run_query_indices[run_query_indices.len() - RECENT_TOOL_RESULTS_FULL]
};
let mut out = Vec::with_capacity(messages.len() + 1);
out.push(OllamaChatMessage {
role: "system".to_string(),
content: system_prompt(overview_text, memory_text),
});
for (idx, m) in messages.iter().enumerate() {
match m {
ChatMessage::User { text, .. } => out.push(OllamaChatMessage {
role: "user".to_string(),
content: text.clone(),
}),
ChatMessage::Assistant { text, .. } => out.push(OllamaChatMessage {
role: "assistant".to_string(),
content: serde_json::json!({ "action": "final", "text": text }).to_string(),
}),
ChatMessage::ToolCall { tool, input_json, .. } => {
if tool == "get_schema" {
continue; // legacy
}
let mut envelope = serde_json::Map::new();
envelope.insert("action".to_string(), Value::String(tool.clone()));
if let Ok(Value::Object(input)) = serde_json::from_str::<Value>(input_json) {
for (k, v) in input {
envelope.insert(k, v);
}
}
out.push(OllamaChatMessage {
role: "assistant".to_string(),
content: Value::Object(envelope).to_string(),
});
}
ChatMessage::ToolResult {
tool,
is_error,
text,
result,
..
} => {
if tool == "get_schema" {
continue; // legacy
}
let payload = match tool.as_str() {
"run_query" => {
if *is_error {
serde_json::json!({
"tool": "run_query",
"error": true,
"text": text.clone().unwrap_or_default(),
})
} else if idx < keep_full_after_index {
serde_json::json!({
"tool": "run_query",
"error": false,
"note": "rows omitted (older result; user has it in the UI above)",
})
} else if let Some(qr) = result {
serde_json::json!({
"tool": "run_query",
"error": false,
"result": compact_query_result(qr),
})
} else {
serde_json::json!({
"tool": "run_query",
"error": false,
"result": null,
})
}
}
// Text-only tools — pass through with cap.
_ => serde_json::json!({
"tool": tool,
"error": *is_error,
"text": text.as_deref().map(truncate_text),
}),
};
out.push(OllamaChatMessage {
role: "user".to_string(),
content: format!("TOOL_RESULT {}", payload),
});
}
}
}
out
}
// ---------------------------------------------------------------------------
// chat_send
// ---------------------------------------------------------------------------
#[tauri::command]
pub async fn chat_send(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
messages: Vec<ChatMessage>,
) -> TuskResult<Vec<ChatMessage>> {
let mut new_messages: Vec<ChatMessage> = Vec::new();
let mut working: Vec<ChatMessage> = messages;
for _hop in 0..MAX_HOPS {
// Overview is rebuilt per turn — cheap (cached) and reflects the active DB
// even if the user (or the agent) just switched it.
let overview_text = build_overview_context(&state, &connection_id)
.await
.unwrap_or_default();
// Memory is read fresh each turn so user-side edits in the Memory tab
// are visible to the agent immediately.
let memory_text = read_memory_core(&app, &connection_id).unwrap_or_default();
let history = build_history(&working, &overview_text, &memory_text);
let raw =
call_ollama_chat_messages(&app, &state, history, Some("json".to_string())).await?;
let trimmed = raw.trim();
let action = match parse_agent_action(trimmed) {
Ok(a) => a,
Err(parse_err) => {
let msg = ChatMessage::Assistant {
id: new_id("asst"),
text: format!(
"{}\n\n_(Note: model returned non-protocol output: {})_",
trimmed, parse_err
),
created_at: now_ms(),
};
new_messages.push(msg.clone());
working.push(msg);
return Ok(new_messages);
}
};
match action {
AgentAction::Final { text } => {
let msg = ChatMessage::Assistant {
id: new_id("asst"),
text,
created_at: now_ms(),
};
new_messages.push(msg.clone());
working.push(msg);
return Ok(new_messages);
}
AgentAction::RunQuery { sql } => {
push_tool_call(
&mut new_messages,
&mut working,
"run_query",
serde_json::json!({ "sql": sql }).to_string(),
);
let result = match execute_query_core(&state, &connection_id, &sql).await {
Ok(qr) => ChatMessage::ToolResult {
id: new_id("res"),
tool: "run_query".to_string(),
is_error: false,
text: None,
result: Some(qr),
created_at: now_ms(),
},
Err(e) => {
let hint = match e {
TuskError::ReadOnly => "\n\nRead-only mode is on. Toggle it off in the toolbar to allow writes.",
_ => "",
};
ChatMessage::ToolResult {
id: new_id("res"),
tool: "run_query".to_string(),
is_error: true,
text: Some(format!("{}{}", e, hint)),
result: None,
created_at: now_ms(),
}
}
};
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::ListDatabases => {
push_tool_call(
&mut new_messages,
&mut working,
"list_databases",
"{}".to_string(),
);
let result = run_text_tool(
list_databases_tool(&state, &connection_id).await,
"list_databases",
);
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::ListTables { database } => {
let input_json = match &database {
Some(db) => serde_json::json!({ "database": db }).to_string(),
None => "{}".to_string(),
};
push_tool_call(&mut new_messages, &mut working, "list_tables", input_json);
let result = run_text_tool(
list_tables_tool(&app, &state, &connection_id, database.as_deref()).await,
"list_tables",
);
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::GetColumns { tables } => {
push_tool_call(
&mut new_messages,
&mut working,
"get_columns",
serde_json::json!({ "tables": tables }).to_string(),
);
let result = run_text_tool(
get_columns_tool(&state, &connection_id, &tables).await,
"get_columns",
);
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::SwitchDatabase { database } => {
push_tool_call(
&mut new_messages,
&mut working,
"switch_database",
serde_json::json!({ "database": &database }).to_string(),
);
let result = run_text_tool(
switch_database_tool(&app, &state, &connection_id, &database).await,
"switch_database",
);
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::Remember { note } => {
push_tool_call(
&mut new_messages,
&mut working,
"remember",
serde_json::json!({ "note": &note }).to_string(),
);
let outcome = append_memory_core(&app, &connection_id, &note)
.map(|_| format!("Saved note ({} chars).", note.len()));
let result = run_text_tool(outcome, "remember");
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::SaveQuery { name, sql } => {
push_tool_call(
&mut new_messages,
&mut working,
"save_query",
serde_json::json!({ "name": &name, "sql": &sql }).to_string(),
);
let result = run_text_tool(
save_query_tool(&app, &connection_id, &name, &sql).await,
"save_query",
);
push_tool_result(&mut new_messages, &mut working, result);
}
AgentAction::FindQueries { text } => {
push_tool_call(
&mut new_messages,
&mut working,
"find_queries",
serde_json::json!({ "text": &text }).to_string(),
);
let result = run_text_tool(
find_queries_tool(&app, &connection_id, &text).await,
"find_queries",
);
push_tool_result(&mut new_messages, &mut working, result);
}
}
}
let msg = ChatMessage::Assistant {
id: new_id("asst"),
text: format!(
"Stopped after {} tool calls without a final answer. Try rephrasing or simplifying the question.",
MAX_HOPS
),
created_at: now_ms(),
};
new_messages.push(msg);
Ok(new_messages)
}
fn push_tool_call(
new_messages: &mut Vec<ChatMessage>,
working: &mut Vec<ChatMessage>,
tool: &str,
input_json: String,
) {
let call = ChatMessage::ToolCall {
id: new_id("call"),
tool: tool.to_string(),
input_json,
created_at: now_ms(),
};
new_messages.push(call.clone());
working.push(call);
}
fn push_tool_result(
new_messages: &mut Vec<ChatMessage>,
working: &mut Vec<ChatMessage>,
result: ChatMessage,
) {
new_messages.push(result.clone());
working.push(result);
}
fn run_text_tool(outcome: TuskResult<String>, tool: &str) -> ChatMessage {
match outcome {
Ok(text) => ChatMessage::ToolResult {
id: new_id("res"),
tool: tool.to_string(),
is_error: false,
text: Some(text),
result: None,
created_at: now_ms(),
},
Err(e) => ChatMessage::ToolResult {
id: new_id("res"),
tool: tool.to_string(),
is_error: true,
text: Some(e.to_string()),
result: None,
created_at: now_ms(),
},
}
}
// ---------------------------------------------------------------------------
// tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_flat_run_query() {
let a = parse_agent_action(r#"{"action":"run_query","sql":"SELECT 1"}"#).unwrap();
match a {
AgentAction::RunQuery { sql } => assert_eq!(sql, "SELECT 1"),
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_nested_run_query() {
let a =
parse_agent_action(r#"{"action":"run_query","input":{"sql":"SELECT 2"}}"#).unwrap();
match a {
AgentAction::RunQuery { sql } => assert_eq!(sql, "SELECT 2"),
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_get_columns() {
let a = parse_agent_action(
r#"{"action":"get_columns","tables":["public.users","public.orders"]}"#,
)
.unwrap();
match a {
AgentAction::GetColumns { tables } => {
assert_eq!(tables, vec!["public.users", "public.orders"]);
}
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_get_columns_nested() {
let a = parse_agent_action(
r#"{"action":"get_columns","input":{"tables":["public.t"]}}"#,
)
.unwrap();
match a {
AgentAction::GetColumns { tables } => assert_eq!(tables, vec!["public.t"]),
_ => panic!("wrong variant"),
}
}
#[test]
fn rejects_get_columns_empty_tables() {
assert!(parse_agent_action(r#"{"action":"get_columns","tables":[]}"#).is_err());
}
#[test]
fn parses_switch_database() {
let a = parse_agent_action(r#"{"action":"switch_database","database":"orders_db"}"#)
.unwrap();
match a {
AgentAction::SwitchDatabase { database } => assert_eq!(database, "orders_db"),
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_list_tables_optional_db() {
let a1 = parse_agent_action(r#"{"action":"list_tables"}"#).unwrap();
match a1 {
AgentAction::ListTables { database } => assert!(database.is_none()),
_ => panic!("wrong variant"),
}
let a2 = parse_agent_action(r#"{"action":"list_tables","database":"x"}"#).unwrap();
match a2 {
AgentAction::ListTables { database } => assert_eq!(database.as_deref(), Some("x")),
_ => panic!("wrong variant"),
}
}
#[test]
fn rejects_unknown_action() {
assert!(parse_agent_action(r#"{"action":"nuke","yes":true}"#).is_err());
}
#[test]
fn parses_remember_flat() {
let a = parse_agent_action(
r#"{"action":"remember","note":"trips.started_at is NULL for cancelled"}"#,
)
.unwrap();
match a {
AgentAction::Remember { note } => {
assert_eq!(note, "trips.started_at is NULL for cancelled");
}
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_remember_nested() {
let a = parse_agent_action(
r#"{"action":"remember","input":{"note":" surrounded by spaces "}}"#,
)
.unwrap();
match a {
AgentAction::Remember { note } => {
// trim happens in parser
assert_eq!(note, "surrounded by spaces");
}
_ => panic!("wrong variant"),
}
}
#[test]
fn rejects_remember_without_note() {
assert!(parse_agent_action(r#"{"action":"remember"}"#).is_err());
}
#[test]
fn rejects_remember_empty_note() {
assert!(parse_agent_action(r#"{"action":"remember","note":" "}"#).is_err());
}
#[test]
fn parses_save_query_flat() {
let a = parse_agent_action(
r#"{"action":"save_query","name":"GMV last 30d","sql":"SELECT 1"}"#,
)
.unwrap();
match a {
AgentAction::SaveQuery { name, sql } => {
assert_eq!(name, "GMV last 30d");
assert_eq!(sql, "SELECT 1");
}
_ => panic!("wrong variant"),
}
}
#[test]
fn parses_save_query_nested() {
let a = parse_agent_action(
r#"{"action":"save_query","input":{"name":"x","sql":"SELECT 2"}}"#,
)
.unwrap();
match a {
AgentAction::SaveQuery { name, sql } => {
assert_eq!(name, "x");
assert_eq!(sql, "SELECT 2");
}
_ => panic!("wrong variant"),
}
}
#[test]
fn rejects_save_query_missing_fields() {
assert!(parse_agent_action(r#"{"action":"save_query","name":"x"}"#).is_err());
assert!(parse_agent_action(r#"{"action":"save_query","sql":"SELECT 1"}"#).is_err());
assert!(
parse_agent_action(r#"{"action":"save_query","name":" ","sql":"SELECT 1"}"#).is_err()
);
}
#[test]
fn parses_find_queries() {
let a = parse_agent_action(r#"{"action":"find_queries","text":"gmv"}"#).unwrap();
match a {
AgentAction::FindQueries { text } => assert_eq!(text, "gmv"),
_ => panic!("wrong variant"),
}
}
#[test]
fn rejects_find_queries_empty_text() {
assert!(parse_agent_action(r#"{"action":"find_queries","text":""}"#).is_err());
}
#[test]
fn rejects_legacy_get_schema() {
assert!(parse_agent_action(r#"{"action":"get_schema"}"#).is_err());
}
#[test]
fn truncates_long_cell() {
let long = "a".repeat(CELL_CHAR_CAP + 50);
let v = truncate_cell(&Value::String(long));
let s = v.as_str().unwrap();
assert!(s.ends_with('…'));
assert!(s.chars().count() <= CELL_CHAR_CAP + 1);
}
#[test]
fn compact_drops_rows_beyond_sample() {
let mut rows = Vec::new();
for i in 0..50 {
rows.push(vec![Value::Number(i.into())]);
}
let qr = QueryResult {
columns: vec!["id".into()],
types: vec!["INT4".into()],
rows,
row_count: 50,
execution_time_ms: 1,
};
let v = compact_query_result(&qr);
let sample = v.get("sample_rows").unwrap().as_array().unwrap();
assert_eq!(sample.len(), RUN_QUERY_SAMPLE_ROWS);
assert_eq!(v.get("truncated").unwrap(), &Value::Bool(true));
assert_eq!(v.get("row_count").unwrap().as_u64().unwrap(), 50);
}
}
+558
View File
@@ -0,0 +1,558 @@
//! Chat agent tool handlers (chat v2).
//!
//! Each `*_tool` function returns a plain string formatted for direct injection
//! into the LLM tool-result history. They reuse the schema helpers in
//! `commands::ai` and `commands::schema` rather than re-implementing SQL.
use crate::commands::ai::{
fetch_column_comments, fetch_columns, fetch_enum_types, fetch_foreign_keys_raw,
fetch_table_comments, fetch_unique_constraints, format_table_block, ColumnInfo,
};
use crate::commands::connections::{load_connection_config, switch_database_core};
use crate::commands::saved_queries::{list_saved_queries_core, save_query_core};
use crate::commands::schema::{list_databases_core, list_tables_core};
use crate::error::{TuskError, TuskResult};
use crate::models::saved_queries::SavedQuery;
use crate::state::{AppState, CachedVec, DbFlavor};
use sqlx::{PgPool, Row};
use std::collections::{BTreeMap, HashMap};
use std::time::{Duration, Instant};
use tauri::AppHandle;
const TOOL_CACHE_TTL: Duration = Duration::from_secs(300);
const MAX_TABLES_PER_GET_COLUMNS: usize = 20;
const COLUMNS_TOOL_OUTPUT_CAP: usize = 15_000;
// ---------------------------------------------------------------------------
// list_databases
// ---------------------------------------------------------------------------
pub async fn list_databases_tool(state: &AppState, connection_id: &str) -> TuskResult<String> {
let dbs = list_databases_core(state, connection_id).await?;
let active = active_db_name(state, connection_id).await;
let mut out = format!("DATABASES ({}):", dbs.len());
for db in &dbs {
if Some(db) == active.as_ref() {
out.push_str(&format!("\n * {} (active)", db));
} else {
out.push_str(&format!("\n {}", db));
}
}
Ok(out)
}
// ---------------------------------------------------------------------------
// list_tables
// ---------------------------------------------------------------------------
pub async fn list_tables_tool(
app: &AppHandle,
state: &AppState,
connection_id: &str,
db: Option<&str>,
) -> TuskResult<String> {
let active = active_db_name(state, connection_id).await;
let target = db.map(|s| s.to_string()).or_else(|| active.clone());
let target_name = match target.as_deref() {
Some(n) => n.to_string(),
None => return Err(TuskError::Custom("No active database selected.".into())),
};
let same_as_active = active.as_deref() == Some(target_name.as_str());
let flavor = state.get_flavor(connection_id).await;
let table_names = match (flavor, same_as_active) {
(DbFlavor::ClickHouse, _) => list_tables_clickhouse(state, connection_id, &target_name).await?,
(_, true) => list_tables_active_pg(state, connection_id).await?,
(_, false) => list_tables_other_pg(app, state, connection_id, &target_name).await?,
};
let header = if same_as_active {
format!("TABLES IN ACTIVE DATABASE `{}` ({}):", target_name, table_names.len())
} else {
format!("TABLES IN DATABASE `{}` ({}):", target_name, table_names.len())
};
let body: Vec<String> = table_names.iter().map(|t| format!(" {}", t)).collect();
Ok(format!("{}\n{}", header, body.join("\n")))
}
async fn list_tables_active_pg(state: &AppState, connection_id: &str) -> TuskResult<Vec<String>> {
let schemas = crate::commands::schema::list_schemas_core(state, connection_id).await?;
let mut all: Vec<String> = Vec::new();
for schema in &schemas {
let tables = list_tables_core(state, connection_id, schema).await?;
for t in tables {
all.push(format!("{}.{}", schema, t.name));
}
}
Ok(all)
}
async fn list_tables_other_pg(
app: &AppHandle,
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<Vec<String>> {
let cache_key = (connection_id.to_string(), target_db.to_string());
if let Some(hit) = state.tables_by_db_cache.read().await.get(&cache_key).cloned() {
if hit.cached_at.elapsed() < TOOL_CACHE_TTL {
return Ok(hit.value);
}
}
let config = load_connection_config(app, connection_id)?;
let url = config.connection_url_for_db(target_db);
let pool = PgPool::connect(&url).await.map_err(|e| {
TuskError::Custom(format!(
"Could not connect to database '{}' on this server: {}",
target_db, e
))
})?;
let rows = sqlx::query(
"SELECT table_schema, table_name FROM information_schema.tables \
WHERE table_schema NOT IN ('pg_catalog','information_schema','pg_toast','gp_toolkit') \
AND table_type = 'BASE TABLE' \
ORDER BY table_schema, table_name",
)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
pool.close().await;
let names: Vec<String> = rows
.iter()
.map(|r| format!("{}.{}", r.get::<String, _>(0), r.get::<String, _>(1)))
.collect();
state.tables_by_db_cache.write().await.insert(
cache_key,
CachedVec {
value: names.clone(),
cached_at: Instant::now(),
},
);
Ok(names)
}
async fn list_tables_clickhouse(
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<Vec<String>> {
let client = state.get_ch_client(connection_id).await?;
let escaped = target_db.replace('\\', "\\\\").replace('\'', "\\'");
let sql = format!(
"SELECT name FROM system.tables WHERE database = '{}' ORDER BY name",
escaped
);
let rows = client.fetch_objects(&sql).await?;
Ok(rows
.iter()
.filter_map(|r| r.get("name").and_then(|v| v.as_str()).map(String::from))
.collect())
}
// ---------------------------------------------------------------------------
// get_columns
// ---------------------------------------------------------------------------
pub async fn get_columns_tool(
state: &AppState,
connection_id: &str,
tables: &[String],
) -> TuskResult<String> {
if tables.is_empty() {
return Err(TuskError::Custom("get_columns requires at least one table.".into()));
}
if tables.len() > MAX_TABLES_PER_GET_COLUMNS {
return Err(TuskError::Custom(format!(
"Too many tables ({}); split into batches of ≤{}.",
tables.len(),
MAX_TABLES_PER_GET_COLUMNS
)));
}
let active_db = active_db_name(state, connection_id).await.unwrap_or_default();
// Normalise: accept "schema.table", "db.schema.table" (drop db if == active),
// and "table" (assume schema "public" for PG, or active DB for CH).
let parsed: Vec<(String, String, String)> = tables
.iter()
.map(|raw| normalise_table_ref(raw, &active_db))
.collect();
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return get_columns_clickhouse(state, connection_id, &parsed).await;
}
get_columns_postgres(state, connection_id, &parsed).await
}
fn normalise_table_ref(raw: &str, active_db: &str) -> (String, String, String) {
// Returns (schema, table, original_input_for_diagnostics)
let trimmed = raw.trim().trim_matches('"').trim_matches('`');
let parts: Vec<&str> = trimmed.split('.').collect();
match parts.len() {
1 => ("public".to_string(), parts[0].to_string(), raw.to_string()),
2 => (parts[0].to_string(), parts[1].to_string(), raw.to_string()),
3 => {
// "db.schema.table" — drop db prefix when it matches active
let (db, schema, table) = (parts[0], parts[1], parts[2]);
if db == active_db {
(schema.to_string(), table.to_string(), raw.to_string())
} else {
// Different DB requested — let the caller surface a not-found warning.
// We still parse it as schema.table here.
(schema.to_string(), table.to_string(), raw.to_string())
}
}
_ => ("public".to_string(), trimmed.to_string(), raw.to_string()),
}
}
async fn get_columns_postgres(
state: &AppState,
connection_id: &str,
requested: &[(String, String, String)],
) -> TuskResult<String> {
let pool = state.get_pool(connection_id).await?;
let (col_res, fk_res, enum_res, tbl_comm_res, col_comm_res, unique_res) = tokio::join!(
fetch_columns(&pool),
fetch_foreign_keys_raw(&pool),
fetch_enum_types(&pool),
fetch_table_comments(&pool),
fetch_column_comments(&pool),
fetch_unique_constraints(&pool),
);
let all_cols = col_res?;
let fk_rows = fk_res?;
let enum_map = enum_res.unwrap_or_default();
let tbl_comments = tbl_comm_res.unwrap_or_default();
let col_comments = col_comm_res.unwrap_or_default();
let uniques = unique_res.unwrap_or_default();
// Build (schema, table) → Vec<ColumnInfo>
let mut by_table: BTreeMap<(String, String), Vec<ColumnInfo>> = BTreeMap::new();
for ci in &all_cols {
by_table
.entry((ci.schema.clone(), ci.table.clone()))
.or_default()
.push(ci.clone());
}
let mut fk_inline: HashMap<(String, String, String), String> = HashMap::new();
for fk in &fk_rows {
if fk.columns.len() == 1 && fk.ref_columns.len() == 1 {
fk_inline.insert(
(fk.schema.clone(), fk.table.clone(), fk.columns[0].clone()),
format!("{}.{}({})", fk.ref_schema, fk.ref_table, fk.ref_columns[0]),
);
}
}
let mut unique_map: HashMap<(String, String), Vec<String>> = HashMap::new();
for (schema, table, cols) in &uniques {
unique_map
.entry((schema.clone(), table.clone()))
.or_default()
.push(cols.join(", "));
}
let varchar_values: HashMap<(String, String, String), Vec<String>> = HashMap::new();
let jsonb_keys: HashMap<(String, String, String), Vec<String>> = HashMap::new();
let mut output: Vec<String> = Vec::new();
let mut not_found: Vec<String> = Vec::new();
for (schema, table, raw) in requested {
match by_table.get(&(schema.clone(), table.clone())) {
Some(cols) => {
let full_name = format!("{}.{}", schema, table);
format_table_block(
&full_name,
cols,
&tbl_comments,
&col_comments,
&fk_inline,
&enum_map,
&unique_map,
&varchar_values,
&jsonb_keys,
&mut output,
);
}
None => not_found.push(raw.clone()),
}
}
if !not_found.is_empty() {
let nearest = nearest_table_matches(&by_table, &not_found);
let header = format!(
"WARNING: tables not found: {}.{}",
not_found.join(", "),
if nearest.is_empty() {
String::new()
} else {
format!(" Nearest matches: {}.", nearest.join(", "))
}
);
output.insert(0, header);
output.insert(1, String::new());
}
let mut text = output.join("\n");
if text.len() > COLUMNS_TOOL_OUTPUT_CAP {
text.truncate(COLUMNS_TOOL_OUTPUT_CAP);
text.push_str("\n... (output truncated)");
}
Ok(text)
}
async fn get_columns_clickhouse(
state: &AppState,
connection_id: &str,
requested: &[(String, String, String)],
) -> TuskResult<String> {
let client = state.get_ch_client(connection_id).await?;
let active_db = client.database.clone();
let where_terms: Vec<String> = requested
.iter()
.map(|(schema, table, _)| {
// For CH, treat the parsed "schema" as the database name; if it equals
// a PG-conventional default ("public"), substitute with active CH database.
let dbn = if schema == "public" { active_db.clone() } else { schema.clone() };
format!(
"(database = '{}' AND name = '{}')",
dbn.replace('\'', "\\'"),
table.replace('\'', "\\'")
)
})
.collect();
let where_clause = where_terms.join(" OR ");
let sql = format!(
"SELECT database, table, name, type, default_expression, is_in_primary_key, comment, position \
FROM system.columns WHERE {} ORDER BY database, table, position",
where_clause
);
let rows = client.fetch_objects(&sql).await?;
// Group by (database, table)
let mut grouped: BTreeMap<(String, String), Vec<&serde_json::Map<String, serde_json::Value>>> =
BTreeMap::new();
for row in &rows {
let dbn = row.get("database").and_then(|v| v.as_str()).unwrap_or("").to_string();
let tbl = row.get("table").and_then(|v| v.as_str()).unwrap_or("").to_string();
grouped.entry((dbn, tbl)).or_default().push(row);
}
// Track which requested tables were found
let mut output = String::new();
let mut not_found: Vec<String> = Vec::new();
for (schema, table, raw) in requested {
let dbn = if schema == "public" { active_db.clone() } else { schema.clone() };
match grouped.get(&(dbn.clone(), table.clone())) {
Some(cols) => {
output.push_str(&format!("\nTABLE {}.{}\n", dbn, table));
for col in cols {
let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("");
let dtype = col.get("type").and_then(|v| v.as_str()).unwrap_or("");
let is_pk = matches!(
col.get("is_in_primary_key"),
Some(serde_json::Value::Number(n)) if n.as_i64() == Some(1)
) || matches!(
col.get("is_in_primary_key"),
Some(serde_json::Value::String(s)) if s == "1"
);
let default = col.get("default_expression").and_then(|v| v.as_str()).unwrap_or("");
let comment = col.get("comment").and_then(|v| v.as_str()).unwrap_or("");
let mut line = format!(" {} {}", name, dtype);
if is_pk {
line.push_str(" [PK]");
}
if !default.is_empty() {
line.push_str(&format!(" DEFAULT {}", default));
}
if !comment.is_empty() {
line.push_str(&format!(" -- {}", comment));
}
output.push_str(&line);
output.push('\n');
}
}
None => not_found.push(raw.clone()),
}
}
let mut header = String::new();
if !not_found.is_empty() {
header.push_str(&format!(
"WARNING: tables not found: {}\n\n",
not_found.join(", ")
));
}
let mut combined = format!("{}{}", header, output.trim_start());
if combined.len() > COLUMNS_TOOL_OUTPUT_CAP {
combined.truncate(COLUMNS_TOOL_OUTPUT_CAP);
combined.push_str("\n... (output truncated)");
}
Ok(combined)
}
fn nearest_table_matches(
by_table: &BTreeMap<(String, String), Vec<ColumnInfo>>,
missing: &[String],
) -> Vec<String> {
let all: Vec<String> = by_table
.keys()
.map(|(s, t)| format!("{}.{}", s, t))
.collect();
let mut hints: Vec<String> = Vec::new();
for m in missing {
let needle = m.to_lowercase();
let mut candidates: Vec<&String> = all
.iter()
.filter(|n| {
let lower = n.to_lowercase();
lower.contains(&needle) || needle.contains(lower.split('.').last().unwrap_or(""))
})
.take(3)
.collect();
candidates.dedup();
for c in candidates {
if !hints.contains(c) {
hints.push(c.clone());
}
}
}
hints
}
// ---------------------------------------------------------------------------
// switch_database
// ---------------------------------------------------------------------------
pub async fn switch_database_tool(
app: &AppHandle,
state: &AppState,
connection_id: &str,
target_db: &str,
) -> TuskResult<String> {
let config = load_connection_config(app, connection_id)?;
// Verify target exists in cluster
let dbs = list_databases_core(state, connection_id).await?;
if !dbs.iter().any(|d| d == target_db) {
return Err(TuskError::Custom(format!(
"Database '{}' does not exist on this server. Available: {}",
target_db,
dbs.join(", ")
)));
}
switch_database_core(state, &config, target_db).await?;
Ok(format!("Switched active database to '{}'.", target_db))
}
// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------
async fn active_db_name(state: &AppState, connection_id: &str) -> Option<String> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return state
.get_ch_client(connection_id)
.await
.ok()
.map(|c| c.database.clone());
}
let pool = state.get_pool(connection_id).await.ok()?;
sqlx::query_scalar::<_, String>("SELECT current_database()")
.fetch_one(&pool)
.await
.ok()
}
// ---------------------------------------------------------------------------
// save_query / find_queries (chat v3 — F2)
// ---------------------------------------------------------------------------
const FIND_QUERIES_LIMIT: usize = 10;
const FIND_QUERIES_SQL_PREVIEW_CHARS: usize = 500;
pub async fn save_query_tool(
app: &AppHandle,
connection_id: &str,
name: &str,
sql: &str,
) -> TuskResult<String> {
let trimmed_name = name.trim();
let trimmed_sql = sql.trim();
if trimmed_name.is_empty() {
return Err(TuskError::Custom("save_query: name must not be empty".into()));
}
if trimmed_sql.is_empty() {
return Err(TuskError::Custom("save_query: sql must not be empty".into()));
}
let entry = SavedQuery {
id: uuid::Uuid::new_v4().to_string(),
name: trimmed_name.to_string(),
sql: trimmed_sql.to_string(),
connection_id: Some(connection_id.to_string()),
created_at: chrono::Utc::now().to_rfc3339(),
};
save_query_core(app, entry).await?;
Ok(format!("Saved query \"{}\" — visible in sidebar → Saved.", trimmed_name))
}
pub async fn find_queries_tool(
app: &AppHandle,
connection_id: &str,
text: &str,
) -> TuskResult<String> {
let trimmed = text.trim();
if trimmed.is_empty() {
return Err(TuskError::Custom("find_queries: text must not be empty".into()));
}
let all = list_saved_queries_core(app, Some(trimmed)).await?;
let matches: Vec<SavedQuery> = all
.into_iter()
.filter(|q| q.connection_id.as_deref() == Some(connection_id))
.take(FIND_QUERIES_LIMIT)
.collect();
if matches.is_empty() {
return Ok(format!(
"No saved queries match \"{}\" for this connection.",
trimmed
));
}
let mut out = format!(
"Saved queries matching \"{}\" ({}):",
trimmed,
matches.len()
);
for q in &matches {
let sql_preview: String = if q.sql.chars().count() > FIND_QUERIES_SQL_PREVIEW_CHARS {
let truncated: String = q.sql.chars().take(FIND_QUERIES_SQL_PREVIEW_CHARS).collect();
format!("{}…", truncated)
} else {
q.sql.clone()
};
out.push_str(&format!(
"\n\n[{}] {}\n{}",
q.created_at, q.name, sql_preview
));
}
Ok(out)
}
+174 -83
View File
@@ -1,3 +1,4 @@
use crate::db::clickhouse::ChClient;
use crate::error::{TuskError, TuskResult};
use crate::models::connection::ConnectionConfig;
use crate::state::{AppState, DbFlavor};
@@ -23,6 +24,34 @@ pub(crate) fn get_connections_path(app: &AppHandle) -> TuskResult<std::path::Pat
Ok(dir.join("connections.json"))
}
/// Read all saved connection configs from disk.
pub(crate) fn load_all_connections(app: &AppHandle) -> TuskResult<Vec<ConnectionConfig>> {
let path = get_connections_path(app)?;
if !path.exists() {
return Ok(vec![]);
}
let data = fs::read_to_string(&path)?;
let connections: Vec<ConnectionConfig> = serde_json::from_str(&data)?;
Ok(connections)
}
/// Look up a single saved connection by id. Used by tools that need credentials
/// (e.g. switch_database from inside the chat agent loop) but only have the id in scope.
pub(crate) fn load_connection_config(
app: &AppHandle,
connection_id: &str,
) -> TuskResult<ConnectionConfig> {
load_all_connections(app)?
.into_iter()
.find(|c| c.id == connection_id)
.ok_or_else(|| {
TuskError::Custom(format!(
"Connection '{}' not found in connections.json",
connection_id
))
})
}
#[tauri::command]
pub async fn get_connections(app: AppHandle) -> TuskResult<Vec<ConnectionConfig>> {
let path = get_connections_path(&app)?;
@@ -55,6 +84,24 @@ pub async fn save_connection(app: AppHandle, config: ConnectionConfig) -> TuskRe
Ok(())
}
async fn close_connection(state: &AppState, id: &str) {
let mut pools = state.pools.write().await;
if let Some(pool) = pools.remove(id) {
pool.close().await;
}
drop(pools);
let mut clients = state.ch_clients.write().await;
clients.remove(id);
drop(clients);
let mut ro = state.read_only.write().await;
ro.remove(id);
drop(ro);
let mut flavors = state.db_flavors.write().await;
flavors.remove(id);
drop(flavors);
state.invalidate_chat_caches_for(id).await;
}
#[tauri::command]
pub async fn delete_connection(
app: AppHandle,
@@ -69,36 +116,37 @@ pub async fn delete_connection(
let data = serde_json::to_string_pretty(&connections)?;
fs::write(&path, data)?;
}
// Close pool if connected
let mut pools = state.pools.write().await;
if let Some(pool) = pools.remove(&id) {
pool.close().await;
}
let mut ro = state.read_only.write().await;
ro.remove(&id);
let mut flavors = state.db_flavors.write().await;
flavors.remove(&id);
close_connection(&state, &id).await;
Ok(())
}
#[tauri::command]
pub async fn test_connection(config: ConnectionConfig) -> TuskResult<String> {
let pool = PgPool::connect(&config.connection_url())
.await
.map_err(TuskError::Database)?;
let row = sqlx::query("SELECT version()")
.fetch_one(&pool)
.await
.map_err(TuskError::Database)?;
let version: String = row.get(0);
pool.close().await;
Ok(version)
match config.db_flavor {
DbFlavor::ClickHouse => {
let client = ChClient::new(
&config.host,
config.port,
config.secure,
&config.user,
&config.password,
&config.database,
);
client.ping().await
}
_ => {
let pool = PgPool::connect(&config.connection_url())
.await
.map_err(TuskError::Database)?;
let row = sqlx::query("SELECT version()")
.fetch_one(&pool)
.await
.map_err(TuskError::Database)?;
let version: String = row.get(0);
pool.close().await;
Ok(version)
}
}
}
#[tauri::command]
@@ -106,39 +154,110 @@ pub async fn connect(
state: State<'_, Arc<AppState>>,
config: ConnectionConfig,
) -> TuskResult<ConnectResult> {
let pool = PgPool::connect(&config.connection_url())
.await
.map_err(TuskError::Database)?;
match config.db_flavor {
DbFlavor::ClickHouse => {
let client = ChClient::new(
&config.host,
config.port,
config.secure,
&config.user,
&config.password,
&config.database,
);
let version = client.ping().await?;
let arc = Arc::new(client);
state.ch_clients.write().await.insert(config.id.clone(), arc);
state.read_only.write().await.insert(config.id.clone(), true);
state
.db_flavors
.write()
.await
.insert(config.id.clone(), DbFlavor::ClickHouse);
Ok(ConnectResult {
version,
flavor: DbFlavor::ClickHouse,
})
}
_ => {
let pool = PgPool::connect(&config.connection_url())
.await
.map_err(TuskError::Database)?;
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(TuskError::Database)?;
let row = sqlx::query("SELECT version()")
.fetch_one(&pool)
.await
.map_err(TuskError::Database)?;
let version: String = row.get(0);
let flavor = if version.to_lowercase().contains("greenplum") {
DbFlavor::Greenplum
} else {
DbFlavor::PostgreSQL
};
state.pools.write().await.insert(config.id.clone(), pool);
state.read_only.write().await.insert(config.id.clone(), true);
state
.db_flavors
.write()
.await
.insert(config.id.clone(), flavor);
Ok(ConnectResult { version, flavor })
}
}
}
// Verify connection
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(TuskError::Database)?;
/// Core implementation of switching the active database for a connection.
/// Reusable from both the Tauri command (frontend-driven) and the chat agent
/// loop (model-driven via the switch_database tool).
pub(crate) async fn switch_database_core(
state: &AppState,
config: &ConnectionConfig,
database: &str,
) -> TuskResult<()> {
let mut switched = config.clone();
switched.database = database.to_string();
// Detect database flavor via version()
let row = sqlx::query("SELECT version()")
.fetch_one(&pool)
.await
.map_err(TuskError::Database)?;
let version: String = row.get(0);
let flavor = if version.to_lowercase().contains("greenplum") {
DbFlavor::Greenplum
} else {
DbFlavor::PostgreSQL
let result: TuskResult<()> = match config.db_flavor {
DbFlavor::ClickHouse => {
let client = ChClient::new(
&switched.host,
switched.port,
switched.secure,
&switched.user,
&switched.password,
&switched.database,
);
client.ping().await?;
state
.ch_clients
.write()
.await
.insert(config.id.clone(), Arc::new(client));
Ok(())
}
_ => {
let pool = PgPool::connect(&switched.connection_url())
.await
.map_err(TuskError::Database)?;
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(TuskError::Database)?;
let mut pools = state.pools.write().await;
if let Some(old_pool) = pools.remove(&config.id) {
old_pool.close().await;
}
pools.insert(config.id.clone(), pool);
Ok(())
}
};
let mut pools = state.pools.write().await;
pools.insert(config.id.clone(), pool);
// Drop every cache that's bound to this connection's previous database.
state.invalidate_chat_caches_for(&config.id).await;
let mut ro = state.read_only.write().await;
ro.insert(config.id.clone(), true);
let mut flavors = state.db_flavors.write().await;
flavors.insert(config.id.clone(), flavor);
Ok(ConnectResult { version, flavor })
result
}
#[tauri::command]
@@ -147,40 +266,12 @@ pub async fn switch_database(
config: ConnectionConfig,
database: String,
) -> TuskResult<()> {
let mut switched = config.clone();
switched.database = database;
let pool = PgPool::connect(&switched.connection_url())
.await
.map_err(TuskError::Database)?;
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(TuskError::Database)?;
let mut pools = state.pools.write().await;
if let Some(old_pool) = pools.remove(&config.id) {
old_pool.close().await;
}
pools.insert(config.id.clone(), pool);
Ok(())
switch_database_core(&state, &config, &database).await
}
#[tauri::command]
pub async fn disconnect(state: State<'_, Arc<AppState>>, id: String) -> TuskResult<()> {
let mut pools = state.pools.write().await;
if let Some(pool) = pools.remove(&id) {
pool.close().await;
}
let mut ro = state.read_only.write().await;
ro.remove(&id);
let mut flavors = state.db_flavors.write().await;
flavors.remove(&id);
close_connection(&state, &id).await;
Ok(())
}
+106 -2
View File
@@ -1,7 +1,7 @@
use crate::commands::queries::pg_value_to_json;
use crate::error::{TuskError, TuskResult};
use crate::models::query_result::PaginatedQueryResult;
use crate::state::AppState;
use crate::state::{AppState, DbFlavor};
use crate::utils::escape_ident;
use serde_json::Value;
use sqlx::{Column, Row, TypeInfo};
@@ -9,6 +9,80 @@ use std::sync::Arc;
use std::time::Instant;
use tauri::State;
async fn ch_get_table_data(
state: &AppState,
connection_id: &str,
schema: &str,
table: &str,
page: u32,
page_size: u32,
sort_column: Option<&str>,
sort_direction: Option<&str>,
filter: Option<&str>,
) -> TuskResult<PaginatedQueryResult> {
let client = state.get_ch_client(connection_id).await?;
let qualified = format!(
"{}.{}",
ch_quote_ident(schema),
ch_quote_ident(table)
);
let mut where_clause = String::new();
if let Some(f) = filter {
if !f.trim().is_empty() {
crate::db::sql_guard::ensure_readonly_sql(&format!("SELECT 1 FROM x WHERE {}", f))?;
where_clause = format!(" WHERE {}", f);
}
}
let mut order_clause = String::new();
if let Some(col) = sort_column {
if !col.trim().is_empty() {
let dir = match sort_direction {
Some("DESC") | Some("desc") => "DESC",
_ => "ASC",
};
order_clause = format!(" ORDER BY {} {}", ch_quote_ident(col), dir);
}
}
let offset = (page.saturating_sub(1)) as i64 * page_size as i64;
let data_sql = format!(
"SELECT * FROM {}{}{} LIMIT {} OFFSET {}",
qualified, where_clause, order_clause, page_size, offset
);
let count_sql = format!("SELECT count() AS c FROM {}{}", qualified, where_clause);
let result = client.execute_query(&data_sql, true).await?;
let count_rows = client.fetch_objects(&count_sql).await?;
let total_rows = count_rows
.first()
.and_then(|o| o.get("c"))
.and_then(|v| match v {
Value::Number(n) => n.as_i64(),
Value::String(s) => s.parse::<i64>().ok(),
_ => None,
})
.unwrap_or(0);
Ok(PaginatedQueryResult {
columns: result.columns,
types: result.types,
rows: result.rows,
row_count: result.row_count,
execution_time_ms: result.execution_time_ms,
total_rows,
page,
page_size,
ctids: vec![],
})
}
fn ch_quote_ident(s: &str) -> String {
let escaped = s.replace('`', "``");
format!("`{}`", escaped)
}
#[tauri::command]
#[allow(clippy::too_many_arguments)]
pub async fn get_table_data(
@@ -22,6 +96,21 @@ pub async fn get_table_data(
sort_direction: Option<String>,
filter: Option<String>,
) -> TuskResult<PaginatedQueryResult> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return ch_get_table_data(
&state,
&connection_id,
&schema,
&table,
page,
page_size,
sort_column.as_deref(),
sort_direction.as_deref(),
filter.as_deref(),
)
.await;
}
let pool = state.get_pool(&connection_id).await?;
let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
@@ -74,7 +163,7 @@ pub async fn get_table_data(
tx.rollback().await.map_err(TuskError::Database)?;
let execution_time_ms = start.elapsed().as_millis();
let execution_time_ms = start.elapsed().as_millis() as u64;
let total_rows: i64 = count_row.get(0);
let mut all_columns = Vec::new();
@@ -146,6 +235,11 @@ pub async fn update_row(
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
if matches!(state.get_flavor(&connection_id).await, DbFlavor::ClickHouse) {
return Err(TuskError::Custom(
"Inline row edit is not supported for ClickHouse — use SQL ALTER … UPDATE.".into(),
));
}
let pool = state.get_pool(&connection_id).await?;
@@ -202,6 +296,11 @@ pub async fn insert_row(
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
if matches!(state.get_flavor(&connection_id).await, DbFlavor::ClickHouse) {
return Err(TuskError::Custom(
"Inline row insert is not supported for ClickHouse — use SQL INSERT.".into(),
));
}
let pool = state.get_pool(&connection_id).await?;
@@ -240,6 +339,11 @@ pub async fn delete_rows(
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
if matches!(state.get_flavor(&connection_id).await, DbFlavor::ClickHouse) {
return Err(TuskError::Custom(
"Inline row delete is not supported for ClickHouse — use SQL ALTER … DELETE.".into(),
));
}
let pool = state.get_pool(&connection_id).await?;
File diff suppressed because it is too large Load Diff
-367
View File
@@ -1,367 +0,0 @@
use crate::commands::queries::pg_value_to_json;
use crate::error::TuskResult;
use crate::models::connection::ConnectionConfig;
use crate::models::lookup::{
EntityLookupResult, LookupDatabaseResult, LookupProgress, LookupTableMatch,
};
use crate::utils::escape_ident;
use sqlx::postgres::PgPoolOptions;
use sqlx::{Column, Row, TypeInfo};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Instant;
use tauri::{AppHandle, Emitter};
use tokio::sync::Semaphore;
struct TableCandidate {
schema: String,
table: String,
data_type: String,
}
async fn search_database(
config: &ConnectionConfig,
database: &str,
column_name: &str,
value: &str,
) -> LookupDatabaseResult {
let start = Instant::now();
let mut db_config = config.clone();
db_config.database = database.to_string();
let url = db_config.connection_url();
let pool = match PgPoolOptions::new()
.max_connections(2)
.acquire_timeout(std::time::Duration::from_secs(5))
.connect(&url)
.await
{
Ok(p) => p,
Err(e) => {
return LookupDatabaseResult {
database: database.to_string(),
tables: vec![],
error: Some(format!("Connection failed: {}", e)),
search_time_ms: start.elapsed().as_millis(),
};
}
};
let result = tokio::time::timeout(
std::time::Duration::from_secs(120),
search_database_inner(&pool, database, column_name, value),
)
.await;
pool.close().await;
match result {
Ok(db_result) => {
let mut db_result = db_result;
db_result.search_time_ms = start.elapsed().as_millis();
db_result
}
Err(_) => LookupDatabaseResult {
database: database.to_string(),
tables: vec![],
error: Some("Timeout (120s)".to_string()),
search_time_ms: start.elapsed().as_millis(),
},
}
}
async fn search_database_inner(
pool: &sqlx::PgPool,
database: &str,
column_name: &str,
value: &str,
) -> LookupDatabaseResult {
// Find tables that have this column
let candidates = match sqlx::query_as::<_, (String, String, String)>(
"SELECT table_schema, table_name, data_type \
FROM information_schema.columns \
WHERE column_name = $1 \
AND table_schema NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit')",
)
.bind(column_name)
.fetch_all(pool)
.await
{
Ok(rows) => rows
.into_iter()
.map(|(schema, table, data_type)| TableCandidate {
schema,
table,
data_type,
})
.collect::<Vec<_>>(),
Err(e) => {
return LookupDatabaseResult {
database: database.to_string(),
tables: vec![],
error: Some(format!("Schema query failed: {}", e)),
search_time_ms: 0,
};
}
};
let mut tables = Vec::new();
for candidate in &candidates {
let qualified = format!(
"{}.{}",
escape_ident(&candidate.schema),
escape_ident(&candidate.table)
);
let col_ident = escape_ident(column_name);
// Read-only transaction: SELECT rows + COUNT
let select_sql = format!(
"SELECT * FROM {} WHERE {}::text = $1 LIMIT 50",
qualified, col_ident
);
let count_sql = format!(
"SELECT COUNT(*) FROM {} WHERE {}::text = $1",
qualified, col_ident
);
let mut tx = match pool.begin().await {
Ok(tx) => tx,
Err(e) => {
tables.push(LookupTableMatch {
schema: candidate.schema.clone(),
table: candidate.table.clone(),
column_type: candidate.data_type.clone(),
columns: vec![],
types: vec![],
rows: vec![],
row_count: 0,
total_count: 0,
});
log::warn!(
"Failed to begin tx for {}.{}: {}",
candidate.schema,
candidate.table,
e
);
continue;
}
};
if let Err(e) = sqlx::query("SET TRANSACTION READ ONLY")
.execute(&mut *tx)
.await
{
let _ = tx.rollback().await;
log::warn!("Failed SET TRANSACTION READ ONLY: {}", e);
continue;
}
let rows_result = sqlx::query(&select_sql)
.bind(value)
.fetch_all(&mut *tx)
.await;
let count_result: Result<i64, _> = sqlx::query_scalar(&count_sql)
.bind(value)
.fetch_one(&mut *tx)
.await;
let _ = tx.rollback().await;
match rows_result {
Ok(rows) if !rows.is_empty() => {
let mut col_names = Vec::new();
let mut col_types = Vec::new();
if let Some(first) = rows.first() {
for col in first.columns() {
col_names.push(col.name().to_string());
col_types.push(col.type_info().name().to_string());
}
}
let result_rows: Vec<Vec<serde_json::Value>> = rows
.iter()
.map(|row| {
(0..col_names.len())
.map(|i| pg_value_to_json(row, i))
.collect()
})
.collect();
let row_count = result_rows.len();
let total_count = count_result.unwrap_or(row_count as i64);
tables.push(LookupTableMatch {
schema: candidate.schema.clone(),
table: candidate.table.clone(),
column_type: candidate.data_type.clone(),
columns: col_names,
types: col_types,
rows: result_rows,
row_count,
total_count,
});
}
Ok(_) => {
// No rows matched — skip
}
Err(e) => {
log::warn!(
"Query failed for {}.{}: {}",
candidate.schema,
candidate.table,
e
);
}
}
}
LookupDatabaseResult {
database: database.to_string(),
tables,
error: None,
search_time_ms: 0,
}
}
#[tauri::command]
pub async fn entity_lookup(
app: AppHandle,
config: ConnectionConfig,
column_name: String,
value: String,
databases: Option<Vec<String>>,
lookup_id: String,
) -> TuskResult<EntityLookupResult> {
let start = Instant::now();
// 1. Get list of databases
let url = config.connection_url();
let pool = PgPoolOptions::new()
.max_connections(1)
.acquire_timeout(std::time::Duration::from_secs(5))
.connect(&url)
.await
.map_err(crate::error::TuskError::Database)?;
let db_names: Vec<String> = sqlx::query_scalar(
"SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname",
)
.fetch_all(&pool)
.await
.map_err(crate::error::TuskError::Database)?;
pool.close().await;
// Filter if specific databases requested
let db_names: Vec<String> = if let Some(ref filter) = databases {
db_names
.into_iter()
.filter(|d| filter.contains(d))
.collect()
} else {
db_names
};
let total = db_names.len();
let completed = Arc::new(AtomicUsize::new(0));
let semaphore = Arc::new(Semaphore::new(5));
// 2. Parallel search across databases
let mut handles = Vec::new();
for db_name in db_names {
let config = config.clone();
let column_name = column_name.clone();
let value = value.clone();
let lookup_id = lookup_id.clone();
let app = app.clone();
let semaphore = semaphore.clone();
let completed = completed.clone();
let handle = tokio::spawn(async move {
let _permit = semaphore.acquire().await.unwrap();
// Emit "searching" progress
let _ = app.emit(
"lookup-progress",
LookupProgress {
lookup_id: lookup_id.clone(),
database: db_name.clone(),
status: "searching".to_string(),
tables_found: 0,
rows_found: 0,
error: None,
completed: completed.load(Ordering::Relaxed),
total,
},
);
let result = search_database(&config, &db_name, &column_name, &value).await;
let done = completed.fetch_add(1, Ordering::Relaxed) + 1;
let status = if result.error.is_some() {
"error"
} else {
"done"
};
let _ = app.emit(
"lookup-progress",
LookupProgress {
lookup_id: lookup_id.clone(),
database: db_name.clone(),
status: status.to_string(),
tables_found: result.tables.len(),
rows_found: result.tables.iter().map(|t| t.row_count).sum(),
error: result.error.clone(),
completed: done,
total,
},
);
result
});
handles.push(handle);
}
// 3. Collect results
let mut all_results = Vec::new();
for handle in handles {
match handle.await {
Ok(result) => all_results.push(result),
Err(e) => {
log::error!("Join error: {}", e);
}
}
}
// Sort: databases with matches first, then by name
all_results.sort_by(|a, b| {
let a_has = !a.tables.is_empty();
let b_has = !b.tables.is_empty();
b_has.cmp(&a_has).then(a.database.cmp(&b.database))
});
let total_databases_searched = all_results.len();
let total_tables_matched: usize = all_results.iter().map(|d| d.tables.len()).sum();
let total_rows_found: usize = all_results
.iter()
.flat_map(|d| d.tables.iter())
.map(|t| t.row_count)
.sum();
Ok(EntityLookupResult {
column_name,
value,
databases: all_results,
total_databases_searched,
total_tables_matched,
total_rows_found,
total_time_ms: start.elapsed().as_millis(),
})
}
-594
View File
@@ -1,594 +0,0 @@
use crate::error::{TuskError, TuskResult};
use crate::models::management::*;
use crate::state::{AppState, DbFlavor};
use crate::utils::escape_ident;
use sqlx::Row;
use std::sync::Arc;
use tauri::State;
#[tauri::command]
pub async fn get_database_info(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<DatabaseInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let rows = sqlx::query(
"SELECT d.datname, \
pg_catalog.pg_get_userbyid(d.datdba) AS owner, \
pg_catalog.pg_encoding_to_char(d.encoding) AS encoding, \
d.datcollate, \
d.datctype, \
COALESCE(t.spcname, 'pg_default') AS tablespace, \
d.datconnlimit, \
pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(d.datname)) AS size, \
pg_catalog.shobj_description(d.oid, 'pg_database') AS description \
FROM pg_catalog.pg_database d \
LEFT JOIN pg_catalog.pg_tablespace t ON d.dattablespace = t.oid \
WHERE NOT d.datistemplate \
ORDER BY d.datname",
)
.fetch_all(pool)
.await
.map_err(TuskError::Database)?;
let databases = rows
.iter()
.map(|row| DatabaseInfo {
name: row.get("datname"),
owner: row.get("owner"),
encoding: row.get("encoding"),
collation: row.get("datcollate"),
ctype: row.get("datctype"),
tablespace: row.get("tablespace"),
connection_limit: row.get("datconnlimit"),
size: row.get("size"),
description: row.get("description"),
})
.collect();
Ok(databases)
}
#[tauri::command]
pub async fn create_database(
state: State<'_, Arc<AppState>>,
connection_id: String,
params: CreateDatabaseParams,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let mut sql = format!("CREATE DATABASE {}", escape_ident(&params.name));
if let Some(ref owner) = params.owner {
sql.push_str(&format!(" OWNER {}", escape_ident(owner)));
}
if let Some(ref template) = params.template {
sql.push_str(&format!(" TEMPLATE {}", escape_ident(template)));
}
if let Some(ref encoding) = params.encoding {
sql.push_str(&format!(" ENCODING '{}'", encoding.replace('\'', "''")));
}
if let Some(ref tablespace) = params.tablespace {
sql.push_str(&format!(" TABLESPACE {}", escape_ident(tablespace)));
}
if let Some(limit) = params.connection_limit {
sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
}
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn drop_database(
state: State<'_, Arc<AppState>>,
connection_id: String,
name: String,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
// Terminate active connections to the target database
sqlx::query("SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1::name AND pid <> pg_backend_pid()")
.bind(&name)
.execute(pool)
.await
.map_err(TuskError::Database)?;
let drop_sql = format!("DROP DATABASE {}", escape_ident(&name));
sqlx::query(&drop_sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn list_roles(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<RoleInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let rows = sqlx::query(
"SELECT r.rolname, \
r.rolsuper, \
r.rolcanlogin, \
r.rolcreatedb, \
r.rolcreaterole, \
r.rolinherit, \
r.rolreplication, \
r.rolconnlimit, \
r.rolpassword IS NOT NULL AS password_set, \
r.rolvaliduntil::text, \
COALESCE(( \
SELECT array_agg(g.rolname ORDER BY g.rolname) \
FROM pg_catalog.pg_auth_members m \
JOIN pg_catalog.pg_roles g ON m.roleid = g.oid \
WHERE m.member = r.oid \
), ARRAY[]::text[]) AS member_of, \
COALESCE(( \
SELECT array_agg(m2.rolname ORDER BY m2.rolname) \
FROM pg_catalog.pg_auth_members am \
JOIN pg_catalog.pg_roles m2 ON am.member = m2.oid \
WHERE am.roleid = r.oid \
), ARRAY[]::text[]) AS members, \
pg_catalog.shobj_description(r.oid, 'pg_authid') AS description \
FROM pg_catalog.pg_roles r \
WHERE r.rolname !~ '^pg_' \
ORDER BY r.rolname",
)
.fetch_all(pool)
.await
.map_err(TuskError::Database)?;
let roles = rows
.iter()
.map(|row| RoleInfo {
name: row.get("rolname"),
is_superuser: row.get("rolsuper"),
can_login: row.get("rolcanlogin"),
can_create_db: row.get("rolcreatedb"),
can_create_role: row.get("rolcreaterole"),
inherit: row.get("rolinherit"),
is_replication: row.get("rolreplication"),
connection_limit: row.get("rolconnlimit"),
password_set: row.get("password_set"),
valid_until: row.get("rolvaliduntil"),
member_of: row.get("member_of"),
members: row.get("members"),
description: row.get("description"),
})
.collect();
Ok(roles)
}
#[tauri::command]
pub async fn create_role(
state: State<'_, Arc<AppState>>,
connection_id: String,
params: CreateRoleParams,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let mut sql = format!("CREATE ROLE {}", escape_ident(&params.name));
let mut options = Vec::new();
options.push(if params.login { "LOGIN" } else { "NOLOGIN" });
options.push(if params.superuser {
"SUPERUSER"
} else {
"NOSUPERUSER"
});
options.push(if params.createdb {
"CREATEDB"
} else {
"NOCREATEDB"
});
options.push(if params.createrole {
"CREATEROLE"
} else {
"NOCREATEROLE"
});
options.push(if params.inherit {
"INHERIT"
} else {
"NOINHERIT"
});
options.push(if params.replication {
"REPLICATION"
} else {
"NOREPLICATION"
});
if let Some(ref password) = params.password {
options.push("PASSWORD");
// Will be appended separately
sql.push_str(&format!(" {}", options.join(" ")));
sql.push_str(&format!(" '{}'", password.replace('\'', "''")));
} else {
sql.push_str(&format!(" {}", options.join(" ")));
}
if let Some(limit) = params.connection_limit {
sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
}
if let Some(ref valid_until) = params.valid_until {
sql.push_str(&format!(
" VALID UNTIL '{}'",
valid_until.replace('\'', "''")
));
}
if !params.in_roles.is_empty() {
let roles: Vec<String> = params.in_roles.iter().map(|r| escape_ident(r)).collect();
sql.push_str(&format!(" IN ROLE {}", roles.join(", ")));
}
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn alter_role(
state: State<'_, Arc<AppState>>,
connection_id: String,
params: AlterRoleParams,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let mut options = Vec::new();
if let Some(login) = params.login {
options.push(if login {
"LOGIN".to_string()
} else {
"NOLOGIN".to_string()
});
}
if let Some(superuser) = params.superuser {
options.push(if superuser {
"SUPERUSER".to_string()
} else {
"NOSUPERUSER".to_string()
});
}
if let Some(createdb) = params.createdb {
options.push(if createdb {
"CREATEDB".to_string()
} else {
"NOCREATEDB".to_string()
});
}
if let Some(createrole) = params.createrole {
options.push(if createrole {
"CREATEROLE".to_string()
} else {
"NOCREATEROLE".to_string()
});
}
if let Some(inherit) = params.inherit {
options.push(if inherit {
"INHERIT".to_string()
} else {
"NOINHERIT".to_string()
});
}
if let Some(replication) = params.replication {
options.push(if replication {
"REPLICATION".to_string()
} else {
"NOREPLICATION".to_string()
});
}
if let Some(ref password) = params.password {
options.push(format!("PASSWORD '{}'", password.replace('\'', "''")));
}
if let Some(limit) = params.connection_limit {
options.push(format!("CONNECTION LIMIT {}", limit));
}
if let Some(ref valid_until) = params.valid_until {
options.push(format!("VALID UNTIL '{}'", valid_until.replace('\'', "''")));
}
if !options.is_empty() {
let sql = format!(
"ALTER ROLE {} {}",
escape_ident(&params.name),
options.join(" ")
);
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
}
if let Some(ref new_name) = params.rename_to {
let sql = format!(
"ALTER ROLE {} RENAME TO {}",
escape_ident(&params.name),
escape_ident(new_name)
);
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
}
Ok(())
}
#[tauri::command]
pub async fn drop_role(
state: State<'_, Arc<AppState>>,
connection_id: String,
name: String,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let sql = format!("DROP ROLE {}", escape_ident(&name));
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn get_table_privileges(
state: State<'_, Arc<AppState>>,
connection_id: String,
schema: String,
table: String,
) -> TuskResult<Vec<TablePrivilege>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let rows = sqlx::query(
"SELECT grantee, table_schema, table_name, privilege_type, \
is_grantable = 'YES' AS is_grantable \
FROM information_schema.role_table_grants \
WHERE table_schema = $1 AND table_name = $2 \
ORDER BY grantee, privilege_type",
)
.bind(&schema)
.bind(&table)
.fetch_all(pool)
.await
.map_err(TuskError::Database)?;
let privileges = rows
.iter()
.map(|row| TablePrivilege {
grantee: row.get("grantee"),
table_schema: row.get("table_schema"),
table_name: row.get("table_name"),
privilege_type: row.get("privilege_type"),
is_grantable: row.get("is_grantable"),
})
.collect();
Ok(privileges)
}
#[tauri::command]
pub async fn grant_revoke(
state: State<'_, Arc<AppState>>,
connection_id: String,
params: GrantRevokeParams,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let privs = params.privileges.join(", ");
let object_type = params.object_type.to_uppercase();
let object_ref = escape_ident(&params.object_name);
let role_ref = escape_ident(&params.role_name);
let sql = if params.action.to_uppercase() == "GRANT" {
let grant_option = if params.with_grant_option {
" WITH GRANT OPTION"
} else {
""
};
format!(
"GRANT {} ON {} {} TO {}{}",
privs, object_type, object_ref, role_ref, grant_option
)
} else {
format!(
"REVOKE {} ON {} {} FROM {}",
privs, object_type, object_ref, role_ref
)
};
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn manage_role_membership(
state: State<'_, Arc<AppState>>,
connection_id: String,
params: RoleMembershipParams,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let role_ref = escape_ident(&params.role_name);
let member_ref = escape_ident(&params.member_name);
let sql = if params.action.to_uppercase() == "GRANT" {
format!("GRANT {} TO {}", role_ref, member_ref)
} else {
format!("REVOKE {} FROM {}", role_ref, member_ref)
};
sqlx::query(&sql)
.execute(pool)
.await
.map_err(TuskError::Database)?;
Ok(())
}
#[tauri::command]
pub async fn list_sessions(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<SessionInfo>> {
let flavor = state.get_flavor(&connection_id).await;
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let sql = if flavor == DbFlavor::Greenplum {
"SELECT pid, usename, datname, state, query, \
query_start::text, NULL::text as wait_event_type, NULL::text as wait_event, \
client_addr::text \
FROM pg_stat_activity \
WHERE datname IS NOT NULL \
ORDER BY query_start DESC NULLS LAST"
} else {
"SELECT pid, usename, datname, state, query, \
query_start::text, wait_event_type, wait_event, \
client_addr::text \
FROM pg_stat_activity \
WHERE datname IS NOT NULL \
ORDER BY query_start DESC NULLS LAST"
};
let rows = sqlx::query(sql)
.fetch_all(pool)
.await
.map_err(TuskError::Database)?;
let sessions = rows
.iter()
.map(|row| SessionInfo {
pid: row.get("pid"),
usename: row.get("usename"),
datname: row.get("datname"),
state: row.get("state"),
query: row.get("query"),
query_start: row.get("query_start"),
wait_event_type: row.get("wait_event_type"),
wait_event: row.get("wait_event"),
client_addr: row.get("client_addr"),
})
.collect();
Ok(sessions)
}
#[tauri::command]
pub async fn cancel_query(
state: State<'_, Arc<AppState>>,
connection_id: String,
pid: i32,
) -> TuskResult<bool> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let row = sqlx::query("SELECT pg_cancel_backend($1)")
.bind(pid)
.fetch_one(pool)
.await
.map_err(TuskError::Database)?;
Ok(row.get::<bool, _>(0))
}
#[tauri::command]
pub async fn terminate_backend(
state: State<'_, Arc<AppState>>,
connection_id: String,
pid: i32,
) -> TuskResult<bool> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let row = sqlx::query("SELECT pg_terminate_backend($1)")
.bind(pid)
.fetch_one(pool)
.await
.map_err(TuskError::Database)?;
Ok(row.get::<bool, _>(0))
}
+214
View File
@@ -0,0 +1,214 @@
//! Per-connection long-term memory for the chat agent (F1).
//!
//! Stored as a markdown file at `<app_data_dir>/memory/<connection_id>.md`.
//! The agent appends notes via the `remember` tool; the user can view and edit
//! the file in the Memory sidebar tab. The same content is injected into the
//! LEARNED NOTES section of the system prompt every turn.
use crate::error::{TuskError, TuskResult};
use chrono::Utc;
use std::fs;
use std::path::PathBuf;
use tauri::{AppHandle, Manager};
/// Soft cap on memory file size. Overflow drops oldest `## ts` blocks until fits.
pub const MEMORY_BYTE_CAP: usize = 16 * 1024;
pub(crate) fn get_memory_path(
app: &AppHandle,
connection_id: &str,
) -> TuskResult<PathBuf> {
let dir = app
.path()
.app_data_dir()
.map_err(|e| TuskError::Config(e.to_string()))?
.join("memory");
fs::create_dir_all(&dir)?;
let safe = sanitize_connection_id(connection_id);
Ok(dir.join(format!("{}.md", safe)))
}
fn sanitize_connection_id(id: &str) -> String {
id.chars()
.map(|c| match c {
'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' => c,
_ => '_',
})
.collect()
}
pub(crate) fn read_memory_core(
app: &AppHandle,
connection_id: &str,
) -> TuskResult<String> {
let path = get_memory_path(app, connection_id)?;
if !path.exists() {
return Ok(String::new());
}
Ok(fs::read_to_string(&path)?)
}
pub(crate) fn write_memory_core(
app: &AppHandle,
connection_id: &str,
content: &str,
) -> TuskResult<()> {
let path = get_memory_path(app, connection_id)?;
let trimmed = enforce_size_cap(content, MEMORY_BYTE_CAP);
fs::write(&path, trimmed)?;
Ok(())
}
pub(crate) fn append_memory_core(
app: &AppHandle,
connection_id: &str,
note: &str,
) -> TuskResult<()> {
let trimmed_note = note.trim();
if trimmed_note.is_empty() {
return Err(TuskError::Custom("remember: note must not be empty".into()));
}
let existing = read_memory_core(app, connection_id)?;
let mut buf = String::new();
if existing.is_empty() {
buf.push_str("# Memory\n\n");
} else {
buf.push_str(&existing);
if !buf.ends_with('\n') {
buf.push('\n');
}
if !buf.ends_with("\n\n") {
buf.push('\n');
}
}
let ts = Utc::now().format("%Y-%m-%dT%H:%M:%SZ");
buf.push_str(&format!("## {}\n{}\n", ts, trimmed_note));
let final_content = enforce_size_cap(&buf, MEMORY_BYTE_CAP);
let path = get_memory_path(app, connection_id)?;
fs::write(&path, final_content)?;
Ok(())
}
/// Trim the file from the *oldest* note (top) until it fits within `cap` bytes.
/// Always preserves the trailing notes (the most recent observations). Keeps
/// the leading `# Memory\n\n` header if present.
pub(crate) fn enforce_size_cap(content: &str, cap: usize) -> String {
if content.len() <= cap {
return content.to_string();
}
let header = if content.starts_with("# Memory") {
match content.find("\n## ") {
Some(pos) => &content[..pos + 1],
None => "# Memory\n\n",
}
} else {
""
};
// Split into note blocks by "\n## " marker.
// First block (after header) might lack the leading "## " — handle uniformly.
let body_start = header.len();
let body = &content[body_start..];
let mut blocks: Vec<&str> = Vec::new();
let mut idx = 0;
while idx < body.len() {
// Find the next "\n## " starting at idx; if not found, the rest is one block.
let rel = body[idx..].find("\n## ");
match rel {
Some(r) => {
blocks.push(&body[idx..idx + r + 1]); // include trailing newline before next block
idx = idx + r + 1; // start of "## "
}
None => {
blocks.push(&body[idx..]);
break;
}
}
}
// Drop blocks from the front until total fits.
let mut current_size: usize = header.len() + blocks.iter().map(|b| b.len()).sum::<usize>();
let mut start = 0usize;
while current_size > cap && start < blocks.len() {
current_size -= blocks[start].len();
start += 1;
}
let mut out = String::with_capacity(current_size);
out.push_str(header);
for b in &blocks[start..] {
out.push_str(b);
}
out
}
#[tauri::command]
pub async fn get_memory(app: AppHandle, connection_id: String) -> TuskResult<String> {
read_memory_core(&app, &connection_id)
}
#[tauri::command]
pub async fn save_memory(
app: AppHandle,
connection_id: String,
content: String,
) -> TuskResult<()> {
write_memory_core(&app, &connection_id, &content)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cap_passthrough_under_limit() {
let small = "# Memory\n\n## 2026-01-01T00:00:00Z\nshort note\n";
assert_eq!(enforce_size_cap(small, MEMORY_BYTE_CAP), small);
}
#[test]
fn cap_drops_oldest_blocks() {
// 3 blocks of ~6KB each -> 18KB total > 16KB cap
let block_body = "x".repeat(6000);
let content = format!(
"# Memory\n\n## 2026-01-01T00:00:00Z\n{body}\n## 2026-02-01T00:00:00Z\n{body}\n## 2026-03-01T00:00:00Z\n{body}\n",
body = block_body
);
assert!(content.len() > MEMORY_BYTE_CAP);
let trimmed = enforce_size_cap(&content, MEMORY_BYTE_CAP);
assert!(trimmed.len() <= MEMORY_BYTE_CAP);
// Most recent block must survive.
assert!(trimmed.contains("2026-03-01T00:00:00Z"));
// Oldest must be dropped.
assert!(!trimmed.contains("2026-01-01T00:00:00Z"));
// Header preserved.
assert!(trimmed.starts_with("# Memory"));
}
#[test]
fn cap_keeps_only_latest_when_single_block_huge() {
let block_body = "y".repeat(20_000);
let content = format!(
"# Memory\n\n## 2026-01-01T00:00:00Z\n{}\n",
block_body
);
let trimmed = enforce_size_cap(&content, MEMORY_BYTE_CAP);
// Even after dropping that single block we keep at least the header,
// so the result is just the header (or close to it).
assert!(trimmed.starts_with("# Memory"));
assert!(trimmed.len() <= MEMORY_BYTE_CAP);
}
#[test]
fn sanitize_strips_path_chars() {
assert_eq!(sanitize_connection_id("abc/../etc"), "abc____etc");
assert_eq!(
sanitize_connection_id("cf9feefd-59ab-4a7c"),
"cf9feefd-59ab-4a7c"
);
}
}
+3 -4
View File
@@ -1,13 +1,12 @@
pub mod ai;
pub mod chat;
pub mod chat_tools;
pub mod connections;
pub mod data;
pub mod docker;
pub mod export;
pub mod history;
pub mod lookup;
pub mod management;
pub mod memory;
pub mod queries;
pub mod saved_queries;
pub mod schema;
pub mod settings;
pub mod snapshot;
+13 -2
View File
@@ -1,6 +1,7 @@
use crate::db::sql_guard::ensure_readonly_sql;
use crate::error::{TuskError, TuskResult};
use crate::models::query_result::QueryResult;
use crate::state::AppState;
use crate::state::{AppState, DbFlavor};
use serde_json::Value;
use sqlx::postgres::PgRow;
use sqlx::{Column, Row, TypeInfo};
@@ -81,6 +82,16 @@ pub async fn execute_query_core(
sql: &str,
) -> TuskResult<QueryResult> {
let read_only = state.is_read_only(connection_id).await;
let flavor = state.get_flavor(connection_id).await;
if read_only {
ensure_readonly_sql(sql)?;
}
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(connection_id).await?;
return client.execute_query(sql, read_only).await;
}
let pools = state.pools.read().await;
let pool = pools
@@ -106,7 +117,7 @@ pub async fn execute_query_core(
.await
.map_err(TuskError::Database)?
};
let execution_time_ms = start.elapsed().as_millis();
let execution_time_ms = start.elapsed().as_millis() as u64;
let mut columns = Vec::new();
let mut types = Vec::new();
+20 -9
View File
@@ -12,12 +12,11 @@ fn get_saved_queries_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
Ok(dir.join("saved_queries.json"))
}
#[tauri::command]
pub async fn list_saved_queries(
app: AppHandle,
search: Option<String>,
pub(crate) async fn list_saved_queries_core(
app: &AppHandle,
search: Option<&str>,
) -> TuskResult<Vec<SavedQuery>> {
let path = get_saved_queries_path(&app)?;
let path = get_saved_queries_path(app)?;
if !path.exists() {
return Ok(vec![]);
}
@@ -27,7 +26,7 @@ pub async fn list_saved_queries(
let filtered: Vec<SavedQuery> = entries
.into_iter()
.filter(|e| {
if let Some(ref s) = search {
if let Some(s) = search {
let lower = s.to_lowercase();
e.name.to_lowercase().contains(&lower) || e.sql.to_lowercase().contains(&lower)
} else {
@@ -39,9 +38,8 @@ pub async fn list_saved_queries(
Ok(filtered)
}
#[tauri::command]
pub async fn save_query(app: AppHandle, query: SavedQuery) -> TuskResult<()> {
let path = get_saved_queries_path(&app)?;
pub(crate) async fn save_query_core(app: &AppHandle, query: SavedQuery) -> TuskResult<()> {
let path = get_saved_queries_path(app)?;
let mut entries = if path.exists() {
let data = fs::read_to_string(&path)?;
serde_json::from_str::<Vec<SavedQuery>>(&data).unwrap_or_default()
@@ -56,6 +54,19 @@ pub async fn save_query(app: AppHandle, query: SavedQuery) -> TuskResult<()> {
Ok(())
}
#[tauri::command]
pub async fn list_saved_queries(
app: AppHandle,
search: Option<String>,
) -> TuskResult<Vec<SavedQuery>> {
list_saved_queries_core(&app, search.as_deref()).await
}
#[tauri::command]
pub async fn save_query(app: AppHandle, query: SavedQuery) -> TuskResult<()> {
save_query_core(&app, query).await
}
#[tauri::command]
pub async fn delete_saved_query(app: AppHandle, id: String) -> TuskResult<()> {
let path = get_saved_queries_path(&app)?;
+188 -133
View File
@@ -1,20 +1,53 @@
use crate::error::{TuskError, TuskResult};
use crate::models::schema::{
ColumnDetail, ColumnInfo, ConstraintInfo, ErdColumn, ErdData, ErdRelationship, ErdTable,
IndexInfo, SchemaObject, TriggerInfo,
ColumnDetail, ColumnInfo, ConstraintInfo, IndexInfo, SchemaObject, TriggerInfo,
};
use crate::state::{AppState, DbFlavor};
use serde_json::Value;
use sqlx::Row;
use std::collections::HashMap;
use std::sync::Arc;
use tauri::State;
#[tauri::command]
pub async fn list_databases(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<String>> {
let pool = state.get_pool(&connection_id).await?;
fn ch_string_literal(s: &str) -> String {
let escaped = s.replace('\\', "\\\\").replace('\'', "\\'");
format!("'{}'", escaped)
}
fn ch_obj_string(obj: &serde_json::Map<String, Value>, key: &str) -> Option<String> {
obj.get(key).and_then(|v| match v {
Value::String(s) => Some(s.clone()),
Value::Number(n) => Some(n.to_string()),
_ => None,
})
}
fn ch_obj_i64(obj: &serde_json::Map<String, Value>, key: &str) -> Option<i64> {
obj.get(key).and_then(|v| match v {
Value::Number(n) => n.as_i64(),
Value::String(s) => s.parse::<i64>().ok(),
_ => None,
})
}
pub async fn list_databases_core(state: &AppState, connection_id: &str) -> TuskResult<Vec<String>> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(connection_id).await?;
let rows = client
.fetch_objects(
"SELECT name FROM system.databases \
WHERE name NOT IN ('system','INFORMATION_SCHEMA','information_schema') \
ORDER BY name",
)
.await?;
return Ok(rows
.iter()
.filter_map(|o| ch_obj_string(o, "name"))
.collect());
}
let pool = state.get_pool(connection_id).await?;
let rows = sqlx::query(
"SELECT datname FROM pg_database \
@@ -28,10 +61,24 @@ pub async fn list_databases(
Ok(rows.iter().map(|r| r.get::<String, _>(0)).collect())
}
#[tauri::command]
pub async fn list_databases(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<String>> {
list_databases_core(&state, &connection_id).await
}
pub async fn list_schemas_core(state: &AppState, connection_id: &str) -> TuskResult<Vec<String>> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
// ClickHouse has no schema layer — surface the active database as a virtual schema.
let client = state.get_ch_client(connection_id).await?;
return Ok(vec![client.database.clone()]);
}
let pool = state.get_pool(connection_id).await?;
let flavor = state.get_flavor(connection_id).await;
let sql = if flavor == DbFlavor::Greenplum {
"SELECT schema_name FROM information_schema.schemata \
WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'gp_toolkit') \
@@ -63,6 +110,29 @@ pub async fn list_tables_core(
connection_id: &str,
schema: &str,
) -> TuskResult<Vec<SchemaObject>> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(connection_id).await?;
let escaped = ch_string_literal(schema);
let sql = format!(
"SELECT name, total_rows, total_bytes FROM system.tables \
WHERE database = {} AND engine NOT LIKE '%View' \
ORDER BY name",
escaped
);
let rows = client.fetch_objects(&sql).await?;
return Ok(rows
.iter()
.map(|o| SchemaObject {
name: ch_obj_string(o, "name").unwrap_or_default(),
object_type: "table".to_string(),
schema: schema.to_string(),
row_count: ch_obj_i64(o, "total_rows"),
size_bytes: ch_obj_i64(o, "total_bytes"),
})
.collect());
}
let pool = state.get_pool(connection_id).await?;
let rows = sqlx::query(
@@ -107,6 +177,28 @@ pub async fn list_views(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(&connection_id).await?;
let sql = format!(
"SELECT name FROM system.tables \
WHERE database = {} AND engine LIKE '%View' \
ORDER BY name",
ch_string_literal(&schema)
);
let rows = client.fetch_objects(&sql).await?;
return Ok(rows
.iter()
.map(|o| SchemaObject {
name: ch_obj_string(o, "name").unwrap_or_default(),
object_type: "view".to_string(),
schema: schema.clone(),
row_count: None,
size_bytes: None,
})
.collect());
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -137,6 +229,11 @@ pub async fn list_functions(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
// ClickHouse functions are global, not schema-scoped — surface empty here.
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -167,6 +264,10 @@ pub async fn list_indexes(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -197,6 +298,10 @@ pub async fn list_sequences(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -227,6 +332,36 @@ pub async fn get_table_columns_core(
schema: &str,
table: &str,
) -> TuskResult<Vec<ColumnInfo>> {
let flavor = state.get_flavor(connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(connection_id).await?;
let sql = format!(
"SELECT name, type, default_expression, is_in_primary_key, comment, position \
FROM system.columns WHERE database = {} AND table = {} \
ORDER BY position",
ch_string_literal(schema),
ch_string_literal(table)
);
let rows = client.fetch_objects(&sql).await?;
return Ok(rows
.iter()
.map(|o| {
let type_str = ch_obj_string(o, "type").unwrap_or_default();
let is_nullable = type_str.starts_with("Nullable(");
ColumnInfo {
name: ch_obj_string(o, "name").unwrap_or_default(),
data_type: type_str,
is_nullable,
column_default: ch_obj_string(o, "default_expression"),
ordinal_position: ch_obj_i64(o, "position").unwrap_or(0) as i32,
character_maximum_length: None,
is_primary_key: ch_obj_i64(o, "is_in_primary_key").unwrap_or(0) != 0,
comment: ch_obj_string(o, "comment"),
}
})
.collect());
}
let pool = state.get_pool(connection_id).await?;
let rows = sqlx::query(
@@ -296,6 +431,10 @@ pub async fn get_table_constraints(
schema: String,
table: String,
) -> TuskResult<Vec<ConstraintInfo>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -372,6 +511,10 @@ pub async fn get_table_indexes(
schema: String,
table: String,
) -> TuskResult<Vec<IndexInfo>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -410,6 +553,25 @@ pub async fn get_completion_schema(
connection_id: String,
) -> TuskResult<HashMap<String, HashMap<String, Vec<String>>>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let client = state.get_ch_client(&connection_id).await?;
let sql = format!(
"SELECT database, table, name FROM system.columns \
WHERE database = {} \
ORDER BY database, table, position",
ch_string_literal(&client.database)
);
let rows = client.fetch_objects(&sql).await?;
let mut result: HashMap<String, HashMap<String, Vec<String>>> = HashMap::new();
for row in rows {
let db = ch_obj_string(&row, "database").unwrap_or_default();
let table = ch_obj_string(&row, "table").unwrap_or_default();
let column = ch_obj_string(&row, "name").unwrap_or_default();
result.entry(db).or_default().entry(table).or_default().push(column);
}
return Ok(result);
}
let pool = state.get_pool(&connection_id).await?;
let sql = if flavor == DbFlavor::Greenplum {
@@ -454,6 +616,19 @@ pub async fn get_column_details(
table: String,
) -> TuskResult<Vec<ColumnDetail>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
let columns = get_table_columns_core(&state, &connection_id, &schema, &table).await?;
return Ok(columns
.into_iter()
.map(|c| ColumnDetail {
column_name: c.name,
data_type: c.data_type,
is_nullable: c.is_nullable,
column_default: c.column_default,
is_identity: false,
})
.collect());
}
let pool = state.get_pool(&connection_id).await?;
let sql = if flavor == DbFlavor::Greenplum {
@@ -500,6 +675,10 @@ pub async fn get_table_triggers(
schema: String,
table: String,
) -> TuskResult<Vec<TriggerInfo>> {
let flavor = state.get_flavor(&connection_id).await;
if matches!(flavor, DbFlavor::ClickHouse) {
return Ok(vec![]);
}
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
@@ -547,127 +726,3 @@ pub async fn get_table_triggers(
.collect())
}
#[tauri::command]
pub async fn get_schema_erd(
state: State<'_, Arc<AppState>>,
connection_id: String,
schema: String,
) -> TuskResult<ErdData> {
let pool = state.get_pool(&connection_id).await?;
// Get all tables with columns
let col_rows = sqlx::query(
"SELECT \
c.table_name, \
c.column_name, \
c.data_type, \
c.is_nullable = 'YES' AS is_nullable, \
COALESCE(( \
SELECT true FROM pg_constraint con \
JOIN pg_class cl ON cl.oid = con.conrelid \
JOIN pg_namespace ns ON ns.oid = cl.relnamespace \
WHERE con.contype = 'p' \
AND ns.nspname = $1 AND cl.relname = c.table_name \
AND EXISTS ( \
SELECT 1 FROM unnest(con.conkey) k \
JOIN pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = k \
WHERE a.attname = c.column_name \
) \
LIMIT 1 \
), false) AS is_pk \
FROM information_schema.columns c \
JOIN information_schema.tables t \
ON t.table_schema = c.table_schema AND t.table_name = c.table_name \
WHERE c.table_schema = $1 AND t.table_type = 'BASE TABLE' \
ORDER BY c.table_name, c.ordinal_position",
)
.bind(&schema)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
// Build tables map
let mut tables_map: HashMap<String, ErdTable> = HashMap::new();
for row in &col_rows {
let table_name: String = row.get(0);
let entry = tables_map
.entry(table_name.clone())
.or_insert_with(|| ErdTable {
schema: schema.clone(),
name: table_name,
columns: Vec::new(),
});
entry.columns.push(ErdColumn {
name: row.get(1),
data_type: row.get(2),
is_nullable: row.get(3),
is_primary_key: row.get(4),
});
}
let tables: Vec<ErdTable> = tables_map.into_values().collect();
// Get all FK relationships
let fk_rows = sqlx::query(
"SELECT \
c.conname AS constraint_name, \
src_ns.nspname AS source_schema, \
src_cl.relname AS source_table, \
ARRAY( \
SELECT a.attname FROM unnest(c.conkey) WITH ORDINALITY AS k(attnum, ord) \
JOIN pg_attribute a ON a.attrelid = c.conrelid AND a.attnum = k.attnum \
ORDER BY k.ord \
)::text[] AS source_columns, \
ref_ns.nspname AS target_schema, \
ref_cl.relname AS target_table, \
ARRAY( \
SELECT a.attname FROM unnest(c.confkey) WITH ORDINALITY AS k(attnum, ord) \
JOIN pg_attribute a ON a.attrelid = c.confrelid AND a.attnum = k.attnum \
ORDER BY k.ord \
)::text[] AS target_columns, \
CASE c.confupdtype \
WHEN 'a' THEN 'NO ACTION' \
WHEN 'r' THEN 'RESTRICT' \
WHEN 'c' THEN 'CASCADE' \
WHEN 'n' THEN 'SET NULL' \
WHEN 'd' THEN 'SET DEFAULT' \
END AS update_rule, \
CASE c.confdeltype \
WHEN 'a' THEN 'NO ACTION' \
WHEN 'r' THEN 'RESTRICT' \
WHEN 'c' THEN 'CASCADE' \
WHEN 'n' THEN 'SET NULL' \
WHEN 'd' THEN 'SET DEFAULT' \
END AS delete_rule \
FROM pg_constraint c \
JOIN pg_class src_cl ON src_cl.oid = c.conrelid \
JOIN pg_namespace src_ns ON src_ns.oid = src_cl.relnamespace \
JOIN pg_class ref_cl ON ref_cl.oid = c.confrelid \
JOIN pg_namespace ref_ns ON ref_ns.oid = ref_cl.relnamespace \
WHERE c.contype = 'f' AND src_ns.nspname = $1 \
ORDER BY c.conname",
)
.bind(&schema)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
let relationships: Vec<ErdRelationship> = fk_rows
.iter()
.map(|r| ErdRelationship {
constraint_name: r.get(0),
source_schema: r.get(1),
source_table: r.get(2),
source_columns: r.get(3),
target_schema: r.get(4),
target_table: r.get(5),
target_columns: r.get(6),
update_rule: r.get(7),
delete_rule: r.get(8),
})
.collect();
Ok(ErdData {
tables,
relationships,
})
}
+1 -10
View File
@@ -1,6 +1,6 @@
use crate::error::{TuskError, TuskResult};
use crate::mcp;
use crate::models::settings::{AppSettings, DockerHost, McpStatus};
use crate::models::settings::{AppSettings, McpStatus};
use crate::state::AppState;
use std::fs;
use std::sync::Arc;
@@ -36,15 +36,6 @@ pub async fn save_app_settings(
let data = serde_json::to_string_pretty(&settings)?;
fs::write(&path, data)?;
// Apply docker host setting
{
let mut docker_host = state.docker_host.write().await;
*docker_host = match settings.docker.host {
DockerHost::Remote => settings.docker.remote_url.clone(),
DockerHost::Local => None,
};
}
// Apply MCP setting: restart or stop
let is_running = *state.mcp_running.read().await;
-362
View File
@@ -1,362 +0,0 @@
use crate::commands::ai::fetch_foreign_keys_raw;
use crate::commands::data::bind_json_value;
use crate::commands::queries::pg_value_to_json;
use crate::error::{TuskError, TuskResult};
use crate::models::snapshot::{
CreateSnapshotParams, RestoreSnapshotParams, Snapshot, SnapshotMetadata, SnapshotProgress,
SnapshotTableData, SnapshotTableMeta,
};
use crate::state::AppState;
use crate::utils::{escape_ident, topological_sort_tables};
use serde_json::Value;
use sqlx::{Column, Row, TypeInfo};
use std::fs;
use std::sync::Arc;
use tauri::{AppHandle, Emitter, Manager, State};
#[tauri::command]
pub async fn create_snapshot(
app: AppHandle,
state: State<'_, Arc<AppState>>,
params: CreateSnapshotParams,
snapshot_id: String,
file_path: String,
) -> TuskResult<SnapshotMetadata> {
let pool = state.get_pool(&params.connection_id).await?;
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "preparing".to_string(),
percent: 5,
message: "Preparing snapshot...".to_string(),
detail: None,
},
);
let mut target_tables: Vec<(String, String)> = params
.tables
.iter()
.map(|t| (t.schema.clone(), t.table.clone()))
.collect();
// Fetch FK info once — used for both dependency expansion and topological sort
let fk_rows = fetch_foreign_keys_raw(&pool).await?;
if params.include_dependencies {
for fk in &fk_rows {
if target_tables
.iter()
.any(|(s, t)| s == &fk.schema && t == &fk.table)
{
let parent = (fk.ref_schema.clone(), fk.ref_table.clone());
if !target_tables.contains(&parent) {
target_tables.push(parent);
}
}
}
}
// FK-based topological sort
let fk_edges: Vec<(String, String, String, String)> = fk_rows
.iter()
.map(|fk| {
(
fk.schema.clone(),
fk.table.clone(),
fk.ref_schema.clone(),
fk.ref_table.clone(),
)
})
.collect();
let sorted_tables = topological_sort_tables(&fk_edges, &target_tables);
let mut tx = pool.begin().await.map_err(TuskError::Database)?;
sqlx::query("SET TRANSACTION READ ONLY")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
let total_tables = sorted_tables.len();
let mut snapshot_tables: Vec<SnapshotTableData> = Vec::new();
let mut table_metas: Vec<SnapshotTableMeta> = Vec::new();
let mut total_rows: u64 = 0;
for (i, (schema, table)) in sorted_tables.iter().enumerate() {
let percent = (10 + (i * 80 / total_tables.max(1))).min(90) as u8;
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "exporting".to_string(),
percent,
message: format!("Exporting {}.{}...", schema, table),
detail: None,
},
);
let qualified = format!("{}.{}", escape_ident(schema), escape_ident(table));
let sql = format!("SELECT * FROM {}", qualified);
let rows = sqlx::query(&sql)
.fetch_all(&mut *tx)
.await
.map_err(TuskError::Database)?;
let mut columns = Vec::new();
let mut column_types = Vec::new();
if let Some(first) = rows.first() {
for col in first.columns() {
columns.push(col.name().to_string());
column_types.push(col.type_info().name().to_string());
}
}
let data_rows: Vec<Vec<Value>> = rows
.iter()
.map(|row| {
(0..columns.len())
.map(|i| pg_value_to_json(row, i))
.collect()
})
.collect();
let row_count = data_rows.len() as u64;
total_rows += row_count;
table_metas.push(SnapshotTableMeta {
schema: schema.clone(),
table: table.clone(),
row_count,
columns: columns.clone(),
column_types: column_types.clone(),
});
snapshot_tables.push(SnapshotTableData {
schema: schema.clone(),
table: table.clone(),
columns,
column_types,
rows: data_rows,
});
}
tx.rollback().await.map_err(TuskError::Database)?;
let metadata = SnapshotMetadata {
id: snapshot_id.clone(),
name: params.name.clone(),
created_at: chrono::Utc::now().to_rfc3339(),
connection_name: String::new(),
database: String::new(),
tables: table_metas,
total_rows,
file_size_bytes: 0,
version: 1,
};
let snapshot = Snapshot {
metadata: metadata.clone(),
tables: snapshot_tables,
};
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "saving".to_string(),
percent: 95,
message: "Saving snapshot file...".to_string(),
detail: None,
},
);
let json = serde_json::to_string_pretty(&snapshot)?;
let file_size = json.len() as u64;
fs::write(&file_path, json)?;
let mut final_metadata = metadata;
final_metadata.file_size_bytes = file_size;
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "done".to_string(),
percent: 100,
message: "Snapshot created successfully".to_string(),
detail: Some(format!("{} rows, {} tables", total_rows, total_tables)),
},
);
Ok(final_metadata)
}
#[tauri::command]
pub async fn restore_snapshot(
app: AppHandle,
state: State<'_, Arc<AppState>>,
params: RestoreSnapshotParams,
snapshot_id: String,
) -> TuskResult<u64> {
if state.is_read_only(&params.connection_id).await {
return Err(TuskError::ReadOnly);
}
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "reading".to_string(),
percent: 5,
message: "Reading snapshot file...".to_string(),
detail: None,
},
);
let data = fs::read_to_string(&params.file_path)?;
let snapshot: Snapshot = serde_json::from_str(&data)?;
let pool = state.get_pool(&params.connection_id).await?;
let mut tx = pool.begin().await.map_err(TuskError::Database)?;
sqlx::query("SET CONSTRAINTS ALL DEFERRED")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
// TRUNCATE in reverse order (children first)
if params.truncate_before_restore {
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "truncating".to_string(),
percent: 15,
message: "Truncating existing data...".to_string(),
detail: None,
},
);
for table_data in snapshot.tables.iter().rev() {
let qualified = format!(
"{}.{}",
escape_ident(&table_data.schema),
escape_ident(&table_data.table)
);
let truncate_sql = format!("TRUNCATE {} CASCADE", qualified);
sqlx::query(&truncate_sql)
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
}
}
// INSERT in forward order (parents first)
let total_tables = snapshot.tables.len();
let mut total_inserted: u64 = 0;
for (i, table_data) in snapshot.tables.iter().enumerate() {
if table_data.columns.is_empty() || table_data.rows.is_empty() {
continue;
}
let percent = (20 + (i * 75 / total_tables.max(1))).min(95) as u8;
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "inserting".to_string(),
percent,
message: format!("Restoring {}.{}...", table_data.schema, table_data.table),
detail: Some(format!("{} rows", table_data.rows.len())),
},
);
let qualified = format!(
"{}.{}",
escape_ident(&table_data.schema),
escape_ident(&table_data.table)
);
let col_list: Vec<String> = table_data.columns.iter().map(|c| escape_ident(c)).collect();
let placeholders: Vec<String> = (1..=table_data.columns.len())
.map(|i| format!("${}", i))
.collect();
let sql = format!(
"INSERT INTO {} ({}) VALUES ({})",
qualified,
col_list.join(", "),
placeholders.join(", ")
);
// Chunked insert
for row in &table_data.rows {
let mut query = sqlx::query(&sql);
for val in row {
query = bind_json_value(query, val);
}
query.execute(&mut *tx).await.map_err(TuskError::Database)?;
total_inserted += 1;
}
}
tx.commit().await.map_err(TuskError::Database)?;
let _ = app.emit(
"snapshot-progress",
SnapshotProgress {
snapshot_id: snapshot_id.clone(),
stage: "done".to_string(),
percent: 100,
message: "Restore completed successfully".to_string(),
detail: Some(format!("{} rows restored", total_inserted)),
},
);
state.invalidate_schema_cache(&params.connection_id).await;
Ok(total_inserted)
}
#[tauri::command]
pub async fn list_snapshots(app: AppHandle) -> TuskResult<Vec<SnapshotMetadata>> {
let dir = app
.path()
.app_data_dir()
.map_err(|e| TuskError::Config(e.to_string()))?
.join("snapshots");
if !dir.exists() {
return Ok(Vec::new());
}
let mut snapshots = Vec::new();
for entry in fs::read_dir(&dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().map(|e| e == "json").unwrap_or(false) {
if let Ok(data) = fs::read_to_string(&path) {
if let Ok(snapshot) = serde_json::from_str::<Snapshot>(&data) {
let mut meta = snapshot.metadata;
meta.file_size_bytes = entry.metadata().map(|m| m.len()).unwrap_or(0);
snapshots.push(meta);
}
}
}
}
snapshots.sort_by(|a, b| b.created_at.cmp(&a.created_at));
Ok(snapshots)
}
#[tauri::command]
pub async fn read_snapshot_metadata(file_path: String) -> TuskResult<SnapshotMetadata> {
let data = fs::read_to_string(&file_path)?;
let snapshot: Snapshot = serde_json::from_str(&data)?;
let mut meta = snapshot.metadata;
meta.file_size_bytes = fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0);
Ok(meta)
}
+168
View File
@@ -0,0 +1,168 @@
use crate::error::{TuskError, TuskResult};
use crate::models::query_result::QueryResult;
use serde::Deserialize;
use serde_json::{Map, Value};
use std::sync::LazyLock;
use std::time::{Duration, Instant};
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(120);
fn http_client() -> &'static reqwest::Client {
static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(5))
.timeout(DEFAULT_TIMEOUT)
.build()
.unwrap_or_default()
});
&CLIENT
}
#[derive(Debug, Clone)]
pub struct ChClient {
pub base_url: String,
pub user: String,
pub password: String,
pub database: String,
}
impl ChClient {
pub fn new(host: &str, port: u16, secure: bool, user: &str, password: &str, database: &str) -> Self {
let scheme = if secure { "https" } else { "http" };
let base_url = format!("{}://{}:{}", scheme, host, port);
Self {
base_url,
user: user.to_string(),
password: password.to_string(),
database: database.to_string(),
}
}
fn endpoint(&self, database: Option<&str>, format: Option<&str>, read_only: bool) -> String {
let db = database.unwrap_or(&self.database);
let mut params = vec![
format!("database={}", urlencode(db)),
format!("user={}", urlencode(&self.user)),
];
if !self.password.is_empty() {
params.push(format!("password={}", urlencode(&self.password)));
}
if let Some(fmt) = format {
params.push(format!("default_format={}", urlencode(fmt)));
}
if read_only {
params.push("readonly=1".to_string());
}
format!("{}/?{}", self.base_url, params.join("&"))
}
/// Execute SQL and return raw response body.
pub async fn execute_raw(&self, sql: &str, format: Option<&str>, read_only: bool) -> TuskResult<String> {
let url = self.endpoint(None, format, read_only);
let resp = http_client()
.post(&url)
.body(sql.to_string())
.send()
.await
.map_err(|e| TuskError::Custom(format!("ClickHouse request failed: {}", e)))?;
let status = resp.status();
let body = resp
.text()
.await
.map_err(|e| TuskError::Custom(format!("Failed to read ClickHouse response: {}", e)))?;
if !status.is_success() {
return Err(TuskError::Custom(format!(
"ClickHouse error ({}): {}",
status,
body.trim()
)));
}
Ok(body)
}
/// Test connection by running `SELECT 1` and return the server version.
pub async fn ping(&self) -> TuskResult<String> {
// Use raw FORMAT TabSeparated to fetch version
let body = self.execute_raw("SELECT version()", Some("TabSeparated"), false).await?;
Ok(body.trim().to_string())
}
/// Execute SQL and parse rows via JSONCompact to preserve column metadata + types.
pub async fn execute_query(&self, sql: &str, read_only: bool) -> TuskResult<QueryResult> {
let start = Instant::now();
let body = self.execute_raw(sql, Some("JSONCompact"), read_only).await?;
let execution_time_ms = start.elapsed().as_millis() as u64;
// Empty body for statements without result set (DDL etc.) — return zero rows
if body.trim().is_empty() {
return Ok(QueryResult {
columns: vec![],
types: vec![],
rows: vec![],
row_count: 0,
execution_time_ms,
});
}
let parsed: ChJsonCompactResponse = serde_json::from_str(&body).map_err(|e| {
TuskError::Custom(format!(
"Failed to parse ClickHouse JSONCompact response: {} (body head: {})",
e,
body.chars().take(200).collect::<String>()
))
})?;
let columns: Vec<String> = parsed.meta.iter().map(|m| m.name.clone()).collect();
let types: Vec<String> = parsed.meta.iter().map(|m| m.r#type.clone()).collect();
let row_count = parsed.data.len();
Ok(QueryResult {
columns,
types,
rows: parsed.data,
row_count,
execution_time_ms,
})
}
/// Execute SQL expecting result rows as objects (for schema introspection helpers).
pub async fn fetch_objects(&self, sql: &str) -> TuskResult<Vec<Map<String, Value>>> {
let body = self.execute_raw(sql, Some("JSONEachRow"), false).await?;
let mut out = Vec::new();
for line in body.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
let value: Value = serde_json::from_str(line).map_err(|e| {
TuskError::Custom(format!("Failed to parse JSONEachRow line: {}", e))
})?;
if let Value::Object(obj) = value {
out.push(obj);
}
}
Ok(out)
}
}
#[derive(Debug, Deserialize)]
struct ChJsonCompactResponse {
meta: Vec<ChMetaEntry>,
data: Vec<Vec<Value>>,
}
#[derive(Debug, Deserialize)]
struct ChMetaEntry {
name: String,
r#type: String,
}
fn urlencode(s: &str) -> String {
s.chars()
.map(|c| match c {
':' | '/' | '?' | '#' | '[' | ']' | '@' | '!' | '$' | '&' | '\'' | '(' | ')' | '*'
| '+' | ',' | ';' | '=' | '%' | ' ' => format!("%{:02X}", c as u8),
_ => c.to_string(),
})
.collect()
}
+2
View File
@@ -0,0 +1,2 @@
pub mod clickhouse;
pub mod sql_guard;
+140
View File
@@ -0,0 +1,140 @@
use crate::error::{TuskError, TuskResult};
/// Cross-flavor whitelist guard for read-only SQL execution.
/// Allows: SELECT, WITH ... SELECT, SHOW, EXPLAIN, DESCRIBE.
/// Rejects: INSERT, UPDATE, DELETE, ALTER, DROP, CREATE, TRUNCATE,
/// RENAME, GRANT, REVOKE, ATTACH, DETACH, OPTIMIZE, SYSTEM.
pub fn ensure_readonly_sql(sql: &str) -> TuskResult<()> {
let normalized = strip_leading_comments(sql).to_ascii_uppercase();
let trimmed = normalized.trim();
if trimmed.is_empty() {
return Err(TuskError::Validation("Empty SQL statement".into()));
}
let allowed_starts = ["SELECT", "WITH", "SHOW", "EXPLAIN", "DESCRIBE", "DESC ", "DESC\n", "VALUES"];
let starts_ok = allowed_starts
.iter()
.any(|p| trimmed.starts_with(p) || trimmed == p.trim());
if !starts_ok {
return Err(TuskError::ReadOnly);
}
// Reject if any forbidden keyword appears as a top-level token
let forbidden = [
"INSERT", "UPDATE", "DELETE", "ALTER", "DROP", "CREATE", "TRUNCATE",
"RENAME", "GRANT", "REVOKE", "ATTACH", "DETACH", "OPTIMIZE", "SYSTEM",
"REPLACE", "MERGE",
];
for kw in forbidden {
if contains_keyword(&normalized, kw) {
return Err(TuskError::ReadOnly);
}
}
Ok(())
}
fn strip_leading_comments(sql: &str) -> &str {
let mut s = sql.trim_start();
loop {
if let Some(rest) = s.strip_prefix("--") {
// line comment — skip to newline
match rest.find('\n') {
Some(idx) => s = rest[idx + 1..].trim_start(),
None => return "",
}
} else if let Some(rest) = s.strip_prefix("/*") {
match rest.find("*/") {
Some(idx) => s = rest[idx + 2..].trim_start(),
None => return "",
}
} else {
break;
}
}
s
}
fn contains_keyword(haystack: &str, kw: &str) -> bool {
let bytes = haystack.as_bytes();
let needle = kw.as_bytes();
let mut i = 0;
while i + needle.len() <= bytes.len() {
if &bytes[i..i + needle.len()] == needle {
let before_ok = i == 0 || !is_word_char(bytes[i - 1]);
let after = i + needle.len();
let after_ok = after == bytes.len() || !is_word_char(bytes[after]);
if before_ok && after_ok {
return true;
}
}
i += 1;
}
false
}
fn is_word_char(b: u8) -> bool {
b.is_ascii_alphanumeric() || b == b'_'
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn allows_select() {
assert!(ensure_readonly_sql("SELECT 1").is_ok());
assert!(ensure_readonly_sql(" SELECT * FROM t").is_ok());
assert!(ensure_readonly_sql("select 1").is_ok());
}
#[test]
fn allows_with_select() {
assert!(ensure_readonly_sql("WITH x AS (SELECT 1) SELECT * FROM x").is_ok());
}
#[test]
fn allows_explain_show() {
assert!(ensure_readonly_sql("EXPLAIN SELECT 1").is_ok());
assert!(ensure_readonly_sql("SHOW TABLES").is_ok());
assert!(ensure_readonly_sql("DESCRIBE t").is_ok());
}
#[test]
fn rejects_dml_ddl() {
assert!(ensure_readonly_sql("INSERT INTO t VALUES (1)").is_err());
assert!(ensure_readonly_sql("UPDATE t SET a=1").is_err());
assert!(ensure_readonly_sql("DELETE FROM t").is_err());
assert!(ensure_readonly_sql("DROP TABLE t").is_err());
assert!(ensure_readonly_sql("CREATE TABLE t(a int)").is_err());
assert!(ensure_readonly_sql("TRUNCATE TABLE t").is_err());
}
#[test]
fn rejects_writable_cte() {
// PG writable CTE — looks like WITH but contains INSERT
assert!(
ensure_readonly_sql("WITH x AS (INSERT INTO t VALUES (1) RETURNING *) SELECT * FROM x")
.is_err()
);
}
#[test]
fn rejects_select_with_drop_chain() {
assert!(ensure_readonly_sql("SELECT 1; DROP TABLE t").is_err());
}
#[test]
fn allows_select_with_keyword_in_string() {
// Real-world: column names containing forbidden keywords should pass; string literals
// containing them should also pass. Our guard is conservative and may reject some
// legitimate queries — that is acceptable for the read-only safety net.
// This test documents the limitation: queries embedding "DROP" as a literal will be rejected.
// The user can disable read-only mode to run them.
}
#[test]
fn strips_leading_comments() {
assert!(ensure_readonly_sql("-- comment\nSELECT 1").is_ok());
assert!(ensure_readonly_sql("/* block */ SELECT 1").is_ok());
assert!(ensure_readonly_sql("/* multi\nline */\n SELECT 1").is_ok());
}
}
-6
View File
@@ -11,9 +11,6 @@ pub enum TuskError {
#[error("Serialization error: {0}")]
Serde(#[from] serde_json::Error),
#[error("Connection not found: {0}")]
ConnectionNotFound(String),
#[error("Not connected: {0}")]
NotConnected(String),
@@ -23,9 +20,6 @@ pub enum TuskError {
#[error("AI error: {0}")]
Ai(String),
#[error("Docker error: {0}")]
Docker(String),
#[error("Configuration error: {0}")]
Config(String),
+7 -49
View File
@@ -1,11 +1,12 @@
mod commands;
mod db;
mod error;
mod mcp;
mod models;
mod state;
mod utils;
use models::settings::{AppSettings, DockerHost};
use models::settings::AppSettings;
use state::AppState;
use std::sync::Arc;
use tauri::Manager;
@@ -37,21 +38,9 @@ pub fn run() {
AppSettings::default()
};
// Apply docker host from settings
let docker_host = match settings.docker.host {
DockerHost::Remote => settings.docker.remote_url.clone(),
DockerHost::Local => None,
};
let mcp_enabled = settings.mcp.enabled;
let mcp_port = settings.mcp.port;
// Set docker host synchronously (state is fresh, no contention)
let state_for_setup = state.clone();
tauri::async_runtime::block_on(async {
*state_for_setup.docker_host.write().await = docker_host;
});
if mcp_enabled {
let shutdown_rx = state.mcp_shutdown_tx.subscribe();
let mcp_state = state.clone();
@@ -101,7 +90,6 @@ pub fn run() {
commands::schema::get_completion_schema,
commands::schema::get_column_details,
commands::schema::get_table_triggers,
commands::schema::get_schema_erd,
// data
commands::data::get_table_data,
commands::data::update_row,
@@ -110,20 +98,6 @@ pub fn run() {
// export
commands::export::export_csv,
commands::export::export_json,
// management
commands::management::get_database_info,
commands::management::create_database,
commands::management::drop_database,
commands::management::list_roles,
commands::management::create_role,
commands::management::alter_role,
commands::management::drop_role,
commands::management::get_table_privileges,
commands::management::grant_revoke,
commands::management::manage_role_membership,
commands::management::list_sessions,
commands::management::cancel_query,
commands::management::terminate_backend,
// history
commands::history::add_history_entry,
commands::history::get_history,
@@ -139,27 +113,11 @@ pub fn run() {
commands::ai::generate_sql,
commands::ai::explain_sql,
commands::ai::fix_sql_error,
commands::ai::generate_validation_sql,
commands::ai::run_validation_rule,
commands::ai::suggest_validation_rules,
commands::ai::generate_test_data_preview,
commands::ai::insert_generated_data,
commands::ai::get_index_advisor_report,
commands::ai::apply_index_recommendation,
// snapshot
commands::snapshot::create_snapshot,
commands::snapshot::restore_snapshot,
commands::snapshot::list_snapshots,
commands::snapshot::read_snapshot_metadata,
// lookup
commands::lookup::entity_lookup,
// docker
commands::docker::check_docker,
commands::docker::list_tusk_containers,
commands::docker::clone_to_docker,
commands::docker::start_container,
commands::docker::stop_container,
commands::docker::remove_container,
// chat
commands::chat::chat_send,
// memory
commands::memory::get_memory,
commands::memory::save_memory,
// settings
commands::settings::get_app_settings,
commands::settings::save_app_settings,
+2 -127
View File
@@ -41,6 +41,8 @@ pub struct OllamaChatRequest {
pub model: String,
pub messages: Vec<OllamaChatMessage>,
pub stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub format: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -57,130 +59,3 @@ pub struct OllamaTagsResponse {
pub struct OllamaModel {
pub name: String,
}
// --- Wave 1: Validation ---
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ValidationStatus {
Pending,
Generating,
Running,
Passed,
Failed,
Error,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationRule {
pub id: String,
pub description: String,
pub generated_sql: String,
pub status: ValidationStatus,
pub violation_count: u64,
pub sample_violations: Vec<Vec<serde_json::Value>>,
pub violation_columns: Vec<String>,
pub error: Option<String>,
}
// --- Wave 2: Data Generator ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateDataParams {
pub connection_id: String,
pub schema: String,
pub table: String,
pub row_count: u32,
pub include_related: bool,
pub custom_instructions: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedDataPreview {
pub tables: Vec<GeneratedTableData>,
pub insert_order: Vec<String>,
pub total_rows: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedTableData {
pub schema: String,
pub table: String,
pub columns: Vec<String>,
pub rows: Vec<Vec<serde_json::Value>>,
pub row_count: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataGenProgress {
pub gen_id: String,
pub stage: String,
pub percent: u8,
pub message: String,
pub detail: Option<String>,
}
// --- Wave 3A: Index Advisor ---
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableStats {
pub schema: String,
pub table: String,
pub seq_scan: i64,
pub idx_scan: i64,
pub n_live_tup: i64,
pub table_size: String,
pub index_size: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexStats {
pub schema: String,
pub table: String,
pub index_name: String,
pub idx_scan: i64,
pub index_size: String,
pub definition: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlowQuery {
pub query: String,
pub calls: i64,
pub total_time_ms: f64,
pub mean_time_ms: f64,
pub rows: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum IndexRecommendationType {
#[serde(rename = "create_index")]
Create,
#[serde(rename = "drop_index")]
Drop,
#[serde(rename = "replace_index")]
Replace,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexRecommendation {
pub id: String,
pub recommendation_type: IndexRecommendationType,
pub table_schema: String,
pub table_name: String,
pub index_name: Option<String>,
pub ddl: String,
pub rationale: String,
pub estimated_impact: String,
pub priority: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexAdvisorReport {
pub table_stats: Vec<TableStats>,
pub index_stats: Vec<IndexStats>,
pub slow_queries: Vec<SlowQuery>,
pub recommendations: Vec<IndexRecommendation>,
pub has_pg_stat_statements: bool,
}
+19
View File
@@ -0,0 +1,19 @@
use crate::models::query_result::QueryResult;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum ChatMessage {
User { id: String, text: String, created_at: i64 },
Assistant { id: String, text: String, created_at: i64 },
ToolCall { id: String, tool: String, input_json: String, created_at: i64 },
ToolResult {
id: String,
tool: String,
is_error: bool,
text: Option<String>,
result: Option<QueryResult>,
created_at: i64,
},
}
+12
View File
@@ -1,3 +1,4 @@
use crate::state::DbFlavor;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -12,6 +13,17 @@ pub struct ConnectionConfig {
pub ssl_mode: Option<String>,
pub color: Option<String>,
pub environment: Option<String>,
/// Database flavor selected by the user. Defaults to PostgreSQL for backwards
/// compatibility with older `connections.json` files written before multi-DB support.
#[serde(default = "default_flavor")]
pub db_flavor: DbFlavor,
/// HTTPS for ClickHouse. Defaults to false.
#[serde(default)]
pub secure: bool,
}
fn default_flavor() -> DbFlavor {
DbFlavor::PostgreSQL
}
impl ConnectionConfig {
-57
View File
@@ -1,57 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DockerStatus {
pub installed: bool,
pub daemon_running: bool,
pub version: Option<String>,
pub error: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CloneMode {
SchemaOnly,
FullClone,
SampleData,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CloneToDockerParams {
pub source_connection_id: String,
pub source_database: String,
pub container_name: String,
pub pg_version: String,
pub host_port: Option<u16>,
pub clone_mode: CloneMode,
pub sample_rows: Option<u32>,
pub postgres_password: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct CloneProgress {
pub clone_id: String,
pub stage: String,
pub percent: u8,
pub message: String,
pub detail: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TuskContainer {
pub container_id: String,
pub name: String,
pub status: String,
pub host_port: u16,
pub pg_version: String,
pub source_database: Option<String>,
pub source_connection: Option<String>,
pub created_at: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct CloneResult {
pub container: TuskContainer,
pub connection_id: String,
pub connection_url: String,
}
-44
View File
@@ -1,44 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LookupTableMatch {
pub schema: String,
pub table: String,
pub column_type: String,
pub columns: Vec<String>,
pub types: Vec<String>,
pub rows: Vec<Vec<serde_json::Value>>,
pub row_count: usize,
pub total_count: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LookupDatabaseResult {
pub database: String,
pub tables: Vec<LookupTableMatch>,
pub error: Option<String>,
pub search_time_ms: u128,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntityLookupResult {
pub column_name: String,
pub value: String,
pub databases: Vec<LookupDatabaseResult>,
pub total_databases_searched: usize,
pub total_tables_matched: usize,
pub total_rows_found: usize,
pub total_time_ms: u128,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LookupProgress {
pub lookup_id: String,
pub database: String,
pub status: String,
pub tables_found: usize,
pub rows_found: usize,
pub error: Option<String>,
pub completed: usize,
pub total: usize,
}
-110
View File
@@ -1,110 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub struct DatabaseInfo {
pub name: String,
pub owner: String,
pub encoding: String,
pub collation: String,
pub ctype: String,
pub tablespace: String,
pub connection_limit: i32,
pub size: String,
pub description: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct CreateDatabaseParams {
pub name: String,
pub owner: Option<String>,
pub template: Option<String>,
pub encoding: Option<String>,
pub tablespace: Option<String>,
pub connection_limit: Option<i32>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct RoleInfo {
pub name: String,
pub is_superuser: bool,
pub can_login: bool,
pub can_create_db: bool,
pub can_create_role: bool,
pub inherit: bool,
pub is_replication: bool,
pub connection_limit: i32,
pub password_set: bool,
pub valid_until: Option<String>,
pub member_of: Vec<String>,
pub members: Vec<String>,
pub description: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct CreateRoleParams {
pub name: String,
pub password: Option<String>,
pub login: bool,
pub superuser: bool,
pub createdb: bool,
pub createrole: bool,
pub inherit: bool,
pub replication: bool,
pub connection_limit: Option<i32>,
pub valid_until: Option<String>,
pub in_roles: Vec<String>,
}
#[derive(Debug, Deserialize)]
pub struct AlterRoleParams {
pub name: String,
pub password: Option<String>,
pub login: Option<bool>,
pub superuser: Option<bool>,
pub createdb: Option<bool>,
pub createrole: Option<bool>,
pub inherit: Option<bool>,
pub replication: Option<bool>,
pub connection_limit: Option<i32>,
pub valid_until: Option<String>,
pub rename_to: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TablePrivilege {
pub grantee: String,
pub table_schema: String,
pub table_name: String,
pub privilege_type: String,
pub is_grantable: bool,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SessionInfo {
pub pid: i32,
pub usename: Option<String>,
pub datname: Option<String>,
pub state: Option<String>,
pub query: Option<String>,
pub query_start: Option<String>,
pub wait_event_type: Option<String>,
pub wait_event: Option<String>,
pub client_addr: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct GrantRevokeParams {
pub action: String,
pub privileges: Vec<String>,
pub object_type: String,
pub object_name: String,
pub role_name: String,
pub with_grant_option: bool,
}
#[derive(Debug, Deserialize)]
pub struct RoleMembershipParams {
pub action: String,
pub role_name: String,
pub member_name: String,
}
+1 -4
View File
@@ -1,11 +1,8 @@
pub mod ai;
pub mod chat;
pub mod connection;
pub mod docker;
pub mod history;
pub mod lookup;
pub mod management;
pub mod query_result;
pub mod saved_queries;
pub mod schema;
pub mod settings;
pub mod snapshot;
+4 -2
View File
@@ -1,13 +1,15 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
// Tauri's IPC layer does not support u128/i128 in command arguments,
// so timings round-trip through frontend → backend as u64 ms.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
pub columns: Vec<String>,
pub types: Vec<String>,
pub rows: Vec<Vec<Value>>,
pub row_count: usize,
pub execution_time_ms: u128,
pub execution_time_ms: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -16,7 +18,7 @@ pub struct PaginatedQueryResult {
pub types: Vec<String>,
pub rows: Vec<Vec<Value>>,
pub row_count: usize,
pub execution_time_ms: u128,
pub execution_time_ms: u64,
pub total_rows: i64,
pub page: u32,
pub page_size: u32,
-34
View File
@@ -60,37 +60,3 @@ pub struct TriggerInfo {
pub is_enabled: bool,
pub definition: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErdColumn {
pub name: String,
pub data_type: String,
pub is_nullable: bool,
pub is_primary_key: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErdTable {
pub schema: String,
pub name: String,
pub columns: Vec<ErdColumn>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErdRelationship {
pub constraint_name: String,
pub source_schema: String,
pub source_table: String,
pub source_columns: Vec<String>,
pub target_schema: String,
pub target_table: String,
pub target_columns: Vec<String>,
pub update_rule: String,
pub delete_rule: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErdData {
pub tables: Vec<ErdTable>,
pub relationships: Vec<ErdRelationship>,
}
-23
View File
@@ -3,7 +3,6 @@ use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AppSettings {
pub mcp: McpSettings,
pub docker: DockerSettings,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -21,28 +20,6 @@ impl Default for McpSettings {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DockerSettings {
pub host: DockerHost,
pub remote_url: Option<String>,
}
impl Default for DockerSettings {
fn default() -> Self {
Self {
host: DockerHost::Local,
remote_url: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum DockerHost {
Local,
Remote,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpStatus {
pub enabled: bool,
-68
View File
@@ -1,68 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotMetadata {
pub id: String,
pub name: String,
pub created_at: String,
pub connection_name: String,
pub database: String,
pub tables: Vec<SnapshotTableMeta>,
pub total_rows: u64,
pub file_size_bytes: u64,
pub version: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotTableMeta {
pub schema: String,
pub table: String,
pub row_count: u64,
pub columns: Vec<String>,
pub column_types: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Snapshot {
pub metadata: SnapshotMetadata,
pub tables: Vec<SnapshotTableData>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotTableData {
pub schema: String,
pub table: String,
pub columns: Vec<String>,
pub column_types: Vec<String>,
pub rows: Vec<Vec<serde_json::Value>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotProgress {
pub snapshot_id: String,
pub stage: String,
pub percent: u8,
pub message: String,
pub detail: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateSnapshotParams {
pub connection_id: String,
pub tables: Vec<TableRef>,
pub name: String,
pub include_dependencies: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableRef {
pub schema: String,
pub table: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RestoreSnapshotParams {
pub connection_id: String,
pub file_path: String,
pub truncate_before_restore: bool,
}
+52 -6
View File
@@ -1,8 +1,10 @@
use crate::db::clickhouse::ChClient;
use crate::error::{TuskError, TuskResult};
use crate::models::ai::AiSettings;
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{watch, RwLock};
@@ -11,22 +13,44 @@ use tokio::sync::{watch, RwLock};
pub enum DbFlavor {
PostgreSQL,
Greenplum,
ClickHouse,
}
#[derive(Clone)]
pub struct SchemaCacheEntry {
pub schema_text: String,
pub cached_at: Instant,
}
#[derive(Clone)]
pub struct CachedString {
pub value: String,
pub cached_at: Instant,
}
#[derive(Clone)]
pub struct CachedVec<T: Clone> {
pub value: Vec<T>,
pub cached_at: Instant,
}
pub struct AppState {
pub pools: RwLock<HashMap<String, PgPool>>,
pub ch_clients: RwLock<HashMap<String, Arc<ChClient>>>,
pub read_only: RwLock<HashMap<String, bool>>,
pub db_flavors: RwLock<HashMap<String, DbFlavor>>,
/// Legacy cache used by generate_sql/explain_sql/fix_sql_error — full DDL.
pub schema_cache: RwLock<HashMap<String, SchemaCacheEntry>>,
/// Chat v2 caches: lite overview per connection.
pub overview_cache: RwLock<HashMap<String, CachedString>>,
/// Chat v2 caches: list of tables per (connection_id, db_name) — used for
/// list_tables on a non-active PG database via temporary pool.
pub tables_by_db_cache: RwLock<HashMap<(String, String), CachedVec<String>>>,
/// Chat v2 caches: column block per (connection_id, db_name, "schema.table").
pub columns_cache: RwLock<HashMap<(String, String, String), CachedString>>,
pub mcp_shutdown_tx: watch::Sender<bool>,
pub mcp_running: RwLock<bool>,
pub docker_host: RwLock<Option<String>>,
pub ai_settings: RwLock<Option<AiSettings>>,
}
@@ -38,16 +62,34 @@ impl AppState {
let (mcp_shutdown_tx, _) = watch::channel(false);
Self {
pools: RwLock::new(HashMap::new()),
ch_clients: RwLock::new(HashMap::new()),
read_only: RwLock::new(HashMap::new()),
db_flavors: RwLock::new(HashMap::new()),
schema_cache: RwLock::new(HashMap::new()),
overview_cache: RwLock::new(HashMap::new()),
tables_by_db_cache: RwLock::new(HashMap::new()),
columns_cache: RwLock::new(HashMap::new()),
mcp_shutdown_tx,
mcp_running: RwLock::new(false),
docker_host: RwLock::new(None),
ai_settings: RwLock::new(None),
}
}
/// Drop every chat-agent cache entry tied to this connection.
/// Called by switch_database_core, disconnect, and on connection delete.
pub async fn invalidate_chat_caches_for(&self, connection_id: &str) {
self.schema_cache.write().await.remove(connection_id);
self.overview_cache.write().await.remove(connection_id);
self.tables_by_db_cache
.write()
.await
.retain(|(cid, _), _| cid != connection_id);
self.columns_cache
.write()
.await
.retain(|(cid, _, _), _| cid != connection_id);
}
pub async fn get_pool(&self, connection_id: &str) -> TuskResult<PgPool> {
let pools = self.pools.read().await;
pools
@@ -56,6 +98,14 @@ impl AppState {
.ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))
}
pub async fn get_ch_client(&self, connection_id: &str) -> TuskResult<Arc<ChClient>> {
let clients = self.ch_clients.read().await;
clients
.get(connection_id)
.cloned()
.ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))
}
pub async fn is_read_only(&self, id: &str) -> bool {
let map = self.read_only.read().await;
map.get(id).copied().unwrap_or(true)
@@ -100,8 +150,4 @@ impl AppState {
);
}
pub async fn invalidate_schema_cache(&self, connection_id: &str) {
let mut cache = self.schema_cache.write().await;
cache.remove(connection_id);
}
}
-150
View File
@@ -1,95 +1,11 @@
use std::collections::{HashMap, HashSet, VecDeque};
pub fn escape_ident(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
/// Topological sort of tables based on foreign key dependencies.
/// Returns tables in insertion order: parents before children.
pub fn topological_sort_tables(
fk_edges: &[(String, String, String, String)], // (schema, table, ref_schema, ref_table)
target_tables: &[(String, String)],
) -> Vec<(String, String)> {
let mut graph: HashMap<(String, String), HashSet<(String, String)>> = HashMap::new();
let mut in_degree: HashMap<(String, String), usize> = HashMap::new();
// Initialize all target tables
for t in target_tables {
graph.entry(t.clone()).or_default();
in_degree.entry(t.clone()).or_insert(0);
}
let target_set: HashSet<(String, String)> = target_tables.iter().cloned().collect();
// Build edges: parent -> child (child depends on parent)
for (schema, table, ref_schema, ref_table) in fk_edges {
let child = (schema.clone(), table.clone());
let parent = (ref_schema.clone(), ref_table.clone());
if child == parent {
continue; // self-referencing
}
if !target_set.contains(&child) || !target_set.contains(&parent) {
continue;
}
if graph
.entry(parent.clone())
.or_default()
.insert(child.clone())
{
*in_degree.entry(child).or_insert(0) += 1;
}
}
// Kahn's algorithm
let mut initial: Vec<(String, String)> = in_degree
.iter()
.filter(|(_, &deg)| deg == 0)
.map(|(k, _)| k.clone())
.collect();
initial.sort(); // deterministic order
let mut queue: VecDeque<(String, String)> = VecDeque::from(initial);
let mut result = Vec::new();
while let Some(node) = queue.pop_front() {
result.push(node.clone());
if let Some(neighbors) = graph.get(&node) {
let mut new_ready: Vec<(String, String)> = neighbors
.iter()
.filter(|neighbor| {
if let Some(deg) = in_degree.get_mut(*neighbor) {
*deg -= 1;
*deg == 0
} else {
false
}
})
.cloned()
.collect();
new_ready.sort();
queue.extend(new_ready);
}
}
// Add any remaining tables (cycles) at the end
for t in target_tables {
if !result.contains(t) {
result.push(t.clone());
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
// ── escape_ident ──────────────────────────────────────────
#[test]
fn escape_ident_simple_name() {
assert_eq!(escape_ident("users"), "\"users\"");
@@ -149,70 +65,4 @@ mod tests {
fn escape_ident_newline() {
assert_eq!(escape_ident("a\nb"), "\"a\nb\"");
}
// ── topological_sort_tables ───────────────────────────────
#[test]
fn topo_sort_no_edges() {
let tables = vec![("public".into(), "b".into()), ("public".into(), "a".into())];
let result = topological_sort_tables(&[], &tables);
assert_eq!(result.len(), 2);
assert!(result.contains(&("public".into(), "a".into())));
assert!(result.contains(&("public".into(), "b".into())));
}
#[test]
fn topo_sort_simple_dependency() {
let edges = vec![(
"public".into(),
"orders".into(),
"public".into(),
"users".into(),
)];
let tables = vec![
("public".into(), "orders".into()),
("public".into(), "users".into()),
];
let result = topological_sort_tables(&edges, &tables);
let user_pos = result.iter().position(|t| t.1 == "users").unwrap();
let order_pos = result.iter().position(|t| t.1 == "orders").unwrap();
assert!(user_pos < order_pos, "users must come before orders");
}
#[test]
fn topo_sort_self_reference() {
let edges = vec![(
"public".into(),
"employees".into(),
"public".into(),
"employees".into(),
)];
let tables = vec![("public".into(), "employees".into())];
let result = topological_sort_tables(&edges, &tables);
assert_eq!(result.len(), 1);
}
#[test]
fn topo_sort_cycle() {
let edges = vec![
("public".into(), "a".into(), "public".into(), "b".into()),
("public".into(), "b".into(), "public".into(), "a".into()),
];
let tables = vec![("public".into(), "a".into()), ("public".into(), "b".into())];
let result = topological_sort_tables(&edges, &tables);
assert_eq!(result.len(), 2);
}
#[test]
fn topo_sort_edge_outside_target_set_ignored() {
let edges = vec![(
"public".into(),
"orders".into(),
"public".into(),
"external".into(),
)];
let tables = vec![("public".into(), "orders".into())];
let result = topological_sort_tables(&edges, &tables);
assert_eq!(result.len(), 1);
}
}
+64
View File
@@ -0,0 +1,64 @@
import { useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import { Send } from "lucide-react";
interface Props {
onSend: (text: string) => void;
disabled?: boolean;
placeholder?: string;
}
export function ChatComposer({ onSend, disabled, placeholder }: Props) {
const [value, setValue] = useState("");
const ref = useRef<HTMLTextAreaElement>(null);
const handleSend = () => {
if (disabled) return;
const text = value.trim();
if (!text) return;
onSend(text);
setValue("");
requestAnimationFrame(() => {
if (ref.current) ref.current.style.height = "auto";
});
};
const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
handleSend();
}
};
const autoresize = (el: HTMLTextAreaElement) => {
el.style.height = "auto";
el.style.height = `${Math.min(el.scrollHeight, 200)}px`;
};
return (
<div className="flex items-end gap-2">
<textarea
ref={ref}
className="flex-1 resize-none rounded-md border border-border/50 bg-background px-3 py-2 text-sm outline-none placeholder:text-muted-foreground/50 focus:border-primary/40 focus:ring-1 focus:ring-primary/20 disabled:opacity-60"
rows={1}
value={value}
placeholder={placeholder}
disabled={disabled}
onChange={(e) => {
setValue(e.target.value);
autoresize(e.target);
}}
onKeyDown={handleKeyDown}
/>
<Button
size="sm"
onClick={handleSend}
disabled={disabled || !value.trim()}
className="h-8 gap-1.5"
>
<Send className="h-3.5 w-3.5" />
Send
</Button>
</div>
);
}
+385
View File
@@ -0,0 +1,385 @@
import { useState } from "react";
import { ResultsTable } from "@/components/results/ResultsTable";
import { ExportDialog } from "@/components/export/ExportDialog";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import {
ChevronDown,
ChevronRight,
AlertCircle,
Sparkles,
User,
Wrench,
Database,
Columns,
Layers,
RefreshCw,
StickyNote,
Bookmark,
BookmarkPlus,
Maximize2,
Download,
} from "lucide-react";
import type { ChatMessage } from "@/types";
interface Props {
message: ChatMessage;
}
export function ChatMessageView({ message }: Props) {
switch (message.role) {
case "user":
return <UserBubble text={message.text} />;
case "assistant":
return <AssistantBubble text={message.text} />;
case "tool_call":
return <ToolCallBlock tool={message.tool} inputJson={message.input_json} />;
case "tool_result":
return (
<ToolResultBlock
tool={message.tool}
isError={message.is_error}
text={message.text ?? null}
result={message.result ?? null}
/>
);
}
}
function UserBubble({ text }: { text: string }) {
return (
<div className="flex gap-2">
<span className="mt-0.5 flex h-6 w-6 shrink-0 items-center justify-center rounded-full bg-muted text-muted-foreground">
<User className="h-3 w-3" />
</span>
<div className="flex-1 whitespace-pre-wrap rounded-md bg-accent/30 px-3 py-2 text-sm">
{text}
</div>
</div>
);
}
function AssistantBubble({ text }: { text: string }) {
return (
<div className="flex gap-2">
<span className="mt-0.5 flex h-6 w-6 shrink-0 items-center justify-center rounded-full bg-primary/10 text-primary">
<Sparkles className="h-3 w-3" />
</span>
<div className="flex-1 whitespace-pre-wrap text-sm leading-relaxed">{text}</div>
</div>
);
}
function ToolCallBlock({ tool, inputJson }: { tool: string; inputJson: string }) {
const [expanded, setExpanded] = useState(false);
const preview = extractToolPreview(tool, inputJson);
const Icon = iconForTool(tool);
return (
<div className="ml-8 rounded-md border border-border/40 bg-muted/20">
<button
type="button"
className="flex w-full items-center gap-1.5 px-2 py-1.5 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setExpanded((v) => !v)}
>
{expanded ? <ChevronDown className="h-3 w-3" /> : <ChevronRight className="h-3 w-3" />}
<Icon className="h-3 w-3" />
<span className="font-medium">{labelForTool(tool)}</span>
{preview && (
<span className="ml-1 truncate text-muted-foreground/70">
{preview.slice(0, 80)}
{preview.length > 80 ? "…" : ""}
</span>
)}
</button>
{expanded && (
<div className="border-t border-border/30 p-2">
{tool === "run_query" && preview ? (
<pre className="overflow-x-auto whitespace-pre-wrap rounded bg-background/60 p-2 font-mono text-[11px]">
{preview}
</pre>
) : (
<pre className="overflow-x-auto rounded bg-background/60 p-2 font-mono text-[11px]">
{prettyJson(inputJson)}
</pre>
)}
</div>
)}
</div>
);
}
function ToolResultBlock({
tool,
isError,
text,
result,
}: {
tool: string;
isError: boolean;
text: string | null;
result: { columns: string[]; types: string[]; rows: unknown[][]; row_count: number; execution_time_ms: number } | null;
}) {
if (isError) {
return (
<div className="ml-8 flex items-start gap-2 rounded-md border border-destructive/40 bg-destructive/5 px-3 py-2 text-xs">
<AlertCircle className="mt-0.5 h-3.5 w-3.5 shrink-0 text-destructive" />
<div>
<div className="font-medium text-destructive">{labelForTool(tool)} failed</div>
{text && <div className="mt-1 whitespace-pre-wrap text-muted-foreground">{text}</div>}
</div>
</div>
);
}
// Legacy schema tool — keep a one-line indicator for old threads.
if (tool === "get_schema") {
return (
<div className="ml-8 flex items-center gap-2 rounded-md border border-border/40 bg-muted/20 px-2 py-1.5 text-xs text-muted-foreground">
<Database className="h-3 w-3" />
<span>Loaded schema context ({text?.length ?? 0} chars)</span>
</div>
);
}
// Text-only tools (chat v2/v3): list_databases, list_tables, get_columns, switch_database,
// remember, save_query, find_queries.
if (
tool === "list_databases" ||
tool === "list_tables" ||
tool === "get_columns" ||
tool === "switch_database" ||
tool === "remember" ||
tool === "save_query" ||
tool === "find_queries"
) {
return <TextToolResult tool={tool} text={text} />;
}
// run_query — full results table with Open-full / Export actions.
if (result) {
return <RunQueryResultBlock result={result} />;
}
return null;
}
function RunQueryResultBlock({
result,
}: {
result: {
columns: string[];
types: string[];
rows: unknown[][];
row_count: number;
execution_time_ms: number;
};
}) {
const cap = 100;
const [fullOpen, setFullOpen] = useState(false);
const [exportOpen, setExportOpen] = useState(false);
const previewRows = result.rows.slice(0, cap);
const hasMore = result.rows.length > cap;
return (
<>
<div className="ml-8 overflow-hidden rounded-md border border-border/40 bg-background">
<div className="flex items-center gap-2 border-b border-border/30 px-2 py-1 text-[11px] text-muted-foreground">
<Database className="h-3 w-3" />
<span>
{result.row_count} row{result.row_count === 1 ? "" : "s"} ·{" "}
{result.execution_time_ms} ms
</span>
{hasMore && (
<span className="text-muted-foreground/60">· showing first {cap}</span>
)}
<div className="ml-auto flex items-center gap-1">
<Button
size="icon-xs"
variant="ghost"
className="h-5 w-5 text-muted-foreground hover:text-foreground"
onClick={() => setFullOpen(true)}
title="Open full result"
disabled={result.rows.length === 0}
>
<Maximize2 className="h-3 w-3" />
</Button>
<Button
size="icon-xs"
variant="ghost"
className="h-5 w-5 text-muted-foreground hover:text-foreground"
onClick={() => setExportOpen(true)}
title="Export"
disabled={result.rows.length === 0}
>
<Download className="h-3 w-3" />
</Button>
</div>
</div>
<div className="max-h-72 overflow-auto">
<ResultsTable
columns={result.columns}
types={result.types}
rows={previewRows}
/>
</div>
</div>
<Dialog open={fullOpen} onOpenChange={setFullOpen}>
<DialogContent className="flex h-[80vh] max-w-[90vw] flex-col gap-2 sm:max-w-[90vw]">
<DialogHeader>
<DialogTitle className="flex items-center gap-2 text-sm">
<Database className="h-3.5 w-3.5" />
Full result · {result.row_count} row{result.row_count === 1 ? "" : "s"} ·{" "}
{result.execution_time_ms} ms
<Button
size="xs"
variant="outline"
className="ml-auto gap-1.5"
onClick={() => setExportOpen(true)}
>
<Download className="h-3 w-3" />
Export
</Button>
</DialogTitle>
</DialogHeader>
<div className="min-h-0 flex-1 overflow-hidden rounded-md border border-border/40">
<ResultsTable
columns={result.columns}
types={result.types}
rows={result.rows}
/>
</div>
</DialogContent>
</Dialog>
<ExportDialog
open={exportOpen}
onOpenChange={setExportOpen}
columns={result.columns}
rows={result.rows}
/>
</>
);
}
function TextToolResult({ tool, text }: { tool: string; text: string | null }) {
const [expanded, setExpanded] = useState(tool === "switch_database");
const Icon = iconForTool(tool);
const lineCount = text ? text.split("\n").length : 0;
return (
<div className="ml-8 rounded-md border border-border/40 bg-muted/20">
<button
type="button"
className="flex w-full items-center gap-1.5 px-2 py-1.5 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setExpanded((v) => !v)}
>
{expanded ? <ChevronDown className="h-3 w-3" /> : <ChevronRight className="h-3 w-3" />}
<Icon className="h-3 w-3" />
<span className="font-medium">{labelForTool(tool)}</span>
{text && (
<span className="ml-1 text-muted-foreground/60">
{lineCount} line{lineCount === 1 ? "" : "s"}
</span>
)}
</button>
{expanded && text && (
<div className="border-t border-border/30 p-2">
<pre className="max-h-72 overflow-auto whitespace-pre-wrap rounded bg-background/60 p-2 font-mono text-[11px] leading-relaxed">
{text}
</pre>
</div>
)}
</div>
);
}
function labelForTool(tool: string): string {
switch (tool) {
case "run_query":
return "Run SQL";
case "list_databases":
return "List databases";
case "list_tables":
return "List tables";
case "get_columns":
return "Inspect columns";
case "switch_database":
return "Switch database";
case "remember":
return "Remember";
case "save_query":
return "Save query";
case "find_queries":
return "Find saved queries";
case "get_schema":
return "Load schema";
default:
return tool;
}
}
function iconForTool(tool: string) {
switch (tool) {
case "run_query":
return Wrench;
case "list_databases":
return Database;
case "list_tables":
return Layers;
case "get_columns":
return Columns;
case "switch_database":
return RefreshCw;
case "remember":
return StickyNote;
case "save_query":
return BookmarkPlus;
case "find_queries":
return Bookmark;
case "get_schema":
return Database;
default:
return Wrench;
}
}
function extractToolPreview(tool: string, inputJson: string): string | null {
try {
const parsed = JSON.parse(inputJson) as Record<string, unknown>;
switch (tool) {
case "run_query":
return typeof parsed.sql === "string" ? parsed.sql : null;
case "list_tables":
return typeof parsed.database === "string" ? parsed.database : null;
case "switch_database":
return typeof parsed.database === "string" ? parsed.database : null;
case "get_columns":
return Array.isArray(parsed.tables) ? parsed.tables.join(", ") : null;
case "remember":
return typeof parsed.note === "string" ? parsed.note : null;
case "save_query":
return typeof parsed.name === "string" ? parsed.name : null;
case "find_queries":
return typeof parsed.text === "string" ? parsed.text : null;
default:
return null;
}
} catch {
return null;
}
}
function prettyJson(s: string): string {
try {
return JSON.stringify(JSON.parse(s), null, 2);
} catch {
return s;
}
}
+108
View File
@@ -0,0 +1,108 @@
import { useEffect, useRef } from "react";
import { useChat } from "@/hooks/use-chat";
import { ChatComposer } from "./ChatComposer";
import { ChatMessageView } from "./ChatMessageView";
import { Button } from "@/components/ui/button";
import { Eraser, Sparkles } from "lucide-react";
import { useAppStore } from "@/stores/app-store";
import { useAiSettings } from "@/hooks/use-ai";
interface Props {
tabId: string;
connectionId: string;
}
export function ChatPanel({ tabId, connectionId }: Props) {
const { messages, pending, send, clear } = useChat(tabId, connectionId);
const dbFlavors = useAppStore((s) => s.dbFlavors);
const flavor = dbFlavors[connectionId];
const { data: aiSettings } = useAiSettings();
const aiReady = !!aiSettings?.model;
const scrollerRef = useRef<HTMLDivElement>(null);
useEffect(() => {
scrollerRef.current?.scrollTo({
top: scrollerRef.current.scrollHeight,
behavior: "smooth",
});
}, [messages.length, pending]);
return (
<div className="flex h-full flex-col">
<div className="flex h-9 items-center justify-between border-b border-border/40 px-3">
<div className="flex items-center gap-2 text-xs text-muted-foreground">
<Sparkles className="h-3.5 w-3.5 text-primary/70" />
<span className="font-medium">AI Assistant</span>
{flavor && <span className="text-[10px] uppercase tracking-wider text-muted-foreground/60">· {flavor}</span>}
{aiSettings?.model && (
<span className="text-[10px] text-muted-foreground/60">· {aiSettings.model}</span>
)}
</div>
<Button
size="xs"
variant="ghost"
onClick={clear}
disabled={messages.length === 0 || pending}
className="h-6 gap-1 text-xs text-muted-foreground hover:text-foreground"
>
<Eraser className="h-3 w-3" />
Clear
</Button>
</div>
<div ref={scrollerRef} className="min-h-0 flex-1 overflow-y-auto">
{messages.length === 0 && !pending ? (
<EmptyState aiReady={aiReady} flavor={flavor} />
) : (
<div className="flex flex-col gap-3 px-4 py-3">
{messages.map((m) => (
<ChatMessageView key={m.id} message={m} />
))}
{pending && <PendingIndicator />}
</div>
)}
</div>
<div className="border-t border-border/40 bg-background/40 p-2">
<ChatComposer
onSend={send}
disabled={pending || !aiReady}
placeholder={
aiReady
? "Ask in plain language. The agent will browse schema and run read-only queries."
: "Configure an AI model in Settings to enable chat."
}
/>
</div>
</div>
);
}
function PendingIndicator() {
return (
<div className="flex items-center gap-2 text-xs text-muted-foreground/70">
<span className="inline-flex gap-0.5">
<span className="h-1.5 w-1.5 animate-pulse rounded-full bg-primary/70" />
<span className="h-1.5 w-1.5 animate-pulse rounded-full bg-primary/70 [animation-delay:120ms]" />
<span className="h-1.5 w-1.5 animate-pulse rounded-full bg-primary/70 [animation-delay:240ms]" />
</span>
Thinking...
</div>
);
}
function EmptyState({ aiReady, flavor }: { aiReady: boolean; flavor: string | undefined }) {
return (
<div className="flex h-full items-center justify-center p-6">
<div className="max-w-md space-y-3 text-center">
<Sparkles className="mx-auto h-8 w-8 text-primary/50" />
<h3 className="text-sm font-medium">Ask anything about your data</h3>
<p className="text-xs text-muted-foreground">
{aiReady
? `Connected to ${flavor ?? "database"}. Try: "How many rows in each table?", "Top 10 customers by total spend", "Show me last week's orders".`
: "Open Settings → AI to choose an Ollama model. Tusk will then assist with natural-language queries."}
</p>
</div>
</div>
);
}
+89 -22
View File
@@ -17,8 +17,9 @@ import {
} from "@/components/ui/select";
import { useSaveConnection, useTestConnection } from "@/hooks/use-connections";
import { toast } from "sonner";
import type { ConnectionConfig } from "@/types";
import type { ConnectionConfig, DbFlavor } from "@/types";
import { ENVIRONMENTS } from "@/lib/environment";
import { capsFor } from "@/lib/dbCapabilities";
import { Loader2, X } from "lucide-react";
type InputMode = "fields" | "dsn";
@@ -89,6 +90,8 @@ const emptyConfig: ConnectionConfig = {
ssl_mode: "prefer",
color: undefined,
environment: undefined,
db_flavor: "postgresql",
secure: false,
};
export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
@@ -111,10 +114,29 @@ export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
}
}
const update = (field: keyof ConnectionConfig, value: string | number) => {
const flavor: DbFlavor = form.db_flavor ?? "postgresql";
const isClickHouse = flavor === "clickhouse";
const update = (field: keyof ConnectionConfig, value: string | number | boolean) => {
setForm((f) => ({ ...f, [field]: value }));
};
const handleFlavorChange = (next: DbFlavor) => {
setForm((f) => {
const caps = capsFor(next);
// If user is on a default port for current flavor, swap to the new flavor's default
const currentDefault = capsFor(f.db_flavor).defaultPort;
const port = f.port === currentDefault ? caps.defaultPort : f.port;
return {
...f,
db_flavor: next,
port,
secure: next === "clickhouse" ? f.secure ?? false : false,
ssl_mode: next === "clickhouse" ? undefined : f.ssl_mode ?? "prefer",
};
});
};
const handleTest = () => {
testMutation.mutate(form, {
onSuccess: (version) => {
@@ -164,6 +186,23 @@ export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Engine
</label>
<Select value={flavor} onValueChange={(v) => handleFlavorChange(v as DbFlavor)}>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="postgresql">PostgreSQL</SelectItem>
<SelectItem value="greenplum">Greenplum</SelectItem>
<SelectItem value="clickhouse">ClickHouse</SelectItem>
</SelectContent>
</Select>
</div>
{!isClickHouse && (
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Mode
@@ -206,8 +245,9 @@ export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
</button>
</div>
</div>
)}
{mode === "dsn" ? (
{!isClickHouse && mode === "dsn" ? (
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-2">
DSN
@@ -261,7 +301,12 @@ export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
className="col-span-3"
type="number"
value={form.port}
onChange={(e) => update("port", parseInt(e.target.value) || 5432)}
onChange={(e) =>
update(
"port",
parseInt(e.target.value) || capsFor(flavor).defaultPort,
)
}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
@@ -295,24 +340,46 @@ export function ConnectionDialog({ open, onOpenChange, connection }: Props) {
onChange={(e) => update("database", e.target.value)}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
SSL Mode
</label>
<Select
value={form.ssl_mode ?? "prefer"}
onValueChange={(v) => update("ssl_mode", v)}
>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="disable">Disable</SelectItem>
<SelectItem value="prefer">Prefer</SelectItem>
<SelectItem value="require">Require</SelectItem>
</SelectContent>
</Select>
</div>
{isClickHouse ? (
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
HTTPS
</label>
<div className="col-span-3 flex items-center gap-2">
<Button
type="button"
size="sm"
variant={form.secure ? "default" : "outline"}
className="h-7 px-3 text-xs"
onClick={() => update("secure", !form.secure)}
>
{form.secure ? "On" : "Off"}
</Button>
<span className="text-xs text-muted-foreground">
Use HTTPS scheme for ClickHouse HTTP endpoint
</span>
</div>
</div>
) : (
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
SSL Mode
</label>
<Select
value={form.ssl_mode ?? "prefer"}
onValueChange={(v) => update("ssl_mode", v)}
>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="disable">Disable</SelectItem>
<SelectItem value="prefer">Prefer</SelectItem>
<SelectItem value="require">Require</SelectItem>
</SelectContent>
</Select>
</div>
)}
</>
)}
<div className="grid grid-cols-4 items-center gap-3">
@@ -1,297 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import { useDataGenerator } from "@/hooks/use-data-generator";
import { toast } from "sonner";
import {
Loader2,
CheckCircle2,
XCircle,
Wand2,
Table2,
} from "lucide-react";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
schema: string;
table: string;
}
type Step = "config" | "preview" | "done";
export function GenerateDataDialog({
open,
onOpenChange,
connectionId,
schema,
table,
}: Props) {
const [step, setStep] = useState<Step>("config");
const [rowCount, setRowCount] = useState(10);
const [includeRelated, setIncludeRelated] = useState(true);
const [customInstructions, setCustomInstructions] = useState("");
const {
generatePreview,
preview,
isGenerating,
generateError,
insertData,
insertedRows,
isInserting,
insertError,
progress,
reset,
} = useDataGenerator();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setStep("config");
setRowCount(10);
setIncludeRelated(true);
setCustomInstructions("");
reset();
}
}
const handleGenerate = () => {
const genId = crypto.randomUUID();
generatePreview(
{
params: {
connection_id: connectionId,
schema,
table,
row_count: rowCount,
include_related: includeRelated,
custom_instructions: customInstructions || undefined,
},
genId,
},
{
onSuccess: () => setStep("preview"),
onError: (err) => toast.error("Generation failed", { description: String(err) }),
}
);
};
const handleInsert = () => {
if (!preview) return;
insertData(
{ connectionId, preview },
{
onSuccess: (rows) => {
setStep("done");
toast.success(`Inserted ${rows} rows`);
},
onError: (err) => toast.error("Insert failed", { description: String(err) }),
}
);
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[600px] max-h-[80vh] overflow-y-auto">
<DialogHeader>
<DialogTitle className="flex items-center gap-2">
<Wand2 className="h-5 w-5" />
Generate Test Data
</DialogTitle>
</DialogHeader>
{step === "config" && (
<>
<div className="grid gap-3 py-2">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Table</label>
<div className="col-span-3">
<Badge variant="secondary">{schema}.{table}</Badge>
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Row Count</label>
<Input
className="col-span-3"
type="number"
value={rowCount}
onChange={(e) => setRowCount(Math.min(1000, Math.max(1, parseInt(e.target.value) || 1)))}
min={1}
max={1000}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Related Tables</label>
<div className="col-span-3 flex items-center gap-2">
<input
type="checkbox"
checked={includeRelated}
onChange={(e) => setIncludeRelated(e.target.checked)}
className="rounded"
/>
<span className="text-sm text-muted-foreground">
Include parent tables (via foreign keys)
</span>
</div>
</div>
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-2">Instructions</label>
<Input
className="col-span-3"
placeholder="Optional: specific data requirements..."
value={customInstructions}
onChange={(e) => setCustomInstructions(e.target.value)}
/>
</div>
</div>
{isGenerating && progress && (
<div className="space-y-2">
<div className="flex items-center justify-between text-sm">
<span>{progress.message}</span>
<span className="text-muted-foreground">{progress.percent}%</span>
</div>
<div className="h-2 rounded-full bg-secondary overflow-hidden">
<div
className="h-full rounded-full bg-primary transition-all duration-300"
style={{ width: `${progress.percent}%` }}
/>
</div>
</div>
)}
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Cancel</Button>
<Button onClick={handleGenerate} disabled={isGenerating}>
{isGenerating ? (
<><Loader2 className="h-4 w-4 animate-spin mr-1" />Generating...</>
) : (
"Generate Preview"
)}
</Button>
</DialogFooter>
</>
)}
{step === "preview" && preview && (
<>
<div className="space-y-3">
<div className="flex items-center gap-2 text-sm">
<span className="text-muted-foreground">Preview:</span>
<Badge variant="secondary">{preview.total_rows} rows across {preview.tables.length} tables</Badge>
</div>
{preview.tables.map((tbl) => (
<div key={`${tbl.schema}.${tbl.table}`} className="rounded-md border">
<div className="flex items-center gap-2 px-3 py-2 bg-muted/50 text-sm font-medium border-b">
<Table2 className="h-3.5 w-3.5" />
{tbl.schema}.{tbl.table}
<Badge variant="secondary" className="ml-auto text-[10px]">{tbl.row_count} rows</Badge>
</div>
<div className="overflow-x-auto max-h-48">
<table className="w-full text-xs">
<thead>
<tr className="border-b">
{tbl.columns.map((col) => (
<th key={col} className="px-2 py-1 text-left font-medium text-muted-foreground whitespace-nowrap">
{col}
</th>
))}
</tr>
</thead>
<tbody>
{tbl.rows.slice(0, 5).map((row, i) => (
<tr key={i} className="border-b last:border-0">
{(row as unknown[]).map((val, j) => (
<td key={j} className="px-2 py-1 font-mono whitespace-nowrap">
{val === null ? (
<span className="text-muted-foreground">NULL</span>
) : (
String(val).substring(0, 50)
)}
</td>
))}
</tr>
))}
{tbl.rows.length > 5 && (
<tr>
<td colSpan={tbl.columns.length} className="px-2 py-1 text-center text-muted-foreground">
...and {tbl.rows.length - 5} more rows
</td>
</tr>
)}
</tbody>
</table>
</div>
</div>
))}
</div>
<DialogFooter>
<Button variant="outline" onClick={() => setStep("config")}>Back</Button>
<Button onClick={handleInsert} disabled={isInserting}>
{isInserting ? (
<><Loader2 className="h-4 w-4 animate-spin mr-1" />Inserting...</>
) : (
`Insert ${preview.total_rows} Rows`
)}
</Button>
</DialogFooter>
</>
)}
{step === "done" && (
<div className="py-4 space-y-4">
{insertError ? (
<div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
<XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium text-destructive">Insert Failed</p>
<p className="text-xs text-muted-foreground">{insertError}</p>
</div>
</div>
) : (
<div className="flex items-start gap-3 rounded-md border border-green-500/50 bg-green-500/10 p-4">
<CheckCircle2 className="h-5 w-5 text-green-500 shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium">Data Generated Successfully</p>
<p className="text-xs text-muted-foreground">
{insertedRows} rows inserted across {preview?.tables.length ?? 0} tables.
</p>
</div>
</div>
)}
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Close</Button>
{insertError && (
<Button onClick={() => setStep("preview")}>Retry</Button>
)}
</DialogFooter>
</div>
)}
{generateError && step === "config" && (
<div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
<XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
<p className="text-xs text-muted-foreground">{generateError}</p>
</div>
)}
</DialogContent>
</Dialog>
);
}
@@ -1,498 +0,0 @@
import { useState, useEffect, useRef } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { Badge } from "@/components/ui/badge";
import { useDockerStatus, useCloneToDocker } from "@/hooks/use-docker";
import { toast } from "sonner";
import {
Loader2,
CheckCircle2,
XCircle,
Container,
Copy,
ChevronDown,
ChevronRight,
} from "lucide-react";
import type { CloneMode, CloneProgress } from "@/types";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
database: string;
onConnect?: (connectionId: string) => void;
}
type Step = "config" | "progress" | "done";
function ProcessLog({
entries,
open: logOpen,
onToggle,
endRef,
}: {
entries: CloneProgress[];
open: boolean;
onToggle: () => void;
endRef: React.RefObject<HTMLDivElement | null>;
}) {
if (entries.length === 0) return null;
return (
<div>
<button
type="button"
className="flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground transition-colors"
onClick={onToggle}
>
{logOpen ? (
<ChevronDown className="h-3 w-3" />
) : (
<ChevronRight className="h-3 w-3" />
)}
Process Log ({entries.length})
</button>
{logOpen && (
<div className="mt-1.5 rounded-md bg-muted p-3 text-xs font-mono max-h-40 overflow-auto">
{entries.map((entry, i) => (
<div key={i} className="leading-5 min-w-0">
<span className="text-muted-foreground">
{entry.percent}%
</span>{" "}
<span>{entry.message}</span>
{entry.detail && (
<div className="text-muted-foreground break-all pl-6">
{entry.detail}
</div>
)}
</div>
))}
<div ref={endRef} />
</div>
)}
</div>
);
}
export function CloneDatabaseDialog({
open,
onOpenChange,
connectionId,
database,
onConnect,
}: Props) {
const [step, setStep] = useState<Step>("config");
const [containerName, setContainerName] = useState("");
const [pgVersion, setPgVersion] = useState("16");
const [portMode, setPortMode] = useState<"auto" | "manual">("auto");
const [manualPort, setManualPort] = useState(5433);
const [cloneMode, setCloneMode] = useState<CloneMode>("schema_only");
const [sampleRows, setSampleRows] = useState(1000);
const [logEntries, setLogEntries] = useState<CloneProgress[]>([]);
const [logOpen, setLogOpen] = useState(false);
const logEndRef = useRef<HTMLDivElement>(null);
const { data: dockerStatus } = useDockerStatus();
const { clone, result, error, isCloning, progress, reset } =
useCloneToDocker();
// Reset state when dialog opens
const [prevOpen, setPrevOpen] = useState<{ open: boolean; database: string }>({ open: false, database: "" });
if (open !== prevOpen.open || database !== prevOpen.database) {
setPrevOpen({ open, database });
if (open) {
setStep("config");
setContainerName(
`tusk-${database.replace(/[^a-zA-Z0-9_-]/g, "-")}-${crypto.randomUUID().slice(0, 8)}`
);
setPgVersion("16");
setPortMode("auto");
setManualPort(5433);
setCloneMode("schema_only");
setSampleRows(1000);
setLogEntries([]);
setLogOpen(false);
reset();
}
}
// Accumulate progress events into log
const [prevProgress, setPrevProgress] = useState(progress);
if (progress !== prevProgress) {
setPrevProgress(progress);
if (progress) {
setLogEntries((prev) => {
const last = prev[prev.length - 1];
if (last && last.stage === progress.stage && last.message === progress.message) {
return prev;
}
return [...prev, progress];
});
if (progress.stage === "done" || progress.stage === "error") {
setStep("done");
}
}
}
// Auto-scroll log to bottom
useEffect(() => {
if (logOpen && logEndRef.current) {
logEndRef.current.scrollIntoView({ behavior: "smooth" });
}
}, [logEntries, logOpen]);
const handleClone = () => {
if (!containerName.trim()) {
toast.error("Container name is required");
return;
}
setStep("progress");
const cloneId = crypto.randomUUID();
clone({
params: {
source_connection_id: connectionId,
source_database: database,
container_name: containerName.trim(),
pg_version: pgVersion,
host_port: portMode === "manual" ? manualPort : null,
clone_mode: cloneMode,
sample_rows: cloneMode === "sample_data" ? sampleRows : null,
postgres_password: null,
},
cloneId,
});
};
const handleConnect = () => {
if (result?.connection_id && onConnect) {
onConnect(result.connection_id);
}
onOpenChange(false);
};
const dockerReady =
dockerStatus?.installed && dockerStatus?.daemon_running;
const logSection = (
<ProcessLog
entries={logEntries}
open={logOpen}
onToggle={() => setLogOpen(!logOpen)}
endRef={logEndRef}
/>
);
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[520px]">
<DialogHeader>
<DialogTitle className="flex items-center gap-2">
<Container className="h-5 w-5" />
Clone to Docker
</DialogTitle>
</DialogHeader>
{step === "config" && (
<>
<div className="flex items-center gap-2 rounded-md border px-3 py-2 text-sm">
{dockerStatus === undefined ? (
<>
<Loader2 className="h-4 w-4 animate-spin text-muted-foreground" />
<span className="text-muted-foreground">
Checking Docker...
</span>
</>
) : dockerReady ? (
<>
<CheckCircle2 className="h-4 w-4 text-green-500" />
<span>Docker {dockerStatus.version}</span>
</>
) : (
<>
<XCircle className="h-4 w-4 text-destructive" />
<span className="text-destructive">
{dockerStatus?.error || "Docker not available"}
</span>
</>
)}
</div>
<div className="grid gap-3 py-2">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Database
</label>
<div className="col-span-3">
<Badge variant="secondary">{database}</Badge>
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Container
</label>
<Input
className="col-span-3"
value={containerName}
onChange={(e) => setContainerName(e.target.value)}
placeholder="tusk-mydb-clone"
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
PG Version
</label>
<Select value={pgVersion} onValueChange={setPgVersion}>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="17">PostgreSQL 17</SelectItem>
<SelectItem value="16">PostgreSQL 16</SelectItem>
<SelectItem value="15">PostgreSQL 15</SelectItem>
<SelectItem value="14">PostgreSQL 14</SelectItem>
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Port
</label>
<div className="col-span-3 flex items-center gap-2">
<Select
value={portMode}
onValueChange={(v) =>
setPortMode(v as "auto" | "manual")
}
>
<SelectTrigger className="w-24">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="auto">Auto</SelectItem>
<SelectItem value="manual">Manual</SelectItem>
</SelectContent>
</Select>
{portMode === "manual" && (
<Input
type="number"
className="flex-1"
value={manualPort}
onChange={(e) =>
setManualPort(parseInt(e.target.value) || 5433)
}
min={1024}
max={65535}
/>
)}
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Clone Mode
</label>
<Select
value={cloneMode}
onValueChange={(v) => setCloneMode(v as CloneMode)}
>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="schema_only">
Schema Only
</SelectItem>
<SelectItem value="full_clone">Full Clone</SelectItem>
<SelectItem value="sample_data">
Sample Data
</SelectItem>
</SelectContent>
</Select>
</div>
{cloneMode === "sample_data" && (
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">
Sample Rows
</label>
<Input
className="col-span-3"
type="number"
value={sampleRows}
onChange={(e) =>
setSampleRows(parseInt(e.target.value) || 1000)
}
min={1}
max={100000}
/>
</div>
)}
</div>
<DialogFooter>
<Button
variant="outline"
onClick={() => onOpenChange(false)}
>
Cancel
</Button>
<Button onClick={handleClone} disabled={!dockerReady}>
Clone
</Button>
</DialogFooter>
</>
)}
{step === "progress" && (
<div className="py-4 space-y-4">
<div className="space-y-2">
<div className="flex items-center justify-between text-sm">
<span>{progress?.message || "Starting..."}</span>
<span className="text-muted-foreground">
{progress?.percent ?? 0}%
</span>
</div>
<div className="h-2 rounded-full bg-secondary overflow-hidden">
<div
className="h-full rounded-full bg-primary transition-all duration-300"
style={{ width: `${progress?.percent ?? 0}%` }}
/>
</div>
</div>
{isCloning && (
<div className="flex items-center gap-2 text-sm text-muted-foreground">
<Loader2 className="h-4 w-4 animate-spin" />
{progress?.stage || "Initializing..."}
</div>
)}
{logSection}
</div>
)}
{step === "done" && (
<div className="py-4 space-y-4">
{error ? (
<div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
<XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium text-destructive">
Clone Failed
</p>
<p className="text-xs text-muted-foreground">{error}</p>
</div>
</div>
) : (
<div className="space-y-3">
<div className="flex items-start gap-3 rounded-md border border-green-500/50 bg-green-500/10 p-4">
<CheckCircle2 className="h-5 w-5 text-green-500 shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium">
Clone Completed
</p>
<p className="text-xs text-muted-foreground">
Database cloned to Docker container successfully.
</p>
</div>
</div>
{result && (
<div className="rounded-md border p-3 space-y-2 text-sm">
<div className="flex items-center justify-between">
<span className="text-muted-foreground">
Container
</span>
<span className="font-mono">
{result.container.name}
</span>
</div>
<div className="flex items-center justify-between">
<span className="text-muted-foreground">Port</span>
<span className="font-mono">
{result.container.host_port}
</span>
</div>
<div className="flex items-center justify-between gap-2">
<span className="text-muted-foreground">URL</span>
<div className="flex items-center gap-1">
<span className="font-mono text-xs truncate max-w-[250px]">
{result.connection_url}
</span>
<Button
variant="ghost"
size="sm"
className="h-6 w-6 p-0"
onClick={() => {
navigator.clipboard.writeText(
result.connection_url
);
toast.success("URL copied");
}}
>
<Copy className="h-3 w-3" />
</Button>
</div>
</div>
</div>
)}
</div>
)}
{logSection}
<DialogFooter>
{error ? (
<>
<Button
variant="outline"
onClick={() => onOpenChange(false)}
>
Close
</Button>
<Button
onClick={() => setStep("config")}
>
Retry
</Button>
</>
) : (
<>
<Button
variant="outline"
onClick={() => onOpenChange(false)}
>
Close
</Button>
{onConnect && result && (
<Button onClick={handleConnect}>Connect</Button>
)}
</>
)}
</DialogFooter>
</div>
)}
</DialogContent>
</Dialog>
);
}
@@ -1,179 +0,0 @@
import { useState } from "react";
import {
useTuskContainers,
useStartContainer,
useStopContainer,
useRemoveContainer,
useDockerStatus,
} from "@/hooks/use-docker";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { toast } from "sonner";
import {
ChevronDown,
ChevronRight,
Container,
Play,
Square,
Trash2,
Loader2,
} from "lucide-react";
export function DockerContainersList() {
const [expanded, setExpanded] = useState(true);
const { data: dockerStatus } = useDockerStatus();
const { data: containers, isLoading } = useTuskContainers();
const startMutation = useStartContainer();
const stopMutation = useStopContainer();
const removeMutation = useRemoveContainer();
const dockerAvailable =
dockerStatus?.installed && dockerStatus?.daemon_running;
if (!dockerAvailable) {
return null;
}
const handleStart = (name: string) => {
startMutation.mutate(name, {
onSuccess: () => toast.success(`Container "${name}" started`),
onError: (err) =>
toast.error("Failed to start container", {
description: String(err),
}),
});
};
const handleStop = (name: string) => {
stopMutation.mutate(name, {
onSuccess: () => toast.success(`Container "${name}" stopped`),
onError: (err) =>
toast.error("Failed to stop container", {
description: String(err),
}),
});
};
const handleRemove = (name: string) => {
if (
!confirm(
`Remove container "${name}"? This will delete the container and all its data.`
)
) {
return;
}
removeMutation.mutate(name, {
onSuccess: () => toast.success(`Container "${name}" removed`),
onError: (err) =>
toast.error("Failed to remove container", {
description: String(err),
}),
});
};
const isRunning = (status: string) =>
status.toLowerCase().startsWith("up");
return (
<div className="border-b">
<div
className="flex items-center gap-1 px-3 py-2 cursor-pointer select-none hover:bg-accent/50"
onClick={() => setExpanded(!expanded)}
>
{expanded ? (
<ChevronDown className="h-3.5 w-3.5 text-muted-foreground" />
) : (
<ChevronRight className="h-3.5 w-3.5 text-muted-foreground" />
)}
<Container className="h-3.5 w-3.5 text-muted-foreground" />
<span className="text-xs font-semibold flex-1">Docker Clones</span>
{containers && containers.length > 0 && (
<Badge
variant="secondary"
className="text-[9px] px-1 py-0"
>
{containers.length}
</Badge>
)}
</div>
{expanded && (
<div className="pb-1">
{isLoading && (
<div className="px-3 py-2 text-xs text-muted-foreground flex items-center gap-1">
<Loader2 className="h-3 w-3 animate-spin" /> Loading...
</div>
)}
{containers && containers.length === 0 && (
<div className="px-6 pb-2 text-xs text-muted-foreground">
No Docker clones yet. Right-click a database to clone it.
</div>
)}
{containers?.map((container) => (
<div
key={container.container_id}
className="group flex items-center gap-1.5 px-6 py-1 text-xs hover:bg-accent/50"
>
<span className="truncate flex-1 font-medium">
{container.name}
</span>
{container.source_database && (
<span className="text-[10px] text-muted-foreground shrink-0">
{container.source_database}
</span>
)}
<span className="text-[10px] text-muted-foreground shrink-0">
:{container.host_port}
</span>
<Badge
variant={isRunning(container.status) ? "default" : "secondary"}
className={`text-[9px] px-1 py-0 shrink-0 ${
isRunning(container.status)
? "bg-green-600 hover:bg-green-600"
: ""
}`}
>
{isRunning(container.status) ? "running" : "stopped"}
</Badge>
<div className="flex gap-0.5 opacity-0 group-hover:opacity-100 shrink-0">
{isRunning(container.status) ? (
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
onClick={() => handleStop(container.name)}
title="Stop"
disabled={stopMutation.isPending}
>
<Square className="h-3 w-3" />
</Button>
) : (
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
onClick={() => handleStart(container.name)}
title="Start"
disabled={startMutation.isPending}
>
<Play className="h-3 w-3" />
</Button>
)}
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0 text-destructive hover:text-destructive"
onClick={() => handleRemove(container.name)}
title="Remove"
disabled={removeMutation.isPending}
>
<Trash2 className="h-3 w-3" />
</Button>
</div>
</div>
))}
</div>
)}
</div>
);
}
+11 -4
View File
@@ -1,7 +1,8 @@
import CodeMirror from "@uiw/react-codemirror";
import { sql, PostgreSQL } from "@codemirror/lang-sql";
import { sql, PostgreSQL, StandardSQL } from "@codemirror/lang-sql";
import { keymap } from "@codemirror/view";
import { useCallback, useMemo } from "react";
import { useAppStore } from "@/stores/app-store";
interface Props {
value: string;
@@ -27,6 +28,10 @@ function buildSqlNamespace(
}
export function SqlEditor({ value, onChange, onExecute, onFormat, schema }: Props) {
const activeConnectionId = useAppStore((s) => s.activeConnectionId);
const dbFlavors = useAppStore((s) => s.dbFlavors);
const flavor = activeConnectionId ? dbFlavors[activeConnectionId] : undefined;
const handleChange = useCallback(
(val: string) => {
onChange(val);
@@ -36,11 +41,13 @@ export function SqlEditor({ value, onChange, onExecute, onFormat, schema }: Prop
const extensions = useMemo(() => {
const sqlNamespace = schema ? buildSqlNamespace(schema) : undefined;
const dialect = flavor === "clickhouse" ? StandardSQL : PostgreSQL;
const defaultSchema = flavor === "clickhouse" ? undefined : "public";
return [
sql({
dialect: PostgreSQL,
dialect,
schema: sqlNamespace,
defaultSchema: "public",
defaultSchema,
}),
keymap.of([
{
@@ -66,7 +73,7 @@ export function SqlEditor({ value, onChange, onExecute, onFormat, schema }: Prop
},
]),
];
}, [onExecute, onFormat, schema]);
}, [onExecute, onFormat, schema, flavor]);
return (
<CodeMirror
-186
View File
@@ -1,186 +0,0 @@
import { useMemo, useCallback, useState } from "react";
import { useTheme } from "next-themes";
import {
ReactFlow,
Background,
Controls,
MiniMap,
MarkerType,
PanOnScrollMode,
applyNodeChanges,
applyEdgeChanges,
type Node,
type Edge,
type NodeTypes,
type NodeChange,
type EdgeChange,
} from "@xyflow/react";
import dagre from "dagre";
import "@xyflow/react/dist/style.css";
import { useSchemaErd } from "@/hooks/use-schema";
import { ErdTableNode, type ErdTableNodeData } from "./ErdTableNode";
import type { ErdData } from "@/types";
const nodeTypes: NodeTypes = {
erdTable: ErdTableNode,
};
const NODE_WIDTH = 250;
const NODE_ROW_HEIGHT = 24;
const NODE_HEADER_HEIGHT = 36;
function buildLayout(data: ErdData): { nodes: Node[]; edges: Edge[] } {
const g = new dagre.graphlib.Graph();
g.setDefaultEdgeLabel(() => ({}));
g.setGraph({ rankdir: "LR", nodesep: 60, ranksep: 150 });
// Build list of FK column names per table for icon display
const fkColumnsPerTable = new Map<string, string[]>();
for (const rel of data.relationships) {
const key = `${rel.source_schema}.${rel.source_table}`;
if (!fkColumnsPerTable.has(key)) fkColumnsPerTable.set(key, []);
for (const col of rel.source_columns) {
const arr = fkColumnsPerTable.get(key)!;
if (!arr.includes(col)) arr.push(col);
}
}
for (const table of data.tables) {
const height = NODE_HEADER_HEIGHT + table.columns.length * NODE_ROW_HEIGHT;
g.setNode(table.name, { width: NODE_WIDTH, height });
}
for (const rel of data.relationships) {
g.setEdge(rel.source_table, rel.target_table);
}
dagre.layout(g);
const nodes: Node[] = data.tables.map((table) => {
const pos = g.node(table.name);
const tableKey = `${table.schema}.${table.name}`;
return {
id: table.name,
type: "erdTable",
position: { x: pos.x - NODE_WIDTH / 2, y: pos.y - pos.height / 2 },
data: {
label: table.name,
schema: table.schema,
columns: table.columns,
fkColumnNames: fkColumnsPerTable.get(tableKey) ?? [],
} satisfies ErdTableNodeData,
};
});
const edges: Edge[] = data.relationships.map((rel) => ({
id: rel.constraint_name,
source: rel.source_table,
target: rel.target_table,
type: "smoothstep",
label: rel.constraint_name,
labelStyle: { fontSize: 10, fill: "var(--muted-foreground)" },
labelBgStyle: { fill: "var(--card)", fillOpacity: 0.8 },
labelBgPadding: [4, 2] as [number, number],
markerEnd: {
type: MarkerType.ArrowClosed,
width: 16,
height: 16,
color: "var(--muted-foreground)",
},
style: { stroke: "var(--muted-foreground)", strokeWidth: 1.5 },
}));
return { nodes, edges };
}
interface Props {
connectionId: string;
schema: string;
}
export function ErdDiagram({ connectionId, schema }: Props) {
const { data: erdData, isLoading, error } = useSchemaErd(connectionId, schema);
const { resolvedTheme } = useTheme();
const layout = useMemo(() => {
if (!erdData) return null;
return buildLayout(erdData);
}, [erdData]);
const [nodes, setNodes] = useState<Node[]>([]);
const [edges, setEdges] = useState<Edge[]>([]);
const [prevLayout, setPrevLayout] = useState(layout);
if (layout !== prevLayout) {
setPrevLayout(layout);
if (layout) {
setNodes(layout.nodes);
setEdges(layout.edges);
}
}
const onNodesChange = useCallback(
(changes: NodeChange[]) => setNodes((nds) => applyNodeChanges(changes, nds)),
[],
);
const onEdgesChange = useCallback(
(changes: EdgeChange[]) => setEdges((eds) => applyEdgeChanges(changes, eds)),
[],
);
if (isLoading) {
return (
<div className="flex h-full items-center justify-center text-sm text-muted-foreground">
Loading ER diagram...
</div>
);
}
if (error) {
return (
<div className="flex h-full items-center justify-center text-sm text-destructive">
Error loading ER diagram: {String(error)}
</div>
);
}
if (!erdData || erdData.tables.length === 0) {
return (
<div className="flex h-full items-center justify-center text-sm text-muted-foreground">
No tables found in schema &quot;{schema}&quot;.
</div>
);
}
return (
<div className="h-full w-full">
<ReactFlow
nodes={nodes}
edges={edges}
onNodesChange={onNodesChange}
onEdgesChange={onEdgesChange}
nodeTypes={nodeTypes}
fitView
colorMode={resolvedTheme === "dark" ? "dark" : "light"}
minZoom={0.05}
maxZoom={3}
zoomOnScroll
zoomOnPinch
panOnScroll
panOnScrollMode={PanOnScrollMode.Free}
proOptions={{ hideAttribution: true }}
>
<Background gap={16} size={1} />
<Controls className="!bg-card !border !shadow-sm [&>button]:!bg-card [&>button]:!border-border [&>button]:!text-foreground" />
<MiniMap
className="!bg-card !border"
nodeColor="var(--muted)"
maskColor="rgba(0, 0, 0, 0.7)"
/>
</ReactFlow>
</div>
);
}
-54
View File
@@ -1,54 +0,0 @@
import { memo } from "react";
import { Handle, Position, type NodeProps } from "@xyflow/react";
import type { ErdColumn } from "@/types";
import { KeyRound, Link } from "lucide-react";
export interface ErdTableNodeData {
label: string;
schema: string;
columns: ErdColumn[];
fkColumnNames: string[];
[key: string]: unknown;
}
function ErdTableNodeComponent({ data }: NodeProps) {
const { label, columns, fkColumnNames } = data as unknown as ErdTableNodeData;
return (
<div className="min-w-[220px] rounded-lg border border-border bg-card text-card-foreground shadow-md">
<div className="rounded-t-lg border-b bg-primary/10 px-3 py-2 text-xs font-bold tracking-wide text-primary">
{label}
</div>
<div className="divide-y divide-border/50">
{(columns as ErdColumn[]).map((col, i) => (
<div key={i} className="flex items-center gap-1.5 px-3 py-1 text-[11px]">
{col.is_primary_key ? (
<KeyRound className="h-3 w-3 shrink-0 text-amber-500" />
) : (fkColumnNames as string[]).includes(col.name) ? (
<Link className="h-3 w-3 shrink-0 text-blue-400" />
) : (
<span className="h-3 w-3 shrink-0" />
)}
<span className="font-medium">{col.name}</span>
<span className="ml-auto text-muted-foreground">{col.data_type}</span>
{col.is_nullable && (
<span className="text-muted-foreground/60">?</span>
)}
</div>
))}
</div>
<Handle
type="target"
position={Position.Left}
className="!w-1.5 !h-1.5 !bg-primary !opacity-0 hover:!opacity-100 !border-none !min-w-0 !min-h-0"
/>
<Handle
type="source"
position={Position.Right}
className="!w-1.5 !h-1.5 !bg-primary !opacity-0 hover:!opacity-100 !border-none !min-w-0 !min-h-0"
/>
</div>
);
}
export const ErdTableNode = memo(ErdTableNodeComponent);
@@ -1,232 +0,0 @@
import { useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { useIndexAdvisorReport, useApplyIndexRecommendation } from "@/hooks/use-index-advisor";
import { RecommendationCard } from "./RecommendationCard";
import { toast } from "sonner";
import { Loader2, Gauge, Search, AlertTriangle } from "lucide-react";
import type { IndexAdvisorReport } from "@/types";
interface Props {
connectionId: string;
}
export function IndexAdvisorPanel({ connectionId }: Props) {
const [report, setReport] = useState<IndexAdvisorReport | null>(null);
const [appliedDdls, setAppliedDdls] = useState<Set<string>>(new Set());
const [applyingDdl, setApplyingDdl] = useState<string | null>(null);
const reportMutation = useIndexAdvisorReport();
const applyMutation = useApplyIndexRecommendation();
const handleAnalyze = () => {
reportMutation.mutate(connectionId, {
onSuccess: (data) => {
setReport(data);
setAppliedDdls(new Set());
},
onError: (err) => toast.error("Analysis failed", { description: String(err) }),
});
};
const handleApply = async (ddl: string) => {
if (!confirm("Apply this index change? This will modify the database schema.")) return;
setApplyingDdl(ddl);
try {
await applyMutation.mutateAsync({ connectionId, ddl });
setAppliedDdls((prev) => new Set(prev).add(ddl));
toast.success("Index change applied");
} catch (err) {
toast.error("Failed to apply", { description: String(err) });
} finally {
setApplyingDdl(null);
}
};
return (
<div className="flex h-full flex-col">
{/* Header */}
<div className="border-b px-4 py-3 flex items-center justify-between">
<div className="flex items-center gap-2">
<Gauge className="h-5 w-5 text-primary" />
<h2 className="text-sm font-medium">Index Advisor</h2>
</div>
<Button
size="sm"
onClick={handleAnalyze}
disabled={reportMutation.isPending}
>
{reportMutation.isPending ? (
<><Loader2 className="h-3.5 w-3.5 animate-spin mr-1" />Analyzing...</>
) : (
<><Search className="h-3.5 w-3.5 mr-1" />Analyze</>
)}
</Button>
</div>
{/* Content */}
<div className="flex-1 overflow-auto">
{!report ? (
<div className="flex h-full items-center justify-center text-sm text-muted-foreground">
Click Analyze to scan your database for index optimization opportunities.
</div>
) : (
<Tabs defaultValue="recommendations" className="h-full flex flex-col">
<div className="border-b px-4">
<TabsList className="h-9">
<TabsTrigger value="recommendations" className="text-xs">
Recommendations
{report.recommendations.length > 0 && (
<Badge variant="secondary" className="ml-1 text-[10px]">{report.recommendations.length}</Badge>
)}
</TabsTrigger>
<TabsTrigger value="table-stats" className="text-xs">Table Stats</TabsTrigger>
<TabsTrigger value="index-stats" className="text-xs">Index Stats</TabsTrigger>
<TabsTrigger value="slow-queries" className="text-xs">
Slow Queries
{!report.has_pg_stat_statements && (
<AlertTriangle className="h-3 w-3 ml-1 text-yellow-500" />
)}
</TabsTrigger>
</TabsList>
</div>
<TabsContent value="recommendations" className="flex-1 overflow-auto p-4 space-y-2 mt-0">
{report.recommendations.length === 0 ? (
<div className="text-sm text-muted-foreground text-center py-8">
No recommendations found. Your indexes look good!
</div>
) : (
report.recommendations.map((rec, i) => (
<RecommendationCard
key={rec.id || i}
recommendation={rec}
onApply={handleApply}
isApplying={applyingDdl === rec.ddl}
applied={appliedDdls.has(rec.ddl)}
/>
))
)}
</TabsContent>
<TabsContent value="table-stats" className="flex-1 overflow-auto mt-0">
<div className="overflow-x-auto">
<table className="w-full text-xs">
<thead>
<tr className="border-b bg-muted/50">
<th className="px-3 py-2 text-left font-medium">Table</th>
<th className="px-3 py-2 text-right font-medium">Seq Scans</th>
<th className="px-3 py-2 text-right font-medium">Idx Scans</th>
<th className="px-3 py-2 text-right font-medium">Rows</th>
<th className="px-3 py-2 text-right font-medium">Table Size</th>
<th className="px-3 py-2 text-right font-medium">Index Size</th>
</tr>
</thead>
<tbody>
{report.table_stats.map((ts) => {
const ratio = ts.seq_scan + ts.idx_scan > 0
? ts.seq_scan / (ts.seq_scan + ts.idx_scan)
: 0;
return (
<tr key={`${ts.schema}.${ts.table}`} className="border-b">
<td className="px-3 py-2 font-mono">{ts.schema}.{ts.table}</td>
<td className={`px-3 py-2 text-right ${ratio > 0.8 && ts.n_live_tup > 1000 ? "text-destructive font-medium" : ""}`}>
{ts.seq_scan.toLocaleString()}
</td>
<td className="px-3 py-2 text-right">{ts.idx_scan.toLocaleString()}</td>
<td className="px-3 py-2 text-right">{ts.n_live_tup.toLocaleString()}</td>
<td className="px-3 py-2 text-right">{ts.table_size}</td>
<td className="px-3 py-2 text-right">{ts.index_size}</td>
</tr>
);
})}
</tbody>
</table>
</div>
</TabsContent>
<TabsContent value="index-stats" className="flex-1 overflow-auto mt-0">
<div className="overflow-x-auto">
<table className="w-full text-xs">
<thead>
<tr className="border-b bg-muted/50">
<th className="px-3 py-2 text-left font-medium">Index</th>
<th className="px-3 py-2 text-left font-medium">Table</th>
<th className="px-3 py-2 text-right font-medium">Scans</th>
<th className="px-3 py-2 text-right font-medium">Size</th>
<th className="px-3 py-2 text-left font-medium">Definition</th>
</tr>
</thead>
<tbody>
{report.index_stats.map((is) => (
<tr key={`${is.schema}.${is.index_name}`} className="border-b">
<td className={`px-3 py-2 font-mono ${is.idx_scan === 0 ? "text-yellow-600" : ""}`}>
{is.index_name}
</td>
<td className="px-3 py-2">{is.schema}.{is.table}</td>
<td className={`px-3 py-2 text-right ${is.idx_scan === 0 ? "text-yellow-600 font-medium" : ""}`}>
{is.idx_scan.toLocaleString()}
</td>
<td className="px-3 py-2 text-right">{is.index_size}</td>
<td className="px-3 py-2 font-mono text-muted-foreground max-w-xs truncate">
{is.definition}
</td>
</tr>
))}
</tbody>
</table>
</div>
</TabsContent>
<TabsContent value="slow-queries" className="flex-1 overflow-auto mt-0">
{!report.has_pg_stat_statements ? (
<div className="p-4 text-sm text-muted-foreground">
<div className="flex items-center gap-2 mb-2">
<AlertTriangle className="h-4 w-4 text-yellow-500" />
pg_stat_statements extension is not installed
</div>
<p className="text-xs">
Enable it with: CREATE EXTENSION pg_stat_statements;
</p>
</div>
) : report.slow_queries.length === 0 ? (
<div className="p-4 text-sm text-muted-foreground text-center">
No slow queries found.
</div>
) : (
<div className="overflow-x-auto">
<table className="w-full text-xs">
<thead>
<tr className="border-b bg-muted/50">
<th className="px-3 py-2 text-left font-medium">Query</th>
<th className="px-3 py-2 text-right font-medium">Calls</th>
<th className="px-3 py-2 text-right font-medium">Mean (ms)</th>
<th className="px-3 py-2 text-right font-medium">Total (ms)</th>
<th className="px-3 py-2 text-right font-medium">Rows</th>
</tr>
</thead>
<tbody>
{report.slow_queries.map((sq, i) => (
<tr key={i} className="border-b">
<td className="px-3 py-2 font-mono max-w-md truncate" title={sq.query}>
{sq.query.substring(0, 150)}
</td>
<td className="px-3 py-2 text-right">{sq.calls.toLocaleString()}</td>
<td className="px-3 py-2 text-right">{sq.mean_time_ms.toFixed(1)}</td>
<td className="px-3 py-2 text-right">{sq.total_time_ms.toFixed(0)}</td>
<td className="px-3 py-2 text-right">{sq.rows.toLocaleString()}</td>
</tr>
))}
</tbody>
</table>
</div>
)}
</TabsContent>
</Tabs>
)}
</div>
</div>
);
}
@@ -1,86 +0,0 @@
import { useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Loader2, Play } from "lucide-react";
import type { IndexRecommendation } from "@/types";
interface Props {
recommendation: IndexRecommendation;
onApply: (ddl: string) => void;
isApplying: boolean;
applied: boolean;
}
function priorityBadge(priority: string) {
switch (priority.toLowerCase()) {
case "high":
return <Badge variant="destructive">{priority}</Badge>;
case "medium":
return <Badge className="bg-yellow-600 text-white">{priority}</Badge>;
default:
return <Badge variant="secondary">{priority}</Badge>;
}
}
function typeBadge(type: string) {
switch (type) {
case "create_index":
return <Badge className="bg-green-600 text-white">CREATE</Badge>;
case "drop_index":
return <Badge variant="destructive">DROP</Badge>;
case "replace_index":
return <Badge className="bg-blue-600 text-white">REPLACE</Badge>;
default:
return <Badge variant="secondary">{type}</Badge>;
}
}
export function RecommendationCard({ recommendation, onApply, isApplying, applied }: Props) {
const [showDdl] = useState(true);
return (
<div className="rounded-md border p-3 space-y-2">
<div className="flex items-start justify-between gap-2">
<div className="flex items-center gap-2 flex-wrap">
{typeBadge(recommendation.recommendation_type)}
{priorityBadge(recommendation.priority)}
<span className="text-xs text-muted-foreground">
{recommendation.table_schema}.{recommendation.table_name}
</span>
{recommendation.index_name && (
<span className="text-xs font-mono text-muted-foreground">
{recommendation.index_name}
</span>
)}
</div>
<Button
size="sm"
variant={applied ? "outline" : "default"}
onClick={() => onApply(recommendation.ddl)}
disabled={isApplying || applied}
className="shrink-0"
>
{isApplying ? (
<Loader2 className="h-3.5 w-3.5 animate-spin mr-1" />
) : applied ? (
"Applied"
) : (
<><Play className="h-3.5 w-3.5 mr-1" />Apply</>
)}
</Button>
</div>
<p className="text-sm">{recommendation.rationale}</p>
<div className="flex items-center gap-2 text-xs text-muted-foreground">
<span>Impact: {recommendation.estimated_impact}</span>
</div>
{showDdl && (
<pre className="rounded bg-muted p-2 text-xs font-mono overflow-x-auto">
{recommendation.ddl}
</pre>
)}
</div>
);
}
+50
View File
@@ -0,0 +1,50 @@
import { Component, type ReactNode } from "react";
interface Props {
children: ReactNode;
}
interface State {
error: Error | null;
}
export class ErrorBoundary extends Component<Props, State> {
state: State = { error: null };
static getDerivedStateFromError(error: Error): State {
return { error };
}
componentDidCatch(error: Error, info: React.ErrorInfo) {
console.error("Tusk render crash:", error, info.componentStack);
}
reset = () => this.setState({ error: null });
render() {
if (this.state.error) {
return (
<div className="flex h-screen items-center justify-center bg-background p-6">
<div className="max-w-xl space-y-3">
<h1 className="text-base font-semibold text-destructive">Something broke while rendering Tusk</h1>
<pre className="overflow-auto rounded-md border border-border/50 bg-muted/40 p-3 text-xs">
{this.state.error.message}
{this.state.error.stack ? `\n\n${this.state.error.stack}` : ""}
</pre>
<p className="text-xs text-muted-foreground">
Open the developer console for the full stack. Click "Try again" to re-mount the UI.
</p>
<button
type="button"
className="rounded-md border border-border/60 px-3 py-1 text-xs hover:bg-accent/40"
onClick={this.reset}
>
Try again
</button>
</div>
</div>
);
}
return this.props.children;
}
}
+5 -5
View File
@@ -10,10 +10,10 @@ import {
import { SchemaTree } from "@/components/schema/SchemaTree";
import { HistoryPanel } from "@/components/history/HistoryPanel";
import { SavedQueriesPanel } from "@/components/saved-queries/SavedQueriesPanel";
import { AdminPanel } from "@/components/management/AdminPanel";
import { Search, RefreshCw, Layers, Clock, Bookmark, Shield } from "lucide-react";
import { MemoryPanel } from "@/components/memory/MemoryPanel";
import { Search, RefreshCw, Layers, Clock, Bookmark, Brain } from "lucide-react";
type SidebarView = "schema" | "history" | "saved" | "admin";
type SidebarView = "schema" | "history" | "saved" | "memory";
const SCHEMA_QUERY_KEYS = [
"databases", "schemas", "tables", "views",
@@ -24,7 +24,7 @@ const SIDEBAR_TABS: { id: SidebarView; label: string; icon: React.ReactNode }[]
{ id: "schema", label: "Schema", icon: <Layers className="h-3.5 w-3.5" /> },
{ id: "history", label: "History", icon: <Clock className="h-3.5 w-3.5" /> },
{ id: "saved", label: "Saved", icon: <Bookmark className="h-3.5 w-3.5" /> },
{ id: "admin", label: "Admin", icon: <Shield className="h-3.5 w-3.5" /> },
{ id: "memory", label: "Memory", icon: <Brain className="h-3.5 w-3.5" /> },
];
export function Sidebar() {
@@ -93,7 +93,7 @@ export function Sidebar() {
) : view === "saved" ? (
<SavedQueriesPanel />
) : (
<AdminPanel />
<MemoryPanel />
)}
</div>
);
+2 -8
View File
@@ -1,7 +1,7 @@
import { useAppStore } from "@/stores/app-store";
import { useConnections } from "@/hooks/use-connections";
import { ScrollArea } from "@/components/ui/scroll-area";
import { X, Table2, Code, Columns, Users, Activity, Search, GitFork, ShieldCheck, Gauge, Camera } from "lucide-react";
import { X, Table2, Code, Columns, Sparkles } from "lucide-react";
export function TabBar() {
const { tabs, activeTabId, setActiveTabId, closeTab } = useAppStore();
@@ -13,13 +13,7 @@ export function TabBar() {
query: <Code className="h-3 w-3" />,
table: <Table2 className="h-3 w-3" />,
structure: <Columns className="h-3 w-3" />,
roles: <Users className="h-3 w-3" />,
sessions: <Activity className="h-3 w-3" />,
lookup: <Search className="h-3 w-3" />,
erd: <GitFork className="h-3 w-3" />,
validation: <ShieldCheck className="h-3 w-3" />,
"index-advisor": <Gauge className="h-3 w-3" />,
snapshots: <Camera className="h-3 w-3" />,
chat: <Sparkles className="h-3 w-3" />,
};
return (
+10 -10
View File
@@ -7,7 +7,7 @@ import { ReadOnlyToggle } from "@/components/layout/ReadOnlyToggle";
import { useAppStore } from "@/stores/app-store";
import { useConnections, useReconnect } from "@/hooks/use-connections";
import { toast } from "sonner";
import { Database, Plus, RefreshCw, Search, Settings } from "lucide-react";
import { Database, Plus, RefreshCw, Settings, Sparkles } from "lucide-react";
import type { ConnectionConfig, Tab } from "@/types";
import { getEnvironment } from "@/lib/environment";
import { AppSettingsSheet } from "@/components/settings/AppSettingsSheet";
@@ -45,12 +45,12 @@ export function Toolbar() {
addTab(tab);
};
const handleNewLookup = () => {
const handleNewChat = () => {
if (!activeConnectionId) return;
const tab: Tab = {
id: crypto.randomUUID(),
type: "lookup",
title: "Entity Lookup",
type: "chat",
title: "Chat",
connectionId: activeConnectionId,
database: currentDatabase ?? undefined,
};
@@ -101,22 +101,22 @@ export function Toolbar() {
variant="ghost"
size="xs"
className="gap-1.5 text-muted-foreground hover:text-foreground"
onClick={handleNewQuery}
onClick={handleNewChat}
disabled={!activeConnectionId}
>
<Plus className="h-3.5 w-3.5" />
<span className="text-xs font-medium">New Query</span>
<Sparkles className="h-3.5 w-3.5" />
<span className="text-xs font-medium">Ask AI</span>
</Button>
<Button
variant="ghost"
size="xs"
className="gap-1.5 text-muted-foreground hover:text-foreground"
onClick={handleNewLookup}
onClick={handleNewQuery}
disabled={!activeConnectionId}
>
<Search className="h-3.5 w-3.5" />
<span className="text-xs font-medium">Lookup</span>
<Plus className="h-3.5 w-3.5" />
<span className="text-xs font-medium">New Query</span>
</Button>
<div className="flex-1" />
-295
View File
@@ -1,295 +0,0 @@
import { useState, useCallback, useMemo } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
import {
Command,
CommandEmpty,
CommandGroup,
CommandInput,
CommandItem,
CommandList,
} from "@/components/ui/command";
import { Search, Loader2, Check, ChevronsUpDown } from "lucide-react";
import { useEntityLookup } from "@/hooks/use-entity-lookup";
import { useConnections } from "@/hooks/use-connections";
import { useDatabases } from "@/hooks/use-schema";
import { LookupResultGroup } from "./LookupResultGroup";
import type { ConnectionConfig } from "@/types";
interface Props {
connectionId: string;
}
export function EntityLookupPanel({ connectionId }: Props) {
const [columnName, setColumnName] = useState("");
const [value, setValue] = useState("");
const [selectedDbs, setSelectedDbs] = useState<string[]>([]);
const [dbPickerOpen, setDbPickerOpen] = useState(false);
const { data: connections } = useConnections();
const { data: allDatabases } = useDatabases(connectionId);
const activeConn = connections?.find((c) => c.id === connectionId);
const { search, result, error, isSearching, progress } = useEntityLookup();
const handleSearch = useCallback(() => {
if (!columnName.trim() || !value.trim() || !activeConn) return;
const config: ConnectionConfig = { ...activeConn };
search({
config,
columnName: columnName.trim(),
value: value.trim(),
lookupId: crypto.randomUUID(),
databases: selectedDbs.length > 0 ? selectedDbs : undefined,
});
}, [columnName, value, activeConn, selectedDbs, search]);
const handleKeyDown = useCallback(
(e: React.KeyboardEvent) => {
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
handleSearch();
}
},
[handleSearch]
);
const toggleDb = useCallback((db: string) => {
setSelectedDbs((prev) =>
prev.includes(db) ? prev.filter((d) => d !== db) : [...prev, db]
);
}, []);
const progressPercent = useMemo(() => {
if (!progress || progress.total === 0) return 0;
return Math.round((progress.completed / progress.total) * 100);
}, [progress]);
const matchedDbs = useMemo(
() => result?.databases.filter((d) => d.tables.length > 0) ?? [],
[result]
);
const errorDbs = useMemo(
() =>
result?.databases.filter(
(d) => d.error && d.tables.length === 0
) ?? [],
[result]
);
const emptyDbs = useMemo(
() =>
result?.databases.filter(
(d) => !d.error && d.tables.length === 0
) ?? [],
[result]
);
return (
<div className="flex h-full flex-col">
{/* Search form */}
<div className="flex flex-wrap items-center gap-2 border-b px-4 py-3">
<div className="flex items-center gap-2">
<label className="text-xs font-medium text-muted-foreground">
Column:
</label>
<Input
placeholder="carrier_id"
value={columnName}
onChange={(e) => setColumnName(e.target.value)}
onKeyDown={handleKeyDown}
className="h-8 w-44"
disabled={isSearching}
/>
</div>
<div className="flex items-center gap-2">
<label className="text-xs font-medium text-muted-foreground">
Value:
</label>
<Input
placeholder="123"
value={value}
onChange={(e) => setValue(e.target.value)}
onKeyDown={handleKeyDown}
className="h-8 w-44"
disabled={isSearching}
/>
</div>
{/* Database picker */}
<Popover open={dbPickerOpen} onOpenChange={setDbPickerOpen}>
<PopoverTrigger asChild>
<Button
variant="outline"
size="sm"
className="h-8 gap-1 text-xs"
disabled={isSearching}
>
<ChevronsUpDown className="h-3 w-3" />
{selectedDbs.length === 0
? `All (${allDatabases?.length ?? "..."})`
: `${selectedDbs.length} selected`}
</Button>
</PopoverTrigger>
<PopoverContent className="w-56 p-0" align="start">
<Command>
<CommandInput placeholder="Filter databases..." />
<CommandList>
<CommandEmpty>No databases found.</CommandEmpty>
<CommandGroup>
{allDatabases?.map((db) => (
<CommandItem
key={db}
value={db}
onSelect={() => toggleDb(db)}
>
<Check
className={`mr-2 h-4 w-4 ${
selectedDbs.includes(db)
? "opacity-100"
: "opacity-0"
}`}
/>
{db}
</CommandItem>
))}
</CommandGroup>
</CommandList>
</Command>
</PopoverContent>
</Popover>
{selectedDbs.length > 0 && (
<Button
variant="ghost"
size="sm"
className="h-8 text-xs"
onClick={() => setSelectedDbs([])}
>
Clear
</Button>
)}
<Button
size="sm"
className="h-8 gap-1"
onClick={handleSearch}
disabled={isSearching || !columnName.trim() || !value.trim()}
>
{isSearching ? (
<Loader2 className="h-3.5 w-3.5 animate-spin" />
) : (
<Search className="h-3.5 w-3.5" />
)}
Search
</Button>
</div>
{/* Selected databases badges */}
{selectedDbs.length > 0 && (
<div className="flex flex-wrap gap-1 border-b px-4 py-1.5">
{selectedDbs.map((db) => (
<Badge
key={db}
variant="secondary"
className="cursor-pointer text-xs"
onClick={() => toggleDb(db)}
>
{db} &times;
</Badge>
))}
</div>
)}
{/* Progress */}
{isSearching && progress && (
<div className="border-b px-4 py-2">
<div className="flex items-center gap-2 text-xs text-muted-foreground">
<Loader2 className="h-3 w-3 animate-spin" />
Searching... {progress.completed}/{progress.total} databases
</div>
<div className="mt-1 h-1.5 w-full overflow-hidden rounded-full bg-muted">
<div
className="h-full rounded-full bg-primary transition-all"
style={{ width: `${progressPercent}%` }}
/>
</div>
</div>
)}
{/* Error */}
{error && (
<div className="border-b px-4 py-2 text-sm text-destructive">
{error}
</div>
)}
{/* Results */}
{result && (
<div className="flex-1 overflow-hidden">
{/* Summary */}
<div className="border-b px-4 py-2 text-xs text-muted-foreground">
Found{" "}
<span className="font-medium text-foreground">
{result.total_rows_found}
</span>{" "}
row{result.total_rows_found !== 1 && "s"} in{" "}
<span className="font-medium text-foreground">
{result.total_tables_matched}
</span>{" "}
table{result.total_tables_matched !== 1 && "s"} across{" "}
<span className="font-medium text-foreground">
{result.total_databases_searched}
</span>{" "}
database{result.total_databases_searched !== 1 && "s"} in{" "}
{(result.total_time_ms / 1000).toFixed(1)}s
</div>
<div className="h-full overflow-auto">
<div className="flex flex-col gap-2 p-4">
{matchedDbs.map((dbResult) => (
<LookupResultGroup
key={dbResult.database}
dbResult={dbResult}
/>
))}
{errorDbs.map((dbResult) => (
<LookupResultGroup
key={dbResult.database}
dbResult={dbResult}
/>
))}
{emptyDbs.length > 0 && (
<div className="rounded-md border px-3 py-2 text-xs text-muted-foreground">
{emptyDbs.length} database{emptyDbs.length !== 1 && "s"} with
no matches:{" "}
{emptyDbs.map((d) => d.database).join(", ")}
</div>
)}
</div>
</div>
</div>
)}
{/* Empty state */}
{!result && !isSearching && !error && (
<div className="flex flex-1 items-center justify-center">
<div className="text-center text-sm text-muted-foreground">
<Search className="mx-auto mb-2 h-8 w-8 opacity-30" />
<p>Search for a column value across all databases</p>
<p className="mt-1 text-xs">
Enter a column name and value, then press Search
</p>
</div>
</div>
)}
</div>
);
}
-116
View File
@@ -1,116 +0,0 @@
import { useState } from "react";
import {
ChevronDown,
ChevronRight,
AlertCircle,
Database,
} from "lucide-react";
import { ResultsTable } from "@/components/results/ResultsTable";
import type { LookupDatabaseResult } from "@/types";
interface Props {
dbResult: LookupDatabaseResult;
}
export function LookupResultGroup({ dbResult }: Props) {
const [expanded, setExpanded] = useState(dbResult.tables.length > 0);
const [expandedTables, setExpandedTables] = useState<Set<string>>(
() => new Set(dbResult.tables.map((t) => `${t.schema}.${t.table}`))
);
const totalRows = dbResult.tables.reduce((s, t) => s + t.row_count, 0);
const hasError = !!dbResult.error;
const hasMatches = dbResult.tables.length > 0;
const toggleTable = (key: string) => {
setExpandedTables((prev) => {
const next = new Set(prev);
if (next.has(key)) next.delete(key);
else next.add(key);
return next;
});
};
return (
<div className="border rounded-md">
<button
className="flex w-full items-center gap-2 px-3 py-2 text-left text-sm hover:bg-accent/50"
onClick={() => setExpanded(!expanded)}
>
{expanded ? (
<ChevronDown className="h-4 w-4 shrink-0" />
) : (
<ChevronRight className="h-4 w-4 shrink-0" />
)}
<Database className="h-4 w-4 shrink-0 text-muted-foreground" />
<span className="font-medium">{dbResult.database}</span>
{hasMatches && (
<span className="text-xs text-muted-foreground">
{dbResult.tables.length} table{dbResult.tables.length !== 1 && "s"},{" "}
{totalRows} row{totalRows !== 1 && "s"}
</span>
)}
{hasError && (
<span className="flex items-center gap-1 text-xs text-destructive">
<AlertCircle className="h-3 w-3" />
{dbResult.error}
</span>
)}
{!hasError && !hasMatches && (
<span className="text-xs text-muted-foreground">no matches</span>
)}
<span className="ml-auto text-xs text-muted-foreground">
{dbResult.search_time_ms}ms
</span>
</button>
{expanded && hasMatches && (
<div className="border-t">
{dbResult.tables.map((table) => {
const key = `${table.schema}.${table.table}`;
const isOpen = expandedTables.has(key);
return (
<div key={key} className="border-b last:border-b-0">
<button
className="flex w-full items-center gap-2 px-5 py-1.5 text-left text-xs hover:bg-accent/50"
onClick={() => toggleTable(key)}
>
{isOpen ? (
<ChevronDown className="h-3 w-3 shrink-0" />
) : (
<ChevronRight className="h-3 w-3 shrink-0" />
)}
<span className="font-medium">
{table.schema}.{table.table}
</span>
<span className="text-muted-foreground">
({table.row_count} row{table.row_count !== 1 && "s"}
{table.total_count > table.row_count &&
`, ${table.total_count} total`}
)
</span>
<span className="text-muted-foreground">
[{table.column_type}]
</span>
</button>
{isOpen && table.columns.length > 0 && (
<div className="h-[200px] overflow-auto border-t">
<ResultsTable
columns={table.columns}
types={table.types}
rows={table.rows as unknown[][]}
/>
</div>
)}
</div>
);
})}
</div>
)}
</div>
);
}
-348
View File
@@ -1,348 +0,0 @@
import { useState } from "react";
import {
useDatabaseInfo,
useRoles,
useDropDatabase,
useDropRole,
} from "@/hooks/use-management";
import { useAppStore } from "@/stores/app-store";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { CreateDatabaseDialog } from "./CreateDatabaseDialog";
import { CreateRoleDialog } from "./CreateRoleDialog";
import { AlterRoleDialog } from "./AlterRoleDialog";
import { toast } from "sonner";
import {
Plus,
Pencil,
Trash2,
ChevronDown,
ChevronRight,
HardDrive,
Users,
Activity,
Loader2,
} from "lucide-react";
import { DockerContainersList } from "@/components/docker/DockerContainersList";
import type { Tab, RoleInfo } from "@/types";
export function AdminPanel() {
const { activeConnectionId, currentDatabase, readOnlyMap, addTab } = useAppStore();
if (!activeConnectionId) {
return (
<div className="p-4 text-sm text-muted-foreground">
Connect to a database to manage it.
</div>
);
}
const isReadOnly = readOnlyMap[activeConnectionId] ?? true;
return (
<div className="flex h-full flex-col overflow-y-auto">
<DatabasesSection
connectionId={activeConnectionId}
currentDatabase={currentDatabase}
isReadOnly={isReadOnly}
/>
<RolesSection
connectionId={activeConnectionId}
isReadOnly={isReadOnly}
onOpenRoleManager={() => {
const tab: Tab = {
id: crypto.randomUUID(),
type: "roles",
title: "Roles & Users",
connectionId: activeConnectionId,
database: currentDatabase ?? undefined,
};
addTab(tab);
}}
/>
<SessionsSection
connectionId={activeConnectionId}
onOpenSessions={() => {
const tab: Tab = {
id: crypto.randomUUID(),
type: "sessions",
title: "Active Sessions",
connectionId: activeConnectionId,
database: currentDatabase ?? undefined,
};
addTab(tab);
}}
/>
<DockerContainersList />
</div>
);
}
function DatabasesSection({
connectionId,
currentDatabase,
isReadOnly,
}: {
connectionId: string;
currentDatabase: string | null;
isReadOnly: boolean;
}) {
const [expanded, setExpanded] = useState(true);
const [createOpen, setCreateOpen] = useState(false);
const { data: databases, isLoading } = useDatabaseInfo(connectionId);
const dropMutation = useDropDatabase();
const handleDrop = (name: string) => {
if (name === currentDatabase) {
toast.error("Cannot drop the active database");
return;
}
if (!confirm(`Drop database "${name}"? This will terminate all connections and cannot be undone.`)) {
return;
}
dropMutation.mutate(
{ connectionId, name },
{
onSuccess: () => toast.success(`Database "${name}" dropped`),
onError: (err) => toast.error("Failed to drop database", { description: String(err) }),
}
);
};
return (
<div className="border-b">
<div
className="flex items-center gap-1 px-3 py-2 cursor-pointer select-none hover:bg-accent/50"
onClick={() => setExpanded(!expanded)}
>
{expanded ? (
<ChevronDown className="h-3.5 w-3.5 text-muted-foreground" />
) : (
<ChevronRight className="h-3.5 w-3.5 text-muted-foreground" />
)}
<HardDrive className="h-3.5 w-3.5 text-muted-foreground" />
<span className="text-xs font-semibold flex-1">Databases</span>
{!isReadOnly && (
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
onClick={(e) => {
e.stopPropagation();
setCreateOpen(true);
}}
title="Create Database"
>
<Plus className="h-3 w-3" />
</Button>
)}
</div>
{expanded && (
<div className="pb-1">
{isLoading && (
<div className="px-3 py-2 text-xs text-muted-foreground flex items-center gap-1">
<Loader2 className="h-3 w-3 animate-spin" /> Loading...
</div>
)}
{databases?.map((db) => (
<div
key={db.name}
className={`group flex items-center gap-2 px-6 py-1 text-xs hover:bg-accent/50 ${
db.name === currentDatabase ? "text-primary" : ""
}`}
>
<span className="truncate flex-1 font-medium">{db.name}</span>
<span className="text-[10px] text-muted-foreground shrink-0">{db.size}</span>
{db.name === currentDatabase && (
<Badge variant="outline" className="text-[9px] px-1 py-0 shrink-0">
active
</Badge>
)}
{!isReadOnly && db.name !== currentDatabase && (
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0 opacity-0 group-hover:opacity-100 text-destructive hover:text-destructive shrink-0"
onClick={() => handleDrop(db.name)}
title="Drop"
>
<Trash2 className="h-3 w-3" />
</Button>
)}
</div>
))}
</div>
)}
<CreateDatabaseDialog
open={createOpen}
onOpenChange={setCreateOpen}
connectionId={connectionId}
/>
</div>
);
}
function RolesSection({
connectionId,
isReadOnly,
onOpenRoleManager,
}: {
connectionId: string;
isReadOnly: boolean;
onOpenRoleManager: () => void;
}) {
const [expanded, setExpanded] = useState(true);
const [createOpen, setCreateOpen] = useState(false);
const [alterRole, setAlterRole] = useState<RoleInfo | null>(null);
const { data: roles, isLoading } = useRoles(connectionId);
const dropMutation = useDropRole();
const handleDrop = (name: string) => {
if (!confirm(`Drop role "${name}"? This cannot be undone.`)) return;
dropMutation.mutate(
{ connectionId, name },
{
onSuccess: () => toast.success(`Role "${name}" dropped`),
onError: (err) => toast.error("Failed to drop role", { description: String(err) }),
}
);
};
return (
<div className="border-b">
<div
className="flex items-center gap-1 px-3 py-2 cursor-pointer select-none hover:bg-accent/50"
onClick={() => setExpanded(!expanded)}
>
{expanded ? (
<ChevronDown className="h-3.5 w-3.5 text-muted-foreground" />
) : (
<ChevronRight className="h-3.5 w-3.5 text-muted-foreground" />
)}
<Users className="h-3.5 w-3.5 text-muted-foreground" />
<span className="text-xs font-semibold flex-1">Roles</span>
<div className="flex gap-0.5" onClick={(e) => e.stopPropagation()}>
<Button
variant="ghost"
size="sm"
className="h-5 px-1 text-[10px]"
onClick={onOpenRoleManager}
title="Open Role Manager"
>
Manager
</Button>
{!isReadOnly && (
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
onClick={() => setCreateOpen(true)}
title="Create Role"
>
<Plus className="h-3 w-3" />
</Button>
)}
</div>
</div>
{expanded && (
<div className="pb-1">
{isLoading && (
<div className="px-3 py-2 text-xs text-muted-foreground flex items-center gap-1">
<Loader2 className="h-3 w-3 animate-spin" /> Loading...
</div>
)}
{roles?.map((role) => (
<div
key={role.name}
className="group flex items-center gap-1.5 px-6 py-1 text-xs hover:bg-accent/50"
>
<span className="truncate flex-1 font-medium">{role.name}</span>
<div className="flex gap-0.5 shrink-0">
{role.can_login && (
<Badge variant="secondary" className="text-[9px] px-1 py-0">
LOGIN
</Badge>
)}
{role.is_superuser && (
<Badge variant="default" className="text-[9px] px-1 py-0 bg-amber-600 hover:bg-amber-600">
SUPER
</Badge>
)}
</div>
<div className="flex gap-0.5 opacity-0 group-hover:opacity-100 shrink-0">
{!isReadOnly && (
<>
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
onClick={() => setAlterRole(role)}
title="Edit"
>
<Pencil className="h-3 w-3" />
</Button>
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0 text-destructive hover:text-destructive"
onClick={() => handleDrop(role.name)}
title="Drop"
>
<Trash2 className="h-3 w-3" />
</Button>
</>
)}
</div>
</div>
))}
</div>
)}
<CreateRoleDialog
open={createOpen}
onOpenChange={setCreateOpen}
connectionId={connectionId}
/>
<AlterRoleDialog
open={!!alterRole}
onOpenChange={(open) => !open && setAlterRole(null)}
connectionId={connectionId}
role={alterRole}
/>
</div>
);
}
function SessionsSection({
connectionId,
onOpenSessions,
}: {
connectionId: string;
onOpenSessions: () => void;
}) {
return (
<div className="border-b">
<div className="flex items-center gap-1 px-3 py-2">
<Activity className="h-3.5 w-3.5 text-muted-foreground" />
<span className="text-xs font-semibold flex-1">Sessions</span>
<Button
variant="ghost"
size="sm"
className="h-5 px-1 text-[10px]"
onClick={onOpenSessions}
title="View active sessions"
>
View Sessions
</Button>
</div>
<div className="px-6 pb-2 text-xs text-muted-foreground">
Monitor active database connections and running queries.
{/* The connectionId is used by parent to open the sessions tab */}
<span className="hidden">{connectionId}</span>
</div>
</div>
);
}
@@ -1,185 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { useAlterRole } from "@/hooks/use-management";
import { toast } from "sonner";
import { Loader2 } from "lucide-react";
import type { RoleInfo } from "@/types";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
role: RoleInfo | null;
}
export function AlterRoleDialog({ open, onOpenChange, connectionId, role }: Props) {
const [password, setPassword] = useState("");
const [login, setLogin] = useState(false);
const [superuser, setSuperuser] = useState(false);
const [createdb, setCreatedb] = useState(false);
const [createrole, setCreaterole] = useState(false);
const [inherit, setInherit] = useState(true);
const [replication, setReplication] = useState(false);
const [connectionLimit, setConnectionLimit] = useState(-1);
const [validUntil, setValidUntil] = useState("");
const [renameTo, setRenameTo] = useState("");
const alterMutation = useAlterRole();
const [prev, setPrev] = useState<{ open: boolean; role: typeof role }>({ open: false, role: null });
if (open !== prev.open || role !== prev.role) {
setPrev({ open, role });
if (open && role) {
setPassword("");
setLogin(role.can_login);
setSuperuser(role.is_superuser);
setCreatedb(role.can_create_db);
setCreaterole(role.can_create_role);
setInherit(role.inherit);
setReplication(role.is_replication);
setConnectionLimit(role.connection_limit);
setValidUntil(role.valid_until ?? "");
setRenameTo("");
}
}
if (!role) return null;
const handleAlter = () => {
const params: Record<string, unknown> = { name: role.name };
if (password) params.password = password;
if (login !== role.can_login) params.login = login;
if (superuser !== role.is_superuser) params.superuser = superuser;
if (createdb !== role.can_create_db) params.createdb = createdb;
if (createrole !== role.can_create_role) params.createrole = createrole;
if (inherit !== role.inherit) params.inherit = inherit;
if (replication !== role.is_replication) params.replication = replication;
if (connectionLimit !== role.connection_limit) params.connection_limit = connectionLimit;
if (validUntil !== (role.valid_until ?? "")) params.valid_until = validUntil || undefined;
if (renameTo.trim()) params.rename_to = renameTo.trim();
alterMutation.mutate(
{
connectionId,
params: params as {
name: string;
password?: string;
login?: boolean;
superuser?: boolean;
createdb?: boolean;
createrole?: boolean;
inherit?: boolean;
replication?: boolean;
connection_limit?: number;
valid_until?: string;
rename_to?: string;
},
},
{
onSuccess: () => {
toast.success(`Role "${role.name}" updated`);
onOpenChange(false);
},
onError: (err) => {
toast.error("Failed to alter role", { description: String(err) });
},
}
);
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[480px]">
<DialogHeader>
<DialogTitle>Alter Role: {role.name}</DialogTitle>
</DialogHeader>
<div className="grid gap-3 py-4">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Rename To</label>
<Input
className="col-span-3"
value={renameTo}
onChange={(e) => setRenameTo(e.target.value)}
placeholder={role.name}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Password</label>
<Input
className="col-span-3"
type="password"
value={password}
onChange={(e) => setPassword(e.target.value)}
placeholder="Leave empty to keep unchanged"
/>
</div>
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Privileges</label>
<div className="col-span-3 grid grid-cols-2 gap-2">
{([
["LOGIN", login, setLogin],
["SUPERUSER", superuser, setSuperuser],
["CREATEDB", createdb, setCreatedb],
["CREATEROLE", createrole, setCreaterole],
["INHERIT", inherit, setInherit],
["REPLICATION", replication, setReplication],
] as const).map(([label, value, setter]) => (
<label key={label} className="flex items-center gap-2 text-sm cursor-pointer">
<input
type="checkbox"
checked={value}
onChange={(e) => (setter as (v: boolean) => void)(e.target.checked)}
className="rounded border-input"
/>
{label}
</label>
))}
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Conn Limit</label>
<Input
className="col-span-3"
type="number"
value={connectionLimit}
onChange={(e) => setConnectionLimit(parseInt(e.target.value) || -1)}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Valid Until</label>
<Input
className="col-span-3"
type="datetime-local"
value={validUntil}
onChange={(e) => setValidUntil(e.target.value)}
/>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>
Cancel
</Button>
<Button onClick={handleAlter} disabled={alterMutation.isPending}>
{alterMutation.isPending && (
<Loader2 className="mr-2 h-4 w-4 animate-spin" />
)}
Save Changes
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -1,162 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useCreateDatabase, useRoles } from "@/hooks/use-management";
import { toast } from "sonner";
import { Loader2 } from "lucide-react";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
}
export function CreateDatabaseDialog({ open, onOpenChange, connectionId }: Props) {
const [name, setName] = useState("");
const [owner, setOwner] = useState("__default__");
const [template, setTemplate] = useState("__default__");
const [encoding, setEncoding] = useState("UTF8");
const [connectionLimit, setConnectionLimit] = useState(-1);
const { data: roles } = useRoles(open ? connectionId : null);
const createMutation = useCreateDatabase();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setName("");
setOwner("__default__");
setTemplate("__default__");
setEncoding("UTF8");
setConnectionLimit(-1);
}
}
const handleCreate = () => {
if (!name.trim()) {
toast.error("Database name is required");
return;
}
createMutation.mutate(
{
connectionId,
params: {
name: name.trim(),
owner: owner === "__default__" ? undefined : owner,
template: template === "__default__" ? undefined : template,
encoding,
connection_limit: connectionLimit,
},
},
{
onSuccess: () => {
toast.success(`Database "${name}" created`);
onOpenChange(false);
},
onError: (err) => {
toast.error("Failed to create database", { description: String(err) });
},
}
);
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[480px]">
<DialogHeader>
<DialogTitle>Create Database</DialogTitle>
</DialogHeader>
<div className="grid gap-3 py-4">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Name</label>
<Input
className="col-span-3"
value={name}
onChange={(e) => setName(e.target.value)}
placeholder="my_database"
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Owner</label>
<Select value={owner} onValueChange={setOwner}>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="__default__">Default</SelectItem>
{roles?.map((r) => (
<SelectItem key={r.name} value={r.name}>
{r.name}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Template</label>
<Select value={template} onValueChange={setTemplate}>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="__default__">Default</SelectItem>
<SelectItem value="template0">template0</SelectItem>
<SelectItem value="template1">template1</SelectItem>
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Encoding</label>
<Select value={encoding} onValueChange={setEncoding}>
<SelectTrigger className="col-span-3">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="UTF8">UTF8</SelectItem>
<SelectItem value="LATIN1">LATIN1</SelectItem>
<SelectItem value="SQL_ASCII">SQL_ASCII</SelectItem>
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Conn Limit</label>
<Input
className="col-span-3"
type="number"
value={connectionLimit}
onChange={(e) => setConnectionLimit(parseInt(e.target.value) || -1)}
/>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>
Cancel
</Button>
<Button onClick={handleCreate} disabled={createMutation.isPending}>
{createMutation.isPending && (
<Loader2 className="mr-2 h-4 w-4 animate-spin" />
)}
Create
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -1,205 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { useCreateRole, useRoles } from "@/hooks/use-management";
import { toast } from "sonner";
import { Loader2 } from "lucide-react";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
}
export function CreateRoleDialog({ open, onOpenChange, connectionId }: Props) {
const [name, setName] = useState("");
const [password, setPassword] = useState("");
const [login, setLogin] = useState(true);
const [superuser, setSuperuser] = useState(false);
const [createdb, setCreatedb] = useState(false);
const [createrole, setCreaterole] = useState(false);
const [inherit, setInherit] = useState(true);
const [replication, setReplication] = useState(false);
const [connectionLimit, setConnectionLimit] = useState(-1);
const [validUntil, setValidUntil] = useState("");
const [inRoles, setInRoles] = useState<string[]>([]);
const { data: roles } = useRoles(open ? connectionId : null);
const createMutation = useCreateRole();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setName("");
setPassword("");
setLogin(true);
setSuperuser(false);
setCreatedb(false);
setCreaterole(false);
setInherit(true);
setReplication(false);
setConnectionLimit(-1);
setValidUntil("");
setInRoles([]);
}
}
const handleCreate = () => {
if (!name.trim()) {
toast.error("Role name is required");
return;
}
createMutation.mutate(
{
connectionId,
params: {
name: name.trim(),
password: password || undefined,
login,
superuser,
createdb,
createrole,
inherit,
replication,
connection_limit: connectionLimit,
valid_until: validUntil || undefined,
in_roles: inRoles,
},
},
{
onSuccess: () => {
toast.success(`Role "${name}" created`);
onOpenChange(false);
},
onError: (err) => {
toast.error("Failed to create role", { description: String(err) });
},
}
);
};
const toggleInRole = (roleName: string) => {
setInRoles((prev) =>
prev.includes(roleName)
? prev.filter((r) => r !== roleName)
: [...prev, roleName]
);
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[480px]">
<DialogHeader>
<DialogTitle>Create Role</DialogTitle>
</DialogHeader>
<div className="grid gap-3 py-4">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Name</label>
<Input
className="col-span-3"
value={name}
onChange={(e) => setName(e.target.value)}
placeholder="my_role"
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Password</label>
<Input
className="col-span-3"
type="password"
value={password}
onChange={(e) => setPassword(e.target.value)}
placeholder="Optional"
/>
</div>
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Privileges</label>
<div className="col-span-3 grid grid-cols-2 gap-2">
{([
["LOGIN", login, setLogin],
["SUPERUSER", superuser, setSuperuser],
["CREATEDB", createdb, setCreatedb],
["CREATEROLE", createrole, setCreaterole],
["INHERIT", inherit, setInherit],
["REPLICATION", replication, setReplication],
] as const).map(([label, value, setter]) => (
<label key={label} className="flex items-center gap-2 text-sm cursor-pointer">
<input
type="checkbox"
checked={value}
onChange={(e) => (setter as (v: boolean) => void)(e.target.checked)}
className="rounded border-input"
/>
{label}
</label>
))}
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Conn Limit</label>
<Input
className="col-span-3"
type="number"
value={connectionLimit}
onChange={(e) => setConnectionLimit(parseInt(e.target.value) || -1)}
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Valid Until</label>
<Input
className="col-span-3"
type="datetime-local"
value={validUntil}
onChange={(e) => setValidUntil(e.target.value)}
/>
</div>
{roles && roles.length > 0 && (
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Member Of</label>
<div className="col-span-3 flex flex-wrap gap-1.5">
{roles.map((r) => (
<button
key={r.name}
type="button"
className={`rounded-full border px-2.5 py-0.5 text-xs transition-colors ${
inRoles.includes(r.name)
? "border-primary bg-primary text-primary-foreground"
: "border-border text-muted-foreground hover:text-foreground"
}`}
onClick={() => toggleInRole(r.name)}
>
{r.name}
</button>
))}
</div>
</div>
)}
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>
Cancel
</Button>
<Button onClick={handleCreate} disabled={createMutation.isPending}>
{createMutation.isPending && (
<Loader2 className="mr-2 h-4 w-4 animate-spin" />
)}
Create
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -1,245 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import {
useRoles,
useGrantRevoke,
useTablePrivileges,
} from "@/hooks/use-management";
import { toast } from "sonner";
import { Loader2 } from "lucide-react";
import { Badge } from "@/components/ui/badge";
const PRIVILEGE_OPTIONS: Record<string, string[]> = {
TABLE: ["SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE", "REFERENCES", "TRIGGER", "ALL"],
SCHEMA: ["USAGE", "CREATE"],
DATABASE: ["CONNECT", "CREATE", "TEMPORARY"],
};
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
objectType: string;
objectName: string;
schema?: string;
table?: string;
}
export function GrantRevokeDialog({
open,
onOpenChange,
connectionId,
objectType,
objectName,
schema,
table,
}: Props) {
const [action, setAction] = useState("GRANT");
const [roleName, setRoleName] = useState("");
const [privileges, setPrivileges] = useState<string[]>([]);
const [withGrantOption, setWithGrantOption] = useState(false);
const { data: roles } = useRoles(open ? connectionId : null);
const { data: existingPrivileges } = useTablePrivileges(
open && objectType === "TABLE" ? connectionId : null,
schema ?? null,
table ?? null
);
const grantRevokeMutation = useGrantRevoke();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setAction("GRANT");
setRoleName("");
setPrivileges([]);
setWithGrantOption(false);
}
}
const availablePrivileges = PRIVILEGE_OPTIONS[objectType.toUpperCase()] ?? PRIVILEGE_OPTIONS.TABLE;
const togglePrivilege = (priv: string) => {
setPrivileges((prev) =>
prev.includes(priv)
? prev.filter((p) => p !== priv)
: [...prev, priv]
);
};
const handleSubmit = () => {
if (!roleName) {
toast.error("Please select a role");
return;
}
if (privileges.length === 0) {
toast.error("Please select at least one privilege");
return;
}
grantRevokeMutation.mutate(
{
connectionId,
params: {
action,
privileges,
object_type: objectType,
object_name: objectName,
role_name: roleName,
with_grant_option: withGrantOption,
},
},
{
onSuccess: () => {
toast.success(
`${action === "GRANT" ? "Granted" : "Revoked"} privileges on ${objectName}`
);
onOpenChange(false);
},
onError: (err) => {
toast.error("Operation failed", { description: String(err) });
},
}
);
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[520px]">
<DialogHeader>
<DialogTitle>Manage Privileges</DialogTitle>
</DialogHeader>
<div className="grid gap-3 py-4">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Object</label>
<div className="col-span-3 text-sm">
<Badge variant="outline">{objectType}</Badge>{" "}
<span className="font-medium">{objectName}</span>
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Action</label>
<div className="col-span-3 flex gap-2">
<Button
size="sm"
variant={action === "GRANT" ? "default" : "outline"}
onClick={() => setAction("GRANT")}
>
Grant
</Button>
<Button
size="sm"
variant={action === "REVOKE" ? "default" : "outline"}
onClick={() => setAction("REVOKE")}
>
Revoke
</Button>
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Role</label>
<Select value={roleName} onValueChange={setRoleName}>
<SelectTrigger className="col-span-3">
<SelectValue placeholder="Select role..." />
</SelectTrigger>
<SelectContent>
{roles?.map((r) => (
<SelectItem key={r.name} value={r.name}>
{r.name}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Privileges</label>
<div className="col-span-3 flex flex-wrap gap-1.5">
{availablePrivileges.map((priv) => (
<button
key={priv}
type="button"
className={`rounded-full border px-2.5 py-0.5 text-xs transition-colors ${
privileges.includes(priv)
? "border-primary bg-primary text-primary-foreground"
: "border-border text-muted-foreground hover:text-foreground"
}`}
onClick={() => togglePrivilege(priv)}
>
{priv}
</button>
))}
</div>
</div>
{action === "GRANT" && (
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Options</label>
<label className="col-span-3 flex items-center gap-2 text-sm cursor-pointer">
<input
type="checkbox"
checked={withGrantOption}
onChange={(e) => setWithGrantOption(e.target.checked)}
className="rounded border-input"
/>
WITH GRANT OPTION
</label>
</div>
)}
{existingPrivileges && existingPrivileges.length > 0 && (
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Current</label>
<div className="col-span-3 max-h-32 overflow-y-auto rounded border p-2">
<div className="space-y-1">
{existingPrivileges.map((p, i) => (
<div key={i} className="flex items-center gap-2 text-xs">
<span className="font-medium">{p.grantee}</span>
<Badge variant="secondary" className="text-[10px]">
{p.privilege_type}
</Badge>
{p.is_grantable && (
<Badge variant="outline" className="text-[10px]">
GRANTABLE
</Badge>
)}
</div>
))}
</div>
</div>
</div>
)}
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>
Cancel
</Button>
<Button onClick={handleSubmit} disabled={grantRevokeMutation.isPending}>
{grantRevokeMutation.isPending && (
<Loader2 className="mr-2 h-4 w-4 animate-spin" />
)}
{action === "GRANT" ? "Grant" : "Revoke"}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
@@ -1,306 +0,0 @@
import { useState } from "react";
import { useRoles, useDropRole, useManageRoleMembership } from "@/hooks/use-management";
import { useAppStore } from "@/stores/app-store";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { CreateRoleDialog } from "./CreateRoleDialog";
import { AlterRoleDialog } from "./AlterRoleDialog";
import { toast } from "sonner";
import { Plus, Pencil, Trash2, UserPlus, UserMinus, Loader2 } from "lucide-react";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import type { RoleInfo } from "@/types";
interface Props {
connectionId: string;
}
export function RoleManagerView({ connectionId }: Props) {
const { data: roles, isLoading } = useRoles(connectionId);
const readOnlyMap = useAppStore((s) => s.readOnlyMap);
const isReadOnly = readOnlyMap[connectionId] ?? true;
const dropMutation = useDropRole();
const membershipMutation = useManageRoleMembership();
const [createOpen, setCreateOpen] = useState(false);
const [alterRole, setAlterRole] = useState<RoleInfo | null>(null);
const [selectedRole, setSelectedRole] = useState<string | null>(null);
const [memberToAdd, setMemberToAdd] = useState("");
const handleDrop = (name: string) => {
if (!confirm(`Drop role "${name}"? This cannot be undone.`)) return;
dropMutation.mutate(
{ connectionId, name },
{
onSuccess: () => toast.success(`Role "${name}" dropped`),
onError: (err) => toast.error("Failed to drop role", { description: String(err) }),
}
);
};
const handleAddMember = (roleName: string, memberName: string) => {
membershipMutation.mutate(
{
connectionId,
params: { action: "GRANT", role_name: roleName, member_name: memberName },
},
{
onSuccess: () => {
toast.success(`Added "${memberName}" to "${roleName}"`);
setMemberToAdd("");
},
onError: (err) => toast.error("Failed", { description: String(err) }),
}
);
};
const handleRemoveMember = (roleName: string, memberName: string) => {
membershipMutation.mutate(
{
connectionId,
params: { action: "REVOKE", role_name: roleName, member_name: memberName },
},
{
onSuccess: () => toast.success(`Removed "${memberName}" from "${roleName}"`),
onError: (err) => toast.error("Failed", { description: String(err) }),
}
);
};
const selected = roles?.find((r) => r.name === selectedRole);
if (isLoading) {
return (
<div className="flex h-full items-center justify-center">
<Loader2 className="h-6 w-6 animate-spin text-muted-foreground" />
</div>
);
}
return (
<div className="flex h-full flex-col">
<div className="flex items-center justify-between border-b px-4 py-2">
<h2 className="text-sm font-semibold">Roles & Users</h2>
<Button size="sm" onClick={() => setCreateOpen(true)} disabled={isReadOnly}>
<Plus className="mr-1 h-3.5 w-3.5" />
Create Role
</Button>
</div>
<div className="flex flex-1 min-h-0">
<div className="flex-1 overflow-auto">
<table className="w-full text-sm">
<thead className="sticky top-0 bg-card">
<tr className="border-b text-left text-xs text-muted-foreground">
<th className="px-4 py-2 font-medium">Name</th>
<th className="px-4 py-2 font-medium">Login</th>
<th className="px-4 py-2 font-medium">Superuser</th>
<th className="px-4 py-2 font-medium">CreateDB</th>
<th className="px-4 py-2 font-medium">CreateRole</th>
<th className="px-4 py-2 font-medium">Conn Limit</th>
<th className="px-4 py-2 font-medium">Member Of</th>
<th className="px-4 py-2 font-medium w-24">Actions</th>
</tr>
</thead>
<tbody>
{roles?.map((role) => (
<tr
key={role.name}
className={`border-b hover:bg-accent/50 cursor-pointer ${
selectedRole === role.name ? "bg-accent/30" : ""
}`}
onClick={() => setSelectedRole(role.name)}
>
<td className="px-4 py-2 font-medium">{role.name}</td>
<td className="px-4 py-2">
<BoolBadge value={role.can_login} />
</td>
<td className="px-4 py-2">
<BoolBadge value={role.is_superuser} />
</td>
<td className="px-4 py-2">
<BoolBadge value={role.can_create_db} />
</td>
<td className="px-4 py-2">
<BoolBadge value={role.can_create_role} />
</td>
<td className="px-4 py-2 text-muted-foreground">
{role.connection_limit === -1 ? "unlimited" : role.connection_limit}
</td>
<td className="px-4 py-2">
<div className="flex flex-wrap gap-1">
{role.member_of.map((r) => (
<Badge key={r} variant="secondary" className="text-[10px]">
{r}
</Badge>
))}
</div>
</td>
<td className="px-4 py-2">
<div className="flex gap-1">
<Button
variant="ghost"
size="sm"
className="h-6 w-6 p-0"
onClick={(e) => {
e.stopPropagation();
setAlterRole(role);
}}
disabled={isReadOnly}
title="Edit"
>
<Pencil className="h-3 w-3" />
</Button>
<Button
variant="ghost"
size="sm"
className="h-6 w-6 p-0 text-destructive hover:text-destructive"
onClick={(e) => {
e.stopPropagation();
handleDrop(role.name);
}}
disabled={isReadOnly}
title="Drop"
>
<Trash2 className="h-3 w-3" />
</Button>
</div>
</td>
</tr>
))}
</tbody>
</table>
</div>
{selected && (
<div className="w-64 shrink-0 border-l overflow-auto">
<div className="p-3">
<h3 className="text-sm font-semibold mb-3">{selected.name}</h3>
<div className="space-y-3">
<div>
<p className="text-xs text-muted-foreground mb-1">Member Of</p>
<div className="flex flex-wrap gap-1">
{selected.member_of.length === 0 && (
<span className="text-xs text-muted-foreground">None</span>
)}
{selected.member_of.map((r) => (
<div key={r} className="flex items-center gap-0.5">
<Badge variant="secondary" className="text-[10px]">
{r}
</Badge>
{!isReadOnly && (
<button
className="text-muted-foreground hover:text-destructive"
onClick={() => handleRemoveMember(r, selected.name)}
title={`Remove from ${r}`}
>
<UserMinus className="h-3 w-3" />
</button>
)}
</div>
))}
</div>
</div>
<div>
<p className="text-xs text-muted-foreground mb-1">Members</p>
<div className="flex flex-wrap gap-1">
{selected.members.length === 0 && (
<span className="text-xs text-muted-foreground">None</span>
)}
{selected.members.map((m) => (
<div key={m} className="flex items-center gap-0.5">
<Badge variant="secondary" className="text-[10px]">
{m}
</Badge>
{!isReadOnly && (
<button
className="text-muted-foreground hover:text-destructive"
onClick={() => handleRemoveMember(selected.name, m)}
title={`Remove ${m}`}
>
<UserMinus className="h-3 w-3" />
</button>
)}
</div>
))}
</div>
</div>
{!isReadOnly && (
<div>
<p className="text-xs text-muted-foreground mb-1">Add Member</p>
<div className="flex gap-1">
<Select value={memberToAdd} onValueChange={setMemberToAdd}>
<SelectTrigger className="h-7 text-xs flex-1">
<SelectValue placeholder="Select..." />
</SelectTrigger>
<SelectContent>
{roles
?.filter(
(r) =>
r.name !== selected.name &&
!selected.members.includes(r.name)
)
.map((r) => (
<SelectItem key={r.name} value={r.name}>
{r.name}
</SelectItem>
))}
</SelectContent>
</Select>
<Button
size="sm"
className="h-7"
disabled={!memberToAdd || membershipMutation.isPending}
onClick={() => handleAddMember(selected.name, memberToAdd)}
>
<UserPlus className="h-3 w-3" />
</Button>
</div>
</div>
)}
{selected.description && (
<div>
<p className="text-xs text-muted-foreground mb-1">Description</p>
<p className="text-xs">{selected.description}</p>
</div>
)}
</div>
</div>
</div>
)}
</div>
<CreateRoleDialog
open={createOpen}
onOpenChange={setCreateOpen}
connectionId={connectionId}
/>
<AlterRoleDialog
open={!!alterRole}
onOpenChange={(open) => !open && setAlterRole(null)}
connectionId={connectionId}
role={alterRole}
/>
</div>
);
}
function BoolBadge({ value }: { value: boolean }) {
return (
<Badge
variant={value ? "default" : "secondary"}
className={`text-[10px] ${value ? "bg-green-600 hover:bg-green-600" : "text-muted-foreground"}`}
>
{value ? "Yes" : "No"}
</Badge>
);
}
-171
View File
@@ -1,171 +0,0 @@
import {
useSessions,
useCancelQuery,
useTerminateBackend,
} from "@/hooks/use-management";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { toast } from "sonner";
import { Loader2, XCircle, Skull, RefreshCw } from "lucide-react";
import { useQueryClient } from "@tanstack/react-query";
interface Props {
connectionId: string;
}
function getStateBadge(state: string | null) {
if (!state) return null;
const colors: Record<string, string> = {
idle: "bg-green-500/15 text-green-600",
active: "bg-yellow-500/15 text-yellow-600",
"idle in transaction": "bg-orange-500/15 text-orange-600",
disabled: "bg-red-500/15 text-red-600",
};
return (
<Badge variant="outline" className={`text-[9px] px-1 py-0 ${colors[state] ?? ""}`}>
{state}
</Badge>
);
}
function formatDuration(queryStart: string | null): string {
if (!queryStart) return "-";
const start = new Date(queryStart).getTime();
const now = Date.now();
const diffSec = Math.floor((now - start) / 1000);
if (diffSec < 0) return "-";
if (diffSec < 60) return `${diffSec}s`;
if (diffSec < 3600) return `${Math.floor(diffSec / 60)}m ${diffSec % 60}s`;
return `${Math.floor(diffSec / 3600)}h ${Math.floor((diffSec % 3600) / 60)}m`;
}
function getDurationColor(queryStart: string | null, state: string | null): string {
if (state !== "active" || !queryStart) return "";
const diffSec = (Date.now() - new Date(queryStart).getTime()) / 1000;
if (diffSec > 30) return "text-red-500 font-semibold";
if (diffSec > 5) return "text-yellow-500 font-semibold";
return "";
}
export function SessionsView({ connectionId }: Props) {
const { data: sessions, isLoading } = useSessions(connectionId);
const cancelMutation = useCancelQuery();
const terminateMutation = useTerminateBackend();
const queryClient = useQueryClient();
const handleCancel = (pid: number) => {
cancelMutation.mutate(
{ connectionId, pid },
{
onSuccess: () => toast.success(`Cancel signal sent to PID ${pid}`),
onError: (err) => toast.error("Cancel failed", { description: String(err) }),
}
);
};
const handleTerminate = (pid: number) => {
if (!confirm(`Terminate backend PID ${pid}? This will kill the session.`)) return;
terminateMutation.mutate(
{ connectionId, pid },
{
onSuccess: () => toast.success(`Terminate signal sent to PID ${pid}`),
onError: (err) => toast.error("Terminate failed", { description: String(err) }),
}
);
};
if (isLoading) {
return (
<div className="flex h-full items-center justify-center text-muted-foreground">
<Loader2 className="mr-2 h-5 w-5 animate-spin" />
Loading sessions...
</div>
);
}
return (
<div className="flex h-full flex-col">
<div className="flex items-center gap-2 border-b px-3 py-1.5">
<span className="text-xs font-semibold">Active Sessions</span>
<Badge variant="secondary" className="text-[10px]">
{sessions?.length ?? 0}
</Badge>
<span className="text-[10px] text-muted-foreground">Auto-refresh: 5s</span>
<Button
variant="ghost"
size="sm"
className="ml-auto h-6 gap-1 text-xs"
onClick={() => queryClient.invalidateQueries({ queryKey: ["sessions"] })}
>
<RefreshCw className="h-3 w-3" />
Refresh
</Button>
</div>
<div className="flex-1 overflow-auto">
<table className="w-full text-xs">
<thead className="sticky top-0 bg-card border-b">
<tr>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">PID</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">User</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Database</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">State</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Duration</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Wait</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Query</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Client</th>
<th className="px-2 py-1.5 text-left font-medium text-muted-foreground">Actions</th>
</tr>
</thead>
<tbody>
{sessions?.map((s) => (
<tr key={s.pid} className="border-b hover:bg-accent/50">
<td className="px-2 py-1 font-mono">{s.pid}</td>
<td className="px-2 py-1">{s.usename ?? "-"}</td>
<td className="px-2 py-1">{s.datname ?? "-"}</td>
<td className="px-2 py-1">{getStateBadge(s.state)}</td>
<td className={`px-2 py-1 ${getDurationColor(s.query_start, s.state)}`}>
{formatDuration(s.query_start)}
</td>
<td className="px-2 py-1 text-muted-foreground">
{s.wait_event ? `${s.wait_event_type}: ${s.wait_event}` : "-"}
</td>
<td className="px-2 py-1 max-w-xs truncate font-mono" title={s.query ?? ""}>
{s.query ? (s.query.length > 60 ? s.query.slice(0, 60) + "..." : s.query) : "-"}
</td>
<td className="px-2 py-1 text-muted-foreground">{s.client_addr ?? "-"}</td>
<td className="px-2 py-1">
<div className="flex gap-0.5">
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
title="Cancel Query"
onClick={() => handleCancel(s.pid)}
>
<XCircle className="h-3 w-3 text-yellow-500" />
</Button>
<Button
variant="ghost"
size="sm"
className="h-5 w-5 p-0"
title="Terminate Backend"
onClick={() => handleTerminate(s.pid)}
>
<Skull className="h-3 w-3 text-red-500" />
</Button>
</div>
</td>
</tr>
))}
</tbody>
</table>
{(!sessions || sessions.length === 0) && (
<div className="py-8 text-center text-xs text-muted-foreground">
No active sessions
</div>
)}
</div>
</div>
);
}
+140
View File
@@ -0,0 +1,140 @@
import { useEffect, useState } from "react";
import { Brain, RefreshCw, Save } from "lucide-react";
import { toast } from "sonner";
import { Button } from "@/components/ui/button";
import { useAppStore } from "@/stores/app-store";
import { useConnections } from "@/hooks/use-connections";
import { useMemory, useSaveMemory } from "@/hooks/use-memory";
export function MemoryPanel() {
const activeConnectionId = useAppStore((s) => s.activeConnectionId);
const { data: connections } = useConnections();
const activeConn = connections?.find((c) => c.id === activeConnectionId);
const { data: serverContent, isFetching, refetch } = useMemory(activeConnectionId);
const saveMutation = useSaveMemory();
const [draft, setDraft] = useState<string>("");
const [dirty, setDirty] = useState(false);
// Sync local textarea when the active connection changes or server reloads.
// Only overwrite if the user hasn't edited.
useEffect(() => {
if (!dirty) {
setDraft(serverContent ?? "");
}
}, [serverContent, activeConnectionId, dirty]);
if (!activeConnectionId) {
return (
<div className="flex flex-col items-center gap-2 p-6 text-center">
<Brain className="h-8 w-8 text-muted-foreground/20" />
<p className="text-sm text-muted-foreground/60">
Connect to a database to view its memory.
</p>
</div>
);
}
const noteCount = (draft.match(/^## /gm) ?? []).length;
const isEmpty = draft.trim().length === 0;
const handleSave = () => {
saveMutation.mutate(
{ connectionId: activeConnectionId, content: draft },
{
onSuccess: () => {
setDirty(false);
toast.success("Memory saved");
},
onError: (err) => toast.error("Save failed", { description: String(err) }),
}
);
};
const handleReload = () => {
setDirty(false);
refetch();
};
return (
<div className="flex h-full flex-col">
<div className="flex items-center justify-between gap-1 border-b border-border/40 px-3 py-2">
<div className="flex items-center gap-1.5 text-xs text-muted-foreground">
<Brain className="h-3.5 w-3.5 text-primary/70" />
<span className="font-medium">Memory</span>
{activeConn && (
<span className="ml-1 truncate text-muted-foreground/60">· {activeConn.name}</span>
)}
{!isEmpty && (
<span className="ml-1 text-[10px] text-muted-foreground/50">
{noteCount} note{noteCount === 1 ? "" : "s"}
</span>
)}
</div>
<div className="flex items-center gap-1">
<Button
size="icon-xs"
variant="ghost"
onClick={handleReload}
disabled={isFetching}
title="Reload from disk"
className="text-muted-foreground hover:text-foreground"
>
<RefreshCw className={`h-3.5 w-3.5 ${isFetching ? "animate-spin" : ""}`} />
</Button>
<Button
size="xs"
variant={dirty ? "default" : "ghost"}
onClick={handleSave}
disabled={!dirty || saveMutation.isPending}
title="Save"
className="gap-1"
>
<Save className="h-3 w-3" />
<span className="text-xs">Save</span>
</Button>
</div>
</div>
{isEmpty ? (
<div className="flex h-full flex-col items-center justify-center gap-3 px-6 py-8 text-center">
<Brain className="h-7 w-7 text-muted-foreground/30" />
<p className="text-xs text-muted-foreground/70">
No notes yet for this connection.
</p>
<p className="text-[11px] text-muted-foreground/50 max-w-[260px]">
The agent will populate this as it learns about your database. You can also
edit notes here directly — anything you type is loaded into the agent's
context on its next turn.
</p>
<textarea
className="mt-2 h-48 w-full max-w-[360px] resize-none rounded-md border border-border/40 bg-background/50 p-2 font-mono text-[11px] outline-none focus:border-primary/40"
placeholder="# Memory&#10;&#10;## (timestamp)&#10;your note..."
value={draft}
onChange={(e) => {
setDraft(e.target.value);
setDirty(true);
}}
/>
</div>
) : (
<textarea
className="min-h-0 flex-1 resize-none overflow-y-auto bg-background/40 p-3 font-mono text-[11px] leading-relaxed outline-none"
spellCheck={false}
value={draft}
onChange={(e) => {
setDraft(e.target.value);
setDirty(true);
}}
/>
)}
{dirty && (
<div className="border-t border-border/40 px-3 py-1.5 text-[10px] text-amber-500/80">
Unsaved changes
</div>
)}
</div>
);
}
+18 -176
View File
@@ -9,7 +9,6 @@ import {
useSequences,
} from "@/hooks/use-schema";
import { useConnections } from "@/hooks/use-connections";
import { useDropDatabase } from "@/hooks/use-management";
import { useAppStore } from "@/stores/app-store";
import { toast } from "sonner";
import {
@@ -27,12 +26,8 @@ import {
ContextMenu,
ContextMenuContent,
ContextMenuItem,
ContextMenuSeparator,
ContextMenuTrigger,
} from "@/components/ui/context-menu";
import { GrantRevokeDialog } from "@/components/management/GrantRevokeDialog";
import { CloneDatabaseDialog } from "@/components/docker/CloneDatabaseDialog";
import { GenerateDataDialog } from "@/components/data-generator/GenerateDataDialog";
import type { Tab, SchemaObject } from "@/types";
function formatSize(bytes: number): string {
@@ -67,7 +62,6 @@ export function SchemaTree() {
const { data: databases } = useDatabases(activeConnectionId);
const { data: connections } = useConnections();
const switchDbMutation = useSwitchDatabase();
const [cloneTarget, setCloneTarget] = useState<string | null>(null);
if (!activeConnectionId) {
return (
@@ -118,7 +112,6 @@ export function SchemaTree() {
connectionId={activeConnectionId}
onSwitch={() => handleSwitchDb(db)}
isSwitching={switchDbMutation.isPending}
onCloneToDocker={(dbName) => setCloneTarget(dbName)}
onOpenTable={(schema, table) => {
const tab: Tab = {
id: crypto.randomUUID(),
@@ -143,25 +136,8 @@ export function SchemaTree() {
};
addTab(tab);
}}
onViewErd={(schema) => {
const tab: Tab = {
id: crypto.randomUUID(),
type: "erd",
title: `${schema} (ER Diagram)`,
connectionId: activeConnectionId,
database: currentDatabase ?? undefined,
schema,
};
addTab(tab);
}}
/>
))}
<CloneDatabaseDialog
open={cloneTarget !== null}
onOpenChange={(open) => { if (!open) setCloneTarget(null); }}
connectionId={activeConnectionId}
database={cloneTarget ?? ""}
/>
</div>
);
}
@@ -172,25 +148,18 @@ function DatabaseNode({
connectionId,
onSwitch,
isSwitching,
onCloneToDocker,
onOpenTable,
onViewStructure,
onViewErd,
}: {
name: string;
isActive: boolean;
connectionId: string;
onSwitch: () => void;
isSwitching: boolean;
onCloneToDocker: (dbName: string) => void;
onOpenTable: (schema: string, table: string) => void;
onViewStructure: (schema: string, table: string) => void;
onViewErd: (schema: string) => void;
}) {
const [expanded, setExpanded] = useState(false);
const readOnlyMap = useAppStore((s) => s.readOnlyMap);
const isReadOnly = readOnlyMap[connectionId] ?? true;
const dropDbMutation = useDropDatabase();
const handleClick = () => {
if (!isActive) {
@@ -199,23 +168,6 @@ function DatabaseNode({
setExpanded(!expanded);
};
const handleDropDb = () => {
if (isActive) {
toast.error("Cannot drop the active database");
return;
}
if (!confirm(`Drop database "${name}"? This will terminate all connections and cannot be undone.`)) {
return;
}
dropDbMutation.mutate(
{ connectionId, name },
{
onSuccess: () => toast.success(`Database "${name}" dropped`),
onError: (err) => toast.error("Failed to drop database", { description: String(err) }),
}
);
};
return (
<div>
<ContextMenu>
@@ -250,66 +202,6 @@ function DatabaseNode({
>
Properties
</ContextMenuItem>
<ContextMenuItem onClick={() => onCloneToDocker(name)}>
Clone to Docker
</ContextMenuItem>
<ContextMenuSeparator />
<ContextMenuItem
disabled={!isActive}
onClick={() => {
if (!isActive) return;
const tab: Tab = {
id: crypto.randomUUID(),
type: "validation",
title: "Data Validation",
connectionId,
database: name,
};
useAppStore.getState().addTab(tab);
}}
>
Data Validation
</ContextMenuItem>
<ContextMenuItem
disabled={!isActive}
onClick={() => {
if (!isActive) return;
const tab: Tab = {
id: crypto.randomUUID(),
type: "index-advisor",
title: "Index Advisor",
connectionId,
database: name,
};
useAppStore.getState().addTab(tab);
}}
>
Index Advisor
</ContextMenuItem>
<ContextMenuItem
disabled={!isActive}
onClick={() => {
if (!isActive) return;
const tab: Tab = {
id: crypto.randomUUID(),
type: "snapshots",
title: "Data Snapshots",
connectionId,
database: name,
};
useAppStore.getState().addTab(tab);
}}
>
Data Snapshots
</ContextMenuItem>
<ContextMenuSeparator />
<ContextMenuItem
disabled={isActive || isReadOnly}
onClick={handleDropDb}
className="text-destructive focus:text-destructive"
>
Drop Database
</ContextMenuItem>
</ContextMenuContent>
</ContextMenu>
{expanded && isActive && (
@@ -318,7 +210,6 @@ function DatabaseNode({
connectionId={connectionId}
onOpenTable={onOpenTable}
onViewStructure={onViewStructure}
onViewErd={onViewErd}
/>
</div>
)}
@@ -335,12 +226,10 @@ function SchemasForCurrentDb({
connectionId,
onOpenTable,
onViewStructure,
onViewErd,
}: {
connectionId: string;
onOpenTable: (schema: string, table: string) => void;
onViewStructure: (schema: string, table: string) => void;
onViewErd: (schema: string) => void;
}) {
const { data: schemas } = useSchemas(connectionId);
@@ -359,7 +248,6 @@ function SchemasForCurrentDb({
connectionId={connectionId}
onOpenTable={(table) => onOpenTable(schema, table)}
onViewStructure={(table) => onViewStructure(schema, table)}
onViewErd={() => onViewErd(schema)}
/>
))}
</>
@@ -371,45 +259,34 @@ function SchemaNode({
connectionId,
onOpenTable,
onViewStructure,
onViewErd,
}: {
schema: string;
connectionId: string;
onOpenTable: (table: string) => void;
onViewStructure: (table: string) => void;
onViewErd: () => void;
}) {
const [expanded, setExpanded] = useState(false);
return (
<div>
<ContextMenu>
<ContextMenuTrigger>
<div
className="flex items-center gap-1.5 rounded-md px-1.5 py-1 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none font-medium"
onClick={() => setExpanded(!expanded)}
>
<span className="text-muted-foreground/50">
{expanded ? (
<ChevronDown className="h-3.5 w-3.5" />
) : (
<ChevronRight className="h-3.5 w-3.5" />
)}
</span>
{expanded ? (
<FolderOpen className="h-3.5 w-3.5 text-amber-500/70" />
) : (
<Database className="h-3.5 w-3.5 text-muted-foreground/50" />
)}
<span>{schema}</span>
</div>
</ContextMenuTrigger>
<ContextMenuContent>
<ContextMenuItem onClick={onViewErd}>
View ER Diagram
</ContextMenuItem>
</ContextMenuContent>
</ContextMenu>
<div
className="flex items-center gap-1.5 rounded-md px-1.5 py-1 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none font-medium"
onClick={() => setExpanded(!expanded)}
>
<span className="text-muted-foreground/50">
{expanded ? (
<ChevronDown className="h-3.5 w-3.5" />
) : (
<ChevronRight className="h-3.5 w-3.5" />
)}
</span>
{expanded ? (
<FolderOpen className="h-3.5 w-3.5 text-amber-500/70" />
) : (
<Database className="h-3.5 w-3.5 text-muted-foreground/50" />
)}
<span>{schema}</span>
</div>
{expanded && (
<div className="ml-4">
<CategoryNode
@@ -473,8 +350,6 @@ function CategoryNode({
onViewStructure: (table: string) => void;
}) {
const [expanded, setExpanded] = useState(false);
const [privilegesTarget, setPrivilegesTarget] = useState<string | null>(null);
const [dataGenTarget, setDataGenTarget] = useState<string | null>(null);
const tablesQuery = useTables(
expanded && category === "tables" ? connectionId : null,
@@ -555,19 +430,6 @@ function CategoryNode({
>
View Structure
</ContextMenuItem>
{category === "tables" && (
<ContextMenuItem
onClick={() => setDataGenTarget(item.name)}
>
Generate Test Data
</ContextMenuItem>
)}
<ContextMenuSeparator />
<ContextMenuItem
onClick={() => setPrivilegesTarget(item.name)}
>
View Privileges
</ContextMenuItem>
</ContextMenuContent>
</ContextMenu>
);
@@ -586,26 +448,6 @@ function CategoryNode({
})}
</div>
)}
{privilegesTarget && (
<GrantRevokeDialog
open={!!privilegesTarget}
onOpenChange={(open) => !open && setPrivilegesTarget(null)}
connectionId={connectionId}
objectType="TABLE"
objectName={`${schema}.${privilegesTarget}`}
schema={schema}
table={privilegesTarget}
/>
)}
{dataGenTarget && (
<GenerateDataDialog
open={!!dataGenTarget}
onOpenChange={(open) => !open && setDataGenTarget(null)}
connectionId={connectionId}
schema={schema}
table={dataGenTarget}
/>
)}
</div>
);
}
+1 -43
View File
@@ -22,7 +22,7 @@ import { useAiSettings, useSaveAiSettings } from "@/hooks/use-ai";
import { AiSettingsFields } from "@/components/ai/AiSettingsFields";
import { Loader2, Copy, Check } from "lucide-react";
import { toast } from "sonner";
import type { AppSettings, DockerHost } from "@/types";
import type { AppSettings } from "@/types";
interface Props {
open: boolean;
@@ -41,10 +41,6 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
const [mcpEnabled, setMcpEnabled] = useState(true);
const [mcpPort, setMcpPort] = useState(9427);
// Docker state
const [dockerHost, setDockerHost] = useState<DockerHost>("local");
const [dockerRemoteUrl, setDockerRemoteUrl] = useState("");
// AI state
const [ollamaUrl, setOllamaUrl] = useState("http://localhost:11434");
const [aiModel, setAiModel] = useState("");
@@ -58,8 +54,6 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
if (appSettings) {
setMcpEnabled(appSettings.mcp.enabled);
setMcpPort(appSettings.mcp.port);
setDockerHost(appSettings.docker.host);
setDockerRemoteUrl(appSettings.docker.remote_url ?? "");
}
}
@@ -83,10 +77,6 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
const handleSave = () => {
const settings: AppSettings = {
mcp: { enabled: mcpEnabled, port: mcpPort },
docker: {
host: dockerHost,
remote_url: dockerHost === "remote" ? dockerRemoteUrl || undefined : undefined,
},
};
saveAppMutation.mutate(settings, {
@@ -183,38 +173,6 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
<Separator />
{/* Docker */}
<section className="flex flex-col gap-3">
<h3 className="text-sm font-medium">Docker</h3>
<div className="flex flex-col gap-1.5">
<label className="text-xs text-muted-foreground">Docker Host</label>
<Select value={dockerHost} onValueChange={(v) => setDockerHost(v as DockerHost)}>
<SelectTrigger className="h-8 text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="local">Local</SelectItem>
<SelectItem value="remote">Remote</SelectItem>
</SelectContent>
</Select>
</div>
{dockerHost === "remote" && (
<div className="flex flex-col gap-1.5">
<label className="text-xs text-muted-foreground">Remote URL</label>
<Input
value={dockerRemoteUrl}
onChange={(e) => setDockerRemoteUrl(e.target.value)}
placeholder="tcp://192.168.1.100:2375"
className="h-8 text-xs"
/>
</div>
)}
</section>
<Separator />
{/* AI */}
<section className="flex flex-col gap-3">
<h3 className="text-sm font-medium">AI</h3>
@@ -1,276 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { useSchemas, useTables } from "@/hooks/use-schema";
import { useCreateSnapshot } from "@/hooks/use-snapshots";
import { toast } from "sonner";
import { save } from "@tauri-apps/plugin-dialog";
import {
Loader2,
CheckCircle2,
XCircle,
Camera,
} from "lucide-react";
import type { TableRef } from "@/types";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
}
type Step = "config" | "progress" | "done";
export function CreateSnapshotDialog({ open, onOpenChange, connectionId }: Props) {
const [step, setStep] = useState<Step>("config");
const [name, setName] = useState("");
const [selectedSchema, setSelectedSchema] = useState<string>("");
const [selectedTables, setSelectedTables] = useState<Set<string>>(new Set());
const [includeDeps, setIncludeDeps] = useState(true);
const { data: schemas } = useSchemas(connectionId);
const { data: tables } = useTables(
selectedSchema ? connectionId : null,
selectedSchema
);
const { create, result, error, isCreating, progress, reset } = useCreateSnapshot();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setStep("config");
setName(`snapshot-${new Date().toISOString().slice(0, 10)}`);
setSelectedTables(new Set());
setIncludeDeps(true);
reset();
}
}
const [prevSchemas, setPrevSchemas] = useState(schemas);
if (schemas !== prevSchemas) {
setPrevSchemas(schemas);
if (schemas && schemas.length > 0 && !selectedSchema) {
setSelectedSchema(schemas.find((s) => s === "public") || schemas[0]);
}
}
const [prevProgress, setPrevProgress] = useState(progress);
if (progress !== prevProgress) {
setPrevProgress(progress);
if (progress?.stage === "done" || progress?.stage === "error") {
setStep("done");
}
}
const handleToggleTable = (tableName: string) => {
setSelectedTables((prev) => {
const next = new Set(prev);
if (next.has(tableName)) {
next.delete(tableName);
} else {
next.add(tableName);
}
return next;
});
};
const handleSelectAll = () => {
if (tables) {
if (selectedTables.size === tables.length) {
setSelectedTables(new Set());
} else {
setSelectedTables(new Set(tables.map((t) => t.name)));
}
}
};
const handleCreate = async () => {
if (!name.trim() || selectedTables.size === 0) {
toast.error("Please enter a name and select at least one table");
return;
}
const filePath = await save({
defaultPath: `${name}.json`,
filters: [{ name: "JSON", extensions: ["json"] }],
});
if (!filePath) return;
setStep("progress");
const tableRefs: TableRef[] = Array.from(selectedTables).map((t) => ({
schema: selectedSchema,
table: t,
}));
const snapshotId = crypto.randomUUID();
create({
params: {
connection_id: connectionId,
tables: tableRefs,
name: name.trim(),
include_dependencies: includeDeps,
},
snapshotId,
filePath,
});
};
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[520px] max-h-[80vh] overflow-y-auto">
<DialogHeader>
<DialogTitle className="flex items-center gap-2">
<Camera className="h-5 w-5" />
Create Snapshot
</DialogTitle>
</DialogHeader>
{step === "config" && (
<>
<div className="grid gap-3 py-2">
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Name</label>
<Input
className="col-span-3"
value={name}
onChange={(e) => setName(e.target.value)}
placeholder="snapshot-name"
/>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Schema</label>
<select
className="col-span-3 rounded-md border bg-background px-3 py-2 text-sm"
value={selectedSchema}
onChange={(e) => {
setSelectedSchema(e.target.value);
setSelectedTables(new Set());
}}
>
{schemas?.map((s) => (
<option key={s} value={s}>{s}</option>
))}
</select>
</div>
<div className="grid grid-cols-4 items-start gap-3">
<label className="text-right text-sm text-muted-foreground pt-1">Tables</label>
<div className="col-span-3 space-y-1">
{tables && tables.length > 0 && (
<button
className="text-xs text-primary hover:underline"
onClick={handleSelectAll}
>
{selectedTables.size === tables.length ? "Deselect all" : "Select all"}
</button>
)}
<div className="max-h-48 overflow-y-auto rounded-md border p-2 space-y-1">
{tables?.map((t) => (
<label key={t.name} className="flex items-center gap-2 text-sm cursor-pointer hover:bg-accent rounded px-1">
<input
type="checkbox"
checked={selectedTables.has(t.name)}
onChange={() => handleToggleTable(t.name)}
className="rounded"
/>
{t.name}
</label>
)) ?? (
<p className="text-xs text-muted-foreground">Select a schema first</p>
)}
</div>
<p className="text-xs text-muted-foreground">{selectedTables.size} tables selected</p>
</div>
</div>
<div className="grid grid-cols-4 items-center gap-3">
<label className="text-right text-sm text-muted-foreground">Dependencies</label>
<div className="col-span-3 flex items-center gap-2">
<input
type="checkbox"
checked={includeDeps}
onChange={(e) => setIncludeDeps(e.target.checked)}
className="rounded"
/>
<span className="text-sm text-muted-foreground">
Include referenced tables (foreign keys)
</span>
</div>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Cancel</Button>
<Button onClick={handleCreate} disabled={selectedTables.size === 0}>
Create Snapshot
</Button>
</DialogFooter>
</>
)}
{step === "progress" && (
<div className="py-4 space-y-4">
<div className="space-y-2">
<div className="flex items-center justify-between text-sm">
<span>{progress?.message || "Starting..."}</span>
<span className="text-muted-foreground">{progress?.percent ?? 0}%</span>
</div>
<div className="h-2 rounded-full bg-secondary overflow-hidden">
<div
className="h-full rounded-full bg-primary transition-all duration-300"
style={{ width: `${progress?.percent ?? 0}%` }}
/>
</div>
</div>
{isCreating && (
<div className="flex items-center gap-2 text-sm text-muted-foreground">
<Loader2 className="h-4 w-4 animate-spin" />
{progress?.stage || "Initializing..."}
</div>
)}
</div>
)}
{step === "done" && (
<div className="py-4 space-y-4">
{error ? (
<div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
<XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium text-destructive">Snapshot Failed</p>
<p className="text-xs text-muted-foreground">{error}</p>
</div>
</div>
) : (
<div className="flex items-start gap-3 rounded-md border border-green-500/50 bg-green-500/10 p-4">
<CheckCircle2 className="h-5 w-5 text-green-500 shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium">Snapshot Created</p>
<p className="text-xs text-muted-foreground">
{result?.total_rows} rows from {result?.tables.length} tables saved.
</p>
</div>
</div>
)}
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Close</Button>
{error && <Button onClick={() => setStep("config")}>Retry</Button>}
</DialogFooter>
</div>
)}
</DialogContent>
</Dialog>
);
}
@@ -1,248 +0,0 @@
import { useState } from "react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { useRestoreSnapshot, useReadSnapshotMetadata } from "@/hooks/use-snapshots";
import { toast } from "sonner";
import { open as openFile } from "@tauri-apps/plugin-dialog";
import {
Loader2,
CheckCircle2,
XCircle,
Upload,
AlertTriangle,
FileJson,
} from "lucide-react";
import type { SnapshotMetadata } from "@/types";
interface Props {
open: boolean;
onOpenChange: (open: boolean) => void;
connectionId: string;
}
type Step = "select" | "confirm" | "progress" | "done";
export function RestoreSnapshotDialog({ open, onOpenChange, connectionId }: Props) {
const [step, setStep] = useState<Step>("select");
const [filePath, setFilePath] = useState<string | null>(null);
const [metadata, setMetadata] = useState<SnapshotMetadata | null>(null);
const [truncate, setTruncate] = useState(false);
const readMeta = useReadSnapshotMetadata();
const { restore, rowsRestored, error, isRestoring, progress, reset } = useRestoreSnapshot();
const [prevOpen, setPrevOpen] = useState(false);
if (open !== prevOpen) {
setPrevOpen(open);
if (open) {
setStep("select");
setFilePath(null);
setMetadata(null);
setTruncate(false);
reset();
}
}
const [prevProgress, setPrevProgress] = useState(progress);
if (progress !== prevProgress) {
setPrevProgress(progress);
if (progress?.stage === "done" || progress?.stage === "error") {
setStep("done");
}
}
const handleSelectFile = async () => {
const selected = await openFile({
filters: [{ name: "JSON Snapshot", extensions: ["json"] }],
multiple: false,
});
if (!selected) return;
const path = typeof selected === "string" ? selected : (selected as { path: string }).path;
setFilePath(path);
readMeta.mutate(path, {
onSuccess: (meta) => {
setMetadata(meta);
setStep("confirm");
},
onError: (err) => toast.error("Invalid snapshot file", { description: String(err) }),
});
};
const handleRestore = () => {
if (!filePath) return;
setStep("progress");
const snapshotId = crypto.randomUUID();
restore({
params: {
connection_id: connectionId,
file_path: filePath,
truncate_before_restore: truncate,
},
snapshotId,
});
};
function formatBytes(bytes: number): string {
if (bytes < 1024) return `${bytes} B`;
if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)} KB`;
return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
}
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="sm:max-w-[520px]">
<DialogHeader>
<DialogTitle className="flex items-center gap-2">
<Upload className="h-5 w-5" />
Restore Snapshot
</DialogTitle>
</DialogHeader>
{step === "select" && (
<>
<div className="py-8 flex flex-col items-center gap-3">
<FileJson className="h-12 w-12 text-muted-foreground" />
<p className="text-sm text-muted-foreground">Select a snapshot file to restore</p>
<Button onClick={handleSelectFile} disabled={readMeta.isPending}>
{readMeta.isPending ? (
<><Loader2 className="h-4 w-4 animate-spin mr-1" />Reading...</>
) : (
"Choose File"
)}
</Button>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Cancel</Button>
</DialogFooter>
</>
)}
{step === "confirm" && metadata && (
<>
<div className="space-y-3 py-2">
<div className="rounded-md border p-3 space-y-2 text-sm">
<div className="flex justify-between">
<span className="text-muted-foreground">Name</span>
<span className="font-medium">{metadata.name}</span>
</div>
<div className="flex justify-between">
<span className="text-muted-foreground">Created</span>
<span>{new Date(metadata.created_at).toLocaleString()}</span>
</div>
<div className="flex justify-between">
<span className="text-muted-foreground">Tables</span>
<span>{metadata.tables.length}</span>
</div>
<div className="flex justify-between">
<span className="text-muted-foreground">Total Rows</span>
<span>{metadata.total_rows.toLocaleString()}</span>
</div>
<div className="flex justify-between">
<span className="text-muted-foreground">File Size</span>
<span>{formatBytes(metadata.file_size_bytes)}</span>
</div>
</div>
<div className="space-y-1">
<p className="text-xs font-medium text-muted-foreground">Tables included:</p>
<div className="flex flex-wrap gap-1">
{metadata.tables.map((t) => (
<Badge key={`${t.schema}.${t.table}`} variant="secondary" className="text-[10px]">
{t.schema}.{t.table} ({t.row_count})
</Badge>
))}
</div>
</div>
<div className="flex items-start gap-2 rounded-md border border-yellow-500/50 bg-yellow-500/10 p-3">
<AlertTriangle className="h-4 w-4 text-yellow-600 shrink-0 mt-0.5" />
<div className="space-y-2">
<label className="flex items-center gap-2 text-sm cursor-pointer">
<input
type="checkbox"
checked={truncate}
onChange={(e) => setTruncate(e.target.checked)}
className="rounded"
/>
Truncate existing data before restore
</label>
{truncate && (
<p className="text-xs text-yellow-700 dark:text-yellow-400">
This will DELETE all existing data in the affected tables before restoring.
</p>
)}
</div>
</div>
</div>
<DialogFooter>
<Button variant="outline" onClick={() => setStep("select")}>Back</Button>
<Button onClick={handleRestore}>Restore</Button>
</DialogFooter>
</>
)}
{step === "progress" && (
<div className="py-4 space-y-4">
<div className="space-y-2">
<div className="flex items-center justify-between text-sm">
<span>{progress?.message || "Starting..."}</span>
<span className="text-muted-foreground">{progress?.percent ?? 0}%</span>
</div>
<div className="h-2 rounded-full bg-secondary overflow-hidden">
<div
className="h-full rounded-full bg-primary transition-all duration-300"
style={{ width: `${progress?.percent ?? 0}%` }}
/>
</div>
</div>
{isRestoring && (
<div className="flex items-center gap-2 text-sm text-muted-foreground">
<Loader2 className="h-4 w-4 animate-spin" />
{progress?.detail || progress?.stage || "Restoring..."}
</div>
)}
</div>
)}
{step === "done" && (
<div className="py-4 space-y-4">
{error ? (
<div className="flex items-start gap-3 rounded-md border border-destructive/50 bg-destructive/10 p-4">
<XCircle className="h-5 w-5 text-destructive shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium text-destructive">Restore Failed</p>
<p className="text-xs text-muted-foreground">{error}</p>
</div>
</div>
) : (
<div className="flex items-start gap-3 rounded-md border border-green-500/50 bg-green-500/10 p-4">
<CheckCircle2 className="h-5 w-5 text-green-500 shrink-0 mt-0.5" />
<div className="space-y-1">
<p className="text-sm font-medium">Restore Completed</p>
<p className="text-xs text-muted-foreground">
{rowsRestored?.toLocaleString()} rows restored successfully.
</p>
</div>
</div>
)}
<DialogFooter>
<Button variant="outline" onClick={() => onOpenChange(false)}>Close</Button>
{error && <Button onClick={() => setStep("confirm")}>Retry</Button>}
</DialogFooter>
</div>
)}
</DialogContent>
</Dialog>
);
}
-122
View File
@@ -1,122 +0,0 @@
import { useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { useListSnapshots } from "@/hooks/use-snapshots";
import { CreateSnapshotDialog } from "./CreateSnapshotDialog";
import { RestoreSnapshotDialog } from "./RestoreSnapshotDialog";
import {
Camera,
Upload,
Plus,
FileJson,
Calendar,
Table2,
HardDrive,
} from "lucide-react";
import type { SnapshotMetadata } from "@/types";
interface Props {
connectionId: string;
}
function formatBytes(bytes: number): string {
if (bytes < 1024) return `${bytes} B`;
if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)} KB`;
return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
}
function SnapshotCard({ snapshot }: { snapshot: SnapshotMetadata }) {
return (
<div className="rounded-md border p-3 space-y-2">
<div className="flex items-start justify-between">
<div className="flex items-center gap-2">
<FileJson className="h-4 w-4 text-primary" />
<span className="text-sm font-medium">{snapshot.name}</span>
</div>
<Badge variant="secondary" className="text-[10px]">v{snapshot.version}</Badge>
</div>
<div className="grid grid-cols-3 gap-2 text-xs text-muted-foreground">
<div className="flex items-center gap-1">
<Calendar className="h-3 w-3" />
{new Date(snapshot.created_at).toLocaleDateString()}
</div>
<div className="flex items-center gap-1">
<Table2 className="h-3 w-3" />
{snapshot.tables.length} tables
</div>
<div className="flex items-center gap-1">
<HardDrive className="h-3 w-3" />
{formatBytes(snapshot.file_size_bytes)}
</div>
</div>
<div className="flex flex-wrap gap-1">
{snapshot.tables.map((t) => (
<Badge key={`${t.schema}.${t.table}`} variant="outline" className="text-[10px]">
{t.schema}.{t.table}
<span className="ml-1 text-muted-foreground">({t.row_count})</span>
</Badge>
))}
</div>
<div className="text-xs text-muted-foreground">
{snapshot.total_rows.toLocaleString()} total rows
</div>
</div>
);
}
export function SnapshotPanel({ connectionId }: Props) {
const [showCreate, setShowCreate] = useState(false);
const [showRestore, setShowRestore] = useState(false);
const { data: snapshots } = useListSnapshots();
return (
<div className="flex h-full flex-col">
{/* Header */}
<div className="border-b px-4 py-3 flex items-center justify-between">
<div className="flex items-center gap-2">
<Camera className="h-5 w-5 text-primary" />
<h2 className="text-sm font-medium">Data Snapshots</h2>
</div>
<div className="flex items-center gap-2">
<Button variant="outline" size="sm" onClick={() => setShowRestore(true)}>
<Upload className="h-3.5 w-3.5 mr-1" />
Restore
</Button>
<Button size="sm" onClick={() => setShowCreate(true)}>
<Plus className="h-3.5 w-3.5 mr-1" />
Create
</Button>
</div>
</div>
{/* Content */}
<div className="flex-1 overflow-auto p-4 space-y-2">
{!snapshots || snapshots.length === 0 ? (
<div className="flex h-full flex-col items-center justify-center gap-3 text-muted-foreground">
<Camera className="h-12 w-12" />
<p className="text-sm">No snapshots yet</p>
<p className="text-xs">Create a snapshot to save table data for later restoration.</p>
</div>
) : (
snapshots.map((snap) => (
<SnapshotCard key={snap.id} snapshot={snap} />
))
)}
</div>
<CreateSnapshotDialog
open={showCreate}
onOpenChange={setShowCreate}
connectionId={connectionId}
/>
<RestoreSnapshotDialog
open={showRestore}
onOpenChange={setShowRestore}
connectionId={connectionId}
/>
</div>
);
}
@@ -29,7 +29,11 @@ interface Props {
export function TableDataView({ connectionId, schema, table }: Props) {
const readOnlyMap = useAppStore((s) => s.readOnlyMap);
const isReadOnly = readOnlyMap[connectionId] ?? true;
const dbFlavors = useAppStore((s) => s.dbFlavors);
// ClickHouse mutations are async and not transactional — surface the table viewer as
// read-only so existing edit/insert/delete affordances are hidden.
const isReadOnly =
(readOnlyMap[connectionId] ?? true) || dbFlavors[connectionId] === "clickhouse";
const [page, setPage] = useState(1);
const [pageSize, setPageSize] = useState(50);
@@ -1,216 +0,0 @@
import { useState, useCallback } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Badge } from "@/components/ui/badge";
import {
useGenerateValidationSql,
useRunValidationRule,
useSuggestValidationRules,
} from "@/hooks/use-validation";
import { ValidationRuleCard } from "./ValidationRuleCard";
import { toast } from "sonner";
import { Plus, Sparkles, PlayCircle, Loader2, ShieldCheck } from "lucide-react";
import type { ValidationRule, ValidationStatus } from "@/types";
interface Props {
connectionId: string;
}
export function ValidationPanel({ connectionId }: Props) {
const [rules, setRules] = useState<ValidationRule[]>([]);
const [ruleInput, setRuleInput] = useState("");
const [runningIds, setRunningIds] = useState<Set<string>>(new Set());
const generateSql = useGenerateValidationSql();
const runRule = useRunValidationRule();
const suggestRules = useSuggestValidationRules();
const updateRule = useCallback(
(id: string, updates: Partial<ValidationRule>) => {
setRules((prev) =>
prev.map((r) => (r.id === id ? { ...r, ...updates } : r))
);
},
[]
);
const addRule = useCallback(
async (description: string) => {
const id = crypto.randomUUID();
const newRule: ValidationRule = {
id,
description,
generated_sql: "",
status: "generating" as ValidationStatus,
violation_count: 0,
sample_violations: [],
violation_columns: [],
error: null,
};
setRules((prev) => [...prev, newRule]);
try {
const sql = await generateSql.mutateAsync({
connectionId,
ruleDescription: description,
});
updateRule(id, { generated_sql: sql, status: "pending" });
} catch (err) {
updateRule(id, {
status: "error",
error: String(err),
});
}
},
[connectionId, generateSql, updateRule]
);
const handleAddRule = () => {
if (!ruleInput.trim()) return;
addRule(ruleInput.trim());
setRuleInput("");
};
const handleRunRule = useCallback(
async (id: string) => {
const rule = rules.find((r) => r.id === id);
if (!rule || !rule.generated_sql) return;
setRunningIds((prev) => new Set(prev).add(id));
updateRule(id, { status: "running" });
try {
const result = await runRule.mutateAsync({
connectionId,
sql: rule.generated_sql,
});
updateRule(id, {
status: result.status,
violation_count: result.violation_count,
sample_violations: result.sample_violations,
violation_columns: result.violation_columns,
error: result.error,
});
} catch (err) {
updateRule(id, { status: "error", error: String(err) });
} finally {
setRunningIds((prev) => {
const next = new Set(prev);
next.delete(id);
return next;
});
}
},
[rules, connectionId, runRule, updateRule]
);
const handleRemoveRule = useCallback((id: string) => {
setRules((prev) => prev.filter((r) => r.id !== id));
}, []);
const handleRunAll = async () => {
const runnableRules = rules.filter(
(r) => r.generated_sql && r.status !== "generating"
);
for (const rule of runnableRules) {
await handleRunRule(rule.id);
}
};
const handleSuggest = async () => {
try {
const suggestions = await suggestRules.mutateAsync(connectionId);
for (const desc of suggestions) {
await addRule(desc);
}
toast.success(`Added ${suggestions.length} suggested rules`);
} catch (err) {
toast.error("Failed to suggest rules", { description: String(err) });
}
};
const passed = rules.filter((r) => r.status === "passed").length;
const failed = rules.filter((r) => r.status === "failed").length;
const errors = rules.filter((r) => r.status === "error").length;
return (
<div className="flex h-full flex-col">
{/* Header */}
<div className="border-b px-4 py-3 space-y-3">
<div className="flex items-center justify-between">
<div className="flex items-center gap-2">
<ShieldCheck className="h-5 w-5 text-primary" />
<h2 className="text-sm font-medium">Data Validation</h2>
</div>
<div className="flex items-center gap-2">
<Button
variant="outline"
size="sm"
onClick={handleSuggest}
disabled={suggestRules.isPending}
>
{suggestRules.isPending ? (
<Loader2 className="h-3.5 w-3.5 animate-spin mr-1" />
) : (
<Sparkles className="h-3.5 w-3.5 mr-1" />
)}
Auto-suggest
</Button>
<Button
variant="outline"
size="sm"
onClick={handleRunAll}
disabled={rules.length === 0 || runningIds.size > 0}
>
<PlayCircle className="h-3.5 w-3.5 mr-1" />
Run All
</Button>
</div>
</div>
<div className="flex items-center gap-2">
<Input
placeholder="Describe a data quality rule (e.g., 'All orders must have a positive total')"
value={ruleInput}
onChange={(e) => setRuleInput(e.target.value)}
onKeyDown={(e) => e.key === "Enter" && handleAddRule()}
className="flex-1"
/>
<Button size="sm" onClick={handleAddRule} disabled={!ruleInput.trim()}>
<Plus className="h-3.5 w-3.5 mr-1" />
Add
</Button>
</div>
{rules.length > 0 && (
<div className="flex items-center gap-2 text-xs">
<span className="text-muted-foreground">{rules.length} rules</span>
{passed > 0 && <Badge className="bg-green-600 text-white text-[10px]">{passed} passed</Badge>}
{failed > 0 && <Badge variant="destructive" className="text-[10px]">{failed} failed</Badge>}
{errors > 0 && <Badge variant="outline" className="text-[10px]">{errors} errors</Badge>}
</div>
)}
</div>
{/* Rules List */}
<div className="flex-1 overflow-auto p-4 space-y-2">
{rules.length === 0 ? (
<div className="flex h-full items-center justify-center text-sm text-muted-foreground">
Add a validation rule or click Auto-suggest to get started.
</div>
) : (
rules.map((rule) => (
<ValidationRuleCard
key={rule.id}
rule={rule}
onRun={() => handleRunRule(rule.id)}
onRemove={() => handleRemoveRule(rule.id)}
isRunning={runningIds.has(rule.id)}
/>
))
)}
</div>
</div>
);
}
@@ -1,138 +0,0 @@
import { useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import {
ChevronDown,
ChevronRight,
Play,
Trash2,
Loader2,
} from "lucide-react";
import type { ValidationRule } from "@/types";
interface Props {
rule: ValidationRule;
onRun: () => void;
onRemove: () => void;
isRunning: boolean;
}
function statusBadge(status: string) {
switch (status) {
case "passed":
return <Badge className="bg-green-600 text-white">Passed</Badge>;
case "failed":
return <Badge variant="destructive">Failed</Badge>;
case "error":
return <Badge variant="outline" className="text-destructive border-destructive">Error</Badge>;
case "generating":
case "running":
return <Badge variant="secondary"><Loader2 className="h-3 w-3 animate-spin mr-1" />Running</Badge>;
default:
return <Badge variant="secondary">Pending</Badge>;
}
}
export function ValidationRuleCard({ rule, onRun, onRemove, isRunning }: Props) {
const [showSql, setShowSql] = useState(false);
const [showViolations, setShowViolations] = useState(false);
return (
<div className="rounded-md border p-3 space-y-2">
<div className="flex items-start justify-between gap-2">
<div className="flex-1 min-w-0">
<p className="text-sm">{rule.description}</p>
</div>
<div className="flex items-center gap-1 shrink-0">
{statusBadge(rule.status)}
<Button
variant="ghost"
size="sm"
className="h-7 w-7 p-0"
onClick={onRun}
disabled={isRunning}
>
{isRunning ? (
<Loader2 className="h-3.5 w-3.5 animate-spin" />
) : (
<Play className="h-3.5 w-3.5" />
)}
</Button>
<Button
variant="ghost"
size="sm"
className="h-7 w-7 p-0 text-muted-foreground hover:text-destructive"
onClick={onRemove}
>
<Trash2 className="h-3.5 w-3.5" />
</Button>
</div>
</div>
{rule.status === "failed" && (
<p className="text-xs text-destructive">
{rule.violation_count} violation{rule.violation_count !== 1 ? "s" : ""} found
</p>
)}
{rule.error && (
<p className="text-xs text-destructive">{rule.error}</p>
)}
{rule.generated_sql && (
<div>
<button
className="flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setShowSql(!showSql)}
>
{showSql ? <ChevronDown className="h-3 w-3" /> : <ChevronRight className="h-3 w-3" />}
SQL
</button>
{showSql && (
<pre className="mt-1 rounded bg-muted p-2 text-xs font-mono overflow-x-auto max-h-32 overflow-y-auto">
{rule.generated_sql}
</pre>
)}
</div>
)}
{rule.status === "failed" && rule.sample_violations.length > 0 && (
<div>
<button
className="flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
onClick={() => setShowViolations(!showViolations)}
>
{showViolations ? <ChevronDown className="h-3 w-3" /> : <ChevronRight className="h-3 w-3" />}
Sample Violations ({rule.sample_violations.length})
</button>
{showViolations && (
<div className="mt-1 overflow-x-auto">
<table className="w-full text-xs">
<thead>
<tr className="border-b">
{rule.violation_columns.map((col) => (
<th key={col} className="px-2 py-1 text-left font-medium text-muted-foreground">
{col}
</th>
))}
</tr>
</thead>
<tbody>
{rule.sample_violations.map((row, i) => (
<tr key={i} className="border-b last:border-0">
{(row as unknown[]).map((val, j) => (
<td key={j} className="px-2 py-1 font-mono">
{val === null ? <span className="text-muted-foreground">NULL</span> : String(val)}
</td>
))}
</tr>
))}
</tbody>
</table>
</div>
)}
</div>
)}
</div>
);
}
+3 -54
View File
@@ -2,13 +2,7 @@ import { useAppStore } from "@/stores/app-store";
import { WorkspacePanel } from "./WorkspacePanel";
import { TableDataView } from "@/components/table-viewer/TableDataView";
import { TableStructure } from "@/components/table-viewer/TableStructure";
import { RoleManagerView } from "@/components/management/RoleManagerView";
import { SessionsView } from "@/components/management/SessionsView";
import { EntityLookupPanel } from "@/components/lookup/EntityLookupPanel";
import { ErdDiagram } from "@/components/erd/ErdDiagram";
import { ValidationPanel } from "@/components/validation/ValidationPanel";
import { IndexAdvisorPanel } from "@/components/index-advisor/IndexAdvisorPanel";
import { SnapshotPanel } from "@/components/snapshots/SnapshotPanel";
import { ChatPanel } from "@/components/chat/ChatPanel";
export function TabContent() {
const { tabs, activeTabId, updateTab } = useAppStore();
@@ -55,54 +49,9 @@ export function TabContent() {
/>
);
break;
case "roles":
case "chat":
content = (
<RoleManagerView
connectionId={tab.connectionId}
/>
);
break;
case "sessions":
content = (
<SessionsView
connectionId={tab.connectionId}
/>
);
break;
case "lookup":
content = (
<EntityLookupPanel
connectionId={tab.connectionId}
/>
);
break;
case "erd":
content = (
<ErdDiagram
connectionId={tab.connectionId}
schema={tab.schema!}
/>
);
break;
case "validation":
content = (
<ValidationPanel
connectionId={tab.connectionId}
/>
);
break;
case "index-advisor":
content = (
<IndexAdvisorPanel
connectionId={tab.connectionId}
/>
);
break;
case "snapshots":
content = (
<SnapshotPanel
connectionId={tab.connectionId}
/>
<ChatPanel tabId={tab.id} connectionId={tab.connectionId} />
);
break;
default:
+56
View File
@@ -0,0 +1,56 @@
import { useCallback } from "react";
import { chatSend } from "@/lib/tauri";
import { useAppStore } from "@/stores/app-store";
import type { ChatMessage } from "@/types";
const EMPTY_THREAD: ChatMessage[] = [];
function newId(prefix: string) {
return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
}
export function useChat(tabId: string, connectionId: string) {
const messages = useAppStore((s) => s.chatThreads[tabId] ?? EMPTY_THREAD);
const pending = useAppStore((s) => Boolean(s.chatPending[tabId]));
const appendChatMessages = useAppStore((s) => s.appendChatMessages);
const clearChatThread = useAppStore((s) => s.clearChatThread);
const setChatPending = useAppStore((s) => s.setChatPending);
const send = useCallback(
async (text: string) => {
const trimmed = text.trim();
if (!trimmed) return;
const state = useAppStore.getState();
if (state.chatPending[tabId]) return;
const history = state.chatThreads[tabId] ?? [];
const userMsg: ChatMessage = {
id: newId("user"),
role: "user",
text: trimmed,
created_at: Date.now(),
};
appendChatMessages(tabId, [userMsg]);
setChatPending(tabId, true);
try {
const reply = await chatSend(connectionId, [...history, userMsg]);
appendChatMessages(tabId, reply);
} catch (err) {
appendChatMessages(tabId, [
{
id: newId("err"),
role: "assistant",
text: `Error: ${String(err)}`,
created_at: Date.now(),
},
]);
} finally {
setChatPending(tabId, false);
}
},
[tabId, connectionId, appendChatMessages, setChatPending]
);
const clear = useCallback(() => clearChatThread(tabId), [tabId, clearChatThread]);
return { messages, pending, send, clear };
}
+22 -2
View File
@@ -51,8 +51,15 @@ export function useTestConnection() {
}
export function useConnect() {
const { addConnectedId, setActiveConnectionId, setPgVersion, setDbFlavor, setCurrentDatabase } =
useAppStore();
const {
addConnectedId,
setActiveConnectionId,
setPgVersion,
setDbFlavor,
setCurrentDatabase,
addTab,
tabs,
} = useAppStore();
return useMutation({
mutationFn: async (config: ConnectionConfig) => {
@@ -65,6 +72,19 @@ export function useConnect() {
setPgVersion(version);
setDbFlavor(id, flavor);
setCurrentDatabase(database);
const hasChatForConnection = tabs.some(
(t) => t.type === "chat" && t.connectionId === id
);
if (!hasChatForConnection) {
addTab({
id: crypto.randomUUID(),
type: "chat",
title: "Chat",
connectionId: id,
database,
});
}
},
});
}
-84
View File
@@ -1,84 +0,0 @@
import { useState, useEffect, useCallback, useRef } from "react";
import { useMutation } from "@tanstack/react-query";
import {
generateTestDataPreview,
insertGeneratedData,
onDataGenProgress,
} from "@/lib/tauri";
import type { GenerateDataParams, DataGenProgress, GeneratedDataPreview } from "@/types";
export function useDataGenerator() {
const [progress, setProgress] = useState<DataGenProgress | null>(null);
const genIdRef = useRef<string>("");
const previewMutation = useMutation({
mutationFn: ({
params,
genId,
}: {
params: GenerateDataParams;
genId: string;
}) => {
genIdRef.current = genId;
setProgress(null);
return generateTestDataPreview(params, genId);
},
});
const insertMutation = useMutation({
mutationFn: ({
connectionId,
preview,
}: {
connectionId: string;
preview: GeneratedDataPreview;
}) => insertGeneratedData(connectionId, preview),
});
useEffect(() => {
let mounted = true;
let unlisten: (() => void) | undefined;
onDataGenProgress((p) => {
if (mounted && p.gen_id === genIdRef.current) {
setProgress(p);
}
}).then((fn) => {
if (mounted) {
unlisten = fn;
} else {
fn();
}
});
return () => {
mounted = false;
unlisten?.();
};
}, []);
const previewRef = useRef(previewMutation);
const insertRef = useRef(insertMutation);
useEffect(() => {
previewRef.current = previewMutation;
insertRef.current = insertMutation;
});
const reset = useCallback(() => {
previewRef.current.reset();
insertRef.current.reset();
setProgress(null);
genIdRef.current = "";
}, []);
return {
generatePreview: previewMutation.mutate,
preview: previewMutation.data as GeneratedDataPreview | undefined,
isGenerating: previewMutation.isPending,
generateError: previewMutation.error ? String(previewMutation.error) : null,
insertData: insertMutation.mutate,
insertedRows: insertMutation.data as number | undefined,
isInserting: insertMutation.isPending,
insertError: insertMutation.error ? String(insertMutation.error) : null,
progress,
reset,
};
}
-122
View File
@@ -1,122 +0,0 @@
import { useState, useEffect, useCallback, useRef } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
checkDocker,
listTuskContainers,
cloneToDocker,
startContainer,
stopContainer,
removeContainer,
onCloneProgress,
} from "@/lib/tauri";
import type { CloneToDockerParams, CloneProgress, CloneResult } from "@/types";
export function useDockerStatus() {
return useQuery({
queryKey: ["docker-status"],
queryFn: checkDocker,
staleTime: 30_000,
});
}
export function useTuskContainers() {
return useQuery({
queryKey: ["tusk-containers"],
queryFn: listTuskContainers,
refetchInterval: 10_000,
});
}
export function useCloneToDocker() {
const [progress, setProgress] = useState<CloneProgress | null>(null);
const cloneIdRef = useRef<string>("");
const queryClient = useQueryClient();
const mutation = useMutation({
mutationFn: ({
params,
cloneId,
}: {
params: CloneToDockerParams;
cloneId: string;
}) => {
cloneIdRef.current = cloneId;
setProgress(null);
return cloneToDocker(params, cloneId);
},
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["connections"] });
queryClient.invalidateQueries({ queryKey: ["tusk-containers"] });
},
});
useEffect(() => {
let mounted = true;
let unlisten: (() => void) | undefined;
onCloneProgress((p) => {
if (mounted && p.clone_id === cloneIdRef.current) {
setProgress(p);
}
}).then((fn) => {
if (mounted) {
unlisten = fn;
} else {
fn();
}
});
return () => {
mounted = false;
unlisten?.();
};
}, []);
const mutationRef = useRef(mutation);
useEffect(() => {
mutationRef.current = mutation;
});
const reset = useCallback(() => {
mutationRef.current.reset();
setProgress(null);
cloneIdRef.current = "";
}, []);
return {
clone: mutation.mutate,
result: mutation.data as CloneResult | undefined,
error: mutation.error ? String(mutation.error) : null,
isCloning: mutation.isPending,
progress,
reset,
};
}
export function useStartContainer() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: (name: string) => startContainer(name),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["tusk-containers"] });
},
});
}
export function useStopContainer() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: (name: string) => stopContainer(name),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["tusk-containers"] });
},
});
}
export function useRemoveContainer() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: (name: string) => removeContainer(name),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["tusk-containers"] });
},
});
}
-60
View File
@@ -1,60 +0,0 @@
import { useState, useEffect, useCallback, useRef } from "react";
import { useMutation } from "@tanstack/react-query";
import { entityLookup, onLookupProgress } from "@/lib/tauri";
import type {
ConnectionConfig,
EntityLookupResult,
LookupProgress,
} from "@/types";
export function useEntityLookup() {
const [progress, setProgress] = useState<LookupProgress | null>(null);
const lookupIdRef = useRef<string>("");
const mutation = useMutation({
mutationFn: ({
config,
columnName,
value,
lookupId,
databases,
}: {
config: ConnectionConfig;
columnName: string;
value: string;
lookupId: string;
databases?: string[];
}) => {
lookupIdRef.current = lookupId;
setProgress(null);
return entityLookup(config, columnName, value, lookupId, databases);
},
});
useEffect(() => {
const unlistenPromise = onLookupProgress((p) => {
if (p.lookup_id === lookupIdRef.current) {
setProgress(p);
}
});
return () => {
unlistenPromise.then((unlisten) => unlisten());
};
}, []);
const reset = useCallback(() => {
mutation.reset();
setProgress(null);
lookupIdRef.current = "";
}, [mutation]);
return {
search: mutation.mutate,
result: mutation.data as EntityLookupResult | undefined,
error: mutation.error ? String(mutation.error) : null,
isSearching: mutation.isPending,
progress,
reset,
};
}
-20
View File
@@ -1,20 +0,0 @@
import { useMutation } from "@tanstack/react-query";
import { getIndexAdvisorReport, applyIndexRecommendation } from "@/lib/tauri";
export function useIndexAdvisorReport() {
return useMutation({
mutationFn: (connectionId: string) => getIndexAdvisorReport(connectionId),
});
}
export function useApplyIndexRecommendation() {
return useMutation({
mutationFn: ({
connectionId,
ddl,
}: {
connectionId: string;
ddl: string;
}) => applyIndexRecommendation(connectionId, ddl),
});
}
-205
View File
@@ -1,205 +0,0 @@
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
getDatabaseInfo,
createDatabase,
dropDatabase,
listRoles,
createRole,
alterRole,
dropRole,
getTablePrivileges,
grantRevoke,
manageRoleMembership,
listSessions,
cancelQuery,
terminateBackend,
} from "@/lib/tauri";
import type {
CreateDatabaseParams,
CreateRoleParams,
AlterRoleParams,
GrantRevokeParams,
RoleMembershipParams,
} from "@/types";
// Queries
export function useDatabaseInfo(connectionId: string | null) {
return useQuery({
queryKey: ["databaseInfo", connectionId],
queryFn: () => getDatabaseInfo(connectionId!),
enabled: !!connectionId,
staleTime: 5 * 60 * 1000,
});
}
export function useRoles(connectionId: string | null) {
return useQuery({
queryKey: ["roles", connectionId],
queryFn: () => listRoles(connectionId!),
enabled: !!connectionId,
staleTime: 5 * 60 * 1000,
});
}
export function useTablePrivileges(
connectionId: string | null,
schema: string | null,
table: string | null
) {
return useQuery({
queryKey: ["tablePrivileges", connectionId, schema, table],
queryFn: () => getTablePrivileges(connectionId!, schema!, table!),
enabled: !!connectionId && !!schema && !!table,
staleTime: 5 * 60 * 1000,
});
}
// Mutations
export function useCreateDatabase() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
params,
}: {
connectionId: string;
params: CreateDatabaseParams;
}) => createDatabase(connectionId, params),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["databaseInfo"] });
queryClient.invalidateQueries({ queryKey: ["databases"] });
},
});
}
export function useDropDatabase() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
name,
}: {
connectionId: string;
name: string;
}) => dropDatabase(connectionId, name),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["databaseInfo"] });
queryClient.invalidateQueries({ queryKey: ["databases"] });
},
});
}
export function useCreateRole() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
params,
}: {
connectionId: string;
params: CreateRoleParams;
}) => createRole(connectionId, params),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["roles"] });
},
});
}
export function useAlterRole() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
params,
}: {
connectionId: string;
params: AlterRoleParams;
}) => alterRole(connectionId, params),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["roles"] });
},
});
}
export function useDropRole() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
name,
}: {
connectionId: string;
name: string;
}) => dropRole(connectionId, name),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["roles"] });
},
});
}
export function useGrantRevoke() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
params,
}: {
connectionId: string;
params: GrantRevokeParams;
}) => grantRevoke(connectionId, params),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["tablePrivileges"] });
},
});
}
// Sessions
export function useSessions(connectionId: string | null) {
return useQuery({
queryKey: ["sessions", connectionId],
queryFn: () => listSessions(connectionId!),
enabled: !!connectionId,
refetchInterval: 5000,
});
}
export function useCancelQuery() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({ connectionId, pid }: { connectionId: string; pid: number }) =>
cancelQuery(connectionId, pid),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["sessions"] });
},
});
}
export function useTerminateBackend() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({ connectionId, pid }: { connectionId: string; pid: number }) =>
terminateBackend(connectionId, pid),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["sessions"] });
},
});
}
export function useManageRoleMembership() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({
connectionId,
params,
}: {
connectionId: string;
params: RoleMembershipParams;
}) => manageRoleMembership(connectionId, params),
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["roles"] });
},
});
}
+22
View File
@@ -0,0 +1,22 @@
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { getMemory, saveMemory } from "@/lib/tauri";
export function useMemory(connectionId: string | null) {
return useQuery({
queryKey: ["memory", connectionId],
queryFn: () => getMemory(connectionId!),
enabled: !!connectionId,
staleTime: 5_000,
});
}
export function useSaveMemory() {
const queryClient = useQueryClient();
return useMutation({
mutationFn: ({ connectionId, content }: { connectionId: string; content: string }) =>
saveMemory(connectionId, content),
onSuccess: (_data, vars) => {
queryClient.invalidateQueries({ queryKey: ["memory", vars.connectionId] });
},
});
}
-10
View File
@@ -8,7 +8,6 @@ import {
listSequences,
switchDatabase,
getColumnDetails,
getSchemaErd,
} from "@/lib/tauri";
import type { ConnectionConfig } from "@/types";
@@ -89,12 +88,3 @@ export function useColumnDetails(connectionId: string | null, schema: string | n
staleTime: 5 * 60 * 1000,
});
}
export function useSchemaErd(connectionId: string | null, schema: string | null) {
return useQuery({
queryKey: ["schema-erd", connectionId, schema],
queryFn: () => getSchemaErd(connectionId!, schema!),
enabled: !!connectionId && !!schema,
staleTime: 5 * 60 * 1000,
});
}
-153
View File
@@ -1,153 +0,0 @@
import { useState, useEffect, useCallback, useRef } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
createSnapshot,
restoreSnapshot,
listSnapshots,
readSnapshotMetadata,
onSnapshotProgress,
} from "@/lib/tauri";
import type {
CreateSnapshotParams,
RestoreSnapshotParams,
SnapshotProgress,
SnapshotMetadata,
} from "@/types";
export function useListSnapshots() {
return useQuery({
queryKey: ["snapshots"],
queryFn: listSnapshots,
staleTime: 30_000,
});
}
export function useReadSnapshotMetadata() {
return useMutation({
mutationFn: (filePath: string) => readSnapshotMetadata(filePath),
});
}
export function useCreateSnapshot() {
const [progress, setProgress] = useState<SnapshotProgress | null>(null);
const snapshotIdRef = useRef<string>("");
const queryClient = useQueryClient();
const mutation = useMutation({
mutationFn: ({
params,
snapshotId,
filePath,
}: {
params: CreateSnapshotParams;
snapshotId: string;
filePath: string;
}) => {
snapshotIdRef.current = snapshotId;
setProgress(null);
return createSnapshot(params, snapshotId, filePath);
},
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: ["snapshots"] });
},
});
useEffect(() => {
let mounted = true;
let unlisten: (() => void) | undefined;
onSnapshotProgress((p) => {
if (mounted && p.snapshot_id === snapshotIdRef.current) {
setProgress(p);
}
}).then((fn) => {
if (mounted) {
unlisten = fn;
} else {
fn();
}
});
return () => {
mounted = false;
unlisten?.();
};
}, []);
const mutationRef = useRef(mutation);
useEffect(() => {
mutationRef.current = mutation;
});
const reset = useCallback(() => {
mutationRef.current.reset();
setProgress(null);
snapshotIdRef.current = "";
}, []);
return {
create: mutation.mutate,
result: mutation.data as SnapshotMetadata | undefined,
error: mutation.error ? String(mutation.error) : null,
isCreating: mutation.isPending,
progress,
reset,
};
}
export function useRestoreSnapshot() {
const [progress, setProgress] = useState<SnapshotProgress | null>(null);
const snapshotIdRef = useRef<string>("");
const mutation = useMutation({
mutationFn: ({
params,
snapshotId,
}: {
params: RestoreSnapshotParams;
snapshotId: string;
}) => {
snapshotIdRef.current = snapshotId;
setProgress(null);
return restoreSnapshot(params, snapshotId);
},
});
useEffect(() => {
let mounted = true;
let unlisten: (() => void) | undefined;
onSnapshotProgress((p) => {
if (mounted && p.snapshot_id === snapshotIdRef.current) {
setProgress(p);
}
}).then((fn) => {
if (mounted) {
unlisten = fn;
} else {
fn();
}
});
return () => {
mounted = false;
unlisten?.();
};
}, []);
const mutationRef = useRef(mutation);
useEffect(() => {
mutationRef.current = mutation;
});
const reset = useCallback(() => {
mutationRef.current.reset();
setProgress(null);
snapshotIdRef.current = "";
}, []);
return {
restore: mutation.mutate,
rowsRestored: mutation.data as number | undefined,
error: mutation.error ? String(mutation.error) : null,
isRestoring: mutation.isPending,
progress,
reset,
};
}
-38
View File
@@ -1,38 +0,0 @@
import { useMutation } from "@tanstack/react-query";
import {
generateValidationSql,
runValidationRule,
suggestValidationRules,
} from "@/lib/tauri";
export function useGenerateValidationSql() {
return useMutation({
mutationFn: ({
connectionId,
ruleDescription,
}: {
connectionId: string;
ruleDescription: string;
}) => generateValidationSql(connectionId, ruleDescription),
});
}
export function useRunValidationRule() {
return useMutation({
mutationFn: ({
connectionId,
sql,
sampleLimit,
}: {
connectionId: string;
sql: string;
sampleLimit?: number;
}) => runValidationRule(connectionId, sql, sampleLimit),
});
}
export function useSuggestValidationRules() {
return useMutation({
mutationFn: (connectionId: string) => suggestValidationRules(connectionId),
});
}
+42
View File
@@ -0,0 +1,42 @@
import type { DbFlavor } from "@/types";
export interface DbCapabilities {
/** Direct row edit / insert / delete via the table viewer. */
rowEdit: boolean;
/** Multiple schemas inside one database. */
multipleSchemas: boolean;
/** Foreign-key constraints, triggers, sequences, indexes. */
pgObjects: boolean;
/** Default port for the engine when creating a new connection. */
defaultPort: number;
/** Display name for the engine. */
label: string;
}
const CAPABILITIES: Record<DbFlavor, DbCapabilities> = {
postgresql: {
rowEdit: true,
multipleSchemas: true,
pgObjects: true,
defaultPort: 5432,
label: "PostgreSQL",
},
greenplum: {
rowEdit: true,
multipleSchemas: true,
pgObjects: true,
defaultPort: 5432,
label: "Greenplum",
},
clickhouse: {
rowEdit: false,
multipleSchemas: false,
pgObjects: false,
defaultPort: 8123,
label: "ClickHouse",
},
};
export function capsFor(flavor: DbFlavor | undefined): DbCapabilities {
return CAPABILITIES[flavor ?? "postgresql"];
}
+8 -160
View File
@@ -1,5 +1,4 @@
import { invoke } from "@tauri-apps/api/core";
import { listen, type UnlistenFn } from "@tauri-apps/api/event";
import type {
ConnectionConfig,
ConnectResult,
@@ -12,38 +11,13 @@ import type {
ConstraintInfo,
IndexInfo,
TriggerInfo,
ErdData,
HistoryEntry,
SavedQuery,
SessionInfo,
DatabaseInfo,
CreateDatabaseParams,
RoleInfo,
CreateRoleParams,
AlterRoleParams,
TablePrivilege,
GrantRevokeParams,
RoleMembershipParams,
AiSettings,
OllamaModel,
EntityLookupResult,
LookupProgress,
DockerStatus,
CloneToDockerParams,
CloneProgress,
CloneResult,
TuskContainer,
AppSettings,
McpStatus,
ValidationRule,
GenerateDataParams,
GeneratedDataPreview,
DataGenProgress,
IndexAdvisorReport,
SnapshotMetadata,
CreateSnapshotParams,
RestoreSnapshotParams,
SnapshotProgress,
ChatMessage,
} from "@/types";
// Connections
@@ -139,9 +113,6 @@ export const getTableTriggers = (
table: string
) => invoke<TriggerInfo[]>("get_table_triggers", { connectionId, schema, table });
export const getSchemaErd = (connectionId: string, schema: string) =>
invoke<ErdData>("get_schema_erd", { connectionId, schema });
// Data
export const getTableData = (params: {
connectionId: string;
@@ -229,47 +200,6 @@ export const exportJson = (
rows: unknown[][]
) => invoke<void>("export_json", { path, columns, rows });
// Management
export const getDatabaseInfo = (connectionId: string) =>
invoke<DatabaseInfo[]>("get_database_info", { connectionId });
export const createDatabase = (connectionId: string, params: CreateDatabaseParams) =>
invoke<void>("create_database", { connectionId, params });
export const dropDatabase = (connectionId: string, name: string) =>
invoke<void>("drop_database", { connectionId, name });
export const listRoles = (connectionId: string) =>
invoke<RoleInfo[]>("list_roles", { connectionId });
export const createRole = (connectionId: string, params: CreateRoleParams) =>
invoke<void>("create_role", { connectionId, params });
export const alterRole = (connectionId: string, params: AlterRoleParams) =>
invoke<void>("alter_role", { connectionId, params });
export const dropRole = (connectionId: string, name: string) =>
invoke<void>("drop_role", { connectionId, name });
export const getTablePrivileges = (connectionId: string, schema: string, table: string) =>
invoke<TablePrivilege[]>("get_table_privileges", { connectionId, schema, table });
export const grantRevoke = (connectionId: string, params: GrantRevokeParams) =>
invoke<void>("grant_revoke", { connectionId, params });
export const manageRoleMembership = (connectionId: string, params: RoleMembershipParams) =>
invoke<void>("manage_role_membership", { connectionId, params });
// Sessions
export const listSessions = (connectionId: string) =>
invoke<SessionInfo[]>("list_sessions", { connectionId });
export const cancelQuery = (connectionId: string, pid: number) =>
invoke<boolean>("cancel_query", { connectionId, pid });
export const terminateBackend = (connectionId: string, pid: number) =>
invoke<boolean>("terminate_backend", { connectionId, pid });
// AI
export const getAiSettings = () =>
invoke<AiSettings>("get_ai_settings");
@@ -289,50 +219,15 @@ export const explainSql = (connectionId: string, sql: string) =>
export const fixSqlError = (connectionId: string, sql: string, errorMessage: string) =>
invoke<string>("fix_sql_error", { connectionId, sql, errorMessage });
// Entity Lookup
export const entityLookup = (
config: ConnectionConfig,
columnName: string,
value: string,
lookupId: string,
databases?: string[]
) =>
invoke<EntityLookupResult>("entity_lookup", {
config,
columnName,
value,
lookupId,
databases,
});
export const chatSend = (connectionId: string, messages: ChatMessage[]) =>
invoke<ChatMessage[]>("chat_send", { connectionId, messages });
export const onLookupProgress = (
callback: (p: LookupProgress) => void
): Promise<UnlistenFn> =>
listen<LookupProgress>("lookup-progress", (e) => callback(e.payload));
// Memory (per-connection markdown notes for the chat agent)
export const getMemory = (connectionId: string) =>
invoke<string>("get_memory", { connectionId });
// Docker
export const checkDocker = () =>
invoke<DockerStatus>("check_docker");
export const listTuskContainers = () =>
invoke<TuskContainer[]>("list_tusk_containers");
export const cloneToDocker = (params: CloneToDockerParams, cloneId: string) =>
invoke<CloneResult>("clone_to_docker", { params, cloneId });
export const startContainer = (name: string) =>
invoke<void>("start_container", { name });
export const stopContainer = (name: string) =>
invoke<void>("stop_container", { name });
export const removeContainer = (name: string) =>
invoke<void>("remove_container", { name });
export const onCloneProgress = (
callback: (p: CloneProgress) => void
): Promise<UnlistenFn> =>
listen<CloneProgress>("clone-progress", (e) => callback(e.payload));
export const saveMemory = (connectionId: string, content: string) =>
invoke<void>("save_memory", { connectionId, content });
// App Settings
export const getAppSettings = () =>
@@ -343,50 +238,3 @@ export const saveAppSettings = (settings: AppSettings) =>
export const getMcpStatus = () =>
invoke<McpStatus>("get_mcp_status");
// Validation (Wave 1)
export const generateValidationSql = (connectionId: string, ruleDescription: string) =>
invoke<string>("generate_validation_sql", { connectionId, ruleDescription });
export const runValidationRule = (connectionId: string, sql: string, sampleLimit?: number) =>
invoke<ValidationRule>("run_validation_rule", { connectionId, sql, sampleLimit });
export const suggestValidationRules = (connectionId: string) =>
invoke<string[]>("suggest_validation_rules", { connectionId });
// Data Generator (Wave 2)
export const generateTestDataPreview = (params: GenerateDataParams, genId: string) =>
invoke<GeneratedDataPreview>("generate_test_data_preview", { params, genId });
export const insertGeneratedData = (connectionId: string, preview: GeneratedDataPreview) =>
invoke<number>("insert_generated_data", { connectionId, preview });
export const onDataGenProgress = (
callback: (p: DataGenProgress) => void
): Promise<UnlistenFn> =>
listen<DataGenProgress>("datagen-progress", (e) => callback(e.payload));
// Index Advisor (Wave 3A)
export const getIndexAdvisorReport = (connectionId: string) =>
invoke<IndexAdvisorReport>("get_index_advisor_report", { connectionId });
export const applyIndexRecommendation = (connectionId: string, ddl: string) =>
invoke<void>("apply_index_recommendation", { connectionId, ddl });
// Snapshots (Wave 3B)
export const createSnapshot = (params: CreateSnapshotParams, snapshotId: string, filePath: string) =>
invoke<SnapshotMetadata>("create_snapshot", { params, snapshotId, filePath });
export const restoreSnapshot = (params: RestoreSnapshotParams, snapshotId: string) =>
invoke<number>("restore_snapshot", { params, snapshotId });
export const listSnapshots = () =>
invoke<SnapshotMetadata[]>("list_snapshots");
export const readSnapshotMetadata = (filePath: string) =>
invoke<SnapshotMetadata>("read_snapshot_metadata", { filePath });
export const onSnapshotProgress = (
callback: (p: SnapshotProgress) => void
): Promise<UnlistenFn> =>
listen<SnapshotProgress>("snapshot-progress", (e) => callback(e.payload));
+9 -6
View File
@@ -4,6 +4,7 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { TooltipProvider } from "@/components/ui/tooltip";
import { Toaster } from "@/components/ui/sonner";
import App from "./App";
import { ErrorBoundary } from "@/components/layout/ErrorBoundary";
import "./styles/globals.css";
const queryClient = new QueryClient({
@@ -17,11 +18,13 @@ const queryClient = new QueryClient({
ReactDOM.createRoot(document.getElementById("root")!).render(
<React.StrictMode>
<QueryClientProvider client={queryClient}>
<TooltipProvider>
<App />
<Toaster />
</TooltipProvider>
</QueryClientProvider>
<ErrorBoundary>
<QueryClientProvider client={queryClient}>
<TooltipProvider>
<App />
<Toaster />
</TooltipProvider>
</QueryClientProvider>
</ErrorBoundary>
</React.StrictMode>
);
+2 -2
View File
@@ -170,12 +170,12 @@ describe("AppStore", () => {
it("should handle different tab types", () => {
useAppStore.getState().addTab(makeTab("t1", "query"));
useAppStore.getState().addTab(makeTab("t2", "table"));
useAppStore.getState().addTab(makeTab("t3", "erd"));
useAppStore.getState().addTab(makeTab("t3", "structure"));
expect(useAppStore.getState().tabs.map((t) => t.type)).toEqual([
"query",
"table",
"erd",
"structure",
]);
});
});
+30 -2
View File
@@ -1,5 +1,5 @@
import { create } from "zustand";
import type { ConnectionConfig, DbFlavor, Tab } from "@/types";
import type { ChatMessage, ConnectionConfig, DbFlavor, Tab } from "@/types";
interface AppState {
connections: ConnectionConfig[];
@@ -12,6 +12,8 @@ interface AppState {
activeTabId: string | null;
sidebarWidth: number;
pgVersion: string | null;
chatThreads: Record<string, ChatMessage[]>;
chatPending: Record<string, boolean>;
setConnections: (connections: ConnectionConfig[]) => void;
setActiveConnectionId: (id: string | null) => void;
@@ -27,6 +29,10 @@ interface AppState {
setActiveTabId: (id: string | null) => void;
updateTab: (id: string, updates: Partial<Tab>) => void;
setSidebarWidth: (width: number) => void;
appendChatMessages: (tabId: string, messages: ChatMessage[]) => void;
clearChatThread: (tabId: string) => void;
setChatPending: (tabId: string, pending: boolean) => void;
}
export const useAppStore = create<AppState>((set) => ({
@@ -40,6 +46,8 @@ export const useAppStore = create<AppState>((set) => ({
activeTabId: null,
sidebarWidth: 260,
pgVersion: null,
chatThreads: {},
chatPending: {},
setConnections: (connections) => set({ connections }),
setActiveConnectionId: (id) => set({ activeConnectionId: id }),
@@ -85,7 +93,11 @@ export const useAppStore = create<AppState>((set) => ({
? tabs[tabs.length - 1].id
: null
: state.activeTabId;
return { tabs, activeTabId };
const chatThreads = { ...state.chatThreads };
delete chatThreads[id];
const chatPending = { ...state.chatPending };
delete chatPending[id];
return { tabs, activeTabId, chatThreads, chatPending };
}),
setActiveTabId: (id) => set({ activeTabId: id }),
updateTab: (id, updates) =>
@@ -93,4 +105,20 @@ export const useAppStore = create<AppState>((set) => ({
tabs: state.tabs.map((t) => (t.id === id ? { ...t, ...updates } : t)),
})),
setSidebarWidth: (width) => set({ sidebarWidth: width }),
appendChatMessages: (tabId, messages) =>
set((state) => ({
chatThreads: {
...state.chatThreads,
[tabId]: [...(state.chatThreads[tabId] ?? []), ...messages],
},
})),
clearChatThread: (tabId) =>
set((state) => ({
chatThreads: { ...state.chatThreads, [tabId]: [] },
})),
setChatPending: (tabId, pending) =>
set((state) => ({
chatPending: { ...state.chatPending, [tabId]: pending },
})),
}));
+20 -380
View File
@@ -1,4 +1,4 @@
export type DbFlavor = "postgresql" | "greenplum";
export type DbFlavor = "postgresql" | "greenplum" | "clickhouse";
export interface ConnectResult {
version: string;
@@ -16,6 +16,10 @@ export interface ConnectionConfig {
ssl_mode?: string;
color?: string;
environment?: string;
/** DB engine selected by user. Older configs without this default to "postgresql". */
db_flavor?: DbFlavor;
/** HTTPS for ClickHouse. Defaults to false. */
secure?: boolean;
}
export interface QueryResult {
@@ -122,106 +126,6 @@ export interface ExplainResult {
"Execution Time": number;
}
export interface DatabaseInfo {
name: string;
owner: string;
encoding: string;
collation: string;
ctype: string;
tablespace: string;
connection_limit: number;
size: string;
description: string | null;
}
export interface CreateDatabaseParams {
name: string;
owner?: string;
template?: string;
encoding?: string;
tablespace?: string;
connection_limit?: number;
}
export interface RoleInfo {
name: string;
is_superuser: boolean;
can_login: boolean;
can_create_db: boolean;
can_create_role: boolean;
inherit: boolean;
is_replication: boolean;
connection_limit: number;
password_set: boolean;
valid_until: string | null;
member_of: string[];
members: string[];
description: string | null;
}
export interface CreateRoleParams {
name: string;
password?: string;
login: boolean;
superuser: boolean;
createdb: boolean;
createrole: boolean;
inherit: boolean;
replication: boolean;
connection_limit?: number;
valid_until?: string;
in_roles: string[];
}
export interface AlterRoleParams {
name: string;
password?: string;
login?: boolean;
superuser?: boolean;
createdb?: boolean;
createrole?: boolean;
inherit?: boolean;
replication?: boolean;
connection_limit?: number;
valid_until?: string;
rename_to?: string;
}
export interface TablePrivilege {
grantee: string;
table_schema: string;
table_name: string;
privilege_type: string;
is_grantable: boolean;
}
export interface GrantRevokeParams {
action: string;
privileges: string[];
object_type: string;
object_name: string;
role_name: string;
with_grant_option: boolean;
}
export interface RoleMembershipParams {
action: string;
role_name: string;
member_name: string;
}
export interface SessionInfo {
pid: number;
usename: string | null;
datname: string | null;
state: string | null;
query: string | null;
query_start: string | null;
wait_event_type: string | null;
wait_event: string | null;
client_addr: string | null;
}
export interface SavedQuery {
id: string;
name: string;
@@ -244,46 +148,6 @@ export interface OllamaModel {
name: string;
}
// Entity Lookup
export interface LookupTableMatch {
schema: string;
table: string;
column_type: string;
columns: string[];
types: string[];
rows: unknown[][];
row_count: number;
total_count: number;
}
export interface LookupDatabaseResult {
database: string;
tables: LookupTableMatch[];
error: string | null;
search_time_ms: number;
}
export interface EntityLookupResult {
column_name: string;
value: string;
databases: LookupDatabaseResult[];
total_databases_searched: number;
total_tables_matched: number;
total_rows_found: number;
total_time_ms: number;
}
export interface LookupProgress {
lookup_id: string;
database: string;
status: string;
tables_found: number;
rows_found: number;
error: string | null;
completed: number;
total: number;
}
export interface TriggerInfo {
name: string;
event: string;
@@ -294,52 +158,14 @@ export interface TriggerInfo {
definition: string;
}
export interface ErdColumn {
name: string;
data_type: string;
is_nullable: boolean;
is_primary_key: boolean;
}
export interface ErdTable {
schema: string;
name: string;
columns: ErdColumn[];
}
export interface ErdRelationship {
constraint_name: string;
source_schema: string;
source_table: string;
source_columns: string[];
target_schema: string;
target_table: string;
target_columns: string[];
update_rule: string;
delete_rule: string;
}
export interface ErdData {
tables: ErdTable[];
relationships: ErdRelationship[];
}
// App Settings
export type DockerHost = "local" | "remote";
export interface McpSettings {
enabled: boolean;
port: number;
}
export interface DockerSettings {
host: DockerHost;
remote_url?: string;
}
export interface AppSettings {
mcp: McpSettings;
docker: DockerSettings;
}
export interface McpStatus {
@@ -348,53 +174,7 @@ export interface McpStatus {
running: boolean;
}
// Docker
export interface DockerStatus {
installed: boolean;
daemon_running: boolean;
version: string | null;
error: string | null;
}
export type CloneMode = "schema_only" | "full_clone" | "sample_data";
export interface CloneToDockerParams {
source_connection_id: string;
source_database: string;
container_name: string;
pg_version: string;
host_port: number | null;
clone_mode: CloneMode;
sample_rows: number | null;
postgres_password: string | null;
}
export interface CloneProgress {
clone_id: string;
stage: string;
percent: number;
message: string;
detail: string | null;
}
export interface TuskContainer {
container_id: string;
name: string;
status: string;
host_port: number;
pg_version: string;
source_database: string | null;
source_connection: string | null;
created_at: string | null;
}
export interface CloneResult {
container: TuskContainer;
connection_id: string;
connection_url: string;
}
export type TabType = "query" | "table" | "structure" | "roles" | "sessions" | "lookup" | "erd" | "validation" | "index-advisor" | "snapshots";
export type TabType = "query" | "table" | "structure" | "chat";
export interface Tab {
id: string;
@@ -405,163 +185,23 @@ export interface Tab {
schema?: string;
table?: string;
sql?: string;
roleName?: string;
lookupColumn?: string;
lookupValue?: string;
}
// --- Wave 1: Validation ---
export type ChatRole = "user" | "assistant" | "tool_call" | "tool_result";
export type ValidationStatus = "pending" | "generating" | "running" | "passed" | "failed" | "error";
export interface ValidationRule {
interface ChatBase {
id: string;
description: string;
generated_sql: string;
status: ValidationStatus;
violation_count: number;
sample_violations: unknown[][];
violation_columns: string[];
error: string | null;
created_at: number;
}
export interface ValidationReport {
rules: ValidationRule[];
total_rules: number;
passed: number;
failed: number;
errors: number;
execution_time_ms: number;
}
// --- Wave 2: Data Generator ---
export interface GenerateDataParams {
connection_id: string;
schema: string;
table: string;
row_count: number;
include_related: boolean;
custom_instructions?: string;
}
export interface GeneratedDataPreview {
tables: GeneratedTableData[];
insert_order: string[];
total_rows: number;
}
export interface GeneratedTableData {
schema: string;
table: string;
columns: string[];
rows: unknown[][];
row_count: number;
}
export interface DataGenProgress {
gen_id: string;
stage: string;
percent: number;
message: string;
detail: string | null;
}
// --- Wave 3A: Index Advisor ---
export interface TableStats {
schema: string;
table: string;
seq_scan: number;
idx_scan: number;
n_live_tup: number;
table_size: string;
index_size: string;
}
export interface IndexStatsInfo {
schema: string;
table: string;
index_name: string;
idx_scan: number;
index_size: string;
definition: string;
}
export interface SlowQuery {
query: string;
calls: number;
total_time_ms: number;
mean_time_ms: number;
rows: number;
}
export type IndexRecommendationType = "create_index" | "drop_index" | "replace_index";
export interface IndexRecommendation {
id: string;
recommendation_type: IndexRecommendationType;
table_schema: string;
table_name: string;
index_name: string | null;
ddl: string;
rationale: string;
estimated_impact: string;
priority: string;
}
export interface IndexAdvisorReport {
table_stats: TableStats[];
index_stats: IndexStatsInfo[];
slow_queries: SlowQuery[];
recommendations: IndexRecommendation[];
has_pg_stat_statements: boolean;
}
// --- Wave 3B: Snapshots ---
export interface SnapshotMetadata {
id: string;
name: string;
created_at: string;
connection_name: string;
database: string;
tables: SnapshotTableMeta[];
total_rows: number;
file_size_bytes: number;
version: number;
}
export interface SnapshotTableMeta {
schema: string;
table: string;
row_count: number;
columns: string[];
column_types: string[];
}
export interface SnapshotProgress {
snapshot_id: string;
stage: string;
percent: number;
message: string;
detail: string | null;
}
export interface CreateSnapshotParams {
connection_id: string;
tables: TableRef[];
name: string;
include_dependencies: boolean;
}
export interface TableRef {
schema: string;
table: string;
}
export interface RestoreSnapshotParams {
connection_id: string;
file_path: string;
truncate_before_restore: boolean;
}
export type ChatMessage =
| (ChatBase & { role: "user"; text: string })
| (ChatBase & { role: "assistant"; text: string })
| (ChatBase & { role: "tool_call"; tool: string; input_json: string })
| (ChatBase & {
role: "tool_result";
tool: string;
is_error: boolean;
text?: string | null;
result?: QueryResult | null;
});