From 88dccd7cded9080cae4f193ea384272582bd5632 Mon Sep 17 00:00:00 2001 From: zipg Date: Sat, 26 Sep 2026 06:06:54 +0800 Subject: [PATCH] fix(import): stream JSON table import and drop the 100 MB cap --- crates/dbx-core/src/data/table_import.rs | 512 ++++++++++++++++++++++- 1 file changed, 490 insertions(+), 22 deletions(-) diff --git a/crates/dbx-core/src/data/table_import.rs b/crates/dbx-core/src/data/table_import.rs index 239bf1f58a..1a6b01590a 100644 --- a/crates/dbx-core/src/data/table_import.rs +++ b/crates/dbx-core/src/data/table_import.rs @@ -14,7 +14,7 @@ use calamine::{ use chrono::{DateTime, NaiveDate, NaiveDateTime}; use quick_xml::events::{BytesStart, Event}; use quick_xml::Reader as XmlReader; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; use sha2::{Digest, Sha256}; use sqlparser::ast::{ DataType, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Insert, ObjectName, ObjectNamePart, @@ -36,11 +36,23 @@ use crate::transfer::{ pub const DEFAULT_PREVIEW_LIMIT: usize = 50; pub const DEFAULT_BATCH_SIZE: usize = 500; pub const CREATE_TABLE_INFERENCE_ROWS: usize = 100; -/// `.sql` 脚本已改为流式解析，不再受体积上限约束；该上限只保留给 JSON 等 -/// 仍然需要整份物化的格式。 +/// 分隔文本、`.sql` 脚本与 JSON 都已改为流式解析，不再受体积上限约束； +/// 该上限只保留给仍然需要整份物化的 Excel（xlsx）。 pub const MAX_NON_STREAMING_IMPORT_BYTES: u64 = 100 * 1024 * 1024; pub const MAX_LEGACY_XLS_IMPORT_BYTES: u64 = 50 * 1024 * 1024; const IMPORT_PROGRESS_INTERVAL: Duration = Duration::from_millis(100); + +// JSON 导入的行形状既决定列清单也决定错误文案，内存版与流式版共用同一批常量， +// 避免两条解析路径对同一种输入给出不同提示。 +const JSON_IMPORT_MUST_BE_OBJECT_OR_ARRAY: &str = "JSON import must be an object or an array"; +const JSON_IMPORT_NO_ROWS: &str = "Import file has no rows"; +const JSON_IMPORT_NO_COLUMNS: &str = "Import file has no columns"; +const JSON_IMPORT_OBJECT_ROWS_REQUIRED: &str = + "JSON import is configured for object rows, but at least one row is not an object"; +const JSON_IMPORT_ARRAY_ROWS_REQUIRED: &str = + "JSON import is configured for array rows, but at least one row is not an array"; +const JSON_IMPORT_MIXED_ROW_SHAPES: &str = + "JSON rows must all be objects or all be arrays; mixed row shapes are not supported"; // Keep preview parsing bounded even when an XLSX dimension declares a huge sparse range. const MAX_FAST_PREVIEW_CELLS: usize = 100_000; // Shared strings stay in memory for small workbooks and spill to an indexed temp file for large ones. @@ -1065,10 +1077,10 @@ pub fn parse_json_bytes_with_options( let items = match value { serde_json::Value::Array(items) => items, serde_json::Value::Object(_) => vec![value], - _ => return Err("JSON import must be an object or an array".to_string()), + _ => return Err(JSON_IMPORT_MUST_BE_OBJECT_OR_ARRAY.to_string()), }; if items.is_empty() { - return Err("Import file has no rows".to_string()); + return Err(JSON_IMPORT_NO_ROWS.to_string()); } let shape = options.json_shape.unwrap_or(TableImportJsonShape::Auto); @@ -1076,10 +1088,10 @@ pub fn parse_json_bytes_with_options( let all_arrays = items.iter().all(|item| item.is_array()); if shape == TableImportJsonShape::Objects && !all_objects { - return Err("JSON import is configured for object rows, but at least one row is not an object".to_string()); + return Err(JSON_IMPORT_OBJECT_ROWS_REQUIRED.to_string()); } if shape == TableImportJsonShape::Arrays && !all_arrays { - return Err("JSON import is configured for array rows, but at least one row is not an array".to_string()); + return Err(JSON_IMPORT_ARRAY_ROWS_REQUIRED.to_string()); } if all_objects { @@ -1094,7 +1106,7 @@ pub fn parse_json_bytes_with_options( } } if columns.is_empty() { - return Err("Import file has no columns".to_string()); + return Err(JSON_IMPORT_NO_COLUMNS.to_string()); } let rows = items .iter() @@ -1113,7 +1125,7 @@ pub fn parse_json_bytes_with_options( if all_arrays { let max_cols = items.iter().filter_map(|item| item.as_array().map(|row| row.len())).max().unwrap_or(0); if max_cols == 0 { - return Err("Import file has no columns".to_string()); + return Err(JSON_IMPORT_NO_COLUMNS.to_string()); } let columns = (0..max_cols).map(|index| format!("column_{}", index + 1)).collect::>(); let rows = items @@ -1129,13 +1141,266 @@ pub fn parse_json_bytes_with_options( return Ok(ParsedImportFile { columns, rows, total_rows: items.len(), effective_encoding: None }); } - Err("JSON rows must all be objects or all be arrays; mixed row shapes are not supported".to_string()) + Err(JSON_IMPORT_MIXED_ROW_SHAPES.to_string()) } pub fn parse_json_bytes(bytes: &[u8], preview_limit: usize) -> Result { parse_json_bytes_with_options(bytes, &TableImportParseOptions::default(), preview_limit) } +// --------------------------------------------------------------------------- +// JSON 文件导入：流式解析 +// +// 顶层数组逐元素反序列化，内存只保留列名与预览行，因此不再有 100 MB 体积上限。 +// 列清单来自全体行（对象行取字段并集、数组行取最大列数），所以预览要完整扫一遍 +// 文件；这与分隔文本预览的做法一致。 +// --------------------------------------------------------------------------- + +/// 对象行按字段名出列，数组行按位置出列（与 [`parse_json_bytes_with_options`] 一致）。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum JsonRowLayout { + Objects, + Arrays, +} + +/// 流式扫描 JSON 文件过程中得到的列清单与行统计。 +#[derive(Debug, Default)] +struct JsonStreamScan { + columns: Vec, + seen_columns: HashSet, + max_columns: usize, + total_rows: usize, + object_rows: usize, + array_rows: usize, + other_rows: usize, +} + +impl JsonStreamScan { + fn push_object_row(&mut self, object: &serde_json::Map) { + for key in object.keys() { + if self.seen_columns.insert(key.clone()) { + self.columns.push(key.clone()); + } + } + } +} + +/// 逐元素流式读取 JSON 文件的顶层数组（或顶层对象）并把每个元素交给 `visit`。 +/// +/// 顶层既不是数组也不是对象（标量文件）时返回与内存版一致的错误。 +fn visit_json_file_rows(path: &str, bytes_read: Option>, visit: &mut F) -> Result<(), String> +where + F: FnMut(serde_json::Value) -> Result<(), String>, +{ + let mut file = File::open(path).map_err(|error| error.to_string())?; + let mut prefix = [0u8; 3]; + let prefix_len = file.read(&mut prefix).map_err(|error| error.to_string())?; + if prefix_len != 3 || prefix != [0xEF, 0xBB, 0xBF] { + file.seek(SeekFrom::Start(0)).map_err(|error| error.to_string())?; + } + let reader = CountingJsonReader { inner: BufReader::with_capacity(256 * 1024, file), bytes_read }; + let mut deserializer = serde_json::Deserializer::from_reader(reader); + // serde_json 会给自定义错误补上「at line 1 column N」，直接透出会让流式与 + // 内存两条路径对同一种输入给出不同提示；这里单独记下原始文案。 + let mut failure: Option = None; + let parsed = { + let visitor = JsonTopLevelVisitor { visit, failure: &mut failure }; + deserializer.deserialize_any(visitor) + }; + if let Err(error) = parsed { + return Err(failure.unwrap_or_else(|| error.to_string())); + } + // 与 `serde_json::from_slice` 保持一致：结尾出现多余内容时报错而不是忽略。 + deserializer.end().map_err(|error| error.to_string()) +} + +/// 计数读取器：JSON 流式导入的字节进度按读取量上报。 +struct CountingJsonReader { + inner: R, + bytes_read: Option>, +} + +impl IoRead for CountingJsonReader { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + let read = self.inner.read(buffer)?; + if let Some(bytes_read) = self.bytes_read.as_ref() { + bytes_read.fetch_add(read as u64, Ordering::Relaxed); + } + Ok(read) + } +} + +/// 顶层值的 visitor：数组逐元素回调；对象整体作为一行回调；标量直接报错。 +struct JsonTopLevelVisitor<'a, F> { + visit: &'a mut F, + failure: &'a mut Option, +} + +impl JsonTopLevelVisitor<'_, F> { + fn reject_scalar(self) -> Result<(), E> { + *self.failure = Some(JSON_IMPORT_MUST_BE_OBJECT_OR_ARRAY.to_string()); + Err(E::custom(JSON_IMPORT_MUST_BE_OBJECT_OR_ARRAY)) + } + + /// 回调自身返回的错误：记下原始文案再中断解析，避免 serde_json 追加位置后缀。 + fn reject_row(failure: &mut Option, message: String) -> Result<(), E> { + *failure = Some(message.clone()); + Err(E::custom(message)) + } +} + +impl<'de, 'a, F> serde::de::Visitor<'de> for JsonTopLevelVisitor<'a, F> +where + F: FnMut(serde_json::Value) -> Result<(), String>, +{ + type Value = (); + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("a JSON object or an array") + } + + fn visit_seq>(self, mut seq: A) -> Result<(), A::Error> { + let Self { visit, failure } = self; + while let Some(value) = seq.next_element::()? { + if let Err(message) = visit(value) { + return Self::reject_row::(failure, message); + } + } + Ok(()) + } + + fn visit_map>(self, mut map: A) -> Result<(), A::Error> { + let Self { visit, failure } = self; + let mut object = serde_json::Map::new(); + while let Some(key) = map.next_key::()? { + let value = map.next_value::()?; + object.insert(key, value); + } + if let Err(message) = visit(serde_json::Value::Object(object)) { + return Self::reject_row::(failure, message); + } + Ok(()) + } + + fn visit_unit(self) -> Result<(), E> { + self.reject_scalar() + } + + fn visit_bool(self, _value: bool) -> Result<(), E> { + self.reject_scalar() + } + + fn visit_i64(self, _value: i64) -> Result<(), E> { + self.reject_scalar() + } + + fn visit_u64(self, _value: u64) -> Result<(), E> { + self.reject_scalar() + } + + fn visit_f64(self, _value: f64) -> Result<(), E> { + self.reject_scalar() + } + + fn visit_str(self, _value: &str) -> Result<(), E> { + self.reject_scalar() + } +} + +/// 单遍扫描 JSON 文件：校验行形状、汇总列清单，并保留前 `preview_limit` 行原始元素。 +fn scan_json_file_rows( + path: &str, + options: &TableImportParseOptions, + preview_limit: usize, + bytes_read: Option>, +) -> Result<(JsonStreamScan, Vec), String> { + let shape = options.json_shape.unwrap_or(TableImportJsonShape::Auto); + let mut scan = JsonStreamScan::default(); + let mut preview_rows: Vec = Vec::new(); + visit_json_file_rows(path, bytes_read, &mut |value| { + scan.total_rows += 1; + if let Some(object) = value.as_object() { + scan.object_rows += 1; + if shape == TableImportJsonShape::Arrays { + return Err(JSON_IMPORT_ARRAY_ROWS_REQUIRED.to_string()); + } + scan.push_object_row(object); + } else if let Some(array) = value.as_array() { + scan.array_rows += 1; + if shape == TableImportJsonShape::Objects { + return Err(JSON_IMPORT_OBJECT_ROWS_REQUIRED.to_string()); + } + scan.max_columns = scan.max_columns.max(array.len()); + } else { + scan.other_rows += 1; + match shape { + TableImportJsonShape::Objects => return Err(JSON_IMPORT_OBJECT_ROWS_REQUIRED.to_string()), + TableImportJsonShape::Arrays => return Err(JSON_IMPORT_ARRAY_ROWS_REQUIRED.to_string()), + TableImportJsonShape::Auto => {} + } + } + if preview_rows.len() < preview_limit { + preview_rows.push(value); + } + Ok(()) + })?; + Ok((scan, preview_rows)) +} + +/// 由扫描结果推导行布局与列清单，错误文案与内存版逐条对齐。 +fn json_scan_layout_and_columns(scan: &JsonStreamScan) -> Result<(JsonRowLayout, Vec), String> { + if scan.total_rows == 0 { + return Err(JSON_IMPORT_NO_ROWS.to_string()); + } + let only_objects = scan.object_rows > 0 && scan.array_rows == 0 && scan.other_rows == 0; + let only_arrays = scan.array_rows > 0 && scan.object_rows == 0 && scan.other_rows == 0; + if only_objects { + if scan.columns.is_empty() { + return Err(JSON_IMPORT_NO_COLUMNS.to_string()); + } + return Ok((JsonRowLayout::Objects, scan.columns.clone())); + } + if only_arrays { + if scan.max_columns == 0 { + return Err(JSON_IMPORT_NO_COLUMNS.to_string()); + } + let columns = (0..scan.max_columns).map(|index| format!("column_{}", index + 1)).collect(); + return Ok((JsonRowLayout::Arrays, columns)); + } + Err(JSON_IMPORT_MIXED_ROW_SHAPES.to_string()) +} + +/// 把一行原始 JSON 元素投影到最终列清单（缺失补 `null`，多余截断）。 +fn project_json_row(value: &serde_json::Value, columns: &[String], layout: JsonRowLayout) -> Vec { + match layout { + JsonRowLayout::Objects => { + let object = value.as_object(); + columns + .iter() + .map(|column| object.and_then(|object| object.get(column)).cloned().unwrap_or(serde_json::Value::Null)) + .collect() + } + JsonRowLayout::Arrays => { + let array = value.as_array(); + (0..columns.len()) + .map(|index| array.and_then(|row| row.get(index)).cloned().unwrap_or(serde_json::Value::Null)) + .collect() + } + } +} + +/// 流式解析 JSON 文件（预览与建表推断共用）。 +fn parse_json_file_streaming( + path: &str, + options: &TableImportParseOptions, + preview_limit: usize, +) -> Result { + let (scan, preview_rows) = scan_json_file_rows(path, options, preview_limit, None)?; + let (layout, columns) = json_scan_layout_and_columns(&scan)?; + let rows = preview_rows.iter().map(|value| project_json_row(value, &columns, layout)).collect::>(); + Ok(ParsedImportFile { columns, rows, total_rows: scan.total_rows, effective_encoding: None }) +} + // --------------------------------------------------------------------------- // SQL 脚本导入：从 .sql 文件中提取 INSERT / REPLACE 语句的数据行 // --------------------------------------------------------------------------- @@ -1744,14 +2009,89 @@ impl SqlImportRowStream { } } +/// 流式 JSON 行来源：先整扫一遍拿到列清单与总行数，再由第二个阻塞线程按批推送行。 +/// +/// 列清单来自全体行（字段并集 / 最大列数），所以第一次扫描必须读完整个文件； +/// 但两次扫描都只保留当前批次的若干行，内存不再随文件体积增长。 +struct JsonImportRowStream { + receiver: tokio::sync::mpsc::Receiver>, String>>, + columns: Vec, + total_rows: usize, + bytes_read: Arc, +} + +impl JsonImportRowStream { + async fn open(path: &str, options: &TableImportParseOptions, batch_size: usize) -> Result { + let bytes_read = Arc::new(AtomicU64::new(0)); + let scan_path = path.to_string(); + let scan_options = options.clone(); + let scan_bytes_read = bytes_read.clone(); + let (scan, _) = tokio::task::spawn_blocking(move || { + scan_json_file_rows(&scan_path, &scan_options, 0, Some(scan_bytes_read)) + }) + .await + .map_err(|error| error.to_string())??; + let (layout, columns) = json_scan_layout_and_columns(&scan)?; + + let (sender, receiver) = tokio::sync::mpsc::channel::>, String>>(2); + let producer_path = path.to_string(); + let producer_columns = columns.clone(); + let producer_bytes_read = bytes_read.clone(); + tokio::task::spawn_blocking(move || { + stream_json_rows_to_channel( + &producer_path, + &producer_columns, + layout, + batch_size, + sender, + producer_bytes_read, + ); + }); + + Ok(Self { receiver, columns, total_rows: scan.total_rows, bytes_read }) + } +} + +/// 第二个扫描遍次：按批把行送进有界通道，消费端跟不上时不会无限堆积内存。 +fn stream_json_rows_to_channel( + path: &str, + columns: &[String], + layout: JsonRowLayout, + batch_size: usize, + sender: tokio::sync::mpsc::Sender>, String>>, + bytes_read: Arc, +) { + let batch_size = batch_size.max(1); + let mut pending: Vec> = Vec::with_capacity(batch_size); + let result = visit_json_file_rows(path, Some(bytes_read), &mut |value| { + pending.push(project_json_row(&value, columns, layout)); + if pending.len() >= batch_size { + sender + .blocking_send(Ok(std::mem::take(&mut pending))) + .map_err(|_| "JSON import consumer closed before the stream finished".to_string())?; + pending = Vec::with_capacity(batch_size); + } + Ok(()) + }); + match result { + Ok(()) if pending.is_empty() => {} + Ok(()) => { + let _ = sender.blocking_send(Ok(pending)); + } + Err(error) => { + let _ = sender.blocking_send(Err(error)); + } + } +} + /// 表导入的行来源。 /// -/// 分隔文本、JSON、Excel 仍然先把行解析进内存；`.sql` 脚本改用 -/// [`SqlImportRowStream`] 增量产出，内存不再随文件体积增长，因此 `.sql` -/// 不再有 100 MB 的体积上限。 +/// 分隔文本与 Excel 仍然先把行解析进内存；`.sql` 脚本与 JSON 改用流式行来源 +/// 增量产出，内存不再随文件体积增长，因此二者都没有 100 MB 的体积上限。 enum ImportRowSource { Materialized { columns: Vec, rows: std::vec::IntoIter>, total_rows: usize }, Sql { stream: Box, bytes_read: Arc, pending: Option>> }, + Json { stream: Box }, } impl ImportRowSource { @@ -1772,6 +2112,10 @@ impl ImportRowSource { let pending = stream.next_batch(first_batch_rows).await?; return Ok(Self::Sql { stream: Box::new(stream), bytes_read, pending }); } + if source_format == TableImportSourceFormat::Json { + let stream = JsonImportRowStream::open(file_path, parse_options, first_batch_rows).await?; + return Ok(Self::Json { stream: Box::new(stream) }); + } let parsed = parse_import_file_with_options_and_text_columns( file_path, Some(source_format), @@ -1789,6 +2133,7 @@ impl ImportRowSource { Self::Sql { stream, .. } => { stream.columns().ok_or_else(|| "No INSERT statements found in SQL file".to_string()) } + Self::Json { stream, .. } => Ok(stream.columns.clone()), } } @@ -1796,12 +2141,16 @@ impl ImportRowSource { match self { Self::Materialized { total_rows, .. } => *total_rows, Self::Sql { stream, .. } => stream.total_rows(), + Self::Json { stream, .. } => stream.total_rows, } } - /// 是否在写入前就已知全部行数。SQL 脚本只有在扫描到 EOF 后才知道总行数。 + /// 是否在写入前就已知全部行数。 + /// + /// SQL 脚本只有在扫描到 EOF 后才知道总行数；JSON 的第一次扫描已经读完整个 + /// 文件，因此在写入前就可以给出确定的总行数。 fn total_rows_known(&self) -> bool { - matches!(self, Self::Materialized { .. }) + matches!(self, Self::Materialized { .. } | Self::Json { .. }) } /// 已读取的源字节数，用于大脚本的按字节进度；物化来源在解析阶段已经读完。 @@ -1809,6 +2158,7 @@ impl ImportRowSource { match self { Self::Materialized { .. } => total_bytes, Self::Sql { bytes_read, .. } => bytes_read.load(Ordering::Relaxed).min(total_bytes), + Self::Json { stream, .. } => stream.bytes_read.load(Ordering::Relaxed).min(total_bytes), } } @@ -1824,6 +2174,7 @@ impl ImportRowSource { } stream.next_batch(max_rows).await } + Self::Json { stream, .. } => stream.receiver.recv().await.transpose(), } } } @@ -3693,8 +4044,8 @@ pub fn parse_xlsx_file(path: &str, preview_limit: usize) -> Result Result<(), String> { - // 分隔文本与 SQL 脚本都是流式解析，内存占用不随文件体积增长，没有体积上限。 - if format.is_delimited() || format == TableImportSourceFormat::Sql { + // 分隔文本、SQL 脚本与 JSON 都是流式解析，内存占用不随文件体积增长，没有体积上限。 + if format.is_delimited() || format == TableImportSourceFormat::Sql || format == TableImportSourceFormat::Json { return Ok(()); } let metadata = std::fs::metadata(path).map_err(|e| e.to_string())?; @@ -3744,8 +4095,11 @@ async fn parse_import_file_with_options_and_text_columns( .map_err(|e| e.to_string())? } TableImportSourceFormat::Json => { - let bytes = tokio::fs::read(path).await.map_err(|e| e.to_string())?; - parse_json_bytes_with_options(&bytes, options, preview_limit) + let path = path.to_string(); + let options = options.clone(); + tokio::task::spawn_blocking(move || parse_json_file_streaming(&path, &options, preview_limit)) + .await + .map_err(|e| e.to_string())? } TableImportSourceFormat::Sql => { // 大脚本走增量解析：只保留前 `preview_limit` 行，内存不随体积增长。 @@ -8786,6 +9140,112 @@ mod tests { assert!(error.contains("configured for object rows")); } + fn json_streaming_fixture(tag: &str, body: &[u8]) -> String { + let path = std::env::temp_dir().join(format!("dbx-table-import-json-{tag}-{}.json", uuid::Uuid::new_v4())); + std::fs::write(&path, body).unwrap(); + path.to_string_lossy().to_string() + } + + /// 流式解析必须与内存版给出完全一致的列清单、行数与行内容。 + #[test] + fn streaming_json_file_parse_matches_the_in_memory_parser() { + let fixtures: [String; 7] = [ + r#"[{"id":1,"name":"Ada"},{"id":2,"active":true}]"#.to_string(), + r#"[["id","name"],[1,"Ada"],[2]]"#.to_string(), + r#"{"id":1,"nested":{"a":[1,2]},"list":[true,null]}"#.to_string(), + "[{\"id\":1,\"name\":\"Ada\"}]".to_string(), + r#"[{"名称":"中文"},{"名称":"emoji 😀"}]"#.to_string(), + r#"[{"n":1.5},{"n":-2},{"n":9007199254740993}]"#.to_string(), + r#"[{"id":1},{"id":2},{"id":3},{"id":4},{"id":5}]"#.to_string(), + ]; + for (index, text) in fixtures.iter().enumerate() { + let mut body = text.as_bytes().to_vec(); + if index == 3 { + body.splice(0..0, [0xEF, 0xBB, 0xBF]); + } + let path = json_streaming_fixture("match", &body); + // 预览行数小于总行数时，两条路径都必须按“全体行推列、只保留前 N 行”处理。 + let expected = parse_json_bytes_with_options(&body, &TableImportParseOptions::default(), 2).unwrap(); + let parsed = parse_json_file_streaming(&path, &TableImportParseOptions::default(), 2).unwrap(); + + assert_eq!(parsed.columns, expected.columns, "columns for {text}"); + assert_eq!(parsed.total_rows, expected.total_rows, "rows for {text}"); + assert_eq!(parsed.rows, expected.rows, "preview for {text}"); + assert_eq!(parsed.effective_encoding, None); + let _ = std::fs::remove_file(path); + } + } + + /// 流式路径的错误文案与内存版逐条对齐，避免同一种文件给出两种提示。 + #[test] + fn streaming_json_file_parse_reports_the_same_errors_as_the_in_memory_parser() { + let cases: [(&[u8], &str); 6] = [ + (b"[]", JSON_IMPORT_NO_ROWS), + (b"{}", JSON_IMPORT_NO_COLUMNS), + (b"[[],[]]", JSON_IMPORT_NO_COLUMNS), + (b"[1,2]", JSON_IMPORT_MIXED_ROW_SHAPES), + (br#"[{"id":1},[1]]"#, JSON_IMPORT_MIXED_ROW_SHAPES), + (b"null", JSON_IMPORT_MUST_BE_OBJECT_OR_ARRAY), + ]; + for (body, needle) in cases { + let path = json_streaming_fixture("error", body); + let expected = parse_json_bytes_with_options(body, &TableImportParseOptions::default(), 10).unwrap_err(); + let error = parse_json_file_streaming(&path, &TableImportParseOptions::default(), 10).unwrap_err(); + + assert!(expected.contains(needle), "{expected} for {}", String::from_utf8_lossy(body)); + assert_eq!(error, expected, "for {}", String::from_utf8_lossy(body)); + let _ = std::fs::remove_file(path); + } + } + + /// 语法错误（截断、多余内容）在流式路径上同样报错，且提示与内存版一致。 + #[test] + fn streaming_json_file_parse_rejects_malformed_input() { + for body in [&b"[{\"id\":1}]{}"[..], b"[{\"id\":1}", b"", b"[{\"id\":1},]"] { + let path = json_streaming_fixture("invalid", body); + let expected = parse_json_bytes_with_options(body, &TableImportParseOptions::default(), 10).unwrap_err(); + let error = parse_json_file_streaming(&path, &TableImportParseOptions::default(), 10).unwrap_err(); + + assert_eq!(error, expected, "for {}", String::from_utf8_lossy(body)); + let _ = std::fs::remove_file(path); + } + } + + #[test] + fn streaming_json_file_parse_honors_the_configured_row_shape() { + let options = TableImportParseOptions { + json_shape: Some(TableImportJsonShape::Arrays), + ..TableImportParseOptions::default() + }; + let path = json_streaming_fixture("shape", br#"[{"id":1}]"#); + let error = parse_json_file_streaming(&path, &options, 10).unwrap_err(); + + assert_eq!(error, JSON_IMPORT_ARRAY_ROWS_REQUIRED); + let _ = std::fs::remove_file(path); + } + + /// 导入行来源按批产出，不把整个 JSON 文件的行同时留在内存里。 + #[tokio::test] + async fn streaming_json_import_source_batches_rows() { + let body = br#"[{"id":1,"name":"Ada"},{"id":2},{"id":3,"name":"Grace"},{"id":4,"name":"Linus"}]"#; + let path = json_streaming_fixture("stream", body); + let expected = parse_json_bytes_with_options(body, &TableImportParseOptions::default(), 10).unwrap(); + + let mut stream = JsonImportRowStream::open(&path, &TableImportParseOptions::default(), 2).await.unwrap(); + assert_eq!(stream.columns, expected.columns); + assert_eq!(stream.total_rows, expected.total_rows); + + let mut rows = Vec::new(); + let mut batch_sizes = Vec::new(); + while let Some(batch) = stream.receiver.recv().await.transpose().unwrap() { + batch_sizes.push(batch.len()); + rows.extend(batch); + } + assert_eq!(batch_sizes, vec![2, 2]); + assert_eq!(rows, expected.rows); + let _ = std::fs::remove_file(path); + } + fn sql_import_options(dialect: DatabaseType) -> TableImportParseOptions { TableImportParseOptions { sql_dialect: Some(dialect), ..TableImportParseOptions::default() } } @@ -9117,17 +9577,25 @@ mod tests { let _ = std::fs::remove_file(path); } - /// `.sql` 已经改成流式解析，不应再受 100 MB 非流式上限约束。 + /// `.sql` 与 JSON 都已经改成流式解析，不应再受 100 MB 非流式上限约束。 #[test] - fn sql_import_is_not_size_capped() { + fn sql_and_json_imports_are_not_size_capped() { let path = std::env::temp_dir().join(format!("dbx-table-import-limit-{}.sql", uuid::Uuid::new_v4())); let file = File::create(&path).unwrap(); file.set_len(MAX_NON_STREAMING_IMPORT_BYTES + 1).unwrap(); drop(file); ensure_non_streaming_file_size(&path.to_string_lossy(), TableImportSourceFormat::Sql).unwrap(); - ensure_non_streaming_file_size(&path.to_string_lossy(), TableImportSourceFormat::Json).unwrap_err(); + ensure_non_streaming_file_size(&path.to_string_lossy(), TableImportSourceFormat::Json).unwrap(); + let xlsx_path = std::env::temp_dir().join(format!("dbx-table-import-limit-{}.xlsx", uuid::Uuid::new_v4())); + let xlsx_file = File::create(&xlsx_path).unwrap(); + xlsx_file.set_len(MAX_NON_STREAMING_IMPORT_BYTES + 1).unwrap(); + drop(xlsx_file); + let error = + ensure_non_streaming_file_size(&xlsx_path.to_string_lossy(), TableImportSourceFormat::Excel).unwrap_err(); + assert!(error.contains("File too large for excel import"), "{error}"); let _ = std::fs::remove_file(path); + let _ = std::fs::remove_file(xlsx_path); } /// 流式 `.sql` 在读完整份脚本之前，`total_rows` 只是已解析的部分行数。