Skip to main content

sentinel_driver/copy/
text.rs

1use crate::error::{Error, Result};
2
3/// Encoder for text COPY IN format.
4///
5/// Text COPY format: tab-separated fields, newline-separated rows.
6/// Special values: `\N` for NULL, backslash escaping for special chars.
7///
8/// # Example
9///
10/// ```rust
11/// use sentinel_driver::copy::text::TextCopyEncoder;
12///
13/// let mut encoder = TextCopyEncoder::new();
14/// encoder.add_row(&[Some("42"), Some("hello world")]);
15/// encoder.add_row(&[Some("7"), None]); // NULL value
16///
17/// let data = encoder.finish();
18/// ```
19pub 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    /// Add a row with the given field values.
31    ///
32    /// `None` represents NULL (encoded as `\N`).
33    /// Values are tab-separated, rows are newline-separated.
34    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    /// Finish encoding and return the text COPY data.
48    pub fn finish(self) -> Vec<u8> {
49        self.buf.into_bytes()
50    }
51
52    /// Get the current buffer size.
53    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
68/// Escape a text value for COPY format.
69///
70/// Backslash, tab, newline, and carriage return need escaping.
71fn 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
83/// Decoder for text COPY OUT format.
84///
85/// Parses tab-separated, newline-separated text data.
86pub struct TextCopyDecoder;
87
88impl TextCopyDecoder {
89    /// Parse a single line of text COPY data into field values.
90    ///
91    /// Returns `None` for NULL fields (`\N`).
92    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    /// Parse multiple lines of text COPY data.
107    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
121/// Unescape a text COPY field value.
122fn 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                    // Should not happen here (handled at field level)
135                    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}