From 68eae4ad4d1d3a7d48da2c4519c3359cede368ff Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 26 Aug 2026 11:03:44 +0800 Subject: [PATCH] =?UTF-8?q?feat(gaussdb):=20=E6=94=AF=E6=8C=81=E5=B8=A6=20?= =?UTF-8?q?psql=20=E6=8E=A7=E5=88=B6=E5=91=BD=E4=BB=A4=E7=9A=84=E5=8C=BF?= =?UTF-8?q?=E5=90=8D=E5=9D=97=E6=89=A7=E8=A1=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/dbx-core/src/query.rs | 51 ++++- crates/dbx-core/src/sql.rs | 274 ++++++++++++++++++++++++- crates/dbx-core/src/sql_file_import.rs | 153 ++++++++++++-- 3 files changed, 444 insertions(+), 34 deletions(-) diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 9d752e8372..9301fb8717 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -2625,10 +2625,12 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( ); } - let statements = db_type.map_or_else( - || split_sql_statements(sql), - |db_type| crate::sql::split_sql_statements_for_database(sql, db_type), + let execution_plan = db_type.map_or_else( + || crate::sql::SqlExecutionPlan { statements: split_sql_statements(sql), stop_on_error: false }, + |db_type| crate::sql::sql_execution_plan_for_database(sql, db_type), ); + let continue_on_error = options.continue_on_error && !execution_plan.stop_on_error; + let statements = execution_plan.statements; if statements.is_empty() { return Ok(vec![empty_query_result(0).into()]); } @@ -2756,7 +2758,7 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( Some(statement_index), backend_error, )); - if !should_continue_batch_after_error(options.continue_on_error, action) { + if !should_continue_batch_after_error(continue_on_error, action) { break; } } @@ -6098,6 +6100,47 @@ for line in sys.stdin: assert_sqlite_batch_error_behavior(false, true).await; } + #[tokio::test] + async fn gaussdb_on_error_stop_overrides_continue_on_error() { + let dir = std::env::temp_dir().join(format!("dbx-query-gaussdb-on-error-stop-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "gaussdb-on-error-stop"; + let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); + state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Gaussdb)); + + let results = execute_multi_core_with_options( + &state, + connection_id, + "", + "\\set ON_ERROR_STOP on\nINSERT INTO missing_table VALUES (1); CREATE TABLE must_not_run (id INTEGER);", + None, + None, + QueryExecutionOptions { continue_on_error: true, ..Default::default() }, + ) + .await + .unwrap(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].columns, vec!["Error"]); + let table_check = execute_sql_statement( + &state, + connection_id, + "", + "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'must_not_run'", + None, + None, + ) + .await + .unwrap(); + assert!(table_check.rows.is_empty()); + + drop(state); + let _ = std::fs::remove_dir_all(dir); + } + #[tokio::test] async fn mysql_batch_stops_after_the_first_statement_error() { let statements = vec!["first".to_string(), "fails".to_string(), "must-not-run".to_string()]; diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index 6c68872d88..6f73e74437 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -119,6 +119,7 @@ struct SqlDialectProfile { keeps_sqlserver_module_batch_at_cursor: bool, preserves_tdsql_leading_directives: bool, requires_whitespace_after_line_comment_dashes: bool, + supports_psql_control_commands: bool, } impl Default for SqlDialectProfile { @@ -136,6 +137,7 @@ impl Default for SqlDialectProfile { keeps_sqlserver_module_batch_at_cursor: false, preserves_tdsql_leading_directives: false, requires_whitespace_after_line_comment_dashes: false, + supports_psql_control_commands: false, } } } @@ -187,7 +189,11 @@ impl SqlDialectProfile { } fn gaussdb() -> Self { - Self { supports_postgres_dollar_quoted_routines: true, ..Self::oracle_like() } + Self { + supports_postgres_dollar_quoted_routines: true, + supports_psql_control_commands: true, + ..Self::oracle_like() + } } fn sql_server() -> Self { @@ -320,15 +326,26 @@ pub struct SqlStatementSplitter { previous: Option, pending_mysql_line_comment_dashes: bool, custom_delimiter: Option, + stop_on_error: bool, options: SqlParsingOptions, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SqlStatementWithControl { + pub sql: String, + pub stop_on_error: bool, +} + impl SqlStatementSplitter { pub fn with_options(options: SqlParsingOptions) -> Self { Self { options, ..Self::default() } } pub fn push_chunk(&mut self, chunk: &str) -> Vec { + self.push_chunk_with_control(chunk).into_iter().map(|statement| statement.sql).collect() + } + + pub(crate) fn push_chunk_with_control(&mut self, chunk: &str) -> Vec { let mut statements = Vec::new(); let chars = chunk.chars().collect::>(); let mut i = 0; @@ -501,10 +518,24 @@ impl SqlStatementSplitter { let buf_end = self.buffer.len() - 1; let last_line_start = self.buffer[..buf_end].rfind('\n').map_or(0, |p| p + 1); let last_line = self.buffer[last_line_start..buf_end].trim(); + if self.options.profile.supports_psql_control_commands { + if let Some(command) = parse_ignorable_psql_control_command(last_line) { + if command == PsqlControlCommand::OnErrorStop { + self.stop_on_error = true; + } + self.buffer.truncate(last_line_start); + self.previous = self.buffer.chars().next_back(); + i += 1; + continue; + } + } if self.options.profile.supports_slash_line_block_delimiter && last_line == "/" { let before = self.buffer[..last_line_start].trim(); if has_executable_sql_with_options(before, self.options) { - statements.push(before.to_string()); + statements.push(SqlStatementWithControl { + sql: before.to_string(), + stop_on_error: self.stop_on_error, + }); } self.buffer.clear(); self.postgres_dollar_quoted_routine = false; @@ -523,7 +554,10 @@ impl SqlStatementSplitter { if last_line_start > 0 { let before = self.buffer[..last_line_start].trim(); if has_executable_sql_with_options(before, self.options) { - statements.push(before.to_string()); + statements.push(SqlStatementWithControl { + sql: before.to_string(), + stop_on_error: self.stop_on_error, + }); } } self.buffer.clear(); @@ -548,7 +582,11 @@ impl SqlStatementSplitter { statements } - pub fn finish(mut self) -> Vec { + pub fn finish(self) -> Vec { + self.finish_with_control().into_iter().map(|statement| statement.sql).collect() + } + + pub(crate) fn finish_with_control(mut self) -> Vec { let mut statements = Vec::new(); if self.pending_mysql_line_comment_dashes { self.in_line_comment = true; @@ -556,16 +594,27 @@ impl SqlStatementSplitter { } let trimmed = self.buffer.trim(); let last_line = trimmed.rsplit('\n').next().unwrap_or(trimmed).trim(); + if self.options.profile.supports_psql_control_commands { + if let Some(command) = parse_ignorable_psql_control_command(last_line) { + if command == PsqlControlCommand::OnErrorStop { + self.stop_on_error = true; + } + let line_start = self.buffer.rfind('\n').map_or(0, |pos| pos + 1); + self.buffer.truncate(line_start); + } + } + let trimmed = self.buffer.trim(); + let last_line = trimmed.rsplit('\n').next().unwrap_or(trimmed).trim(); if self.options.profile.supports_custom_delimiter_commands && parse_delimiter_command(last_line).is_some() { let before = trimmed.rsplit_once('\n').map(|x| x.0).unwrap_or("").trim(); if has_executable_sql_with_options(before, self.options) { - statements.push(before.to_string()); + statements.push(SqlStatementWithControl { sql: before.to_string(), stop_on_error: self.stop_on_error }); } self.buffer.clear(); } else if self.options.profile.supports_slash_line_block_delimiter && last_line == "/" { let before = trimmed.rsplit_once('\n').map(|x| x.0).unwrap_or("").trim(); if has_executable_sql_with_options(before, self.options) { - statements.push(before.to_string()); + statements.push(SqlStatementWithControl { sql: before.to_string(), stop_on_error: self.stop_on_error }); } self.buffer.clear(); } else if let Some(ref delim) = self.custom_delimiter { @@ -577,10 +626,14 @@ impl SqlStatementSplitter { statements } - fn push_current_statement(&mut self, statements: &mut Vec) { + pub fn stop_on_error(&self) -> bool { + self.stop_on_error + } + + fn push_current_statement(&mut self, statements: &mut Vec) { let statement = self.buffer.trim(); if has_executable_sql_with_options(statement, self.options) { - statements.push(statement.to_string()); + statements.push(SqlStatementWithControl { sql: statement.to_string(), stop_on_error: self.stop_on_error }); } self.buffer.clear(); self.postgres_dollar_quoted_routine = false; @@ -604,7 +657,25 @@ pub fn split_sql_statements(sql: &str) -> Vec { } pub fn split_sql_statements_for_database(sql: &str, db_type: DatabaseType) -> Vec { - split_sql_statements_with_options(sql, SqlParsingOptions::for_database_type(db_type)) + sql_execution_plan_for_database(sql, db_type).statements +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SqlExecutionPlan { + pub statements: Vec, + pub stop_on_error: bool, +} + +pub fn sql_execution_plan_for_database(sql: &str, db_type: DatabaseType) -> SqlExecutionPlan { + let options = SqlParsingOptions::for_database_type(db_type); + if !options.profile.supports_psql_control_commands { + return SqlExecutionPlan { statements: split_sql_statements_with_options(sql, options), stop_on_error: false }; + } + + let (sql, stop_on_error) = preprocess_psql_control_commands(sql, options.profile); + let mut parsing_options = options; + parsing_options.profile.supports_psql_control_commands = false; + SqlExecutionPlan { statements: split_sql_statements_with_options(&sql, parsing_options), stop_on_error } } pub fn split_sql_statements_with_options(sql: &str, options: SqlParsingOptions) -> Vec { @@ -629,7 +700,13 @@ pub fn find_statement_at_cursor_for_database(sql: &str, cursor_pos: usize, db_ty if db_type == DatabaseType::SqlServer { return find_sqlserver_statement_at_cursor(sql, cursor_pos); } - find_statement_at_cursor_with_options(sql, cursor_pos, SqlParsingOptions::for_database_type(db_type)) + let mut options = SqlParsingOptions::for_database_type(db_type); + if options.profile.supports_psql_control_commands { + let (sql, _) = preprocess_psql_control_commands(sql, options.profile); + options.profile.supports_psql_control_commands = false; + return find_statement_at_cursor_with_options(&sql, cursor_pos, options); + } + find_statement_at_cursor_with_options(sql, cursor_pos, options) } pub fn find_statement_at_cursor_with_options(sql: &str, cursor_pos: usize, options: SqlParsingOptions) -> String { @@ -1417,6 +1494,70 @@ fn parse_delimiter_command(line: &str) -> Option<&str> { rest.map(|r| r.trim()).filter(|r| !r.is_empty()) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PsqlControlCommand { + OnErrorStop, + ClientDisplay, +} + +fn parse_ignorable_psql_control_command(line: &str) -> Option { + let mut parts = line.split_whitespace(); + let command = parts.next()?; + + if command.eq_ignore_ascii_case("\\timing") { + let valid = + parts.next().is_none_or(|value| value.eq_ignore_ascii_case("on") || value.eq_ignore_ascii_case("off")); + return (valid && parts.next().is_none()).then_some(PsqlControlCommand::ClientDisplay); + } + + if !command.eq_ignore_ascii_case("\\set") { + return None; + } + let variable = parts.next()?; + let value = parts.next()?; + if parts.next().is_some() { + return None; + } + + match variable.to_ascii_uppercase().as_str() { + "ON_ERROR_STOP" if matches!(value.to_ascii_lowercase().as_str(), "on" | "true" | "1") => { + Some(PsqlControlCommand::OnErrorStop) + } + "VERBOSITY" if matches!(value.to_ascii_lowercase().as_str(), "default" | "verbose" | "terse" | "sqlstate") => { + Some(PsqlControlCommand::ClientDisplay) + } + _ => None, + } +} + +fn preprocess_psql_control_commands(sql: &str, profile: SqlDialectProfile) -> (String, bool) { + let mut output = String::with_capacity(sql.len()); + let mut scanner = SqlScanner::with_profile(profile); + let mut offset = 0usize; + let mut stop_on_error = false; + + for segment in sql.split_inclusive('\n') { + let line = segment.strip_suffix('\n').unwrap_or(segment); + let command = (!scanner.is_masked()).then(|| parse_ignorable_psql_control_command(line.trim())).flatten(); + if let Some(command) = command { + stop_on_error |= command == PsqlControlCommand::OnErrorStop; + output.extend(line.chars().map(|ch| if ch.is_whitespace() { ch } else { ' ' })); + if segment.ends_with('\n') { + output.push('\n'); + scanner.step(sql, offset + line.len(), '\n'); + } + } else { + output.push_str(segment); + for (relative_idx, ch) in segment.char_indices() { + scanner.step(sql, offset + relative_idx, ch); + } + } + offset += segment.len(); + } + + (output, stop_on_error) +} + pub fn statement_summary(statement: &str) -> String { const MAX_LEN: usize = 120; @@ -3482,6 +3623,119 @@ END;"; assert_eq!(split_sql_statements_for_database(sql, DatabaseType::Xugu), vec![sql.to_string()]); } + #[test] + fn gaussdb_split_ignores_psql_controls_before_anonymous_block_per_issue_6468() { + let block = "\ +/* 将逻辑放到匿名块中，便于捕获异常 pl/pgsql */ +DECLARE + v_date_str VARCHAR(20) := '20240824'; + v_date_dt DATE := TO_DATE(v_date_str, 'YYYYMMDD'); + v_eff_num NUMERIC := 0; + v_start_time TIMESTAMP := CURRENT_TIMESTAMP; + v_step TEXT := ''; + v_job_name VARCHAR := 'F_XY_DG_IBDISCOUNT_DET'; +BEGIN + RAISE NOTICE 'issue 6468'; +END;"; + let sql = format!("\\set ON_ERROR_STOP on\n\\set VERBOSITY verbose\n\\timing\n\n{block}\n/"); + + assert_eq!(split_sql_statements_for_database(&sql, DatabaseType::Gaussdb), vec![block]); + + let cursor = sql[..sql.find("RAISE NOTICE").unwrap()].encode_utf16().count(); + let executable_block = &block[block.find("DECLARE").unwrap()..]; + assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, DatabaseType::Gaussdb), executable_block); + } + + #[test] + fn gaussdb_streaming_split_ignores_chunked_psql_controls() { + let mut splitter = + SqlStatementSplitter::with_options(SqlParsingOptions::for_database_type(DatabaseType::Gaussdb)); + + assert!(splitter.push_chunk("\\set ON_ERROR_").is_empty()); + assert!(splitter.push_chunk("STOP on\n\\tim").is_empty()); + assert!(splitter.stop_on_error()); + assert_eq!( + splitter.push_chunk("ing\nDECLARE\n value NUMBER := 1;\nBEGIN\n value := value + 1;\nEND;\n/\n"), + vec!["DECLARE\n value NUMBER := 1;\nBEGIN\n value := value + 1;\nEND;"] + ); + assert!(splitter.finish().is_empty()); + } + + #[test] + fn gaussdb_streaming_split_tracks_on_error_stop_by_statement_position() { + let mut splitter = + SqlStatementSplitter::with_options(SqlParsingOptions::for_database_type(DatabaseType::Gaussdb)); + + let mut statements = splitter.push_chunk_with_control("SELECT 0;\nSELECT\n\\set ON_ERROR_"); + statements.extend(splitter.push_chunk_with_control("STOP on\n 1;\nSELECT 2;")); + statements.extend(splitter.finish_with_control()); + + assert_eq!( + statements, + vec![ + super::SqlStatementWithControl { sql: "SELECT 0".to_string(), stop_on_error: false }, + super::SqlStatementWithControl { sql: "SELECT\n 1".to_string(), stop_on_error: true }, + super::SqlStatementWithControl { sql: "SELECT 2".to_string(), stop_on_error: true }, + ] + ); + } + + #[test] + fn gaussdb_psql_control_does_not_submit_unfinished_query_buffer() { + let sql = "SELECT\n\\timing\n 1;\nSELECT 2;"; + let statements = split_sql_statements_for_database(sql, DatabaseType::Gaussdb); + + assert_eq!(statements.len(), 2); + assert_eq!(statements[0].split_whitespace().collect::>(), vec!["SELECT", "1"]); + assert_eq!(statements[1], "SELECT 2"); + + let mut splitter = + SqlStatementSplitter::with_options(SqlParsingOptions::for_database_type(DatabaseType::Gaussdb)); + assert!(splitter.push_chunk("SELECT\n\\tim").is_empty()); + assert_eq!(splitter.push_chunk("ing\n 1;"), vec!["SELECT\n 1"]); + } + + #[test] + fn gaussdb_split_does_not_ignore_psql_variable_assignment() { + for control in ["\\set tx_date 20240824", "\\set MY_ON_ERROR_STOP on", "\\set ON_ERROR_STOP_EXTRA on"] { + let sql = format!("{control}\nDECLARE\n value TEXT := '$tx_date';\nBEGIN\n NULL;\nEND;\n/"); + let statements = split_sql_statements_for_database(&sql, DatabaseType::Gaussdb); + assert!(statements[0].starts_with(control)); + } + } + + #[test] + fn gaussdb_ignores_only_safe_psql_controls() { + assert_eq!( + super::parse_ignorable_psql_control_command("\\set ON_ERROR_STOP on"), + Some(super::PsqlControlCommand::OnErrorStop) + ); + assert_eq!( + super::parse_ignorable_psql_control_command("\\set VERBOSITY verbose"), + Some(super::PsqlControlCommand::ClientDisplay) + ); + assert_eq!( + super::parse_ignorable_psql_control_command("\\timing"), + Some(super::PsqlControlCommand::ClientDisplay) + ); + assert_eq!(super::parse_ignorable_psql_control_command("\\set ON_ERROR_STOP off"), None); + assert_eq!(super::parse_ignorable_psql_control_command("\\set tx_date 20240824"), None); + assert_eq!(super::parse_ignorable_psql_control_command("\\set MY_ON_ERROR_STOP on"), None); + assert_eq!(super::parse_ignorable_psql_control_command("\\set ON_ERROR_STOP_EXTRA on"), None); + assert_eq!(super::parse_ignorable_psql_control_command("\\set ON_ERROR_STOP sometimes"), None); + } + + #[test] + fn gaussdb_execution_plan_preserves_on_error_stop() { + let plan = super::sql_execution_plan_for_database( + "\\set ON_ERROR_STOP on\nSELECT 1;\nSELECT 2;", + DatabaseType::Gaussdb, + ); + + assert!(plan.stop_on_error); + assert_eq!(plan.statements, vec!["SELECT 1", "SELECT 2"]); + } + #[test] fn gaussdb_split_separates_dollar_quoted_function_from_following_statements() { let sql = "\ diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 8a3b79fba4..7511ab79f9 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -14,7 +14,7 @@ use crate::query::{ use crate::sql::{ optimize_sql_file_import_statements, prepare_sql_file_statement, split_sql_batches, statement_summary, SqlFileImportStatement, SqlFileImportStatementKind, SqlFileProgress, SqlFileRequest, SqlFileStatementAction, - SqlFileStatus, SqlParsingOptions, SqlStatementSplitter, + SqlFileStatus, SqlParsingOptions, SqlStatementSplitter, SqlStatementWithControl, }; use crate::types::QueryResult; @@ -31,6 +31,12 @@ struct StatementErrorDecision { result: Result, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct ControlledSqlFileImportStatement { + statement: SqlFileImportStatement, + stop_on_error: bool, +} + const SQL_FILE_READ_CHUNK_BYTES: usize = 256 * 1024; const SQL_FILE_STATEMENT_BATCH_SIZE: usize = 256; const SQL_FILE_PREVIEW_ENCODING_SAMPLE_BYTES: usize = 1024 * 1024; @@ -325,10 +331,12 @@ pub async fn execute_sql_file_content( mut emit: impl FnMut(SqlFileProgress), ) -> Result<(), String> { let import_target = sql_file_import_target(state, &request.connection_id).await; - let statements = - split_sql_file_import_statements(file_content, import_target.as_ref().map(|target| target.db_type)); + let statements = split_sql_file_import_statements_with_control( + file_content, + import_target.as_ref().map(|target| target.db_type), + ); - let planned_statements = optimize_sql_file_import_statements( + let planned_statements = optimize_controlled_sql_file_import_statements( &statements, import_target.as_ref().map(|target| target.db_type), import_target.as_ref().and_then(|target| target.driver_profile.as_deref()), @@ -652,17 +660,23 @@ impl StreamingSqlFileSplitter { } } - fn push_chunk(&mut self, chunk: &str) -> Vec { + fn push_chunk(&mut self, chunk: &str) -> Vec { match self { - Self::Statements(splitter) => splitter.push_chunk(chunk), - Self::SqlServerBatches(splitter) => splitter.push_chunk(chunk), + Self::Statements(splitter) => splitter.push_chunk_with_control(chunk), + Self::SqlServerBatches(splitter) => splitter + .push_chunk(chunk) + .into_iter() + .map(|sql| SqlStatementWithControl { sql, stop_on_error: false }) + .collect(), } } - fn finish(self) -> Vec { + fn finish(self) -> Vec { match self { - Self::Statements(splitter) => splitter.finish(), - Self::SqlServerBatches(splitter) => splitter.finish(), + Self::Statements(splitter) => splitter.finish_with_control(), + Self::SqlServerBatches(splitter) => { + splitter.finish().into_iter().map(|sql| SqlStatementWithControl { sql, stop_on_error: false }).collect() + } } } } @@ -719,7 +733,7 @@ async fn execute_sql_file_statement_batch( request: &SqlFileRequest, token: &CancellationToken, started_at: Instant, - statements: &mut Vec, + statements: &mut Vec, import_target: Option<&SqlFileImportTarget>, mysql_executor: Option<&mut MySqlSqlFileExecutor>, progress: &mut SqlFileExecutionProgress, @@ -729,7 +743,7 @@ async fn execute_sql_file_statement_batch( return Ok(()); } let statements = std::mem::take(statements); - let planned_statements = optimize_sql_file_import_statements( + let planned_statements = optimize_controlled_sql_file_import_statements( &statements, import_target.map(|target| target.db_type), import_target.and_then(|target| target.driver_profile.as_deref()), @@ -767,20 +781,58 @@ fn emit_sql_file_terminal_progress( )); } +#[cfg(test)] fn split_sql_file_import_statements(file_content: &str, db_type: Option) -> Vec { + split_sql_file_import_statements_with_control(file_content, db_type) + .into_iter() + .map(|statement| statement.sql) + .collect() +} + +fn split_sql_file_import_statements_with_control( + file_content: &str, + db_type: Option, +) -> Vec { if db_type == Some(DatabaseType::SqlServer) { // GO is a client-side batch delimiter, not T-SQL. SQL Server module DDL // must also remain a complete batch because procedure bodies contain semicolons. - return split_sql_batches(file_content); + return split_sql_batches(file_content) + .into_iter() + .map(|sql| SqlStatementWithControl { sql, stop_on_error: false }) + .collect(); } let options = db_type.map(SqlParsingOptions::for_database_type).unwrap_or_default(); let mut splitter = SqlStatementSplitter::with_options(options); - let mut statements = splitter.push_chunk(file_content); - statements.extend(splitter.finish()); + let mut statements = splitter.push_chunk_with_control(file_content); + statements.extend(splitter.finish_with_control()); statements } +fn optimize_controlled_sql_file_import_statements( + statements: &[SqlStatementWithControl], + db_type: Option, + driver_profile: Option<&str>, +) -> Vec { + let mut controlled = Vec::new(); + let mut start = 0; + while start < statements.len() { + let stop_on_error = statements[start].stop_on_error; + let mut end = start + 1; + while end < statements.len() && statements[end].stop_on_error == stop_on_error { + end += 1; + } + let sql = statements[start..end].iter().map(|statement| statement.sql.clone()).collect::>(); + controlled.extend( + optimize_sql_file_import_statements(&sql, db_type, driver_profile) + .into_iter() + .map(|statement| ControlledSqlFileImportStatement { statement, stop_on_error }), + ); + start = end; + } + controlled +} + #[allow(clippy::too_many_arguments)] pub fn sql_file_progress( execution_id: &str, @@ -1056,7 +1108,7 @@ async fn execute_planned_statements_with_progress( request: &SqlFileRequest, token: &CancellationToken, started_at: Instant, - planned_statements: &[SqlFileImportStatement], + planned_statements: &[ControlledSqlFileImportStatement], mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, progress: &mut SqlFileExecutionProgress, emit: &mut impl FnMut(SqlFileProgress), @@ -1066,7 +1118,7 @@ async fn execute_planned_statements_with_progress( return Ok(()); } - let next_statement_index = progress.statement_index + planned_statement.source_statement_count; + let next_statement_index = progress.statement_index + planned_statement.statement.source_statement_count; if execute_statement_with_progress( state, request, @@ -1096,13 +1148,15 @@ async fn execute_statement_with_progress( token: &CancellationToken, started_at: Instant, statement_index: usize, - statement: &SqlFileImportStatement, + controlled_statement: &ControlledSqlFileImportStatement, success_count: &mut usize, failure_count: &mut usize, affected_rows: &mut u64, mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, emit: &mut impl FnMut(SqlFileProgress), ) -> Result { + let statement = &controlled_statement.statement; + let continue_on_error = request.continue_on_error && !controlled_statement.stop_on_error; if token.is_cancelled() { let summary = statement_summary(&statement.sql); emit(sql_file_progress( @@ -1192,6 +1246,7 @@ async fn execute_statement_with_progress( started_at, statement_index + 1 - statement.source_statement_count, statement, + continue_on_error, success_count, failure_count, affected_rows, @@ -1204,7 +1259,7 @@ async fn execute_statement_with_progress( let decision = statement_error_decision( &request.execution_id, token, - request.continue_on_error, + continue_on_error, started_at, statement_index, *success_count, @@ -1231,6 +1286,7 @@ async fn execute_merged_statement_fallback_with_progress( started_at: Instant, first_statement_index: usize, statement: &SqlFileImportStatement, + continue_on_error: bool, success_count: &mut usize, failure_count: &mut usize, affected_rows: &mut u64, @@ -1296,7 +1352,7 @@ async fn execute_merged_statement_fallback_with_progress( let decision = statement_error_decision( &request.execution_id, token, - request.continue_on_error, + continue_on_error, started_at, statement_index, *success_count, @@ -1659,6 +1715,63 @@ mod tests { ); } + #[tokio::test] + async fn streaming_gaussdb_on_error_stop_overrides_continue_on_error_at_script_position() { + let dir = std::env::temp_dir().join(format!("dbx-sql-file-stop-on-error-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = crate::connection::AppState::new(storage); + let config: crate::models::connection::ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": "gauss-stream", + "name": "GaussDB stream test", + "db_type": "gaussdb", + "host": "localhost", + "port": 5432, + "username": "", + "password": "", + "database": null + })) + .unwrap(); + state.configs.write().await.insert(config.id.clone(), config); + + let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); + state + .connections + .write() + .await + .insert("gauss-stream".to_string(), crate::connection::PoolKind::Sqlite(pool.clone())); + + let path = temporary_sql_file( + b"CREATE TABLE side_effects(value INTEGER);\nINSERT INTO missing_before_control VALUES (1);\nINSERT INTO side_effects VALUES (1);\n\\set ON_ERROR_STOP on\nINSERT INTO missing_after_control VALUES (1);\nINSERT INTO side_effects VALUES (2);", + ) + .await; + let request = SqlFileRequest { + execution_id: "gauss-stop-on-error".to_string(), + connection_id: "gauss-stream".to_string(), + database: String::new(), + file_path: path.to_string_lossy().to_string(), + continue_on_error: true, + }; + let mut progress = Vec::new(); + + let result = + execute_sql_file_path(&state, &request, &path, CancellationToken::new(), Instant::now(), |event| { + progress.push(event) + }) + .await; + + assert!(result.is_err()); + let count = crate::db::sqlite::execute_query(&pool, "SELECT COUNT(*) FROM side_effects").await.unwrap().rows[0] + [0] + .as_i64(); + assert_eq!(count, Some(1)); + assert!(progress.iter().any(|event| event.status == SqlFileStatus::Error)); + assert!(!progress.iter().any(|event| event.status == SqlFileStatus::Done)); + + let _ = tokio::fs::remove_file(path).await; + let _ = tokio::fs::remove_dir_all(dir).await; + } + #[test] fn postgres_pg_dump_guards_are_removed_from_import_plan() { let statements = split_sql_file_import_statements(