diff --git a/scripts/table-page-cache.test.mjs b/scripts/table-page-cache.test.mjs new file mode 100644 index 0000000..09158e4 --- /dev/null +++ b/scripts/table-page-cache.test.mjs @@ -0,0 +1,74 @@ +import assert from 'node:assert/strict'; +import { test } from 'node:test'; +import { build } from 'esbuild'; + +const { outputFiles } = await build({ + entryPoints: ['src/utils/tablePageCache.ts'], + bundle: true, + platform: 'node', + format: 'esm', + write: false, +}); +const { createTablePageCache, integerKeysetCursor, tablePageKeys } = await import( + `data:text/javascript;base64,${Buffer.from(outputFiles[0].contents).toString('base64')}` +); + +test('keyset cursor is used only for a safe integer primary key in default order', () => { + const page = { columns: ['id', 'name'], rows: [[2, 'two'], [3, 'three']], time: 1 }; + assert.equal(integerKeysetCursor(page, ['id'], { id: 'INTEGER' }, null), 3); + assert.equal(integerKeysetCursor(page, ['id'], { id: 'INTEGER' }, 'name'), undefined); + assert.equal(integerKeysetCursor(page, ['id', 'name'], { id: 'INTEGER' }, null), undefined); + assert.equal(integerKeysetCursor(page, ['id'], { id: 'TEXT' }, null), undefined); + assert.equal(integerKeysetCursor({ ...page, rows: [[Number.MAX_SAFE_INTEGER + 1, 'x']] }, + ['id'], { id: 'BIGINT' }, null), undefined); +}); + +const base = { + connectionId: 'a', database: 'app', tableName: 'users', filters: [], + sortColumn: null, sortDirection: 'ASC', limit: 100, offset: 0, generation: 0, +}; + +test('page identity includes pagination, filters, sort, target, and refresh generation', () => { + const key = tablePageKeys(base); + assert.notEqual(key.page, tablePageKeys({ ...base, offset: 100 }).page); + assert.equal(key.count, tablePageKeys({ ...base, offset: 100 }).count); + assert.equal(key.count, tablePageKeys({ ...base, sortColumn: 'id' }).count); + for (const change of [ + { filters: [{ column: 'name', value: 'x', enabled: true }] }, + { database: 'other' }, { connectionId: 'b' }, { tableName: 'orders' }, + { generation: 1 }, + ]) { + assert.notEqual(key.count, tablePageKeys({ ...base, ...change }).count); + } +}); + +test('cached pages remain visible when stale and refresh only when due', () => { + const cache = createTablePageCache(2, 4, 1000); + const keys = tablePageKeys(base); + cache.putPage('tab', keys.page, { columns: ['id'], rows: [[1]], time: 5 }, 100); + assert.equal(cache.getPage('tab', keys.page, 500)?.fresh, true); + assert.equal(cache.getPage('tab', keys.page, 1200)?.fresh, false); + assert.deepEqual(cache.getPage('tab', keys.page, 1200)?.value.rows, [[1]]); +}); + +test('count is shared across pages but not across filters or refreshes', () => { + const cache = createTablePageCache(2, 4, 1000); + const key = tablePageKeys(base).count; + cache.putCount('tab', key, 350, 100); + assert.equal(cache.getCount('tab', tablePageKeys({ ...base, offset: 200 }).count, 500)?.value, 350); + assert.equal(cache.getCount('tab', tablePageKeys({ ...base, generation: 1 }).count, 500), undefined); + assert.equal(cache.getCount('tab', key, 1200)?.fresh, false); +}); + +test('evicts oldest pages per tab and clears sensitive rows when tab closes', () => { + const cache = createTablePageCache(2, 4, 1000); + for (let offset = 0; offset < 300; offset += 100) { + const key = tablePageKeys({ ...base, offset }).page; + cache.putPage('tab', key, { columns: ['secret'], rows: [[offset]], time: 1 }, 100); + } + assert.equal(cache.getPage('tab', tablePageKeys(base).page, 200), undefined); + cache.putCount('tab', tablePageKeys(base).count, 3, 100); + cache.dropTab('tab'); + assert.equal(cache.getPage('tab', tablePageKeys({ ...base, offset: 200 }).page, 200), undefined); + assert.equal(cache.getCount('tab', tablePageKeys(base).count, 200), undefined); +}); diff --git a/scripts/table-structure-cache.test.mjs b/scripts/table-structure-cache.test.mjs new file mode 100644 index 0000000..2929a87 --- /dev/null +++ b/scripts/table-structure-cache.test.mjs @@ -0,0 +1,40 @@ +import assert from 'node:assert/strict'; +import { test } from 'node:test'; +import { build } from 'esbuild'; + +const { outputFiles } = await build({ + entryPoints: ['src/utils/tableStructureCache.ts'], + bundle: true, platform: 'node', format: 'esm', write: false, +}); +const { createTableStructureCache } = await import( + `data:text/javascript;base64,${Buffer.from(outputFiles[0].contents).toString('base64')}` +); + +test('table structure loads once across tab switches and refreshes after expiry', async () => { + const cache = createTableStructureCache(2, 10); + let calls = 0; + const load = async () => ++calls; + assert.equal(await cache.get('conn', 'db', 'users', load), 1); + assert.equal(await cache.get('conn', 'db', 'users', load), 1); + assert.equal(calls, 1); + await new Promise(resolve => setTimeout(resolve, 12)); + assert.equal(await cache.get('conn', 'db', 'users', load), 2); +}); + +test('concurrent structure requests share one flight and invalidation rejects stale cache', async () => { + const cache = createTableStructureCache(); + let resolve; + let calls = 0; + const load = () => { + calls += 1; + return new Promise(done => { resolve = done; }); + }; + const first = cache.get('conn', 'db', 'users', load); + const second = cache.get('conn', 'db', 'users', load); + assert.equal(calls, 1); + cache.invalidate('conn'); + resolve('old'); + assert.equal(await first, 'old'); + assert.equal(await second, 'old'); + assert.equal(await cache.get('conn', 'db', 'users', async () => 'new'), 'new'); +}); diff --git a/src-tauri/src/core/mod.rs b/src-tauri/src/core/mod.rs index 17d19bd..5cbf14f 100644 --- a/src-tauri/src/core/mod.rs +++ b/src-tauri/src/core/mod.rs @@ -85,6 +85,8 @@ pub struct QueryResult { pub total_count: Option, pub page: Option, pub page_size: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub has_more: Option, } #[derive(Debug, Serialize)] @@ -201,6 +203,7 @@ pub struct StreamingComplete { pub execution_time_ms: u64, pub total_rows: u64, pub affected_rows: u64, + pub truncated: bool, } #[cfg(test)] diff --git a/src-tauri/src/core/query_engine.rs b/src-tauri/src/core/query_engine.rs index 3c0a065..3fe36be 100644 --- a/src-tauri/src/core/query_engine.rs +++ b/src-tauri/src/core/query_engine.rs @@ -6,18 +6,65 @@ use crate::core::{ TableConstraintStructure, TableIndexStructure, TableMetadata, TableStructure, }; use anyhow::{anyhow, Result}; -use futures::StreamExt; +use futures::{StreamExt, TryStreamExt}; use serde_json::Value; use sqlx::{Acquire, Column, ConnectOptions, Executor, Row, Statement, TypeInfo, ValueRef}; use std::collections::{BTreeMap, BTreeSet}; -use std::fs::File; -use std::io::Write; +use std::fs::{File, OpenOptions}; +use std::io::{BufWriter, Write}; +use std::path::PathBuf; use std::time::Instant; use tauri::Emitter; use tokio::time::{sleep, Duration}; use tokio_util::sync::CancellationToken; use uuid::Uuid; +const MAX_QUERY_RESULT_ROWS: u64 = 10_000; + +struct PendingExport { + temporary: PathBuf, + destination: PathBuf, +} + +impl PendingExport { + fn create(file_path: &str) -> Result<(Self, File)> { + let destination = PathBuf::from(file_path); + let filename = destination + .file_name() + .ok_or_else(|| anyhow!("Invalid export path"))?; + let temporary = destination.with_file_name(format!( + ".{}.{}.tmp", + filename.to_string_lossy(), + Uuid::new_v4() + )); + let file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary)?; + Ok(( + Self { + temporary, + destination, + }, + file, + )) + } + + fn commit(self) -> Result<()> { + #[cfg(windows)] + std::fs::copy(&self.temporary, &self.destination)?; + #[cfg(not(windows))] + std::fs::rename(&self.temporary, &self.destination)?; + Ok(()) + } +} + +impl Drop for PendingExport { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.temporary); + } +} + fn mysql_type_is_boolean(type_name: &str, type_display: &str) -> bool { type_name == "bool" || type_name == "boolean" @@ -28,6 +75,134 @@ fn json_boolean_number(value: bool) -> Value { Value::Number(serde_json::Number::from(u8::from(value))) } +enum TableExportWriter { + Csv(csv::Writer), + Json { + writer: W, + columns: Vec, + first: bool, + }, + Sql { + writer: W, + insert_prefix: String, + }, +} + +impl TableExportWriter { + fn new( + format: &str, + mut writer: W, + columns: Vec, + table: &str, + db_type: &str, + ) -> Result { + match format { + "csv" => { + let mut csv = csv::Writer::from_writer(writer); + csv.write_record(&columns)?; + Ok(Self::Csv(csv)) + } + "json" => { + writer.write_all(b"[")?; + Ok(Self::Json { + writer, + columns, + first: true, + }) + } + "sql" => { + let quote = if db_type == "mysql" { '`' } else { '"' }; + let quoted = |name: &str| { + format!( + "{quote}{}{quote}", + name.replace(quote, "e.to_string().repeat(2)) + ) + }; + let insert_prefix = format!( + "INSERT INTO {} ({}) VALUES (", + quoted(table), + columns + .iter() + .map(|column| quoted(column)) + .collect::>() + .join(", "), + ); + Ok(Self::Sql { + writer, + insert_prefix, + }) + } + _ => Err(anyhow!("Unsupported export format")), + } + } + + fn write_row(&mut self, row: Vec) -> Result<()> { + match self { + Self::Csv(writer) => { + let record = row + .into_iter() + .map(|value| match value { + Value::Null => String::new(), + Value::String(value) => value, + Value::Number(value) => value.to_string(), + Value::Bool(value) => value.to_string(), + value => value.to_string(), + }) + .collect::>(); + writer.write_record(&record)?; + } + Self::Json { + writer, + columns, + first, + } => { + if !*first { + writer.write_all(b",")?; + } + let object = columns + .iter() + .cloned() + .zip(row) + .collect::>(); + serde_json::to_writer(writer, &object)?; + *first = false; + } + Self::Sql { + writer, + insert_prefix, + } => { + writer.write_all(insert_prefix.as_bytes())?; + for (index, value) in row.into_iter().enumerate() { + if index > 0 { + writer.write_all(b", ")?; + } + match value { + Value::Null => writer.write_all(b"NULL")?, + Value::String(value) => write!(writer, "'{}'", value.replace("'", "''"))?, + Value::Number(value) => write!(writer, "{value}")?, + Value::Bool(value) => write!(writer, "{value}")?, + value => write!(writer, "'{}'", value.to_string().replace("'", "''"))?, + } + } + writer.write_all(b");\n")?; + } + } + Ok(()) + } + + fn finish(mut self) -> Result<()> { + match &mut self { + Self::Csv(writer) => writer.flush()?, + Self::Json { writer, .. } => { + writer.write_all(b"]")?; + writer.flush()?; + } + Self::Sql { writer, .. } => writer.flush()?, + } + Ok(()) + } +} + fn group_foreign_keys( columns: impl IntoIterator, ) -> Vec { @@ -67,15 +242,277 @@ fn build_explain_sql(sql: &str, database_type: &DatabaseType) -> Result #[cfg(test)] mod tests { use super::{ - build_explain_sql, build_where_clause, group_foreign_keys, json_boolean_number, - mysql_type_is_boolean, QueryEngine, + build_browse_order_clause, build_explain_sql, build_seek_predicate, build_where_clause, + group_foreign_keys, json_boolean_number, mysql_type_is_boolean, with_seek_predicate, + PendingExport, QueryEngine, TableExportWriter, }; use crate::core::{ connection_manager::ConnectionManager, ConnectionConfig, DatabaseType, FilterConfig, }; - use serde_json::json; + use serde_json::{json, Value}; use uuid::Uuid; + #[test] + fn table_browse_uses_primary_key_as_stable_tie_breaker() { + let keys = vec!["id".to_string()]; + assert_eq!( + build_browse_order_clause(None, None, &keys, "mysql"), + "ORDER BY `id` ASC" + ); + assert_eq!( + build_browse_order_clause(Some("name"), Some("DESC"), &keys, "postgres"), + "ORDER BY \"name\" DESC, \"id\" ASC" + ); + assert_eq!( + build_browse_order_clause(Some("id"), Some("DESC"), &keys, "sqlite"), + "ORDER BY \"id\" DESC" + ); + assert_eq!( + with_seek_predicate("WHERE \"status\" = 1".into(), Some("\"id\" > 9".into())), + "WHERE \"status\" = 1 AND \"id\" > 9" + ); + assert_eq!( + build_seek_predicate(&json!(9), "id", "sqlite").unwrap(), + "\"id\" > 9" + ); + assert!(build_seek_predicate(&json!("9 OR 1=1"), "id", "mysql").is_err()); + } + + #[test] + fn streamed_table_export_preserves_csv_json_and_sql_values() { + let columns = vec!["id".to_string(), "name".to_string(), "note".to_string()]; + let row = vec![json!(1), json!("O'Brien"), Value::Null]; + for format in ["csv", "json", "sql"] { + let mut output = Vec::new(); + let mut writer = + TableExportWriter::new(format, &mut output, columns.clone(), "people", "mysql") + .unwrap(); + writer.write_row(row.clone()).unwrap(); + writer.finish().unwrap(); + let content = String::from_utf8(output).unwrap(); + match format { + "csv" => assert_eq!(content, "id,name,note\n1,O'Brien,\n"), + "json" => assert_eq!( + serde_json::from_str::(&content).unwrap(), + json!([{"id": 1, "name": "O'Brien", "note": null}]) + ), + "sql" => assert_eq!( + content, + "INSERT INTO `people` (`id`, `name`, `note`) VALUES (1, 'O''Brien', NULL);\n" + ), + _ => unreachable!(), + } + } + } + + #[test] + fn empty_json_export_is_valid() { + let mut output = Vec::new(); + TableExportWriter::new("json", &mut output, vec!["id".into()], "people", "sqlite") + .unwrap() + .finish() + .unwrap(); + assert_eq!(String::from_utf8(output).unwrap(), "[]"); + } + + #[test] + fn failed_stream_keeps_the_previous_export_and_removes_partial_file() { + use std::io::Write; + + let path = + std::env::temp_dir().join(format!("sqlmate-export-abort-{}.json", Uuid::new_v4())); + std::fs::write(&path, b"previous export").unwrap(); + let (pending, mut file) = PendingExport::create(path.to_str().unwrap()).unwrap(); + let temporary = pending.temporary.clone(); + file.write_all(b"incomplete").unwrap(); + drop(file); + drop(pending); + assert_eq!(std::fs::read(&path).unwrap(), b"previous export"); + assert!(!temporary.exists()); + std::fs::remove_file(path).unwrap(); + } + + #[tokio::test] + async fn sqlite_progress_handler_interrupts_a_cancelled_statement() { + use sqlx::ConnectOptions; + use tokio_util::sync::CancellationToken; + + let mut conn = sqlx::sqlite::SqliteConnectOptions::new() + .filename(":memory:") + .connect() + .await + .unwrap(); + let token = CancellationToken::new(); + { + let progress_token = token.clone(); + let mut handle = conn.lock_handle().await.unwrap(); + handle.set_progress_handler(1000, move || !progress_token.is_cancelled()); + } + token.cancel(); + let result = sqlx::query_scalar::<_, i64>( + "WITH RECURSIVE sequence(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM sequence WHERE n < 1000000) SELECT sum(n) FROM sequence", + ) + .fetch_one(&mut conn) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn sqlite_browse_keyset_paging_survives_an_insert_before_the_cursor() { + use sqlx::{ConnectOptions, Executor}; + + let id = Uuid::new_v4(); + let path = std::env::temp_dir().join(format!("sqlmate-browse-{id}.db")); + let mut setup = sqlx::sqlite::SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .connect() + .await + .unwrap(); + setup + .execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)") + .await + .unwrap(); + setup + .execute("INSERT INTO items VALUES (2, 'second'), (3, 'third'), (4, 'fourth')") + .await + .unwrap(); + drop(setup); + + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": id, "name": "test", "db_type": "Sqlite", + "database": path.to_string_lossy(), "ssl_enabled": false, "ssh_enabled": false, + "safety_mode": "Silent" + })) + .unwrap(); + let manager = ConnectionManager::new(); + manager.connect(config, None).await.unwrap(); + let keys = vec!["id".to_string()]; + let first = QueryEngine::get_table_data( + &manager, + &id, + "items", + 2, + 0, + vec![], + None, + None, + keys.clone(), + None, + ) + .await + .unwrap(); + assert_eq!(first.has_more, Some(true)); + assert_eq!(first.rows[0][0], json!(2)); + assert_eq!(first.rows[1][0], json!(3)); + + QueryEngine::execute_query( + &manager, + &id, + "INSERT INTO items VALUES (1, 'first')", + None, + None, + ) + .await + .unwrap(); + let second = QueryEngine::get_table_data( + &manager, + &id, + "items", + 2, + 2, + vec![], + None, + None, + keys, + Some(json!(3)), + ) + .await + .unwrap(); + assert_eq!(second.rows.len(), 1); + assert_eq!(second.rows[0][0], json!(4)); + assert_eq!(second.has_more, Some(false)); + + manager.disconnect(&id).await.unwrap(); + std::fs::remove_file(path).unwrap(); + } + + #[tokio::test] + async fn sqlite_table_export_streams_rows_and_keeps_empty_csv_headers() { + use sqlx::{ConnectOptions, Executor}; + + let id = Uuid::new_v4(); + let db_path = std::env::temp_dir().join(format!("sqlmate-export-{id}.db")); + let export_path = std::env::temp_dir().join(format!("sqlmate-export-{id}.csv")); + let mut setup = sqlx::sqlite::SqliteConnectOptions::new() + .filename(&db_path) + .create_if_missing(true) + .connect() + .await + .unwrap(); + setup + .execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)") + .await + .unwrap(); + setup + .execute("INSERT INTO items (name) VALUES ('first'), ('second')") + .await + .unwrap(); + drop(setup); + + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": id, "name": "test", "db_type": "Sqlite", + "database": db_path.to_string_lossy(), "ssl_enabled": false, "ssh_enabled": false, + "safety_mode": "Silent" + })) + .unwrap(); + let manager = ConnectionManager::new(); + manager.connect(config, None).await.unwrap(); + + let count = QueryEngine::export_table_data( + &manager, + &id, + "items", + vec![], + None, + None, + "csv", + export_path.to_str().unwrap(), + ) + .await + .unwrap(); + assert_eq!(count, 2); + let csv = std::fs::read_to_string(&export_path).unwrap(); + assert!(csv.starts_with("id,name\n")); + assert!(csv.contains("1,first\n")); + + let empty_filter = FilterConfig { + id: "none".into(), + column: "id".into(), + operator: "=".into(), + value: "999".into(), + enabled: true, + }; + let count = QueryEngine::export_table_data( + &manager, + &id, + "items", + vec![empty_filter], + None, + None, + "csv", + export_path.to_str().unwrap(), + ) + .await + .unwrap(); + assert_eq!(count, 0); + assert_eq!(std::fs::read_to_string(&export_path).unwrap(), "id,name\n"); + + manager.disconnect(&id).await.unwrap(); + std::fs::remove_file(export_path).unwrap(); + std::fs::remove_file(db_path).unwrap(); + } + #[test] fn recognizes_mariadb_boolean_types() { assert!(mysql_type_is_boolean("boolean", "BOOLEAN")); @@ -766,6 +1203,84 @@ fn build_order_clause( } } +fn quote_browse_column(name: &str, db_type: &str) -> String { + if db_type == "mysql" { + format!("`{}`", name.replace('`', "``")) + } else { + format!("\"{}\"", name.replace('"', "\"\"")) + } +} + +fn build_browse_order_clause( + sort_column: Option<&str>, + sort_direction: Option<&str>, + primary_keys: &[String], + db_type: &str, +) -> String { + let mut terms = Vec::new(); + if let Some(column) = sort_column.filter(|column| !column.is_empty()) { + let direction = if sort_direction == Some("DESC") { + "DESC" + } else { + "ASC" + }; + terms.push(format!( + "{} {direction}", + quote_browse_column(column, db_type) + )); + } + for column in primary_keys { + if sort_column != Some(column.as_str()) { + terms.push(format!("{} ASC", quote_browse_column(column, db_type))); + } + } + if terms.is_empty() { + String::new() + } else { + format!("ORDER BY {}", terms.join(", ")) + } +} + +fn build_seek_predicate(value: &Value, primary_key: &str, db_type: &str) -> Result { + // JSON integers are the only cursors accepted here. Strings need bound parameters + // and collation-specific ordering before they can be used safely. + let literal = match value { + Value::Number(value) if value.is_i64() || value.is_u64() => value.to_string(), + _ => { + return Err(anyhow!( + "This primary-key value cannot be used as a page cursor" + )) + } + }; + Ok(format!( + "{} > {literal}", + quote_browse_column(primary_key, db_type) + )) +} + +fn with_seek_predicate(where_clause: String, predicate: Option) -> String { + match predicate { + Some(predicate) if where_clause.is_empty() => format!("WHERE {predicate}"), + Some(predicate) => format!("{where_clause} AND {predicate}"), + None => where_clause, + } +} + +async fn cancel_postgres_statement(pool: &sqlx::PgPool, backend_pid: i32) { + let _ = tokio::time::timeout( + Duration::from_secs(2), + sqlx::query_scalar::<_, bool>("SELECT pg_cancel_backend($1)") + .bind(backend_pid) + .fetch_one(pool), + ) + .await; +} + +async fn cancel_mysql_statement(pool: &sqlx::MySqlPool, connection_id: u64) { + let sql = format!("KILL QUERY {connection_id}"); + let _ = tokio::time::timeout(Duration::from_secs(2), sqlx::query(&sql).execute(pool)).await; +} + pub struct QueryEngine; impl QueryEngine { @@ -884,20 +1399,24 @@ impl QueryEngine { } macro_rules! stream_db { - ($pool:expr, $db_macro:ident, $stream:expr) => {{ + ($pool:expr, $db_macro:ident, $stream:expr, $cancel:expr) => {{ use sqlx::Either; let mut stream = $stream; let mut columns_sent = false; let mut batch = Vec::new(); let mut total_rows = 0u64; + let mut displayed_rows = 0u64; let mut affected_rows = 0u64; + let mut truncated = false; let batch_size = 1000; - while let Some(res_result) = StreamExt::next(&mut stream).await { - if token.is_cancelled() { + while let Some(res_result) = tokio::select! { + _ = token.cancelled() => { + $cancel.await; return Ok(()); - } - + }, + result = StreamExt::next(&mut stream) => result, + } { match res_result? { Either::Left(result) => { affected_rows += result.rows_affected(); @@ -916,18 +1435,26 @@ impl QueryEngine { columns_sent = true; } - batch.push($db_macro!(&row)); total_rows += 1; + if displayed_rows >= MAX_QUERY_RESULT_ROWS { + truncated = true; + // Only database-enforced read-only queries may stop early. + if read_only { + break; + } + continue; + } + batch.push($db_macro!(&row)); + displayed_rows += 1; if batch.len() >= batch_size { window.emit( "query-batch", StreamingBatch { query_id, - rows: batch.clone(), + rows: std::mem::take(&mut batch), }, )?; - batch.clear(); sleep(Duration::from_millis(5)).await; } } @@ -967,6 +1494,7 @@ impl QueryEngine { execution_time_ms: start.elapsed().as_millis() as u64, total_rows, affected_rows, + truncated, }, )?; @@ -980,39 +1508,54 @@ impl QueryEngine { if let Some(pool) = pools.get(connection_id) { if read_only { let mut tx = pool.begin().await?; + let backend_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *tx) + .await?; sqlx::query("SET TRANSACTION READ ONLY") .execute(&mut *tx) .await?; if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); + let cancel_pool = (*pool).clone(); stream_db!( &mut *tx, postgres_row_to_values, query.fetch(&mut *tx).map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async { cancel_postgres_statement(&cancel_pool, backend_pid).await } ); } + let cancel_pool = (*pool).clone(); stream_db!( &mut *tx, postgres_row_to_values, - sqlx::raw_sql(sql).fetch_many(&mut *tx) + sqlx::raw_sql(sql).fetch_many(&mut *tx), + async { cancel_postgres_statement(&cancel_pool, backend_pid).await } ); } + let mut conn = pool.acquire().await?; + let backend_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *conn) + .await?; if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); + let cancel_pool = (*pool).clone(); stream_db!( - pool, + &mut *conn, postgres_row_to_values, query - .fetch(pool) + .fetch(&mut *conn) .map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async { cancel_postgres_statement(&cancel_pool, backend_pid).await } ); } + let cancel_pool = (*pool).clone(); stream_db!( - pool, + &mut *conn, postgres_row_to_values, - sqlx::raw_sql(sql).fetch_many(pool) + sqlx::raw_sql(sql).fetch_many(&mut *conn), + async { cancel_postgres_statement(&cancel_pool, backend_pid).await } ); } } @@ -1023,40 +1566,55 @@ impl QueryEngine { if let Some(pool) = pools.get(connection_id) { if read_only { let mut conn = pool.acquire().await?; + let backend_id: u64 = sqlx::query_scalar("SELECT CONNECTION_ID()") + .fetch_one(&mut *conn) + .await?; sqlx::query("SET TRANSACTION READ ONLY") .execute(&mut *conn) .await?; let mut tx = conn.begin().await?; if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); + let cancel_pool = (*pool).clone(); stream_db!( &mut *tx, mysql_row_to_values, query.fetch(&mut *tx).map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async { cancel_mysql_statement(&cancel_pool, backend_id).await } ); } + let cancel_pool = (*pool).clone(); stream_db!( &mut *tx, mysql_row_to_values, - sqlx::raw_sql(sql).fetch_many(&mut *tx) + sqlx::raw_sql(sql).fetch_many(&mut *tx), + async { cancel_mysql_statement(&cancel_pool, backend_id).await } ); } + let mut conn = pool.acquire().await?; + let backend_id: u64 = sqlx::query_scalar("SELECT CONNECTION_ID()") + .fetch_one(&mut *conn) + .await?; if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); + let cancel_pool = (*pool).clone(); stream_db!( - pool, + &mut *conn, mysql_row_to_values, query - .fetch(pool) + .fetch(&mut *conn) .map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async { cancel_mysql_statement(&cancel_pool, backend_id).await } ); } + let cancel_pool = (*pool).clone(); stream_db!( - pool, + &mut *conn, mysql_row_to_values, - sqlx::raw_sql(sql).fetch_many(pool) + sqlx::raw_sql(sql).fetch_many(&mut *conn), + async { cancel_mysql_statement(&cancel_pool, backend_id).await } ); } } @@ -1065,49 +1623,58 @@ impl QueryEngine { { let pools = manager.get_sqlite_pools().await; if let Some(pool) = pools.get(connection_id) { - if read_only { - let config = manager.get_config(connection_id).await?; - let path = config - .database - .ok_or_else(|| anyhow!("SQLite path is missing"))?; - if path == ":memory:" { - return Err(anyhow!( - "Read-only SQL, including parameterized SELECT, is unavailable for in-memory SQLite databases" - )); - } + let config = manager.get_config(connection_id).await?; + let path = config + .database + .ok_or_else(|| anyhow!("SQLite path is missing"))?; + if path != ":memory:" { let mut conn = sqlx::sqlite::SqliteConnectOptions::new() .filename(path) - .read_only(true) + .read_only(read_only) .connect() .await?; + { + let progress_token = token.clone(); + let mut handle = conn.lock_handle().await?; + handle.set_progress_handler(1000, move || !progress_token.is_cancelled()); + } if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); stream_db!( &mut conn, sqlite_row_to_values, query.fetch(&mut conn).map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async {} ); } stream_db!( &mut conn, sqlite_row_to_values, - sqlx::raw_sql(sql).fetch_many(&mut conn) + sqlx::raw_sql(sql).fetch_many(&mut conn), + async {} ); } + if read_only { + return Err(anyhow!( + "Read-only SQL, including parameterized SELECT, is unavailable for in-memory SQLite databases" + )); + } if let Some(values) = parameter_values { let query = bind_values!(sqlx::query::(sql), values); stream_db!( pool, sqlite_row_to_values, query.fetch(pool).map(|row| row - .map(sqlx::Either::::Right)) + .map(sqlx::Either::::Right)), + async {} ); } stream_db!( pool, sqlite_row_to_values, - sqlx::raw_sql(sql).fetch_many(pool) + sqlx::raw_sql(sql).fetch_many(pool), + async {} ); } } @@ -1159,6 +1726,7 @@ impl QueryEngine { total_count: None, page, page_size, + has_more: None, }); }}; } @@ -1285,6 +1853,7 @@ impl QueryEngine { total_count, page, page_size, + has_more: None, }); } } @@ -1348,6 +1917,7 @@ impl QueryEngine { total_count, page, page_size, + has_more: None, }); } } @@ -1411,6 +1981,7 @@ impl QueryEngine { total_count, page, page_size, + has_more: None, }); } } @@ -1560,7 +2131,15 @@ impl QueryEngine { filters: Vec, sort_column: Option, sort_direction: Option, + primary_keys: Vec, + seek_after: Option, ) -> Result { + let limit = limit.clamp(1, 1000); + if seek_after.is_some() && (sort_column.is_some() || primary_keys.len() != 1) { + return Err(anyhow!( + "Keyset paging requires one primary key and default sort" + )); + } let db_type = { if manager .get_postgres_pools() @@ -1591,44 +2170,86 @@ impl QueryEngine { .bind(table_name) .fetch_one(&pool) .await?; - let where_clause = build_where_clause(filters, "postgres"); - let order_clause = build_order_clause(sort_column, sort_direction, "postgres"); + let predicate = seek_after + .as_ref() + .map(|value| build_seek_predicate(value, &primary_keys[0], "postgres")) + .transpose()?; + let where_clause = + with_seek_predicate(build_where_clause(filters, "postgres"), predicate); + let order_clause = build_browse_order_clause( + sort_column.as_deref(), + sort_direction.as_deref(), + &primary_keys, + "postgres", + ); let sql = format!( "SELECT * FROM {}\"{}\" {} {} LIMIT {} OFFSET {};", if public_exists { "\"public\"." } else { "" }, table_name.replace("\"", "\"\""), where_clause, order_clause, - limit, - offset + limit + 1, + if seek_after.is_some() { 0 } else { offset } ); - Self::execute_query(manager, connection_id, &sql, None, None).await + let mut result = + Self::execute_query(manager, connection_id, &sql, None, None).await?; + result.has_more = Some(result.rows.len() > limit as usize); + result.rows.truncate(limit as usize); + Ok(result) } Some("mysql") => { - let where_clause = build_where_clause(filters, "mysql"); - let order_clause = build_order_clause(sort_column, sort_direction, "mysql"); + let predicate = seek_after + .as_ref() + .map(|value| build_seek_predicate(value, &primary_keys[0], "mysql")) + .transpose()?; + let where_clause = + with_seek_predicate(build_where_clause(filters, "mysql"), predicate); + let order_clause = build_browse_order_clause( + sort_column.as_deref(), + sort_direction.as_deref(), + &primary_keys, + "mysql", + ); let sql = format!( "SELECT * FROM `{}` {} {} LIMIT {} OFFSET {};", table_name.replace("`", "``"), where_clause, order_clause, - limit, - offset + limit + 1, + if seek_after.is_some() { 0 } else { offset } ); - Self::execute_query(manager, connection_id, &sql, None, None).await + let mut result = + Self::execute_query(manager, connection_id, &sql, None, None).await?; + result.has_more = Some(result.rows.len() > limit as usize); + result.rows.truncate(limit as usize); + Ok(result) } Some("sqlite") => { - let where_clause = build_where_clause(filters, "sqlite"); - let order_clause = build_order_clause(sort_column, sort_direction, "sqlite"); + let predicate = seek_after + .as_ref() + .map(|value| build_seek_predicate(value, &primary_keys[0], "sqlite")) + .transpose()?; + let where_clause = + with_seek_predicate(build_where_clause(filters, "sqlite"), predicate); + let order_clause = build_browse_order_clause( + sort_column.as_deref(), + sort_direction.as_deref(), + &primary_keys, + "sqlite", + ); let sql = format!( "SELECT * FROM \"{}\" {} {} LIMIT {} OFFSET {};", table_name.replace("\"", "\"\""), where_clause, order_clause, - limit, - offset + limit + 1, + if seek_after.is_some() { 0 } else { offset } ); - Self::execute_query(manager, connection_id, &sql, None, None).await + let mut result = + Self::execute_query(manager, connection_id, &sql, None, None).await?; + result.has_more = Some(result.rows.len() > limit as usize); + result.rows.truncate(limit as usize); + Ok(result) } Some(_) => Err(anyhow!("Unknown database type")), None => Err(anyhow!("Connection not found")), @@ -2365,88 +2986,61 @@ impl QueryEngine { _ => return Err(anyhow!("Unknown database type")), }; - let result = Self::execute_query(manager, connection_id, &sql, None, None).await?; - let rows_count = result.rows.len() as u64; - - let mut file = File::create(file_path)?; - - match format { - "csv" => { - let mut wtr = csv::Writer::from_writer(file); - // Write headers - wtr.write_record(&result.columns)?; - // Write rows - for row in result.rows { - let record: Vec = row - .into_iter() - .map(|v| match v { - Value::Null => "".to_string(), - Value::String(s) => s, - Value::Number(n) => n.to_string(), - Value::Bool(b) => b.to_string(), - _ => v.to_string(), - }) - .collect(); - wtr.write_record(&record)?; + macro_rules! export_rows { + ($pool:expr, $row_macro:ident) => {{ + let statement = Executor::prepare($pool, sql.as_str()).await?; + let columns = statement + .columns() + .iter() + .map(|column| column.name().to_string()) + .collect(); + let (pending, file) = PendingExport::create(file_path)?; + let file = BufWriter::new(file); + let mut writer = + TableExportWriter::new(format, file, columns, table_name, db_type)?; + let mut rows_count = 0; + let mut stream = sqlx::query(&sql).fetch($pool); + while let Some(row) = stream.try_next().await? { + writer.write_row($row_macro!(&row))?; + rows_count += 1; } - wtr.flush()?; + drop(stream); + writer.finish()?; + pending.commit()?; + return Ok(rows_count); + }}; + } + + match db_type { + "postgres" => { + let pool = manager + .get_postgres_pools() + .await + .get(connection_id) + .cloned() + .ok_or_else(|| anyhow!("Connection not found"))?; + export_rows!(&pool, postgres_row_to_values); } - "json" => { - let mut json_rows = Vec::new(); - for row in result.rows { - let mut obj = serde_json::Map::new(); - for (i, col) in result.columns.iter().enumerate() { - obj.insert(col.clone(), row[i].clone()); - } - json_rows.push(Value::Object(obj)); - } - let json_data = serde_json::to_string_pretty(&json_rows)?; - file.write_all(json_data.as_bytes())?; + "mysql" => { + let pool = manager + .get_mysql_pools() + .await + .get(connection_id) + .cloned() + .ok_or_else(|| anyhow!("Connection not found"))?; + export_rows!(&pool, mysql_row_to_values); } - "sql" => { - for row in result.rows { - let values: Vec = row - .into_iter() - .map(|v| match v { - Value::Null => "NULL".to_string(), - Value::String(s) => format!("'{}'", s.replace("'", "''")), - Value::Number(n) => n.to_string(), - Value::Bool(b) => if b { "true" } else { "false" }.to_string(), - _ => format!("'{}'", v.to_string().replace("'", "''")), - }) - .collect(); - - let insert_sql = match db_type { - "mysql" => format!( - "INSERT INTO `{}` ({}) VALUES ({});\n", - table_name.replace("`", "``"), - result - .columns - .iter() - .map(|c| format!("`{}`", c.replace("`", "``"))) - .collect::>() - .join(", "), - values.join(", ") - ), - _ => format!( - "INSERT INTO \"{}\" ({}) VALUES ({});\n", - table_name.replace("\"", "\"\""), - result - .columns - .iter() - .map(|c| format!("\"{}\"", c.replace("\"", "\"\""))) - .collect::>() - .join(", "), - values.join(", ") - ), - }; - file.write_all(insert_sql.as_bytes())?; - } + "sqlite" => { + let pool = manager + .get_sqlite_pools() + .await + .get(connection_id) + .cloned() + .ok_or_else(|| anyhow!("Connection not found"))?; + export_rows!(&pool, sqlite_row_to_values); } - _ => return Err(anyhow!("Unsupported export format")), + _ => Err(anyhow!("Unknown database type")), } - - Ok(rows_count) } pub async fn get_sidebar_items( diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 604f853..ce863c7 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -810,6 +810,8 @@ async fn get_table_data( filters: Option>, sort_column: Option, sort_direction: Option, + primary_keys: Option>, + seek_after: Option, ) -> Result { let filters = filters.unwrap_or_default(); QueryEngine::get_table_data( @@ -821,6 +823,8 @@ async fn get_table_data( filters, sort_column, sort_direction, + primary_keys.unwrap_or_default(), + seek_after, ) .await .map_err(|e| e.to_string()) diff --git a/src/App.tsx b/src/App.tsx index c074e63..ecdf357 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -68,7 +68,7 @@ function App() { // ⌘R or Ctrl+R for Refresh if ((e.metaKey || e.ctrlKey) && e.key === 'r') { e.preventDefault(); - useDatabaseStore.getState().triggerRefresh(); + useDatabaseStore.getState().triggerRefresh(true); } }; @@ -124,7 +124,7 @@ function App() { >
triggerRefresh(true)} onCheckForUpdates={() => void updater.checkForUpdates()} isCheckingForUpdates={updater.status === 'checking'} onCommit={() => console.log('Commit active tab')} diff --git a/src/components/DatabaseExplorer.tsx b/src/components/DatabaseExplorer.tsx index e3f20e5..b0c1926 100644 --- a/src/components/DatabaseExplorer.tsx +++ b/src/components/DatabaseExplorer.tsx @@ -81,7 +81,7 @@ const TableItem = ({ table, active, onClick, connectionId }: { const sql = `CREATE TABLE ${quoteId(cloneName.trim())} AS SELECT * FROM ${quoteId(table)}`; try { await invokeWithExecutionPolicy('execute_query', connectionId, expectedDatabase, { kind: 'sql', payload: sql }, { sql }); - triggerRefresh(); + triggerRefresh(true); } catch (err: any) { alert(`Clone failed: ${err}`); } @@ -105,7 +105,7 @@ const TableItem = ({ table, active, onClick, connectionId }: { try { await invokeWithExecutionPolicy('execute_query', connectionId, expectedDatabase, { kind: 'sql', payload: sql }, { sql }); tabs.filter(t => t.type === 'table' && t.tableName === table && t.connectionId === connectionId).forEach(t => closeTab(t.id)); - triggerRefresh(); + triggerRefresh(true); } catch (err: any) { alert(`Delete failed: ${err}`); } diff --git a/src/components/DatabaseSelectorModal.tsx b/src/components/DatabaseSelectorModal.tsx index b89f6f7..6f02e73 100644 --- a/src/components/DatabaseSelectorModal.tsx +++ b/src/components/DatabaseSelectorModal.tsx @@ -84,7 +84,7 @@ export const DatabaseSelectorModal = () => { // Select the new database await setActiveDatabase(newDbName.trim()); - useDatabaseStore.getState().triggerRefresh(); + useDatabaseStore.getState().triggerRefresh(true); handleClose(); } catch (err: any) { setError(err.toString()); diff --git a/src/components/SidebarTree.tsx b/src/components/SidebarTree.tsx index 8324ac2..2507ecf 100644 --- a/src/components/SidebarTree.tsx +++ b/src/components/SidebarTree.tsx @@ -119,7 +119,6 @@ export const SidebarTree = () => { setShowDbName, showConnectionName, setShowConnectionName, - refreshTrigger, fetchSidebarItems, appearanceSettings, } = useDatabaseStore(); @@ -128,7 +127,7 @@ export const SidebarTree = () => { if (selectedConnectionId) { fetchSidebarItems(selectedConnectionId); } - }, [selectedConnectionId, refreshTrigger, fetchSidebarItems]); + }, [selectedConnectionId, fetchSidebarItems]); const [openSections, setOpenSections] = useState>({ Table: true, diff --git a/src/components/TabContentQuery.tsx b/src/components/TabContentQuery.tsx index e8a1ae0..eb2905c 100644 --- a/src/components/TabContentQuery.tsx +++ b/src/components/TabContentQuery.tsx @@ -84,6 +84,7 @@ export const TabContentQuery = ({ id, initialQuery = '', connectionId }: TabCont columns: tab?.columns, stats: tab?.stats }); + const handledStatsRef = useRef(null); // Targeted sync back to store useEffect(() => { @@ -154,17 +155,21 @@ export const TabContentQuery = ({ id, initialQuery = '', connectionId }: TabCont setMessages(prev => [...prev, `[${new Date().toLocaleTimeString()}] Error: ${error}`]); setViewMode('message'); } - if (streamingStats && !isLoading && !error && executedQuery) { + if (streamingStats && !isLoading && !error && executedQuery && handledStatsRef.current !== streamingStats) { + handledStatsRef.current = streamingStats; const queryForResult = executedQuery || query; const cleanQuery = queryForResult.replace(/\/\*[\s\S]*?\*\/|--.*$/gm, '').trim(); const isDdl = /^\s*(CREATE|DROP|ALTER|RENAME|TRUNCATE)\b/i.test(cleanQuery); - const isDml = /^\s*(INSERT|UPDATE|DELETE)\b/i.test(queryForResult); + const isDml = /^\s*(INSERT|UPDATE|DELETE)\b/i.test(cleanQuery); - if (isDdl) { - triggerRefresh(); + if (isDdl || isDml) { + triggerRefresh(isDdl); } let completionMsg = `[${new Date().toLocaleTimeString()}] Query complete: ${streamingStats.rows} rows fetched in ${streamingStats.time}ms`; + if (streamingStats.truncated) { + completionMsg += '. Showing the first 10,000 rows only; use LIMIT/OFFSET to inspect another slice'; + } if (streamingStats.affectedRows && streamingStats.affectedRows > 0) { completionMsg += `. ${streamingStats.affectedRows} rows affected.`; } else if (isDdl || (isDml && streamingStats.affectedRows === 0)) { @@ -498,7 +503,9 @@ export const TabContentQuery = ({ id, initialQuery = '', connectionId }: TabCont {rows.length > 0 ? (
- ✓ {rows.length} {streamingStats.totalRows ? `/ ${streamingStats.totalRows}` : ''} rows + {streamingStats.truncated + ? `Showing first ${rows.length.toLocaleString()} rows (limit reached)` + : `✓ ${rows.length} ${streamingStats.totalRows ? `/ ${streamingStats.totalRows}` : ''} rows`}
) : streamingStats.affectedRows ? (
diff --git a/src/components/TabContentStructure.tsx b/src/components/TabContentStructure.tsx index 5ff02b9..fef1cb9 100644 --- a/src/components/TabContentStructure.tsx +++ b/src/components/TabContentStructure.tsx @@ -2,6 +2,7 @@ import { useState, useEffect } from 'react'; import { invoke } from '@tauri-apps/api/core'; import { useDatabaseStore, TableStructure } from '../store/databaseStore'; import { Package, Key, Hash, ShieldAlert } from 'lucide-react'; +import { tableStructureCache } from '../utils/tableStructureCache'; interface TabContentStructureProps { tableName: string; @@ -21,10 +22,10 @@ export const TabContentStructure = ({ tableName, connectionId, tabId }: TabConte setLoading(true); try { - const result = await invoke('get_table_structure', { - connectionId, - tableName, - }); + const result = await tableStructureCache.get( + connectionId, activeTab?.database || '', tableName, + () => invoke('get_table_structure', { connectionId, tableName }), + ); setTableStructure(tabId, result); } catch (err) { console.error("[ERROR] Failed to fetch table structure:", err); @@ -34,7 +35,7 @@ export const TabContentStructure = ({ tableName, connectionId, tabId }: TabConte }; fetchStructure(); - }, [connectionId, tableName, tabId, structure, setTableStructure]); + }, [connectionId, tableName, tabId, structure, setTableStructure, activeTab?.database]); if (loading) { return ( diff --git a/src/components/TabContentTable.tsx b/src/components/TabContentTable.tsx index 06b26e6..9aa203c 100644 --- a/src/components/TabContentTable.tsx +++ b/src/components/TabContentTable.tsx @@ -15,6 +15,8 @@ import { createDefaultTableFilter } from '../utils/tableFilters'; import { filtersForForeignKey } from '../utils/foreignKeyNavigation'; import { invokeWithExecutionPolicy } from '../utils/executionPolicy'; import { buildRowMutationBatch } from '../utils/rowMutations'; +import { integerKeysetCursor, tablePageCache, tablePageKeys } from '../utils/tablePageCache'; +import { tableStructureCache } from '../utils/tableStructureCache'; import { blockDraftPersistence, getDraftSchemaError, @@ -57,8 +59,26 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable } = useDatabaseStore(); const activeTab = tabs.find(t => t.id === id); const originTab = tabs.find(t => t.id === activeTab?.originTabId && t.connectionId === connectionId); - const [tableData, setTableData] = useState(activeTab?.rows || []); - const [tableColumns, setTableColumns] = useState(activeTab?.columns || []); + const initialKeys = tablePageKeys({ + connectionId, database: activeTab?.database || activeDatabase || '', tableName, + filters: activeTab?.filters?.filter(f => f.enabled) || [], + sortColumn: activeTab?.sortConfig?.column || null, + sortDirection: activeTab?.sortConfig?.direction || 'ASC', + limit: activeTab?.pageSize || 100, offset: activeTab?.offset || 0, + generation: refreshTrigger, + }); + const currentPageIdentity = initialKeys.page; + const initialPage = tablePageCache.getPage(id, initialKeys.page) + || (activeTab?.rowsIdentity === currentPageIdentity && activeTab.rows && activeTab.columns + ? { value: { rows: activeTab.rows, columns: activeTab.columns, time: activeTab.stats?.time || 0 }, fresh: false } + : undefined); + const initialCount = tablePageCache.getCount(id, initialKeys.count); + const [tableData, setTableData] = useState(initialPage?.value.rows || []); + const [tableColumns, setTableColumns] = useState(initialPage?.value.columns || activeTab?.columns || []); + const [loadedPageIdentity, setLoadedPageIdentity] = useState(initialPage ? currentPageIdentity : ''); + const [totalRows, setTotalRows] = useState(initialCount?.value ?? null); + const [hasMore, setHasMore] = useState(initialPage && 'hasMore' in initialPage.value + ? initialPage.value.hasMore ?? false : false); const [loading, setLoading] = useState(false); const [committing, setCommitting] = useState(false); const connection = savedConnections.find(saved => saved.id === connectionId); @@ -69,16 +89,19 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable && (activeTab?.columns?.length || 0) > 0); const [draftStorageError, setDraftStorageError] = useState(null); const [draftSchemaError, setDraftSchemaError] = useState(() => getDraftSchemaError(draftContext)); - const [dataLoaded, setDataLoaded] = useState(() => (activeTab?.columns?.length || 0) > 0); + const [dataLoaded, setDataLoaded] = useState(() => !!initialPage); const [tableLoadError, setTableLoadError] = useState(null); + const [countLoadError, setCountLoadError] = useState(null); + const [countingRows, setCountingRows] = useState(false); const draftLoadStarted = useRef(false); - const [executionTime, setExecutionTime] = useState(activeTab?.stats?.time); + const [executionTime, setExecutionTime] = useState(initialPage?.value.time); const [pkColumns, setPkColumns] = useState([]); const [structureColumns, setStructureColumns] = useState([]); const [nullableColumns, setNullableColumns] = useState([]); const [columnTypes, setColumnTypes] = useState>({}); const [structureLoaded, setStructureLoaded] = useState(false); const [structureError, setStructureError] = useState(null); + const [structureVersion, setStructureVersion] = useState(-1); const [foreignKeys, setForeignKeys] = useState([]); const [showExportModal, setShowExportModal] = useState(false); const [commitError, setCommitError] = useState(null); @@ -88,7 +111,7 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable const isSafeMode = connection?.safetyMode === 'Safe'; const dirtyRef = useRef(false); const commitInFlightRef = useRef(false); - const loadedRowsRef = useRef(activeTab?.rows || []); + const loadedRowsRef = useRef(initialPage?.value.rows || []); const tableDataRef = useRef(tableData); tableDataRef.current = tableData; dirtyRef.current = mutations.state.hasChanges; @@ -164,29 +187,37 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable setStructureLoaded(false); setStructureError(null); try { - const result = await invoke('get_table_structure', { connectionId, tableName }); + const result = await tableStructureCache.get( + connectionId, activeTab?.database || activeDatabase || '', tableName, + () => invoke('get_table_structure', { connectionId, tableName }), + ); if (cancelled) return; - if (result && result.columns) { + if (result && Array.isArray(result.columns)) { setPkColumns(result.columns.filter((c: any) => c.is_primary_key).map((c: any) => c.name)); setStructureColumns(result.columns.map((c: any) => c.name)); setNullableColumns(result.columns.filter((c: any) => c.is_nullable).map((c: any) => c.name)); setColumnTypes(Object.fromEntries(result.columns.map((c: any) => [c.name, c.data_type]))); setForeignKeys(result.foreign_keys || []); setStructureLoaded(true); + setStructureVersion(refreshTrigger); + } else { + throw new Error('Database returned invalid table structure.'); } } catch (err) { if (cancelled) return; console.error("[ERROR] Failed to fetch table structure:", err); setStructureError(String(err)); + setStructureVersion(refreshTrigger); } }; if (tableName && connectionId) { fetchStructure(); } return () => { cancelled = true; }; - }, [tableName, connectionId]); + }, [tableName, connectionId, refreshTrigger]); useEffect(() => { + if (structureVersion !== refreshTrigger || (!structureLoaded && !structureError)) return; let cancelled = false; const fetchData = async () => { if (dirtyRef.current && dataLoaded) { @@ -194,87 +225,141 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable setLoading(false); return; } - setLoading(true); + const limit = activeTab?.pageSize || 100; + const offset = activeTab?.offset || 0; + const filters = activeTab?.filters?.filter(f => f.enabled) || []; + const sortConfig = activeTab?.sortConfig; + const keys = tablePageKeys({ + connectionId, database: activeTab?.database || activeDatabase || '', tableName, + filters, sortColumn: sortConfig?.column || null, + sortDirection: sortConfig?.direction || 'ASC', limit, offset, + generation: refreshTrigger, + }); + const cachedPage = tablePageCache.getPage(id, keys.page); + const retainedPage = !cachedPage && activeTab?.rowsIdentity === keys.page && activeTab.rows && activeTab.columns + ? { value: { columns: activeTab.columns, rows: activeTab.rows, time: activeTab.stats?.time || 0 }, fresh: false } + : undefined; + const visiblePage = cachedPage || retainedPage; + const cachedCount = tablePageCache.getCount(id, keys.count); setTableLoadError(null); - try { - const limit = activeTab?.pageSize || 100; - const offset = activeTab?.offset || 0; - const filters = activeTab?.filters?.filter(f => f.enabled) || []; - const sortConfig = activeTab?.sortConfig; - - const [result, count] = await Promise.all([ - invoke('get_table_data', { - connectionId, - tableName, - limit, - offset, - filters, - sortColumn: sortConfig?.column, - sortDirection: sortConfig?.direction - }), - invoke('get_table_count', { - connectionId, - tableName, - filters - }) - ]); + setCountLoadError(null); + setTotalRows(cachedCount?.fresh ? cachedCount.value : null); + if (visiblePage) { + setTableColumns(visiblePage.value.columns); + setTableData(visiblePage.value.rows); + setDataLoaded(true); + setLoadedPageIdentity(currentPageIdentity); + setExecutionTime(visiblePage.value.time); + setHasMore(cachedPage?.value.hasMore ?? false); + tableDataRef.current = visiblePage.value.rows; + loadedRowsRef.current = visiblePage.value.rows; + updateTab(id, { columns: visiblePage.value.columns, rows: visiblePage.value.rows, rowsIdentity: currentPageIdentity }); + if (cachedPage?.value.hasMore === false && !cachedCount?.fresh) { + setTotalRows(offset + visiblePage.value.rows.length); + } + } + setLoading(!visiblePage); + if (!visiblePage) setHasMore(false); - if (cancelled) return; - if (dirtyRef.current && dataLoaded) { - setCommitError('Apply or discard pending changes before reloading this table.'); - return; - } + if (cachedPage?.fresh) { + return; + } - if (result && result.columns && result.rows) { - setTableColumns(result.columns); - setTableData(result.rows); - setDataLoaded(true); - tableDataRef.current = result.rows; - loadedRowsRef.current = result.rows; - if (result.execution_time_ms) { - setExecutionTime(result.execution_time_ms); - } - - // Construct SQL for history - const sql = `SELECT * FROM "${tableName}"${filters.length ? ' WHERE ...' : ''} LIMIT ${limit} OFFSET ${offset};`; - setLastExecutedSql(sql); - addToHistory({ - sql, - connectionId, - database: activeDatabase || undefined, - executionTimeMs: result.execution_time_ms || 0, - rowsAffected: result.rows.length - }); + try { + const previousKeys = offset >= limit ? tablePageKeys({ + connectionId, database: activeTab?.database || activeDatabase || '', tableName, + filters, sortColumn: null, sortDirection: 'ASC', limit, + offset: offset - limit, generation: refreshTrigger, + }) : null; + const previous = previousKeys ? tablePageCache.getPage(id, previousKeys.page) : undefined; + const seekAfter = previous?.fresh + ? integerKeysetCursor(previous.value, pkColumns, columnTypes, sortConfig?.column || null) + : undefined; + const result = await invoke('get_table_data', { + connectionId, tableName, limit, offset, filters, + sortColumn: sortConfig?.column, sortDirection: sortConfig?.direction, + primaryKeys: pkColumns, seekAfter, + }); + if (cancelled) return; + if (dirtyRef.current && dataLoaded) { + setCommitError('Apply or discard pending changes before reloading this table.'); + return; + } + if (!Array.isArray(result?.columns) || !Array.isArray(result?.rows)) { + throw new Error('Database returned an invalid table page.'); + } + const time = result.execution_time_ms || 0; + const pageHasMore = result.has_more === true; + tablePageCache.putPage(id, keys.page, { columns: result.columns, rows: result.rows, time, hasMore: pageHasMore }); + setTableColumns(result.columns); + setTableData(result.rows); + setDataLoaded(true); + setLoadedPageIdentity(currentPageIdentity); + tableDataRef.current = result.rows; + loadedRowsRef.current = result.rows; + setExecutionTime(time); + setHasMore(pageHasMore); + if (!pageHasMore) { + const count = offset + result.rows.length; + tablePageCache.putCount(id, keys.count, count); + setTotalRows(count); + updateTab(id, { totalRows: count }); + } - // Update store - if (id) { - updateTab(id, { - columns: result.columns, - rows: result.rows, - totalRows: count, - stats: { - time: result.execution_time_ms || 0, - rows: result.rows.length, - totalRows: count - } - }); - } - } + const sql = `SELECT * FROM "${tableName}"${filters.length ? ' WHERE ...' : ''} LIMIT ${limit} OFFSET ${offset};`; + setLastExecutedSql(sql); + addToHistory({ + sql, connectionId, database: activeDatabase || undefined, + executionTimeMs: time, rowsAffected: result.rows.length, + }); + updateTab(id, { + columns: result.columns, rows: result.rows, rowsIdentity: currentPageIdentity, + stats: { time, rows: result.rows.length }, + }); } catch (err) { - if (cancelled) return; - console.error("[ERROR] Failed to fetch table data:", err); - setTableLoadError(String(err)); - setCommitError(`Could not load table data: ${err}`); + if (cancelled) return; + console.error('[ERROR] Failed to fetch table data:', err); + setTableLoadError(String(err)); + setCommitError(`Could not load table data: ${err}`); } finally { - if (!cancelled) setLoading(false); + if (!cancelled) setLoading(false); } }; fetchData(); return () => { cancelled = true; }; - }, [connectionId, tableName, refreshTrigger, activeTab?.offset, activeTab?.pageSize, activeTab?.sortConfig]); + }, [connectionId, tableName, refreshTrigger, activeTab?.offset, activeTab?.pageSize, + activeTab?.sortConfig, structureLoaded, structureError, structureVersion, pkColumns, columnTypes]); // Note: activeTab?.filters is implicitly covered by refreshTrigger when user clicks "Apply" + const handleCountRows = async () => { + if (countingRows || !activeTab) return; + const filters = activeTab.filters?.filter(filter => filter.enabled) || []; + const keys = tablePageKeys({ + connectionId, database: activeTab.database || activeDatabase || '', tableName, filters, + sortColumn: activeTab.sortConfig?.column || null, + sortDirection: activeTab.sortConfig?.direction || 'ASC', + limit: activeTab.pageSize || 100, offset: activeTab.offset || 0, + generation: refreshTrigger, + }); + setCountingRows(true); + setCountLoadError(null); + try { + const count = await invoke('get_table_count', { connectionId, tableName, filters }); + const current = useDatabaseStore.getState(); + const currentTab = current.tabs.find(tab => tab.id === id); + if (!currentTab || current.refreshTrigger !== refreshTrigger || currentTab.connectionId !== connectionId + || JSON.stringify(currentTab.filters?.filter(filter => filter.enabled) || []) !== JSON.stringify(filters)) return; + tablePageCache.putCount(id, keys.count, count); + setTotalRows(count); + updateTab(id, { totalRows: count }); + } catch (err) { + setCountLoadError(String(err)); + } finally { + setCountingRows(false); + } + }; + const handleCommit = async () => { if (!connectionId || activeTab?.originTabId || !draftReady || !metadataMatches || draftSchemaError || !mutations.state.hasChanges || commitInFlightRef.current) return; commitInFlightRef.current = true; @@ -433,12 +518,17 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable Cannot restore row drafts until table structure loads: {structureError}
)} - {tableLoadError && activeTab?.originTabId && ( + {tableLoadError && loadedPageIdentity === currentPageIdentity && (
- Could not load referenced row: {tableLoadError} + Could not refresh table data. Showing the last loaded page: {tableLoadError}
)} - {activeTab?.originTabId && dataLoaded && !loading && !tableLoadError && tableData.length === 0 && ( + {countLoadError && ( +
+ Could not count rows. Pagination still works: {countLoadError} +
+ )} + {activeTab?.originTabId && dataLoaded && loadedPageIdentity === currentPageIdentity && !loading && !tableLoadError && tableData.length === 0 && (
No matching row. It may have been deleted or may not be visible to this connection.
@@ -454,10 +544,15 @@ export const TabContentTable = ({ id, tableName, connectionId }: TabContentTable {viewMode === 'data' ? (
- {tableLoadError && !dataLoaded ? ( + {loading && loadedPageIdentity === currentPageIdentity && ( +
+ Refreshing data... +
+ )} + {tableLoadError && loadedPageIdentity !== currentPageIdentity ? (
Could not load table data: {tableLoadError}
- ) : loading || (!draftReady && !structureError) ? ( -
{loading ? 'Loading data...' : 'Restoring saved draft...'}
+ ) : loadedPageIdentity !== currentPageIdentity || (!draftReady && !structureError) ? ( +
{loadedPageIdentity !== currentPageIdentity ? 'Loading data...' : 'Restoring saved draft...'}
) : ( id && setViewMode(id, type === 'Structure' ? 'structure' : 'data')} onAddRow={viewMode === 'data' && draftReady && metadataMatches && !draftSchemaError && !isSafeMode && !activeTab?.originTabId && !committing && !loading ? handleAddRow : undefined} offset={activeTab?.offset || 0} - pageSize={activeTab?.pageSize || 100} - totalRows={activeTab?.totalRows || 0} + totalRows={totalRows} + displayedRows={tableData.length} + hasMore={hasMore} + countingRows={countingRows} + onCountRows={handleCountRows} onPageChange={handlePageChange} executionTime={executionTime} onToggleFilters={handleToggleFilters} diff --git a/src/components/TableFooter.tsx b/src/components/TableFooter.tsx index f5c7645..650605e 100644 --- a/src/components/TableFooter.tsx +++ b/src/components/TableFooter.tsx @@ -6,8 +6,11 @@ interface TableFooterProps { onTypeChange: (type: 'Data' | 'Structure') => void; onAddRow?: () => void; offset: number; - pageSize: number; - totalRows: number; + totalRows: number | null; + displayedRows: number; + hasMore: boolean; + countingRows: boolean; + onCountRows: () => void; onPageChange: (direction: 'next' | 'prev') => void; executionTime?: number; onToggleFilters?: () => void; @@ -22,8 +25,11 @@ export const TableFooter = ({ onTypeChange, onAddRow, offset, - pageSize, totalRows, + displayedRows, + hasMore, + countingRows, + onCountRows, onPageChange, executionTime, onToggleFilters, @@ -32,8 +38,8 @@ export const TableFooter = ({ isColumnsVisible, onExport }: TableFooterProps) => { - const start = totalRows > 0 ? offset + 1 : 0; - const end = Math.min(offset + pageSize, totalRows); + const start = displayedRows > 0 ? offset + 1 : 0; + const end = offset + displayedRows; return (
@@ -76,8 +82,15 @@ export const TableFooter = ({ {executionTime}ms
)} -
- {start}-{end} of {totalRows} {totalRows === 1 ? 'row' : 'rows'} +
+ {start}-{end}{totalRows === null ? ` rows${hasMore ? '+' : ''}` : ` of ${totalRows} ${totalRows === 1 ? 'row' : 'rows'}`} + {totalRows === null && ( + + )}
@@ -123,7 +136,7 @@ export const TableFooter = ({