feat: add AI data validation, test data generator, index advisor, and snapshots

Four new killer features leveraging AI (Ollama) and PostgreSQL internals:

- Data Validation: describe quality rules in natural language, AI generates
  SQL to find violations, run with pass/fail results and sample violations
- Test Data Generator: right-click table to generate realistic FK-aware test
  data with AI, preview before inserting in a transaction
- Index Advisor: analyze pg_stat tables + AI recommendations for CREATE/DROP
  INDEX with one-click apply
- Data Snapshots: export selected tables to JSON (FK-ordered), restore from
  file with optional truncate in a transaction

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-02-21 13:27:41 +03:00
parent d507162377
commit a3b05b0328
26 changed files with 3438 additions and 17 deletions

View File

@@ -1,15 +1,21 @@
use crate::commands::data::bind_json_value;
use crate::commands::queries::pg_value_to_json;
use crate::error::{TuskError, TuskResult};
use crate::models::ai::{
AiProvider, AiSettings, OllamaChatMessage, OllamaChatRequest, OllamaChatResponse,
OllamaModel, OllamaTagsResponse,
AiProvider, AiSettings, GenerateDataParams, GeneratedDataPreview, GeneratedTableData,
IndexAdvisorReport, IndexRecommendation, IndexStats,
OllamaChatMessage, OllamaChatRequest, OllamaChatResponse, OllamaModel, OllamaTagsResponse,
SlowQuery, TableStats, ValidationRule, ValidationStatus, DataGenProgress,
};
use crate::state::AppState;
use sqlx::Row;
use crate::utils::{escape_ident, topological_sort_tables};
use serde_json::Value;
use sqlx::{Column, Row};
use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::sync::Arc;
use std::time::Duration;
use tauri::{AppHandle, Manager, State};
use std::time::{Duration, Instant};
use tauri::{AppHandle, Emitter, Manager, State};
const MAX_RETRIES: u32 = 2;
const RETRY_DELAY_MS: u64 = 1000;
@@ -386,7 +392,7 @@ pub async fn fix_sql_error(
// Schema context builder
// ---------------------------------------------------------------------------
async fn build_schema_context(
pub(crate) async fn build_schema_context(
state: &AppState,
connection_id: &str,
) -> TuskResult<String> {
@@ -665,16 +671,16 @@ async fn fetch_columns(pool: &sqlx::PgPool) -> TuskResult<Vec<ColumnInfo>> {
.collect())
}
struct ForeignKeyInfo {
schema: String,
table: String,
columns: Vec<String>,
ref_schema: String,
ref_table: String,
ref_columns: Vec<String>,
pub(crate) struct ForeignKeyInfo {
pub(crate) schema: String,
pub(crate) table: String,
pub(crate) columns: Vec<String>,
pub(crate) ref_schema: String,
pub(crate) ref_table: String,
pub(crate) ref_columns: Vec<String>,
}
async fn fetch_foreign_keys_raw(pool: &sqlx::PgPool) -> TuskResult<Vec<ForeignKeyInfo>> {
pub(crate) async fn fetch_foreign_keys_raw(pool: &sqlx::PgPool) -> TuskResult<Vec<ForeignKeyInfo>> {
let rows = sqlx::query(
"SELECT \
cn.nspname AS schema_name, cl.relname AS table_name, \
@@ -1043,3 +1049,609 @@ fn clean_sql_response(raw: &str) -> String {
};
without_fences.trim().to_string()
}
// ---------------------------------------------------------------------------
// Wave 1: AI Data Assertions (Validation)
// ---------------------------------------------------------------------------
#[tauri::command]
pub async fn generate_validation_sql(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
rule_description: String,
) -> TuskResult<String> {
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are an expert PostgreSQL data quality validator. Given a database schema and a natural \
language data quality rule, generate a SELECT query that finds ALL rows violating the rule.\n\
\n\
OUTPUT FORMAT:\n\
- Raw SQL only. No explanations, no markdown code fences, no comments.\n\
- The query MUST be a SELECT statement.\n\
- Return violating rows with enough context columns to identify them.\n\
\n\
VALIDATION PATTERNS:\n\
- NULL checks: SELECT * FROM table WHERE required_column IS NULL\n\
- Format checks: WHERE column !~ 'pattern'\n\
- Range checks: WHERE column < min OR column > max\n\
- FK integrity: LEFT JOIN parent ON ... WHERE parent.id IS NULL\n\
- Uniqueness: GROUP BY ... HAVING COUNT(*) > 1\n\
- Date consistency: WHERE start_date > end_date\n\
- Enum validity: WHERE column NOT IN ('val1', 'val2', ...)\n\
\n\
ONLY reference tables and columns that exist in the schema.\n\
\n\
{}\n",
schema_text
);
let raw = call_ollama_chat(&app, &state, system_prompt, rule_description).await?;
Ok(clean_sql_response(&raw))
}
#[tauri::command]
pub async fn run_validation_rule(
state: State<'_, Arc<AppState>>,
connection_id: String,
sql: String,
sample_limit: Option<u32>,
) -> TuskResult<ValidationRule> {
let sql_upper = sql.trim().to_uppercase();
if !sql_upper.starts_with("SELECT") {
return Err(TuskError::Custom(
"Validation query must be a SELECT statement".to_string(),
));
}
let pool = state.get_pool(&connection_id).await?;
let limit = sample_limit.unwrap_or(10);
let _start = Instant::now();
let mut tx = (&pool).begin().await.map_err(TuskError::Database)?;
sqlx::query("SET TRANSACTION READ ONLY")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
sqlx::query("SET statement_timeout = '30s'")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
// Count violations
let count_sql = format!("SELECT COUNT(*) FROM ({}) AS _v", sql);
let count_row = sqlx::query(&count_sql)
.fetch_one(&mut *tx)
.await
.map_err(TuskError::Database)?;
let violation_count: i64 = count_row.get(0);
// Sample violations
let sample_sql = format!("SELECT * FROM ({}) AS _v LIMIT {}", sql, limit);
let sample_rows = sqlx::query(&sample_sql)
.fetch_all(&mut *tx)
.await
.map_err(TuskError::Database)?;
tx.rollback().await.map_err(TuskError::Database)?;
let mut violation_columns = Vec::new();
let mut sample_violations = Vec::new();
if let Some(first) = sample_rows.first() {
for col in first.columns() {
violation_columns.push(col.name().to_string());
}
}
for row in &sample_rows {
let vals: Vec<Value> = (0..violation_columns.len())
.map(|i| pg_value_to_json(row, i))
.collect();
sample_violations.push(vals);
}
let status = if violation_count > 0 {
ValidationStatus::Failed
} else {
ValidationStatus::Passed
};
Ok(ValidationRule {
id: String::new(),
description: String::new(),
generated_sql: sql,
status,
violation_count: violation_count as u64,
sample_violations,
violation_columns,
error: None,
})
}
#[tauri::command]
pub async fn suggest_validation_rules(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<String>> {
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are a data quality expert. Given a database schema, suggest 5-10 data quality \
validation rules as natural language descriptions.\n\
\n\
OUTPUT FORMAT:\n\
- Return ONLY a JSON array of strings, each string being a validation rule.\n\
- No markdown, no explanations, no code fences.\n\
- Example: [\"All users must have a non-empty email address\", \"Order total must be positive\"]\n\
\n\
RULE CATEGORIES TO COVER:\n\
- NOT NULL checks for critical columns\n\
- Business logic (positive amounts, valid ranges, consistent dates)\n\
- Referential integrity (orphaned foreign keys)\n\
- Format validation (emails, phone numbers, codes)\n\
- Enum/status field validity\n\
- Date consistency (start before end, not in future where inappropriate)\n\
\n\
{}\n",
schema_text
);
let raw = call_ollama_chat(&app, &state, system_prompt, "Suggest validation rules".to_string()).await?;
let cleaned = raw.trim();
let json_start = cleaned.find('[').unwrap_or(0);
let json_end = cleaned.rfind(']').map(|i| i + 1).unwrap_or(cleaned.len());
let json_str = &cleaned[json_start..json_end];
let rules: Vec<String> = serde_json::from_str(json_str).map_err(|e| {
TuskError::Ai(format!("Failed to parse AI response as JSON array: {}. Response: {}", e, cleaned))
})?;
Ok(rules)
}
// ---------------------------------------------------------------------------
// Wave 2: AI Data Generator
// ---------------------------------------------------------------------------
#[tauri::command]
pub async fn generate_test_data_preview(
app: AppHandle,
state: State<'_, Arc<AppState>>,
params: GenerateDataParams,
gen_id: String,
) -> TuskResult<GeneratedDataPreview> {
let pool = state.get_pool(&params.connection_id).await?;
let _ = app.emit("datagen-progress", DataGenProgress {
gen_id: gen_id.clone(),
stage: "schema".to_string(),
percent: 10,
message: "Building schema context...".to_string(),
detail: None,
});
let schema_text = build_schema_context(&state, &params.connection_id).await?;
// Get FK info for topological sort
let fk_rows = fetch_foreign_keys_raw(&pool).await?;
let mut target_tables = vec![(params.schema.clone(), params.table.clone())];
if params.include_related {
// Add parent tables (tables referenced by FKs from target)
for fk in &fk_rows {
if fk.schema == params.schema && fk.table == params.table {
let parent = (fk.ref_schema.clone(), fk.ref_table.clone());
if !target_tables.contains(&parent) {
target_tables.push(parent);
}
}
}
}
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 insert_order: Vec<String> = sorted_tables
.iter()
.map(|(s, t)| format!("{}.{}", s, t))
.collect();
let row_count = params.row_count.min(1000);
let _ = app.emit("datagen-progress", DataGenProgress {
gen_id: gen_id.clone(),
stage: "generating".to_string(),
percent: 30,
message: "AI is generating test data...".to_string(),
detail: None,
});
let tables_desc: Vec<String> = sorted_tables
.iter()
.map(|(s, t)| {
let count = if s == &params.schema && t == &params.table {
row_count
} else {
(row_count / 3).max(1)
};
format!("{}.{}: {} rows", s, t, count)
})
.collect();
let custom = params
.custom_instructions
.as_deref()
.unwrap_or("Generate realistic sample data");
let system_prompt = format!(
"You are a PostgreSQL test data generator. Generate realistic test data as JSON.\n\
\n\
OUTPUT FORMAT:\n\
- Return ONLY a JSON object where keys are \"schema.table\" and values are arrays of row objects.\n\
- Each row object has column names as keys and values matching the column types.\n\
- No markdown, no explanations, no code fences.\n\
- Example: {{\"public.users\": [{{\"name\": \"Alice\", \"email\": \"alice@example.com\"}}]}}\n\
\n\
RULES:\n\
1. Respect column types exactly (text, integer, boolean, timestamp, uuid, etc.)\n\
2. Use valid foreign key values - parent tables are generated first, reference their IDs\n\
3. Respect enum types - use only valid enum values\n\
4. Omit auto-increment/serial/identity columns (they have DEFAULT auto-increment)\n\
5. Generate realistic data: real names, valid emails, plausible dates, etc.\n\
6. Respect NOT NULL constraints\n\
7. For UUID columns, generate valid UUIDs\n\
8. For timestamp columns, use ISO 8601 format\n\
\n\
Tables to generate (in this exact order):\n\
{}\n\
\n\
Custom instructions: {}\n\
\n\
{}\n",
tables_desc.join("\n"),
custom,
schema_text
);
let raw = call_ollama_chat(
&app,
&state,
system_prompt,
format!("Generate test data for {} tables", sorted_tables.len()),
)
.await?;
let _ = app.emit("datagen-progress", DataGenProgress {
gen_id: gen_id.clone(),
stage: "parsing".to_string(),
percent: 80,
message: "Parsing generated data...".to_string(),
detail: None,
});
// Parse JSON response
let cleaned = raw.trim();
let json_start = cleaned.find('{').unwrap_or(0);
let json_end = cleaned.rfind('}').map(|i| i + 1).unwrap_or(cleaned.len());
let json_str = &cleaned[json_start..json_end];
let data_map: HashMap<String, Vec<HashMap<String, Value>>> =
serde_json::from_str(json_str).map_err(|e| {
TuskError::Ai(format!("Failed to parse generated data: {}. Response: {}", e, &cleaned[..cleaned.len().min(500)]))
})?;
let mut tables = Vec::new();
let mut total_rows: u32 = 0;
for (schema, table) in &sorted_tables {
let key = format!("{}.{}", schema, table);
if let Some(rows_data) = data_map.get(&key) {
let columns: Vec<String> = if let Some(first) = rows_data.first() {
first.keys().cloned().collect()
} else {
Vec::new()
};
let rows: Vec<Vec<Value>> = rows_data
.iter()
.map(|row_map| columns.iter().map(|col| row_map.get(col).cloned().unwrap_or(Value::Null)).collect())
.collect();
let count = rows.len() as u32;
total_rows += count;
tables.push(GeneratedTableData {
schema: schema.clone(),
table: table.clone(),
columns,
rows,
row_count: count,
});
}
}
let _ = app.emit("datagen-progress", DataGenProgress {
gen_id: gen_id.clone(),
stage: "done".to_string(),
percent: 100,
message: "Data generation complete".to_string(),
detail: Some(format!("{} rows across {} tables", total_rows, tables.len())),
});
Ok(GeneratedDataPreview {
tables,
insert_order,
total_rows,
})
}
#[tauri::command]
pub async fn insert_generated_data(
state: State<'_, Arc<AppState>>,
connection_id: String,
preview: GeneratedDataPreview,
) -> TuskResult<u64> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let pool = state.get_pool(&connection_id).await?;
let mut tx = (&pool).begin().await.map_err(TuskError::Database)?;
// Defer constraints for circular FKs
sqlx::query("SET CONSTRAINTS ALL DEFERRED")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
let mut total_inserted: u64 = 0;
for table_data in &preview.tables {
if table_data.columns.is_empty() || table_data.rows.is_empty() {
continue;
}
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(", ")
);
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)?;
// Invalidate schema cache since data changed
state.invalidate_schema_cache(&connection_id).await;
Ok(total_inserted)
}
// ---------------------------------------------------------------------------
// Wave 3A: Smart Index Advisor
// ---------------------------------------------------------------------------
#[tauri::command]
pub async fn get_index_advisor_report(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<IndexAdvisorReport> {
let pool = state.get_pool(&connection_id).await?;
// Fetch table stats
let table_stats_rows = sqlx::query(
"SELECT schemaname, relname, seq_scan, idx_scan, n_live_tup, \
pg_size_pretty(pg_total_relation_size(schemaname || '.' || relname)) AS table_size, \
pg_size_pretty(pg_indexes_size(quote_ident(schemaname) || '.' || quote_ident(relname))) AS index_size \
FROM pg_stat_user_tables \
ORDER BY seq_scan DESC \
LIMIT 50"
)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
let table_stats: Vec<TableStats> = table_stats_rows
.iter()
.map(|r| TableStats {
schema: r.get(0),
table: r.get(1),
seq_scan: r.get(2),
idx_scan: r.get(3),
n_live_tup: r.get(4),
table_size: r.get(5),
index_size: r.get(6),
})
.collect();
// Fetch index stats
let index_stats_rows = sqlx::query(
"SELECT schemaname, relname, indexrelname, idx_scan, \
pg_size_pretty(pg_relation_size(indexrelid)) AS index_size, \
pg_get_indexdef(indexrelid) AS definition \
FROM pg_stat_user_indexes \
ORDER BY idx_scan ASC \
LIMIT 50"
)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
let index_stats: Vec<IndexStats> = index_stats_rows
.iter()
.map(|r| IndexStats {
schema: r.get(0),
table: r.get(1),
index_name: r.get(2),
idx_scan: r.get(3),
index_size: r.get(4),
definition: r.get(5),
})
.collect();
// Fetch slow queries (graceful if pg_stat_statements not available)
let (slow_queries, has_pg_stat_statements) = match sqlx::query(
"SELECT query, calls, total_exec_time, mean_exec_time, rows \
FROM pg_stat_statements \
WHERE calls > 0 \
ORDER BY mean_exec_time DESC \
LIMIT 20"
)
.fetch_all(&pool)
.await
{
Ok(rows) => {
let queries: Vec<SlowQuery> = rows
.iter()
.map(|r| SlowQuery {
query: r.get(0),
calls: r.get(1),
total_time_ms: r.get(2),
mean_time_ms: r.get(3),
rows: r.get(4),
})
.collect();
(queries, true)
}
Err(_) => (Vec::new(), false),
};
// Build AI prompt for recommendations
let schema_text = build_schema_context(&state, &connection_id).await?;
let mut stats_text = String::from("TABLE STATISTICS:\n");
for ts in &table_stats {
stats_text.push_str(&format!(
" {}.{}: seq_scan={}, idx_scan={}, rows={}, size={}, idx_size={}\n",
ts.schema, ts.table, ts.seq_scan, ts.idx_scan, ts.n_live_tup, ts.table_size, ts.index_size
));
}
stats_text.push_str("\nINDEX STATISTICS:\n");
for is in &index_stats {
stats_text.push_str(&format!(
" {}.{}.{}: scans={}, size={}, def={}\n",
is.schema, is.table, is.index_name, is.idx_scan, is.index_size, is.definition
));
}
if !slow_queries.is_empty() {
stats_text.push_str("\nSLOW QUERIES:\n");
for sq in &slow_queries {
stats_text.push_str(&format!(
" calls={}, mean={:.1}ms, total={:.1}ms, rows={}: {}\n",
sq.calls, sq.mean_time_ms, sq.total_time_ms, sq.rows,
sq.query.chars().take(200).collect::<String>()
));
}
}
let system_prompt = format!(
"You are a PostgreSQL performance expert. Analyze the database statistics and recommend index changes.\n\
\n\
OUTPUT FORMAT:\n\
- Return ONLY a JSON array of recommendation objects.\n\
- No markdown, no explanations, no code fences.\n\
- Each object: {{\"recommendation_type\": \"create_index\"|\"drop_index\"|\"replace_index\", \
\"table_schema\": \"...\", \"table_name\": \"...\", \"index_name\": \"...\"|null, \
\"ddl\": \"CREATE INDEX CONCURRENTLY ...\", \"rationale\": \"...\", \
\"estimated_impact\": \"high\"|\"medium\"|\"low\", \"priority\": \"high\"|\"medium\"|\"low\"}}\n\
\n\
RULES:\n\
1. Prefer CREATE INDEX CONCURRENTLY to avoid locking\n\
2. Never suggest dropping PRIMARY KEY or UNIQUE indexes\n\
3. Suggest dropping indexes with 0 scans on tables with many rows\n\
4. Suggest composite indexes for commonly co-filtered columns\n\
5. Suggest partial indexes for low-cardinality boolean columns\n\
6. Consider covering indexes for frequently selected columns\n\
7. High seq_scan + high row count = strong candidate for new index\n\
\n\
{}\n\
\n\
{}\n",
stats_text, schema_text
);
let raw = call_ollama_chat(
&app,
&state,
system_prompt,
"Analyze indexes and provide recommendations".to_string(),
)
.await?;
let cleaned = raw.trim();
let json_start = cleaned.find('[').unwrap_or(0);
let json_end = cleaned.rfind(']').map(|i| i + 1).unwrap_or(cleaned.len());
let json_str = &cleaned[json_start..json_end];
let recommendations: Vec<IndexRecommendation> =
serde_json::from_str(json_str).unwrap_or_default();
Ok(IndexAdvisorReport {
table_stats,
index_stats,
slow_queries,
recommendations,
has_pg_stat_statements,
})
}
#[tauri::command]
pub async fn apply_index_recommendation(
state: State<'_, Arc<AppState>>,
connection_id: String,
ddl: String,
) -> TuskResult<()> {
if state.is_read_only(&connection_id).await {
return Err(TuskError::ReadOnly);
}
let ddl_upper = ddl.trim().to_uppercase();
if !ddl_upper.starts_with("CREATE INDEX") && !ddl_upper.starts_with("DROP INDEX") {
return Err(TuskError::Custom(
"Only CREATE INDEX and DROP INDEX statements are allowed".to_string(),
));
}
let pool = state.get_pool(&connection_id).await?;
// CONCURRENTLY cannot run inside a transaction, execute directly
sqlx::query(&ddl)
.execute(&pool)
.await
.map_err(TuskError::Database)?;
// Invalidate schema cache
state.invalidate_schema_cache(&connection_id).await;
Ok(())
}

View File

@@ -283,7 +283,7 @@ pub async fn delete_rows(
Ok(total_affected)
}
fn bind_json_value<'q>(
pub(crate) fn bind_json_value<'q>(
query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
value: &'q Value,
) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {

View File

@@ -9,4 +9,5 @@ pub mod management;
pub mod queries;
pub mod saved_queries;
pub mod schema;
pub mod snapshot;
pub mod settings;

View File

@@ -0,0 +1,349 @@
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();
if params.include_dependencies {
let fk_rows = fetch_foreign_keys_raw(&pool).await?;
for fk in &fk_rows {
for (schema, table) in &params.tables.iter().map(|t| (t.schema.clone(), t.table.clone())).collect::<Vec<_>>() {
if &fk.schema == schema && &fk.table == 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_rows = fetch_foreign_keys_raw(&pool).await?;
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 as u8) * 80 / total_tables.max(1) 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 as u8) * 75 / total_tables.max(1) 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::Custom(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)
}

View File

@@ -135,6 +135,18 @@ 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

View File

@@ -57,3 +57,137 @@ 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>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationReport {
pub rules: Vec<ValidationRule>,
pub total_rules: usize,
pub passed: usize,
pub failed: usize,
pub errors: usize,
pub execution_time_ms: u128,
}
// --- 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 {
CreateIndex,
DropIndex,
ReplaceIndex,
}
#[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,
}

View File

@@ -8,3 +8,4 @@ pub mod query_result;
pub mod saved_queries;
pub mod schema;
pub mod settings;
pub mod snapshot;

View File

@@ -0,0 +1,68 @@
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,
}

View File

@@ -1,3 +1,75 @@
use std::collections::{HashMap, HashSet};
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 queue: Vec<(String, String)> = in_degree
.iter()
.filter(|(_, &deg)| deg == 0)
.map(|(k, _)| k.clone())
.collect();
queue.sort(); // deterministic order
let mut result = Vec::new();
while let Some(node) = queue.pop() {
result.push(node.clone());
if let Some(neighbors) = graph.get(&node) {
for neighbor in neighbors {
if let Some(deg) = in_degree.get_mut(neighbor) {
*deg -= 1;
if *deg == 0 {
queue.push(neighbor.clone());
queue.sort();
}
}
}
}
}
// Add any remaining tables (cycles) at the end
for t in target_tables {
if !result.contains(t) {
result.push(t.clone());
}
}
result
}

View File

@@ -0,0 +1,295 @@
import { useState, useEffect } 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();
useEffect(() => {
if (open) {
setStep("config");
setRowCount(10);
setIncludeRelated(true);
setCustomInstructions("");
reset();
}
}, [open, 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>
);
}

View File

@@ -0,0 +1,232 @@
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>
);
}

View File

@@ -0,0 +1,86 @@
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>
);
}

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 } from "lucide-react";
import { X, Table2, Code, Columns, Users, Activity, Search, GitFork, ShieldCheck, Gauge, Camera } from "lucide-react";
export function TabBar() {
const { tabs, activeTabId, setActiveTabId, closeTab } = useAppStore();
@@ -17,6 +17,9 @@ export function TabBar() {
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" />,
};
return (

View File

@@ -32,6 +32,7 @@ import {
} 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 {
@@ -246,6 +247,55 @@ function DatabaseNode({
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}
@@ -415,6 +465,7 @@ function CategoryNode({
}) {
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,
@@ -489,6 +540,13 @@ function CategoryNode({
>
View Structure
</ContextMenuItem>
{category === "tables" && (
<ContextMenuItem
onClick={() => setDataGenTarget(item.name)}
>
Generate Test Data
</ContextMenuItem>
)}
<ContextMenuSeparator />
<ContextMenuItem
onClick={() => setPrivilegesTarget(item.name)}
@@ -524,6 +582,15 @@ function CategoryNode({
table={privilegesTarget}
/>
)}
{dataGenTarget && (
<GenerateDataDialog
open={!!dataGenTarget}
onOpenChange={(open) => !open && setDataGenTarget(null)}
connectionId={connectionId}
schema={schema}
table={dataGenTarget}
/>
)}
</div>
);
}

View File

@@ -0,0 +1,271 @@
import { useState, useEffect } 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();
useEffect(() => {
if (open) {
setStep("config");
setName(`snapshot-${new Date().toISOString().slice(0, 10)}`);
setSelectedTables(new Set());
setIncludeDeps(true);
reset();
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [open, reset]);
useEffect(() => {
if (schemas && schemas.length > 0 && !selectedSchema) {
setSelectedSchema(schemas.find((s) => s === "public") || schemas[0]);
}
}, [schemas, selectedSchema]);
useEffect(() => {
if (progress?.stage === "done" || progress?.stage === "error") {
setStep("done");
}
}, [progress]);
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>
);
}

View File

@@ -0,0 +1,244 @@
import { useState, useEffect } 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();
useEffect(() => {
if (open) {
setStep("select");
setFilePath(null);
setMetadata(null);
setTruncate(false);
reset();
}
}, [open, reset]);
useEffect(() => {
if (progress?.stage === "done" || progress?.stage === "error") {
setStep("done");
}
}, [progress]);
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>
);
}

View File

@@ -0,0 +1,122 @@
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>
);
}

View File

@@ -0,0 +1,216 @@
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>
);
}

View File

@@ -0,0 +1,138 @@
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>
);
}

View File

@@ -6,6 +6,9 @@ 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";
export function TabContent() {
const { tabs, activeTabId, updateTab } = useAppStore();
@@ -81,6 +84,27 @@ export function TabContent() {
/>
);
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}
/>
);
break;
default:
content = null;
}

View File

@@ -0,0 +1,73 @@
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(() => {
const unlistenPromise = onDataGenProgress((p) => {
if (p.gen_id === genIdRef.current) {
setProgress(p);
}
});
return () => {
unlistenPromise.then((unlisten) => unlisten());
};
}, []);
const previewRef = useRef(previewMutation);
previewRef.current = previewMutation;
const insertRef = useRef(insertMutation);
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,
};
}

View File

@@ -0,0 +1,20 @@
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),
});
}

131
src/hooks/use-snapshots.ts Normal file
View File

@@ -0,0 +1,131 @@
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(() => {
const unlistenPromise = onSnapshotProgress((p) => {
if (p.snapshot_id === snapshotIdRef.current) {
setProgress(p);
}
});
return () => {
unlistenPromise.then((unlisten) => unlisten());
};
}, []);
const mutationRef = useRef(mutation);
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(() => {
const unlistenPromise = onSnapshotProgress((p) => {
if (p.snapshot_id === snapshotIdRef.current) {
setProgress(p);
}
});
return () => {
unlistenPromise.then((unlisten) => unlisten());
};
}, []);
const mutationRef = useRef(mutation);
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,
};
}

View File

@@ -0,0 +1,38 @@
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),
});
}

View File

@@ -35,6 +35,15 @@ import type {
TuskContainer,
AppSettings,
McpStatus,
ValidationRule,
GenerateDataParams,
GeneratedDataPreview,
DataGenProgress,
IndexAdvisorReport,
SnapshotMetadata,
CreateSnapshotParams,
RestoreSnapshotParams,
SnapshotProgress,
} from "@/types";
// Connections
@@ -334,3 +343,50 @@ 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));

View File

@@ -394,7 +394,7 @@ export interface CloneResult {
connection_url: string;
}
export type TabType = "query" | "table" | "structure" | "roles" | "sessions" | "lookup" | "erd";
export type TabType = "query" | "table" | "structure" | "roles" | "sessions" | "lookup" | "erd" | "validation" | "index-advisor" | "snapshots";
export interface Tab {
id: string;
@@ -409,3 +409,159 @@ export interface Tab {
lookupColumn?: string;
lookupValue?: string;
}
// --- Wave 1: Validation ---
export type ValidationStatus = "pending" | "generating" | "running" | "passed" | "failed" | "error";
export interface ValidationRule {
id: string;
description: string;
generated_sql: string;
status: ValidationStatus;
violation_count: number;
sample_violations: unknown[][];
violation_columns: string[];
error: string | null;
}
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;
}