1use bytes::BufMut;
6use std::{
7 fmt::{self},
8 marker::PhantomData,
9};
10
11const COMMA_EQ_SPACE: [char; 3] = [',', '=', ' '];
13const COMMA_SPACE: [char; 2] = [',', ' '];
14const DOUBLE_QUOTE: [char; 1] = ['"'];
15
16#[doc(hidden)]
17#[derive(Clone, Copy, Debug, Default)]
18pub struct BeforeMeasurement;
19#[doc(hidden)]
20#[derive(Clone, Copy, Debug)]
21pub struct AfterMeasurement;
22#[doc(hidden)]
23#[derive(Clone, Copy, Debug)]
24pub struct AfterTag;
25#[doc(hidden)]
26#[derive(Clone, Copy, Debug)]
27pub struct AfterField;
28#[doc(hidden)]
29#[derive(Clone, Copy, Debug)]
30pub struct AfterTimestamp;
31
32#[derive(Debug, Default)]
105pub struct LineProtocolBuilder<B, S = BeforeMeasurement>
106where
107 B: BufMut,
108{
109 buf: B,
110 _marker: PhantomData<S>,
111}
112
113impl LineProtocolBuilder<Vec<u8>, BeforeMeasurement> {
114 pub fn new() -> Self {
116 Self::new_with(vec![])
117 }
118}
119
120impl<B> LineProtocolBuilder<B, BeforeMeasurement>
121where
122 B: BufMut,
123{
124 pub fn new_with(buf: B) -> Self {
126 Self {
127 buf,
128 _marker: PhantomData,
129 }
130 }
131
132 pub fn measurement(self, measurement: &str) -> LineProtocolBuilder<B, AfterMeasurement> {
136 let measurement = escape(measurement, COMMA_SPACE);
137 self.write(format_args!("{measurement}"))
138 }
139
140 pub fn build(self) -> B {
142 self.buf
143 }
144}
145
146impl<B> LineProtocolBuilder<B, AfterMeasurement>
147where
148 B: BufMut,
149{
150 pub fn tag(self, tag_key: &str, tag_value: &str) -> Self {
156 let tag_key = escape(tag_key, COMMA_EQ_SPACE);
157 let tag_value = escape(tag_value, COMMA_EQ_SPACE);
158 self.write(format_args!(",{tag_key}={tag_value}"))
159 }
160
161 pub fn field<F>(self, field_key: &str, field_value: F) -> LineProtocolBuilder<B, AfterField>
172 where
173 F: FieldValue,
174 {
175 self.write(format_args!(" {}", format_field(field_key, &field_value)))
176 }
177}
178
179impl<B> LineProtocolBuilder<B, AfterField>
180where
181 B: BufMut,
182{
183 pub fn field<F: FieldValue>(self, field_key: &str, field_value: F) -> Self {
187 self.write(format_args!(",{}", format_field(field_key, &field_value)))
188 }
189
190 pub fn timestamp(self, ts: i64) -> LineProtocolBuilder<B, AfterTimestamp> {
198 self.write(format_args!(" {ts}"))
199 }
200
201 pub fn close_line(self) -> LineProtocolBuilder<B, BeforeMeasurement> {
203 self.close()
204 }
205}
206
207impl<B> LineProtocolBuilder<B, AfterTimestamp>
208where
209 B: BufMut,
210{
211 pub fn close_line(self) -> LineProtocolBuilder<B, BeforeMeasurement> {
213 self.close()
214 }
215}
216
217impl<B, S> LineProtocolBuilder<B, S>
218where
219 B: BufMut,
220{
221 fn close(self) -> LineProtocolBuilder<B, BeforeMeasurement> {
222 self.write(format_args!("\n"))
223 }
224
225 fn write<S2>(self, args: fmt::Arguments<'_>) -> LineProtocolBuilder<B, S2> {
226 use std::io::Write;
227 let mut writer = self.buf.writer();
229 write!(&mut writer, "{args}").unwrap();
230 LineProtocolBuilder {
231 buf: writer.into_inner(),
232 _marker: PhantomData,
233 }
234 }
235}
236
237fn escape<const N: usize>(src: &str, special_characters: [char; N]) -> Escaped<'_, N> {
240 Escaped {
241 src,
242 special_characters,
243 }
244}
245
246struct Escaped<'a, const N: usize> {
247 src: &'a str,
248 special_characters: [char; N],
249}
250
251impl<'a, const N: usize> fmt::Display for Escaped<'a, N> {
252 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253 for ch in self.src.chars() {
254 if self.special_characters.contains(&ch) || ch == '\\' {
255 write!(f, "\\")?;
256 }
257 write!(f, "{ch}")?;
258 }
259 Ok(())
260 }
261}
262
263fn format_field<'a, F>(field_key: &'a str, field_value: &'a F) -> impl fmt::Display + 'a
266where
267 F: FieldValue,
268{
269 FormattedField {
270 field_key,
271 field_value,
272 }
273}
274
275struct FormattedField<'a, F>
276where
277 F: FieldValue,
278{
279 field_key: &'a str,
280 field_value: &'a F,
281}
282
283impl<'a, F> fmt::Display for FormattedField<'a, F>
284where
285 F: FieldValue,
286{
287 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
288 write!(f, "{}=", escape(self.field_key, COMMA_EQ_SPACE))?;
289 self.field_value.fmt(f)
290 }
291}
292
293pub trait FieldValue {
297 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
298}
299
300impl FieldValue for &str {
301 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
302 write!(f, "\"{}\"", escape(self, DOUBLE_QUOTE))
303 }
304}
305
306impl FieldValue for f64 {
307 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
308 write!(f, "{self}")
309 }
310}
311
312impl FieldValue for bool {
313 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
314 write!(f, "{self}")
315 }
316}
317
318impl FieldValue for i64 {
319 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
320 write!(f, "{self}i")
321 }
322}
323
324impl FieldValue for u64 {
325 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
326 write!(f, "{self}u")
327 }
328}
329
330#[cfg(test)]
331mod tests {
332 use crate::{parse_lines, FieldSet, ParsedLine};
333
334 use super::*;
335
336 #[test]
337 fn test_string_escape() {
338 assert_eq!(
339 format!("\"{}\"", escape(r#"foo"#, DOUBLE_QUOTE)),
340 r#""foo""#
341 );
342 assert_eq!(
343 format!("\"{}\"", escape(r"foo \ bar", DOUBLE_QUOTE)),
344 r#""foo \\ bar""#
345 );
346 assert_eq!(
347 format!("\"{}\"", escape(r#"foo " bar"#, DOUBLE_QUOTE)),
348 r#""foo \" bar""#
349 );
350 assert_eq!(
351 format!("\"{}\"", escape(r#"foo \" bar"#, DOUBLE_QUOTE)),
352 r#""foo \\\" bar""#
353 );
354 }
355
356 #[test]
357 fn test_lp_builder() {
358 const PLAIN: &str = "plain";
359 const WITH_SPACE: &str = "with space";
360 const WITH_COMMA: &str = "with,comma";
361 const WITH_EQ: &str = "with=eq";
362 const WITH_DOUBLE_QUOTE: &str = r#"with"doublequote"#;
363 const WITH_SINGLE_QUOTE: &str = "with'singlequote";
364 const WITH_BACKSLASH: &str = r"with\ backslash";
365
366 let builder = LineProtocolBuilder::new()
367 .measurement("tag_keys")
369 .tag(PLAIN, "dummy")
370 .tag(WITH_SPACE, "dummy")
371 .tag(WITH_COMMA, "dummy")
372 .tag(WITH_EQ, "dummy")
373 .tag(WITH_DOUBLE_QUOTE, "dummy")
374 .tag(WITH_SINGLE_QUOTE, "dummy")
375 .tag(WITH_BACKSLASH, "dummy")
376 .field("dummy", true)
377 .close_line()
378 .measurement("tag_values")
380 .tag("plain", PLAIN)
381 .tag("withspace", WITH_SPACE)
382 .tag("withcomma", WITH_COMMA)
383 .tag("witheq", WITH_EQ)
384 .tag("withdoublequote", WITH_DOUBLE_QUOTE)
385 .tag("withsinglaquote", WITH_SINGLE_QUOTE)
386 .tag("withbackslash", WITH_BACKSLASH)
387 .field("dummy", true)
388 .close_line()
389 .measurement("field keys")
391 .field(PLAIN, true)
392 .field(WITH_SPACE, true)
393 .field(WITH_COMMA, true)
394 .field(WITH_EQ, true)
395 .field(WITH_DOUBLE_QUOTE, true)
396 .field(WITH_SINGLE_QUOTE, true)
397 .field(WITH_BACKSLASH, true)
398 .close_line()
399 .measurement("field values")
401 .field("mybool", false)
402 .field("mysigned", 51_i64)
403 .field("myunsigned", 51_u64)
404 .field("myfloat", 51.0)
405 .field("mystring", "some value")
406 .field("mystringwithquotes", "some \" value")
407 .close_line()
408 .measurement(PLAIN)
410 .field("dummy", true)
411 .close_line()
412 .measurement(WITH_SPACE)
414 .field("dummy", true)
415 .close_line()
416 .measurement(WITH_COMMA)
418 .field("dummy", true)
419 .close_line()
420 .measurement(WITH_EQ)
422 .field("dummy", true)
423 .close_line()
424 .measurement(WITH_DOUBLE_QUOTE)
426 .field("dummy", true)
427 .close_line()
428 .measurement(WITH_SINGLE_QUOTE)
430 .field("dummy", true)
431 .close_line()
432 .measurement(WITH_BACKSLASH)
434 .field("dummy", true)
435 .close_line()
436 .measurement("without timestamp")
438 .field("dummy", true)
439 .close_line()
440 .measurement("with timestamp")
442 .field("dummy", true)
443 .timestamp(1234)
444 .close_line();
445
446 let lp = String::from_utf8(builder.build()).unwrap();
447 println!("-----\n{lp}-----");
448
449 let parsed_lines = parse_lines(&lp)
450 .collect::<Result<Vec<ParsedLine<'_>>, _>>()
451 .unwrap();
452
453 let get_tag_key = |n: usize, f: usize| {
454 format!("{}", parsed_lines[n].series.tag_set.as_ref().unwrap()[f].0)
455 };
456 let row = 0;
457 assert_eq!(get_tag_key(row, 0), PLAIN);
458 assert_eq!(get_tag_key(row, 1), WITH_SPACE);
459 assert_eq!(get_tag_key(row, 2), WITH_COMMA);
460 assert_eq!(get_tag_key(row, 3), WITH_EQ);
461 assert_eq!(get_tag_key(row, 4), WITH_DOUBLE_QUOTE);
462 assert_eq!(get_tag_key(row, 5), WITH_SINGLE_QUOTE);
463 assert_eq!(get_tag_key(row, 6), WITH_BACKSLASH);
464
465 let get_tag_value = |n: usize, f: usize| {
466 format!("{}", parsed_lines[n].series.tag_set.as_ref().unwrap()[f].1)
467 };
468 let row = 1;
469 assert_eq!(get_tag_value(row, 0), PLAIN);
470 assert_eq!(get_tag_value(row, 1), WITH_SPACE);
471 assert_eq!(get_tag_value(row, 2), WITH_COMMA);
472 assert_eq!(get_tag_value(row, 3), WITH_EQ);
473 assert_eq!(get_tag_value(row, 4), WITH_DOUBLE_QUOTE);
474 assert_eq!(get_tag_value(row, 5), WITH_SINGLE_QUOTE);
475 assert_eq!(get_tag_value(row, 6), WITH_BACKSLASH);
476
477 let get_field_key = |n: usize, f: usize| format!("{}", parsed_lines[n].field_set[f].0);
478 let row = 2;
479 assert_eq!(get_field_key(row, 0), PLAIN);
480 assert_eq!(get_field_key(row, 1), WITH_SPACE);
481 assert_eq!(get_field_key(row, 2), WITH_COMMA);
482 assert_eq!(get_field_key(row, 3), WITH_EQ);
483 assert_eq!(get_field_key(row, 4), WITH_DOUBLE_QUOTE);
484 assert_eq!(get_field_key(row, 5), WITH_SINGLE_QUOTE);
485 assert_eq!(get_field_key(row, 6), WITH_BACKSLASH);
486
487 let get_field_value = |n: usize, f: usize| format!("{}", parsed_lines[n].field_set[f].1);
488 let row = 3;
489 assert_eq!(get_field_value(row, 0), "false");
490 assert_eq!(get_field_value(row, 1), "51i");
491 assert_eq!(get_field_value(row, 2), "51u");
492 assert_eq!(get_field_value(row, 3), "51");
493 assert_eq!(get_field_value(row, 4), "some value");
494 let get_measurement = |n: usize| format!("{}", parsed_lines[n].series.measurement);
498 assert_eq!(get_measurement(4), PLAIN);
499 assert_eq!(get_measurement(5), WITH_SPACE);
500 assert_eq!(get_measurement(6), WITH_COMMA);
501 assert_eq!(get_measurement(7), WITH_EQ);
502 assert_eq!(get_measurement(8), WITH_DOUBLE_QUOTE);
503 assert_eq!(get_measurement(9), WITH_SINGLE_QUOTE);
504 assert_eq!(get_measurement(10), WITH_BACKSLASH);
505
506 let get_timestamp = |n: usize| parsed_lines[n].timestamp;
507 assert_eq!(get_timestamp(11), None);
508 assert_eq!(get_timestamp(12), Some(1234));
509 }
510
511 #[test]
512 fn test_float_formatting() {
513 let builder = LineProtocolBuilder::new()
516 .measurement("tag_keys")
517 .tag("foo", "bar")
518 .field("my_float", 3.0)
519 .close_line();
520
521 let lp = String::from_utf8(builder.build()).unwrap();
522 println!("-----\n{lp}-----");
523
524 let parsed_lines = parse_lines(&lp)
525 .collect::<Result<Vec<ParsedLine<'_>>, _>>()
526 .unwrap();
527
528 assert_eq!(parsed_lines.len(), 1);
529 let parsed_line = &parsed_lines[0];
530
531 let expected_fields = vec![("my_float".into(), crate::FieldValue::F64(3.0))]
532 .into_iter()
533 .collect::<FieldSet<'_>>();
534
535 assert_eq!(parsed_line.field_set, expected_fields)
536 }
537}