diff --git a/Cargo.toml b/Cargo.toml index fcc72a3..ba6ca87 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,7 +39,14 @@ serde = { version = "1", features = ["derive"] } # intercepts and serde_yaml does not — so every `-o yaml` output carrying a # `Value` number turns into a nested mapping. That breaks commands unrelated to # this one, which costs more than the rounding it buys. -serde_json = "1" +# +# `raw_value` is on, for the json load's reshape (src/commands/json_rows.rs): +# it hands each row over as its own JSON text, so a row is rewritten +# byte-for-byte instead of round-tripping through `Value` — which would sort +# its keys (`BTreeMap`, since `preserve_order` is off) and round a number +# wider than i64/u64/f64. Unlike `arbitrary_precision` above, this feature +# changes nothing for code that does not name `RawValue`. +serde_json = { version = "1", features = ["raw_value"] } arrow = { version = "59", default-features = false, features = [ "ipc", # Timestamps carrying a named IANA zone ("UTC", "America/New_York") diff --git a/README.md b/README.md index e16c4db..253c544 100644 --- a/README.md +++ b/README.md @@ -42,14 +42,15 @@ PostgreSQL-dialect SQL. Everything else builds on that. ## Getting your data in -**Upload a file** directly — csv, newline-delimited json, or parquet: +**Upload a file** directly — csv, json, or parquet: ```sh hotdata databases load --catalog demo --table listings --file ./listings.csv ``` The format comes from the extension; pass `--format csv|json|parquet` when the -extension is missing or misleading. +extension is missing or misleading. json is read whatever shape it arrives in — +an array of objects, a pretty-printed document, or one object per line. A load **replaces** the table by default. `--mode append` adds rows to an existing table instead: diff --git a/skills/hotdata/SKILL.md b/skills/hotdata/SKILL.md index 3682eac..ed17605 100644 --- a/skills/hotdata/SKILL.md +++ b/skills/hotdata/SKILL.md @@ -87,7 +87,7 @@ Returns workspaces with `public_id`, `name`, `active`, `favorite`, `provision_st **Instant databases** are Hotdata-owned catalogs you create and populate yourself — no remote source to sync. Query them in SQL as **`..`**. Prefer **`hotdata databases`** for this workflow. -**File formats:** `databases tables load` accepts **csv**, **newline-delimited json**, and **parquet** (local `--file`, remote `--url`, or a pre-staged `--upload-id`). The format is read from the file's extension; `--format csv|json|parquet` overrides it, and an unrecognised extension is left to the server to resolve from the bytes. +**File formats:** `databases tables load` accepts **csv**, **json**, and **parquet** (local `--file`, remote `--url`, or a pre-staged `--upload-id`). The format is read from the file's extension; `--format csv|json|parquet` overrides it, and an unrecognised extension is left to the server to resolve from the bytes. **json is read in any shape** — an array of objects, a pretty-printed document, or one object per line: `--file`/`--url` reshape it to newline-delimited json before upload, so no `jq` step is needed. Every row must be an object, and the reshape is byte-faithful per row — key order and number digits are preserved. Two inputs are refused locally, before anything uploads: a row that is not an object, and a source with no rows at all (`[]`, or an empty file). Note the **load itself** reads json numbers as f64, whatever the shape: an integer wider than i64/u64 or a decimal past ~17 significant digits lands rounded — use parquet, or a string column, where the exact digits matter. **Active database:** `hotdata databases use ` saves the active database to config. `databases tables list`/`load`/`remove`, `databases queries`/`results`, and all `databases context` commands default to the active database; pass **`--database `** to override per-command. (`databases tables show` instead takes a fully-qualified `catalog.schema.table`.) @@ -109,7 +109,7 @@ hotdata databases attach [--database ] [--alias ] hotdata databases detach [--database ] # Preferred: load by catalog alias (server declares the table/schema if missing). -# Loads csv, newline-delimited json, or parquet — format read from the extension. +# Loads csv, json, or parquet — format read from the extension. hotdata databases load --catalog --table
[--schema public] (--file | --url | --upload-id | --result-id ) [--mode replace|append|delete|update|upsert] [--append] [--format csv|json|parquet] [--key ]... [--workspace-id ] # Also available via tables subcommand @@ -128,7 +128,7 @@ hotdata databases tables remove
[--database ] [--schema public] [--w - `unset` — clears the active database from config. - `` — inspect one database (returns id, catalog, name, expires_at; a fork also shows its `forked_from` record). - `remove` — removes the instant database; clears the active-database config if it matched. -- `load` (top-level shorthand) — loads a file into `--catalog.--schema.--table`. Accepts `--file`, `--url`, `--upload-id`, or `--result-id` (load a saved query result by id — from `hotdata databases results` or a query's `[result-id: …]` footer — instead of a file; the result must belong to the target database). **Formats:** csv, newline-delimited json (`.json`/`.jsonl`/`.ndjson`), and parquet; the format comes from the file's extension, and `--format` overrides it (needed when the extension is absent or misleading). An unrecognised extension is not rejected — the server reads the bytes. A table or schema that was never declared is declared by the server as part of the load, so no up-front `--table` is required. +- `load` (top-level shorthand) — loads a file into `--catalog.--schema.--table`. Accepts `--file`, `--url`, `--upload-id`, or `--result-id` (load a saved query result by id — from `hotdata databases results` or a query's `[result-id: …]` footer — instead of a file; the result must belong to the target database). **Formats:** csv, json (`.json`/`.jsonl`/`.ndjson`), and parquet; the format comes from the file's extension, and `--format` overrides it (needed when the extension is absent or misleading). An unrecognised extension is not rejected — the server reads the bytes, and a file that plainly opens a json array is taken as json even without an extension. A json source in any shape (array, pretty-printed, one object per line) is reshaped locally to newline-delimited json before upload; an already-newline-delimited file is uploaded untouched. A table or schema that was never declared is declared by the server as part of the load, so no up-front `--table` is required. - **Load modes** (`--mode`, default `replace`) — `replace` supersedes the table's contents; `append` adds rows; `delete`, `update`, and `upsert` match existing rows **by key**. `--append` is the old shorthand for `--mode append` and still works, but the two cannot be combined. The keyed modes need a key: declare one with `databases tables add --key`, or name it per-load with `--key` (repeat for a composite key). `delete` uploads only the key columns; `update` replaces matching rows and ignores unmatched ones; `upsert` inserts the unmatched instead. Keyed modes are not available with `--result-id`. - `tables list` — lists tables with `TABLE` (`..
`), `SYNCED`, `LAST_SYNC`. Uses active database when `--database` is omitted. - `tables add` — declares a table **with its key and storage layout**, which a load cannot infer. `--key` (repeatable) is what enables the `delete`/`update`/`upsert` load modes on that table. `--sorted-by ` or `=desc` sets sort order; `--partition-by ` partitions on the value, `=month` (or `year`/`day`/`hour`) on a calendar part — one partition per calendar month needs **both** `=year` and `=month`, or every March shares a partition. Sort and partition are fixed once the table exists. `--key-determines` (repeatable) asserts a column's value is fixed by the key: it prunes keyed loads harder, and is **correctness-affecting** — declare it only where the invariant really holds, or a keyed load can leave a duplicate key behind. Re-adding an existing table is a conflict (409), and `tables remove` does not clear the declaration — the table leaves the listing but the name stays declared and still conflicts. So **a key cannot be retrofitted onto a table declared without one**: declare it with `--key` up front, or use a new table name. @@ -146,6 +146,15 @@ hotdata databases load --catalog airbnb --table listings --url https://example.c hotdata query "SELECT count(*) FROM airbnb.public.listings" ``` +csv and json load on the same command, with nothing to convert first — +`reviews.json` may hold `[{…}, {…}]`, one object per line, or a single +pretty-printed object: + +``` +hotdata databases load --catalog airbnb --table hosts --file hosts.csv +hotdata databases load --catalog airbnb --table reviews --file reviews.json +``` + Keeping a table in sync by key. Declare the key **before the table's first load** — `tables add` on a table that already exists returns 409, and a key cannot be added afterwards: diff --git a/skills/hotdata/references/WORKFLOWS.md b/skills/hotdata/references/WORKFLOWS.md index 9c5d417..07c9f93 100644 --- a/skills/hotdata/references/WORKFLOWS.md +++ b/skills/hotdata/references/WORKFLOWS.md @@ -94,7 +94,7 @@ A `hotdata query` runs inside **one** instant database; its scope sees that data | | **Instant databases** | |---|------------------------| -| **Best for** | Files you own (csv, newline-delimited json, parquet); catalog-style `alias.schema.table` | +| **Best for** | Files you own (csv, json, parquet); catalog-style `alias.schema.table` | | **SQL prefix** | `..
` where catalog = `--catalog` alias | | **CLI** | `hotdata databases create --catalog` + `databases load` | | **Declare schema up front** | Optional — the load declares a missing table/schema. Declare with `databases tables add --key` when you need the keyed load modes; a key cannot be added later | diff --git a/src/commands.rs b/src/commands.rs index 204020a..b2820af 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -9,6 +9,7 @@ pub mod indexes; pub mod ingest; pub mod ingest_common; pub mod jobs; +pub mod json_rows; pub mod prompt; pub mod queries; pub mod query; diff --git a/src/commands/databases.rs b/src/commands/databases.rs index 2480f8b..fb28145 100644 --- a/src/commands/databases.rs +++ b/src/commands/databases.rs @@ -218,9 +218,10 @@ pub enum DatabasesCommands { #[arg(long, value_parser = ["replace", "append", "delete", "update", "upsert"])] mode: Option, - /// Format of the uploaded file: csv, json (newline-delimited), or - /// parquet. Detected from the file's extension when omitted; pass it - /// when the extension is absent or misleading. Not valid with + /// Format of the uploaded file: csv, json, or parquet. Detected from + /// the file's extension when omitted; pass it when the extension is + /// absent or misleading. json reads an array of objects, a + /// pretty-printed document, or one object per line. Not valid with /// `--result-id`, which is always parquet. #[arg(long, value_parser = ["csv", "json", "parquet"], conflicts_with = "result_id")] format: Option, @@ -444,9 +445,10 @@ pub enum DatabaseTablesCommands { #[arg(long, value_parser = ["replace", "append", "delete", "update", "upsert"])] mode: Option, - /// Format of the uploaded file: csv, json (newline-delimited), or - /// parquet. Detected from the file's extension when omitted; pass it - /// when the extension is absent or misleading. Not valid with + /// Format of the uploaded file: csv, json, or parquet. Detected from + /// the file's extension when omitted; pass it when the extension is + /// absent or misleading. json reads an array of objects, a + /// pretty-printed document, or one object per line. Not valid with /// `--result-id`, which is always parquet. #[arg(long, value_parser = ["csv", "json", "parquet"], conflicts_with = "result_id")] format: Option, @@ -1245,31 +1247,156 @@ fn upload_data_path( result } -/// Upload `path`, announcing the content type its extension implies. No -/// extension is rejected: an unrecognised one announces -/// `application/octet-stream`, which the load sniffs. -fn upload_data_file(api: &Api, path: &str) -> String { - let file_size = match std::fs::metadata(path) { - Ok(m) => m.len(), - Err(e) => { - eprintln!("error opening file '{path}': {e}"); - std::process::exit(1); +/// Upload `path`, announcing the content type its extension implies, and +/// return the upload id with the `format` the load should carry. No extension +/// is rejected: an unrecognised one announces `application/octet-stream`, +/// which the load sniffs. +/// +/// `format` is the one already resolved from `--format` and the source's own +/// extension. It comes back unchanged unless the file was a JSON shape the +/// load cannot read and was rewritten — see [`json_source_rewrite`] — in which +/// case the upload is newline-delimited json and says so. +fn upload_data_file<'a>( + api: &Api, + path: &str, + format: Option<&'a str>, +) -> (String, Option<&'a str>) { + let source = Path::new(path); + match json_source_rewrite(source, format).unwrap_or_else(|e| e.exit()) { + Some(rewritten) => ( + upload_rewritten_json(api, rewritten).unwrap_or_else(|e| e.exit()), + Some("json"), + ), + None => { + let file_size = match std::fs::metadata(path) { + Ok(m) => m.len(), + Err(e) => { + eprintln!("error opening file '{path}': {e}"); + std::process::exit(1); + } + }; + + let id = upload_data_path( + api, + source, + crate::client::sdk::content_type_for_path(path), + file_size, + ) + .unwrap_or_else(|e| e.exit()); + (id, format) } - }; + } +} - upload_data_path( - api, - Path::new(path), - crate::client::sdk::content_type_for_path(path), - file_size, - ) - .unwrap_or_else(|e| e.exit()) +/// The newline-delimited rewrite of `path` a load needs, or `None` when the +/// file can be uploaded as it stands. +/// +/// A `json` load reads one JSON value per line, and a `.json` file on disk is +/// as often an array of objects or a single pretty-printed document. Both are +/// rewritten here, so the shape inside the file is not something the caller has +/// to know or convert first. Everything else is untouched: parquet and csv are +/// not JSON, and already-newline-delimited json is what the load wants. +/// +/// Errors are [`ApiError::Transport`] so the caller can clean up a staged +/// download before surfacing them. +fn json_source_rewrite( + path: &Path, + format: Option<&str>, +) -> Result, ApiError> { + use crate::commands::json_rows; + use std::io::{BufReader, BufWriter, Read, Seek}; + + let fail = |msg: String| ApiError::Transport(msg); + // Only a json load reads rows off lines, so a source that names another + // format is left alone without being opened at all. + if matches!(format, Some(f) if f != "json") { + return Ok(None); + } + + let mut file = std::fs::File::open(path) + .map_err(|e| fail(format!("error opening file '{}': {e}", path.display())))?; + + // The shape is read off the first record, never the whole file, so a large + // newline-delimited source costs one short read to leave alone. + let mut head = Vec::new(); + Read::by_ref(&mut file) + .take(json_rows::SNIFF_BYTES as u64) + .read_to_end(&mut head) + .map_err(|e| fail(format!("error reading '{}': {e}", path.display())))?; + + if !load_reads_json(format, &head) { + return Ok(None); + } + let shape = json_rows::shape_of(&head); + if !shape.needs_rewrite() { + return Ok(None); + } + + file.rewind() + .map_err(|e| fail(format!("error reading '{}': {e}", path.display())))?; + let temp = tempfile::Builder::new() + .prefix("hotdata-upload-") + // Names the rewrite for what it holds: the upload announces the ndjson + // content type this extension implies. + .suffix(".jsonl") + .tempfile() + .map_err(|e| fail(format!("error creating a temp file: {e}")))?; + + let spinner = crate::util::spinner("Reading json..."); + let rows = json_rows::to_ndjson(shape, BufReader::new(file), BufWriter::new(temp.as_file())); + spinner.finish_and_clear(); + + match rows { + Ok(0) => Err(fail(format!( + "error: '{}' carries no json rows", + path.display() + ))), + Ok(_) => Ok(Some(temp)), + Err(msg) => Err(fail(format!( + "error reading json from '{}': {msg}", + path.display() + ))), + } +} + +/// Whether the load will read this source as JSON, given the `format` already +/// resolved from `--format` and the source's extension. +/// +/// An unresolved format sends no `format` at all and leaves the server to read +/// the bytes, which it reads as csv — a leading `[` then fails as a one-column +/// csv rather than as json. So a source that plainly opens a JSON array is +/// claimed as json here even with no extension to say so, which is the +/// `--url`-without-a-file-name case. +fn load_reads_json(format: Option<&str>, head: &[u8]) -> bool { + match format { + Some(f) => f == "json", + None => { + crate::commands::json_rows::shape_of(head) == crate::commands::json_rows::Shape::Array + } + } +} + +/// Upload a newline-delimited rewrite, announcing the ndjson content type its +/// `.jsonl` name implies, and delete it before returning — on both arms, for +/// the reason [`upload_temp_file`] documents. +fn upload_rewritten_json(api: &Api, temp: tempfile::NamedTempFile) -> Result { + upload_temp_file(temp, |path| { + let size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0); + let name = path.to_string_lossy(); + upload_data_path( + api, + path, + crate::client::sdk::content_type_for_path(&name), + size, + ) + }) } /// Download `url` and upload it, announcing the content type the URL's own /// extension implies — the staged temp file's name is generated and says -/// nothing about the bytes. -fn upload_data_url(api: &Api, url: &str) -> String { +/// nothing about the bytes. Returns the upload id and the `format` the load +/// should carry, which [`upload_data_file`] documents. +fn upload_data_url<'a>(api: &Api, url: &str, format: Option<&'a str>) -> (String, Option<&'a str>) { // The presigned upload needs a seekable, size-known source (the SDK opens // the path, declares its byte count, and PUTs it directly to storage), so // download the URL to a temp file first, then upload that file on the same @@ -1327,10 +1454,61 @@ fn upload_data_url(api: &Api, url: &str) -> String { }; dl_pb.finish_and_clear(); - let size = std::fs::metadata(temp.path()).map(|m| m.len()).unwrap_or(0); - let content_type = crate::client::sdk::content_type_for_path(url); - upload_temp_file(temp, |path| upload_data_path(api, path, content_type, size)) - .unwrap_or_else(|e| e.exit()) + let url_content_type = crate::client::sdk::content_type_for_path(url); + // The download is reshaped on the same terms as a local `--file`: a json + // array or a pretty-printed document becomes newline-delimited json before + // it goes up, and the rewrite is what gets uploaded — with the download + // deleted the moment the rewrite exists, so only one copy is on disk for + // the length of the upload. + let (source, rewritten) = download_to_upload(temp, |path| json_source_rewrite(path, format)) + .unwrap_or_else(|e| e.exit()); + + let id = upload_temp_file(source, |path| { + let size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0); + // The rewrite's own `.jsonl` name announces ndjson; an untouched + // download is announced by the URL, whose extension is authoritative + // where the staged name is only advisory. + let staged = path.to_string_lossy(); + let content_type = if rewritten { + crate::client::sdk::content_type_for_path(&staged) + } else { + url_content_type + }; + upload_data_path(api, path, content_type, size) + }) + .unwrap_or_else(|e| e.exit()); + + (id, if rewritten { Some("json") } else { format }) +} + +/// The file a staged download should upload — itself, or the rewrite that +/// replaced it — and whether it was rewritten. +/// +/// A rewrite makes the download dead weight, so it is deleted here rather than +/// held until the upload returns: a 5 GB json array would otherwise need 10 GB +/// of temp space for the length of the upload, and fail with ENOSPC on a small +/// runner. A download that needed no rewrite is still the thing to upload, so +/// it survives. Either way nothing outlives an error, for the reason +/// [`upload_temp_file`] documents: the caller exits without unwinding. +fn download_to_upload( + download: tempfile::NamedTempFile, + rewrite: R, +) -> Result<(tempfile::NamedTempFile, bool), ApiError> +where + R: FnOnce(&Path) -> Result, ApiError>, +{ + match rewrite(download.path()) { + Ok(Some(rewritten)) => { + // Deleted before the upload starts, not after it ends. + drop(download); + Ok((rewritten, true)) + } + Ok(None) => Ok((download, false)), + Err(e) => { + drop(download); + Err(e) + } + } } /// Upload an already-downloaded temp file, guaranteeing the file is deleted @@ -1341,9 +1519,9 @@ fn upload_data_url(api: &Api, url: &str) -> String { /// [`ApiError::exit`] (`std::process::exit`) on the `Err` arm, and /// `process::exit` runs no destructors. Owning `temp` in this function means it /// drops (deleting a potentially multi-GB download) before the caller can exit. -fn upload_temp_file(temp: tempfile::NamedTempFile, upload: F) -> Result +fn upload_temp_file(temp: tempfile::NamedTempFile, upload: F) -> Result where - F: FnOnce(&Path) -> Result, + F: FnOnce(&Path) -> Result, { let result = upload(temp.path()); // Delete now, while still inside this function, so cleanup precedes any @@ -1855,8 +2033,8 @@ pub fn create( "{}", format!( concat!( - "Load a table:\n", - " hotdata databases load --catalog {0} --table
--file \n", + "Load a table (csv, json, or parquet):\n", + " hotdata databases load --catalog {0} --table
--file \n", " hotdata databases load --catalog {0} --table
--url \n", "\nQuery with:\n", " hotdata query \"SELECT * FROM {0}.public.
LIMIT 10\"\n", @@ -2662,21 +2840,22 @@ pub fn tables_load( // names the format, and an unrecognised one sends nothing so the server // sniffs. A staged --upload-id carries the content type it was created // with, so it too sends only what --format asked for. + // + // A file or URL can come back with a format the upload settled rather than + // the one resolved here: a json source that was not newline-delimited is + // uploaded as a rewrite that is. let body = match (result_id, upload_id, file, url) { (Some(rid), None, None, None) => load_table_request_from_result(rid, mode, key), (None, Some(id), None, None) => load_table_request(id, mode, format, key), - (None, None, Some(path), None) => load_table_request( - &upload_data_file(&api, path), - mode, - format.or_else(|| format_for_path(path)), - key, - ), - (None, None, None, Some(u)) => load_table_request( - &upload_data_url(&api, u), - mode, - format.or_else(|| format_for_path(u)), - key, - ), + (None, None, Some(path), None) => { + let (upload, format) = + upload_data_file(&api, path, format.or_else(|| format_for_path(path))); + load_table_request(&upload, mode, format, key) + } + (None, None, None, Some(u)) => { + let (upload, format) = upload_data_url(&api, u, format.or_else(|| format_for_path(u))); + load_table_request(&upload, mode, format, key) + } (None, None, None, None) => { eprintln!( "error: one of --file , --url , --upload-id , or --result-id is required" @@ -3532,6 +3711,111 @@ mod tests { ); } + /// Write `body` to a temp file named with `suffix`, for the rewrite tests. + fn source_file(suffix: &str, body: &str) -> tempfile::NamedTempFile { + use std::io::Write; + let mut f = tempfile::Builder::new().suffix(suffix).tempfile().unwrap(); + f.write_all(body.as_bytes()).unwrap(); + f.flush().unwrap(); + f + } + + #[test] + fn a_json_array_file_is_rewritten_to_newline_delimited_json() { + // The shape a `.json` file most often carries: what an HTTP API returns + // and what `jq` writes. A `json` load reads one value per line, so the + // array is rewritten before it goes up. + let src = source_file(".json", "[{\"id\":1},{\"id\":2}]"); + let rewritten = json_source_rewrite(src.path(), Some("json")) + .unwrap() + .expect("an array is not newline-delimited, so it is rewritten"); + assert_eq!( + std::fs::read_to_string(rewritten.path()).unwrap(), + "{\"id\":1}\n{\"id\":2}\n" + ); + assert_eq!( + crate::client::sdk::content_type_for_path(&rewritten.path().to_string_lossy()), + "application/x-ndjson", + "the rewrite's own name has to announce what it holds" + ); + } + + #[test] + fn a_pretty_printed_json_file_is_rewritten() { + let src = source_file(".json", "{\n \"id\": 1\n}\n"); + let rewritten = json_source_rewrite(src.path(), Some("json")) + .unwrap() + .expect("a value spanning lines is not newline-delimited"); + assert_eq!( + std::fs::read_to_string(rewritten.path()).unwrap(), + "{\"id\":1}\n" + ); + } + + #[test] + fn newline_delimited_json_is_uploaded_untouched() { + // The fast path that keeps a large `.jsonl` from being rewritten to say + // exactly the same thing. + let src = source_file(".jsonl", "{\"id\":1}\n{\"id\":2}\n"); + assert!( + json_source_rewrite(src.path(), Some("json")) + .unwrap() + .is_none() + ); + } + + #[test] + fn a_csv_or_parquet_source_is_never_reshaped() { + // Only a json load reads rows off lines; nothing else is even sniffed + // for shape, whatever its first byte looks like. + let csv = source_file(".csv", "id,name\n1,a\n"); + assert!( + json_source_rewrite(csv.path(), Some("csv")) + .unwrap() + .is_none() + ); + let looks_like_json = source_file(".csv", "[{\"id\":1}]"); + assert!( + json_source_rewrite(looks_like_json.path(), Some("csv")) + .unwrap() + .is_none() + ); + } + + #[test] + fn an_extensionless_json_array_is_claimed_as_json() { + // With no extension the load sends no format and the server reads the + // bytes — as csv, which a leading `[` fails on. A plain array is taken + // as json here instead, and comes back rewritten. + let src = source_file("", "[{\"id\":1}]"); + assert!(load_reads_json(None, b" [{\"id\":1}]")); + assert!(json_source_rewrite(src.path(), None).unwrap().is_some()); + // Newline-delimited bytes with no extension still go to the server to + // sniff, which reads them correctly. + assert!(!load_reads_json(None, b"{\"id\":1}\n{\"id\":2}\n")); + } + + #[test] + fn a_json_source_with_no_rows_is_refused() { + let src = source_file(".json", " \n"); + let err = json_source_rewrite(src.path(), Some("json")).unwrap_err(); + assert!( + format!("{err:?}").contains("no json rows"), + "an empty json file should say so here, not fail schema inference server-side: {err:?}" + ); + } + + #[test] + fn a_json_row_that_is_not_an_object_is_refused_before_upload() { + let src = source_file(".json", "[{\"id\":1}, 7]"); + let err = json_source_rewrite(src.path(), Some("json")).unwrap_err(); + let msg = format!("{err:?}"); + assert!( + msg.contains("one object per row"), + "the error should name the shape a load needs: {msg}" + ); + } + #[test] fn format_for_path_ignores_a_url_query_string() { // A presigned storage URL carries its signature in the query, so the @@ -4038,7 +4322,7 @@ mod tests { let path = temp.path().to_path_buf(); assert!(path.exists()); - let result = upload_temp_file(temp, |p| { + let result: Result = upload_temp_file(temp, |p| { assert!(p.exists(), "file present while the upload runs"); Err(ApiError::Transport("upload boom".into())) }); @@ -4064,6 +4348,72 @@ mod tests { assert!(!path.exists(), "temp file must be removed on success"); } + /// A staged download plus the rewrite that may replace it, for the + /// lifetime tests below. + fn staged_pair() -> (tempfile::NamedTempFile, tempfile::NamedTempFile) { + ( + tempfile::Builder::new().suffix(".json").tempfile().unwrap(), + tempfile::Builder::new() + .suffix(".jsonl") + .tempfile() + .unwrap(), + ) + } + + #[test] + fn a_rewritten_download_is_deleted_before_the_upload() { + // Once the rewrite holds the rows the download is dead weight, and the + // upload is the long part: keeping both would put two copies of a + // multi-gigabyte source on disk for its whole duration (a 5 GB array + // needing 10 GB of temp, then ENOSPC on a small runner). + let (download, rewrite) = staged_pair(); + let download_path = download.path().to_path_buf(); + let rewrite_path = rewrite.path().to_path_buf(); + + let (source, rewritten) = download_to_upload(download, |_| Ok(Some(rewrite))).unwrap(); + + assert!(rewritten); + assert_eq!(source.path(), rewrite_path); + assert!( + !download_path.exists(), + "the download must be gone before the upload starts" + ); + assert!(rewrite_path.exists(), "the rewrite is what gets uploaded"); + } + + #[test] + fn a_download_that_needs_no_rewrite_is_the_upload() { + // Nothing replaced it, so it has to survive to be uploaded — the case + // that stops this from being a plain "delete the download first". + let (download, _) = staged_pair(); + let download_path = download.path().to_path_buf(); + + let (source, rewritten) = download_to_upload(download, |_| Ok(None)).unwrap(); + + assert!(!rewritten); + assert_eq!(source.path(), download_path); + assert!(download_path.exists()); + } + + #[test] + fn a_failed_reshape_deletes_the_download() { + // The caller exits on this error without unwinding, so cleanup has to + // happen here — the rule `upload_temp_file` documents. + let (download, _) = staged_pair(); + let download_path = download.path().to_path_buf(); + + let err = download_to_upload(download, |_| { + Err(ApiError::Transport("reshape boom".into())) + }) + .unwrap_err(); + + assert!(matches!(err, ApiError::Transport(_))); + assert!( + !download_path.exists(), + "temp file must be removed before the failure is returned" + ); + } + // --- `set`'s advisory existence check ----------------------------------- /// A database API token is denied `GET /v1/databases/{id}` by its diff --git a/src/commands/json_rows.rs b/src/commands/json_rows.rs new file mode 100644 index 0000000..d3b72ae --- /dev/null +++ b/src/commands/json_rows.rs @@ -0,0 +1,304 @@ +//! Reshaping a local JSON file into the newline-delimited JSON a load reads. +//! +//! A `json` load reads one JSON value per line. The JSON people actually have +//! on disk is often something else — an array of objects (what an HTTP API +//! returns, what `jq` writes), or a single pretty-printed document — and the +//! server refuses both with a parser message about the bytes ("failed to infer +//! JSON schema: EOF while parsing a list"), not about the shape. So the load +//! rewrites those shapes to newline-delimited JSON before uploading, and +//! `--file data.json` works whatever is inside it. +//! +//! Newline-delimited input is uploaded untouched: [`shape_of`] reads the first +//! record, not the file, so a large `.jsonl` is never rewritten to say the same +//! thing. +//! +//! A row is rewritten as **its own JSON text**, with only the whitespace +//! between tokens dropped — never re-serialized from a parsed value. Both +//! things a `serde_json::Value` round-trip would change are load-visible: it +//! sorts an object's keys (a `BTreeMap`, since `preserve_order` is off), which +//! reorders the columns the load infers, and it rounds a number wider than +//! i64/u64 through f64. So the rows that go up are the rows that were on disk. + +use std::io::{Read, Write}; + +use serde::de::{DeserializeSeed, Deserializer, SeqAccess, Visitor}; +use serde_json::Value; +use serde_json::value::RawValue; + +/// How many leading bytes [`shape_of`] wants to decide. +/// +/// Large enough that the first record of any plausible newline-delimited file +/// is whole within it. A first record longer than this reads as +/// [`Shape::Values`] and is rewritten — a slower path to the same rows, since a +/// stream of values is exactly what newline-delimited JSON is and a row is +/// rewritten as the text it already was. +pub const SNIFF_BYTES: usize = 64 * 1024; + +/// The JSON shape a source carries, as read off its first bytes. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum Shape { + /// One whole JSON value per line — what a `json` load already reads. + Ndjson, + /// A single JSON array whose elements are the rows. + Array, + /// Whitespace-separated JSON values: one pretty-printed document, several + /// concatenated, or records spanning lines. + Values, +} + +impl Shape { + /// Whether a load needs this shape rewritten before upload. + pub fn needs_rewrite(self) -> bool { + self != Shape::Ndjson + } +} + +/// The shape of a JSON source, read from `head` — its first [`SNIFF_BYTES`] +/// bytes, or the whole thing when it is shorter. +/// +/// Decided on two questions, in order: does it open an array, and is its first +/// line a whole value on its own? Only the second is a fast path worth +/// protecting, so an inconclusive answer errs towards rewriting. +pub fn shape_of(head: &[u8]) -> Shape { + // Lossy is right for a sniff: `head` can end mid-character when it is a + // truncated read, and a replacement char there changes no answer below. + let text = String::from_utf8_lossy(head); + let text = text.trim_start(); + if text.starts_with('[') { + return Shape::Array; + } + let first_line = text.split('\n').next().unwrap_or_default(); + if serde_json::from_str::(first_line).is_ok() { + Shape::Ndjson + } else { + Shape::Values + } +} + +/// Read every row `src` carries in `shape` and write it to `out` as one +/// compact JSON object per line. Returns the row count. +/// +/// Streams: rows are read and written one at a time, so a multi-gigabyte array +/// costs one row of memory rather than the whole document. +pub fn to_ndjson(shape: Shape, src: R, out: W) -> Result { + let mut rows = Rows { out, count: 0 }; + let mut de = serde_json::Deserializer::from_reader(src); + match shape { + // The elements of one array, read through the array rather than into it. + Shape::Array => { + (&mut rows).deserialize(&mut de).map_err(describe)?; + // Reject anything after the array's close: a file holding `[…] [b` + // is a truncated write, not two tables. + de.end().map_err(describe)?; + } + // A stream of whole values, which is what both other shapes are. + Shape::Ndjson | Shape::Values => { + for row in de.into_iter::>() { + rows.write(&row.map_err(describe)?)?; + } + } + } + rows.out.flush().map_err(|e| format!("{e}"))?; + Ok(rows.count) +} + +/// The rows written so far, and where they go. +struct Rows { + out: W, + count: u64, +} + +impl Rows { + /// Write one row as the JSON text it already is, refusing a value that is + /// not an object: a load needs named columns, and a bare scalar or list + /// names none. + fn write(&mut self, row: &RawValue) -> Result<(), String> { + let text = row.get(); + // `RawValue` holds one whole JSON value, so its first character settles + // the type without a parse. + if !text.trim_start().starts_with('{') { + return Err(format!( + "a json load reads one object per row, and this file carries {} \ + among its rows — reshape it so every row is an object", + type_of(text) + )); + } + write_compact(&mut self.out, text).map_err(|e| format!("{e}"))?; + self.out.write_all(b"\n").map_err(|e| format!("{e}"))?; + self.count += 1; + Ok(()) + } +} + +/// Write one row's JSON text as a single line: whitespace between tokens is +/// dropped, and every other byte is copied through untouched — so the row keeps +/// its key order, its number digits, and anything inside its strings. +/// +/// Only the string state has to be tracked, since a `"` is the one delimiter +/// whitespace can hide behind, and a backslash escape is the one thing that can +/// hide a `"`. +fn write_compact(out: &mut W, text: &str) -> std::io::Result<()> { + let bytes = text.as_bytes(); + let mut copy_from = 0; + let mut i = 0; + let mut in_string = false; + + while i < bytes.len() { + match bytes[i] { + b'\\' if in_string => i += 2, // the escaped byte is whatever it is + b'"' => { + in_string = !in_string; + i += 1; + } + b if !in_string && b.is_ascii_whitespace() => { + // Flush what precedes the gap, then step over the whole gap. + out.write_all(&bytes[copy_from..i])?; + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + copy_from = i; + } + _ => i += 1, + } + } + // Whatever follows the last gap. `copy_from` only ever takes a position + // inside the row or its end, so it indexes even when a trailing escape has + // carried `i` past the end — which valid JSON cannot carry anyway. + out.write_all(&bytes[copy_from..]) +} + +// Reads the elements of a top-level array without holding the array: serde +// hands each element over as it is parsed, and each is written and dropped +// before the next is read. +impl<'de, W: Write> DeserializeSeed<'de> for &mut Rows { + type Value = (); + + fn deserialize>(self, de: D) -> Result<(), D::Error> { + de.deserialize_seq(self) + } +} + +impl<'de, W: Write> Visitor<'de> for &mut Rows { + type Value = (); + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "a json array of objects") + } + + fn visit_seq>(self, mut seq: A) -> Result<(), A::Error> { + while let Some(row) = seq.next_element::>()? { + // Our own message, carried out through serde so the row's line and + // column travel with it. + self.write(&row).map_err(serde::de::Error::custom)?; + } + Ok(()) + } +} + +/// The JSON type name of a row, read off the first character of its text, for +/// the not-an-object message. +fn type_of(text: &str) -> &'static str { + match text.trim_start().as_bytes().first() { + Some(b'[') => "a list", + Some(b'"') => "a string", + Some(b't' | b'f') => "a boolean", + Some(b'n') => "a null", + Some(b'{') => "an object", + _ => "a number", + } +} + +/// A parse failure, as the CLI reports it. `serde_json` already appends the +/// line and column, including for the messages [`Rows::write`] raises. +fn describe(err: serde_json::Error) -> String { + format!("{err}") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ndjson(src: &str) -> Result<(u64, String), String> { + let shape = shape_of(src.as_bytes()); + let mut out: Vec = Vec::new(); + let rows = to_ndjson(shape, src.as_bytes(), &mut out)?; + Ok((rows, String::from_utf8(out).unwrap())) + } + + #[test] + fn newline_delimited_json_is_left_alone() { + assert_eq!( + shape_of(b"{\"id\":1}\n{\"id\":2}\n"), + Shape::Ndjson, + "a first line that is a whole value means the file is already ndjson" + ); + // A single-line record with no trailing newline is still ndjson. + assert_eq!(shape_of(b"{\"id\":1}"), Shape::Ndjson); + } + + #[test] + fn an_array_of_objects_becomes_one_row_per_line() { + let (rows, out) = + ndjson("[\n {\"id\":1,\"n\":\"a\"},\n {\"id\":2,\"n\":\"b\"}\n]\n").unwrap(); + assert_eq!(rows, 2); + assert_eq!(out, "{\"id\":1,\"n\":\"a\"}\n{\"id\":2,\"n\":\"b\"}\n"); + } + + #[test] + fn a_pretty_printed_object_becomes_one_row() { + let (rows, out) = ndjson("{\n \"id\": 1,\n \"n\": \"a\"\n}\n").unwrap(); + assert_eq!(rows, 1); + assert_eq!(out, "{\"id\":1,\"n\":\"a\"}\n"); + } + + #[test] + fn concatenated_objects_become_one_row_each() { + let (rows, out) = ndjson("{\"id\":1} {\"id\":2}{\"id\":3}").unwrap(); + assert_eq!(rows, 3); + assert_eq!(out, "{\"id\":1}\n{\"id\":2}\n{\"id\":3}\n"); + } + + #[test] + fn a_row_keeps_its_own_number_text_and_key_order() { + // The rewrite is byte-faithful per row. A round-trip through `Value` + // would not be: it sorts an object's keys (`BTreeMap`, since + // `preserve_order` is off), which reorders the columns the load infers, + // and it rounds a number wider than i64/u64 to f64 + // (123456789012345678901 comes back as 1.2345678901234568e20). + let src = r#"[{"z":1,"big":123456789012345678901,"d":0.12345678901234567890}]"#; + let (rows, out) = ndjson(src).unwrap(); + assert_eq!(rows, 1); + assert_eq!( + out, + "{\"z\":1,\"big\":123456789012345678901,\"d\":0.12345678901234567890}\n" + ); + } + + #[test] + fn compaction_drops_whitespace_between_tokens_and_nothing_else() { + // Whitespace inside a string is part of the value, and an escaped quote + // does not end the string. + let src = "[\n {\"note\": \"a b\\n c\",\n \"esc\": \"q\\\" \\\\\"}\n]"; + let (_, out) = ndjson(src).unwrap(); + assert_eq!(out, "{\"note\":\"a b\\n c\",\"esc\":\"q\\\" \\\\\"}\n"); + } + + #[test] + fn a_row_that_is_not_an_object_is_refused() { + let err = ndjson("[{\"id\":1}, 7]").unwrap_err(); + assert!( + err.contains("object"), + "the error should say a row must be an object, got: {err}" + ); + } + + #[test] + fn malformed_json_is_refused() { + assert!(ndjson("[{\"id\":1},").is_err()); + } + + #[test] + fn an_empty_source_yields_no_rows() { + assert_eq!(ndjson(" \n").unwrap().0, 0); + } +}