diesel/pg/connection/
copy.rs

1use core::ffi;
2use std::io::BufRead;
3use std::io::Read;
4use std::io::Write;
5
6use super::raw::RawConnection;
7use super::result::PgResult;
8use crate::QueryResult;
9
10#[allow(missing_debug_implementations)] // `PgConnection` is not debug
11pub struct CopyFromSink<'conn> {
12    conn: &'conn mut RawConnection,
13    finished: bool,
14}
15
16impl<'conn> CopyFromSink<'conn> {
17    pub(super) fn new(conn: &'conn mut RawConnection) -> Self {
18        Self {
19            conn,
20            finished: false,
21        }
22    }
23
24    pub(super) fn finish(mut self, err: Option<String>) -> QueryResult<()> {
25        self.conn
26            .finish_copy_from(err)
27            .inspect(|_| self.finished = true)
28    }
29}
30
31impl Write for CopyFromSink<'_> {
32    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
33        self.conn
34            .put_copy_data(buf)
35            .map_err(std::io::Error::other)?;
36        Ok(buf.len())
37    }
38
39    fn flush(&mut self) -> std::io::Result<()> {
40        Ok(())
41    }
42}
43
44impl Drop for CopyFromSink<'_> {
45    fn drop(&mut self) {
46        if !self.finished {
47            let _ = self.conn.consume_results();
48        }
49    }
50}
51
52#[allow(missing_debug_implementations)] // `PgConnection` is not debug
53pub struct CopyToBuffer<'conn> {
54    conn: &'conn mut RawConnection,
55    ptr: *mut ffi::c_char,
56    offset: usize,
57    len: usize,
58    result: PgResult,
59    done: bool,
60}
61
62impl<'conn> CopyToBuffer<'conn> {
63    pub(super) fn new(conn: &'conn mut RawConnection, result: PgResult) -> Self {
64        Self {
65            conn,
66            ptr: core::ptr::null_mut(),
67            offset: 0,
68            len: 0,
69            result,
70            done: false,
71        }
72    }
73
74    #[allow(unsafe_code)] // construct a slice from a raw ptr
75    pub(crate) fn data_slice(&self) -> &[u8] {
76        if !self.ptr.is_null() && self.offset < self.len {
77            let slice = unsafe { core::slice::from_raw_parts(self.ptr as *const u8, self.len - 1) };
78            &slice[self.offset..]
79        } else {
80            &[]
81        }
82    }
83
84    pub(crate) fn get_result(&self) -> &PgResult {
85        &self.result
86    }
87}
88
89impl Drop for CopyToBuffer<'_> {
90    #[allow(unsafe_code)] // ffi code
91    fn drop(&mut self) {
92        let _ = self.conn.consume_results();
93        if !self.ptr.is_null() {
94            unsafe { pq_sys::PQfreemem(self.ptr as *mut ffi::c_void) };
95            self.ptr = core::ptr::null_mut();
96        }
97    }
98}
99
100impl Read for CopyToBuffer<'_> {
101    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
102        let data = self.fill_buf()?;
103        let len = usize::min(buf.len(), data.len());
104        buf[..len].copy_from_slice(&data[..len]);
105        self.consume(len);
106        Ok(len)
107    }
108}
109
110impl BufRead for CopyToBuffer<'_> {
111    #[allow(unsafe_code)] // ffi code + ptr arithmetic
112    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
113        if self.data_slice().is_empty() && !self.done {
114            let len = unsafe {
115                if !self.ptr.is_null() {
116                    pq_sys::PQfreemem(self.ptr as *mut ffi::c_void);
117                    self.ptr = core::ptr::null_mut();
118                }
119                pq_sys::PQgetCopyData(self.conn.internal_connection.as_ptr(), &mut self.ptr, 0)
120            };
121            match len {
122                len if len >= 0 => {
123                    self.len = 1 + usize::try_from(len).map_err(std::io::Error::other)?
124                }
125                -1 => {
126                    self.done = true;
127                    self.len = 0;
128                    self.conn.consume_results().map_err(std::io::Error::other)?;
129                }
130                _ => {
131                    let error = self.conn.last_error_message();
132                    self.done = true;
133                    self.len = 0;
134                    return Err(std::io::Error::other(error));
135                }
136            }
137            self.offset = 0;
138        }
139        Ok(self.data_slice())
140    }
141
142    fn consume(&mut self, amt: usize) {
143        self.offset = usize::min(self.len, self.offset + amt);
144    }
145}
diesel/pg/connection/copy.rs

diesel/pg/connection/
copy.rs
