mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-23 03:07:13 +00:00
corpus-core: make the stream contract hold under hostile input
This commit is contained in:
@@ -13,6 +13,5 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[lints.rust]
|
||||
unsafe_code = "forbid"
|
||||
missing_docs = "warn"
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -26,11 +26,23 @@ pub struct Chunk {
|
||||
/// What part of the work the passage belongs to.
|
||||
pub role: Role,
|
||||
/// Seconds from the start of the recording, for timed sources.
|
||||
#[serde(
|
||||
default,
|
||||
serialize_with = "crate::float::serialize_option",
|
||||
deserialize_with = "crate::float::deserialize_option"
|
||||
)]
|
||||
pub start_time: Option<f64>,
|
||||
/// Seconds from the start of the recording, for timed sources.
|
||||
#[serde(
|
||||
default,
|
||||
serialize_with = "crate::float::serialize_option",
|
||||
deserialize_with = "crate::float::deserialize_option"
|
||||
)]
|
||||
pub end_time: Option<f64>,
|
||||
/// Set tags, sorted and deduplicated; see [`flags`].
|
||||
#[serde(default)]
|
||||
/// Set tags; see [`flags`]. Sorted and deduplicated on the way in and on
|
||||
/// the way out, so the same set of flags is always the same bytes
|
||||
/// whatever order the producer set them in.
|
||||
#[serde(default, deserialize_with = "sorted_flags")]
|
||||
pub flags: Vec<String>,
|
||||
}
|
||||
|
||||
@@ -62,12 +74,25 @@ impl Chunk {
|
||||
|
||||
/// Whether a flag is set.
|
||||
pub fn has_flag(&self, flag: &str) -> bool {
|
||||
self.flags
|
||||
.binary_search_by(|f| f.as_str().cmp(flag))
|
||||
.is_ok()
|
||||
self.flags.iter().any(|set| set == flag)
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::Record for Chunk {
|
||||
fn version(&self) -> u32 {
|
||||
self.v
|
||||
}
|
||||
}
|
||||
|
||||
/// Sort and deduplicate flags as they are read, so a record written by
|
||||
/// another implementation gets the same treatment as one built here.
|
||||
fn sorted_flags<'de, D: serde::Deserializer<'de>>(input: D) -> Result<Vec<String>, D::Error> {
|
||||
let mut flags = Vec::<String>::deserialize(input)?;
|
||||
flags.sort_unstable();
|
||||
flags.dedup();
|
||||
Ok(flags)
|
||||
}
|
||||
|
||||
/// The flag names the tools in this workspace set.
|
||||
pub mod flags {
|
||||
/// The passage's footnote markers were resolved against a citation list.
|
||||
@@ -122,6 +147,25 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flags_read_from_a_stream_are_sorted_and_unique_too() {
|
||||
let chunk: Chunk = serde_json::from_str(
|
||||
r#"{"idx":0,"text":"x","tokens":1,"role":"body","flags":["zebra","alpha","alpha"]}"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(chunk.flags, ["alpha", "zebra"]);
|
||||
assert!(chunk.has_flag("zebra"));
|
||||
assert!(chunk.has_flag("alpha"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_timestamp_that_is_not_a_number_is_refused_rather_than_dropped() {
|
||||
let mut chunk = Chunk::new(0, "x", 1);
|
||||
chunk.start_time = Some(f64::NAN);
|
||||
let e = serde_json::to_string(&chunk).unwrap_err();
|
||||
assert!(e.to_string().contains("cannot be written"), "{e}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flags_stay_sorted_and_unique() {
|
||||
let mut chunk = Chunk::new(0, "x", 1);
|
||||
|
||||
@@ -18,8 +18,9 @@ pub struct Doc {
|
||||
pub kind: Kind,
|
||||
/// Extension the name claimed, lowercased, without the dot.
|
||||
pub declared_ext: Option<String>,
|
||||
/// Size on disk.
|
||||
pub bytes: u64,
|
||||
/// Size on disk. Signed, because this is a database column as often as
|
||||
/// it is a number, and every database that will hold it counts in `i64`.
|
||||
pub bytes: i64,
|
||||
/// Page count, for paginated formats.
|
||||
pub pages: Option<u32>,
|
||||
/// ISO 639-2/B language code, when one was detected.
|
||||
@@ -29,12 +30,13 @@ pub struct Doc {
|
||||
/// Whether the file has to be OCRed to be searchable.
|
||||
pub needs_ocr: Option<bool>,
|
||||
/// Human-readable observations, in the order they were made.
|
||||
#[serde(default)]
|
||||
pub notes: Vec<String>,
|
||||
}
|
||||
|
||||
impl Doc {
|
||||
/// A doc with only what inspection of the bytes always yields.
|
||||
pub fn new(path: impl Into<String>, kind: Kind, bytes: u64) -> Self {
|
||||
pub fn new(path: impl Into<String>, kind: Kind, bytes: i64) -> Self {
|
||||
Self {
|
||||
v: crate::V,
|
||||
path: path.into(),
|
||||
@@ -55,6 +57,12 @@ impl Doc {
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::Record for Doc {
|
||||
fn version(&self) -> u32 {
|
||||
self.v
|
||||
}
|
||||
}
|
||||
|
||||
/// The format a file actually is.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
@@ -137,8 +145,13 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn a_record_written_without_a_version_reads_as_version_one() {
|
||||
let doc: Doc =
|
||||
serde_json::from_str(r#"{"path":"a","kind":"txt","bytes":1,"notes":[]}"#).unwrap();
|
||||
let doc: Doc = serde_json::from_str(r#"{"path":"a","kind":"txt","bytes":1}"#).unwrap();
|
||||
assert_eq!(doc.v, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_omitted_empty_list_reads_as_an_empty_list() {
|
||||
let doc: Doc = serde_json::from_str(r#"{"path":"a","kind":"txt","bytes":1}"#).unwrap();
|
||||
assert!(doc.notes.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,22 +4,32 @@
|
||||
//! without parsing anything. That only works if "no" and "broke" are
|
||||
//! different numbers, which is why the failure codes below start where
|
||||
//! `sysexits.h` puts them and leave 0, 1 and 2 to the answer.
|
||||
//!
|
||||
//! One case has no code of its own on purpose: a consumer closing the pipe.
|
||||
//! `… | head -1` is a normal way to use a filter, and a filter that reports
|
||||
//! failure because someone stopped reading would make every such pipeline
|
||||
//! look broken. A tool that sees a broken pipe stops and exits [`OK`].
|
||||
|
||||
use std::process::ExitCode;
|
||||
|
||||
/// Success, or a predicate answering yes.
|
||||
pub const OK: i32 = 0;
|
||||
pub const OK: u8 = 0;
|
||||
/// A predicate answering no.
|
||||
pub const NO: i32 = 1;
|
||||
pub const NO: u8 = 1;
|
||||
/// A predicate that could not tell — a damaged file, a format it does not
|
||||
/// know. Distinct from no, because a caller may want to retry or escalate.
|
||||
pub const INDETERMINATE: i32 = 2;
|
||||
pub const INDETERMINATE: u8 = 2;
|
||||
/// The command line was wrong.
|
||||
pub const USAGE: i32 = 64;
|
||||
/// The input was not what the tool can work on.
|
||||
pub const DATA: i32 = 65;
|
||||
pub const USAGE: u8 = 64;
|
||||
/// The input was not what the tool can work on: bad JSON, a record from a
|
||||
/// version this build does not implement, bytes that are not text.
|
||||
pub const DATA: u8 = 65;
|
||||
/// An external program the tool needs is not on `PATH`.
|
||||
pub const UNAVAILABLE: i32 = 69;
|
||||
pub const UNAVAILABLE: u8 = 69;
|
||||
/// The tool failed at something that is nobody's fault but its own.
|
||||
pub const SOFTWARE: u8 = 70;
|
||||
/// Reading or writing failed.
|
||||
pub const IO: i32 = 74;
|
||||
pub const IO: u8 = 74;
|
||||
|
||||
/// The answer a predicate tool gives.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -34,7 +44,7 @@ pub enum Verdict {
|
||||
|
||||
impl Verdict {
|
||||
/// The process exit code for this answer.
|
||||
pub const fn code(self) -> i32 {
|
||||
pub const fn code(self) -> u8 {
|
||||
match self {
|
||||
Verdict::Yes => OK,
|
||||
Verdict::No => NO,
|
||||
@@ -49,6 +59,12 @@ impl From<bool> for Verdict {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Verdict> for ExitCode {
|
||||
fn from(verdict: Verdict) -> Self {
|
||||
ExitCode::from(verdict.code())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -62,7 +78,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn failure_codes_never_collide_with_an_answer() {
|
||||
for code in [USAGE, DATA, UNAVAILABLE, IO] {
|
||||
for code in [USAGE, DATA, UNAVAILABLE, SOFTWARE, IO] {
|
||||
assert!(code > INDETERMINATE);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
//! Serde plumbing for floats that JSON can actually carry.
|
||||
//!
|
||||
//! JSON has no spelling for NaN or infinity, and `serde_json` writes them as
|
||||
//! `null`. That is a silent corruption in both directions: a score written as
|
||||
//! `null` fails to read back and kills the stream at the consumer, and a
|
||||
//! timestamp written as `null` reads back as "no timestamp" with nothing to
|
||||
//! say it was ever anything else. So a non-finite number is refused where it
|
||||
//! is produced, which is the only place anyone can still fix it.
|
||||
|
||||
use serde::{Deserialize, Deserializer, Serializer, de::Error as _, ser::Error as _};
|
||||
|
||||
pub(crate) fn serialize<S: Serializer>(value: &f64, out: S) -> Result<S::Ok, S::Error> {
|
||||
if !value.is_finite() {
|
||||
return Err(S::Error::custom(format!(
|
||||
"{value} cannot be written as JSON"
|
||||
)));
|
||||
}
|
||||
out.serialize_f64(*value)
|
||||
}
|
||||
|
||||
pub(crate) fn serialize_option<S: Serializer>(
|
||||
value: &Option<f64>,
|
||||
out: S,
|
||||
) -> Result<S::Ok, S::Error> {
|
||||
match value {
|
||||
Some(value) => serialize(value, out),
|
||||
None => out.serialize_none(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn deserialize<'de, D: Deserializer<'de>>(input: D) -> Result<f64, D::Error> {
|
||||
let value = f64::deserialize(input)?;
|
||||
if !value.is_finite() {
|
||||
return Err(D::Error::custom(format!("{value} is not a finite number")));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub(crate) fn deserialize_option<'de, D: Deserializer<'de>>(
|
||||
input: D,
|
||||
) -> Result<Option<f64>, D::Error> {
|
||||
let value = Option::<f64>::deserialize(input)?;
|
||||
match value {
|
||||
Some(value) if !value.is_finite() => {
|
||||
Err(D::Error::custom(format!("{value} is not a finite number")))
|
||||
}
|
||||
value => Ok(value),
|
||||
}
|
||||
}
|
||||
@@ -22,8 +22,18 @@ pub struct Hit {
|
||||
/// Position within the channel's own ranking, from one.
|
||||
pub rank: u32,
|
||||
/// The channel's raw score, on whatever scale that channel uses.
|
||||
#[serde(
|
||||
serialize_with = "crate::float::serialize",
|
||||
deserialize_with = "crate::float::deserialize"
|
||||
)]
|
||||
pub score: f64,
|
||||
/// Anything the producer wants to carry through fusion untouched.
|
||||
///
|
||||
/// Keys serialise in sorted order, which is what makes a fused stream
|
||||
/// byte-reproducible. That holds because `serde_json`'s map is a
|
||||
/// `BTreeMap`; enabling its `preserve_order` feature anywhere in the
|
||||
/// dependency graph would switch it to insertion order and break it. The
|
||||
/// test below is what would catch that.
|
||||
#[serde(default)]
|
||||
pub payload: Map<String, Value>,
|
||||
}
|
||||
@@ -42,6 +52,12 @@ impl Hit {
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::Record for Hit {
|
||||
fn version(&self) -> u32 {
|
||||
self.v
|
||||
}
|
||||
}
|
||||
|
||||
/// A retrieval channel.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
@@ -80,6 +96,22 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_score_that_is_not_a_number_is_refused_rather_than_written_as_null() {
|
||||
let mut hit = Hit::new("a", Channel::Fts, 1, 0.0);
|
||||
for bad in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
|
||||
hit.score = bad;
|
||||
let e = serde_json::to_string(&hit).unwrap_err();
|
||||
assert!(e.to_string().contains("cannot be written"), "{bad}: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_score_that_overflows_a_float_is_refused_on_the_way_in() {
|
||||
let line = r#"{"id":"a","channel":"fts","rank":1,"score":1e999}"#;
|
||||
assert!(serde_json::from_str::<Hit>(line).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn payload_keys_serialise_in_a_fixed_order() {
|
||||
let mut hit = Hit::new("abc", Channel::Vector, 2, 0.25);
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
|
||||
pub mod chunk;
|
||||
pub mod doc;
|
||||
mod float;
|
||||
pub mod exit;
|
||||
pub mod hit;
|
||||
pub mod ndjson;
|
||||
@@ -22,6 +23,16 @@ pub use ndjson::{Error, read, write_line};
|
||||
/// The record version these types serialise as.
|
||||
pub const V: u32 = 1;
|
||||
|
||||
/// A record that travels between tools as one line of NDJSON.
|
||||
///
|
||||
/// The version is on the trait rather than read off the struct so that
|
||||
/// [`ndjson::read`] can refuse a stream from a future version of the suite
|
||||
/// instead of parsing it into whatever happens to line up.
|
||||
pub trait Record: serde::Serialize + serde::de::DeserializeOwned {
|
||||
/// The version this record announced.
|
||||
fn version(&self) -> u32;
|
||||
}
|
||||
|
||||
/// Serde default for the `v` field, so a record written by an older tool that
|
||||
/// omitted it still reads back as version 1 rather than 0.
|
||||
pub(crate) const fn v() -> u32 {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! Prints the JSON Schema for the record types, so a consumer in another
|
||||
//! language can generate its own bindings instead of reading Rust.
|
||||
|
||||
use std::io::Write;
|
||||
use std::process::ExitCode;
|
||||
|
||||
use corpus_core::{Chunk, Doc, Hit, exit};
|
||||
@@ -14,28 +15,41 @@ Prints the JSON Schema for one record type, or an object of all three.
|
||||
";
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let mut args = std::env::args().skip(1);
|
||||
match args.next().as_deref() {
|
||||
Some("schema") => {
|
||||
let which = args.next().unwrap_or_else(|| "all".into());
|
||||
match schema(&which) {
|
||||
Some(json) => {
|
||||
println!("{json}");
|
||||
ExitCode::from(exit::OK as u8)
|
||||
}
|
||||
None => {
|
||||
eprintln!("corpus-core: unknown record type {which:?}\n\n{USAGE}");
|
||||
ExitCode::from(exit::USAGE as u8)
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("--help" | "-h") => {
|
||||
println!("{USAGE}");
|
||||
ExitCode::from(exit::OK as u8)
|
||||
}
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let args: Vec<&str> = args.iter().map(String::as_str).collect();
|
||||
|
||||
match args.as_slice() {
|
||||
["--help" | "-h"] => print(USAGE),
|
||||
["--version"] => print(&format!("corpus-core {}\n", env!("CARGO_PKG_VERSION"))),
|
||||
["schema"] => emit("all"),
|
||||
["schema", which] => emit(which),
|
||||
_ => {
|
||||
eprintln!("{USAGE}");
|
||||
ExitCode::from(exit::USAGE as u8)
|
||||
eprint!("{USAGE}");
|
||||
ExitCode::from(exit::USAGE)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn emit(which: &str) -> ExitCode {
|
||||
match schema(which) {
|
||||
Some(json) => print(&format!("{json}\n")),
|
||||
None => {
|
||||
eprintln!("corpus-core: unknown record type {which:?}\n");
|
||||
eprint!("{USAGE}");
|
||||
ExitCode::from(exit::USAGE)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write to stdout, treating a consumer that stopped reading as the end of
|
||||
/// the job rather than as a failure. `println!` would panic there.
|
||||
fn print(text: &str) -> ExitCode {
|
||||
match std::io::stdout().write_all(text.as_bytes()) {
|
||||
Ok(()) => ExitCode::from(exit::OK),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => ExitCode::from(exit::OK),
|
||||
Err(e) => {
|
||||
eprintln!("corpus-core: {e}");
|
||||
ExitCode::from(exit::IO)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,22 +3,64 @@
|
||||
//! Every list-shaped stream between these tools is NDJSON: one record per
|
||||
//! line, produced and consumed one at a time, so a slow stage never has to
|
||||
//! hold the whole list and never blocks the pipe behind it.
|
||||
//!
|
||||
//! The reader is deliberately strict. It refuses a record whose version it
|
||||
//! does not know, it stops at the first bad line rather than skipping past a
|
||||
//! disagreement about the contract, and it bounds how much one line may cost.
|
||||
//! A filter pointed at a binary file should fail, not allocate the file.
|
||||
|
||||
use std::io::{BufRead, Write};
|
||||
use std::io::{BufRead, ErrorKind, Read, Write};
|
||||
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::Record;
|
||||
|
||||
/// How long one line may be before the reader gives up, in bytes.
|
||||
///
|
||||
/// Generous for a record — a chunk of a thousand tokens is a few kilobytes —
|
||||
/// and small enough that a file with no newlines in it cannot be read into
|
||||
/// memory by accident.
|
||||
pub const MAX_LINE: usize = 16 * 1024 * 1024;
|
||||
|
||||
/// A UTF-8 byte-order mark, which Windows tooling puts at the front of files
|
||||
/// and which is not whitespace, not JSON, and not the producer's fault.
|
||||
const BOM: &[u8] = &[0xef, 0xbb, 0xbf];
|
||||
|
||||
/// Something went wrong reading or writing a stream.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
/// A line was not the record the reader expected.
|
||||
#[error("line {line}: {source}")]
|
||||
#[error("stream line {line}: {source}")]
|
||||
Parse {
|
||||
/// One-based line number within the stream, counting blank lines.
|
||||
line: u64,
|
||||
/// What serde made of it.
|
||||
source: serde_json::Error,
|
||||
},
|
||||
/// A line was not UTF-8 at all.
|
||||
#[error("stream line {line}: not valid UTF-8")]
|
||||
Encoding {
|
||||
/// One-based line number within the stream.
|
||||
line: u64,
|
||||
},
|
||||
/// A record announced a version this build does not implement.
|
||||
#[error("stream line {line}: record version {found}, expected {expected}")]
|
||||
Version {
|
||||
/// One-based line number within the stream.
|
||||
line: u64,
|
||||
/// The version the record announced.
|
||||
found: u32,
|
||||
/// The version this build understands.
|
||||
expected: u32,
|
||||
},
|
||||
/// A line exceeded [`MAX_LINE`].
|
||||
#[error("stream line {line}: longer than the {limit} byte line limit")]
|
||||
TooLong {
|
||||
/// One-based line number within the stream.
|
||||
line: u64,
|
||||
/// The limit that was exceeded.
|
||||
limit: usize,
|
||||
},
|
||||
/// A record could not be turned into JSON at all.
|
||||
#[error("encoding a record: {0}")]
|
||||
Encode(#[source] serde_json::Error),
|
||||
@@ -27,27 +69,127 @@ pub enum Error {
|
||||
Io(#[from] std::io::Error),
|
||||
}
|
||||
|
||||
/// Read a stream of records, one per line.
|
||||
///
|
||||
/// Blank lines are skipped, so a stream that was concatenated or hand-edited
|
||||
/// still reads. The iterator stops at the first error rather than skipping
|
||||
/// it: a malformed line means the producer and the consumer disagree about
|
||||
/// the contract, and continuing would bury that.
|
||||
impl Error {
|
||||
/// Whether this is the consumer downstream having closed the pipe.
|
||||
///
|
||||
/// Not an error in a filter: `… | head -1` is a normal way to use one.
|
||||
/// A tool that sees this should stop and exit successfully, the way every
|
||||
/// other program in a pipeline does.
|
||||
pub fn is_broken_pipe(&self) -> bool {
|
||||
matches!(self, Error::Io(e) if e.kind() == ErrorKind::BrokenPipe)
|
||||
}
|
||||
|
||||
/// The line this error happened on, when it happened on one.
|
||||
pub fn line(&self) -> Option<u64> {
|
||||
match self {
|
||||
Error::Parse { line, .. }
|
||||
| Error::Encoding { line }
|
||||
| Error::Version { line, .. }
|
||||
| Error::TooLong { line, .. } => Some(*line),
|
||||
Error::Encode(_) | Error::Io(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a stream of records, one per line, with the default line limit.
|
||||
pub fn read<T, R>(reader: R) -> impl Iterator<Item = Result<T, Error>>
|
||||
where
|
||||
T: DeserializeOwned,
|
||||
T: Record,
|
||||
R: BufRead,
|
||||
{
|
||||
read_with_limit(reader, MAX_LINE)
|
||||
}
|
||||
|
||||
/// Read a stream of records, one per line, refusing any line longer than
|
||||
/// `limit` bytes.
|
||||
///
|
||||
/// Blank lines are skipped, so a stream that was concatenated or hand-edited
|
||||
/// still reads, and a byte-order mark at the very front is ignored. The
|
||||
/// iterator ends at the first error: a bad line means the producer and the
|
||||
/// consumer disagree about the contract, and reading on would bury that.
|
||||
pub fn read_with_limit<T, R>(mut reader: R, limit: usize) -> impl Iterator<Item = Result<T, Error>>
|
||||
where
|
||||
T: Record,
|
||||
R: BufRead,
|
||||
{
|
||||
let mut line_no = 0u64;
|
||||
reader.lines().filter_map(move |line| {
|
||||
line_no += 1;
|
||||
match line {
|
||||
Err(e) => Some(Err(Error::Io(e))),
|
||||
Ok(line) if line.trim().is_empty() => None,
|
||||
Ok(line) => Some(serde_json::from_str(&line).map_err(|source| Error::Parse {
|
||||
line: line_no,
|
||||
source,
|
||||
})),
|
||||
let mut buf = Vec::new();
|
||||
let mut done = false;
|
||||
let mut at_start = true;
|
||||
|
||||
std::iter::from_fn(move || {
|
||||
if done {
|
||||
return None;
|
||||
}
|
||||
loop {
|
||||
buf.clear();
|
||||
// Bounded so that a stream with no newline in it costs `limit`
|
||||
// bytes and not the whole stream.
|
||||
let read = (&mut reader)
|
||||
.take(limit as u64 + 1)
|
||||
.read_until(b'\n', &mut buf);
|
||||
let read = match read {
|
||||
Ok(read) => read,
|
||||
Err(e) => {
|
||||
done = true;
|
||||
return Some(Err(Error::Io(e)));
|
||||
}
|
||||
};
|
||||
if read == 0 {
|
||||
done = true;
|
||||
return None;
|
||||
}
|
||||
line_no += 1;
|
||||
|
||||
let mut line = buf.as_slice();
|
||||
if line.last() == Some(&b'\n') {
|
||||
line = &line[..line.len() - 1];
|
||||
if line.last() == Some(&b'\r') {
|
||||
line = &line[..line.len() - 1];
|
||||
}
|
||||
} else if read > limit {
|
||||
done = true;
|
||||
return Some(Err(Error::TooLong {
|
||||
line: line_no,
|
||||
limit,
|
||||
}));
|
||||
}
|
||||
if at_start {
|
||||
at_start = false;
|
||||
if let Some(rest) = line.strip_prefix(BOM) {
|
||||
line = rest;
|
||||
}
|
||||
}
|
||||
|
||||
let line = match std::str::from_utf8(line) {
|
||||
Ok(line) => line,
|
||||
Err(_) => {
|
||||
done = true;
|
||||
return Some(Err(Error::Encoding { line: line_no }));
|
||||
}
|
||||
};
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
return Some(match serde_json::from_str::<T>(line) {
|
||||
Err(source) => {
|
||||
done = true;
|
||||
Err(Error::Parse {
|
||||
line: line_no,
|
||||
source,
|
||||
})
|
||||
}
|
||||
Ok(record) if record.version() != crate::V => {
|
||||
done = true;
|
||||
Err(Error::Version {
|
||||
line: line_no,
|
||||
found: record.version(),
|
||||
expected: crate::V,
|
||||
})
|
||||
}
|
||||
Ok(record) => Ok(record),
|
||||
});
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -61,7 +203,15 @@ where
|
||||
T: Serialize,
|
||||
W: Write,
|
||||
{
|
||||
serde_json::to_writer(&mut *writer, record).map_err(Error::Encode)?;
|
||||
// A write failure arrives wrapped in a serde error; unwrapping it here is
|
||||
// what lets a caller tell a closed pipe from a record it could not encode.
|
||||
serde_json::to_writer(&mut *writer, record).map_err(|e| {
|
||||
if e.is_io() {
|
||||
Error::Io(e.into())
|
||||
} else {
|
||||
Error::Encode(e)
|
||||
}
|
||||
})?;
|
||||
writer.write_all(b"\n")?;
|
||||
writer.flush()?;
|
||||
Ok(())
|
||||
@@ -70,27 +220,90 @@ where
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{Channel, Hit};
|
||||
use crate::{Channel, Chunk, Hit};
|
||||
|
||||
#[test]
|
||||
fn reads_a_record_per_line_and_skips_blanks() {
|
||||
let input = "{\"id\":\"a\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}\n\n\
|
||||
{\"id\":\"b\",\"channel\":\"vec\",\"rank\":2,\"score\":0.5}\n";
|
||||
let hits: Vec<Hit> = read(input.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(hits.len(), 2);
|
||||
assert_eq!(hits[1].id, "b");
|
||||
assert_eq!(hits[1].channel, Channel::Vector);
|
||||
fn hit_line(id: &str) -> String {
|
||||
format!("{{\"id\":\"{id}\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}}")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bad_line_reports_its_line_number_and_stops_the_stream() {
|
||||
let input = "{\"id\":\"a\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}\nnot json\n";
|
||||
fn reads_a_record_per_line_and_skips_blanks() {
|
||||
let input = format!("{}\n\n{}\n", hit_line("a"), hit_line("b"));
|
||||
let hits: Vec<Hit> = read(input.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(hits.len(), 2);
|
||||
assert_eq!(hits[1].id, "b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_last_line_without_a_newline_still_reads() {
|
||||
let hits: Vec<Hit> = read(hit_line("a").as_bytes())
|
||||
.collect::<Result<_, _>>()
|
||||
.unwrap();
|
||||
assert_eq!(hits.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crlf_endings_read() {
|
||||
let input = format!("{}\r\n{}\r\n", hit_line("a"), hit_line("b"));
|
||||
let hits: Vec<Hit> = read(input.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(hits.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_byte_order_mark_does_not_eat_the_first_record() {
|
||||
let input = format!("\u{feff}{}\n", hit_line("a"));
|
||||
let hits: Vec<Hit> = read(input.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(hits[0].id, "a");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bad_line_reports_its_line_number_and_ends_the_stream() {
|
||||
let input = format!("\n\n{}\nnot json\n{}\n", hit_line("a"), hit_line("b"));
|
||||
let mut stream = read::<Hit, _>(input.as_bytes());
|
||||
assert!(stream.next().unwrap().is_ok());
|
||||
match stream.next().unwrap() {
|
||||
Err(Error::Parse { line, .. }) => assert_eq!(line, 2),
|
||||
Err(e) => assert_eq!(e.line(), Some(4)),
|
||||
other => panic!("expected a parse error, got {other:?}"),
|
||||
}
|
||||
assert!(
|
||||
stream.next().is_none(),
|
||||
"reading continued past a broken contract"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_line_that_is_not_utf8_is_a_data_error_not_an_io_error() {
|
||||
let input: &[u8] = b"{\"id\":\"\xff\xfe\"}\n";
|
||||
match read::<Hit, _>(input).next().unwrap() {
|
||||
Err(Error::Encoding { line }) => assert_eq!(line, 1),
|
||||
other => panic!("expected an encoding error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_endless_line_is_refused_rather_than_buffered() {
|
||||
let input = vec![b'x'; 4096];
|
||||
match read_with_limit::<Hit, _>(input.as_slice(), 128)
|
||||
.next()
|
||||
.unwrap()
|
||||
{
|
||||
Err(Error::TooLong { line, limit }) => {
|
||||
assert_eq!(line, 1);
|
||||
assert_eq!(limit, 128);
|
||||
}
|
||||
other => panic!("expected a length error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_record_from_a_future_version_is_refused() {
|
||||
let input = "{\"v\":2,\"id\":\"a\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}\n";
|
||||
match read::<Hit, _>(input.as_bytes()).next().unwrap() {
|
||||
Err(Error::Version { found, expected, .. }) => {
|
||||
assert_eq!((found, expected), (2, 1));
|
||||
}
|
||||
other => panic!("expected a version error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -101,7 +314,16 @@ mod tests {
|
||||
let text = String::from_utf8(out).unwrap();
|
||||
assert_eq!(text.lines().count(), 2);
|
||||
assert!(text.ends_with('\n'));
|
||||
assert!(!text.contains("\n "));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_newline_inside_a_record_stays_one_line() {
|
||||
let mut out = Vec::new();
|
||||
write_line(&mut out, &Chunk::new(0, "two\nlines", 2)).unwrap();
|
||||
let text = String::from_utf8(out).unwrap();
|
||||
assert_eq!(text.lines().count(), 1);
|
||||
let back: Vec<Chunk> = read(text.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(back[0].text, "two\nlines");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -112,4 +334,41 @@ mod tests {
|
||||
let back: Vec<Hit> = read(out.as_slice()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(back, vec![hit]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_record_is_flushed_as_it_is_written() {
|
||||
#[derive(Default)]
|
||||
struct Counting {
|
||||
flushes: usize,
|
||||
}
|
||||
impl Write for Counting {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
Ok(buf.len())
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
self.flushes += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
let mut out = Counting::default();
|
||||
for i in 0..3 {
|
||||
write_line(&mut out, &Hit::new("a", Channel::Meta, i, 1.0)).unwrap();
|
||||
}
|
||||
assert_eq!(out.flushes, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_closed_consumer_is_reported_as_a_broken_pipe_not_as_bad_data() {
|
||||
struct Closed;
|
||||
impl Write for Closed {
|
||||
fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
|
||||
Err(std::io::Error::new(ErrorKind::BrokenPipe, "broken pipe"))
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
let e = write_line(&mut Closed, &Hit::new("a", Channel::Meta, 1, 1.0)).unwrap_err();
|
||||
assert!(e.is_broken_pipe(), "got {e:?}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
//! The schema command, as a consumer in another language would meet it.
|
||||
|
||||
use std::process::{Command, Output};
|
||||
|
||||
fn run(args: &[&str]) -> Output {
|
||||
Command::new(env!("CARGO_BIN_EXE_corpus-core"))
|
||||
.args(args)
|
||||
.output()
|
||||
.expect("the binary should run")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_named_record_prints_its_schema() {
|
||||
let out = run(&["schema", "doc"]);
|
||||
assert_eq!(out.status.code(), Some(0));
|
||||
let json: serde_json::Value = serde_json::from_slice(&out.stdout).expect("valid JSON");
|
||||
assert_eq!(json["title"], "Doc");
|
||||
assert!(out.stderr.is_empty(), "diagnostics leaked to stdout's side");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_record_named_means_all_of_them() {
|
||||
let out = run(&["schema"]);
|
||||
assert_eq!(out.status.code(), Some(0));
|
||||
let json: serde_json::Value = serde_json::from_slice(&out.stdout).expect("valid JSON");
|
||||
for record in ["doc", "chunk", "hit"] {
|
||||
assert!(json.get(record).is_some(), "missing {record}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unknown_record_is_a_usage_error_with_nothing_on_stdout() {
|
||||
let out = run(&["schema", "nope"]);
|
||||
assert_eq!(out.status.code(), Some(64));
|
||||
assert!(out.stdout.is_empty());
|
||||
assert!(!out.stderr.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn arguments_that_mean_nothing_are_refused_rather_than_ignored() {
|
||||
for args in [
|
||||
vec![],
|
||||
vec!["schema", "doc", "junk"],
|
||||
vec!["nonsense"],
|
||||
vec!["--json"],
|
||||
] {
|
||||
let out = run(&args);
|
||||
assert_eq!(out.status.code(), Some(64), "for {args:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn it_says_what_version_it_is() {
|
||||
let out = run(&["--version"]);
|
||||
assert_eq!(out.status.code(), Some(0));
|
||||
let said = String::from_utf8(out.stdout).unwrap();
|
||||
assert!(said.contains(env!("CARGO_PKG_VERSION")), "{said:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_same_schema_comes_out_byte_identical_every_time() {
|
||||
assert_eq!(run(&["schema", "all"]).stdout, run(&["schema", "all"]).stdout);
|
||||
}
|
||||
Reference in New Issue
Block a user