Merge pull request #702 from DengQingNian/optimize/export-pagination-orderby-batchsize

perf(ui): 优化数据导出的ui体验和性能

crates/dbx-core/src/table_export.rs

@@ -6,10 +6,10 @@ use crate::connection::AppState;
 use crate::csv_export::{escape_csv, format_csv, value_to_csv_text};
 use crate::database_export::is_export_cancelled;
 pub use crate::database_export::ExportStatus;
-use crate::transfer::{count_sql, execute_on_pool, pagination_sql};
+use crate::transfer::{count_sql, execute_on_pool, keyset_pagination_sql, pagination_sql};
 use crate::xlsx_export::{build_xlsx_workbook, XlsxWorksheetData};
 
-const DEFAULT_BATCH_SIZE: usize = 500;
+const DEFAULT_BATCH_SIZE: usize = 2000;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -24,6 +24,8 @@ pub struct TableExportRequest {
     pub format: String,
     #[serde(default)]
     pub columns: Option<Vec<String>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub batch_size: Option<usize>,
 }
 
 #[derive(Debug, Clone, Serialize)]
@@ -76,7 +78,10 @@ pub async fn export_table_data_core(
 
     let col_names: Vec<String> = columns.iter().map(|c| c.name.clone()).collect();
 
-    // 4. Optionally filter to requested columns
+    // 4. Extract primary keys from column metadata (before column filtering)
+    let primary_keys: Vec<String> = columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.clone()).collect();
+
+    // 5. Optionally filter to requested columns
     let col_names = if let Some(requested_cols) = &request.columns {
         if requested_cols.is_empty() {
             col_names
@@ -95,7 +100,19 @@ pub async fn export_table_data_core(
         return Err("No columns found for table".to_string());
     }
 
-    // 5. Get total row count for progress estimation
+    // Use keyset pagination when all PKs are in the selected (filtered) columns.
+    // This avoids the OFFSET performance penalty for large tables.
+    // When no PK is available, falls back to offset-based pagination.
+    let use_keyset = !primary_keys.is_empty() && primary_keys.iter().all(|pk| col_names.contains(pk));
+
+    // PK column indices within result rows (for extracting last-row values)
+    let pk_indices: Vec<usize> = if use_keyset {
+        primary_keys.iter().map(|pk| col_names.iter().position(|c| c == pk).unwrap()).collect()
+    } else {
+        Vec::new()
+    };
+
+    // 6. Get total row count for progress estimation
     let count_query = count_sql(&request.table_name, request.schema.as_deref().unwrap_or(""), &db_type);
     let total_rows = match execute_on_pool(state, &pool_key, &count_query).await {
         Ok(result) => result.rows.first().and_then(|r| r.first()).and_then(|v| match v {
@@ -106,7 +123,7 @@ pub async fn export_table_data_core(
         Err(_) => None,
     };
 
-    // 6. Emit initial Running progress
+    // 7. Emit initial Running progress
     on_progress(TableExportProgress {
         export_id: request.export_id.clone(),
         table_name: request.table_name.clone(),
@@ -116,12 +133,15 @@ pub async fn export_table_data_core(
         error_message: None,
     });
 
-    // 7. Create output file
+    // 8. Create output file
     let mut file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?;
 
     let mut rows_exported: u64 = 0;
+    let batch_size = request.batch_size.unwrap_or(DEFAULT_BATCH_SIZE).max(1);
     let mut offset: u64 = 0;
-    let batch_size = DEFAULT_BATCH_SIZE;
+
+    // Track last primary key values for keyset pagination
+    let mut last_pk_values: Vec<Value> = Vec::new();
 
     match request.format.to_lowercase().as_str() {
         "csv" => {
@@ -144,14 +164,26 @@ pub async fn export_table_data_core(
                     return Ok(());
                 }
 
-                let sql = pagination_sql(
-                    &col_names,
-                    &request.table_name,
-                    request.schema.as_deref().unwrap_or(""),
-                    &db_type,
-                    offset,
-                    batch_size,
-                );
+                let sql = if use_keyset {
+                    keyset_pagination_sql(
+                        &col_names,
+                        &request.table_name,
+                        request.schema.as_deref().unwrap_or(""),
+                        &db_type,
+                        &primary_keys,
+                        &last_pk_values,
+                        batch_size,
+                    )
+                } else {
+                    pagination_sql(
+                        &col_names,
+                        &request.table_name,
+                        request.schema.as_deref().unwrap_or(""),
+                        &db_type,
+                        offset,
+                        batch_size,
+                    )
+                };
 
                 let result = execute_on_pool(state, &pool_key, &sql).await?;
                 let row_count = result.rows.len();
@@ -173,7 +205,15 @@ pub async fn export_table_data_core(
                 }
 
                 rows_exported += row_count as u64;
-                offset += row_count as u64;
+
+                if use_keyset {
+                    // Keyset pagination: track last PK values for next batch
+                    if let Some(last_row) = result.rows.last() {
+                        last_pk_values = pk_indices.iter().map(|&i| last_row[i].clone()).collect();
+                    }
+                } else {
+                    offset += row_count as u64;
+                }
 
                 on_progress(TableExportProgress {
                     export_id: request.export_id.clone(),
@@ -206,14 +246,26 @@ pub async fn export_table_data_core(
                     return Ok(());
                 }
 
-                let sql = pagination_sql(
-                    &col_names,
-                    &request.table_name,
-                    request.schema.as_deref().unwrap_or(""),
-                    &db_type,
-                    offset,
-                    batch_size,
-                );
+                let sql = if use_keyset {
+                    keyset_pagination_sql(
+                        &col_names,
+                        &request.table_name,
+                        request.schema.as_deref().unwrap_or(""),
+                        &db_type,
+                        &primary_keys,
+                        &last_pk_values,
+                        batch_size,
+                    )
+                } else {
+                    pagination_sql(
+                        &col_names,
+                        &request.table_name,
+                        request.schema.as_deref().unwrap_or(""),
+                        &db_type,
+                        offset,
+                        batch_size,
+                    )
+                };
 
                 let result = execute_on_pool(state, &pool_key, &sql).await?;
                 let row_count = result.rows.len();
@@ -223,7 +275,15 @@ pub async fn export_table_data_core(
 
                 all_rows.extend(result.rows);
                 rows_exported += row_count as u64;
-                offset += row_count as u64;
+
+                if use_keyset {
+                    // Keyset pagination: track last PK values for next batch
+                    if let Some(last_row) = all_rows.last() {
+                        last_pk_values = pk_indices.iter().map(|&i| last_row[i].clone()).collect();
+                    }
+                } else {
+                    offset += row_count as u64;
+                }
 
                 on_progress(TableExportProgress {
                     export_id: request.export_id.clone(),

crates/dbx-core/src/transfer.rs

@@ -1253,6 +1253,87 @@ pub fn count_sql(table: &str, schema: &str, db_type: &DatabaseType) -> String {
     format!("SELECT COUNT(*) FROM {full_table}")
 }
 
+pub fn keyset_pagination_sql(
+    columns: &[String],
+    table: &str,
+    schema: &str,
+    db_type: &DatabaseType,
+    primary_keys: &[String],
+    last_pk_values: &[serde_json::Value],
+    limit: usize,
+) -> String {
+    let full_table = qualified_table(table, schema, db_type);
+    let col_list = columns.iter().map(|c| quote_identifier(c, db_type)).collect::<Vec<_>>().join(", ");
+    let order =
+        primary_keys.iter().map(|pk| format!("{} ASC", quote_identifier(pk, db_type))).collect::<Vec<_>>().join(", ");
+
+    let where_clause = keyset_where_clause(primary_keys, last_pk_values, db_type);
+
+    match db_type {
+        DatabaseType::SqlServer | DatabaseType::Oracle => {
+            format!(
+                "SELECT {col_list} FROM {full_table}{where_clause} ORDER BY {order} OFFSET 0 ROWS FETCH NEXT {limit} ROWS ONLY"
+            )
+        }
+        _ => {
+            format!("SELECT {col_list} FROM {full_table}{where_clause} ORDER BY {order} LIMIT {limit}")
+        }
+    }
+}
+
+fn keyset_where_clause(
+    primary_keys: &[String],
+    last_pk_values: &[serde_json::Value],
+    db_type: &DatabaseType,
+) -> String {
+    if primary_keys.is_empty() || last_pk_values.is_empty() {
+        return String::new();
+    }
+
+    let quoted_keys = primary_keys.iter().map(|pk| quote_identifier(pk, db_type)).collect::<Vec<_>>();
+    let literals = last_pk_values.iter().map(|v| value_to_sql_literal(v, db_type)).collect::<Vec<_>>();
+    let comparison_count = quoted_keys.len().min(literals.len());
+    if comparison_count == 0 {
+        return String::new();
+    }
+
+    let mut clauses = Vec::with_capacity(comparison_count);
+    for index in 0..comparison_count {
+        let mut parts = Vec::with_capacity(index + 1);
+        for prefix_index in 0..index {
+            parts.push(format!("{} = {}", quoted_keys[prefix_index], literals[prefix_index]));
+        }
+        parts.push(format!("{} > {}", quoted_keys[index], literals[index]));
+        if parts.len() == 1 {
+            clauses.push(parts.remove(0));
+        } else {
+            clauses.push(format!("({})", parts.join(" AND ")));
+        }
+    }
+
+    if clauses.len() == 1 {
+        format!(" WHERE {}", clauses[0])
+    } else {
+        format!(" WHERE ({})", clauses.join(" OR "))
+    }
+}
+
+fn value_to_sql_literal(value: &serde_json::Value, _db_type: &DatabaseType) -> String {
+    match value {
+        serde_json::Value::Null => "NULL".to_string(),
+        serde_json::Value::Bool(b) => {
+            if *b {
+                "TRUE".to_string()
+            } else {
+                "FALSE".to_string()
+            }
+        }
+        serde_json::Value::Number(n) => n.to_string(),
+        serde_json::Value::String(s) => quote_string_literal(s),
+        _ => quote_string_literal(&value.to_string()),
+    }
+}
+
 pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
     let connections = state.connections.read().await;
     let pool = connections.get(pool_key).ok_or("Connection not found")?;
@@ -2746,6 +2827,42 @@ mod tests {
         assert_eq!(sql, "SELECT \"id\", \"name\" FROM \"public\".\"users\" ORDER BY \"id\" LIMIT 100 OFFSET 200");
     }
 
+    #[test]
+    fn sqlserver_keyset_pagination_includes_offset_fetch() {
+        let sql = keyset_pagination_sql(
+            &[String::from("id"), String::from("name")],
+            "users",
+            "dbo",
+            &DatabaseType::SqlServer,
+            &[String::from("id")],
+            &[],
+            100,
+        );
+
+        assert_eq!(
+            sql,
+            "SELECT [id], [name] FROM [dbo].[users] ORDER BY [id] ASC OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"
+        );
+    }
+
+    #[test]
+    fn composite_keyset_pagination_uses_portable_lexicographic_predicate() {
+        let sql = keyset_pagination_sql(
+            &[String::from("tenant_id"), String::from("id"), String::from("name")],
+            "users",
+            "dbo",
+            &DatabaseType::SqlServer,
+            &[String::from("tenant_id"), String::from("id")],
+            &[json!(10), json!(25)],
+            100,
+        );
+
+        assert_eq!(
+            sql,
+            "SELECT [tenant_id], [id], [name] FROM [dbo].[users] WHERE ([tenant_id] > 10 OR ([tenant_id] = 10 AND [id] > 25)) ORDER BY [tenant_id] ASC, [id] ASC OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"
+        );
+    }
+
     #[test]
     fn postgres_generates_index_and_foreign_key_sql() {
         let indexes = vec![db::IndexInfo {