Enable optional table locking for SQL Server bulk loads

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> {
     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(|| {

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<Compat<TcpStream>>;
 type TestResult<T> = Result<T, Box<dyn std::error::Error>>;
 
+#[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<dyn std::error::Error>>(())
+    }
+    .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<T>(actual: T, expected: T, context: &str) -> TestResult<()>
 where
     T: Debug + PartialEq,