diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 636066d99..4858e5cb3 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -11,7 +11,7 @@ use crate::object_source_sql::build_export_object_source_sql; use crate::sql_dialect::{qualified_table_name, quote_table_identifier, uses_single_row_insert_statements}; use crate::transfer::{ format_ch_array_sql_literal, format_pg_array_sql_literal, is_identity_column_extra, quote_identifier, - selected_columns_include_identity_extras, wrap_dameng_identity_insert_sql, + quote_postgres_string_literal, selected_columns_include_identity_extras, wrap_dameng_identity_insert_sql, wrap_dameng_identity_insert_sql_for_table, }; use crate::types::ObjectSourceKind; @@ -361,7 +361,7 @@ fn format_postgres_json_export_literal(value: &Value) -> String { let text = value.as_str().map_or_else(|| value.to_string(), ToString::to_string); // PostgreSQL standard strings keep backslashes literal; JSON text needs its // own escape sequences, so only SQL-escape the surrounding string delimiter. - postgres_string_literal(&text) + quote_postgres_string_literal(&text) } fn format_postgres_vector_export_literal(value: &Value) -> String { @@ -375,7 +375,7 @@ fn format_postgres_vector_export_literal(value: &Value) -> String { Value::String(text) => text.to_string(), _ => value.to_string(), }; - postgres_string_literal(&text) + quote_postgres_string_literal(&text) } fn format_postgres_vector_export_text(arr: &[Value]) -> String { @@ -400,6 +400,7 @@ fn quote_export_sql_string(text: &str) -> String { fn quote_export_sql_string_for_database(text: &str, database_type: Option) -> String { match database_type { Some(DatabaseType::Dameng) => quote_dameng_export_sql_string(text), + Some(DatabaseType::Postgres) => quote_postgres_string_literal(text), database_type if is_mysql_compatible_export_literal_target(database_type) => { quote_mysql_compatible_export_sql_string(text) } @@ -988,10 +989,6 @@ fn postgres_sequence_qualified_name(schema: &str, sequence_name: &str) -> String } } -fn postgres_string_literal(value: &str) -> String { - format!("'{}'", value.replace('\'', "''")) -} - fn generate_postgres_sequence_create_ddl(sequence: &PostgresExportSequence, schema: &str) -> String { let qualified_name = postgres_sequence_qualified_name(schema, &sequence.name); let cycle = if sequence.cycle { "CYCLE" } else { "NO CYCLE" }; @@ -1023,7 +1020,7 @@ fn generate_postgres_sequence_setval_sql(sequence: &PostgresExportSequence, sche return None; } - let sequence_literal = postgres_string_literal(&postgres_sequence_qualified_name(schema, &sequence.name)); + let sequence_literal = quote_postgres_string_literal(&postgres_sequence_qualified_name(schema, &sequence.name)); match (sequence.owner_table.as_deref(), sequence.owner_column.as_deref()) { (Some(owner_table), Some(owner_column)) => { let owner_table = crate::transfer::qualified_table(owner_table, schema, &DatabaseType::Postgres); @@ -2439,7 +2436,7 @@ mod tests { } #[test] - fn postgres_export_inserts_keep_literal_control_characters() { + fn postgres_export_inserts_escape_control_characters() { let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { database_type: Some(DatabaseType::Postgres), schema: Some("public".to_string()), @@ -2453,7 +2450,35 @@ mod tests { }) .unwrap(); - assert_eq!(statements, vec!["INSERT INTO \"public\".\"notes\" (\"body\") VALUES ('line1\nline2\tend');"]); + assert_eq!(statements, vec!["INSERT INTO \"public\".\"notes\" (\"body\") VALUES (E'line1\\nline2\\tend');"]); + } + + #[test] + fn postgres_export_inserts_escape_quotes_and_backslashes_without_changing_plain_strings() { + let statements = build_export_insert_statements(BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Postgres), + schema: Some("public".to_string()), + table_name: Some("notes".to_string()), + qualified_table_name: None, + columns: vec!["carriage_return".to_string(), "quote".to_string(), "path".to_string(), "plain".to_string()], + column_types: vec![ + Some("text".to_string()), + Some("text".to_string()), + Some("text".to_string()), + Some("text".to_string()), + ], + column_extras: Vec::new(), + rows: vec![vec![json!("line1\rline2"), json!("O'Hara"), json!(r"C:\tmp"), json!("plain")]], + batch_size: Some(10), + }) + .unwrap(); + + assert_eq!( + statements, + vec![ + r#"INSERT INTO "public"."notes" ("carriage_return", "quote", "path", "plain") VALUES (E'line1\rline2', 'O''Hara', E'C:\\tmp', 'plain');"# + ] + ); } #[test] @@ -2474,7 +2499,7 @@ mod tests { assert_eq!( statements, vec![ - r#"INSERT INTO "public"."events" ("payload") VALUES ('{"text":"say \"hi\"","path":"C:\\tmp","quote":"O''Hara"}');"# + r#"INSERT INTO "public"."events" ("payload") VALUES (E'{"text":"say \\"hi\\"","path":"C:\\\\tmp","quote":"O''Hara"}');"# ] ); } diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 046a97e71..a8e7ab99d 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -22,8 +22,7 @@ use crate::query_result_sql::{ use crate::table_export::TableExportProgress; use crate::transfer::keyset_pagination_sql; use crate::xlsx_export::{ - finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook_with_trailing_sheets, StreamingXlsxWriter, - XlsxWorksheetData, + finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook_with_options, StreamingXlsxWriter, XlsxWorksheetData, }; use serde_json::Value; use sqlparser::ast::{ @@ -155,7 +154,14 @@ fn start_query_result_xlsx_workbook( column_types: &[String], ) -> Result, String> { let trailing_sheets = query_sql_worksheets(request); - start_streaming_xlsx_workbook_with_trailing_sheets(writer, Some("Result"), columns, column_types, &trailing_sheets) + start_streaming_xlsx_workbook_with_options( + writer, + Some("Result"), + columns, + column_types, + &trailing_sheets, + request.date_time_format.as_deref(), + ) } fn progress( diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index c4b87a6c7..542884705 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -19,7 +19,7 @@ use crate::connection::{task_client_session_id, AppState, PoolKind}; use crate::models::connection::DatabaseType; use crate::transfer::{ execute_on_pool, generate_insert_typed, generate_insert_typed_sql_batches, get_columns_for_transfer, - qualified_table, quote_identifier, + normalize_postgres_integer_literal, qualified_table, quote_identifier, }; pub const DEFAULT_PREVIEW_LIMIT: usize = 50; @@ -3113,7 +3113,18 @@ fn normalize_import_value( kingbase_oracle_mode: bool, date_time_format: Option<&str>, ) -> serde_json::Value { - normalize_import_temporal_value(value, data_type, db_type, kingbase_oracle_mode, date_time_format) + let normalized = normalize_import_temporal_value(value, data_type, db_type, kingbase_oracle_mode, date_time_format); + let integer_text = + normalized.as_str().map(str::to_owned).or_else(|| normalized.as_number().map(ToString::to_string)); + if let Some(integer_text) = integer_text + .as_deref() + .and_then(|value| normalize_postgres_integer_literal(value, db_type, data_type)) + .and_then(|value| value.parse::().ok()) + { + // Normalize before both INSERT and COPY paths; COPY does not pass through SQL literal escaping. + return serde_json::Value::Number(integer_text.into()); + } + normalized } pub fn build_import_insert_batches( @@ -7144,6 +7155,29 @@ mod tests { assert_eq!(String::from_utf8(data).unwrap(), "1\ta\\\\b\\tline\\nnext\\v\n\\N\t\\\\N\n"); } + #[test] + fn postgres_copy_normalizes_zero_fraction_integer_values_for_integer_targets() { + let plan = CompiledImportPlan { + mapped_source_indexes: vec![0, 1, 2], + target_columns: vec!["small_value".to_string(), "big_value".to_string(), "label".to_string()], + column_types: vec![Some("smallint".to_string()), Some("bigint".to_string()), Some("text".to_string())], + }; + let (sql, data) = build_postgres_copy_text_batch( + &[vec![serde_json::json!("1.0"), serde_json::json!(2.0), serde_json::json!("3.0")]], + &plan, + "numbers", + "public", + None, + ) + .unwrap(); + + assert_eq!( + sql, + "COPY \"public\".\"numbers\" (\"small_value\", \"big_value\", \"label\") FROM STDIN WITH (FORMAT text)" + ); + assert_eq!(String::from_utf8(data).unwrap(), "1\t2\t3.0\n"); + } + #[test] fn postgres_copy_eligibility_requires_plain_table_without_rls_or_rules() { assert_eq!( diff --git a/crates/dbx-core/src/temporal_format.rs b/crates/dbx-core/src/temporal_format.rs index 3d16264c7..64ac6b3c6 100644 --- a/crates/dbx-core/src/temporal_format.rs +++ b/crates/dbx-core/src/temporal_format.rs @@ -1,4 +1,4 @@ -use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime}; +use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, Timelike}; use serde_json::Value; use std::fmt::Write as _; @@ -17,6 +17,12 @@ enum ParsedTemporal { Time(NaiveTime), } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ExcelTemporalKind { + Date, + DateTime, +} + fn temporal_kind(data_type: Option<&str>) -> Option { let normalized = data_type?.trim().to_ascii_lowercase().replace(char::is_whitespace, " "); let base = normalized.split(['(', ':', ' ']).next().unwrap_or(""); @@ -122,6 +128,40 @@ fn parse_temporal(value: &str, preferred_pattern: Option<&str>) -> Option, + preferred_pattern: Option<&str>, +) -> Option<(f64, ExcelTemporalKind)> { + let kind = match temporal_kind(data_type)? { + TemporalKind::Date => ExcelTemporalKind::Date, + TemporalKind::DateTime => ExcelTemporalKind::DateTime, + // Excel cannot retain a timezone in a numeric date cell, so preserve + // timezone-bearing and time-only values as text. + TemporalKind::DateTimeWithTimeZone | TemporalKind::Time => return None, + }; + let parsed = parse_temporal(value.trim(), preferred_pattern)?; + let (date, time) = match (kind, parsed) { + (ExcelTemporalKind::Date, ParsedTemporal::Date(date)) => (date, NaiveTime::default()), + (ExcelTemporalKind::Date, ParsedTemporal::DateTime(value)) => (value.date(), value.time()), + (ExcelTemporalKind::DateTime, ParsedTemporal::Date(date)) => (date, NaiveTime::default()), + (ExcelTemporalKind::DateTime, ParsedTemporal::DateTime(value)) => (value.date(), value.time()), + (_, ParsedTemporal::Zoned(_) | ParsedTemporal::Time(_)) => return None, + }; + let excel_min_date = NaiveDate::from_ymd_opt(1900, 1, 1)?; + if date < excel_min_date { + return None; + } + let epoch = NaiveDate::from_ymd_opt(1899, 12, 31)?; + let mut serial = date.signed_duration_since(epoch).num_days() as f64; + // Excel's 1900 date system retains the historical fake 1900-02-29. + if date >= NaiveDate::from_ymd_opt(1900, 3, 1)? { + serial += 1.0; + } + let seconds = time.num_seconds_from_midnight() as f64 + f64::from(time.nanosecond()) / 1_000_000_000.0; + Some((serial + seconds / 86_400.0, kind)) +} + fn format_parsed(parsed: ParsedTemporal, pattern: &str) -> Option { let pattern = dayjs_to_chrono_pattern(pattern)?; let mut output = String::new(); diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index f90b8f0f6..4e7349bbd 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -234,6 +234,37 @@ fn quote_string_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } +pub(crate) fn quote_postgres_string_literal(value: &str) -> String { + if !value.contains('\\') && !value.chars().any(|character| character.is_ascii_control()) { + return quote_string_literal(value); + } + + let mut escaped = String::with_capacity(value.len()); + for character in value.chars() { + match character { + '\\' => escaped.push_str("\\\\"), + '\n' => escaped.push_str("\\n"), + '\r' => escaped.push_str("\\r"), + '\t' => escaped.push_str("\\t"), + '\x08' => escaped.push_str("\\b"), + '\x0c' => escaped.push_str("\\f"), + '\'' => escaped.push_str("''"), + character if character.is_ascii_control() => { + const HEX: &[u8; 16] = b"0123456789ABCDEF"; + let byte = character as u8; + escaped.push_str("\\x"); + escaped.push(HEX[(byte >> 4) as usize] as char); + escaped.push(HEX[(byte & 0x0F) as usize] as char); + } + character => escaped.push(character), + } + } + + // Escape string constants keep control characters out of the physical + // script and remain correct regardless of standard_conforming_strings. + format!("E'{escaped}'") +} + fn postgres_schema_exists_sql(schema: &str) -> String { format!("SELECT 1 FROM pg_catalog.pg_namespace WHERE nspname = {} LIMIT 1", quote_string_literal(schema)) } @@ -263,12 +294,45 @@ fn is_postgres_transfer_dialect(db_type: &DatabaseType) -> bool { matches!(db_type, DatabaseType::Postgres | DatabaseType::Kingbase) } -fn is_postgres_integer_like_type(data_type: &str) -> bool { +fn postgres_integer_bounds(data_type: &str) -> Option<(i128, i128)> { let normalized = data_type.trim().to_ascii_lowercase(); - matches!( - normalized.split(['(', ' ']).next().unwrap_or(""), - "smallint" | "integer" | "bigint" | "int2" | "int4" | "int8" - ) + match normalized.split(['(', ' ']).next().unwrap_or("") { + "smallint" | "int2" => Some((i128::from(i16::MIN), i128::from(i16::MAX))), + "integer" | "int4" => Some((i128::from(i32::MIN), i128::from(i32::MAX))), + "bigint" | "int8" => Some((i128::from(i64::MIN), i128::from(i64::MAX))), + _ => None, + } +} + +fn is_postgres_integer_like_type(data_type: &str) -> bool { + postgres_integer_bounds(data_type).is_some() +} + +pub(crate) fn normalize_postgres_integer_literal( + value: &str, + db_type: &DatabaseType, + column_type: Option<&str>, +) -> Option { + let bounds = column_type.filter(|_| is_postgres_transfer_dialect(db_type)).and_then(postgres_integer_bounds)?; + + // Excel numeric cells arrive as f64; normalize only an explicit zero fraction so real decimals, + // scientific notation, and values outside the target integer range stay untouched. + if value.bytes().any(|byte| matches!(byte, b'e' | b'E')) { + return None; + } + let (integer, fraction) = value.split_once('.')?; + if fraction.is_empty() || !fraction.bytes().all(|byte| byte == b'0') { + return None; + } + let digits = integer.strip_prefix('-').or_else(|| integer.strip_prefix('+')).unwrap_or(integer); + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + let parsed = integer.parse::().ok()?; + if parsed < bounds.0 || parsed > bounds.1 { + return None; + } + Some(integer.to_string()) } fn is_postgres_sequence_default(default_value: Option<&str>) -> bool { @@ -1050,17 +1114,25 @@ pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, colum } } }, - serde_json::Value::Number(n) => match db_type { - DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => { - if column_type.is_some_and(is_mysql_bit_type) { - format!("b'{}'", n) - } else { - n.to_string() - } + serde_json::Value::Number(n) => { + if let Some(integer_literal) = normalize_postgres_integer_literal(&n.to_string(), db_type, column_type) { + return integer_literal; } - _ => n.to_string(), - }, + match db_type { + DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => { + if column_type.is_some_and(is_mysql_bit_type) { + format!("b'{}'", n) + } else { + n.to_string() + } + } + _ => n.to_string(), + } + } serde_json::Value::String(s) => { + if let Some(integer_literal) = normalize_postgres_integer_literal(s, db_type, column_type) { + return integer_literal; + } if let Some(binary_literal) = format_postgres_binary_sql_literal(s, db_type, column_type) { return binary_literal; } @@ -1075,6 +1147,9 @@ pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, colum } let literal = format_literal_string(s, db_type, column_type); + if *db_type == DatabaseType::Postgres { + return quote_postgres_string_literal(&literal); + } let escaped = if is_postgres_family_target(db_type) { literal.replace('\'', "''") } else { @@ -6261,7 +6336,7 @@ mod tests { assert_eq!( sql, r#"INSERT INTO "public"."events" ("payload") VALUES -('{"message":"hello\nworld"}')"# +(E'{"message":"hello\\nworld"}')"# ); } @@ -6279,7 +6354,7 @@ mod tests { assert_eq!( sql, r#"INSERT INTO "public"."files" ("path") VALUES -('C:\tmp\file.txt')"# +(E'C:\\tmp\\file.txt')"# ); } @@ -6360,6 +6435,80 @@ mod tests { ); } + #[test] + fn postgres_insert_escapes_control_characters_quotes_and_backslashes() { + let sql = generate_insert_typed( + &[ + String::from("line_break"), + String::from("carriage_return"), + String::from("quote"), + String::from("path"), + String::from("plain"), + ], + &[ + Some(String::from("text")), + Some(String::from("text")), + Some(String::from("text")), + Some(String::from("text")), + Some(String::from("text")), + ], + &[vec![json!("line1\nline2"), json!("line1\rline2"), json!("O'Hara"), json!(r"C:\tmp"), json!("plain")]], + "notes", + "public", + &DatabaseType::Postgres, + ); + + assert_eq!( + sql, + r#"INSERT INTO "public"."notes" ("line_break", "carriage_return", "quote", "path", "plain") VALUES +(E'line1\nline2', E'line1\rline2', 'O''Hara', E'C:\\tmp', 'plain')"# + ); + } + + #[test] + fn postgres_insert_formats_whole_number_values_as_integer_literals() { + let sql = generate_insert_typed( + &[String::from("small_value"), String::from("integer_value"), String::from("big_value")], + &[Some(String::from("smallint")), Some(String::from("integer")), Some(String::from("bigint"))], + &[vec![json!(1.0), json!(-2.0), json!(3.0)]], + "numbers", + "public", + &DatabaseType::Postgres, + ); + + assert_eq!( + sql, + "INSERT INTO \"public\".\"numbers\" (\"small_value\", \"integer_value\", \"big_value\") VALUES\n(1, -2, 3)" + ); + } + + #[test] + fn postgres_integer_literal_normalization_preserves_non_whole_and_out_of_range_values() { + let scientific: serde_json::Value = serde_json::from_str("1e20").unwrap(); + let out_of_range: serde_json::Value = serde_json::from_str("9223372036854775808.0").unwrap(); + + assert_eq!(escape_value_typed(&json!("1.0"), &DatabaseType::Postgres, Some("smallint")), "1"); + assert_eq!(escape_value_typed(&json!("-2.0"), &DatabaseType::Postgres, Some("integer")), "-2"); + assert_eq!(escape_value_typed(&json!("3.0"), &DatabaseType::Postgres, Some("bigint")), "3"); + assert_eq!(escape_value_typed(&json!(1.25), &DatabaseType::Postgres, Some("smallint")), "1.25"); + assert_eq!(escape_value_typed(&json!("1.25"), &DatabaseType::Postgres, Some("smallint")), "'1.25'"); + assert_eq!(escape_value_typed(&scientific, &DatabaseType::Postgres, Some("bigint")), "1e+20"); + assert_eq!(escape_value_typed(&json!("1e3"), &DatabaseType::Postgres, Some("bigint")), "'1e3'"); + assert_eq!(escape_value_typed(&out_of_range, &DatabaseType::Postgres, Some("bigint")), "9223372036854775808.0"); + assert_eq!( + escape_value_typed(&json!("9223372036854775808.0"), &DatabaseType::Postgres, Some("bigint")), + "'9223372036854775808.0'" + ); + assert_eq!(escape_value_typed(&json!("32768.0"), &DatabaseType::Postgres, Some("smallint")), "'32768.0'"); + assert_eq!( + escape_value_typed(&json!("2147483648.0"), &DatabaseType::Postgres, Some("integer")), + "'2147483648.0'" + ); + assert_eq!(escape_value_typed(&json!("1.0"), &DatabaseType::Postgres, Some("text")), "'1.0'"); + assert_eq!(escape_value_typed(&json!("1.0"), &DatabaseType::Postgres, Some("numeric")), "'1.0'"); + assert_eq!(escape_value_typed(&json!("1.0"), &DatabaseType::Postgres, None), "'1.0'"); + } + #[test] fn oracle_single_row_insert_keeps_values_shape() { let sql = generate_insert_typed( diff --git a/crates/dbx-core/src/xlsx_export.rs b/crates/dbx-core/src/xlsx_export.rs index a58a69bf4..8b3ba2954 100644 --- a/crates/dbx-core/src/xlsx_export.rs +++ b/crates/dbx-core/src/xlsx_export.rs @@ -2,6 +2,11 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::io::{Cursor, Seek, Write}; +use crate::temporal_format::{excel_temporal_serial, ExcelTemporalKind}; + +const XLSX_DATE_STYLE: usize = 2; +const XLSX_DATETIME_STYLE: usize = 3; + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct XlsxWorksheetData { @@ -21,6 +26,7 @@ pub struct StreamingXlsxWriter { column_types: Vec, next_row_number: usize, trailing_sheets: Vec, + date_time_format: Option, } /// Estimate column widths from header names only (used by the streaming path @@ -53,13 +59,25 @@ pub(crate) fn header_row_xml(columns: &[String]) -> String { ) } -/// Build a single `` XML fragment for a data row. -pub(crate) fn data_row_xml(row_number: usize, columns: &[String], column_types: &[String], row: &[Value]) -> String { +fn data_row_xml_with_date_time_format( + row_number: usize, + columns: &[String], + column_types: &[String], + row: &[Value], + date_time_format: Option<&str>, +) -> String { let cells = columns .iter() .enumerate() .map(|(col_index, _)| { - typed_cell_xml(row.get(col_index), column_types.get(col_index), row_number - 1, col_index, None) + typed_cell_xml( + row.get(col_index), + column_types.get(col_index), + row_number - 1, + col_index, + None, + date_time_format, + ) }) .collect::(); format!("{cells}") @@ -87,15 +105,27 @@ pub(crate) fn start_streaming_xlsx_workbook( columns: &[String], column_types: &[String], ) -> Result, String> { - start_streaming_xlsx_workbook_with_trailing_sheets(writer, sheet_name, columns, column_types, &[]) + start_streaming_xlsx_workbook_with_options(writer, sheet_name, columns, column_types, &[], None) } +#[cfg(test)] pub(crate) fn start_streaming_xlsx_workbook_with_trailing_sheets( writer: W, sheet_name: Option<&str>, columns: &[String], column_types: &[String], trailing_sheets: &[XlsxWorksheetData], +) -> Result, String> { + start_streaming_xlsx_workbook_with_options(writer, sheet_name, columns, column_types, trailing_sheets, None) +} + +pub(crate) fn start_streaming_xlsx_workbook_with_options( + writer: W, + sheet_name: Option<&str>, + columns: &[String], + column_types: &[String], + trailing_sheets: &[XlsxWorksheetData], + date_time_format: Option<&str>, ) -> Result, String> { let primary_sheet = XlsxWorksheetData { sheet_name: sheet_name.map(str::to_string), @@ -113,7 +143,7 @@ pub(crate) fn start_streaming_xlsx_workbook_with_trailing_sheets StreamingXlsxWriter { /// Append a single data row to the worksheet. pub fn write_row(&mut self, row: &[Value]) -> Result<(), String> { self.zip - .write_all(data_row_xml(self.next_row_number, &self.columns, &self.column_types, row).as_bytes()) + .write_all( + data_row_xml_with_date_time_format( + self.next_row_number, + &self.columns, + &self.column_types, + row, + self.date_time_format.as_deref(), + ) + .as_bytes(), + ) .map_err(|err| err.to_string())?; self.next_row_number += 1; Ok(()) @@ -345,7 +385,20 @@ fn typed_cell_xml( row_index: usize, col_index: usize, style: Option, + date_time_format: Option<&str>, ) -> String { + if let Some(Value::String(value)) = value { + if let Some((serial, temporal_kind)) = + excel_temporal_serial(value, column_type.map(String::as_str), date_time_format) + { + let reference = cell_ref(row_index, col_index); + let style = match temporal_kind { + ExcelTemporalKind::Date => XLSX_DATE_STYLE, + ExcelTemporalKind::DateTime => XLSX_DATETIME_STYLE, + }; + return format!("{serial}"); + } + } if is_numeric_column_type(column_type) { if let Some(Value::String(value)) = value { if let Some(number) = safe_excel_number(value) { @@ -391,7 +444,14 @@ fn worksheet_xml(data: &XlsxWorksheetData) -> String { .iter() .enumerate() .map(|(col_index, _)| { - typed_cell_xml(row.get(col_index), data.column_types.get(col_index), excel_row - 1, col_index, None) + typed_cell_xml( + row.get(col_index), + data.column_types.get(col_index), + excel_row - 1, + col_index, + None, + None, + ) }) .collect::(); format!("{cells}") @@ -490,17 +550,96 @@ fn workbook_rels_xml_for_sheet_count(sheet_count: usize) -> String { ) } -fn styles_xml() -> &'static str { - concat!( - "", - "", - "", - "", - "", - "", - "", - "", - "" +fn dayjs_to_excel_number_format(pattern: &str) -> Option<(String, bool, bool)> { + let pattern = pattern.trim(); + if pattern.is_empty() || pattern.len() > 100 || pattern.contains('%') { + return None; + } + let tokens = [ + ("YYYY", "yyyy", true, false), + ("SSS", "000", false, true), + ("ZZ", "", false, true), + ("MM", "mm", true, false), + ("DD", "dd", true, false), + ("HH", "hh", false, true), + ("mm", "mm", false, true), + ("ss", "ss", false, true), + ("M", "m", true, false), + ("D", "d", true, false), + ("H", "h", false, true), + ("m", "m", false, true), + ("s", "s", false, true), + ("Z", "", false, true), + ]; + let mut output = String::with_capacity(pattern.len()); + let mut has_date = false; + let mut has_time = false; + let mut index = 0; + while index < pattern.len() { + let remaining = &pattern[index..]; + if remaining.starts_with('[') { + let close = remaining.find(']')?; + let literal = &remaining[1..close]; + output.push('"'); + output.push_str(&literal.replace('"', "\"\"")); + output.push('"'); + index += close + 1; + continue; + } + if let Some((token, replacement, is_date, is_time)) = + tokens.iter().find(|(token, ..)| remaining.starts_with(token)) + { + if replacement.is_empty() { + return None; + } + output.push_str(replacement); + has_date |= *is_date; + has_time |= *is_time; + index += token.len(); + continue; + } + let character = remaining.chars().next()?; + if character.is_ascii_alphabetic() { + // Keep the accepted Day.js token set aligned with temporal_format.rs; + // unsupported tokens must not silently alter the exported display. + return None; + } + output.push(character); + index += character.len_utf8(); + } + (has_date && !output.is_empty()).then_some((output, has_date, has_time)) +} + +fn styles_xml(date_time_format: Option<&str>) -> String { + let default_date = "yyyy-mm-dd".to_string(); + let default_datetime = "yyyy-mm-dd hh:mm:ss".to_string(); + let (date_format, datetime_format) = date_time_format + .and_then(dayjs_to_excel_number_format) + .map(|(format, has_date, has_time)| { + if has_time { + (default_date.clone(), format) + } else if has_date { + (format.clone(), format) + } else { + (default_date.clone(), default_datetime.clone()) + } + }) + .unwrap_or((default_date, default_datetime)); + format!( + concat!( + "", + "", + "", + "", + "", + "", + "", + "", + "", + "" + ), + escape_xml(&date_format), + escape_xml(&datetime_format) ) } @@ -535,7 +674,7 @@ pub fn build_xlsx_workbook_multi(sheets: &[XlsxWorksheetData]) -> Result ("_rels/.rels", root_rels_xml().to_string()), ("xl/workbook.xml", workbook_xml_for_sheets(&sheet_names)), ("xl/_rels/workbook.xml.rels", workbook_rels_xml_for_sheet_count(sheets.len())), - ("xl/styles.xml", styles_xml().to_string()), + ("xl/styles.xml", styles_xml(None)), ]; let cursor = Cursor::new(Vec::::new()); @@ -559,7 +698,8 @@ pub fn build_xlsx_workbook_multi(sheets: &[XlsxWorksheetData]) -> Result mod tests { use super::{ build_xlsx_workbook, build_xlsx_workbook_multi, start_streaming_xlsx_workbook, - start_streaming_xlsx_workbook_with_trailing_sheets, XlsxWorksheetData, + start_streaming_xlsx_workbook_with_options, start_streaming_xlsx_workbook_with_trailing_sheets, + XlsxWorksheetData, }; use calamine::{open_workbook_auto, Reader}; use serde_json::json; @@ -630,6 +770,49 @@ mod tests { assert!(sheet.contains("00123")); } + #[test] + fn writes_temporal_columns_as_excel_dates_without_retyping_other_values() { + let workbook = build_xlsx_workbook(&XlsxWorksheetData { + sheet_name: Some("Typed values".to_string()), + columns: vec![ + "day".to_string(), + "created_at".to_string(), + "label".to_string(), + "invalid_day".to_string(), + "amount".to_string(), + "zoned_at".to_string(), + ], + column_types: vec![ + "date".to_string(), + "timestamp without time zone".to_string(), + "text".to_string(), + "date".to_string(), + "numeric".to_string(), + "timestamp with time zone".to_string(), + ], + rows: vec![vec![ + json!("2024-02-25"), + json!("2024-02-25 13:02:15"), + json!("2024-02-25"), + json!("not-a-date"), + json!("2800.000000"), + json!("2024-02-25T13:02:15+08:00"), + ]], + }) + .expect("build workbook"); + + let sheet = read_zip_entry(&workbook, "xl/worksheets/sheet1.xml"); + let styles = read_zip_entry(&workbook, "xl/styles.xml"); + assert!(sheet.contains("45347")); + assert!(sheet.contains("45347.543229166666")); + assert!(sheet.contains("2024-02-25")); + assert!(sheet.contains("not-a-date")); + assert!(sheet.contains("2800.000000")); + assert!(sheet.contains("2024-02-25T13:02:15+08:00")); + assert!(styles.contains("numFmtId=\"164\" formatCode=\"yyyy-mm-dd\"")); + assert!(styles.contains("numFmtId=\"165\" formatCode=\"yyyy-mm-dd hh:mm:ss\"")); + } + #[test] fn writes_mysql_57_numeric_strings_as_numeric_cells() { let workbook = build_xlsx_workbook(&XlsxWorksheetData { @@ -782,6 +965,34 @@ mod tests { let _ = fs::remove_file(&path); } + #[test] + fn streaming_temporal_cells_keep_the_configured_excel_display_format() { + let path = std::env::temp_dir().join(format!("dbx-temporal-stream-test-{}.xlsx", uuid::Uuid::new_v4())); + { + let file = fs::File::create(&path).expect("create temp xlsx"); + let columns = ["created_at".to_string()]; + let column_types = ["timestamp without time zone".to_string()]; + let mut writer = start_streaming_xlsx_workbook_with_options( + file, + Some("Temporal"), + &columns, + &column_types, + &[], + Some("YYYY/MM/DD HH:mm:ss.SSS"), + ) + .expect("start workbook"); + writer.write_row(&[json!("2024/02/25 13:02:15.125")]).expect("write temporal row"); + drop(writer.finish().expect("finish workbook")); + } + + let bytes = fs::read(&path).expect("read workbook"); + let sheet = read_zip_entry(&bytes, "xl/worksheets/sheet1.xml"); + let styles = read_zip_entry(&bytes, "xl/styles.xml"); + assert!(sheet.contains(""), "sheet={sheet}"); + assert!(styles.contains("numFmtId=\"165\" formatCode=\"yyyy/mm/dd hh:mm:ss.000\"")); + let _ = fs::remove_file(&path); + } + #[test] fn streams_xlsx_rows_with_a_trailing_sql_worksheet() { let path = std::env::temp_dir().join(format!("dbx-stream-sql-test-{}.xlsx", uuid::Uuid::new_v4())); diff --git a/crates/dbx-core/tests/live_postgres_query_result_export.rs b/crates/dbx-core/tests/live_postgres_query_result_export.rs index 5e4810a3d..1ab5bc58b 100644 --- a/crates/dbx-core/tests/live_postgres_query_result_export.rs +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -1,3 +1,4 @@ +use std::io::Read; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::{Duration, Instant}; @@ -156,6 +157,75 @@ async fn live_postgres_query_result_export_uses_single_streamed_query() { assert_eq!(csv.lines().count(), 2051, "unexpected csv row count"); } +#[tokio::test] +#[ignore = "requires DBX_LIVE_POSTGRES_HOST/PORT/USER/PASSWORD/DATABASE pointing at a writable PostgreSQL database"] +async fn live_postgres_query_result_xlsx_preserves_temporal_cell_types() { + let host = std::env::var("DBX_LIVE_POSTGRES_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()); + let port = std::env::var("DBX_LIVE_POSTGRES_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(5432); + let user = std::env::var("DBX_LIVE_POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string()); + let password = std::env::var("DBX_LIVE_POSTGRES_PASSWORD").unwrap_or_default(); + let database = std::env::var("DBX_LIVE_POSTGRES_DATABASE").unwrap_or_else(|_| "postgres".to_string()); + let url = format!("postgresql://{user}:{password}@{host}:{port}/{database}"); + let setup_pool = postgres::connect(&url, Duration::from_secs(10)).await.expect("connect PostgreSQL"); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let schema = format!("dbx_xlsx_temporal_{}", &suffix[..8]); + let setup = vec![ + format!("CREATE SCHEMA \"{schema}\""), + format!("CREATE TABLE \"{schema}\".events (day date, created_at timestamp without time zone, label text)"), + format!("INSERT INTO \"{schema}\".events VALUES ('2024-02-25', '2024-02-25 13:02:15', '2024-02-25')"), + ]; + let cleanup = vec![format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE")]; + let _ = postgres::execute_batch(&setup_pool, &cleanup).await; + postgres::execute_batch(&setup_pool, &setup).await.expect("create temporal export fixture"); + + let dir = std::env::temp_dir().join(format!("dbx-live-postgres-xlsx-temporal-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "live-postgres-xlsx-temporal"; + let config = live_postgres_config(connection_id, &host, port, &user, &password, &database); + state.configs.write().await.insert(config.id.clone(), config); + + let file_path = dir.join("result.xlsx"); + let request = QueryResultExportRequest { + export_id: format!("live-postgres-xlsx-temporal-{suffix}"), + connection_id: connection_id.to_string(), + database: database.clone(), + schema: Some(schema.clone()), + sql: format!("SELECT day, created_at, label FROM \"{schema}\".events"), + query_base_sql: format!("SELECT day, created_at, label FROM \"{schema}\".events"), + setup_sql: Vec::new(), + database_type: DatabaseType::Postgres, + use_agent_cursor: false, + file_path: file_path.to_string_lossy().to_string(), + format: "xlsx".to_string(), + include_sql_sheet: false, + page_size: 100, + row_limit: None, + total_rows: None, + timeout_secs: Some(30), + keyset_optimization_enabled: false, + client_session_id: Some(format!("live-postgres-xlsx-temporal-{suffix}")), + execution_id: Some(format!("live-postgres-xlsx-temporal-{suffix}")), + date_time_format: None, + }; + + export_query_result_core(&state, &request, None, |_| {}).await.expect("export temporal XLSX"); + + let workbook = std::fs::read(&file_path).unwrap(); + let mut archive = zip::ZipArchive::new(std::io::Cursor::new(workbook)).expect("open generated XLSX"); + let mut sheet = String::new(); + archive.by_name("xl/worksheets/sheet1.xml").unwrap().read_to_string(&mut sheet).unwrap(); + assert!(sheet.contains(""), "sheet={sheet}"); + assert!(sheet.contains(""), "sheet={sheet}"); + assert!(sheet.contains("2024-02-25"), "sheet={sheet}"); + + let cleanup_result = postgres::execute_batch(&setup_pool, &cleanup).await; + let _ = std::fs::remove_dir_all(&dir); + cleanup_result.expect("cleanup temporal export fixture"); +} + #[tokio::test] #[ignore = "requires DBX_LIVE_POSTGRES_* env vars for temporary-table CSV/XLSX export"] async fn live_postgres_truncated_batch_result_export_replays_safe_temp_setup() {