use serde::{Deserialize, Serialize}; use serde_json::Value; use std::borrow::Cow; use std::io::{BufWriter, Write}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; use tokio_util::sync::CancellationToken; use crate::connection::MysqlMode; use crate::connection::{config_for_pool_key, task_client_session_id, AppState, PoolKind}; #[cfg(test)] use crate::csv_export::format_tsv_rows; use crate::csv_export::{format_csv, format_tsv, push_table_csv_row, push_tsv_row}; pub use crate::database_export::ExportStatus; use crate::database_export::{ build_export_insert_statements, is_export_cancelled, is_internal_export_column, BuildExportInsertStatementsOptions, }; use crate::db::agent_driver::AgentTableReadStartParams; use crate::models::connection::DatabaseType; use crate::query::{close_query_session, execute_sql_statement_with_options, QueryExecutionOptions}; use crate::transfer::{ count_sql_with_where_and_identifier_quote, execute_read_on_pool, execute_read_on_pool_with_max_rows, keyset_pagination_sql_with_identifier_quote, pagination_sql_with_filter_order_and_identifier_quote, }; use crate::types::QueryResult; use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook_with_options}; const DEFAULT_BATCH_SIZE: usize = 10_000; const SQL_INSERT_BATCH_SIZE: usize = 100; pub fn table_export_client_session_id(export_id: &str) -> String { task_client_session_id("table-export", export_id) } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TableExportRequest { pub export_id: String, pub connection_id: String, pub database: String, pub schema: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub identifier_quote: Option, pub table_name: String, pub file_path: String, /// "csv", "xlsx", "json", "markdown", "sql", or "txt" pub format: String, #[serde(default)] pub columns: Option>, #[serde(default)] pub column_types: Option>>, #[serde(default)] pub primary_keys: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub where_input: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub order_by: Option, #[serde(default)] pub skip_count: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub batch_size: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub row_limit: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub date_time_format: Option, #[serde(default)] pub numeric_column_right_align: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub column_comments: Option>>, } #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct TableExportProgress { pub export_id: String, pub table_name: String, pub rows_exported: u64, pub total_rows: Option, pub status: ExportStatus, #[serde(default, skip_serializing_if = "Option::is_none")] pub error_message: Option, } /// Format rows as CSV text without a header row. /// Used for streaming subsequent pagination batches. #[cfg_attr(not(test), allow(dead_code))] fn format_csv_rows(rows: &[Vec]) -> String { // 注意：该无表头批次路径的 Null 输出为 ""（带引号空串），与查询结果导出 // 及首批 format_csv 的裸空单元格语义不同，直写化必须保留该差异 let mut out = String::with_capacity(crate::csv_export::estimated_rows_capacity(rows)); for (row_index, row) in rows.iter().enumerate() { if row_index > 0 { out.push('\n'); } push_table_csv_row(&mut out, row); } out } fn write_table_text_row(file: &mut W, csv: bool, row: &[Value], buffer: &mut String) -> Result<(), String> { buffer.clear(); buffer.push('\n'); if csv { push_table_csv_row(buffer, row); } else { push_tsv_row(buffer, row); } file.write_all(buffer.as_bytes()).map_err(|error| format!("Failed to write export rows: {error}")) } fn write_table_text_rows( file: &mut W, csv: bool, rows: &[Vec], buffer: &mut String, ) -> Result<(), String> { buffer.clear(); for row in rows { buffer.push('\n'); if csv { push_table_csv_row(buffer, row); } else { push_tsv_row(buffer, row); } } file.write_all(buffer.as_bytes()).map_err(|error| format!("Failed to write export rows: {error}")) } fn export_column_types(request: &TableExportRequest) -> Vec { request .column_types .as_deref() .unwrap_or_default() .iter() .map(|column_type| column_type.clone().unwrap_or_default()) .collect() } fn resolve_requested_export_columns( database_type: DatabaseType, columns: &[String], column_types: Option<&[Option]>, primary_keys: Option<&[String]>, ) -> (Vec, Vec>, Vec) { let mut resolved_columns = Vec::with_capacity(columns.len()); let mut resolved_column_types = column_types.map(|_| Vec::with_capacity(columns.len())).unwrap_or_default(); // Filter names and their index-aligned type metadata together so the // fetched rows and later SQL literal formatting keep the same positions. for (index, column) in columns.iter().enumerate() { if is_internal_export_column(Some(database_type), column) { continue; } resolved_columns.push(column.clone()); if let Some(column_types) = column_types { resolved_column_types.push(column_types.get(index).cloned().unwrap_or(None)); } } let resolved_primary_keys = primary_keys .unwrap_or_default() .iter() .filter(|column| !is_internal_export_column(Some(database_type), column)) .cloned() .collect(); (resolved_columns, resolved_column_types, resolved_primary_keys) } fn requested_mysql_sql_export_needs_column_metadata(database_type: DatabaseType, format: &str) -> bool { database_type == DatabaseType::Mysql && format.eq_ignore_ascii_case("sql") } fn resolve_requested_export_column_types( requested_columns: &[String], requested_column_types: &[Option], table_columns: &[crate::db::ColumnInfo], ) -> Vec> { requested_columns .iter() .enumerate() .map(|(index, requested)| { requested_column_types .get(index) .cloned() .flatten() .filter(|column_type| !column_type.trim().is_empty()) .or_else(|| { table_columns .iter() .find(|column| column.name.eq_ignore_ascii_case(requested)) .map(|column| column.data_type.clone()) }) }) .collect() } fn resolve_requested_export_column_extras( requested_columns: &[String], table_columns: &[crate::db::ColumnInfo], ) -> Vec> { requested_columns .iter() .map(|requested| { table_columns .iter() .find(|column| column.name.eq_ignore_ascii_case(requested)) .and_then(|column| column.extra.clone()) }) .collect() } fn write_json_row_object(writer: &mut W, columns: &[String], row: &[Value]) -> Result<(), String> { writer.write_all(b"{\n").map_err(|e| format!("Failed to write JSON: {e}"))?; let mut first = true; for (index, column) in columns.iter().enumerate() { let Some(value) = row.get(index) else { continue; }; if !first { writer.write_all(b",\n").map_err(|e| format!("Failed to write JSON: {e}"))?; } writer.write_all(b" ").map_err(|e| format!("Failed to write JSON: {e}"))?; serde_json::to_writer(&mut *writer, column).map_err(|e| format!("Failed to write JSON: {e}"))?; writer.write_all(b": ").map_err(|e| format!("Failed to write JSON: {e}"))?; serde_json::to_writer(&mut *writer, value).map_err(|e| format!("Failed to write JSON: {e}"))?; first = false; } writer.write_all(b"\n}").map_err(|e| format!("Failed to write JSON: {e}")) } fn display_cell(value: &Value) -> String { match value { Value::Null => "NULL".to_string(), Value::Bool(value) => value.to_string(), Value::Number(value) => value.to_string(), Value::String(value) => value.clone(), other => other.to_string(), } } fn markdown_cell(value: &str) -> String { value.replace('|', "\\|").replace("\r\n", "
").replace('\n', "
") } fn format_markdown_header(columns: &[String]) -> String { let header = columns.iter().map(|column| markdown_cell(column)).collect::>().join(" | "); let separator = columns.iter().map(|_| "---").collect::>().join(" | "); format!("| {header} |\n| {separator} |\n") } fn format_markdown_rows(rows: &[Vec]) -> String { rows.iter() .map(|row| { let cells = row.iter().map(|cell| markdown_cell(&display_cell(cell))).collect::>().join(" | "); format!("| {cells} |") }) .collect::>() .join("\n") } #[derive(Clone, Copy, Debug, PartialEq)] struct TableExportSqlContext<'a> { database_type: DatabaseType, schema: Option<&'a str>, } fn table_export_sql_context<'a>( database_type: DatabaseType, driver_profile: Option<&str>, schema: Option<&'a str>, ) -> TableExportSqlContext<'a> { // GBase8s selects from the active database and rejects the owner-qualified // names produced for MySQL-family GBase connections. if database_type == DatabaseType::Gbase && driver_profile.is_some_and(|profile| profile.eq_ignore_ascii_case("gbase8s")) { TableExportSqlContext { database_type: DatabaseType::Informix, schema: None } } else { TableExportSqlContext { database_type, schema } } } fn table_export_query_columns<'a>( request: &TableExportRequest, sql_context: &TableExportSqlContext<'_>, columns: &'a [String], ) -> Result, String> { if sql_context.database_type != DatabaseType::Iotdb { return Ok(Cow::Borrowed(columns)); } let full_table = crate::sql_dialect::table_data_qualified_table_name( Some(sql_context.database_type), sql_context.schema, &request.table_name, request.identifier_quote.as_deref(), ); let measurement_prefix = format!("{full_table}."); if !columns.iter().any(|column| column.eq_ignore_ascii_case("Time") || column.starts_with(&measurement_prefix)) { return Ok(Cow::Borrowed(columns)); } // IoTDB returns absolute timeseries labels, but its SELECT list accepts // only paths relative to the queried device. Keep the labels unchanged // for export output and normalize only the query projection here. let query_columns = columns .iter() .filter_map(|column| { if column.eq_ignore_ascii_case("Time") { None } else { Some(column.strip_prefix(&measurement_prefix).unwrap_or(column).to_string()) } }) .collect::>(); if query_columns.is_empty() { return Err("IoTDB table export requires at least one non-Time column".to_string()); } Ok(Cow::Owned(query_columns)) } #[allow(clippy::too_many_arguments)] fn table_page_sql( request: &TableExportRequest, sql_context: &TableExportSqlContext<'_>, col_names: &[String], column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], offset: u64, batch_size: usize, ) -> String { let sql = if use_keyset { keyset_pagination_sql_with_identifier_quote( col_names, &request.table_name, sql_context.schema.unwrap_or(""), &sql_context.database_type, primary_keys, last_pk_values, batch_size, request.identifier_quote.as_deref(), ) } else { pagination_sql_with_filter_order_and_identifier_quote( col_names, &request.table_name, sql_context.schema.unwrap_or(""), &sql_context.database_type, offset, batch_size, request.where_input.as_deref(), request.order_by.as_deref(), primary_keys, request.identifier_quote.as_deref(), ) }; replace_mysql_spatial_export_select_list(sql, request, &sql_context.database_type, col_names, column_types) } fn mysql_spatial_export_column_expression(column: &str, identifier_quote: Option<&str>) -> String { let quoted = crate::sql_dialect::quote_table_data_identifier(Some(DatabaseType::Mysql), column, identifier_quote); format!( "CASE WHEN {quoted} IS NULL THEN NULL ELSE CONCAT('DBX_WKB:', ST_SRID({quoted}), ':', HEX(ST_AsWKB({quoted}))) END AS {quoted}" ) } fn mysql_spatial_export_select_list( request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], column_types: &[Option], ) -> Option { if *db_type != DatabaseType::Mysql || !request.format.eq_ignore_ascii_case("sql") { return None; } let mut has_spatial_column = false; let expressions = col_names .iter() .enumerate() .map(|(index, column)| { if column_types .get(index) .and_then(|column_type| column_type.as_deref()) .is_some_and(crate::database_export::is_mysql_spatial_export_type) { has_spatial_column = true; mysql_spatial_export_column_expression(column, request.identifier_quote.as_deref()) } else { crate::sql_dialect::quote_table_data_identifier( Some(*db_type), column, request.identifier_quote.as_deref(), ) } }) .collect::>(); has_spatial_column.then(|| expressions.join(", ")) } fn replace_mysql_spatial_export_select_list( sql: String, request: &TableExportRequest, db_type: &DatabaseType, col_names: &[String], column_types: &[Option], ) -> String { let Some(replacement) = mysql_spatial_export_select_list(request, db_type, col_names, column_types) else { return sql; }; let original = col_names .iter() .map(|column| { crate::sql_dialect::quote_table_data_identifier(Some(*db_type), column, request.identifier_quote.as_deref()) }) .collect::>() .join(", "); let prefix = format!("SELECT {original}"); if sql.starts_with(&prefix) { format!("SELECT {replacement}{}", &sql[prefix.len()..]) } else { log::warn!( "MySQL spatial table export could not replace its SELECT list; geometry columns will be exported as WKT" ); sql } } fn table_cursor_sql( request: &TableExportRequest, sql_context: &TableExportSqlContext<'_>, col_names: &[String], column_types: &[Option], primary_keys: &[String], ) -> String { let full_table = crate::sql_dialect::table_data_qualified_table_name( Some(sql_context.database_type), sql_context.schema, &request.table_name, request.identifier_quote.as_deref(), ); let col_list = mysql_spatial_export_select_list(request, &sql_context.database_type, col_names, column_types) .unwrap_or_else(|| { col_names .iter() .map(|column| { crate::sql_dialect::quote_table_data_identifier( Some(sql_context.database_type), column, request.identifier_quote.as_deref(), ) }) .collect::>() .join(", ") }); let predicate = crate::sql_dialect::normalize_where_input(request.where_input.as_deref()); let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") }; let order_by = request .order_by .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) .map(str::to_string) .or_else(|| { if primary_keys.is_empty() { None } else { Some( primary_keys .iter() .map(|column| { format!( "{} ASC", crate::sql_dialect::quote_table_data_identifier( Some(sql_context.database_type), column, request.identifier_quote.as_deref(), ) ) }) .collect::>() .join(", "), ) } }) .map(|value| format!(" ORDER BY {value}")) .unwrap_or_default(); format!("SELECT {col_list} FROM {full_table}{where_clause}{order_by}") } fn is_agent_table_read_unsupported(error: &str) -> bool { let lower = error.to_ascii_lowercase(); lower.contains("unknown method") || lower.contains("method not found") } #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum TableExportCursorKind { Agent, ExternalDriver, } #[derive(Clone, Debug, PartialEq, Eq)] enum TableExportCursorSession { Agent(String), ExternalDriver(String), } async fn table_export_cursor_kind(state: &AppState, pool_key: &str) -> Option { let connections = state.connections.read().await; match connections.get(pool_key) { Some(PoolKind::Agent(_)) => Some(TableExportCursorKind::Agent), Some(PoolKind::ExternalDriver { .. }) => Some(TableExportCursorKind::ExternalDriver), _ => None, } } async fn table_export_query_timeout_secs(state: &AppState, pool_key: &str) -> u64 { let configs = state.configs.read().await; config_for_pool_key(pool_key, &configs).map(|config| config.query_timeout_secs).unwrap_or(0) } async fn execute_external_driver_export_page( state: &AppState, pool_key: &str, request: &TableExportRequest, sql_context: &TableExportSqlContext<'_>, query_col_names: &[String], column_types: &[Option], primary_keys: &[String], active_batch_size: usize, result_session_id: Option, cancel_token: CancellationToken, ) -> Result { let sql = table_cursor_sql(request, sql_context, query_col_names, column_types, primary_keys); let max_rows = request.row_limit.unwrap_or(i32::MAX as usize).min(i32::MAX as usize).max(1); let timeout_secs = table_export_query_timeout_secs(state, pool_key).await; execute_sql_statement_with_options( state, &request.connection_id, &request.database, &sql, request.schema.as_deref(), Some(cancel_token), QueryExecutionOptions { max_rows: Some(max_rows), fetch_size: Some(active_batch_size), page_size: Some(active_batch_size), result_session_id, client_session_id: Some(table_export_client_session_id(&request.export_id)), timeout_secs: Some(timeout_secs), ..Default::default() }, ) .await } async fn execute_table_export_count( state: &AppState, pool_key: &str, request: &TableExportRequest, sql: &str, cancel_token: CancellationToken, ) -> Result { if table_export_cursor_kind(state, pool_key).await != Some(TableExportCursorKind::ExternalDriver) { return execute_read_on_pool(state, pool_key, sql).await; } let timeout_secs = table_export_query_timeout_secs(state, pool_key).await; execute_sql_statement_with_options( state, &request.connection_id, &request.database, sql, request.schema.as_deref(), Some(cancel_token), QueryExecutionOptions { max_rows: Some(1), fetch_size: Some(1), client_session_id: Some(table_export_client_session_id(&request.export_id)), timeout_secs: Some(timeout_secs), ..Default::default() }, ) .await } #[allow(clippy::too_many_arguments)] async fn fetch_table_export_batch( state: &AppState, pool_key: &str, request: &TableExportRequest, db_type: &DatabaseType, sql_context: &TableExportSqlContext<'_>, col_names: &[String], query_col_names: &[String], column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], offset: u64, active_batch_size: usize, cursor_session: &mut Option, table_read_attempted: &mut bool, table_read_completed: &mut bool, cancel_token: CancellationToken, ) -> Result { if *table_read_completed { return Ok(QueryResult { columns: col_names.to_vec(), column_types: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], rows: Vec::new(), affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }); } // VictoriaMetrics metrics are read through MetricsQL rather than SQL table // pagination. Execute the range query once and let the normal writers emit // the returned matrix/vector rows. if *db_type == DatabaseType::VictoriaMetrics { *table_read_attempted = true; *table_read_completed = true; let query = crate::db::victoriametrics_driver::metric_range_query(&request.table_name, "1h"); return execute_sql_statement_with_options( state, &request.connection_id, &request.database, &query, request.schema.as_deref(), Some(cancel_token), QueryExecutionOptions { max_rows: request.row_limit, client_session_id: Some(table_export_client_session_id(&request.export_id)), ..Default::default() }, ) .await; } if !*table_read_attempted { match table_export_cursor_kind(state, pool_key).await { Some(TableExportCursorKind::Agent) => { *table_read_attempted = true; let sql = table_cursor_sql(request, sql_context, query_col_names, column_types, primary_keys); let max_rows = request.row_limit.unwrap_or(i32::MAX as usize); let query_timeout = table_export_query_timeout_secs(state, pool_key).await; let params = AgentTableReadStartParams { sql, database: Some(request.database.clone()), schema: request.schema.clone(), page_size: active_batch_size, max_rows, fetch_size: Some(active_batch_size), timeout_secs: (query_timeout > 0).then_some(query_timeout), }; let connections = state.connections.read().await; let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { return Err("Agent table read requires an agent connection".to_string()); }; let client = client.clone(); drop(connections); let mut client = client.lock().await; match client.start_table_read::(params).await { Ok(result) => { *cursor_session = result.session_id.clone().map(TableExportCursorSession::Agent); if result.session_id.is_none() && !result.has_more { *table_read_completed = true; } return Ok(result); } Err(error) if is_agent_table_read_unsupported(&error) => { log::debug!("Agent table-read cursor unsupported, falling back to paginated export: {error}"); } Err(error) => return Err(error), } } Some(TableExportCursorKind::ExternalDriver) => { *table_read_attempted = true; let result = execute_external_driver_export_page( state, pool_key, request, sql_context, query_col_names, column_types, primary_keys, active_batch_size, None, cancel_token.clone(), ) .await?; if result.has_more { let session_id = result .session_id .clone() .ok_or("JDBC export cursor did not return a session id for additional rows")?; *cursor_session = Some(TableExportCursorSession::ExternalDriver(session_id)); } else { *table_read_completed = true; } return Ok(result); } None => {} } } if let Some(session) = cursor_session.clone() { return match session { TableExportCursorSession::Agent(session_id) => { let connections = state.connections.read().await; let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { return Err("Table read session requires an agent connection".to_string()); }; let client = client.clone(); drop(connections); let mut client = client.lock().await; match client.fetch_table_read_page::(&session_id, active_batch_size).await { Ok(result) => { *cursor_session = result.session_id.clone().or(Some(session_id)).map(TableExportCursorSession::Agent); if !result.has_more { *cursor_session = None; *table_read_completed = true; } Ok(result) } Err(error) => { let _ = client.close_table_read_session::(&session_id).await; *cursor_session = None; Err(error) } } } TableExportCursorSession::ExternalDriver(session_id) => { match execute_external_driver_export_page( state, pool_key, request, sql_context, query_col_names, column_types, primary_keys, active_batch_size, Some(session_id.clone()), cancel_token.clone(), ) .await { Ok(result) => { if result.has_more { let next_session_id = result.session_id.clone().unwrap_or(session_id); *cursor_session = Some(TableExportCursorSession::ExternalDriver(next_session_id)); } else { *cursor_session = None; *table_read_completed = true; } Ok(result) } Err(error) => { if cancel_token.is_cancelled() { cursor_session.take(); } else { close_table_export_cursor_if_open(state, pool_key, request, cursor_session).await; } Err(error) } } } }; } fetch_paginated_table_export_batch( state, pool_key, request, sql_context, query_col_names, column_types, primary_keys, use_keyset, last_pk_values, offset, active_batch_size, ) .await } #[allow(clippy::too_many_arguments)] async fn fetch_paginated_table_export_batch( state: &AppState, pool_key: &str, request: &TableExportRequest, sql_context: &TableExportSqlContext<'_>, query_col_names: &[String], column_types: &[Option], primary_keys: &[String], use_keyset: bool, last_pk_values: &[Value], offset: u64, active_batch_size: usize, ) -> Result { let sql = table_page_sql( request, sql_context, query_col_names, column_types, primary_keys, use_keyset, last_pk_values, offset, active_batch_size, ); execute_read_on_pool_with_max_rows(state, pool_key, &sql, Some(active_batch_size)).await } async fn close_table_export_cursor_if_open( state: &AppState, pool_key: &str, request: &TableExportRequest, cursor_session: &mut Option, ) { let Some(session) = cursor_session.take() else { return; }; match session { TableExportCursorSession::Agent(session_id) => { let connections = state.connections.read().await; let Some(PoolKind::Agent(client)) = connections.get(pool_key) else { return; }; let client = client.clone(); drop(connections); let mut client = client.lock().await; let _ = client.close_table_read_session::(&session_id).await; } TableExportCursorSession::ExternalDriver(session_id) => { let client_session_id = table_export_client_session_id(&request.export_id); let _ = close_query_session( state, &request.connection_id, &request.database, &session_id, Some(&client_session_id), None, ) .await; } } } async fn start_export_cancel_watcher(export_id: String, cancelled: Arc, token: CancellationToken) { loop { if is_export_cancelled(&export_id).await { cancelled.store(true, Ordering::SeqCst); token.cancel(); break; } tokio::time::sleep(Duration::from_millis(100)).await; } } async fn stream_native_table_rows( state: &AppState, pool_key: &str, db_type: &DatabaseType, sql: &str, row_limit: Option, cancelled: &AtomicBool, cancel_token: CancellationToken, on_row: impl FnMut(&[Value]) -> Result<(), String>, ) -> Result { let connections = state.connections.read().await; match connections.get(pool_key) { Some(PoolKind::Mysql(pool, mode)) => { let pool = pool.clone(); let bare = *mode == MysqlMode::Bare; drop(connections); crate::db::mysql::stream_query_rows( &pool, sql, bare, row_limit, crate::db::mysql::MySqlQueryDialect::for_connection(*db_type, None), cancelled, on_row, ) .await?; Ok(true) } Some(PoolKind::Postgres(pool)) => { let pool = pool.clone(); drop(connections); crate::db::postgres::stream_query_rows(&pool, sql, row_limit, cancelled, on_row).await?; Ok(true) } Some(PoolKind::SqlServer(client)) => { let client = client.clone(); drop(connections); let mut on_row = on_row; let mut client = client.lock().await; crate::db::sqlserver::stream_first_result_set(&mut client, sql, row_limit, Some(cancel_token), |item| { if let crate::db::sqlserver::SqlServerStreamItem::Row(row) = item { on_row(row)?; } Ok(()) }) .await?; Ok(true) } _ => Ok(false), } } #[allow(clippy::too_many_arguments)] async fn try_export_native_table_stream( state: &AppState, pool_key: &str, request: &TableExportRequest, db_type: &DatabaseType, sql_context: &TableExportSqlContext<'_>, col_names: &[String], query_col_names: &[String], column_types: &[Option], column_extras: &[Option], primary_keys: &[String], total_rows: Option, row_limit: Option, batch_size: usize, on_progress: &impl Fn(TableExportProgress), cancelled: Arc, cancel_token: CancellationToken, ) -> Result { let sql = table_cursor_sql(request, sql_context, query_col_names, column_types, primary_keys); let mut rows_exported = 0_u64; let progress_interval = batch_size.max(1) as u64; let stream_result = match request.format.to_lowercase().as_str() { "csv" => { let mut file = BufWriter::new( std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?, ); file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?; let header = format_csv(col_names, &[]); let header = header.strip_suffix('\n').unwrap_or(&header); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?; let mut row_buffer = String::new(); let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { let formatted = crate::temporal_format::format_temporal_export_row_cow( row, column_types, request.date_time_format.as_deref(), ); write_table_text_row(&mut file, true, formatted.as_ref(), &mut row_buffer)?; rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; } result } "txt" => { let mut file = BufWriter::new( std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?, ); let header = format_tsv(col_names, &[]); let header = header.strip_suffix('\n').unwrap_or(&header); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write TXT: {e}"))?; let mut row_buffer = String::new(); let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { let formatted = crate::temporal_format::format_temporal_export_row_cow( row, column_types, request.date_time_format.as_deref(), ); write_table_text_row(&mut file, false, formatted.as_ref(), &mut row_buffer)?; rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; } result } "xlsx" => { let xlsx_column_types = export_column_types(request); let column_comments: Vec> = request.column_comments.clone().unwrap_or_default(); let xlsx_file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; let mut writer = start_streaming_xlsx_workbook_with_options( BufWriter::new(xlsx_file), Some(&request.table_name), col_names, &xlsx_column_types, &column_comments, &[], request.date_time_format.as_deref(), request.numeric_column_right_align, )?; let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { let formatted = crate::temporal_format::format_temporal_export_row_cow( row, column_types, request.date_time_format.as_deref(), ); writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Writing, error_message: None, }); let mut xlsx_buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; xlsx_buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; } result } "json" => { let mut file = BufWriter::new( std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?, ); file.write_all(b"[\n").map_err(|e| format!("Failed to write JSON: {e}"))?; let mut is_first_row = true; let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { if !is_first_row { file.write_all(b",\n").map_err(|e| format!("Failed to write JSON: {e}"))?; } let formatted = crate::temporal_format::format_temporal_export_row_cow( row, column_types, request.date_time_format.as_deref(), ); write_json_row_object(&mut file, col_names, formatted.as_ref())?; is_first_row = false; rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { file.write_all(b"\n]\n").map_err(|e| format!("Failed to write JSON: {e}"))?; file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; } result } "markdown" | "md" => { let mut file = BufWriter::new( std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?, ); file.write_all(format_markdown_header(col_names).as_bytes()) .map_err(|e| format!("Failed to write Markdown: {e}"))?; let mut wrote_rows = false; let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { let formatted = crate::temporal_format::format_temporal_export_row( row, column_types, request.date_time_format.as_deref(), ); let rows_markdown = format_markdown_rows(&[formatted]); if !rows_markdown.is_empty() { if wrote_rows { file.write_all(b"\n").map_err(|e| format!("Failed to write Markdown: {e}"))?; } file.write_all(rows_markdown.as_bytes()) .map_err(|e| format!("Failed to write Markdown: {e}"))?; wrote_rows = true; } rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { file.write_all(b"\n").map_err(|e| format!("Failed to write Markdown: {e}"))?; file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; } result } "sql" => { let mut file = BufWriter::new( std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?, ); let mut pending_rows: Vec> = Vec::new(); let mut wrote_statements = false; let mut flush_pending = |file: &mut BufWriter, pending_rows: &mut Vec>| -> Result<(), String> { if pending_rows.is_empty() { return Ok(()); } let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { database_type: Some(*db_type), identifier_quote: request.identifier_quote.clone(), schema: request.schema.clone(), table_name: Some(request.table_name.clone()), qualified_table_name: None, columns: col_names.to_vec(), column_types: column_types.to_vec(), column_extras: column_extras.to_vec(), spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: std::mem::take(pending_rows), batch_size: Some(SQL_INSERT_BATCH_SIZE), })?; if !statements.is_empty() { if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; } file.write_all(statements.join("\n").as_bytes()) .map_err(|e| format!("Failed to write SQL: {e}"))?; wrote_statements = true; } Ok(()) }; let result = stream_native_table_rows( state, pool_key, db_type, &sql, row_limit, &cancelled, cancel_token.clone(), |row| { pending_rows.push(row.to_vec()); if pending_rows.len() >= SQL_INSERT_BATCH_SIZE { flush_pending(&mut file, &mut pending_rows)?; } rows_exported += 1; if rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } Ok(()) }, ) .await; if result.is_ok() { flush_pending(&mut file, &mut pending_rows)?; if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; } file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; } result } _ => Ok(false), }; match stream_result { Ok(false) => Ok(false), Ok(true) => { if !rows_exported.is_multiple_of(progress_interval) { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); } on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Done, error_message: None, }); Ok(true) } Err(error) if cancelled.load(Ordering::SeqCst) || error == crate::query::canceled_error() => { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); Ok(true) } Err(error) => Err(error), } } fn next_export_batch_size(row_limit: Option, rows_exported: u64, batch_size: usize) -> Option { let remaining = row_limit.map(|limit| limit.saturating_sub(rows_exported as usize)); if matches!(remaining, Some(0)) { return None; } Some(remaining.map_or(batch_size, |value| value.min(batch_size)).max(1)) } pub async fn export_table_data_core( state: &AppState, request: &TableExportRequest, on_progress: impl Fn(TableExportProgress), ) -> Result<(), String> { let cancelled = Arc::new(AtomicBool::new(false)); let cancel_token = CancellationToken::new(); let cancel_watcher = tokio::spawn(start_export_cancel_watcher(request.export_id.clone(), cancelled.clone(), cancel_token.clone())); let last_rows_exported = std::sync::atomic::AtomicU64::new(0); let tracked_progress = |progress: TableExportProgress| { last_rows_exported.store(progress.rows_exported, Ordering::SeqCst); on_progress(progress); }; let result = export_table_data_core_inner(state, request, &tracked_progress, cancelled.clone(), cancel_token.clone()).await; cancel_watcher.abort(); let client_session_id = table_export_client_session_id(&request.export_id); let _ = state.close_client_session_pool(&request.connection_id, Some(&request.database), &client_session_id).await; match result { Err(error) if cancelled.load(Ordering::SeqCst) || error == crate::query::canceled_error() => { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported: last_rows_exported.load(Ordering::SeqCst), total_rows: None, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); Ok(()) } result => result, } } async fn export_table_data_core_inner( state: &AppState, request: &TableExportRequest, on_progress: &impl Fn(TableExportProgress), cancelled: Arc, cancel_token: CancellationToken, ) -> Result<(), String> { // 1. Get database type and the profile that can refine its SQL dialect. let (db_type, driver_profile) = { let configs = state.configs.read().await; let config = configs .get(&request.connection_id) .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?; (config.db_type, config.driver_profile.clone()) }; let sql_context = table_export_sql_context(db_type, driver_profile.as_deref(), request.schema.as_deref()); // 2. Get pool let client_session_id = table_export_client_session_id(&request.export_id); let pool_key = state .get_or_create_pool_for_session(&request.connection_id, Some(&request.database), Some(&client_session_id)) .await?; // 3. Resolve columns. Data grid exports can provide columns/primary keys // directly, which avoids expensive metadata round-trips on JDBC drivers. let requested_columns = request.columns.as_ref().filter(|columns| !columns.is_empty()); let (col_names, column_types, column_extras, primary_keys) = if let Some(requested_columns) = requested_columns { let (col_names, requested_column_types, primary_keys) = resolve_requested_export_columns( db_type, requested_columns, request.column_types.as_deref(), request.primary_keys.as_deref(), ); let (column_types, column_extras) = if requested_mysql_sql_export_needs_column_metadata(db_type, &request.format) { let table_columns = crate::schema::get_columns_core( state, &request.connection_id, &request.database, request.schema.as_deref().unwrap_or(""), &request.table_name, ) .await?; ( resolve_requested_export_column_types(&col_names, &requested_column_types, &table_columns), resolve_requested_export_column_extras(&col_names, &table_columns), ) } else { (requested_column_types, Vec::new()) }; (col_names, column_types, column_extras, primary_keys) } else { let columns = crate::schema::get_columns_core( state, &request.connection_id, &request.database, request.schema.as_deref().unwrap_or(""), &request.table_name, ) .await?; let col_names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let column_types: Vec> = columns.iter().map(|c| Some(c.data_type.clone())).collect(); let column_extras: Vec> = columns.iter().map(|c| c.extra.clone()).collect(); let primary_keys: Vec = columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.clone()).collect(); (col_names, column_types, column_extras, primary_keys) }; if col_names.is_empty() { return Err("No columns found for table".to_string()); } let query_col_names = table_export_query_columns(request, &sql_context, &col_names)?; // 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 has_custom_filter_or_order = request.where_input.as_ref().is_some_and(|value| !value.trim().is_empty()) || request.order_by.as_ref().is_some_and(|value| !value.trim().is_empty()); let use_keyset = !has_custom_filter_or_order && !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 = 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 when requested. Data // grid exports skip this by default because COUNT can be the slowest query // on large HANA/JDBC tables, especially with filters. let row_limit = request.row_limit; let total_rows = if request.skip_count || db_type == DatabaseType::VictoriaMetrics { None } else { let count_query = count_sql_with_where_and_identifier_quote( &request.table_name, sql_context.schema.unwrap_or(""), &sql_context.database_type, request.where_input.as_deref(), None, request.identifier_quote.as_deref(), ); match execute_table_export_count(state, &pool_key, request, &count_query, cancel_token.clone()).await { Ok(result) => result .rows .first() .and_then(|r| r.first()) .and_then(|v| match v { Value::Number(n) => n.as_u64(), Value::String(s) => s.parse::().ok(), _ => None, }) .map(|total| row_limit.map_or(total, |limit| total.min(limit as u64))), Err(_) => None, } }; // 7. Emit initial Running progress on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported: 0, total_rows, status: ExportStatus::Running, error_message: None, }); if try_export_native_table_stream( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &column_extras, &primary_keys, total_rows, row_limit, request.batch_size.unwrap_or(DEFAULT_BATCH_SIZE).max(1), &on_progress, cancelled, cancel_token.clone(), ) .await? { return Ok(()); } // 8. Create output file let file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?; let mut file = BufWriter::new(file); let mut text_buffer = String::new(); 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 mut cursor_session: Option = None; let mut table_read_attempted = false; let mut table_read_completed = false; // Track last primary key values for keyset pagination let mut last_pk_values: Vec = Vec::new(); match request.format.to_lowercase().as_str() { "csv" => { // Write UTF-8 BOM file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?; let mut is_first_batch = true; loop { // Check cancellation between batches if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } let formatted_rows = crate::temporal_format::format_temporal_export_rows_cow( &result.rows, &column_types, request.date_time_format.as_deref(), ); if is_first_batch { // First batch: write header + rows via format_csv let csv_content = format_csv(&col_names, formatted_rows.as_ref()); file.write_all(csv_content.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?; is_first_batch = false; } else { // Subsequent batches: write rows only (prepend newline for separation) write_table_text_rows(&mut file, true, formatted_rows.as_ref(), &mut text_buffer)?; } rows_exported += 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } } "txt" => { let mut is_first_batch = true; let header = format_tsv(&col_names, &[]); let header = header.strip_suffix('\n').unwrap_or(&header); file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write TXT: {e}"))?; loop { if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } let formatted_rows = crate::temporal_format::format_temporal_export_rows_cow( &result.rows, &column_types, request.date_time_format.as_deref(), ); if is_first_batch { write_table_text_rows(&mut file, false, formatted_rows.as_ref(), &mut text_buffer)?; is_first_batch = false; } else { write_table_text_rows(&mut file, false, formatted_rows.as_ref(), &mut text_buffer)?; } rows_exported += row_count as u64; if use_keyset { 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } } "xlsx" => { let xlsx_column_types = export_column_types(request); let column_comments: Vec> = request.column_comments.clone().unwrap_or_default(); // Create a dedicated file handle for the streaming XLSX writer // instead of cloning the outer BufWriter's handle. This avoids // sharing a file descriptor between two independent buffers. let xlsx_file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; let mut writer = start_streaming_xlsx_workbook_with_options( BufWriter::new(xlsx_file), Some(&request.table_name), &col_names, &xlsx_column_types, &column_comments, &[], request.date_time_format.as_deref(), request.numeric_column_right_align, )?; loop { // Check cancellation between batches if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } for row in &result.rows { let formatted = crate::temporal_format::format_temporal_export_row_cow( row, &column_types, request.date_time_format.as_deref(), ); writer.write_row(formatted.as_ref()).map_err(|e| format!("Failed to write XLSX row: {e}"))?; } rows_exported += 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } // Emit Writing progress before building XLSX on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Writing, error_message: None, }); // Explicitly flush the XLSX writer's BufWriter so IO errors // (e.g. disk-full) are surfaced rather than silently swallowed // by Drop. let mut xlsx_buf = finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; xlsx_buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; } "json" => { file.write_all(b"[\n").map_err(|e| format!("Failed to write JSON: {e}"))?; let mut is_first_row = true; loop { if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } for row in &result.rows { if !is_first_row { file.write_all(b",\n").map_err(|e| format!("Failed to write JSON: {e}"))?; } let formatted = crate::temporal_format::format_temporal_export_row_cow( row, &column_types, request.date_time_format.as_deref(), ); write_json_row_object(&mut file, &col_names, formatted.as_ref())?; is_first_row = false; } rows_exported += row_count as u64; if use_keyset { 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } file.write_all(b"\n]\n").map_err(|e| format!("Failed to write JSON: {e}"))?; } "markdown" | "md" => { file.write_all(format_markdown_header(&col_names).as_bytes()) .map_err(|e| format!("Failed to write Markdown: {e}"))?; let mut wrote_rows = false; loop { if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } let formatted_rows = crate::temporal_format::format_temporal_export_rows_cow( &result.rows, &column_types, request.date_time_format.as_deref(), ); let rows_markdown = format_markdown_rows(formatted_rows.as_ref()); if !rows_markdown.is_empty() { if wrote_rows { file.write_all(b"\n").map_err(|e| format!("Failed to write Markdown: {e}"))?; } file.write_all(rows_markdown.as_bytes()).map_err(|e| format!("Failed to write Markdown: {e}"))?; wrote_rows = true; } rows_exported += row_count as u64; if use_keyset { 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } file.write_all(b"\n").map_err(|e| format!("Failed to write Markdown: {e}"))?; } "sql" => { let mut wrote_statements = false; loop { if is_export_cancelled(&request.export_id).await { on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Cancelled, error_message: Some("Export cancelled".to_string()), }); close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; return Ok(()); } let Some(active_batch_size) = next_export_batch_size(row_limit, rows_exported, batch_size) else { break; }; let result = fetch_table_export_batch( state, &pool_key, request, &db_type, &sql_context, &col_names, query_col_names.as_ref(), &column_types, &primary_keys, use_keyset, &last_pk_values, offset, active_batch_size, &mut cursor_session, &mut table_read_attempted, &mut table_read_completed, cancel_token.clone(), ) .await?; let row_count = result.rows.len(); if row_count == 0 { break; } let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { database_type: Some(db_type), identifier_quote: request.identifier_quote.clone(), schema: request.schema.clone(), table_name: Some(request.table_name.clone()), qualified_table_name: None, columns: col_names.clone(), column_types: column_types.clone(), column_extras: column_extras.clone(), spatial_columns: result.spatial_columns.clone(), spatial_values: result.spatial_values.clone(), rows: result.rows.clone(), batch_size: Some(100), })?; if !statements.is_empty() { if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; } file.write_all(statements.join("\n").as_bytes()) .map_err(|e| format!("Failed to write SQL: {e}"))?; wrote_statements = true; } rows_exported += row_count as u64; if use_keyset { 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(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Running, error_message: None, }); if row_count < active_batch_size { break; } } if wrote_statements { file.write_all(b"\n").map_err(|e| format!("Failed to write SQL: {e}"))?; } } other => { return Err(format!("Unsupported export format: {other}")); } } close_table_export_cursor_if_open(state, &pool_key, request, &mut cursor_session).await; file.flush().map_err(|e| format!("Failed to flush export file: {e}"))?; // 8. Emit Done progress on_progress(TableExportProgress { export_id: request.export_id.clone(), table_name: request.table_name.clone(), rows_exported, total_rows, status: ExportStatus::Done, error_message: None, }); Ok(()) } #[cfg(test)] mod tests { use super::*; use crate::database_export::{clear_export_cancelled, set_export_cancelled}; #[cfg(unix)] use crate::models::connection::ConnectionConfig; #[cfg(unix)] use crate::plugins::{ InstalledPlugin, PluginDriverManifest, PluginDriverSession, PluginManifest, PluginRuntimeEnv, }; #[cfg(unix)] use crate::storage::Storage; use crate::xlsx_export::{build_xlsx_workbook, XlsxWorksheetData}; use serde_json::json; use std::io::Read; #[cfg(unix)] use std::os::unix::fs::PermissionsExt; #[cfg(unix)] struct ExternalDriverExportFixture { state: AppState, request: TableExportRequest, calls: std::path::PathBuf, output: std::path::PathBuf, dir: std::path::PathBuf, } #[cfg(unix)] async fn external_driver_export_fixture( rpc_body: &str, batch_size: usize, row_limit: Option, skip_count: bool, ) -> ExternalDriverExportFixture { let dir = std::env::temp_dir().join(format!("dbx-jdbc-table-export-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let executable = dir.join("plugin.sh"); let calls = dir.join("calls.log"); let script = format!( "#!/bin/sh\nCALLS='{}'\nwhile IFS= read -r line; do\n id=$(printf '%s' \"$line\" | sed -E 's/.*\"id\":([0-9]+).*/\\1/')\n{}\ndone\n", calls.display(), rpc_body ); std::fs::write(&executable, script).unwrap(); let mut permissions = std::fs::metadata(&executable).unwrap().permissions(); permissions.set_mode(0o755); std::fs::set_permissions(&executable, permissions).unwrap(); let plugin = InstalledPlugin { manifest: PluginManifest { id: "jdbc".to_string(), name: "JDBC".to_string(), version: "test".to_string(), protocol_version: 1, description: String::new(), executable: Some("plugin.sh".to_string()), drivers: vec![PluginDriverManifest { id: "jdbc".to_string(), label: "JDBC".to_string(), kind: "external".to_string(), database_type: Some("jdbc".to_string()), }], }, path: dir.clone(), }; let session = Arc::new( PluginDriverSession::start_for_test(plugin, "jdbc".to_string(), PluginRuntimeEnv::default()) .await .expect("test JDBC plugin should start"), ); let config: ConnectionConfig = serde_json::from_value(json!({ "id": "conn-1", "name": "JDBC", "db_type": "jdbc", "host": "", "port": 0, "username": "", "password": "", "database": "PUBLIC", "query_timeout_secs": 30 })) .unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); state.configs.write().await.insert(config.id.clone(), config.clone()); let export_id = format!("export-{}", uuid::Uuid::new_v4()); let pool_key = format!("{}:session:{}", config.id, table_export_client_session_id(&export_id).replace(':', "_")); state.connections.write().await.insert( pool_key, PoolKind::ExternalDriver { driver_id: "jdbc".to_string(), config: Arc::new(config), session }, ); let output = dir.join("export.csv"); let request = TableExportRequest { export_id, connection_id: "conn-1".to_string(), database: "PUBLIC".to_string(), schema: Some("PUBLIC".to_string()), identifier_quote: None, table_name: "EXPORT_SAMPLE".to_string(), file_path: output.to_string_lossy().into_owned(), format: "csv".to_string(), columns: Some(vec!["id".to_string(), "name".to_string()]), column_types: Some(vec![Some("INTEGER".to_string()), Some("VARCHAR".to_string())]), primary_keys: Some(vec!["id".to_string()]), where_input: None, order_by: None, skip_count, batch_size: Some(batch_size), row_limit, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; ExternalDriverExportFixture { state, request, calls, output, dir } } #[cfg(unix)] async fn run_external_driver_export( fixture: &ExternalDriverExportFixture, ) -> Result, String> { let progress = Arc::new(std::sync::Mutex::new(Vec::new())); let captured = progress.clone(); let result = export_table_data_core(&fixture.state, &fixture.request, move |event| { captured.lock().unwrap().push(event); }) .await; let events = progress.lock().unwrap().clone(); result.map(|_| events) } #[cfg(unix)] async fn wait_for_external_driver_call(calls: &std::path::Path, expected: &str) { tokio::time::timeout(Duration::from_secs(5), async { loop { if std::fs::read_to_string(calls).unwrap_or_default().lines().any(|line| line == expected) { return; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await .unwrap_or_else(|_| panic!("timed out waiting for plugin call: {expected}")); } #[cfg(unix)] fn cleanup_external_driver_export_fixture(fixture: ExternalDriverExportFixture) { let _ = std::fs::remove_dir_all(fixture.dir); } /// Read and decompress a single entry from an in-memory XLSX (ZIP) buffer. fn read_zip_entry(bytes: &[u8], path: &str) -> String { let mut archive = zip::ZipArchive::new(std::io::Cursor::new(bytes.to_vec())).expect("open xlsx as zip archive"); let mut entry = archive.by_name(path).unwrap_or_else(|_| panic!("missing zip entry: {path}")); let mut content = String::new(); entry.read_to_string(&mut content).expect("read zip entry"); content } // ----------------------------------------------------------------------- // Helper: check that two CSV strings are equivalent by splitting lines // ----------------------------------------------------------------------- fn csv_lines_equal(actual: &str, expected: &str) -> bool { let actual_lines: Vec<&str> = actual.lines().collect(); let expected_lines: Vec<&str> = expected.lines().collect(); actual_lines == expected_lines } // ----------------------------------------------------------------------- // format_csv_rows // ----------------------------------------------------------------------- #[test] fn formats_csv_rows_with_multiple_columns() { let rows = vec![vec![json!(1), json!("Alice")], vec![json!(2), json!("Bob \"Builder\"")]]; let out = format_csv_rows(&rows); assert!(csv_lines_equal(&out, "\"1\",\"Alice\"\n\"2\",\"Bob \"\"Builder\"\"\"")); } #[test] fn formats_csv_rows_with_null_values() { let rows = vec![vec![json!(1), Value::Null, json!("active")], vec![json!(2), json!("some notes"), Value::Null]]; let out = format_csv_rows(&rows); assert!(csv_lines_equal(&out, "\"1\",\"\",\"active\"\n\"2\",\"some notes\",\"\"")); } #[test] fn formats_csv_rows_with_boolean_and_number_values() { let rows = vec![vec![json!(true), json!(2.75)], vec![json!(false), json!(-42)]]; let out = format_csv_rows(&rows); assert!(csv_lines_equal(&out, "\"true\",\"2.75\"\n\"false\",\"-42\"")); } #[test] fn formats_csv_rows_returns_empty_string_for_empty_rows() { let rows: Vec> = vec![]; let out = format_csv_rows(&rows); assert_eq!(out, ""); } #[test] fn formats_csv_rows_single_row() { let rows = vec![vec![json!("just"), json!("one")]]; let out = format_csv_rows(&rows); assert_eq!(out, "\"just\",\"one\""); } #[test] fn reusable_text_row_buffer_preserves_table_null_semantics() { let row = vec![Value::Null, json!(""), json!("line\n\"two\"")]; let mut output = Vec::new(); let mut buffer = String::new(); write_table_text_row(&mut output, true, &row, &mut buffer).expect("write csv row"); assert_eq!(String::from_utf8(output).expect("utf8 csv"), "\n\"\",\"\",\"line\n\"\"two\"\"\""); } // ----------------------------------------------------------------------- // format_tsv (Navicat-style TXT export) // ----------------------------------------------------------------------- #[test] fn formats_tsv_with_header_and_tab_separated_values() { let columns = vec!["id".to_string(), "name".to_string()]; let rows = vec![vec![json!(1), json!("Alice")], vec![json!(2), json!("Bob")]]; assert_eq!(format_tsv(&columns, &rows), "id\tname\n1\tAlice\n2\tBob"); } #[test] fn formats_tsv_renders_null_as_empty() { let columns = vec!["id".to_string(), "note".to_string()]; let rows = vec![vec![json!(1), Value::Null]]; assert_eq!(format_tsv(&columns, &rows), "id\tnote\n1\t"); } #[test] fn formats_tsv_quotes_fields_containing_tab_or_newline() { let columns = vec!["a".to_string(), "b".to_string()]; let rows = vec![vec![json!("x\ty"), json!("line1\nline2")]]; assert_eq!(format_tsv(&columns, &rows), "a\tb\n\"x\ty\"\t\"line1\nline2\""); } #[test] fn formats_tsv_escapes_embedded_quotes() { let columns = vec!["name".to_string()]; let rows = vec![vec![json!(r#"Bob "Builder""#)]]; assert_eq!(format_tsv(&columns, &rows), "name\n\"Bob \"\"Builder\"\"\""); } #[test] fn formats_tsv_rows_returns_empty_for_empty_rows() { let rows: Vec> = vec![]; assert_eq!(format_tsv_rows(&rows), ""); } #[test] fn export_batch_size_respects_row_limit_remaining_rows() { assert_eq!(next_export_batch_size(None, 12_000, 10_000), Some(10_000)); assert_eq!(next_export_batch_size(Some(15_000), 0, 10_000), Some(10_000)); assert_eq!(next_export_batch_size(Some(15_000), 10_000, 10_000), Some(5_000)); assert_eq!(next_export_batch_size(Some(15_000), 15_000, 10_000), None); } #[test] fn iotdb_table_export_omits_implicit_time_from_all_query_paths() { let request = TableExportRequest { export_id: "export-iotdb".to_string(), connection_id: "conn-1".to_string(), database: "root.test".to_string(), schema: Some("root.test".to_string()), identifier_quote: None, table_name: "device2".to_string(), file_path: "device2.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: Some("WHERE temperature > 1".to_string()), order_by: Some("Time DESC".to_string()), skip_count: true, batch_size: Some(50), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Iotdb, None, request.schema.as_deref()); let columns = vec!["Time".to_string(), "root.test.device2.temperature".to_string()]; let query_columns = table_export_query_columns(&request, &context, &columns).unwrap(); assert_eq!(query_columns.as_ref(), &["temperature".to_string()]); assert_eq!( table_cursor_sql(&request, &context, query_columns.as_ref(), &[], &[]), "SELECT temperature FROM root.test.device2 WHERE (temperature > 1) ORDER BY Time DESC" ); assert_eq!( table_page_sql(&request, &context, query_columns.as_ref(), &[], &[], false, &[], 100, 50), "SELECT temperature FROM \"root.test\".\"device2\" WHERE (temperature > 1) ORDER BY Time DESC LIMIT 50 OFFSET 100" ); let csv = format_csv(&columns, &[vec![json!(1_700_000_000_000_i64), json!(21.5)]]); assert!(csv.starts_with("\"Time\",\"root.test.device2.temperature\"\n")); assert!(csv.contains("\"1700000000000\",\"21.5\"")); let workbook = build_xlsx_workbook(&XlsxWorksheetData { sheet_name: Some("device2".to_string()), columns, column_types: vec!["INT64".to_string(), "DOUBLE".to_string()], column_comments: vec![], rows: vec![vec![json!(1_700_000_000_000_i64), json!(21.5)]], numeric_column_right_align: false, }) .unwrap(); let sheet = read_zip_entry(&workbook, "xl/worksheets/sheet1.xml"); assert!(sheet.contains("Time")); assert!(sheet.contains("root.test.device2.temperature")); } #[test] fn iotdb_table_export_matches_time_case_insensitively() { let request = TableExportRequest { export_id: "export-iotdb-case".to_string(), connection_id: "conn-1".to_string(), database: "root.test".to_string(), schema: Some("root.test".to_string()), identifier_quote: None, table_name: "device2".to_string(), file_path: "device2.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: true, batch_size: Some(50), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Iotdb, None, request.schema.as_deref()); let columns = vec!["tImE".to_string(), "temperature".to_string()]; let query_columns = table_export_query_columns(&request, &context, &columns).unwrap(); assert_eq!( table_cursor_sql(&request, &context, query_columns.as_ref(), &[], &[]), "SELECT temperature FROM root.test.device2" ); } #[test] fn iotdb_table_export_rejects_only_implicit_time() { let request = TableExportRequest { export_id: "export-iotdb-time".to_string(), connection_id: "conn-1".to_string(), database: "root.test".to_string(), schema: Some("root.test".to_string()), identifier_quote: None, table_name: "device2".to_string(), file_path: "device2.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: true, batch_size: Some(50), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Iotdb, None, request.schema.as_deref()); let error = table_export_query_columns(&request, &context, &["TIME".to_string()]).unwrap_err(); assert_eq!(error, "IoTDB table export requires at least one non-Time column"); } #[test] fn iotdb_table_export_does_not_guess_other_device_paths() { let request = TableExportRequest { export_id: "export-iotdb-other-device".to_string(), connection_id: "conn-1".to_string(), database: "root.test".to_string(), schema: Some("root.test".to_string()), identifier_quote: None, table_name: "device2".to_string(), file_path: "device2.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: true, batch_size: Some(50), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Iotdb, None, request.schema.as_deref()); let columns = vec![ "Time".to_string(), "root.test.device2.temperature".to_string(), "root.other.device.temperature".to_string(), ]; assert_eq!( table_export_query_columns(&request, &context, &columns).unwrap().as_ref(), &["temperature".to_string(), "root.other.device.temperature".to_string()] ); } #[test] fn non_iotdb_table_export_preserves_columns_named_time() { let request = TableExportRequest { export_id: "export-time-column".to_string(), connection_id: "conn-1".to_string(), database: "app".to_string(), schema: None, identifier_quote: None, table_name: "samples".to_string(), file_path: "samples.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: true, batch_size: Some(25), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let columns = vec!["Time".to_string(), "value".to_string()]; for (database_type, cursor_sql, page_sql) in [ ( DatabaseType::Mysql, "SELECT `Time`, `value` FROM `samples`", "SELECT `Time`, `value` FROM `samples` LIMIT 25 OFFSET 10", ), ( DatabaseType::Postgres, "SELECT \"Time\", \"value\" FROM \"samples\"", "SELECT \"Time\", \"value\" FROM \"samples\" LIMIT 25 OFFSET 10", ), ] { let context = table_export_sql_context(database_type, None, None); let query_columns = table_export_query_columns(&request, &context, &columns).unwrap(); assert!(matches!(query_columns, Cow::Borrowed(_))); assert_eq!(table_cursor_sql(&request, &context, query_columns.as_ref(), &[], &[]), cursor_sql); assert_eq!( table_page_sql(&request, &context, query_columns.as_ref(), &[], &[], false, &[], 10, 25), page_sql ); } } #[test] fn oracle_table_cursor_sql_builds_single_ordered_select() { let request = TableExportRequest { export_id: "export-1".to_string(), connection_id: "conn-1".to_string(), database: "ORCL".to_string(), schema: Some("APP".to_string()), identifier_quote: None, table_name: "events".to_string(), file_path: "events.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: Some("WHERE status = 'active'".to_string()), order_by: None, skip_count: false, batch_size: Some(500), row_limit: Some(1000), date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Oracle, None, request.schema.as_deref()); let sql = table_cursor_sql( &request, &context, &[String::from("id"), String::from("status")], &[], &[String::from("id")], ); assert_eq!( sql, "SELECT \"id\", \"status\" FROM \"APP\".\"events\" WHERE (status = 'active') ORDER BY \"id\" ASC" ); assert!(!sql.contains("OFFSET")); assert!(!sql.contains("FETCH NEXT")); assert!(!sql.contains("ROWNUM")); } #[test] fn gbase8s_table_export_uses_owner_free_informix_queries() { let request = TableExportRequest { export_id: "export-gbase8s".to_string(), connection_id: "conn-1".to_string(), database: "appdb".to_string(), schema: Some("gbasedbt".to_string()), identifier_quote: Some(String::new()), table_name: "orders".to_string(), file_path: "orders.txt".to_string(), format: "txt".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: false, batch_size: Some(50), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let columns = vec!["id".to_string(), "payload".to_string()]; let primary_keys = vec!["id".to_string()]; let context = table_export_sql_context(DatabaseType::Gbase, Some("gbase8s"), request.schema.as_deref()); assert_eq!(context.database_type, DatabaseType::Informix); assert_eq!(context.schema, None); assert_eq!( table_cursor_sql(&request, &context, &columns, &[], &primary_keys), "SELECT id, payload FROM orders ORDER BY id ASC" ); assert_eq!( table_page_sql(&request, &context, &columns, &[], &primary_keys, false, &[], 100, 50), "SELECT SKIP 100 FIRST 50 id, payload FROM orders ORDER BY id" ); assert_eq!( table_page_sql(&request, &context, &columns, &[], &primary_keys, true, &[json!(10)], 0, 50), "SELECT FIRST 50 id, payload FROM orders WHERE id > 10 ORDER BY id ASC" ); assert_eq!( count_sql_with_where_and_identifier_quote( &request.table_name, context.schema.unwrap_or(""), &context.database_type, None, None, request.identifier_quote.as_deref(), ), "SELECT COUNT(*) FROM orders" ); let regular_gbase = table_export_sql_context(DatabaseType::Gbase, Some("gbase8a"), request.schema.as_deref()); assert_eq!(regular_gbase.database_type, DatabaseType::Gbase); assert_eq!(regular_gbase.schema, Some("gbasedbt")); assert_eq!( table_page_sql(&request, ®ular_gbase, &columns, &[], &primary_keys, false, &[], 100, 50), "SELECT \"id\", \"payload\" FROM \"gbasedbt\".\"orders\" ORDER BY \"id\" LIMIT 50 OFFSET 100" ); let informix = table_export_sql_context(DatabaseType::Informix, None, request.schema.as_deref()); assert_eq!(informix.database_type, DatabaseType::Informix); assert_eq!(informix.schema, Some("gbasedbt")); assert_eq!( table_page_sql(&request, &informix, &columns, &[], &primary_keys, false, &[], 0, 50), "SELECT FIRST 50 id, payload FROM gbasedbt.orders ORDER BY id" ); } #[test] fn gaussdb_m_table_export_uses_backticks_across_all_query_paths() { let request = TableExportRequest { export_id: "export-gaussdb-m".to_string(), connection_id: "conn-1".to_string(), database: "app".to_string(), schema: Some("app_schema".to_string()), identifier_quote: Some("`".to_string()), table_name: "order".to_string(), file_path: "order.csv".to_string(), format: "csv".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: false, batch_size: Some(100), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let columns = vec!["id".to_string(), "DisplayName".to_string()]; let primary_keys = vec!["id".to_string()]; let context = table_export_sql_context(DatabaseType::Gaussdb, None, request.schema.as_deref()); assert_eq!( table_cursor_sql(&request, &context, &columns, &[], &primary_keys), "SELECT id, `DisplayName` FROM app_schema.`order` ORDER BY id ASC" ); assert_eq!( table_page_sql(&request, &context, &columns, &[], &primary_keys, false, &[], 100, 100), "SELECT id, `DisplayName` FROM app_schema.`order` ORDER BY id LIMIT 100 OFFSET 100" ); assert_eq!( table_page_sql(&request, &context, &columns, &[], &primary_keys, true, &[json!(10)], 0, 100,), "SELECT id, `DisplayName` FROM app_schema.`order` WHERE id > 10 ORDER BY id ASC LIMIT 100" ); assert_eq!( count_sql_with_where_and_identifier_quote( &request.table_name, request.schema.as_deref().unwrap(), &DatabaseType::Gaussdb, None, None, request.identifier_quote.as_deref(), ), "SELECT COUNT(*) FROM app_schema.`order`" ); } #[test] fn mysql_sql_table_export_selects_spatial_columns_as_wkb_markers() { let request = TableExportRequest { export_id: "export-mysql-spatial".to_string(), connection_id: "conn-1".to_string(), database: "app".to_string(), schema: None, identifier_quote: None, table_name: "spatial_data".to_string(), file_path: "spatial_data.sql".to_string(), format: "sql".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: true, batch_size: Some(100), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let columns = vec!["id".to_string(), "geom".to_string(), "name".to_string()]; let column_types = vec![Some("int".to_string()), Some("geometry".to_string()), Some("varchar".to_string())]; let primary_keys = vec!["id".to_string()]; let context = table_export_sql_context(DatabaseType::Mysql, None, None); let cursor_sql = table_cursor_sql(&request, &context, &columns, &column_types, &primary_keys); assert!(cursor_sql.contains("ST_SRID(`geom`), ':', HEX(ST_AsWKB(`geom`))")); assert!(cursor_sql.contains("AS `geom`")); assert!(!cursor_sql.contains("SELECT `id`, `geom`, `name`")); let page_sql = table_page_sql(&request, &context, &columns, &column_types, &primary_keys, false, &[], 0, 100); assert!(page_sql.contains("ST_AsWKB(`geom`)")); assert!(page_sql.contains("LIMIT 100 OFFSET 0")); let mut csv_request = request; csv_request.format = "csv".to_string(); let csv_sql = table_cursor_sql(&csv_request, &context, &columns, &column_types, &primary_keys); assert_eq!(csv_sql, "SELECT `id`, `geom`, `name` FROM `spatial_data` ORDER BY `id` ASC"); } #[test] fn oracle_requested_export_columns_omit_synthetic_rowid_and_keep_metadata_aligned() { let columns = vec!["__DBX_ROWID".to_string(), "ID".to_string(), "NAME".to_string()]; let column_types = vec![Some("VARCHAR2".to_string()), Some("NUMBER".to_string()), Some("VARCHAR2".to_string())]; let primary_keys = vec!["__DBX_ROWID".to_string()]; let (columns, column_types, primary_keys) = resolve_requested_export_columns(DatabaseType::Oracle, &columns, Some(&column_types), Some(&primary_keys)); assert_eq!(columns, vec!["ID", "NAME"]); assert_eq!(column_types, vec![Some("NUMBER".to_string()), Some("VARCHAR2".to_string())]); assert!(primary_keys.is_empty()); let request = TableExportRequest { export_id: "export-rowid".to_string(), connection_id: "conn-1".to_string(), database: "ORCL".to_string(), schema: Some("APP".to_string()), identifier_quote: None, table_name: "USERS".to_string(), file_path: "users.sql".to_string(), format: "sql".to_string(), columns: None, column_types: None, primary_keys: None, where_input: None, order_by: None, skip_count: false, batch_size: Some(100), row_limit: None, date_time_format: None, numeric_column_right_align: false, column_comments: None, }; let context = table_export_sql_context(DatabaseType::Oracle, None, request.schema.as_deref()); let sql = table_cursor_sql(&request, &context, &columns, &[], &primary_keys); assert_eq!(sql, "SELECT \"ID\", \"NAME\" FROM \"APP\".\"USERS\""); let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { database_type: Some(DatabaseType::Oracle), identifier_quote: request.identifier_quote, schema: request.schema, table_name: Some(request.table_name), qualified_table_name: None, columns, column_types, column_extras: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![vec![json!(1), json!("Ada")]], batch_size: Some(100), }) .unwrap(); assert_eq!(statements, vec!["INSERT INTO \"APP\".\"USERS\" (\"ID\", \"NAME\") VALUES (1, 'Ada');"]); } #[test] fn requested_export_columns_preserve_regular_oracle_and_non_oracle_columns() { let oracle_columns = vec!["ROW_ID".to_string(), "NAME".to_string()]; let (resolved_oracle, _, _) = resolve_requested_export_columns(DatabaseType::Oracle, &oracle_columns, None, None); assert_eq!(resolved_oracle, oracle_columns); let mysql_columns = vec!["__DBX_ROWID".to_string(), "name".to_string()]; let (resolved_mysql, _, _) = resolve_requested_export_columns(DatabaseType::Mysql, &mysql_columns, None, None); assert_eq!(resolved_mysql, mysql_columns); } #[test] fn requested_mysql_sql_export_resolves_column_metadata_only_for_sql() { let table_columns = vec![ crate::db::ColumnInfo { name: "ID".to_string(), data_type: "int".to_string(), extra: Some("auto_increment".to_string()), ..Default::default() }, crate::db::ColumnInfo { name: "virtual_total".to_string(), data_type: "geometry".to_string(), extra: Some("VIRTUAL GENERATED".to_string()), ..Default::default() }, ]; let requested_columns = vec!["virtual_total".to_string(), "id".to_string(), "missing".to_string()]; assert!(requested_mysql_sql_export_needs_column_metadata(DatabaseType::Mysql, "SQL")); for format in ["csv", "json", "xlsx"] { assert!(!requested_mysql_sql_export_needs_column_metadata(DatabaseType::Mysql, format)); } assert!(!requested_mysql_sql_export_needs_column_metadata(DatabaseType::Postgres, "sql")); assert_eq!( resolve_requested_export_column_types( &requested_columns, &[Some("".to_string()), Some("bigint".to_string())], &table_columns, ), vec![Some("geometry".to_string()), Some("bigint".to_string()), None] ); assert_eq!( resolve_requested_export_column_extras(&requested_columns, &table_columns), vec![Some("VIRTUAL GENERATED".to_string()), Some("auto_increment".to_string()), None] ); } #[test] fn agent_table_read_unsupported_detects_old_agent_errors() { assert!(is_agent_table_read_unsupported("Agent RPC error (-1): unknown method: start_table_read")); assert!(is_agent_table_read_unsupported("Agent RPC error (-32601): Method not found")); assert!(!is_agent_table_read_unsupported("ORA-00933: SQL command not properly ended")); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_reads_all_cursor_pages() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[1,"Ada"],[2,"Grace"]],"affected_rows":0,"execution_time_ms":1,"session_id":"cursor-1","has_more":true}}\n' "$id" ;; *'"method":"fetchQueryPage"'*) echo fetchQueryPage >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[3,"Linus"]],"affected_rows":0,"execution_time_ms":1,"session_id":null,"has_more":false}}\n' "$id" ;; *'"method":"executeQuery"'*) echo executeQuery >> "$CALLS" printf '{"id":%s,"result":{"columns":["count"],"rows":[[3]],"affected_rows":0,"execution_time_ms":1}}\n' "$id" ;; *'"method":"closeQuerySession"'*) echo closeQuerySession >> "$CALLS" printf '{"id":%s,"result":{"ok":true}}\n' "$id" ;; esac"#, 2, None, false, ) .await; let progress = run_external_driver_export(&fixture).await.expect("multi-page JDBC export should succeed"); let csv = std::fs::read_to_string(&fixture.output).unwrap(); assert!(csv.contains("\"1\",\"Ada\"")); assert!(csv.contains("\"2\",\"Grace\"")); assert!(csv.contains("\"3\",\"Linus\"")); assert_eq!(csv.matches("\"Ada\"").count(), 1); assert_eq!( std::fs::read_to_string(&fixture.calls).unwrap(), "executeQuery\nexecuteQueryPage\nfetchQueryPage\n" ); assert_eq!(progress.last().and_then(|event| event.total_rows), Some(3)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Done))); assert!(fixture.state.connections.read().await.is_empty()); cleanup_external_driver_export_fixture(fixture); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_does_not_repeat_legacy_one_shot_results() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" printf '{"id":%s,"error":{"message":"Unsupported JDBC plugin method: executeQueryPage"}}\n' "$id" ;; *'"method":"executeQuery"'*) echo executeQuery >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[1,"Ada"],[2,"Grace"],[3,"Linus"]],"affected_rows":0,"execution_time_ms":1}}\n' "$id" ;; esac"#, 2, None, true, ) .await; run_external_driver_export(&fixture).await.expect("legacy JDBC export should succeed"); let csv = std::fs::read_to_string(&fixture.output).unwrap(); assert_eq!(csv.matches("\"Ada\"").count(), 1); assert_eq!(csv.matches("\"Grace\"").count(), 1); assert_eq!(csv.matches("\"Linus\"").count(), 1); assert_eq!(std::fs::read_to_string(&fixture.calls).unwrap(), "executeQueryPage\nexecuteQuery\n"); cleanup_external_driver_export_fixture(fixture); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_closes_cursor_at_row_limit() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[1,"Ada"]],"affected_rows":0,"execution_time_ms":1,"session_id":"cursor-1","has_more":true}}\n' "$id" ;; *'"method":"closeQuerySession"'*) echo closeQuerySession >> "$CALLS" printf '{"id":%s,"result":{"ok":true}}\n' "$id" ;; esac"#, 1, Some(1), true, ) .await; run_external_driver_export(&fixture).await.expect("row-limited JDBC export should succeed"); assert_eq!(std::fs::read_to_string(&fixture.calls).unwrap(), "executeQueryPage\ncloseQuerySession\n"); assert!(fixture.state.connections.read().await.is_empty()); cleanup_external_driver_export_fixture(fixture); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_closes_cursor_after_fetch_error() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[1,"Ada"],[2,"Grace"]],"affected_rows":0,"execution_time_ms":1,"session_id":"cursor-1","has_more":true}}\n' "$id" ;; *'"method":"fetchQueryPage"'*) echo fetchQueryPage >> "$CALLS" printf '{"id":%s,"error":{"message":"simulated fetch failure"}}\n' "$id" ;; *'"method":"closeQuerySession"'*) echo closeQuerySession >> "$CALLS" printf '{"id":%s,"result":{"ok":true}}\n' "$id" ;; esac"#, 2, None, true, ) .await; let error = run_external_driver_export(&fixture).await.expect_err("fetch errors must fail the export"); assert!(error.starts_with("simulated fetch failure")); assert_eq!( std::fs::read_to_string(&fixture.calls).unwrap(), "executeQueryPage\nfetchQueryPage\ncloseQuerySession\n" ); assert!(fixture.state.connections.read().await.is_empty()); cleanup_external_driver_export_fixture(fixture); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_cancels_blocked_execute() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" sleep 30 ;; esac"#, 2, None, true, ) .await; let export = Box::pin(run_external_driver_export(&fixture)); let cancel = async { wait_for_external_driver_call(&fixture.calls, "executeQueryPage").await; set_export_cancelled(&fixture.request.export_id).await; tokio::time::Instant::now() }; let (result, cancel_requested_at) = tokio::time::timeout(Duration::from_secs(7), async { tokio::join!(export, cancel) }) .await .expect("blocked JDBC execute should be interrupted promptly"); let progress = result.expect("cancelled JDBC export should complete without an error"); assert!(cancel_requested_at.elapsed() < Duration::from_secs(2)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Cancelled))); assert!(fixture.state.connections.read().await.is_empty()); clear_export_cancelled(&fixture.request.export_id).await; cleanup_external_driver_export_fixture(fixture); } #[cfg(unix)] #[tokio::test] async fn external_driver_table_export_cancels_blocked_fetch() { let fixture = external_driver_export_fixture( r#" case "$line" in *'"method":"executeQueryPage"'*) echo executeQueryPage >> "$CALLS" printf '{"id":%s,"result":{"columns":["id","name"],"rows":[[1,"Ada"],[2,"Grace"]],"affected_rows":0,"execution_time_ms":1,"session_id":"cursor-1","has_more":true}}\n' "$id" ;; *'"method":"fetchQueryPage"'*) echo fetchQueryPage >> "$CALLS" sleep 30 ;; esac"#, 2, None, true, ) .await; let export = Box::pin(run_external_driver_export(&fixture)); let cancel = async { wait_for_external_driver_call(&fixture.calls, "fetchQueryPage").await; set_export_cancelled(&fixture.request.export_id).await; tokio::time::Instant::now() }; let (result, cancel_requested_at) = tokio::time::timeout(Duration::from_secs(7), async { tokio::join!(export, cancel) }) .await .expect("blocked JDBC fetch should be interrupted promptly"); let progress = result.expect("cancelled JDBC export should complete without an error"); assert!(cancel_requested_at.elapsed() < Duration::from_secs(2)); assert!(matches!(progress.last().map(|event| &event.status), Some(ExportStatus::Cancelled))); assert!(fixture.state.connections.read().await.is_empty()); clear_export_cancelled(&fixture.request.export_id).await; cleanup_external_driver_export_fixture(fixture); } #[test] fn writes_json_row_without_allocating_object_map() { let mut out = Vec::new(); write_json_row_object( &mut out, &["id".to_string(), "name".to_string(), "missing".to_string()], &[json!(1), json!("Ada")], ) .unwrap(); assert_eq!(String::from_utf8(out).unwrap(), "{\n \"id\": 1,\n \"name\": \"Ada\"\n}"); } #[test] fn formats_csv_rows_escapes_embedded_commas_and_newlines() { let rows = vec![vec![json!("hello,world"), json!("line1\nline2")]]; let out = format_csv_rows(&rows); assert!(out.contains("\"hello,world\"")); assert!(out.contains("\"line1\nline2\"")); let records: Vec> = csv::ReaderBuilder::new() .has_headers(false) .from_reader(out.as_bytes()) .records() .map(|record| record.unwrap().iter().map(str::to_string).collect()) .collect(); assert_eq!(records, vec![vec!["hello,world".to_string(), "line1\nline2".to_string()]]); } // ----------------------------------------------------------------------- // Cancellation flow // ----------------------------------------------------------------------- #[test] fn cancellation_set_and_cleared_correctly() { let export_id = "test-cancel-1"; assert!(!poll_is_cancelled(export_id)); block_on(set_export_cancelled(export_id)); assert!(poll_is_cancelled(export_id)); block_on(clear_export_cancelled(export_id)); assert!(!poll_is_cancelled(export_id)); } #[test] fn cancellation_is_id_scoped() { let id_a = "cancel-scope-a"; let id_b = "cancel-scope-b"; block_on(set_export_cancelled(id_a)); assert!(poll_is_cancelled(id_a)); assert!(!poll_is_cancelled(id_b)); block_on(clear_export_cancelled(id_a)); } // ----------------------------------------------------------------------- // XLSX workbook integration // ----------------------------------------------------------------------- #[test] fn builds_xlsx_workbook_with_table_export_data() { let data = XlsxWorksheetData { sheet_name: Some("employees".to_string()), columns: vec!["id".to_string(), "name".to_string(), "salary".to_string()], column_types: vec![], column_comments: vec![], rows: vec![ vec![json!(1), json!("Alice"), json!(75000.50)], vec![json!(2), json!("Bob"), json!(82000)], vec![json!(3), Value::Null, json!(0)], ], numeric_column_right_align: false, }; let workbook = build_xlsx_workbook(&data).expect("XLSX build should succeed"); assert_eq!(workbook[0], 0x50, "Should be a ZIP (PK) archive"); assert_eq!(workbook[1], 0x4b); // Entries are Deflate-compressed; assert on their decompressed contents. let workbook_xml = read_zip_entry(&workbook, "xl/workbook.xml"); let sheet = read_zip_entry(&workbook, "xl/worksheets/sheet1.xml"); assert!(workbook_xml.contains("name=\"employees\"")); assert!(sheet.contains("75000.5")); assert!(sheet.contains("Alice")); } // ----------------------------------------------------------------------- // CSV header + rows (format_csv) — basic integration check // ----------------------------------------------------------------------- #[test] fn format_csv_produces_header_and_rows() { let out = format_csv( &["col1".to_string(), "col2".to_string()], &[vec![json!("a"), json!("b")], vec![json!("c"), json!("d")]], ); let lines: Vec<&str> = out.lines().collect(); assert_eq!(lines.len(), 3, "header + 2 data rows = 3 lines"); assert_eq!(lines[0], "\"col1\",\"col2\""); assert_eq!(lines[1], "\"a\",\"b\""); assert_eq!(lines[2], "\"c\",\"d\""); } // ----------------------------------------------------------------------- // Helpers for async cancellation in tests // ----------------------------------------------------------------------- fn poll_is_cancelled(export_id: &str) -> bool { block_on(is_export_cancelled(export_id)) } fn block_on(future: F) -> F::Output { tokio::runtime::Runtime::new().expect("create tokio runtime").block_on(future) } }