postgres/
copy_in_writer.rs

Help
1use crate::connection::ConnectionRef;
2use crate::lazy_pin::LazyPin;
3use bytes::{Bytes, BytesMut};
4use futures_util::SinkExt;
5use std::io;
6use std::io::Write;
7use tokio_postgres::{CopyInSink, Error};
8
9/// The writer returned by the `copy_in` method.
10///
11/// The copy *must* be explicitly completed via the `finish` method. If it is not, the copy will be aborted.
12pub struct CopyInWriter<'a> {
13    pub(crate) connection: ConnectionRef<'a>,
14    pub(crate) sink: LazyPin<CopyInSink<Bytes>>,
15    buf: BytesMut,
16}
17
18impl<'a> CopyInWriter<'a> {
19    pub(crate) fn new(connection: ConnectionRef<'a>, sink: CopyInSink<Bytes>) -> CopyInWriter<'a> {
20        CopyInWriter {
21            connection,
22            sink: LazyPin::new(sink),
23            buf: BytesMut::new(),
24        }
25    }
26
27    /// Completes the copy, returning the number of rows written.
28    ///
29    /// If this is not called, the copy will be aborted.
30    pub fn finish(mut self) -> Result<u64, Error> {
31        self.flush_inner()?;
32        self.connection.block_on(self.sink.pinned().finish())
33    }
34
35    fn flush_inner(&mut self) -> Result<(), Error> {
36        if self.buf.is_empty() {
37            return Ok(());
38        }
39
40        self.connection
41            .block_on(self.sink.pinned().send(self.buf.split().freeze()))
42    }
43}
44
45impl Write for CopyInWriter<'_> {
46    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
47        if self.buf.len() > 4096 {
48            self.flush()?;
49        }
50
51        self.buf.extend_from_slice(buf);
52        Ok(buf.len())
53    }
54
55    fn flush(&mut self) -> io::Result<()> {
56        self.flush_inner()
57            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
58    }
59}
postgres/copy_in_writer.rs

postgres/
copy_in_writer.rs
