sentinel_driver/copy/
text.rs1use crate::error::{Error, Result};
2
3pub struct TextCopyEncoder {
20 buf: String,
21}
22
23impl TextCopyEncoder {
24 pub fn new() -> Self {
25 Self {
26 buf: String::with_capacity(8192),
27 }
28 }
29
30 pub fn add_row(&mut self, fields: &[Option<&str>]) {
35 for (i, field) in fields.iter().enumerate() {
36 if i > 0 {
37 self.buf.push('\t');
38 }
39 match field {
40 Some(val) => escape_text_value(&mut self.buf, val),
41 None => self.buf.push_str("\\N"),
42 }
43 }
44 self.buf.push('\n');
45 }
46
47 pub fn finish(self) -> Vec<u8> {
49 self.buf.into_bytes()
50 }
51
52 pub fn len(&self) -> usize {
54 self.buf.len()
55 }
56
57 pub fn is_empty(&self) -> bool {
58 self.buf.is_empty()
59 }
60}
61
62impl Default for TextCopyEncoder {
63 fn default() -> Self {
64 Self::new()
65 }
66}
67
68fn escape_text_value(buf: &mut String, val: &str) {
72 for ch in val.chars() {
73 match ch {
74 '\\' => buf.push_str("\\\\"),
75 '\t' => buf.push_str("\\t"),
76 '\n' => buf.push_str("\\n"),
77 '\r' => buf.push_str("\\r"),
78 other => buf.push(other),
79 }
80 }
81}
82
83pub struct TextCopyDecoder;
87
88impl TextCopyDecoder {
89 pub fn parse_line(line: &str) -> Result<Vec<Option<String>>> {
93 let mut fields = Vec::new();
94
95 for raw_field in line.split('\t') {
96 if raw_field == "\\N" {
97 fields.push(None);
98 } else {
99 fields.push(Some(unescape_text_value(raw_field)?));
100 }
101 }
102
103 Ok(fields)
104 }
105
106 pub fn parse_all(data: &str) -> Result<Vec<Vec<Option<String>>>> {
108 let mut rows = Vec::new();
109
110 for line in data.lines() {
111 if line.is_empty() {
112 continue;
113 }
114 rows.push(Self::parse_line(line)?);
115 }
116
117 Ok(rows)
118 }
119}
120
121fn unescape_text_value(val: &str) -> Result<String> {
123 let mut result = String::with_capacity(val.len());
124 let mut chars = val.chars();
125
126 while let Some(ch) = chars.next() {
127 if ch == '\\' {
128 match chars.next() {
129 Some('\\') | None => result.push('\\'),
130 Some('t') => result.push('\t'),
131 Some('n') => result.push('\n'),
132 Some('r') => result.push('\r'),
133 Some('N') => {
134 return Err(Error::Copy("unexpected \\N inside field".into()));
136 }
137 Some(other) => {
138 result.push('\\');
139 result.push(other);
140 }
141 }
142 } else {
143 result.push(ch);
144 }
145 }
146
147 Ok(result)
148}