wowlab_types/
copy.rs

1//! Postgres `COPY ... FROM STDIN` text-format serialization.
2
3use compact_str::CompactString;
4use serde::Serialize;
5
6/// Encodes rows in Postgres `COPY ... FROM STDIN` text format.
7#[derive(Debug)]
8pub struct CopyRow {
9    buf: Vec<u8>,
10    row_start: bool,
11}
12
13impl CopyRow {
14    /// Create an empty COPY buffer.
15    #[must_use]
16    pub fn new() -> Self {
17        Self {
18            buf: Vec::new(),
19            row_start: true,
20        }
21    }
22
23    /// Encode `value` as the next column of the current row.
24    pub fn push_bind(&mut self, value: impl CopyField) -> &mut Self {
25        if !self.row_start {
26            self.buf.push(b'\t');
27        }
28
29        self.row_start = false;
30        value.encode(&mut self.buf);
31
32        self
33    }
34
35    /// Terminate the current row with a newline.
36    pub fn finish_row(&mut self) {
37        self.buf.push(b'\n');
38        self.row_start = true;
39    }
40
41    /// Current buffered byte length.
42    #[must_use]
43    pub fn len(&self) -> usize {
44        self.buf.len()
45    }
46
47    /// Whether the buffer is empty.
48    #[must_use]
49    pub fn is_empty(&self) -> bool {
50        self.buf.is_empty()
51    }
52
53    /// Take the buffered bytes, leaving the buffer empty.
54    pub fn take(&mut self) -> Vec<u8> {
55        std::mem::take(&mut self.buf)
56    }
57}
58
59impl Default for CopyRow {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65/// A value encodable as a `COPY` text-format field.
66pub trait CopyField {
67    /// Append this value's text-format encoding to `buf`.
68    fn encode(&self, buf: &mut Vec<u8>);
69}
70
71fn encode_text(s: &str, buf: &mut Vec<u8>) {
72    for &b in s.as_bytes() {
73        match b {
74            b'\\' => buf.extend_from_slice(b"\\\\"),
75            b'\t' => buf.extend_from_slice(b"\\t"),
76            b'\n' => buf.extend_from_slice(b"\\n"),
77            b'\r' => buf.extend_from_slice(b"\\r"),
78            _ => buf.push(b),
79        }
80    }
81}
82
83#[expect(
84    clippy::expect_used,
85    reason = "std::io::Write for Vec<u8> is infallible because extending a vector reports allocation failure by aborting"
86)]
87fn encode_display(value: impl std::fmt::Display, buf: &mut Vec<u8>) {
88    use std::io::Write;
89
90    write!(buf, "{value}").expect("writing to Vec<u8> is infallible");
91}
92
93impl<T> CopyField for &T
94where
95    T: CopyField + ?Sized,
96{
97    fn encode(&self, buf: &mut Vec<u8>) {
98        (**self).encode(buf);
99    }
100}
101
102impl<T> CopyField for Option<T>
103where
104    T: CopyField,
105{
106    fn encode(&self, buf: &mut Vec<u8>) {
107        match self {
108            Some(v) => v.encode(buf),
109            None => buf.extend_from_slice(b"\\N"),
110        }
111    }
112}
113
114impl CopyField for bool {
115    fn encode(&self, buf: &mut Vec<u8>) {
116        buf.push(if *self { b't' } else { b'f' });
117    }
118}
119
120macro_rules! copy_field_int {
121    ($($ty:ty),+) => {$(
122        impl CopyField for $ty {
123            fn encode(&self, buf: &mut Vec<u8>) {
124                encode_display(self, buf);
125            }
126        }
127    )+};
128}
129
130copy_field_int!(i16, i32, i64);
131
132macro_rules! copy_field_float {
133    ($($ty:ty),+) => {$(
134        impl CopyField for $ty {
135            fn encode(&self, buf: &mut Vec<u8>) {
136                if self.is_nan() {
137                    buf.extend_from_slice(b"NaN");
138                } else if self.is_infinite() {
139                    buf.extend_from_slice(if self.is_sign_positive() {
140                        b"Infinity"
141                    } else {
142                        b"-Infinity"
143                    });
144                } else {
145                    encode_display(self, buf);
146                }
147            }
148        }
149    )+};
150}
151
152copy_field_float!(f32, f64);
153
154impl CopyField for str {
155    fn encode(&self, buf: &mut Vec<u8>) {
156        encode_text(self, buf);
157    }
158}
159
160impl CopyField for String {
161    fn encode(&self, buf: &mut Vec<u8>) {
162        encode_text(self, buf);
163    }
164}
165
166impl CopyField for CompactString {
167    fn encode(&self, buf: &mut Vec<u8>) {
168        encode_text(self.as_str(), buf);
169    }
170}
171
172impl CopyField for serde_json::Value {
173    fn encode(&self, buf: &mut Vec<u8>) {
174        encode_text(&self.to_string(), buf);
175    }
176}
177
178impl CopyField for [i32] {
179    fn encode(&self, buf: &mut Vec<u8>) {
180        buf.push(b'{');
181
182        for (i, v) in self.iter().enumerate() {
183            if i > 0 {
184                buf.push(b',');
185            }
186
187            encode_display(v, buf);
188        }
189
190        buf.push(b'}');
191    }
192}
193
194impl CopyField for Vec<i32> {
195    fn encode(&self, buf: &mut Vec<u8>) {
196        self.as_slice().encode(buf);
197    }
198}
199
200/// Serialize a value to a JSON `serde_json::Value`, falling back to `Null` on error.
201// #t(rust_pub_api_foreign_types) COPY serialization deliberately exposes serde_json's canonical value tree
202pub fn to_json<T>(value: &T) -> serde_json::Value
203where
204    T: Serialize,
205{
206    serde_json::to_value(value).unwrap_or(serde_json::Value::Null)
207}
208
209/// A row type that can stream itself into a Postgres `COPY ... FROM STDIN` text sink.
210pub trait CopyInsert {
211    /// Column names in bind order (includes the injected `patch_version`).
212    const COLUMNS: &'static [&'static str];
213    /// Push every column value onto `w` in `COLUMNS` order. `patch` fills the injected `patch_version` column.
214    fn copy_row(&self, w: &mut CopyRow, patch: &str);
215}
216
217#[cfg(test)]
218mod tests {
219    use googletest::prelude::*;
220    use rstest::rstest;
221
222    use super::*;
223
224    fn encoded(field: impl CopyField) -> Result<String> {
225        let mut buf = Vec::new();
226
227        field.encode(&mut buf);
228
229        String::from_utf8(buf).or_fail()
230    }
231
232    #[gtest]
233    #[rstest]
234    #[case::plain("abc", "abc")]
235    #[case::backslash("a\\b", "a\\\\b")]
236    #[case::tab("a\tb", "a\\tb")]
237    #[case::newline("a\nb", "a\\nb")]
238    #[case::carriage("a\rb", "a\\rb")]
239    #[case::all("\\\t\n\r", "\\\\\\t\\n\\r")]
240    #[case::empty("", "")]
241    fn str_encode_cases(#[case] input: &str, #[case] expected: &str) -> Result<()> {
242        verify_that!(encoded(input)?, eq(expected))
243    }
244
245    #[gtest]
246    fn bool_encode() -> Result<()> {
247        verify_that!(encoded(true)?, eq("t"))?;
248
249        verify_that!(encoded(false)?, eq("f"))
250    }
251
252    #[gtest]
253    #[rstest]
254    #[case::normal(1.5, "1.5")]
255    #[case::nan(f64::NAN, "NaN")]
256    #[case::pos_inf(f64::INFINITY, "Infinity")]
257    #[case::neg_inf(f64::NEG_INFINITY, "-Infinity")]
258    fn f64_encode_cases(#[case] input: f64, #[case] expected: &str) -> Result<()> {
259        verify_that!(encoded(input)?, eq(expected))
260    }
261
262    #[gtest]
263    fn f64_encode_neg_zero() -> Result<()> {
264        verify_that!(encoded(-0.0f64)?, eq(&format!("{}", -0.0f64)))
265    }
266
267    #[gtest]
268    #[rstest]
269    #[case::normal(1.5, "1.5")]
270    #[case::nan(f32::NAN, "NaN")]
271    #[case::pos_inf(f32::INFINITY, "Infinity")]
272    #[case::neg_inf(f32::NEG_INFINITY, "-Infinity")]
273    fn f32_encode_cases(#[case] input: f32, #[case] expected: &str) -> Result<()> {
274        verify_that!(encoded(input)?, eq(expected))
275    }
276
277    #[gtest]
278    fn option_encode() -> Result<()> {
279        verify_that!(encoded(Some(5i32))?, eq("5"))?;
280
281        verify_that!(encoded(None::<i32>)?, eq("\\N"))
282    }
283
284    #[gtest]
285    #[rstest]
286    #[case::empty(vec![], "{}")]
287    #[case::single(vec![7], "{7}")]
288    #[case::many(vec![1, -2, 3], "{1,-2,3}")]
289    fn vec_i32_encode_cases(#[case] input: Vec<i32>, #[case] expected: &str) -> Result<()> {
290        verify_that!(encoded(input.as_slice())?, eq(expected))?;
291
292        verify_that!(encoded(input)?, eq(expected))
293    }
294
295    #[gtest]
296    fn copy_row_assembly() -> Result<()> {
297        let mut row = CopyRow::new();
298
299        row.push_bind(1i32)
300            .push_bind("a")
301            .push_bind(Option::<i32>::None);
302        row.finish_row();
303        let bytes = row.take();
304
305        verify_that!(String::from_utf8(bytes).or_fail()?, eq("1\ta\t\\N\n"))?;
306        verify_that!(row.is_empty(), eq(true))?;
307
308        verify_that!(row.len(), eq(0))
309    }
310
311    #[gtest]
312    fn copy_row_two_rows() -> Result<()> {
313        let mut row = CopyRow::new();
314
315        row.push_bind(1i32).push_bind("a");
316        row.finish_row();
317        row.push_bind(2i32).push_bind("b");
318        row.finish_row();
319        let out = String::from_utf8(row.take()).or_fail()?;
320
321        verify_that!(out, eq("1\ta\n2\tb\n"))
322    }
323}
wowlab_types/copy.rs

wowlab_types/
copy.rs
