feat(gaussdb): 支持带 psql 控制命令的匿名块执行
This commit is contained in:
3 files changed
+444
-34
No files matched your search
+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 = "\
|
||||
|
||||
Reference in new issue
Block a user