postgresql/
copy.rs

1//! Typed binary `COPY ... FROM STDIN` / `TO STDOUT`.
2//!
3//! Rows are written and read positionally, in the order of the `COPY` column
4//! list. [`ToRow`] is implemented for tuples and — via
5//! `#[derive(ToRow)]` — for structs, whose fields are encoded in declaration
6//! order. [`FromCopyRow`] is implemented for scalars and tuples; a struct is
7//! read field by field with [`FromCopyRow::from_copy_row`] or by decoding into
8//! a tuple first.
9
10use std::marker::PhantomData;
11
12use bytes::{Buf, Bytes, BytesMut};
13use futures_util::StreamExt;
14use postgresql_protocol::oid::Oid;
15
16use crate::{
17    connection::{COPY_HEADER, CopyInWriter, CopyOutStream, encode_copy_row},
18    decode::FromSql,
19    encode::ToSql,
20    error::Error,
21};
22
23/// A value that can be written as one binary `COPY` row.
24pub trait ToRow {
25    /// Appends the row (field count and framed fields) to `buf`.
26    fn encode_row(&self, buf: &mut BytesMut) -> Result<(), Error>;
27}
28
29impl<T: ToSql> ToRow for T {
30    fn encode_row(&self, buf: &mut BytesMut) -> Result<(), Error> {
31        encode_copy_row(&[self], buf)
32    }
33}
34
35macro_rules! to_row_tuple {
36    ($($idx:tt $name:ident),+) => {
37        impl<$($name: ToSql),+> ToRow for ($($name,)+) {
38            fn encode_row(&self, buf: &mut BytesMut) -> Result<(), Error> {
39                encode_copy_row(&[$( &self.$idx as &dyn ToSql ),+], buf)
40            }
41        }
42    };
43}
44
45to_row_tuple!(0 A);
46to_row_tuple!(0 A, 1 B);
47to_row_tuple!(0 A, 1 B, 2 C);
48to_row_tuple!(0 A, 1 B, 2 C, 3 D);
49to_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E);
50to_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F);
51to_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F, 6 G);
52to_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F, 6 G, 7 H);
53
54/// A value that can be decoded from one binary `COPY` row.
55pub trait FromCopyRow: Sized {
56    /// Decodes a row from its fields and their column OIDs.
57    fn from_copy_row(oids: &[Oid], fields: &[Option<Bytes>]) -> Result<Self, Error>;
58}
59
60fn decode_field<T: FromSql>(oids: &[Oid], fields: &[Option<Bytes>], index: usize) -> Result<T, Error> {
61    let oid = oids.get(index).copied().unwrap_or(0);
62    if oid != 0 && !T::accepts(oid) {
63        return Err(Error::Decode(format!(
64            "COPY column {index} has type OID {oid} which cannot be decoded as the requested type"
65        )));
66    }
67    match fields.get(index) {
68        None => Err(Error::Decode(format!("COPY row is missing field {index}"))),
69        Some(None) => T::from_sql_null(),
70        Some(Some(bytes)) => T::from_sql(oid, bytes),
71    }
72}
73
74impl<T: FromSql> FromCopyRow for T {
75    fn from_copy_row(oids: &[Oid], fields: &[Option<Bytes>]) -> Result<Self, Error> {
76        decode_field(oids, fields, 0)
77    }
78}
79
80macro_rules! from_copy_row_tuple {
81    ($($idx:tt $name:ident),+) => {
82        impl<$($name: FromSql),+> FromCopyRow for ($($name,)+) {
83            fn from_copy_row(oids: &[Oid], fields: &[Option<Bytes>]) -> Result<Self, Error> {
84                Ok(($( decode_field::<$name>(oids, fields, $idx)?, )+))
85            }
86        }
87    };
88}
89
90from_copy_row_tuple!(0 A);
91from_copy_row_tuple!(0 A, 1 B);
92from_copy_row_tuple!(0 A, 1 B, 2 C);
93from_copy_row_tuple!(0 A, 1 B, 2 C, 3 D);
94from_copy_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E);
95from_copy_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F);
96from_copy_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F, 6 G);
97from_copy_row_tuple!(0 A, 1 B, 2 C, 3 D, 4 E, 5 F, 6 G, 7 H);
98
99/// A writer for a binary `COPY ... FROM STDIN` decoded from `T`.
100pub struct CopyIn<T> {
101    inner: CopyInWriter,
102    _marker: PhantomData<fn() -> T>,
103}
104
105impl<T: ToRow> CopyIn<T> {
106    pub(crate) fn new(inner: CopyInWriter) -> Self {
107        CopyIn {
108            inner,
109            _marker: PhantomData,
110        }
111    }
112
113    /// Writes one row.
114    ///
115    /// Rows are encoded into a reusable buffer and coalesced into larger
116    /// `CopyData` messages, so a bulk load does not cost one allocation and one
117    /// network write per row.
118    ///
119    /// # Errors
120    ///
121    /// Returns an error if encoding fails or the connection is closed.
122    pub async fn write_row(&self, row: &T) -> Result<(), Error> {
123        self.inner.append(|buf| row.encode_row(buf)).await
124    }
125
126    /// Completes the COPY and returns the number of rows written.
127    ///
128    /// # Errors
129    ///
130    /// Returns an error if the server rejects the COPY.
131    pub async fn finish(self) -> Result<u64, Error> {
132        self.inner.finish().await
133    }
134}
135
136/// A reader for a binary `COPY ... TO STDOUT` decoded into `T`.
137pub struct CopyOut<T> {
138    inner: CopyOutStream,
139    oids: Vec<Oid>,
140    /// Unconsumed bytes. Chunks are appended (amortized O(1)); a completed row
141    /// is split off with `split_to`, so a row straddling many `CopyData`
142    /// messages is never re-copied.
143    buf: BytesMut,
144    /// Largest row this reader will buffer, in bytes, before erroring.
145    max_row_len: usize,
146    header_skipped: bool,
147    trailer_seen: bool,
148    done: bool,
149    _marker: PhantomData<fn() -> T>,
150}
151
152impl<T: FromCopyRow> CopyOut<T> {
153    pub(crate) fn new(inner: CopyOutStream, oids: Vec<Oid>, max_row_len: usize) -> Self {
154        CopyOut {
155            inner,
156            oids,
157            buf: BytesMut::new(),
158            max_row_len,
159            header_skipped: false,
160            trailer_seen: false,
161            done: false,
162            _marker: PhantomData,
163        }
164    }
165
166    /// Reads the next row, or `None` at the end of the COPY stream.
167    ///
168    /// # Errors
169    ///
170    /// Returns a decode error for a malformed or truncated stream (a stream
171    /// that ends without the binary `COPY` trailer is an error, not a clean
172    /// end), a row larger than the configured limit, a server error, or an I/O
173    /// error.
174    pub async fn next(&mut self) -> Result<Option<T>, Error> {
175        loop {
176            if let Some(row) = self.try_parse()? {
177                return Ok(Some(row));
178            }
179            if self.done {
180                if !self.buf.is_empty() {
181                    return Err(Error::Decode(format!(
182                        "binary COPY stream has {} trailing bytes after its trailer",
183                        self.buf.len()
184                    )));
185                }
186                return Ok(None);
187            }
188            match self.inner.next().await {
189                Some(Ok(chunk)) => {
190                    // A single row is bounded by `max_row_len` inside
191                    // `try_parse` (which checks the cumulative row offset), and
192                    // each `CopyData` message is already bounded by the
193                    // connection's message limit, so appending here cannot grow
194                    // the buffer without bound.
195                    self.buf.extend_from_slice(&chunk);
196                }
197                Some(Err(err)) => {
198                    self.done = true;
199                    return Err(err);
200                }
201                None => {
202                    self.done = true;
203                    if !self.trailer_seen {
204                        return Err(Error::Decode(
205                            "binary COPY stream ended before its trailer (truncated or malformed)".into(),
206                        ));
207                    }
208                    if !self.buf.is_empty() {
209                        return Err(Error::Decode(format!(
210                            "binary COPY stream has {} trailing bytes after its trailer",
211                            self.buf.len()
212                        )));
213                    }
214                }
215            }
216        }
217    }
218
219    /// Reads every remaining row.
220    ///
221    /// # Errors
222    ///
223    /// Returns a decode, server or I/O error.
224    pub async fn collect_all(mut self) -> Result<Vec<T>, Error> {
225        let mut out = Vec::new();
226        while let Some(row) = self.next().await? {
227            out.push(row);
228        }
229        Ok(out)
230    }
231
232    /// Parses one row if a complete one is buffered, without copying its
233    /// fields.
234    fn try_parse(&mut self) -> Result<Option<T>, Error> {
235        if !self.header_skipped {
236            if self.buf.len() < COPY_HEADER.len() {
237                return Ok(None);
238            }
239            if &self.buf[..COPY_HEADER.len()] != COPY_HEADER {
240                return Err(Error::Decode("not a binary COPY stream (bad file header)".into()));
241            }
242            self.buf.advance(COPY_HEADER.len());
243            self.header_skipped = true;
244        }
245        if self.buf.len() < 2 {
246            return Ok(None);
247        }
248        let raw_count = i16::from_be_bytes([self.buf[0], self.buf[1]]);
249        if raw_count == -1 {
250            // The end-of-stream trailer.
251            self.buf.advance(2);
252            self.trailer_seen = true;
253            self.done = true;
254            return Ok(None);
255        }
256        let count = usize::try_from(raw_count)
257            .map_err(|_| Error::Decode(format!("COPY row has an invalid field count {raw_count}")))?;
258
259        // First pass: find the end of the row, waiting for more data when the
260        // row is incomplete. Nothing is consumed until the whole row is here.
261        let mut pos = 2usize;
262        for _ in 0..count {
263            if self.buf.len() < pos + 4 {
264                return Ok(None);
265            }
266            let raw_len = i32::from_be_bytes([self.buf[pos], self.buf[pos + 1], self.buf[pos + 2], self.buf[pos + 3]]);
267            pos += 4;
268            if raw_len == -1 {
269                continue;
270            }
271            let len = usize::try_from(raw_len)
272                .map_err(|_| Error::Decode(format!("COPY field has an invalid length {raw_len}")))?;
273            let end = pos
274                .checked_add(len)
275                .ok_or_else(|| Error::Decode("binary COPY field length overflows the address space".into()))?;
276            // Enforce the row limit whether or not the bytes have arrived yet:
277            // a fully buffered row must obey it too, not only one being waited
278            // for.
279            if end > self.max_row_len {
280                return Err(Error::Decode(format!(
281                    "binary COPY field of {len} bytes exceeds the {} byte row limit",
282                    self.max_row_len
283                )));
284            }
285            if self.buf.len() < end {
286                return Ok(None);
287            }
288            pos = end;
289        }
290
291        // Second pass: hand out `Bytes` slices of the row, sharing one
292        // allocation instead of copying every field.
293        let row = self.buf.split_to(pos).freeze();
294        let mut fields: Vec<Option<Bytes>> = Vec::with_capacity(count);
295        let mut p = 2usize;
296        for _ in 0..count {
297            let raw_len = i32::from_be_bytes([row[p], row[p + 1], row[p + 2], row[p + 3]]);
298            p += 4;
299            if raw_len == -1 {
300                fields.push(None);
301            } else {
302                let len = usize::try_from(raw_len)
303                    .map_err(|_| Error::Decode(format!("COPY field has an invalid length {raw_len}")))?;
304                fields.push(Some(row.slice(p..p + len)));
305                p += len;
306            }
307        }
308        Ok(Some(T::from_copy_row(&self.oids, &fields)?))
309    }
310}
311
312impl<T> std::fmt::Debug for CopyOut<T> {
313    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        f.debug_struct("CopyOut").finish_non_exhaustive()
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use tokio::sync::mpsc;
321
322    use super::*;
323    use crate::connection::COPY_HEADER;
324
325    type Row = (i32, Option<String>);
326
327    fn stream() -> (mpsc::Sender<Result<Bytes, Error>>, CopyOut<Row>) {
328        // Large enough that the tests can push the whole stream before reading.
329        let (tx, rx) = mpsc::channel(256);
330        (
331            tx,
332            CopyOut::new(
333                CopyOutStream {
334                    rx,
335                },
336                vec![23, 25],
337                128 * 1024 * 1024,
338            ),
339        )
340    }
341
342    fn row(fields: &[Option<&[u8]>]) -> Vec<u8> {
343        let mut out = Vec::new();
344        out.extend_from_slice(&(fields.len() as i16).to_be_bytes());
345        for field in fields {
346            match field {
347                None => out.extend_from_slice(&(-1i32).to_be_bytes()),
348                Some(bytes) => {
349                    out.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
350                    out.extend_from_slice(bytes);
351                }
352            }
353        }
354        out
355    }
356
357    fn header() -> Vec<u8> {
358        COPY_HEADER.to_vec()
359    }
360
361    #[tokio::test]
362    async fn reads_rows_and_nulls_across_chunks() {
363        let (tx, mut out) = stream();
364        let mut data = header();
365        data.extend_from_slice(&row(&[Some(&1i32.to_be_bytes()), Some(b"a")]));
366        data.extend_from_slice(&row(&[Some(&2i32.to_be_bytes()), None]));
367        data.extend_from_slice(&(-1i16).to_be_bytes()); // trailer
368        // Feed it one byte at a time: the parser must wait, not truncate.
369        for byte in data {
370            tx.send(Ok(Bytes::from(vec![byte]))).await.unwrap();
371        }
372        drop(tx);
373        assert_eq!(out.next().await.unwrap().unwrap(), (1, Some("a".to_string())));
374        assert_eq!(out.next().await.unwrap().unwrap(), (2, None));
375        assert!(out.next().await.unwrap().is_none());
376    }
377
378    #[tokio::test]
379    async fn truncated_stream_is_an_error_not_a_clean_end() {
380        let (tx, mut out) = stream();
381        let mut data = header();
382        let mut partial = row(&[Some(&1i32.to_be_bytes()), Some(b"abcdefgh")]);
383        partial.truncate(partial.len() - 4);
384        data.extend_from_slice(&partial);
385        tx.send(Ok(Bytes::from(data))).await.unwrap();
386        drop(tx); // stream ends without the -1 trailer
387        let err = out.next().await.unwrap_err();
388        assert!(matches!(err, Error::Decode(_)), "{err}");
389    }
390
391    #[tokio::test]
392    async fn trailing_garbage_after_the_trailer_is_an_error() {
393        let (tx, mut out) = stream();
394        let mut data = header();
395        data.extend_from_slice(&(-1i16).to_be_bytes());
396        data.extend_from_slice(&[0xde, 0xad]);
397        tx.send(Ok(Bytes::from(data))).await.unwrap();
398        drop(tx);
399        let err = out.next().await.unwrap_err();
400        assert!(matches!(err, Error::Decode(_)), "{err}");
401    }
402
403    #[tokio::test]
404    async fn bad_file_header_is_an_error() {
405        let (tx, mut out) = stream();
406        tx.send(Ok(Bytes::from(b"NOTACOPYFILE\n\xff\r\n\0\0\0\0\0\0\0\0\0".to_vec())))
407            .await
408            .unwrap();
409        drop(tx);
410        let err = out.next().await.unwrap_err();
411        assert!(matches!(err, Error::Decode(_)), "{err}");
412    }
413
414    #[tokio::test]
415    async fn invalid_field_counts_and_lengths_are_errors() {
416        for payload in [
417            // field count -2 (only -1 is the trailer)
418            {
419                let mut d = header();
420                d.extend_from_slice(&(-2i16).to_be_bytes());
421                d
422            },
423            // field length -2 (only -1 is NULL)
424            {
425                let mut d = header();
426                d.extend_from_slice(&row(&[Some(&1i32.to_be_bytes())]));
427                d.truncate(d.len() - 4);
428                d.extend_from_slice(&(-2i32).to_be_bytes());
429                d
430            },
431        ] {
432            let (tx, mut out) = stream();
433            tx.send(Ok(Bytes::from(payload))).await.unwrap();
434            drop(tx);
435            assert!(out.next().await.is_err());
436        }
437    }
438
439    #[tokio::test]
440    async fn a_fully_buffered_row_over_the_limit_is_rejected() {
441        // The row limit must be enforced even when the whole row is already
442        // buffered, not only while waiting for bytes.
443        let (tx, rx) = mpsc::channel(256);
444        let mut out = CopyOut::<Row>::new(
445            CopyOutStream {
446                rx,
447            },
448            vec![23, 25],
449            8,
450        );
451        let mut data = header();
452        data.extend_from_slice(&row(&[Some(&1i32.to_be_bytes()), Some(b"way-too-long")]));
453        tx.send(Ok(Bytes::from(data))).await.unwrap();
454        drop(tx);
455        let err = out.next().await.unwrap_err();
456        assert!(matches!(err, Error::Decode(_)), "{err}");
457    }
458
459    #[tokio::test]
460    async fn fields_borrow_from_the_same_buffer() {
461        let (tx, mut out) = stream();
462        let mut data = header();
463        data.extend_from_slice(&row(&[Some(&7i32.to_be_bytes()), Some(b"xyz")]));
464        data.extend_from_slice(&(-1i16).to_be_bytes());
465        tx.send(Ok(Bytes::from(data))).await.unwrap();
466        drop(tx);
467        let row = out.next().await.unwrap().unwrap();
468        assert_eq!(row.0, 7);
469        assert_eq!(row.1.as_deref(), Some("xyz"));
470        assert!(out.next().await.unwrap().is_none());
471    }
472}
postgresql/copy.rs

postgresql/
copy.rs
