Skip to main content

http_streams_core/
csv_record_codec.rs

1//! Framing CSV records.
2//!
3//! Built directly on `csv-core` rather than on line splitting, for three reasons:
4//!
5//! 1. **Correctness.** A CSV field may contain a newline when quoted, and this crate's own
6//!    encoder emits exactly that. Splitting on `\n` truncates such a field and yields the
7//!    remainder as a bogus record: silent data loss on the row that *did* decode.
8//! 2. **Allocation.** `csv::ReaderBuilder::from_reader` allocates an 8 KiB buffer, so parsing
9//!    one record at a time through it allocates 8 KiB per row.
10//! 3. **Types.** This framer is not generic over the item type. A decoder that structurally
11//!    mentions `T` forces a `T: 'b` bound onto every caller's public signature; deserialisation
12//!    therefore happens in a separate step, where `T` appears only as a return type.
13
14use crate::error::{StreamError, StreamErrorKind};
15use bytes::{Buf, BytesMut};
16use csv_core::{ReadRecordResult, Reader as CoreReader};
17use tokio_util::codec::Decoder;
18
19/// Initial size of the field and offset buffers; both grow on demand.
20const INITIAL_FIELDS: usize = 512;
21const INITIAL_ENDS: usize = 16;
22
23/// How to frame CSV records.
24#[derive(Debug, Clone, Copy)]
25pub struct CsvFrameConfig {
26    /// Field delimiter.
27    pub delimiter: u8,
28    /// Quote character.
29    pub quote: u8,
30    /// Whether a doubled quote character is an escaped quote.
31    pub double_quote: bool,
32    /// Escape character, if escaping is by prefix rather than by doubling.
33    pub escape: Option<u8>,
34    /// Record terminator.
35    pub terminator: csv_core::Terminator,
36}
37
38impl CsvFrameConfig {
39    fn build(&self) -> CoreReader {
40        csv_core::ReaderBuilder::new()
41            .delimiter(self.delimiter)
42            .quote(self.quote)
43            .double_quote(self.double_quote)
44            .escape(self.escape)
45            .terminator(self.terminator)
46            .build()
47    }
48}
49
50/// A [`Decoder`] that yields one [`csv::ByteRecord`] per CSV record.
51///
52/// Not generic over the item type: see the module docs.
53#[derive(Debug)]
54pub struct CsvRecordCodec {
55    core: CoreReader,
56    output: Vec<u8>,
57    ends: Vec<usize>,
58    outlen: usize,
59    endlen: usize,
60    header_pending: bool,
61    max_len: usize,
62}
63
64impl CsvRecordCodec {
65    /// A framer reading records per `config`, skipping a leading header row if `has_headers`.
66    pub fn new(config: CsvFrameConfig, has_headers: bool, max_len: usize) -> Self {
67        Self {
68            core: config.build(),
69            output: vec![0; INITIAL_FIELDS],
70            ends: vec![0; INITIAL_ENDS],
71            outlen: 0,
72            endlen: 0,
73            header_pending: has_headers,
74            max_len,
75        }
76    }
77
78    /// Builds the finished record and resets the accumulators for the next one.
79    fn take_record(&mut self) -> csv::ByteRecord {
80        let mut fields: Vec<&[u8]> = Vec::with_capacity(self.endlen);
81        let mut start = 0;
82        for &end in &self.ends[..self.endlen] {
83            fields.push(&self.output[start..end]);
84            start = end;
85        }
86        let record = csv::ByteRecord::from(fields);
87        self.outlen = 0;
88        self.endlen = 0;
89        record
90    }
91
92    /// One pass of the framing loop, mirroring `csv`'s own reader.
93    ///
94    /// `at_eof` says whether the caller may signal end of input, which `csv-core` recognises as
95    /// an empty input slice and which is the only way to flush a final record that has no
96    /// trailing terminator.
97    fn next_record(
98        &mut self,
99        buf: &mut BytesMut,
100        at_eof: bool,
101    ) -> Result<Option<csv::ByteRecord>, StreamError> {
102        loop {
103            let input_was_empty = buf.is_empty();
104            if input_was_empty && !at_eof {
105                return Ok(None);
106            }
107
108            let (res, nin, nout, nend) = self.core.read_record(
109                &buf[..],
110                &mut self.output[self.outlen..],
111                &mut self.ends[self.endlen..],
112            );
113
114            buf.advance(nin);
115            self.outlen += nout;
116            self.endlen += nend;
117
118            // Counts the offset vector as well as the field bytes. `csv-core` records one end
119            // offset per field whether or not that field wrote any bytes, so a record like
120            // `a,,,,,,,...` keeps `outlen` near zero while `ends` grows without ever tripping a
121            // bytes-only check.
122            let record_bytes = self
123                .outlen
124                .saturating_add(self.endlen.saturating_mul(std::mem::size_of::<usize>()));
125            if record_bytes > self.max_len {
126                return Err(StreamError::new(
127                    StreamErrorKind::MaxLenReachedError,
128                    None,
129                    Some("Max record length reached".into()),
130                ));
131            }
132
133            match res {
134                // Not a whole record yet. When the body has ended, going round again passes an
135                // empty slice, which is how `csv-core` is told there will be no more input.
136                ReadRecordResult::InputEmpty => {
137                    if input_was_empty {
138                        return Ok(None);
139                    }
140                    continue;
141                }
142                ReadRecordResult::OutputFull => {
143                    self.output.resize(self.output.len().saturating_mul(2).max(1), 0);
144                    continue;
145                }
146                ReadRecordResult::OutputEndsFull => {
147                    self.ends.resize(self.ends.len().saturating_mul(2).max(1), 0);
148                    continue;
149                }
150                ReadRecordResult::Record => {
151                    let record = self.take_record();
152                    // The header slot is consumed here rather than by a `.skip(1)` downstream:
153                    // skipping the first *yielded* item would swallow a header that failed to
154                    // frame, and the stream would report itself as cleanly completed.
155                    if self.header_pending {
156                        self.header_pending = false;
157                        continue;
158                    }
159                    return Ok(Some(record));
160                }
161                ReadRecordResult::End => return Ok(None),
162            }
163        }
164    }
165}
166
167impl Decoder for CsvRecordCodec {
168    /// A framed record, not yet deserialised. Errors here are framing errors and are terminal.
169    type Item = csv::ByteRecord;
170    type Error = StreamError;
171
172    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
173        self.next_record(buf, false)
174    }
175
176    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
177        self.next_record(buf, true)
178    }
179}