From a3b05b03281581e6a3436d35ff5284fc53597c83 Mon Sep 17 00:00:00 2001 From: "A.Shakhmatov" Date: Sat, 21 Feb 2026 13:27:41 +0300 Subject: [PATCH] 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 --- src-tauri/src/commands/ai.rs | 640 +++++++++++++++++- src-tauri/src/commands/data.rs | 2 +- src-tauri/src/commands/mod.rs | 1 + src-tauri/src/commands/snapshot.rs | 349 ++++++++++ src-tauri/src/lib.rs | 12 + src-tauri/src/models/ai.rs | 134 ++++ src-tauri/src/models/mod.rs | 1 + src-tauri/src/models/snapshot.rs | 68 ++ src-tauri/src/utils.rs | 72 ++ .../data-generator/GenerateDataDialog.tsx | 295 ++++++++ .../index-advisor/IndexAdvisorPanel.tsx | 232 +++++++ .../index-advisor/RecommendationCard.tsx | 86 +++ src/components/layout/TabBar.tsx | 5 +- src/components/schema/SchemaTree.tsx | 67 ++ .../snapshots/CreateSnapshotDialog.tsx | 271 ++++++++ .../snapshots/RestoreSnapshotDialog.tsx | 244 +++++++ src/components/snapshots/SnapshotPanel.tsx | 122 ++++ src/components/validation/ValidationPanel.tsx | 216 ++++++ .../validation/ValidationRuleCard.tsx | 138 ++++ src/components/workspace/TabContent.tsx | 24 + src/hooks/use-data-generator.ts | 73 ++ src/hooks/use-index-advisor.ts | 20 + src/hooks/use-snapshots.ts | 131 ++++ src/hooks/use-validation.ts | 38 ++ src/lib/tauri.ts | 56 ++ src/types/index.ts | 158 ++++- 26 files changed, 3438 insertions(+), 17 deletions(-) create mode 100644 src-tauri/src/commands/snapshot.rs create mode 100644 src-tauri/src/models/snapshot.rs create mode 100644 src/components/data-generator/GenerateDataDialog.tsx create mode 100644 src/components/index-advisor/IndexAdvisorPanel.tsx create mode 100644 src/components/index-advisor/RecommendationCard.tsx create mode 100644 src/components/snapshots/CreateSnapshotDialog.tsx create mode 100644 src/components/snapshots/RestoreSnapshotDialog.tsx create mode 100644 src/components/snapshots/SnapshotPanel.tsx create mode 100644 src/components/validation/ValidationPanel.tsx create mode 100644 src/components/validation/ValidationRuleCard.tsx create mode 100644 src/hooks/use-data-generator.ts create mode 100644 src/hooks/use-index-advisor.ts create mode 100644 src/hooks/use-snapshots.ts create mode 100644 src/hooks/use-validation.ts diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index 8fe6eb0..99eca49 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -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 { @@ -665,16 +671,16 @@ async fn fetch_columns(pool: &sqlx::PgPool) -> TuskResult> { .collect()) } -struct ForeignKeyInfo { - schema: String, - table: String, - columns: Vec, - ref_schema: String, - ref_table: String, - ref_columns: Vec, +pub(crate) struct ForeignKeyInfo { + pub(crate) schema: String, + pub(crate) table: String, + pub(crate) columns: Vec, + pub(crate) ref_schema: String, + pub(crate) ref_table: String, + pub(crate) ref_columns: Vec, } -async fn fetch_foreign_keys_raw(pool: &sqlx::PgPool) -> TuskResult> { +pub(crate) async fn fetch_foreign_keys_raw(pool: &sqlx::PgPool) -> TuskResult> { 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>, + connection_id: String, + rule_description: String, +) -> TuskResult { + 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>, + connection_id: String, + sql: String, + sample_limit: Option, +) -> TuskResult { + 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 = (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>, + connection_id: String, +) -> TuskResult> { + 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 = 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>, + params: GenerateDataParams, + gen_id: String, +) -> TuskResult { + let pool = state.get_pool(¶ms.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, ¶ms.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 = 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 = sorted_tables + .iter() + .map(|(s, t)| { + let count = if s == ¶ms.schema && t == ¶ms.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>> = + 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 = if let Some(first) = rows_data.first() { + first.keys().cloned().collect() + } else { + Vec::new() + }; + + let rows: Vec> = 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>, + connection_id: String, + preview: GeneratedDataPreview, +) -> TuskResult { + 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 = table_data.columns.iter().map(|c| escape_ident(c)).collect(); + let placeholders: Vec = (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>, + connection_id: String, +) -> TuskResult { + 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 = 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 = 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 = 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::() + )); + } + } + + 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 = + 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>, + 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(()) +} diff --git a/src-tauri/src/commands/data.rs b/src-tauri/src/commands/data.rs index 51e41cf..9a6aca6 100644 --- a/src-tauri/src/commands/data.rs +++ b/src-tauri/src/commands/data.rs @@ -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> { diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index 7cf60e4..b75c8f9 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -9,4 +9,5 @@ pub mod management; pub mod queries; pub mod saved_queries; pub mod schema; +pub mod snapshot; pub mod settings; diff --git a/src-tauri/src/commands/snapshot.rs b/src-tauri/src/commands/snapshot.rs new file mode 100644 index 0000000..8f93eaf --- /dev/null +++ b/src-tauri/src/commands/snapshot.rs @@ -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>, + params: CreateSnapshotParams, + snapshot_id: String, + file_path: String, +) -> TuskResult { + let pool = state.get_pool(¶ms.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 ¶ms.tables.iter().map(|t| (t.schema.clone(), t.table.clone())).collect::>() { + 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 = Vec::new(); + let mut table_metas: Vec = 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> = 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>, + params: RestoreSnapshotParams, + snapshot_id: String, +) -> TuskResult { + if state.is_read_only(¶ms.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(¶ms.file_path)?; + let snapshot: Snapshot = serde_json::from_str(&data)?; + + let pool = state.get_pool(¶ms.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 = table_data.columns.iter().map(|c| escape_ident(c)).collect(); + let placeholders: Vec = (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(¶ms.connection_id).await; + + Ok(total_inserted) +} + +#[tauri::command] +pub async fn list_snapshots(app: AppHandle) -> TuskResult> { + 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::(&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 { + 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) +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 6a794b5..e0dc796 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -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 diff --git a/src-tauri/src/models/ai.rs b/src-tauri/src/models/ai.rs index 3b312dc..b26d3a0 100644 --- a/src-tauri/src/models/ai.rs +++ b/src-tauri/src/models/ai.rs @@ -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>, + pub violation_columns: Vec, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ValidationReport { + pub rules: Vec, + 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, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GeneratedDataPreview { + pub tables: Vec, + pub insert_order: Vec, + pub total_rows: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GeneratedTableData { + pub schema: String, + pub table: String, + pub columns: Vec, + pub rows: Vec>, + 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, +} + +// --- 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, + 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, + pub index_stats: Vec, + pub slow_queries: Vec, + pub recommendations: Vec, + pub has_pg_stat_statements: bool, +} diff --git a/src-tauri/src/models/mod.rs b/src-tauri/src/models/mod.rs index e0815a1..988c820 100644 --- a/src-tauri/src/models/mod.rs +++ b/src-tauri/src/models/mod.rs @@ -8,3 +8,4 @@ pub mod query_result; pub mod saved_queries; pub mod schema; pub mod settings; +pub mod snapshot; diff --git a/src-tauri/src/models/snapshot.rs b/src-tauri/src/models/snapshot.rs new file mode 100644 index 0000000..9089051 --- /dev/null +++ b/src-tauri/src/models/snapshot.rs @@ -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, + 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, + pub column_types: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Snapshot { + pub metadata: SnapshotMetadata, + pub tables: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SnapshotTableData { + pub schema: String, + pub table: String, + pub columns: Vec, + pub column_types: Vec, + pub rows: Vec>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SnapshotProgress { + pub snapshot_id: String, + pub stage: String, + pub percent: u8, + pub message: String, + pub detail: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CreateSnapshotParams { + pub connection_id: String, + pub tables: Vec, + 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, +} diff --git a/src-tauri/src/utils.rs b/src-tauri/src/utils.rs index 4bb82d1..df85413 100644 --- a/src-tauri/src/utils.rs +++ b/src-tauri/src/utils.rs @@ -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 == 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 +} diff --git a/src/components/data-generator/GenerateDataDialog.tsx b/src/components/data-generator/GenerateDataDialog.tsx new file mode 100644 index 0000000..0d8a606 --- /dev/null +++ b/src/components/data-generator/GenerateDataDialog.tsx @@ -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("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 ( + + + + + + Generate Test Data + + + + {step === "config" && ( + <> +
+
+ +
+ {schema}.{table} +
+
+ +
+ + setRowCount(Math.min(1000, Math.max(1, parseInt(e.target.value) || 1)))} + min={1} + max={1000} + /> +
+ +
+ +
+ setIncludeRelated(e.target.checked)} + className="rounded" + /> + + Include parent tables (via foreign keys) + +
+
+ +
+ + setCustomInstructions(e.target.value)} + /> +
+
+ + {isGenerating && progress && ( +
+
+ {progress.message} + {progress.percent}% +
+
+
+
+
+ )} + + + + + + + )} + + {step === "preview" && preview && ( + <> +
+
+ Preview: + {preview.total_rows} rows across {preview.tables.length} tables +
+ + {preview.tables.map((tbl) => ( +
+
+ + {tbl.schema}.{tbl.table} + {tbl.row_count} rows +
+
+ + + + {tbl.columns.map((col) => ( + + ))} + + + + {tbl.rows.slice(0, 5).map((row, i) => ( + + {(row as unknown[]).map((val, j) => ( + + ))} + + ))} + {tbl.rows.length > 5 && ( + + + + )} + +
+ {col} +
+ {val === null ? ( + NULL + ) : ( + String(val).substring(0, 50) + )} +
+ ...and {tbl.rows.length - 5} more rows +
+
+
+ ))} +
+ + + + + + + )} + + {step === "done" && ( +
+ {insertError ? ( +
+ +
+

Insert Failed

+

{insertError}

+
+
+ ) : ( +
+ +
+

Data Generated Successfully

+

+ {insertedRows} rows inserted across {preview?.tables.length ?? 0} tables. +

+
+
+ )} + + + + {insertError && ( + + )} + +
+ )} + + {generateError && step === "config" && ( +
+ +

{generateError}

+
+ )} + +
+ ); +} diff --git a/src/components/index-advisor/IndexAdvisorPanel.tsx b/src/components/index-advisor/IndexAdvisorPanel.tsx new file mode 100644 index 0000000..cb6c727 --- /dev/null +++ b/src/components/index-advisor/IndexAdvisorPanel.tsx @@ -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(null); + const [appliedDdls, setAppliedDdls] = useState>(new Set()); + const [applyingDdl, setApplyingDdl] = useState(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 ( +
+ {/* Header */} +
+
+ +

Index Advisor

+
+ +
+ + {/* Content */} +
+ {!report ? ( +
+ Click Analyze to scan your database for index optimization opportunities. +
+ ) : ( + +
+ + + Recommendations + {report.recommendations.length > 0 && ( + {report.recommendations.length} + )} + + Table Stats + Index Stats + + Slow Queries + {!report.has_pg_stat_statements && ( + + )} + + +
+ + + {report.recommendations.length === 0 ? ( +
+ No recommendations found. Your indexes look good! +
+ ) : ( + report.recommendations.map((rec, i) => ( + + )) + )} + + + +
+ + + + + + + + + + + + + {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 ( + + + + + + + + + ); + })} + +
TableSeq ScansIdx ScansRowsTable SizeIndex Size
{ts.schema}.{ts.table} 0.8 && ts.n_live_tup > 1000 ? "text-destructive font-medium" : ""}`}> + {ts.seq_scan.toLocaleString()} + {ts.idx_scan.toLocaleString()}{ts.n_live_tup.toLocaleString()}{ts.table_size}{ts.index_size}
+
+ + + +
+ + + + + + + + + + + + {report.index_stats.map((is) => ( + + + + + + + + ))} + +
IndexTableScansSizeDefinition
+ {is.index_name} + {is.schema}.{is.table} + {is.idx_scan.toLocaleString()} + {is.index_size} + {is.definition} +
+
+ + + + {!report.has_pg_stat_statements ? ( +
+
+ + pg_stat_statements extension is not installed +
+

+ Enable it with: CREATE EXTENSION pg_stat_statements; +

+
+ ) : report.slow_queries.length === 0 ? ( +
+ No slow queries found. +
+ ) : ( +
+ + + + + + + + + + + + {report.slow_queries.map((sq, i) => ( + + + + + + + + ))} + +
QueryCallsMean (ms)Total (ms)Rows
+ {sq.query.substring(0, 150)} + {sq.calls.toLocaleString()}{sq.mean_time_ms.toFixed(1)}{sq.total_time_ms.toFixed(0)}{sq.rows.toLocaleString()}
+
+ )} + + + )} +
+
+ ); +} diff --git a/src/components/index-advisor/RecommendationCard.tsx b/src/components/index-advisor/RecommendationCard.tsx new file mode 100644 index 0000000..5cc5223 --- /dev/null +++ b/src/components/index-advisor/RecommendationCard.tsx @@ -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 {priority}; + case "medium": + return {priority}; + default: + return {priority}; + } +} + +function typeBadge(type: string) { + switch (type) { + case "create_index": + return CREATE; + case "drop_index": + return DROP; + case "replace_index": + return REPLACE; + default: + return {type}; + } +} + +export function RecommendationCard({ recommendation, onApply, isApplying, applied }: Props) { + const [showDdl] = useState(true); + + return ( +
+
+
+ {typeBadge(recommendation.recommendation_type)} + {priorityBadge(recommendation.priority)} + + {recommendation.table_schema}.{recommendation.table_name} + + {recommendation.index_name && ( + + {recommendation.index_name} + + )} +
+ +
+ +

{recommendation.rationale}

+ +
+ Impact: {recommendation.estimated_impact} +
+ + {showDdl && ( + 
+          {recommendation.ddl}
+
+ )} +
+ ); +} diff --git a/src/components/layout/TabBar.tsx b/src/components/layout/TabBar.tsx index f357d46..a538203 100644 --- a/src/components/layout/TabBar.tsx +++ b/src/components/layout/TabBar.tsx @@ -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: , lookup: , erd: , + validation: , + "index-advisor": , + snapshots: , }; return ( diff --git a/src/components/schema/SchemaTree.tsx b/src/components/schema/SchemaTree.tsx index 9d047bf..4415831 100644 --- a/src/components/schema/SchemaTree.tsx +++ b/src/components/schema/SchemaTree.tsx @@ -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 + { + if (!isActive) return; + const tab: Tab = { + id: crypto.randomUUID(), + type: "validation", + title: "Data Validation", + connectionId, + database: name, + }; + useAppStore.getState().addTab(tab); + }} + > + Data Validation + + { + if (!isActive) return; + const tab: Tab = { + id: crypto.randomUUID(), + type: "index-advisor", + title: "Index Advisor", + connectionId, + database: name, + }; + useAppStore.getState().addTab(tab); + }} + > + Index Advisor + + { + if (!isActive) return; + const tab: Tab = { + id: crypto.randomUUID(), + type: "snapshots", + title: "Data Snapshots", + connectionId, + database: name, + }; + useAppStore.getState().addTab(tab); + }} + > + Data Snapshots + + (null); + const [dataGenTarget, setDataGenTarget] = useState(null); const tablesQuery = useTables( expanded && category === "tables" ? connectionId : null, @@ -489,6 +540,13 @@ function CategoryNode({ > View Structure + {category === "tables" && ( + setDataGenTarget(item.name)} + > + Generate Test Data + + )} setPrivilegesTarget(item.name)} @@ -524,6 +582,15 @@ function CategoryNode({ table={privilegesTarget} /> )} + {dataGenTarget && ( + !open && setDataGenTarget(null)} + connectionId={connectionId} + schema={schema} + table={dataGenTarget} + /> + )} ); } diff --git a/src/components/snapshots/CreateSnapshotDialog.tsx b/src/components/snapshots/CreateSnapshotDialog.tsx new file mode 100644 index 0000000..17b44bf --- /dev/null +++ b/src/components/snapshots/CreateSnapshotDialog.tsx @@ -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("config"); + const [name, setName] = useState(""); + const [selectedSchema, setSelectedSchema] = useState(""); + const [selectedTables, setSelectedTables] = useState>(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 ( + + + + + + Create Snapshot + + + + {step === "config" && ( + <> +
+
+ + setName(e.target.value)} + placeholder="snapshot-name" + /> +
+ +
+ + +
+ +
+ +
+ {tables && tables.length > 0 && ( + + )} +
+ {tables?.map((t) => ( + + )) ?? ( +

Select a schema first

+ )} +
+

{selectedTables.size} tables selected

+
+
+ +
+ +
+ setIncludeDeps(e.target.checked)} + className="rounded" + /> + + Include referenced tables (foreign keys) + +
+
+
+ + + + + + + )} + + {step === "progress" && ( +
+
+
+ {progress?.message || "Starting..."} + {progress?.percent ?? 0}% +
+
+
+
+
+ {isCreating && ( +
+ + {progress?.stage || "Initializing..."} +
+ )} +
+ )} + + {step === "done" && ( +
+ {error ? ( +
+ +
+

Snapshot Failed

+

{error}

+
+
+ ) : ( +
+ +
+

Snapshot Created

+

+ {result?.total_rows} rows from {result?.tables.length} tables saved. +

+
+
+ )} + + + + {error && } + +
+ )} + +
+ ); +} diff --git a/src/components/snapshots/RestoreSnapshotDialog.tsx b/src/components/snapshots/RestoreSnapshotDialog.tsx new file mode 100644 index 0000000..f713ba5 --- /dev/null +++ b/src/components/snapshots/RestoreSnapshotDialog.tsx @@ -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("select"); + const [filePath, setFilePath] = useState(null); + const [metadata, setMetadata] = useState(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 ( + + + + + + Restore Snapshot + + + + {step === "select" && ( + <> +
+ +

Select a snapshot file to restore

+ +
+ + + + + )} + + {step === "confirm" && metadata && ( + <> +
+
+
+ Name + {metadata.name} +
+
+ Created + {new Date(metadata.created_at).toLocaleString()} +
+
+ Tables + {metadata.tables.length} +
+
+ Total Rows + {metadata.total_rows.toLocaleString()} +
+
+ File Size + {formatBytes(metadata.file_size_bytes)} +
+
+ +
+

Tables included:

+
+ {metadata.tables.map((t) => ( + + {t.schema}.{t.table} ({t.row_count}) + + ))} +
+
+ +
+ +
+ + {truncate && ( +

+ This will DELETE all existing data in the affected tables before restoring. +

+ )} +
+
+
+ + + + + + + )} + + {step === "progress" && ( +
+
+
+ {progress?.message || "Starting..."} + {progress?.percent ?? 0}% +
+
+
+
+
+ {isRestoring && ( +
+ + {progress?.detail || progress?.stage || "Restoring..."} +
+ )} +
+ )} + + {step === "done" && ( +
+ {error ? ( +
+ +
+

Restore Failed

+

{error}

+
+
+ ) : ( +
+ +
+

Restore Completed

+

+ {rowsRestored?.toLocaleString()} rows restored successfully. +

+
+
+ )} + + + + {error && } + +
+ )} + +
+ ); +} diff --git a/src/components/snapshots/SnapshotPanel.tsx b/src/components/snapshots/SnapshotPanel.tsx new file mode 100644 index 0000000..b975183 --- /dev/null +++ b/src/components/snapshots/SnapshotPanel.tsx @@ -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 ( +
+
+
+ + {snapshot.name} +
+ v{snapshot.version} +
+ +
+
+ + {new Date(snapshot.created_at).toLocaleDateString()} +
+
+ + {snapshot.tables.length} tables +
+
+ + {formatBytes(snapshot.file_size_bytes)} +
+
+ +
+ {snapshot.tables.map((t) => ( + + {t.schema}.{t.table} + ({t.row_count}) + + ))} +
+ +
+ {snapshot.total_rows.toLocaleString()} total rows +
+
+ ); +} + +export function SnapshotPanel({ connectionId }: Props) { + const [showCreate, setShowCreate] = useState(false); + const [showRestore, setShowRestore] = useState(false); + const { data: snapshots } = useListSnapshots(); + + return ( +
+ {/* Header */} +
+
+ +

Data Snapshots

+
+
+ + +
+
+ + {/* Content */} +
+ {!snapshots || snapshots.length === 0 ? ( +
+ +

No snapshots yet

+

Create a snapshot to save table data for later restoration.

+
+ ) : ( + snapshots.map((snap) => ( + + )) + )} +
+ + + +
+ ); +} diff --git a/src/components/validation/ValidationPanel.tsx b/src/components/validation/ValidationPanel.tsx new file mode 100644 index 0000000..2124f48 --- /dev/null +++ b/src/components/validation/ValidationPanel.tsx @@ -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([]); + const [ruleInput, setRuleInput] = useState(""); + const [runningIds, setRunningIds] = useState>(new Set()); + + const generateSql = useGenerateValidationSql(); + const runRule = useRunValidationRule(); + const suggestRules = useSuggestValidationRules(); + + const updateRule = useCallback( + (id: string, updates: Partial) => { + 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 ( +
+ {/* Header */} +
+
+
+ +

Data Validation

+
+
+ + +
+
+ +
+ setRuleInput(e.target.value)} + onKeyDown={(e) => e.key === "Enter" && handleAddRule()} + className="flex-1" + /> + +
+ + {rules.length > 0 && ( +
+ {rules.length} rules + {passed > 0 && {passed} passed} + {failed > 0 && {failed} failed} + {errors > 0 && {errors} errors} +
+ )} +
+ + {/* Rules List */} +
+ {rules.length === 0 ? ( +
+ Add a validation rule or click Auto-suggest to get started. +
+ ) : ( + rules.map((rule) => ( + handleRunRule(rule.id)} + onRemove={() => handleRemoveRule(rule.id)} + isRunning={runningIds.has(rule.id)} + /> + )) + )} +
+
+ ); +} diff --git a/src/components/validation/ValidationRuleCard.tsx b/src/components/validation/ValidationRuleCard.tsx new file mode 100644 index 0000000..8796c66 --- /dev/null +++ b/src/components/validation/ValidationRuleCard.tsx @@ -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 Passed; + case "failed": + return Failed; + case "error": + return Error; + case "generating": + case "running": + return Running; + default: + return Pending; + } +} + +export function ValidationRuleCard({ rule, onRun, onRemove, isRunning }: Props) { + const [showSql, setShowSql] = useState(false); + const [showViolations, setShowViolations] = useState(false); + + return ( +
+
+
+

{rule.description}

+
+
+ {statusBadge(rule.status)} + + +
+
+ + {rule.status === "failed" && ( +

+ {rule.violation_count} violation{rule.violation_count !== 1 ? "s" : ""} found +

+ )} + + {rule.error && ( +

{rule.error}

+ )} + + {rule.generated_sql && ( +
+ + {showSql && ( + 
+              {rule.generated_sql}
+
+ )} +
+ )} + + {rule.status === "failed" && rule.sample_violations.length > 0 && ( +
+ + {showViolations && ( +
+ + + + {rule.violation_columns.map((col) => ( + + ))} + + + + {rule.sample_violations.map((row, i) => ( + + {(row as unknown[]).map((val, j) => ( + + ))} + + ))} + +
+ {col} +
+ {val === null ? NULL : String(val)} +
+
+ )} +
+ )} +
+ ); +} diff --git a/src/components/workspace/TabContent.tsx b/src/components/workspace/TabContent.tsx index 7ac0b29..5657b71 100644 --- a/src/components/workspace/TabContent.tsx +++ b/src/components/workspace/TabContent.tsx @@ -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 = ( + + ); + break; + case "index-advisor": + content = ( + + ); + break; + case "snapshots": + content = ( + + ); + break; default: content = null; } diff --git a/src/hooks/use-data-generator.ts b/src/hooks/use-data-generator.ts new file mode 100644 index 0000000..50f1e30 --- /dev/null +++ b/src/hooks/use-data-generator.ts @@ -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(null); + const genIdRef = useRef(""); + + 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, + }; +} diff --git a/src/hooks/use-index-advisor.ts b/src/hooks/use-index-advisor.ts new file mode 100644 index 0000000..ba50706 --- /dev/null +++ b/src/hooks/use-index-advisor.ts @@ -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), + }); +} diff --git a/src/hooks/use-snapshots.ts b/src/hooks/use-snapshots.ts new file mode 100644 index 0000000..7c54376 --- /dev/null +++ b/src/hooks/use-snapshots.ts @@ -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(null); + const snapshotIdRef = useRef(""); + 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(null); + const snapshotIdRef = useRef(""); + + 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, + }; +} diff --git a/src/hooks/use-validation.ts b/src/hooks/use-validation.ts new file mode 100644 index 0000000..1ecc767 --- /dev/null +++ b/src/hooks/use-validation.ts @@ -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), + }); +} diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 526108b..47734ea 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -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("get_mcp_status"); + +// Validation (Wave 1) +export const generateValidationSql = (connectionId: string, ruleDescription: string) => + invoke("generate_validation_sql", { connectionId, ruleDescription }); + +export const runValidationRule = (connectionId: string, sql: string, sampleLimit?: number) => + invoke("run_validation_rule", { connectionId, sql, sampleLimit }); + +export const suggestValidationRules = (connectionId: string) => + invoke("suggest_validation_rules", { connectionId }); + +// Data Generator (Wave 2) +export const generateTestDataPreview = (params: GenerateDataParams, genId: string) => + invoke("generate_test_data_preview", { params, genId }); + +export const insertGeneratedData = (connectionId: string, preview: GeneratedDataPreview) => + invoke("insert_generated_data", { connectionId, preview }); + +export const onDataGenProgress = ( + callback: (p: DataGenProgress) => void +): Promise => + listen("datagen-progress", (e) => callback(e.payload)); + +// Index Advisor (Wave 3A) +export const getIndexAdvisorReport = (connectionId: string) => + invoke("get_index_advisor_report", { connectionId }); + +export const applyIndexRecommendation = (connectionId: string, ddl: string) => + invoke("apply_index_recommendation", { connectionId, ddl }); + +// Snapshots (Wave 3B) +export const createSnapshot = (params: CreateSnapshotParams, snapshotId: string, filePath: string) => + invoke("create_snapshot", { params, snapshotId, filePath }); + +export const restoreSnapshot = (params: RestoreSnapshotParams, snapshotId: string) => + invoke("restore_snapshot", { params, snapshotId }); + +export const listSnapshots = () => + invoke("list_snapshots"); + +export const readSnapshotMetadata = (filePath: string) => + invoke("read_snapshot_metadata", { filePath }); + +export const onSnapshotProgress = ( + callback: (p: SnapshotProgress) => void +): Promise => + listen("snapshot-progress", (e) => callback(e.payload)); diff --git a/src/types/index.ts b/src/types/index.ts index 7855ac2..5d48f6f 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -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; +}