csv: add additional csv builder settings

src/csv_format.rs

@@ -6,17 +6,87 @@ use http::HeaderMap;
 use serde::Serialize;
 
 pub struct CsvStreamFormat {
-    with_header: bool,
+    has_headers: bool,
     delimiter: u8,
+    flexible: bool,
+    quote_style: csv::QuoteStyle,
+    quote: u8,
+    double_quote: bool,
+    escape: u8,
+    terminator: csv::Terminator,
+}
+
+impl Default for CsvStreamFormat {
+    fn default() -> Self {
+        Self {
+            has_headers: true,
+            delimiter: b',',
+            flexible: false,
+            quote_style: csv::QuoteStyle::Necessary,
+            quote: b'"',
+            double_quote: true,
+            escape: b'\\',
+            terminator: csv::Terminator::Any(b'\n'),
+        }
+    }
 }
 
 impl CsvStreamFormat {
-    pub fn new(with_header: bool, delimiter: u8) -> Self {
+    pub fn new(has_headers: bool, delimiter: u8) -> Self {
         Self {
-            with_header,
+            has_headers,
             delimiter,
+            ..Default::default()
         }
     }
+
+    /// Sets whether to use flexible serialize.
+    pub fn with_flexible(mut self, flexible: bool) -> Self {
+        self.flexible = flexible;
+        self
+    }
+
+    /// Sets the quote style to use.
+    pub fn with_quote_style(mut self, quote_style: csv::QuoteStyle) -> Self {
+        self.quote_style = quote_style;
+        self
+    }
+
+    /// Sets the quote character to use.
+    pub fn with_quote(mut self, quote: u8) -> Self {
+        self.quote = quote;
+        self
+    }
+
+    /// Sets whether to double quote.
+    pub fn with_double_quote(mut self, double_quote: bool) -> Self {
+        self.double_quote = double_quote;
+        self
+    }
+
+    /// Sets the escape character to use.
+    pub fn with_escape(mut self, escape: u8) -> Self {
+        self.escape = escape;
+        self
+    }
+
+    /// Sets the line terminator to use.
+    pub fn with_terminator(mut self, terminator: csv::Terminator) -> Self {
+        self.terminator = terminator;
+        self
+    }
+
+    /// Set the field delimiter to use.
+    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
+        self.delimiter = delimiter;
+        self
+    }
+
+    /// Set whether to write headers.
+    pub fn with_has_headers(mut self, has_headers: bool) -> Self {
+        self.has_headers = has_headers;
+        self
+    }
 }
 
 impl<T> StreamingFormat<T> for CsvStreamFormat
@@ -27,30 +97,34 @@ where
         &'a self,
         stream: BoxStream<'b, T>,
     ) -> BoxStream<'b, Result<axum::body::Bytes, axum::Error>> {
-        fn write_csv_record<T>(obj: T, header: bool, delimiter: u8) -> Result<Vec<u8>, axum::Error>
-        where
-            T: Serialize,
-        {
-            let mut writer = csv::WriterBuilder::new()
-                .has_headers(header)
-                .delimiter(delimiter)
-                .from_writer(vec![]);
-            writer.serialize(obj).map_err(axum::Error::new)?;
-            writer.flush().map_err(axum::Error::new)?;
-            writer.into_inner().map_err(axum::Error::new)
-        }
-
-        let stream_with_header = self.with_header;
+        let stream_with_header = self.has_headers;
         let stream_delimiter = self.delimiter;
+        let stream_flexible = self.flexible;
+        let stream_quote_style = self.quote_style;
+        let stream_quote = self.quote;
+        let stream_double_quote = self.double_quote;
+        let stream_escape = self.escape;
+        let terminator = self.terminator;
+
         let stream_bytes: BoxStream<Result<axum::body::Bytes, axum::Error>> = Box::pin({
             stream.enumerate().map(move |(index, obj)| {
-                let write_csv_res = if index == 0 {
-                    write_csv_record(obj, stream_with_header, stream_delimiter)
-                } else {
-                    write_csv_record(obj, false, stream_delimiter)
-                };
-
-                write_csv_res.map(axum::body::Bytes::from)
+                let mut writer = csv::WriterBuilder::new()
+                    .has_headers(index == 0 && stream_with_header)
+                    .delimiter(stream_delimiter)
+                    .flexible(stream_flexible)
+                    .quote_style(stream_quote_style)
+                    .quote(stream_quote)
+                    .double_quote(stream_double_quote)
+                    .escape(stream_escape)
+                    .terminator(terminator)
+                    .from_writer(vec![]);
+
+                writer.serialize(obj).map_err(axum::Error::new)?;
+                writer.flush().map_err(axum::Error::new)?;
+                writer
+                    .into_inner()
+                    .map_err(axum::Error::new)
+                    .map(axum::body::Bytes::from)
             })
         });
 
@@ -103,7 +177,12 @@ mod tests {
 
         let app = Router::new().route(
             "/",
-            get(|| async { StreamBodyAs::new(CsvStreamFormat::new(false, b','), test_stream) }),
+            get(|| async {
+                StreamBodyAs::new(
+                    CsvStreamFormat::new(false, b'.').with_delimiter(b','),
+                    test_stream,
+                )
+            }),
         );
 
         let client = TestClient::new(app);