fix(postgres): preserve exported and imported values
This commit is contained in:
t8y2
parent
47d48a09e9
commit
634f604cd6
|
|
@ -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<DatabaseType>) -> 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"}');"#
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<W: Write + Seek>(
|
|||
column_types: &[String],
|
||||
) -> Result<StreamingXlsxWriter<W>, 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(
|
||||
|
|
|
|||
|
|
@ -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::<i64>().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!(
|
||||
|
|
|
|||
|
|
@ -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<TemporalKind> {
|
||||
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<Parsed
|
|||
.or_else(|| parse_known_temporal(value))
|
||||
}
|
||||
|
||||
pub(crate) fn excel_temporal_serial(
|
||||
value: &str,
|
||||
data_type: Option<&str>,
|
||||
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<String> {
|
||||
let pattern = dayjs_to_chrono_pattern(pattern)?;
|
||||
let mut output = String::new();
|
||||
|
|
|
|||
|
|
@ -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<String> {
|
||||
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::<i128>().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(
|
||||
|
|
|
|||
|
|
@ -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<W: Write + Seek> {
|
|||
column_types: Vec<String>,
|
||||
next_row_number: usize,
|
||||
trailing_sheets: Vec<XlsxWorksheetData>,
|
||||
date_time_format: Option<String>,
|
||||
}
|
||||
|
||||
/// 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 `<row>` 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::<String>();
|
||||
format!("<row r=\"{row_number}\">{cells}</row>")
|
||||
|
|
@ -87,15 +105,27 @@ pub(crate) fn start_streaming_xlsx_workbook<W: Write + Seek>(
|
|||
columns: &[String],
|
||||
column_types: &[String],
|
||||
) -> Result<StreamingXlsxWriter<W>, 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<W: Write + Seek>(
|
||||
writer: W,
|
||||
sheet_name: Option<&str>,
|
||||
columns: &[String],
|
||||
column_types: &[String],
|
||||
trailing_sheets: &[XlsxWorksheetData],
|
||||
) -> Result<StreamingXlsxWriter<W>, 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<W: Write + Seek>(
|
||||
writer: W,
|
||||
sheet_name: Option<&str>,
|
||||
columns: &[String],
|
||||
column_types: &[String],
|
||||
trailing_sheets: &[XlsxWorksheetData],
|
||||
date_time_format: Option<&str>,
|
||||
) -> Result<StreamingXlsxWriter<W>, 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<W: Write + Seek
|
|||
write_zip_entry(&mut zip, "_rels/.rels", root_rels_xml())?;
|
||||
write_zip_entry(&mut zip, "xl/workbook.xml", &workbook_xml_for_sheets(&sheet_names))?;
|
||||
write_zip_entry(&mut zip, "xl/_rels/workbook.xml.rels", &workbook_rels_xml_for_sheet_count(sheet_count))?;
|
||||
write_zip_entry(&mut zip, "xl/styles.xml", styles_xml())?;
|
||||
write_zip_entry(&mut zip, "xl/styles.xml", &styles_xml(date_time_format))?;
|
||||
|
||||
// Begin the sheet1.xml entry with header, frozen pane, column widths and
|
||||
// the header row.
|
||||
|
|
@ -142,6 +172,7 @@ pub(crate) fn start_streaming_xlsx_workbook_with_trailing_sheets<W: Write + Seek
|
|||
column_types: column_types.to_vec(),
|
||||
next_row_number: 2,
|
||||
trailing_sheets: trailing_sheets.to_vec(),
|
||||
date_time_format: date_time_format.map(str::to_string),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -149,7 +180,16 @@ impl<W: Write + Seek> StreamingXlsxWriter<W> {
|
|||
/// 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<usize>,
|
||||
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!("<c r=\"{reference}\" s=\"{style}\"><v>{serial}</v></c>");
|
||||
}
|
||||
}
|
||||
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::<String>();
|
||||
format!("<row r=\"{excel_row}\">{cells}</row>")
|
||||
|
|
@ -490,17 +550,96 @@ fn workbook_rels_xml_for_sheet_count(sheet_count: usize) -> String {
|
|||
)
|
||||
}
|
||||
|
||||
fn styles_xml() -> &'static str {
|
||||
concat!(
|
||||
"<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>",
|
||||
"<styleSheet xmlns=\"http://schemas.openxmlformats.org/spreadsheetml/2006/main\">",
|
||||
"<fonts count=\"2\"><font><sz val=\"11\"/><name val=\"Calibri\"/></font><font><b/><sz val=\"11\"/><name val=\"Calibri\"/></font></fonts>",
|
||||
"<fills count=\"2\"><fill><patternFill patternType=\"none\"/></fill><fill><patternFill patternType=\"gray125\"/></fill></fills>",
|
||||
"<borders count=\"1\"><border><left/><right/><top/><bottom/><diagonal/></border></borders>",
|
||||
"<cellStyleXfs count=\"1\"><xf numFmtId=\"0\" fontId=\"0\" fillId=\"0\" borderId=\"0\"/></cellStyleXfs>",
|
||||
"<cellXfs count=\"2\"><xf numFmtId=\"0\" fontId=\"0\" fillId=\"0\" borderId=\"0\" xfId=\"0\"/><xf numFmtId=\"0\" fontId=\"1\" fillId=\"0\" borderId=\"0\" xfId=\"0\" applyFont=\"1\"/></cellXfs>",
|
||||
"<cellStyles count=\"1\"><cellStyle name=\"Normal\" xfId=\"0\" builtinId=\"0\"/></cellStyles>",
|
||||
"</styleSheet>"
|
||||
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!(
|
||||
"<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>",
|
||||
"<styleSheet xmlns=\"http://schemas.openxmlformats.org/spreadsheetml/2006/main\">",
|
||||
"<numFmts count=\"2\"><numFmt numFmtId=\"164\" formatCode=\"{}\"/><numFmt numFmtId=\"165\" formatCode=\"{}\"/></numFmts>",
|
||||
"<fonts count=\"2\"><font><sz val=\"11\"/><name val=\"Calibri\"/></font><font><b/><sz val=\"11\"/><name val=\"Calibri\"/></font></fonts>",
|
||||
"<fills count=\"2\"><fill><patternFill patternType=\"none\"/></fill><fill><patternFill patternType=\"gray125\"/></fill></fills>",
|
||||
"<borders count=\"1\"><border><left/><right/><top/><bottom/><diagonal/></border></borders>",
|
||||
"<cellStyleXfs count=\"1\"><xf numFmtId=\"0\" fontId=\"0\" fillId=\"0\" borderId=\"0\"/></cellStyleXfs>",
|
||||
"<cellXfs count=\"4\"><xf numFmtId=\"0\" fontId=\"0\" fillId=\"0\" borderId=\"0\" xfId=\"0\"/><xf numFmtId=\"0\" fontId=\"1\" fillId=\"0\" borderId=\"0\" xfId=\"0\" applyFont=\"1\"/><xf numFmtId=\"164\" fontId=\"0\" fillId=\"0\" borderId=\"0\" xfId=\"0\" applyNumberFormat=\"1\"/><xf numFmtId=\"165\" fontId=\"0\" fillId=\"0\" borderId=\"0\" xfId=\"0\" applyNumberFormat=\"1\"/></cellXfs>",
|
||||
"<cellStyles count=\"1\"><cellStyle name=\"Normal\" xfId=\"0\" builtinId=\"0\"/></cellStyles>",
|
||||
"</styleSheet>"
|
||||
),
|
||||
escape_xml(&date_format),
|
||||
escape_xml(&datetime_format)
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -535,7 +674,7 @@ pub fn build_xlsx_workbook_multi(sheets: &[XlsxWorksheetData]) -> Result<Vec<u8>
|
|||
("_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::<u8>::new());
|
||||
|
|
@ -559,7 +698,8 @@ pub fn build_xlsx_workbook_multi(sheets: &[XlsxWorksheetData]) -> Result<Vec<u8>
|
|||
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("<c r=\"C2\" t=\"inlineStr\"><is><t>00123</t></is></c>"));
|
||||
}
|
||||
|
||||
#[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("<c r=\"A2\" s=\"2\"><v>45347</v></c>"));
|
||||
assert!(sheet.contains("<c r=\"B2\" s=\"3\"><v>45347.543229166666</v></c>"));
|
||||
assert!(sheet.contains("<c r=\"C2\" t=\"inlineStr\"><is><t>2024-02-25</t></is></c>"));
|
||||
assert!(sheet.contains("<c r=\"D2\" t=\"inlineStr\"><is><t>not-a-date</t></is></c>"));
|
||||
assert!(sheet.contains("<c r=\"E2\"><v>2800.000000</v></c>"));
|
||||
assert!(sheet.contains("<c r=\"F2\" t=\"inlineStr\"><is><t>2024-02-25T13:02:15+08:00</t></is></c>"));
|
||||
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("<c r=\"A2\" s=\"3\"><v>"), "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()));
|
||||
|
|
|
|||
|
|
@ -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("<c r=\"A2\" s=\"2\"><v>"), "sheet={sheet}");
|
||||
assert!(sheet.contains("<c r=\"B2\" s=\"3\"><v>"), "sheet={sheet}");
|
||||
assert!(sheet.contains("<c r=\"C2\" t=\"inlineStr\"><is><t>2024-02-25</t></is></c>"), "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() {
|
||||
|
|
|
|||
Loading…
Reference in New Issue