1use crate::content_type::ContentType;
4use crate::envelope::StreamFormatEnvelope;
5use crate::error::{StreamError, StreamErrorKind};
6use crate::format::{
7 DecodeOptions, DefaultFormat, IdentityParser, ItemEncoder, StreamFormat, StreamFormatDecode,
8 StreamFormatEncode,
9};
10use crate::json_array_codec::JsonArrayCodec;
11use crate::json_nl_codec::JsonNewLineCodec;
12use bytes::{BufMut, BytesMut};
13use serde::{Deserialize, Serialize};
14use std::io::Write;
15
16const JSON_ARRAY_BEGIN: &[u8] = b"[";
17const JSON_ARRAY_END: &[u8] = b"]";
18const JSON_ARRAY_ENVELOPE_END: &[u8] = b"]}";
19const JSON_SEP: &[u8] = b",";
20const JSON_NL_SEP: &[u8] = b"\n";
21
22const JSON_ARRAY_CONTENT_TYPE: &str = "application/json";
24
25const JSON_NL_CONTENT_TYPE: &str = "application/jsonstream";
30const JSON_NL_ALIASES: &[&str] = &[
31 "application/jsonstream",
32 "application/x-ndjson",
33 "application/ndjson",
34 "application/jsonl",
35 "application/x-jsonl",
36];
37
38#[derive(Debug, Clone)]
40pub struct JsonArrayStreamFormat<E = ()>
41where
42 E: Serialize,
43{
44 envelope: Option<StreamFormatEnvelope<E>>,
45}
46
47impl JsonArrayStreamFormat {
48 pub fn new() -> JsonArrayStreamFormat<()> {
50 JsonArrayStreamFormat { envelope: None }
51 }
52
53 pub fn with_envelope<E>(envelope: E, array_field: &str) -> JsonArrayStreamFormat<E>
57 where
58 E: Serialize,
59 {
60 JsonArrayStreamFormat {
61 envelope: Some(StreamFormatEnvelope {
62 object: envelope,
63 array_field: array_field.to_string(),
64 }),
65 }
66 }
67}
68
69impl Default for JsonArrayStreamFormat<()> {
70 fn default() -> Self {
71 JsonArrayStreamFormat { envelope: None }
72 }
73}
74
75impl DefaultFormat for JsonArrayStreamFormat<()> {
76 fn default_format() -> Self {
77 Self::default()
78 }
79}
80
81impl<E> StreamFormat for JsonArrayStreamFormat<E>
82where
83 E: Serialize,
84{
85 fn format_name(&self) -> &'static str {
86 "json_array"
87 }
88
89 fn default_content_type(&self) -> &'static str {
90 JSON_ARRAY_CONTENT_TYPE
91 }
92
93 fn accepts_content_type(&self, ct: &ContentType<'_>) -> bool {
94 ct.matches(JSON_ARRAY_CONTENT_TYPE) || ct.has_suffix("json")
95 }
96}
97
98pub struct JsonArrayEncoder {
105 prologue: Result<Vec<u8>, Option<StreamError>>,
106 epilogue: &'static [u8],
107}
108
109impl JsonArrayEncoder {
110 fn bare() -> Self {
111 Self {
112 prologue: Ok(JSON_ARRAY_BEGIN.to_vec()),
113 epilogue: JSON_ARRAY_END,
114 }
115 }
116
117 fn enveloped<E: Serialize>(envelope: &StreamFormatEnvelope<E>) -> Self {
118 let prologue = match serde_json::to_vec(&envelope.object) {
122 Ok(bytes) if bytes.len() > 1 => {
123 let mut buf = Vec::with_capacity(bytes.len() + envelope.array_field.len() + 4);
124 buf.extend_from_slice(&bytes[0..bytes.len() - 1]);
125 if bytes.len() > 2 {
128 buf.extend_from_slice(JSON_SEP);
129 }
130 buf.extend_from_slice(format!("\"{}\":", envelope.array_field).as_bytes());
131 buf.extend_from_slice(JSON_ARRAY_BEGIN);
132 Ok(buf)
133 }
134 Ok(bytes) => Err(Some(StreamError::new(
135 StreamErrorKind::CodecError,
136 None,
137 Some(format!("Too short envelope: {bytes:?}")),
138 ))),
139 Err(err) => Err(Some(StreamError::new(
140 StreamErrorKind::CodecError,
141 Some(Box::new(err)),
142 None,
143 ))),
144 };
145
146 Self {
147 prologue,
148 epilogue: JSON_ARRAY_ENVELOPE_END,
149 }
150 }
151}
152
153impl<T> ItemEncoder<T> for JsonArrayEncoder
154where
155 T: Serialize,
156{
157 fn prologue(&mut self, buf: &mut BytesMut) -> Result<(), StreamError> {
158 match &mut self.prologue {
159 Ok(bytes) => {
160 buf.extend_from_slice(bytes);
161 Ok(())
162 }
163 Err(err) => Err(err.take().unwrap_or_else(|| {
166 StreamError::new(StreamErrorKind::CodecError, None, Some("Bad envelope".into()))
167 })),
168 }
169 }
170
171 fn encode(&mut self, item: &T, index: u64, buf: &mut BytesMut) -> Result<(), StreamError> {
172 let mut writer = buf.writer();
173 if index != 0 {
174 writer
175 .write_all(JSON_SEP)
176 .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))?;
177 }
178 serde_json::to_writer(&mut writer, item)
179 .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))
180 }
181
182 fn epilogue(&mut self, buf: &mut BytesMut) -> Result<(), StreamError> {
183 buf.extend_from_slice(self.epilogue);
184 Ok(())
185 }
186}
187
188impl<T, E> StreamFormatEncode<T> for JsonArrayStreamFormat<E>
189where
190 T: Serialize,
191 E: Serialize,
192{
193 type Encoder = JsonArrayEncoder;
194
195 fn encoder(&self) -> Self::Encoder {
196 match &self.envelope {
197 Some(envelope) => JsonArrayEncoder::enveloped(envelope),
198 None => JsonArrayEncoder::bare(),
199 }
200 }
201}
202
203impl<T, E> StreamFormatDecode<T> for JsonArrayStreamFormat<E>
204where
205 T: for<'de> Deserialize<'de>,
206 E: Serialize,
207{
208 type Frame = Result<T, StreamError>;
209 type Framer = JsonArrayCodec<T>;
210 type Parser = IdentityParser;
211
212 fn framer(&self, options: &DecodeOptions) -> Self::Framer {
213 JsonArrayCodec::new_with_max_length(options.max_obj_len)
214 }
215
216 fn parser(&self) -> Self::Parser {
217 IdentityParser
218 }
219}
220
221#[derive(Debug, Clone, Copy, Default)]
223pub struct JsonNewLineStreamFormat;
224
225impl JsonNewLineStreamFormat {
226 pub fn new() -> Self {
228 Self
229 }
230}
231
232impl DefaultFormat for JsonNewLineStreamFormat {
233 fn default_format() -> Self {
234 Self
235 }
236}
237
238impl StreamFormat for JsonNewLineStreamFormat {
239 fn format_name(&self) -> &'static str {
240 "json_nl"
241 }
242
243 fn default_content_type(&self) -> &'static str {
244 JSON_NL_CONTENT_TYPE
245 }
246
247 fn accepts_content_type(&self, ct: &ContentType<'_>) -> bool {
248 ct.matches_any(JSON_NL_ALIASES)
249 }
250}
251
252#[derive(Debug, Clone, Copy, Default)]
254pub struct JsonNewLineEncoder;
255
256impl<T> ItemEncoder<T> for JsonNewLineEncoder
257where
258 T: Serialize,
259{
260 fn encode(&mut self, item: &T, _index: u64, buf: &mut BytesMut) -> Result<(), StreamError> {
261 let mut writer = buf.writer();
262 serde_json::to_writer(&mut writer, item)
263 .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))?;
264 writer
265 .write_all(JSON_NL_SEP)
266 .map_err(|err| StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None))
267 }
268}
269
270impl<T> StreamFormatEncode<T> for JsonNewLineStreamFormat
271where
272 T: Serialize,
273{
274 type Encoder = JsonNewLineEncoder;
275
276 fn encoder(&self) -> Self::Encoder {
277 JsonNewLineEncoder
278 }
279}
280
281impl<T> StreamFormatDecode<T> for JsonNewLineStreamFormat
282where
283 T: for<'de> Deserialize<'de>,
284{
285 type Frame = Result<T, StreamError>;
286 type Framer = JsonNewLineCodec<T>;
287 type Parser = IdentityParser;
288
289 fn framer(&self, options: &DecodeOptions) -> Self::Framer {
290 JsonNewLineCodec::new_with_max_length(options.max_obj_len)
291 }
292
293 fn parser(&self) -> Self::Parser {
294 IdentityParser
295 }
296}