diff --git a/packages/python/e2e-attestations/claude-1326-stream-rows.json b/packages/python/e2e-attestations/claude-1326-stream-rows.json new file mode 100644 index 000000000..48674468b --- /dev/null +++ b/packages/python/e2e-attestations/claude-1326-stream-rows.json @@ -0,0 +1,7 @@ +{ + "command": "just test-e2e", + "ran_at": 1791241697, + "exit_code": 0, + "commit": "55ef5b88f40ba4a5568f1fed86a784caf8129cd7", + "branch": "claude/1326-stream-rows" +} diff --git a/packages/rust/changelog.d/2026-10-05-cli-stream-rows.md b/packages/rust/changelog.d/2026-10-05-cli-stream-rows.md new file mode 100644 index 000000000..27e8e494c --- /dev/null +++ b/packages/rust/changelog.d/2026-10-05-cli-stream-rows.md @@ -0,0 +1,4 @@ +**Changed** `dirsql query` writes JSON rows as SQLite yields them instead of +building every row in memory first. `SELECT path FROM './*'` over a 1M-file +directory drops from 0.81s to 0.44s, and `SELECT path FROM './**/*.md'` over a +1.6M-file tree from 1.33s to 0.80s. The bytes printed are unchanged. diff --git a/packages/rust/migrations.d/2026-10-05-cli-stream-rows.md b/packages/rust/migrations.d/2026-10-05-cli-stream-rows.md new file mode 100644 index 000000000..caa8686f8 --- /dev/null +++ b/packages/rust/migrations.d/2026-10-05-cli-stream-rows.md @@ -0,0 +1,27 @@ +### CLI: `dirsql query` streams its JSON rows + +#### Summary + +`dirsql query` builds less in memory before printing JSON. The bytes printed +are unchanged; nothing breaks. + +#### Required changes + +_None._ + +#### Deprecations removed + +_None._ + +#### Behavior changes without code changes + +_None._ + +#### Verification + +```bash +cd "$(mktemp -d)" +touch b.md a.md +dirsql query "SELECT path, path FROM './*'" +# expected: [{"path":"a.md"},{"path":"b.md"}] +``` diff --git a/packages/rust/src/cli/execute.rs b/packages/rust/src/cli/execute.rs index 2cb2cbe27..89313fe17 100644 --- a/packages/rust/src/cli/execute.rs +++ b/packages/rust/src/cli/execute.rs @@ -15,7 +15,7 @@ use serde::Deserialize; use serde_json::Value; use super::AppState; -use super::serialize::result_to_json; +use super::serialize::{JsonRows, result_to_json}; use crate::{DirSQL, DirSqlError}; /// Why a query failed, classified independently of transport. The HTTP @@ -79,6 +79,28 @@ pub async fn execute_query( } } +/// [`execute_query`] for the one-shot CLI's JSON output: the same intake and +/// errors, with the rows written as JSON while SQLite yields them. The text +/// is the `{value}\n` the CLI printed from [`execute_query`]'s array, held +/// until the statement ends so a failing query still prints no rows. +pub async fn execute_query_json( + state: &AppState, + raw_body: String, +) -> Result, QueryFailure> { + let sql = parse_sql_body(&raw_body)?; + let db = require_ready(state)?; + let join = tokio::task::spawn_blocking(move || { + db.query_each(&sql, JsonRows::new, JsonRows::row) + .map(JsonRows::finish) + }) + .await; + match join { + Ok(Ok(text)) => Ok(text), + Ok(Err(err)) => Err(classify_query_error(err)), + Err(join_err) => Err(QueryFailure::Internal(join_err.to_string())), + } +} + /// Await `task`, bounded by `timeout` when one is given ([`QueryFailure::Timeout`] /// carrying the bound on expiry) and to completion with `None`. async fn run_bounded( diff --git a/packages/rust/src/cli/run.rs b/packages/rust/src/cli/run.rs index b9db5f050..1d6f103ca 100644 --- a/packages/rust/src/cli/run.rs +++ b/packages/rust/src/cli/run.rs @@ -15,11 +15,14 @@ //! The `dirsql` binary is a shim over [`run_cli`], so `cargo install dirsql //! --features cli` and every other entry path run the same code. -use std::io::IsTerminal; +use std::io::{IsTerminal, Write}; use std::path::PathBuf; use super::{ - AppState, ServerConfig, execute::execute_query, init::InitOptions, repl::run_repl, + AppState, ServerConfig, + execute::{execute_query, execute_query_json}, + init::InitOptions, + repl::run_repl, serve_with_state, table, }; use crate::{DirSQL, Extension, Row, Table}; @@ -336,7 +339,7 @@ pub fn run_cli(argv: Vec) -> i32 { let code = runtime.block_on(async { match cli.command.take() { Some(Command::Init(args)) => run_init(args), - Some(Command::Query(args)) => run_query(args).await, + Some(Command::Query(args)) => run_query(args, &mut std::io::stdout()).await, Some(Command::Server(args)) => run_server(args).await, Some(Command::Context) => { print!("{}", super::context::guide()); @@ -359,12 +362,15 @@ pub fn run_cli(argv: Vec) -> i32 { async fn run_default(cli: Cli) -> u8 { match cli.sql { Some(sql) => { - run_query(QueryArgs { - sql, - on_file: cli.on_file, - format: cli.format, - common: cli.common, - }) + run_query( + QueryArgs { + sql, + on_file: cli.on_file, + format: cli.format, + common: cli.common, + }, + &mut std::io::stdout(), + ) .await } None => { @@ -393,11 +399,11 @@ async fn run_default(cli: Cli) -> u8 { /// One-shot `dirsql query`: build the index exactly as server mode would /// (same `load_state` / hook loading), run the SQL through the shared -/// [`execute_query`] pipeline, print the result JSON on stdout, and exit. +/// [`execute_query`] pipeline, write the result to `out`, and exit. /// Any [`QueryFailure`](super::execute::QueryFailure) prints its /// message — the same string the HTTP `{"error": …}` body carries — to /// stderr with a non-zero exit. -async fn run_query(args: QueryArgs) -> u8 { +async fn run_query(args: QueryArgs, out: &mut dyn Write) -> u8 { let parser = match resolve_on_file(&args.on_file) { Ok(parser) => parser, Err(message) => { @@ -411,9 +417,17 @@ async fn run_query(args: QueryArgs) -> u8 { // natively; only the long-lived server enforces `query_timeout` (408). let format = args.format.resolve(std::io::stdout().is_terminal()); - match execute_query(&state, query_body(&args.sql), None).await { - Ok(value) => { - print!("{}", render_rows(&value, format)); + let printed = match format { + Format::Table => execute_query(&state, query_body(&args.sql), None) + .await + .map(|value| render_rows(&value, format).into_bytes()), + _ => execute_query_json(&state, query_body(&args.sql)).await, + }; + match printed { + Ok(text) => { + // A closed pipe (`| head`) is the reader's choice, not a failure. + let _ = out.write_all(&text); + let _ = out.flush(); 0 } Err(failure) => { @@ -664,6 +678,43 @@ mod tests { assert_eq!(query_body(" "), r#"{"sql":" "}"#); } + fn query_args(argv: &[&str]) -> QueryArgs { + match Cli::parse_from(argv).command { + Some(Command::Query(args)) => args, + other => panic!("expected a query subcommand, got {other:?}"), + } + } + + #[tokio::test] + async fn run_query_writes_the_rows_as_json_and_exits_zero() { + let mut out = Vec::new(); + let args = query_args(&[ + "dirsql", + "query", + "--format", + "json", + "SELECT 1 AS a, 'x' AS b", + ]); + assert_eq!(run_query(args, &mut out).await, 0); + assert_eq!(out, b"[{\"a\":1,\"b\":\"x\"}]\n"); + } + + #[tokio::test] + async fn run_query_writes_the_rows_as_a_table_when_asked() { + let mut out = Vec::new(); + let args = query_args(&["dirsql", "query", "--format", "table", "SELECT 1 AS a"]); + assert_eq!(run_query(args, &mut out).await, 0); + assert_eq!(out, b"a\n-\n1\n\n1 row\n"); + } + + #[tokio::test] + async fn run_query_exits_one_and_writes_nothing_when_the_query_fails() { + let mut out = Vec::new(); + let args = query_args(&["dirsql", "query", "--format", "json", "SELECT nope"]); + assert_eq!(run_query(args, &mut out).await, 1); + assert!(out.is_empty()); + } + /// The `ConfigArgs` parsed from a `query` subcommand invocation (#609: /// config flags are subcommand-local, so they live on the Query variant). fn query_common(argv: &[&str]) -> ConfigArgs { diff --git a/packages/rust/src/cli/serialize.rs b/packages/rust/src/cli/serialize.rs index 5ca8cbc5c..4841a15b6 100644 --- a/packages/rust/src/cli/serialize.rs +++ b/packages/rust/src/cli/serialize.rs @@ -27,6 +27,73 @@ fn ordered_row_to_json(columns: &[String], row: &Row) -> Value { Value::Object(map) } +/// A query's rows written one at a time as the JSON array [`result_to_json`] +/// renders, with no map per row. +pub(super) struct JsonRows { + out: Vec, + /// Each key as JSON with the position of the cell it takes. + keys: Vec<(Vec, usize)>, + rows: usize, +} + +impl JsonRows { + pub(super) fn new(columns: &[String]) -> Self { + // A name the projection repeats is one key, at its first position, + // holding its last cell, as a JSON object built cell by cell is. + let mut keys: Vec<(Vec, usize)> = Vec::new(); + let mut firsts: Vec<&String> = Vec::new(); + for (i, name) in columns.iter().enumerate() { + match firsts.iter().position(|first| *first == name) { + Some(at) => keys[at].1 = i, + None => { + firsts.push(name); + keys.push((json_bytes(name), i)); + } + } + } + Self { + out: b"[".to_vec(), + keys, + rows: 0, + } + } + + pub(super) fn row(&mut self, cells: &[CellValue]) { + if self.rows > 0 { + self.out.push(b','); + } + self.rows += 1; + self.out.push(b'{'); + for (n, (key, i)) in self.keys.iter().enumerate() { + if n > 0 { + self.out.push(b','); + } + self.out.extend_from_slice(key); + self.out.push(b':'); + write_cell(&mut self.out, &cells[*i]); + } + self.out.push(b'}'); + } + + /// The array with the newline the CLI prints after it. + pub(super) fn finish(mut self) -> Vec { + self.out.extend_from_slice(b"]\n"); + self.out + } +} + +fn json_bytes(text: &str) -> Vec { + serde_json::to_vec(text).expect("a string always serializes") +} + +fn write_cell(out: &mut Vec, value: &CellValue) { + let written = match value { + CellValue::Text(text) => serde_json::to_writer(&mut *out, text), + other => serde_json::to_writer(&mut *out, &cell_to_json(other)), + }; + written.expect("writing JSON to memory cannot fail"); +} + pub(super) fn row_to_json(row: &Row) -> Value { let mut map = Map::with_capacity(row.len()); for (k, v) in row { @@ -120,6 +187,73 @@ mod tests { assert_eq!(json[0].to_string(), r#"{"title":"Hello"}"#); } + fn rendered(result: &QueryResult) -> String { + format!("{}\n", Value::Array(result_to_json(result))) + } + + fn streamed(result: &QueryResult) -> String { + let mut rows = JsonRows::new(&result.columns); + for row in &result.rows { + let cells: Vec = result + .columns + .iter() + .map(|column| row[column].clone()) + .collect(); + rows.row(&cells); + } + String::from_utf8(rows.finish()).unwrap() + } + + fn row_of(cells: &[(&str, CellValue)]) -> Row { + cells + .iter() + .map(|(name, value)| ((*name).to_string(), value.clone())) + .collect() + } + + #[test] + fn streamed_rows_match_the_rendered_array_cell_for_cell() { + let result = QueryResult { + columns: vec!["path".into(), "size".into(), "ratio".into(), "body".into()], + rows: vec![ + row_of(&[ + ("path", CellValue::Text("a \"quoted\" \u{e9}\n.md".into())), + ("size", CellValue::Integer(-3)), + ("ratio", CellValue::Real(0.1)), + ("body", CellValue::Blob(vec![0, 255])), + ]), + row_of(&[ + ("path", CellValue::Null), + ("size", CellValue::Integer(i64::MAX)), + ("ratio", CellValue::Real(f64::NAN)), + ("body", CellValue::Text(String::new())), + ]), + ], + }; + assert_eq!(streamed(&result), rendered(&result)); + } + + #[test] + fn streamed_rows_keep_one_key_per_repeated_column_name() { + let result = QueryResult { + columns: vec!["a".into(), "b".into(), "a".into()], + rows: vec![row_of(&[ + ("a", CellValue::Integer(1)), + ("b", CellValue::Integer(2)), + ])], + }; + assert_eq!(streamed(&result), rendered(&result)); + } + + #[test] + fn no_rows_stream_as_an_empty_array() { + let result = QueryResult { + columns: vec!["path".into()], + rows: Vec::new(), + }; + assert_eq!(streamed(&result), "[]\n"); + } + #[test] fn insert_event_emits_expected_shape() { let mut row: Row = HashMap::new(); diff --git a/packages/rust/src/db.rs b/packages/rust/src/db.rs index abdf5f0b7..268d89f71 100644 --- a/packages/rust/src/db.rs +++ b/packages/rust/src/db.rs @@ -995,15 +995,79 @@ impl Db { self.run_statement(sql) } + /// [`query_ordered`](Self::query_ordered) folding each row's cells, in + /// projection order, into the state `start` makes from the projection, + /// as SQLite yields them rather than collected first. + pub fn query_each( + &self, + sql: &str, + start: impl FnOnce(&[String]) -> S, + mut on_row: impl FnMut(&mut S, &[Value]), + ) -> Result { + if !self.batched.is_empty() { + // A batched worker may force a second run, so the rows of the + // first are not final until it ends. + let result = self.query_ordered(sql)?; + let mut state = start(&result.columns); + for row in &result.rows { + let cells: Vec = result + .columns + .iter() + .map(|column| row.get(column).cloned().unwrap_or(Value::Null)) + .collect(); + on_row(&mut state, &cells); + } + return Ok(state); + } + let _calls = self.calls.phase(); + let _scope = self.scope.enter(); + let (mut stmt, column_names) = self.prepare_statement(sql)?; + let mut state = start(&column_names); + let mut rows = stmt.query([])?; + let mut cells = Vec::with_capacity(column_names.len()); + while let Some(row) = rows.next()? { + cells.clear(); + for i in 0..column_names.len() { + let val: rusqlite::types::Value = row.get(i)?; + cells.push(Value::from(val)); + } + on_row(&mut state, &cells); + } + Ok(state) + } + /// Prepare and run one read-only statement, registering the path-tables /// it names along the way. fn run_statement(&self, sql: &str) -> Result { + let (mut stmt, column_names) = self.prepare_statement(sql)?; + let rows = stmt.query_map([], |row| { + let mut map = HashMap::new(); + for (i, name) in column_names.iter().enumerate() { + let val: rusqlite::types::Value = row.get(i)?; + map.insert(name.clone(), Value::from(val)); + } + Ok(map) + })?; + + let mut results = Vec::new(); + for row in rows { + results.push(row?); + } + Ok(QueryResult { + columns: column_names, + rows: results, + }) + } + + /// Prepare one read-only statement, registering the path-tables it names + /// along the way, with its column names. + fn prepare_statement(&self, sql: &str) -> Result<(rusqlite::Statement<'_>, Vec)> { // Each iteration must register a table no earlier iteration did; a // repeat means the fallback is not making progress, so the SQLite // error stands. That is what bounds the loop. let mut attempted: HashSet = HashSet::new(); - let (mut stmt, reads) = loop { + let (stmt, reads) = loop { let error = match self.prepare_guarded(sql) { Ok(prepared) => break prepared, Err(e) => e, @@ -1022,24 +1086,7 @@ impl Db { } self.scope.warm(&reads); let column_names: Vec = stmt.column_names().iter().map(|s| s.to_string()).collect(); - - let rows = stmt.query_map([], |row| { - let mut map = HashMap::new(); - for (i, name) in column_names.iter().enumerate() { - let val: rusqlite::types::Value = row.get(i)?; - map.insert(name.clone(), Value::from(val)); - } - Ok(map) - })?; - - let mut results = Vec::new(); - for row in rows { - results.push(row?); - } - Ok(QueryResult { - columns: column_names, - rows: results, - }) + Ok((stmt, column_names)) } /// Prepare `sql` with the internal-table / ATTACH authorizer installed for @@ -1392,6 +1439,65 @@ mod tests { ); } + fn folded(db: &Db, sql: &str) -> (Vec, Vec>) { + db.query_each( + sql, + |columns| (columns.to_vec(), Vec::new()), + |(_, rows), cells| rows.push(cells.to_vec()), + ) + .unwrap() + } + + #[test] + fn query_each_folds_each_row_in_projection_order() { + let db = titled_docs(&["a", "b"]); + + let (columns, rows) = folded(&db, "SELECT 7 AS n, title FROM docs ORDER BY title"); + + assert_eq!(columns, ["n", "title"]); + assert_eq!( + rows, + [ + [Value::Integer(7), Value::Text("a".into())], + [Value::Integer(7), Value::Text("b".into())] + ] + ); + } + + #[test] + fn query_each_folds_the_rows_a_batched_function_serves() { + let mut db = titled_docs(&["a", "b"]); + let (worker, _) = functions::test_support::batched_worker( + "up", + 8, + Arc::clone(&db.calls), + vec![r#"{"ok": "A"}"#, r#"{"results": [{"ok": "B"}]}"#], + ); + db.register_batched_worker(worker, &[1]); + + let (columns, rows) = folded(&db, "SELECT title, up(title) AS v FROM docs ORDER BY title"); + + assert_eq!(columns, ["title", "v"]); + assert_eq!( + rows, + [ + [Value::Text("a".into()), Value::Text("A".into())], + [Value::Text("b".into()), Value::Text("B".into())] + ] + ); + } + + #[test] + fn query_each_refuses_a_write() { + let db = titled_docs(&["a"]); + + let err = db + .query_each("DELETE FROM docs", |_| (), |_, _| {}) + .unwrap_err(); + + assert!(matches!(err, DbError::WriteForbidden)); + } + /// A statement that never reaches the function asks the worker nothing /// and keeps its first run's rows. #[test] diff --git a/packages/rust/src/lib.rs b/packages/rust/src/lib.rs index ae8136579..9965779cb 100644 --- a/packages/rust/src/lib.rs +++ b/packages/rust/src/lib.rs @@ -490,6 +490,19 @@ impl DirSQL { db.query_ordered(sql).map_err(map_db_error) } + /// [`query_ordered`](Self::query_ordered) folding each row's cells, in + /// projection order, into the state `start` makes from the projection, + /// as they are read rather than collected first. + pub fn query_each( + &self, + sql: &str, + start: impl FnOnce(&[String]) -> S, + on_row: impl FnMut(&mut S, &[Value]), + ) -> Result { + let db = self.inner.db.lock().map_err(DirSqlError::lock)?; + db.query_each(sql, start, on_row).map_err(map_db_error) + } + /// Lazily create the filesystem watcher. Idempotent; subsequent calls are /// no-ops. Called implicitly by [`poll_events`](Self::poll_events) and /// [`watch`](Self::watch). diff --git a/packages/ts/e2e-attestations/claude-1326-stream-rows.json b/packages/ts/e2e-attestations/claude-1326-stream-rows.json new file mode 100644 index 000000000..41aa7eef6 --- /dev/null +++ b/packages/ts/e2e-attestations/claude-1326-stream-rows.json @@ -0,0 +1,7 @@ +{ + "command": "pnpm test:e2e", + "ran_at": 1791241759, + "exit_code": 0, + "commit": "d3344083e2f3449179f98349076dc4985e6755da", + "branch": "claude/1326-stream-rows" +}