From 9cae3656defdd77766963b538776c2fd068fae69 Mon Sep 17 00:00:00 2001 From: mag1cfrog Date: Mon, 3 Aug 2026 23:00:58 -0700 Subject: [PATCH] feat: add bulk-load table lock control --- src/connection.rs | 44 ++++++++++++++++++++++++++ tests/compatibility_sqlserver.rs | 54 +++++++++++++++++++++++++++++++- 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/src/connection.rs b/src/connection.rs index 203cbc7..b77d645 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -136,6 +136,26 @@ impl ConnectedMssqlClient { }) } + /// Enables or disables SQL Server's persistent `table lock on bulk load` + /// option for a table. + /// + /// Enabling requires `ALTER` permission. Callers may ignore a known + /// nonfatal enable failure and continue without this optimization. Other + /// failures should be propagated. + /// + /// If enabling succeeds for a temporary load table, disabling must succeed + /// before the table is published. A disable failure should prevent + /// publication and trigger cleanup of the temporary table. + pub async fn set_bulk_load_table_lock( + &mut self, + table: &TableName, + enabled: bool, + ) -> Result<()> { + self.execute_statement(&bulk_load_table_lock_sql(table, enabled)) + .await?; + Ok(()) + } + /// Starts a bulk writer on this same SQL Server connection. /// /// The returned writer borrows the connected client, so lifecycle SQL and @@ -214,6 +234,15 @@ fn target_row_count_query(table: &TableName) -> String { ) } +fn bulk_load_table_lock_sql(table: &TableName, enabled: bool) -> String { + let value = if enabled { "ON" } else { "OFF" }; + + format!( + "EXEC sys.sp_tableoption {}, 'table lock on bulk load', '{value}';", + sql_string_literal(&table.quoted_sql()) + ) +} + fn count_big_i64_to_u64(count: i64) -> Result { u64::try_from(count).map_err(|_| Error::TargetRowCountUnexpectedResult { reason: "target row count was outside the supported range".to_owned(), @@ -282,6 +311,21 @@ mod tests { Ok(()) } + #[test] + fn bulk_load_table_lock_sql_uses_quoted_table_name_and_requested_state() -> crate::Result<()> { + let table = crate::TableName::new("tenant's", "people's")?; + + assert_eq!( + super::bulk_load_table_lock_sql(&table, true), + "EXEC sys.sp_tableoption N'[tenant''s].[people''s]', 'table lock on bulk load', 'ON';" + ); + assert_eq!( + super::bulk_load_table_lock_sql(&table, false), + "EXEC sys.sp_tableoption N'[tenant''s].[people''s]', 'table lock on bulk load', 'OFF';" + ); + Ok(()) + } + #[test] fn count_big_conversion_rejects_negative_values_without_panicking() { let error = super::count_big_i64_to_u64(-1).err().unwrap_or_else(|| { diff --git a/tests/compatibility_sqlserver.rs b/tests/compatibility_sqlserver.rs index 6a37f20..8ea2cbc 100644 --- a/tests/compatibility_sqlserver.rs +++ b/tests/compatibility_sqlserver.rs @@ -11,7 +11,8 @@ use arrow_array::{ArrayRef, Int32Array, RecordBatch, TimestampMicrosecondArray}; use arrow_schema::{DataType, Field, Schema, TimeUnit}; use arrow_sql_server::{ BulkWriter, CompatibilityLevel, MssqlProfile, MssqlVersion, PlanOptions, TableName, - TimestampPolicy, WriteBackend, WriteOptions, create_table_sql_from_mappings, + TimestampPolicy, WriteBackend, WriteOptions, connect_mssql_client_from_ado_string, + create_table_sql_from_mappings, }; use tokio::net::TcpStream; use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt}; @@ -25,6 +26,48 @@ static TABLE_COUNTER: AtomicU64 = AtomicU64::new(0); type TestClient = tiberius::Client>; type TestResult = Result>; +#[tokio::test] +async fn bulk_load_table_lock_can_be_enabled_and_disabled() -> TestResult<()> { + let Some((connection_string, database)) = integration_config() else { + eprintln!( + "skipping SQL Server bulk-load table-lock compatibility probe: {CONNECTION_STRING_ENV} or {TEST_DATABASE_ENV} is not set" + ); + return Ok(()); + }; + + let connection_string = format!("{connection_string};database={database}"); + let mut client = connect_mssql_client_from_ado_string(&connection_string).await?; + let table = unique_table_name()?; + client + .execute_statement(&format!( + "CREATE TABLE {} ([value] int NOT NULL)", + table.quoted_sql() + )) + .await?; + + let result = async { + client.set_bulk_load_table_lock(&table, true).await?; + client + .execute_statement(&bulk_load_table_lock_assertion_sql(&table, true)) + .await?; + client.set_bulk_load_table_lock(&table, false).await?; + client + .execute_statement(&bulk_load_table_lock_assertion_sql(&table, false)) + .await?; + + Ok::<(), Box>(()) + } + .await; + + let drop_result = client + .execute_statement(&format!("DROP TABLE IF EXISTS {}", table.quoted_sql())) + .await; + result?; + drop_result?; + + Ok(()) +} + #[tokio::test] async fn datetime_rounding_matches_sql_server_casts() -> TestResult<()> { let Some((connection_string, database)) = integration_config() else { @@ -297,6 +340,15 @@ fn integration_config() -> Option<(String, String)> { Some((connection_string, database)) } +fn bulk_load_table_lock_assertion_sql(table: &TableName, expected: bool) -> String { + let expected = u8::from(expected); + let table = table.quoted_sql().replace('\'', "''"); + + format!( + "IF NOT EXISTS (SELECT 1 FROM sys.tables WHERE [object_id] = OBJECT_ID(N'{table}') AND [lock_on_bulk_load] = {expected}) RAISERROR('unexpected table lock on bulk load state', 16, 1);" + ) +} + fn ensure_eq(actual: T, expected: T, context: &str) -> TestResult<()> where T: Debug + PartialEq,