feat(gaussdb): 支持带 psql 控制命令的匿名块执行
This commit is contained in:
3 files changed
+444
-34
No files matched your search
@@ -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()];
|
||||
|
||||
+264
-10
@@ -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<char>,
|
||||
pending_mysql_line_comment_dashes: bool,
|
||||
custom_delimiter: Option<String>,
|
||||
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<String> {
|
||||
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<SqlStatementWithControl> {
|
||||
let mut statements = Vec::new();
|
||||
let chars = chunk.chars().collect::<Vec<_>>();
|
||||
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<String> {
|
||||
pub fn finish(self) -> Vec<String> {
|
||||
self.finish_with_control().into_iter().map(|statement| statement.sql).collect()
|
||||
}
|
||||
|
||||
pub(crate) fn finish_with_control(mut self) -> Vec<SqlStatementWithControl> {
|
||||
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<String>) {
|
||||
pub fn stop_on_error(&self) -> bool {
|
||||
self.stop_on_error
|
||||
}
|
||||
|
||||
fn push_current_statement(&mut self, statements: &mut Vec<SqlStatementWithControl>) {
|
||||
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<String> {
|
||||
}
|
||||
|
||||
pub fn split_sql_statements_for_database(sql: &str, db_type: DatabaseType) -> Vec<String> {
|
||||
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<String>,
|
||||
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<String> {
|
||||
@@ -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<PsqlControlCommand> {
|
||||
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<_>>(), 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 = "\
|
||||
|
||||
@@ -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<bool, String>,
|
||||
}
|
||||
|
||||
#[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<String> {
|
||||
fn push_chunk(&mut self, chunk: &str) -> Vec<SqlStatementWithControl> {
|
||||
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<String> {
|
||||
fn finish(self) -> Vec<SqlStatementWithControl> {
|
||||
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<String>,
|
||||
statements: &mut Vec<SqlStatementWithControl>,
|
||||
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<DatabaseType>) -> Vec<String> {
|
||||
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<DatabaseType>,
|
||||
) -> Vec<SqlStatementWithControl> {
|
||||
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<DatabaseType>,
|
||||
driver_profile: Option<&str>,
|
||||
) -> Vec<ControlledSqlFileImportStatement> {
|
||||
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::<Vec<_>>();
|
||||
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<bool, String> {
|
||||
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(
|
||||
|
||||
Reference in new issue
Block a user