Skip to main content

http_streams_core/
csv_format.rs

1//! CSV, both directions.
2
3use crate::content_type::ContentType;
4use crate::csv_record_codec::{CsvFrameConfig, CsvRecordCodec};
5use crate::error::{StreamError, StreamErrorKind};
6use crate::format::{
7    DecodeOptions, DefaultFormat, FrameParser, ItemEncoder, StreamFormat, StreamFormatDecode,
8    StreamFormatEncode,
9};
10use bytes::BytesMut;
11use serde::{Deserialize, Serialize};
12
13const CSV_CONTENT_TYPE: &str = "text/csv";
14const CSV_ALIASES: &[&str] = &["text/csv", "application/csv"];
15
16/// CSV rows, with an optional header row.
17#[derive(Debug, Clone)]
18pub struct CsvStreamFormat {
19    has_headers: bool,
20    delimiter: u8,
21    flexible: bool,
22    quote_style: csv::QuoteStyle,
23    quote: u8,
24    double_quote: bool,
25    escape: u8,
26    terminator: csv::Terminator,
27}
28
29impl Default for CsvStreamFormat {
30    fn default() -> Self {
31        Self {
32            has_headers: true,
33            delimiter: b',',
34            flexible: false,
35            quote_style: csv::QuoteStyle::Necessary,
36            quote: b'"',
37            double_quote: true,
38            escape: b'\\',
39            terminator: csv::Terminator::Any(b'\n'),
40        }
41    }
42}
43
44impl DefaultFormat for CsvStreamFormat {
45    fn default_format() -> Self {
46        Self::default()
47    }
48}
49
50impl CsvStreamFormat {
51    /// CSV with the given header behaviour and field delimiter, everything else default.
52    pub fn new(has_headers: bool, delimiter: u8) -> Self {
53        Self {
54            has_headers,
55            delimiter,
56            ..Default::default()
57        }
58    }
59
60    /// Sets whether to use flexible serialize.
61    ///
62    /// Encode only. On the way back in, records are framed and deserialised one at a time, so
63    /// there is no "first record" to compare a field count against. That was already true of
64    /// the previous decoder, where a fresh reader was built per row.
65    pub fn with_flexible(mut self, flexible: bool) -> Self {
66        self.flexible = flexible;
67        self
68    }
69
70    /// Sets the quote style to use.
71    pub fn with_quote_style(mut self, quote_style: csv::QuoteStyle) -> Self {
72        self.quote_style = quote_style;
73        self
74    }
75
76    /// Sets the quote character to use.
77    pub fn with_quote(mut self, quote: u8) -> Self {
78        self.quote = quote;
79        self
80    }
81
82    /// Sets whether to double quote.
83    pub fn with_double_quote(mut self, double_quote: bool) -> Self {
84        self.double_quote = double_quote;
85        self
86    }
87
88    /// Sets the escape character to use.
89    pub fn with_escape(mut self, escape: u8) -> Self {
90        self.escape = escape;
91        self
92    }
93
94    /// Sets the record terminator to use.
95    ///
96    /// Honoured in both directions: the framer is configured from the same value.
97    pub fn with_terminator(mut self, terminator: csv::Terminator) -> Self {
98        self.terminator = terminator;
99        self
100    }
101
102    /// Set the field delimiter to use.
103    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
104        self.delimiter = delimiter;
105        self
106    }
107
108    /// Set whether to write headers.
109    pub fn with_has_headers(mut self, has_headers: bool) -> Self {
110        self.has_headers = has_headers;
111        self
112    }
113
114    fn writer_builder(&self, write_headers: bool) -> csv::WriterBuilder {
115        let mut builder = csv::WriterBuilder::new();
116        builder
117            .has_headers(write_headers)
118            .delimiter(self.delimiter)
119            .flexible(self.flexible)
120            .quote_style(self.quote_style)
121            .quote(self.quote)
122            .double_quote(self.double_quote)
123            .escape(self.escape)
124            .terminator(self.terminator);
125        builder
126    }
127
128    fn frame_config(&self) -> CsvFrameConfig {
129        CsvFrameConfig {
130            delimiter: self.delimiter,
131            quote: self.quote,
132            double_quote: self.double_quote,
133            // `csv`'s writer escapes by doubling unless `double_quote` is off, so the escape
134            // character only applies in the other case.
135            escape: if self.double_quote {
136                None
137            } else {
138                Some(self.escape)
139            },
140            terminator: match self.terminator {
141                csv::Terminator::CRLF => csv_core::Terminator::CRLF,
142                csv::Terminator::Any(b) => csv_core::Terminator::Any(b),
143                _ => csv_core::Terminator::Any(b'\n'),
144            },
145        }
146    }
147}
148
149impl StreamFormat for CsvStreamFormat {
150    fn format_name(&self) -> &'static str {
151        "csv"
152    }
153
154    fn default_content_type(&self) -> &'static str {
155        CSV_CONTENT_TYPE
156    }
157
158    fn accepts_content_type(&self, ct: &ContentType<'_>) -> bool {
159        ct.matches_any(CSV_ALIASES)
160    }
161}
162
163/// Per-stream state for [`CsvStreamFormat`].
164///
165/// A fresh `csv::Writer` is built for every row, because `csv` writes its header from the
166/// first serialised record and there is no way to ask an existing writer to stop. The header
167/// therefore has to be produced by a writer configured for it, and every later row by one that
168/// is not.
169pub struct CsvEncoder {
170    format: CsvStreamFormat,
171}
172
173impl<T> ItemEncoder<T> for CsvEncoder
174where
175    T: Serialize,
176{
177    fn encode(&mut self, item: &T, index: u64, buf: &mut BytesMut) -> Result<(), StreamError> {
178        let write_headers = index == 0 && self.format.has_headers;
179        let mut writer = self
180            .format
181            .writer_builder(write_headers)
182            .from_writer(Vec::new());
183
184        writer
185            .serialize(item)
186            .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))?;
187        writer
188            .flush()
189            .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))?;
190
191        let bytes = writer
192            .into_inner()
193            .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))?;
194
195        buf.extend_from_slice(&bytes);
196        Ok(())
197    }
198}
199
200impl<T> StreamFormatEncode<T> for CsvStreamFormat
201where
202    T: Serialize,
203{
204    type Encoder = CsvEncoder;
205
206    fn encoder(&self) -> Self::Encoder {
207        CsvEncoder {
208            format: self.clone(),
209        }
210    }
211}
212
213/// Deserialises one framed CSV record.
214///
215/// Positional, not by header name: the header row is consumed by the framer and discarded,
216/// which is the behaviour this pair of crates has always had.
217#[derive(Debug, Clone, Copy, Default)]
218pub struct CsvParser;
219
220impl<T> FrameParser<csv::ByteRecord, T> for CsvParser
221where
222    T: for<'de> Deserialize<'de>,
223{
224    fn parse(&self, frame: csv::ByteRecord) -> Result<T, StreamError> {
225        frame.deserialize(None).map_err(|err| {
226            StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None)
227        })
228    }
229}
230
231impl<T> StreamFormatDecode<T> for CsvStreamFormat
232where
233    T: for<'de> Deserialize<'de>,
234{
235    type Frame = csv::ByteRecord;
236    type Framer = CsvRecordCodec;
237    type Parser = CsvParser;
238
239    fn framer(&self, options: &DecodeOptions) -> Self::Framer {
240        CsvRecordCodec::new(self.frame_config(), self.has_headers, options.max_obj_len)
241    }
242
243    fn parser(&self) -> Self::Parser {
244        CsvParser
245    }
246}