fix(sqlserver): support GO in SQL file execution
This commit is contained in:
t8y2
parent
cb5ed9baf8
commit
b5f462f8e0
|
|
@ -974,6 +974,7 @@ fn split_sql_batch_ranges(sql: &str, profile: SqlDialectProfile) -> Vec<SqlState
|
|||
let mut current_start = 0;
|
||||
let lines: Vec<&str> = sql.split('\n').collect();
|
||||
let mut offset = 0;
|
||||
let mut scanner = SqlScanner::with_profile(profile);
|
||||
|
||||
for line in &lines {
|
||||
let line_start = offset;
|
||||
|
|
@ -981,15 +982,23 @@ fn split_sql_batch_ranges(sql: &str, profile: SqlDialectProfile) -> Vec<SqlState
|
|||
offset = line_end + 1; // +1 for the '\n'
|
||||
|
||||
let trimmed = line.trim();
|
||||
if profile.supports_go_batch_separator
|
||||
let is_batch_separator = profile.supports_go_batch_separator
|
||||
&& !scanner.is_masked()
|
||||
&& (trimmed.eq_ignore_ascii_case("go")
|
||||
|| trimmed.to_ascii_lowercase().starts_with("go ") && trimmed[2..].trim().is_empty())
|
||||
{
|
||||
|| trimmed.to_ascii_lowercase().starts_with("go ") && trimmed[2..].trim().is_empty());
|
||||
if is_batch_separator {
|
||||
push_batch_range(&mut batches, sql, current_start, line_start);
|
||||
current_start = line_end.min(sql.len());
|
||||
if current_start < sql.len() && sql.as_bytes()[current_start] == b'\n' {
|
||||
current_start += 1;
|
||||
}
|
||||
} else {
|
||||
for (relative_idx, ch) in line.char_indices() {
|
||||
scanner.step(sql, line_start + relative_idx, ch);
|
||||
}
|
||||
}
|
||||
if line_end < sql.len() {
|
||||
scanner.step(sql, line_end, '\n');
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2741,6 +2750,22 @@ mod tests {
|
|||
assert_eq!(super::split_sql_batches("SELECT 1\nGO"), vec!["SELECT 1"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_batches_keeps_go_inside_multiline_string() {
|
||||
assert_eq!(
|
||||
super::split_sql_batches("SELECT 'first line\nGO\nlast line';\nGO\nSELECT 2"),
|
||||
vec!["SELECT 'first line\nGO\nlast line';", "SELECT 2"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_batches_keeps_go_inside_block_comment() {
|
||||
assert_eq!(
|
||||
super::split_sql_batches("SELECT 1;\n/*\nGO\n*/\nSELECT 2;\nGO\nSELECT 3;"),
|
||||
vec!["SELECT 1;\n/*\nGO\n*/\nSELECT 2;", "SELECT 3;"]
|
||||
);
|
||||
}
|
||||
|
||||
// --- DELIMITER support ---
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -10,9 +10,9 @@ use crate::query::{
|
|||
QueryExecutionOptions,
|
||||
};
|
||||
use crate::sql::{
|
||||
optimize_sql_file_import_statements, prepare_sql_file_statement, statement_summary, SqlFileImportStatement,
|
||||
SqlFileImportStatementKind, SqlFileProgress, SqlFileRequest, SqlFileStatementAction, SqlFileStatus,
|
||||
SqlParsingOptions, SqlStatementSplitter,
|
||||
optimize_sql_file_import_statements, prepare_sql_file_statement, split_sql_batches, statement_summary,
|
||||
SqlFileImportStatement, SqlFileImportStatementKind, SqlFileProgress, SqlFileRequest, SqlFileStatementAction,
|
||||
SqlFileStatus, SqlParsingOptions, SqlStatementSplitter,
|
||||
};
|
||||
use crate::types::QueryResult;
|
||||
|
||||
|
|
@ -232,11 +232,8 @@ 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 options =
|
||||
import_target.as_ref().map(|target| SqlParsingOptions::for_database_type(target.db_type)).unwrap_or_default();
|
||||
let mut splitter = SqlStatementSplitter::with_options(options);
|
||||
let mut statements = splitter.push_chunk(file_content);
|
||||
statements.extend(splitter.finish());
|
||||
let statements =
|
||||
split_sql_file_import_statements(file_content, import_target.as_ref().map(|target| target.db_type));
|
||||
|
||||
let planned_statements = optimize_sql_file_import_statements(
|
||||
&statements,
|
||||
|
|
@ -258,6 +255,20 @@ pub async fn execute_sql_file_content(
|
|||
.await
|
||||
}
|
||||
|
||||
fn split_sql_file_import_statements(file_content: &str, db_type: Option<DatabaseType>) -> Vec<String> {
|
||||
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);
|
||||
}
|
||||
|
||||
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());
|
||||
statements
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn sql_file_progress(
|
||||
execution_id: &str,
|
||||
|
|
@ -900,6 +911,42 @@ mod tests {
|
|||
use super::*;
|
||||
use crate::models::connection::DatabaseType;
|
||||
|
||||
#[test]
|
||||
fn sqlserver_sql_file_splits_go_batches_without_sending_delimiters() {
|
||||
let statements = split_sql_file_import_statements(
|
||||
"CREATE TABLE dbo.items (id INT);\nGO\nINSERT INTO dbo.items VALUES (1);\nGO\nSELECT * FROM dbo.items;",
|
||||
Some(DatabaseType::SqlServer),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
statements,
|
||||
vec!["CREATE TABLE dbo.items (id INT);", "INSERT INTO dbo.items VALUES (1);", "SELECT * FROM dbo.items;"]
|
||||
);
|
||||
assert!(statements
|
||||
.iter()
|
||||
.all(|statement| !statement.lines().any(|line| line.trim().eq_ignore_ascii_case("go"))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sqlserver_sql_file_keeps_module_body_in_one_batch() {
|
||||
let statements = split_sql_file_import_statements(
|
||||
"CREATE PROCEDURE dbo.demo AS\nBEGIN\n SELECT 1;\n SELECT 2;\nEND\nGO\nSELECT 3;",
|
||||
Some(DatabaseType::SqlServer),
|
||||
);
|
||||
|
||||
assert_eq!(statements.len(), 2);
|
||||
assert_eq!(statements[0], "CREATE PROCEDURE dbo.demo AS\nBEGIN\n SELECT 1;\n SELECT 2;\nEND");
|
||||
assert_eq!(statements[1], "SELECT 3;");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_sqlserver_sql_file_keeps_statement_splitting_behavior() {
|
||||
assert_eq!(
|
||||
split_sql_file_import_statements("SELECT 1; SELECT 2;", Some(DatabaseType::Postgres)),
|
||||
vec!["SELECT 1", "SELECT 2"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_on_error_returns_err_with_terminal_error_progress() {
|
||||
let decision = statement_error_decision(
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
use dbx_core::connection::{AppState, PoolKind};
|
||||
use dbx_core::models::connection::DatabaseType;
|
||||
use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest};
|
||||
use dbx_core::sql::{SqlFileRequest, SqlFileStatus};
|
||||
use dbx_core::sql_file_import::execute_sql_file_content;
|
||||
use dbx_core::storage::Storage;
|
||||
use dbx_core::table_structure_sql::{
|
||||
build_table_structure_change_sql, ColumnInfo, EditableStructureColumn, TableStructureSqlOptions,
|
||||
};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn live_sqlserver_config(id: &str, database: &str) -> dbx_core::models::connection::ConnectionConfig {
|
||||
dbx_core::models::connection::ConnectionConfig {
|
||||
|
|
@ -341,6 +344,84 @@ async fn live_sqlserver_query_result_export_streams_cte_query_to_csv() {
|
|||
assert!(!csv.contains("\n\n"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at a writable SQL Server database"]
|
||||
async fn live_sqlserver_sql_file_import_executes_go_batches() {
|
||||
let database = std::env::var("DBX_LIVE_SQLSERVER_DATABASE").unwrap_or_else(|_| "tempdb".to_string());
|
||||
let host = std::env::var("DBX_LIVE_SQLSERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
|
||||
let port = std::env::var("DBX_LIVE_SQLSERVER_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(1433);
|
||||
let user = std::env::var("DBX_LIVE_SQLSERVER_USER").unwrap_or_else(|_| "sa".to_string());
|
||||
let password = std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD");
|
||||
let client =
|
||||
dbx_core::db::sqlserver::connect(&host, port, &user, &password, Some(&database), None, Duration::from_secs(10))
|
||||
.await
|
||||
.expect("connect SQL Server");
|
||||
|
||||
let suffix = uuid::Uuid::new_v4().simple().to_string();
|
||||
let table = format!("dbx_sql_file_{suffix}");
|
||||
let procedure = format!("dbx_sql_file_proc_{suffix}");
|
||||
let dir = std::env::temp_dir().join(format!("dbx-live-sqlserver-file-{suffix}"));
|
||||
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 = "live-sqlserver-file";
|
||||
let mut config = live_sqlserver_config(connection_id, &database);
|
||||
config.host = host;
|
||||
config.port = port;
|
||||
config.username = user;
|
||||
config.password = password;
|
||||
state.configs.write().await.insert(connection_id.to_string(), config);
|
||||
state.connections.write().await.insert(
|
||||
format!("{connection_id}:{database}"),
|
||||
PoolKind::SqlServer(std::sync::Arc::new(tokio::sync::Mutex::new(client))),
|
||||
);
|
||||
|
||||
let script = format!(
|
||||
"CREATE TABLE [dbo].[{table}] (id INT NOT NULL);\n\
|
||||
GO\n\
|
||||
INSERT INTO [dbo].[{table}] (id) VALUES (1);\n\
|
||||
GO\n\
|
||||
CREATE PROCEDURE [dbo].[{procedure}] AS\n\
|
||||
BEGIN\n\
|
||||
SELECT COUNT(*) AS item_count FROM [dbo].[{table}];\n\
|
||||
END\n\
|
||||
GO"
|
||||
);
|
||||
let request = SqlFileRequest {
|
||||
execution_id: format!("live-sqlserver-file-{suffix}"),
|
||||
connection_id: connection_id.to_string(),
|
||||
database: database.clone(),
|
||||
file_path: "fixture.sql".to_string(),
|
||||
continue_on_error: false,
|
||||
};
|
||||
let done_seen = AtomicBool::new(false);
|
||||
|
||||
execute_sql_file_content(&state, &request, &script, CancellationToken::new(), Instant::now(), |progress| {
|
||||
if progress.status == SqlFileStatus::Done {
|
||||
done_seen.store(true, Ordering::Relaxed);
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("execute SQL Server file with GO batches");
|
||||
|
||||
let pool_key = format!("{connection_id}:{database}");
|
||||
let connections = state.connections.read().await;
|
||||
let PoolKind::SqlServer(client) = connections.get(&pool_key).expect("SQL Server pool") else {
|
||||
panic!("expected SQL Server pool");
|
||||
};
|
||||
let mut client = client.lock().await;
|
||||
let rows = dbx_core::db::sqlserver::execute_query(&mut client, &format!("EXEC [dbo].[{procedure}]")).await;
|
||||
let cleanup = format!("DROP PROCEDURE [dbo].[{procedure}]; DROP TABLE [dbo].[{table}];");
|
||||
let _ = dbx_core::db::sqlserver::execute_batch(&mut client, &cleanup).await;
|
||||
drop(client);
|
||||
drop(connections);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
|
||||
assert!(done_seen.load(Ordering::Relaxed));
|
||||
let rows = rows.expect("execute imported procedure");
|
||||
assert_eq!(rows.rows.first().and_then(|row| row.first()), Some(&serde_json::json!(1)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at a writable SQL Server database"]
|
||||
async fn live_sqlserver_transfer_table_skips_rowversion_insert_column() {
|
||||
|
|
|
|||
Loading…
Reference in New Issue