tools/tusk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
Alexey Shakhmatov
2026-05-06 19:30:44 +03:00
parent 652937f7f5
commit 4f7afc17f4
89 changed files with 4151 additions and 10756 deletions
Download Patch File Download Diff File Expand all files Collapse all files

321
src-tauri/src/commands/schema.rs
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,
})
}