diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index 410ffe7dad..98f3044a90 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -40,6 +40,8 @@ const XLSX_CANCELLABLE_READ_CHUNK_BYTES: usize = 64 * 1024; const XLSX_CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(25); // INSERT ALL has a practical statement-size limit on Oracle, even when the requested batch is larger. const MAX_ORACLE_IMPORT_BATCH_ROWS: usize = 500; +const SQLITE_APPEND_COMMIT_ROWS: usize = 10_000; +const SQLITE_APPEND_COMMIT_SQL_BYTES: usize = 8 * 1024 * 1024; const POSTGRES_COPY_TARGET_BYTES: usize = 8 * 1024 * 1024; const POSTGRES_COPY_MAX_ROWS: usize = 50_000; // Bound the additional memory used while converting one source row into owned @@ -3266,6 +3268,7 @@ fn effective_import_batch_size(db_type: &DatabaseType, requested: usize) -> usiz DatabaseType::OceanbaseOracle | DatabaseType::Iris => 1, DatabaseType::CloudflareD1 => 100, DatabaseType::SqlServer => 1000, + DatabaseType::Sqlite => SQLITE_APPEND_COMMIT_ROWS, _ => usize::MAX, }; requested.max(1).min(max_rows) @@ -4146,6 +4149,61 @@ struct ImportBatchExecutionPolicy { allow_postgres_copy: bool, } +#[derive(Debug)] +struct SqliteAppendTransaction { + statements: Vec, + rows: usize, + sql_bytes: usize, + max_rows: usize, + max_sql_bytes: usize, +} + +impl SqliteAppendTransaction { + fn new() -> Self { + Self::with_limits(SQLITE_APPEND_COMMIT_ROWS, SQLITE_APPEND_COMMIT_SQL_BYTES) + } + + fn with_limits(max_rows: usize, max_sql_bytes: usize) -> Self { + Self { + statements: Vec::new(), + rows: 0, + sql_bytes: 0, + max_rows: max_rows.max(1), + max_sql_bytes: max_sql_bytes.max(1), + } + } + + fn should_flush_before(&self, batch: &ImportSqlBatch) -> bool { + !self.statements.is_empty() + && (self.rows.saturating_add(batch.row_count) > self.max_rows + || self.sql_bytes.saturating_add(batch.sql.len()) > self.max_sql_bytes) + } + + fn push(&mut self, batch: ImportSqlBatch) { + self.rows = self.rows.saturating_add(batch.row_count); + self.sql_bytes = self.sql_bytes.saturating_add(batch.sql.len()); + self.statements.push(batch.sql); + } + + fn is_ready(&self) -> bool { + self.rows >= self.max_rows || self.sql_bytes >= self.max_sql_bytes + } + + fn take(&mut self) -> (Vec, usize) { + let statements = std::mem::take(&mut self.statements); + let rows = std::mem::take(&mut self.rows); + self.sql_bytes = 0; + (statements, rows) + } +} + +fn sqlite_append_transaction_for_import( + mode: &TableImportMode, + db_type: &DatabaseType, +) -> Option { + (matches!(mode, TableImportMode::Append) && *db_type == DatabaseType::Sqlite).then(SqliteAppendTransaction::new) +} + fn supports_transactional_import_truncate(db_type: &DatabaseType) -> bool { matches!( db_type, @@ -4204,6 +4262,92 @@ async fn execute_import_transaction( result } +#[allow(clippy::too_many_arguments)] +async fn flush_sqlite_append_transaction( + state: &AppState, + pool_key: &str, + import_id: &str, + is_cancelled: &impl Fn(&str) -> std::pin::Pin + Send + '_>>, + connection_id: &str, + database: &str, + schema: &str, + transaction: &mut SqliteAppendTransaction, + rows_imported: usize, + db_write_ms: &mut u128, + statement_count: &mut usize, +) -> Result { + if transaction.statements.is_empty() { + return Ok(rows_imported); + } + ensure_import_write_allowed(import_id, is_cancelled, rows_imported).await?; + let (statements, rows) = transaction.take(); + execute_import_transaction( + state, + pool_key, + connection_id, + database, + schema, + &statements, + db_write_ms, + statement_count, + ) + .await + .map_err(|message| ImportRowsBatchError::with_rows_imported(rows_imported, message))?; + Ok(rows_imported.saturating_add(rows)) +} + +async fn finish_sqlite_append_transaction( + state: &AppState, + pool_key: &str, + request: &TableImportRequest, + is_cancelled: &impl Fn(&str) -> std::pin::Pin + Send + '_>>, + transaction: &mut Option, + rows_imported: usize, + total_rows: usize, + started_at: Instant, + db_write_ms: &mut u128, + statement_count: &mut usize, + progress_callback: &mut F, +) -> Result +where + F: FnMut(TableImportProgress), +{ + let Some(transaction) = transaction.as_mut() else { + return Ok(rows_imported); + }; + match flush_sqlite_append_transaction( + state, + pool_key, + &request.import_id, + is_cancelled, + &request.connection_id, + &request.database, + &request.schema, + transaction, + rows_imported, + db_write_ms, + statement_count, + ) + .await + { + Ok(rows) => Ok(rows), + Err(error) if error.cancelled => { + progress_callback(import_progress( + &request.import_id, + TableImportStatus::Cancelled, + rows_imported, + total_rows, + started_at, + None, + )); + Err(error.message) + } + Err(error) => { + Err(emit_import_error(progress_callback, request, rows_imported, total_rows, started_at, error.message)) + } + } +} + #[allow(clippy::too_many_arguments)] async fn execute_import_rows_batch( state: &AppState, @@ -4224,6 +4368,7 @@ async fn execute_import_rows_batch( mode: &TableImportMode, pending_truncate: bool, postgres_copy_accumulator: &mut Option, + sqlite_append_transaction: &mut Option, kingbase_oracle_mode: bool, date_time_format: Option<&str>, hard_sql_bytes: Option, @@ -4296,6 +4441,45 @@ async fn execute_import_rows_batch( hard_sql_bytes, ) .map_err(|message| ImportRowsBatchError::with_rows_imported(rows_imported, message))?; + if let Some(transaction) = sqlite_append_transaction.as_mut() { + for batch in batches { + ensure_import_write_allowed(import_id, is_cancelled, rows_imported).await?; + if transaction.should_flush_before(&batch) { + rows_imported = flush_sqlite_append_transaction( + state, + pool_key, + import_id, + is_cancelled, + connection_id, + database, + schema, + transaction, + rows_imported, + db_write_ms, + statement_count, + ) + .await?; + } + transaction.push(batch); + if transaction.is_ready() { + rows_imported = flush_sqlite_append_transaction( + state, + pool_key, + import_id, + is_cancelled, + connection_id, + database, + schema, + transaction, + rows_imported, + db_write_ms, + statement_count, + ) + .await?; + } + } + return Ok(rows_imported); + } if execution_policy.transactional { ensure_import_write_allowed(import_id, is_cancelled, rows_imported).await?; let mut statements = Vec::with_capacity(batches.len() + usize::from(execution_policy.include_truncate)); @@ -5367,6 +5551,7 @@ where &request.table, &request.schema, ); + let mut sqlite_append_transaction = sqlite_append_transaction_for_import(&request.mode, db_type); let mut pending_truncate = matches!(request.mode, TableImportMode::Truncate) && supports_transactional_import_truncate(db_type); if matches!(request.mode, TableImportMode::Truncate) && !pending_truncate { @@ -5427,6 +5612,7 @@ where &request.mode, pending_truncate, &mut postgres_copy_accumulator, + &mut sqlite_append_transaction, kingbase_oracle_mode, request.date_time_format.as_deref(), import_sql_hard_limit, @@ -5524,6 +5710,20 @@ where )); } } + rows_imported = finish_sqlite_append_transaction( + state, + pool_key, + request, + &is_cancelled, + &mut sqlite_append_transaction, + rows_imported, + total_rows, + started_at, + &mut db_write_ms, + &mut statement_count, + &mut progress_callback, + ) + .await?; let flushed_rows = match flush_pending_postgres_copy( state, pool_key, @@ -5830,6 +6030,7 @@ where &request.table, &request.schema, ); + let mut sqlite_append_transaction = sqlite_append_transaction_for_import(&request.mode, db_type); let mut pending_truncate = matches!(request.mode, TableImportMode::Truncate) && supports_transactional_import_truncate(db_type); if matches!(request.mode, TableImportMode::Truncate) && !pending_truncate { @@ -5913,6 +6114,7 @@ where &request.mode, pending_truncate, &mut postgres_copy_accumulator, + &mut sqlite_append_transaction, kingbase_oracle_mode, request.date_time_format.as_deref(), import_sql_hard_limit, @@ -6015,6 +6217,20 @@ where )); } } + rows_imported = finish_sqlite_append_transaction( + state, + pool_key, + request, + &is_cancelled, + &mut sqlite_append_transaction, + rows_imported, + 0, + started_at, + &mut db_write_ms, + &mut statement_count, + &mut progress_callback, + ) + .await?; let flushed_rows = match flush_pending_postgres_copy( state, pool_key, @@ -6162,6 +6378,7 @@ where &request.table, &request.schema, ); + let mut sqlite_append_transaction = sqlite_append_transaction_for_import(&request.mode, db_type); let mut pending_truncate = matches!(request.mode, TableImportMode::Truncate) && supports_transactional_import_truncate(db_type); @@ -6207,6 +6424,7 @@ where &request.mode, pending_truncate, &mut postgres_copy_accumulator, + &mut sqlite_append_transaction, kingbase_oracle_mode, request.date_time_format.as_deref(), import_sql_hard_limit, @@ -6254,6 +6472,22 @@ where } } + rows_imported = finish_sqlite_append_transaction( + state, + pool_key, + request, + &is_cancelled, + &mut sqlite_append_transaction, + rows_imported, + total_rows, + started_at, + &mut db_write_ms, + &mut statement_count, + &mut progress_callback, + ) + .await? + .min(total_rows); + let flushed_rows = flush_pending_postgres_copy( state, pool_key, @@ -9251,6 +9485,7 @@ mod tests { let cancellation_checks = Arc::new(AtomicUsize::new(0)); let checks_for_import = cancellation_checks.clone(); let mut postgres_copy_accumulator = None; + let mut sqlite_append_transaction = None; let mut db_write_ms = 0; let mut statement_count = 0; @@ -9276,6 +9511,7 @@ mod tests { &TableImportMode::Append, false, &mut postgres_copy_accumulator, + &mut sqlite_append_transaction, false, None, None, @@ -9423,6 +9659,220 @@ mod tests { assert!(!policy.allow_postgres_copy); } + fn sqlite_append_test_plan() -> CompiledImportPlan { + CompiledImportPlan { + mapped_source_indexes: vec![0], + target_columns: vec!["id".to_string()], + column_types: vec![Some("integer".to_string())], + } + } + + struct SqliteAppendTestContext { + _dir: tempfile::TempDir, + state: AppState, + sqlite: crate::db::sqlite::SqliteHandle, + pool_key: String, + plan: CompiledImportPlan, + postgres_copy_accumulator: Option, + transaction: Option, + db_write_ms: u128, + statement_count: usize, + } + + impl SqliteAppendTestContext { + async fn new(test_name: &str, max_rows: usize) -> Self { + let dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let pool_key = format!("{test_name}:session:import"); + let database_path = dir.path().join("target.db"); + let sqlite = + crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); + crate::db::sqlite::execute_query(&sqlite, "CREATE TABLE items (id INTEGER PRIMARY KEY)").await.unwrap(); + state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + Self { + _dir: dir, + state, + sqlite, + pool_key, + plan: sqlite_append_test_plan(), + postgres_copy_accumulator: None, + transaction: Some(SqliteAppendTransaction::with_limits(max_rows, usize::MAX)), + db_write_ms: 0, + statement_count: 0, + } + } + + async fn append(&mut self, ids: &[i64]) -> Result { + let rows = ids.iter().map(|id| vec![serde_json::json!(id)]).collect::>(); + execute_import_rows_batch( + &self.state, + &self.pool_key, + &self.pool_key, + &|_| Box::pin(async { false }), + &self.pool_key, + "", + &rows, + Some(&self.plan), + None, + &[], + &[], + &[], + "items", + "", + &DatabaseType::Sqlite, + &TableImportMode::Append, + false, + &mut self.postgres_copy_accumulator, + &mut self.transaction, + false, + None, + None, + &mut self.db_write_ms, + &mut self.statement_count, + ) + .await + } + + async fn ids(&self) -> Vec> { + crate::db::sqlite::execute_query(&self.sqlite, "SELECT id FROM items ORDER BY id").await.unwrap().rows + } + } + + #[tokio::test] + async fn sqlite_append_commits_only_bounded_row_windows() { + let mut context = SqliteAppendTestContext::new("sqlite-append-window", 3).await; + let first = context.append(&[1, 2]).await.unwrap(); + assert_eq!(first, 0); + assert!(context.ids().await.is_empty()); + + let second = context.append(&[3]).await.unwrap(); + assert_eq!(second, 3); + assert_eq!(context.ids().await.len(), 3); + assert_eq!(context.statement_count, 2); + } + + #[tokio::test] + async fn sqlite_append_failure_keeps_prior_window_and_rolls_back_current_window() { + let mut context = SqliteAppendTestContext::new("sqlite-append-failure", 2).await; + let committed = context.append(&[1, 2]).await.unwrap(); + assert_eq!(committed, 2); + context.transaction.as_mut().unwrap().max_rows = 4; + + let pending = context.append(&[3, 4]).await.unwrap(); + assert_eq!(pending, 0); + let error = context.append(&[5, 1]).await.unwrap_err(); + assert_eq!(error.rows_imported, 0); + assert_eq!(context.ids().await, vec![vec![serde_json::json!(1)], vec![serde_json::json!(2)]]); + } + + #[tokio::test] + async fn sqlite_append_cancellation_drops_the_uncommitted_window() { + let mut context = SqliteAppendTestContext::new("sqlite-append-cancel", 10).await; + context.append(&[1, 2]).await.unwrap(); + + let error = flush_sqlite_append_transaction( + &context.state, + &context.pool_key, + "sqlite-append-cancel", + &|_| Box::pin(async { true }), + "sqlite-append-cancel", + "", + "", + context.transaction.as_mut().unwrap(), + 0, + &mut context.db_write_ms, + &mut context.statement_count, + ) + .await + .unwrap_err(); + assert!(error.cancelled); + assert_eq!(error.rows_imported, 0); + assert!(context.ids().await.is_empty()); + } + + #[tokio::test] + async fn delimited_sqlite_append_import_flushes_the_final_window() { + let dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "sqlite-delimited-append"; + let pool_key = format!("{connection_id}:session:import"); + let database_path = dir.path().join("target.db"); + let sqlite = crate::db::sqlite::connect_path_create_if_missing(database_path.to_str().unwrap()).await.unwrap(); + crate::db::sqlite::execute_query(&sqlite, "CREATE TABLE items (id INTEGER, name TEXT)").await.unwrap(); + state.connections.write().await.insert(pool_key.clone(), PoolKind::Sqlite(sqlite.clone())); + let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": connection_id, + "name": "SQLite delimited append test", + "db_type": "sqlite", + "host": "", + "port": 0, + "username": "", + "password": "", + "database": database_path.to_string_lossy() + })) + .unwrap(); + state.configs.write().await.insert(connection_id.to_string(), config); + let data_path = dir.path().join("rows.txt"); + std::fs::write(&data_path, b"id%name\n1%Ada\n2%Grace\n3%Linus\n").unwrap(); + let request = TableImportRequest { + import_id: "sqlite-delimited-append".to_string(), + connection_id: connection_id.to_string(), + database: String::new(), + schema: String::new(), + table: "items".to_string(), + file_path: data_path.to_string_lossy().to_string(), + source_ref: None, + source_format: Some(TableImportSourceFormat::Delimited), + parse_options: TableImportParseOptions { + delimiter: Some("%".to_string()), + ..TableImportParseOptions::default() + }, + mappings: vec![ + TableImportColumnMapping { + source_column: "id".to_string(), + target_column: "id".to_string(), + target_data_type: None, + }, + TableImportColumnMapping { + source_column: "name".to_string(), + target_column: "name".to_string(), + target_data_type: None, + }, + ], + mode: TableImportMode::Append, + create_table: false, + batch_size: 2, + date_time_format: None, + prepared_source: None, + retain_source: false, + }; + + let summary = import_table_file_core( + &state, + &request, + &DatabaseType::Sqlite, + &pool_key, + |_| Box::pin(async { false }), + |_| {}, + ) + .await + .unwrap(); + + assert_eq!(summary.rows_imported, 3); + let rows = + crate::db::sqlite::execute_query(&sqlite, "SELECT id, name FROM items ORDER BY id").await.unwrap().rows; + assert_eq!( + rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("Ada")], + vec![serde_json::json!(2), serde_json::json!("Grace")], + vec![serde_json::json!(3), serde_json::json!("Linus")] + ] + ); + } + #[tokio::test] async fn preview_missing_source_fails_before_parsing() { let path = std::env::temp_dir().join(format!("dbx-missing-import-{}.csv", uuid::Uuid::new_v4()));