feat(backup): add gzip compression and history filtering

crates/dbx-core/src/connection.rs

@@ -193,7 +193,9 @@ enum ConnectionDatabaseInfoSource {
 /// Held connection for a manual transaction session
 pub enum TxnConnection {
     Postgres(Box<deadpool_postgres::Object>),
-    Mysql(mysql_async::Conn),
+    /// A dedicated MySQL connection. Cancellation may consume and discard this
+    /// connection instead of trying to reuse it after an interrupted result set.
+    Mysql(Option<mysql_async::Conn>),
     /// Dedicated agent multi_session workload client with an open sticky TX.
     Agent {
         client: Arc<db::agent_driver::PooledAgentClient>,

crates/dbx-core/src/database_export.rs

@@ -8,6 +8,9 @@ use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use tokio_util::sync::CancellationToken;
 
+use flate2::write::GzEncoder;
+use flate2::Compression;
+
 use crate::connection::task_client_session_id;
 use crate::models::connection::DatabaseType;
 use crate::mysql_ddl_normalize::DdlNormalizeOptions;
@@ -30,6 +33,14 @@ pub fn database_export_client_session_id(export_id: &str) -> String {
     task_client_session_id("database-export", export_id)
 }
 
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub enum DatabaseExportOutputCompression {
+    #[default]
+    None,
+    Gzip,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct DatabaseExportRequest {
@@ -60,10 +71,43 @@ pub struct DatabaseExportRequest {
     #[serde(default)]
     pub fail_on_error: bool,
     #[serde(default)]
+    pub output_compression: DatabaseExportOutputCompression,
+    #[serde(default)]
     pub snapshot_session_id: Option<String>,
     pub batch_size: usize,
 }
 
+enum DatabaseExportWriter {
+    Plain(BufWriter<std::fs::File>),
+    Gzip(Box<GzEncoder<BufWriter<std::fs::File>>>),
+}
+
+impl Write for DatabaseExportWriter {
+    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
+        match self {
+            Self::Plain(writer) => writer.write(buffer),
+            Self::Gzip(writer) => writer.write(buffer),
+        }
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        match self {
+            Self::Plain(writer) => writer.flush(),
+            Self::Gzip(writer) => writer.flush(),
+        }
+    }
+}
+
+impl DatabaseExportWriter {
+    fn finish(self) -> Result<(), String> {
+        match self {
+            Self::Plain(mut writer) => writer.flush(),
+            Self::Gzip(writer) => writer.finish().and_then(|mut output| output.flush()),
+        }
+        .map_err(|error| format!("Failed to finalize export file: {error}"))
+    }
+}
+
 #[derive(Debug, Default, PartialEq, Eq)]
 struct DatabaseExportObjectCounts {
     tables: usize,
@@ -1624,10 +1668,19 @@ async fn get_export_table_ddl_isolated(
     database: String,
     schema: String,
     table: String,
+    client_session_id: String,
 ) -> Result<String, String> {
     let task = tokio::spawn(async move {
-        crate::schema::get_table_relation_export_ddl_core(&state, &connection_id, &database, &schema, &table, None)
-            .await
+        crate::schema::get_table_relation_export_ddl_core_for_session(
+            &state,
+            &connection_id,
+            &database,
+            &schema,
+            &table,
+            None,
+            Some(&client_session_id),
+        )
+        .await
     });
     let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
     task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
@@ -1639,9 +1692,18 @@ async fn get_export_table_columns_isolated(
     database: String,
     schema: String,
     table: String,
+    client_session_id: String,
 ) -> Result<Vec<crate::db::ColumnInfo>, String> {
     let task = tokio::spawn(async move {
-        crate::schema::get_columns_core(&state, &connection_id, &database, &schema, &table).await
+        crate::schema::get_columns_core_for_session(
+            &state,
+            &connection_id,
+            &database,
+            &schema,
+            &table,
+            Some(&client_session_id),
+        )
+        .await
     });
     let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
     task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
@@ -1691,6 +1753,20 @@ pub async fn begin_database_backup_snapshot_core(
     state: &crate::connection::AppState,
     connection_id: &str,
     database: &str,
+) -> Result<DatabaseBackupSnapshot, String> {
+    begin_database_backup_snapshot_core_for_export(state, connection_id, database, None).await
+}
+
+/// Opens the consistent-snapshot transaction used by a scheduled backup.
+///
+/// When the snapshot is being created for an export run, `export_id` keeps
+/// pool checkout and transaction creation cancellable even before a child
+/// database export has been created.
+pub async fn begin_database_backup_snapshot_core_for_export(
+    state: &crate::connection::AppState,
+    connection_id: &str,
+    database: &str,
+    export_id: Option<&str>,
 ) -> Result<DatabaseBackupSnapshot, String> {
     let db_type = state
         .configs
@@ -1703,22 +1779,51 @@ pub async fn begin_database_backup_snapshot_core(
         return Err("Consistent database backup snapshots are only supported for MySQL and PostgreSQL".to_string());
     }
 
-    let session_id = crate::query::begin_database_backup_snapshot(state, connection_id, database).await?;
-    let schemas = if matches!(db_type, DatabaseType::Postgres) {
+    let session_id = if let Some(export_id) = export_id {
+        // Do not use `await_export_operation` here: if the transaction is
+        // created at exactly the same time as cancellation, we still need the
+        // returned session id to roll it back rather than leaking it.
+        let operation = crate::query::begin_database_backup_snapshot(state, connection_id, database);
+        tokio::pin!(operation);
+        loop {
+            if is_export_cancelled_now(export_id) {
+                return Err(EXPORT_CANCELLED_ERROR.to_string());
+            }
+            tokio::select! {
+                biased;
+                result = &mut operation => {
+                    let session_id = result?;
+                    if is_export_cancelled_now(export_id) {
+                        let _ = crate::query::rollback_manual_transaction(state, &session_id).await;
+                        return Err(EXPORT_CANCELLED_ERROR.to_string());
+                    }
+                    break session_id;
+                },
+                _ = tokio::time::sleep(EXPORT_CANCEL_POLL_INTERVAL) => {}
+            }
+        }
+    } else {
+        crate::query::begin_database_backup_snapshot(state, connection_id, database).await?
+    };
+
+    let schemas_result = if matches!(db_type, DatabaseType::Postgres) {
         const POSTGRES_BACKUP_SCHEMAS_SQL: &str = "SELECT n.nspname FROM pg_catalog.pg_namespace n \
              WHERE n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast') \
              AND n.nspname NOT LIKE 'pg_toast_temp_%' \
              AND n.nspname NOT LIKE 'pg_temp_%' ORDER BY n.nspname";
-        let results = crate::query::execute_in_manual_transaction(
+        let operation = crate::query::execute_in_manual_transaction(
             state,
             &session_id,
             POSTGRES_BACKUP_SCHEMAS_SQL,
             database,
             None,
             Some(10_000),
-        )
-        .await?;
-        results
+        );
+        let results = match export_id {
+            Some(export_id) => await_export_operation(export_id, Box::pin(operation)).await,
+            None => operation.await,
+        }?;
+        Ok(results
             .into_iter()
             .next()
             .map(|result| {
@@ -1729,9 +1834,20 @@ pub async fn begin_database_backup_snapshot_core(
                     .filter_map(|value| value.as_str().map(str::to_string))
                     .collect()
             })
-            .unwrap_or_default()
+            .unwrap_or_default())
     } else {
-        vec![database.to_string()]
+        Ok(vec![database.to_string()])
+    };
+    let schemas = match schemas_result {
+        Ok(schemas) => schemas,
+        Err(error) => {
+            let _ = crate::query::rollback_manual_transaction(state, &session_id).await;
+            return Err(error);
+        }
+    };
+    if export_id.map(is_export_cancelled_now).unwrap_or(false) {
+        let _ = crate::query::rollback_manual_transaction(state, &session_id).await;
+        return Err(EXPORT_CANCELLED_ERROR.to_string());
     };
     if schemas.is_empty() {
         let _ = crate::query::rollback_manual_transaction(state, &session_id).await;
@@ -2152,6 +2268,16 @@ pub async fn export_database_sql_core(
     // caller future small, this prevents the metadata-prefetch locals from
     // exhausting the bounded stack used by test and runtime worker threads.
     let result = Box::pin(export_database_sql_core_inner(state, request, &on_progress)).await;
+    let metadata_session_id = database_export_client_session_id(&request.export_id);
+    if let Err(error) =
+        state.close_metadata_session_pool(&request.connection_id, Some(&request.database), &metadata_session_id).await
+    {
+        log::warn!(
+            "[database-export] failed to close metadata session '{}' for '{}': {error}",
+            metadata_session_id,
+            request.connection_id
+        );
+    }
     if result.as_ref().err().is_some_and(|error| error == EXPORT_CANCELLED_ERROR) {
         // Every caller (Tauri and web SSE) needs a terminal Cancelled event;
         // returning the marker as an error would leave the web EventSource
@@ -2231,7 +2357,12 @@ async fn export_database_sql_core_inner(
             request.file_path
         ));
     }
-    let mut file = BufWriter::new(file);
+    let mut file = match request.output_compression {
+        DatabaseExportOutputCompression::None => DatabaseExportWriter::Plain(BufWriter::new(file)),
+        DatabaseExportOutputCompression::Gzip => {
+            DatabaseExportWriter::Gzip(Box::new(GzEncoder::new(BufWriter::new(file), Compression::default())))
+        }
+    };
 
     let create_database_preamble = if request.include_create_database && matches!(db_type, DatabaseType::Mysql) {
         Some(mysql_database_export_preamble_for_request(state, request).await)
@@ -2455,6 +2586,7 @@ async fn export_database_sql_core_inner(
         let prefetch_targets: Vec<(usize, String)> =
             tables.iter().enumerate().map(|(index, table_info)| (index, table_info.name.clone())).collect();
         let mut prefetch_stream = futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| {
+            let client_session_id = client_session_id.clone();
             Box::pin(async move {
                 if is_export_cancelled_now(&request.export_id) {
                     return (index, PrefetchedTableMetadata { ddl: None, columns: None });
@@ -2469,6 +2601,7 @@ async fn export_database_sql_core_inner(
                                 request.database.clone(),
                                 request.schema.clone(),
                                 table_name.clone(),
+                                client_session_id.clone(),
                             )),
                         )
                         .await,
@@ -2489,6 +2622,7 @@ async fn export_database_sql_core_inner(
                                 request.database.clone(),
                                 request.schema.clone(),
                                 table_name.clone(),
+                                client_session_id.clone(),
                             )),
                         )
                         .await,
@@ -2641,6 +2775,7 @@ async fn export_database_sql_core_inner(
                             request.database.clone(),
                             request.schema.clone(),
                             table_name.clone(),
+                            client_session_id.clone(),
                         )),
                     )
                     .await
@@ -2688,6 +2823,7 @@ async fn export_database_sql_core_inner(
                             request.database.clone(),
                             request.schema.clone(),
                             table_name.clone(),
+                            client_session_id.clone(),
                         )),
                     )
                     .await
@@ -2726,9 +2862,8 @@ async fn export_database_sql_core_inner(
                             batch_size,
                             Some(cancel_token.clone()),
                             |rows| {
-                                // PostgreSQL cancellation must unwind through its CancelRequest path so the
-                                // connection reaches ReadyForQuery before rollback. Other snapshot streams do
-                                // not consume the token yet, so retain their existing next-batch cancellation.
+                                // PostgreSQL and MySQL consume the token while waiting for rows. Retain this
+                                // batch-boundary check so cancellation also wins while a batch is being written.
                                 if snapshot_batch_cancelled(&db_type, &request.export_id) {
                                     return Err(EXPORT_CANCELLED_ERROR.to_string());
                                 }
@@ -3062,7 +3197,7 @@ async fn export_database_sql_core_inner(
         writeln!(file, "SET FOREIGN_KEY_CHECKS = 1;").map_err(|e| format!("Failed to write file: {e}"))?;
     }
 
-    file.flush().map_err(|e| format!("Failed to finalize export file: {e}"))?;
+    file.finish()?;
 
     // Emit Done progress
     on_progress(ExportProgress {
@@ -3139,9 +3274,9 @@ mod tests {
         mysql_view_dependencies_sql, normalize_export_table_ddl, record_export_destination_identity,
         record_export_error, replace_database_export_select_list, sort_export_views_by_dependencies,
         split_postgres_export_table_triggers, write_database_export_rows, BuildDatabaseSqlExportOptions,
-        BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, DatabaseExportRequest, DdlNormalizeOptions,
-        ExportedTableSql, PostgresExportExtension, PostgresExportSequence, PostgresExtensionMembers,
-        DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT,
+        BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, DatabaseExportRequest, DatabaseExportWriter,
+        DdlNormalizeOptions, ExportedTableSql, PostgresExportExtension, PostgresExportSequence,
+        PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT,
     };
     use crate::connection::AppState;
     use crate::models::connection::DatabaseType;
@@ -3149,6 +3284,7 @@ mod tests {
     use crate::types::SpatialColumn;
     use crate::types::{ObjectInfo, ObjectSourceKind, TableInfo};
     use serde_json::{json, Value};
+    use std::io::{Read, Write};
     use std::sync::{
         atomic::{AtomicBool, Ordering},
         Arc, Mutex,
@@ -3324,11 +3460,29 @@ mod tests {
             drop_table_if_exists: false,
             omit_auto_increment: false,
             fail_on_error: false,
+            output_compression: Default::default(),
             snapshot_session_id: None,
             batch_size: 1000,
         }
     }
 
+    #[test]
+    fn gzip_export_writer_finishes_a_readable_stream() {
+        let directory = tempfile::tempdir().unwrap();
+        let path = directory.path().join("backup.sql.gz");
+        let file = std::fs::File::create(&path).unwrap();
+        let mut writer = DatabaseExportWriter::Gzip(Box::new(flate2::write::GzEncoder::new(
+            std::io::BufWriter::new(file),
+            flate2::Compression::default(),
+        )));
+        writer.write_all(b"SELECT 1;\n").unwrap();
+        writer.finish().unwrap();
+
+        let mut output = String::new();
+        flate2::read::GzDecoder::new(std::fs::File::open(path).unwrap()).read_to_string(&mut output).unwrap();
+        assert_eq!(output, "SELECT 1;\n");
+    }
+
     #[test]
     fn export_progress_total_counts_only_requested_object_categories() {
         let counts = DatabaseExportObjectCounts {

crates/dbx-core/src/query.rs

@@ -4750,7 +4750,7 @@ async fn begin_transaction_session(
             if !syntax_errors.is_empty() {
                 return Err(format!("START TRANSACTION failed for all compatible forms: {}", syntax_errors.join("; ")));
             }
-            (TxnConnection::Mysql(conn), probe_pool_key.clone())
+            (TxnConnection::Mysql(Some(conn)), probe_pool_key.clone())
         }
         TxnPoolHandle::Agent => {
             let client_session_id = format!("manual-txn-{}", uuid::Uuid::new_v4());
@@ -5072,7 +5072,14 @@ pub async fn execute_in_manual_transaction_with_options(
             TxnConnection::Postgres(conn) => {
                 execute_manual_txn_postgres_statement(conn.as_ref(), statement, row_limit).await
             }
-            TxnConnection::Mysql(conn) => execute_manual_txn_mysql_statement(conn, statement, row_limit).await,
+            TxnConnection::Mysql(conn) => {
+                execute_manual_txn_mysql_statement(
+                    conn.as_mut().ok_or_else(|| MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR.to_string())?,
+                    statement,
+                    row_limit,
+                )
+                .await
+            }
             TxnConnection::Agent { client, .. } => {
                 execute_manual_txn_agent_statement(
                     client,
@@ -5235,45 +5242,85 @@ where
                 Err(error) => Err(error),
             }
         }
-        TxnConnection::Mysql(conn) => match conn.query_iter(sql).await {
-            Ok(mut result) => match result.stream::<mysql_async::Row>().await {
-                Ok(Some(mut stream)) => {
-                    let mut batch = Vec::with_capacity(batch_size);
-                    let mut total_rows = 0_u64;
-                    let mut error = None;
-                    while let Some(row_result) = stream.next().await {
-                        match row_result {
-                            Ok(row) => {
-                                batch.push(
-                                    (0..row.len()).map(|index| db::mysql::mysql_value_to_json(&row, index)).collect(),
-                                );
-                                total_rows += 1;
-                                if batch.len() >= batch_size {
-                                    if let Err(err) = on_batch(std::mem::take(&mut batch)) {
-                                        error = Some(err);
+        TxnConnection::Mysql(Some(conn)) => {
+            let query_result = match cancel_token.as_ref() {
+                Some(cancel_token) => {
+                    tokio::select! {
+                        _ = cancel_token.cancelled() => Err(QUERY_CANCELED.to_string()),
+                        result = conn.query_iter(sql) => result.map_err(|error| format!("Query failed: {error}")),
+                    }
+                }
+                None => conn.query_iter(sql).await.map_err(|error| format!("Query failed: {error}")),
+            };
+            match query_result {
+                Ok(mut result) => {
+                    let stream_result = match cancel_token.as_ref() {
+                        Some(cancel_token) => {
+                            tokio::select! {
+                                _ = cancel_token.cancelled() => Err(QUERY_CANCELED.to_string()),
+                                stream = result.stream::<mysql_async::Row>() => stream.map_err(|error| format!("Query failed: {error}")),
+                            }
+                        }
+                        None => {
+                            result.stream::<mysql_async::Row>().await.map_err(|error| format!("Query failed: {error}"))
+                        }
+                    };
+                    match stream_result {
+                        Ok(Some(mut stream)) => {
+                            let mut batch = Vec::with_capacity(batch_size);
+                            let mut total_rows = 0_u64;
+                            let mut error = None;
+                            loop {
+                                let next_row = match cancel_token.as_ref() {
+                                    Some(cancel_token) => {
+                                        tokio::select! {
+                                            _ = cancel_token.cancelled() => {
+                                                error = Some(QUERY_CANCELED.to_string());
+                                                break;
+                                            }
+                                            row = stream.next() => row,
+                                        }
+                                    }
+                                    None => stream.next().await,
+                                };
+                                let Some(row_result) = next_row else { break };
+                                match row_result {
+                                    Ok(row) => {
+                                        batch.push(
+                                            (0..row.len())
+                                                .map(|index| db::mysql::mysql_value_to_json(&row, index))
+                                                .collect(),
+                                        );
+                                        total_rows += 1;
+                                        if batch.len() >= batch_size {
+                                            if let Err(err) = on_batch(std::mem::take(&mut batch)) {
+                                                error = Some(err);
+                                                break;
+                                            }
+                                            batch = Vec::with_capacity(batch_size);
+                                        }
+                                    }
+                                    Err(err) => {
+                                        error = Some(format!("Query failed: {err}"));
                                         break;
                                     }
-                                    batch = Vec::with_capacity(batch_size);
                                 }
                             }
-                            Err(err) => {
-                                error = Some(format!("Query failed: {err}"));
-                                break;
+                            if error.is_none() && !batch.is_empty() {
+                                if let Err(err) = on_batch(batch) {
+                                    error = Some(err);
+                                }
                             }
+                            error.map_or(Ok(total_rows), Err)
                         }
+                        Ok(None) => Err("Empty result set stream".to_string()),
+                        Err(error) => Err(error),
                     }
-                    if error.is_none() && !batch.is_empty() {
-                        if let Err(err) = on_batch(batch) {
-                            error = Some(err);
-                        }
-                    }
-                    error.map_or(Ok(total_rows), Err)
                 }
-                Ok(None) => Err("Empty result set stream".to_string()),
-                Err(err) => Err(format!("Query failed: {err}")),
-            },
-            Err(err) => Err(format!("Query failed: {err}")),
-        },
+                Err(error) => Err(error),
+            }
+        }
+        TxnConnection::Mysql(None) => Err(MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR.to_string()),
         TxnConnection::Agent { .. } => {
             Err("Streaming rows inside an agent manual transaction is not supported".to_string())
         }
@@ -5288,9 +5335,12 @@ where
             sessions.remove(txn_session_id)
         };
         if let Some(session) = removed {
-            let rollback_result =
+            let rollback_result = if err == QUERY_CANCELED {
+                discard_mysql_manual_txn_connection(&mut conn, operation_budget.cleanup_timeout).await
+            } else {
                 rollback_manual_txn_connection_with_postgres_timeout(&mut conn, Some(operation_budget.cleanup_timeout))
-                    .await;
+                    .await
+            };
             release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await;
             if let Err(rollback_error) = rollback_result {
                 return Err(format!("{err}. Transaction cleanup failed: {rollback_error}"));
@@ -5334,9 +5384,10 @@ async fn rollback_manual_txn_connection_with_postgres_timeout(
                 conn.execute_typed("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?;
             }
         }
-        TxnConnection::Mysql(conn) => {
+        TxnConnection::Mysql(Some(conn)) => {
             conn.query_drop("ROLLBACK").await.map_err(|e| format!("ROLLBACK failed: {e}"))?;
         }
+        TxnConnection::Mysql(None) => return Ok(()),
         TxnConnection::Agent { client, .. } => {
             let mut locked = client.lock().await;
             match locked.rollback_manual_transaction::<serde_json::Value>().await {
@@ -5366,6 +5417,24 @@ async fn rollback_manual_txn_connection_with_postgres_timeout(
     Ok(())
 }
 
+/// A cancelled MySQL row stream may still have unread result packets.  Do not
+/// send ROLLBACK on that connection: mysql_async only consumes a dropped stream
+/// when the next result set is requested.  Discarding the dedicated connection
+/// makes MySQL roll the transaction back and prevents protocol desynchronization.
+async fn discard_mysql_manual_txn_connection(conn: &mut TxnConnection, timeout: Duration) -> Result<(), String> {
+    let TxnConnection::Mysql(conn) = conn else {
+        return rollback_manual_txn_connection_with_postgres_timeout(conn, Some(timeout)).await;
+    };
+    let Some(conn) = conn.take() else {
+        return Ok(());
+    };
+    match tokio::time::timeout(timeout, conn.disconnect()).await {
+        Ok(Ok(())) => Ok(()),
+        Ok(Err(error)) => Err(format!("Failed to discard cancelled MySQL connection: {error}")),
+        Err(_) => Err(format!("Discarding cancelled MySQL connection timed out after {} seconds", timeout.as_secs())),
+    }
+}
+
 async fn release_manual_txn_session_pool(state: &AppState, connection_id: &str, conn: &mut TxnConnection) {
     let (client_session_id, database, cleanup_guard) = match conn {
         TxnConnection::Agent { client_session_id, database, cleanup_guard, .. }
@@ -5628,9 +5697,10 @@ pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) -
         TxnConnection::Postgres(conn) => {
             conn.execute_typed("COMMIT", &[]).await.map_err(|e| format!("COMMIT failed: {e}"))?;
         }
-        TxnConnection::Mysql(conn) => {
+        TxnConnection::Mysql(Some(conn)) => {
             conn.query_drop("COMMIT").await.map_err(|e| format!("COMMIT failed: {e}"))?;
         }
+        TxnConnection::Mysql(None) => return Err(MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR.to_string()),
         TxnConnection::Agent { client, .. } => {
             let mut locked = client.lock().await;
             locked

crates/dbx-core/src/schema.rs

@@ -7443,6 +7443,7 @@ pub async fn get_table_ddl_core(
         table,
         object_type,
         TableDdlOptions::SINGLE_RELATION,
+        None,
     )
     .await
 }
@@ -7455,17 +7456,27 @@ pub async fn get_table_export_ddl_core(
     table: &str,
     object_type: Option<db::ObjectSourceKind>,
 ) -> Result<String, String> {
-    get_table_ddl_core_with_options(state, connection_id, database, schema, table, object_type, TableDdlOptions::EXPORT)
-        .await
+    get_table_ddl_core_with_options(
+        state,
+        connection_id,
+        database,
+        schema,
+        table,
+        object_type,
+        TableDdlOptions::EXPORT,
+        None,
+    )
+    .await
 }
 
-pub(crate) async fn get_table_relation_export_ddl_core(
+pub(crate) async fn get_table_relation_export_ddl_core_for_session(
     state: &AppState,
     connection_id: &str,
     database: &str,
     schema: &str,
     table: &str,
     object_type: Option<db::ObjectSourceKind>,
+    client_session_id: Option<&str>,
 ) -> Result<String, String> {
     get_table_ddl_core_with_options(
         state,
@@ -7475,6 +7486,7 @@ pub(crate) async fn get_table_relation_export_ddl_core(
         table,
         object_type,
         TableDdlOptions::RELATION_EXPORT,
+        client_session_id,
     )
     .await
 }
@@ -7495,6 +7507,7 @@ pub async fn get_table_display_ddl_core(
         table,
         object_type,
         TableDdlOptions::DISPLAY,
+        None,
     )
     .await
 }
@@ -7507,6 +7520,7 @@ async fn get_table_ddl_core_with_options(
     table: &str,
     object_type: Option<db::ObjectSourceKind>,
     options: TableDdlOptions,
+    client_session_id: Option<&str>,
 ) -> Result<String, String> {
     if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() {
         return Err("DDL is not supported for SQL Server linked server tables".to_string());
@@ -7551,8 +7565,8 @@ async fn get_table_ddl_core_with_options(
         return Ok(source.source);
     }
 
-    retry_metadata_connection(state, connection_id, Some(database), || {
-        get_table_ddl_once(state, connection_id, database, schema, table, options)
+    retry_metadata_connection_for_session(state, connection_id, Some(database), client_session_id, || {
+        get_table_ddl_once(state, connection_id, database, schema, table, options, client_session_id)
     })
     .await
 }
@@ -7579,8 +7593,10 @@ async fn get_table_ddl_once(
     schema: &str,
     table: &str,
     options: TableDdlOptions,
+    client_session_id: Option<&str>,
 ) -> Result<String, String> {
-    let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
+    let pool_key =
+        state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
     let db_config = connection_config(state, connection_id).await;
 
     {

crates/dbx-core/src/sql_file_import.rs

@@ -1,4 +1,6 @@
+use std::io::Read as StdRead;
 use std::path::Path;
+use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
 use tokio::io::{AsyncReadExt, BufReader};
@@ -535,12 +537,68 @@ pub async fn read_sql_file_preview(file_path: &Path, max_chars: usize) -> Result
 }
 
 struct SqlFileStreamDecoder {
-    reader: BufReader<tokio::fs::File>,
+    reader: SqlFileByteReader,
     decoder: encoding_rs::Decoder,
     pending_bytes: Vec<u8>,
     reached_eof: bool,
 }
 
+enum SqlFileByteReader {
+    Plain(BufReader<tokio::fs::File>),
+    Gzip(Arc<Mutex<flate2::read::GzDecoder<std::io::BufReader<std::fs::File>>>>),
+}
+
+impl SqlFileByteReader {
+    async fn open(file_path: &Path) -> Result<Self, String> {
+        if is_gzip_sql_file_path(file_path) {
+            let path = file_path.to_path_buf();
+            let reader = tokio::task::spawn_blocking(move || {
+                let file = std::fs::File::open(&path).map_err(|error| error.to_string())?;
+                Ok::<_, String>(flate2::read::GzDecoder::new(std::io::BufReader::new(file)))
+            })
+            .await
+            .map_err(|error| format!("Failed to open compressed SQL file: {error}"))??;
+            return Ok(Self::Gzip(Arc::new(Mutex::new(reader))));
+        }
+
+        let file = tokio::fs::File::open(file_path).await.map_err(|error| error.to_string())?;
+        Ok(Self::Plain(BufReader::with_capacity(SQL_FILE_READ_CHUNK_BYTES, file)))
+    }
+
+    async fn read(&mut self, buffer: &mut [u8]) -> Result<usize, String> {
+        match self {
+            Self::Plain(reader) => reader.read(buffer).await.map_err(|error| error.to_string()),
+            Self::Gzip(reader) => {
+                let reader = reader.clone();
+                let capacity = buffer.len();
+                let (chunk, read) = tokio::task::spawn_blocking(move || {
+                    let mut chunk = vec![0; capacity];
+                    let mut reader =
+                        reader.lock().map_err(|_| "Compressed SQL reader lock was poisoned".to_string())?;
+                    let read =
+                        reader.read(&mut chunk).map_err(|error| format!("Failed to decompress SQL file: {error}"))?;
+                    Ok::<_, String>((chunk, read))
+                })
+                .await
+                .map_err(|error| format!("Failed to read compressed SQL file: {error}"))??;
+                buffer[..read].copy_from_slice(&chunk[..read]);
+                Ok(read)
+            }
+        }
+    }
+}
+
+fn is_gzip_sql_file_path(path: &Path) -> bool {
+    path.extension()
+        .and_then(|extension| extension.to_str())
+        .is_some_and(|extension| extension.eq_ignore_ascii_case("gz"))
+        && path
+            .file_stem()
+            .and_then(|stem| Path::new(stem).extension())
+            .and_then(|extension| extension.to_str())
+            .is_some_and(|extension| extension.eq_ignore_ascii_case("sql"))
+}
+
 impl SqlFileStreamDecoder {
     async fn open(file_path: &Path) -> Result<Self, String> {
         Self::open_with_detection_limit(file_path, None).await
@@ -548,18 +606,13 @@ impl SqlFileStreamDecoder {
 
     async fn open_with_detection_limit(file_path: &Path, detection_limit: Option<usize>) -> Result<Self, String> {
         let (encoding, bom_len) = detect_sql_file_encoding(file_path, detection_limit).await?;
-        let mut file = tokio::fs::File::open(file_path).await.map_err(|error| error.to_string())?;
+        let mut reader = SqlFileByteReader::open(file_path).await?;
         let mut prefix = [0u8; 3];
-        let prefix_len = file.read(&mut prefix).await.map_err(|error| error.to_string())?;
+        let prefix_len = reader.read(&mut prefix).await?;
         let prefix = &prefix[..prefix_len];
         let mut pending_bytes = prefix[bom_len..].to_vec();
         pending_bytes.reserve(SQL_FILE_READ_CHUNK_BYTES);
-        Ok(Self {
-            reader: BufReader::with_capacity(SQL_FILE_READ_CHUNK_BYTES, file),
-            decoder: encoding.new_decoder_without_bom_handling(),
-            pending_bytes,
-            reached_eof: false,
-        })
+        Ok(Self { reader, decoder: encoding.new_decoder_without_bom_handling(), pending_bytes, reached_eof: false })
     }
 
     async fn next_chunk(&mut self) -> Result<Option<String>, String> {
@@ -568,7 +621,7 @@ impl SqlFileStreamDecoder {
         }
         while !self.reached_eof && self.pending_bytes.len() < SQL_FILE_READ_CHUNK_BYTES {
             let mut buffer = vec![0u8; SQL_FILE_READ_CHUNK_BYTES - self.pending_bytes.len()];
-            let read = self.reader.read(&mut buffer).await.map_err(|error| error.to_string())?;
+            let read = self.reader.read(&mut buffer).await?;
             if read == 0 {
                 self.reached_eof = true;
                 break;
@@ -596,7 +649,7 @@ async fn detect_sql_file_encoding(
     file_path: &Path,
     detection_limit: Option<usize>,
 ) -> Result<(&'static encoding_rs::Encoding, usize), String> {
-    let mut file = tokio::fs::File::open(file_path).await.map_err(|error| error.to_string())?;
+    let mut file = SqlFileByteReader::open(file_path).await?;
     let mut prefix = [0u8; 3];
     let prefix_len = file.read(&mut prefix).await.map_err(|error| error.to_string())?;
     let prefix = &prefix[..prefix_len];
@@ -623,7 +676,7 @@ async fn detect_sql_file_encoding(
             let remaining =
                 detection_limit.map(|limit| limit.saturating_sub(inspected_bytes)).unwrap_or(SQL_FILE_READ_CHUNK_BYTES);
             let mut buffer = vec![0u8; SQL_FILE_READ_CHUNK_BYTES.min(remaining.max(1))];
-            let read = file.read(&mut buffer).await.map_err(|error| error.to_string())?;
+            let read = file.read(&mut buffer).await?;
             if read == 0 {
                 reached_eof = true;
             } else {
@@ -1512,6 +1565,27 @@ mod tests {
         path
     }
 
+    async fn temporary_gzip_sql_file(bytes: &[u8]) -> std::path::PathBuf {
+        let path = std::env::temp_dir().join(format!(
+            "dbx-sql-file-{}-{}.sql.gz",
+            std::process::id(),
+            TEMP_SQL_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed)
+        ));
+        let bytes = bytes.to_vec();
+        let output = path.clone();
+        tokio::task::spawn_blocking(move || {
+            use std::io::Write;
+
+            let file = std::fs::File::create(&output).unwrap();
+            let mut writer = flate2::write::GzEncoder::new(file, flate2::Compression::default());
+            writer.write_all(&bytes).unwrap();
+            writer.finish().unwrap();
+        })
+        .await
+        .unwrap();
+        path
+    }
+
     fn test_progress(status: SqlFileStatus, statement_index: usize) -> SqlFileProgress {
         SqlFileProgress {
             execution_id: "test-execution".to_string(),
@@ -1707,6 +1781,19 @@ mod tests {
         assert_eq!(decoded, "SELECT '中文';");
     }
 
+    #[tokio::test]
+    async fn streaming_decoder_reads_gzip_sql_files() {
+        let path = temporary_gzip_sql_file(b"SELECT 'compressed';\n").await;
+        let mut decoder = SqlFileStreamDecoder::open(&path).await.unwrap();
+        let mut decoded = String::new();
+        while let Some(chunk) = decoder.next_chunk().await.unwrap() {
+            decoded.push_str(&chunk);
+        }
+        tokio::fs::remove_file(path).await.unwrap();
+
+        assert_eq!(decoded, "SELECT 'compressed';\n");
+    }
+
     #[test]
     fn non_sqlserver_sql_file_keeps_statement_splitting_behavior() {
         assert_eq!(