tokio_postgres/
copy_in.rs

Help
1use crate::client::{InnerClient, Responses};
2use crate::codec::FrontendMessage;
3use crate::connection::RequestMessages;
4use crate::query::extract_row_affected;
5use crate::{query, simple_query, slice_iter, Error, Statement};
6use bytes::{Buf, BufMut, Bytes, BytesMut};
7use futures_channel::mpsc;
8use futures_util::{future, ready, Sink, SinkExt, Stream, StreamExt};
9use log::debug;
10use pin_project_lite::pin_project;
11use postgres_protocol::message::backend::Message;
12use postgres_protocol::message::frontend;
13use postgres_protocol::message::frontend::CopyData;
14use std::marker::{PhantomData, PhantomPinned};
15use std::pin::Pin;
16use std::task::{Context, Poll};
17
18enum CopyInMessage {
19    Message(FrontendMessage),
20    Done,
21}
22
23pub struct CopyInReceiver {
24    receiver: mpsc::Receiver<CopyInMessage>,
25    done: bool,
26}
27
28impl CopyInReceiver {
29    fn new(receiver: mpsc::Receiver<CopyInMessage>) -> CopyInReceiver {
30        CopyInReceiver {
31            receiver,
32            done: false,
33        }
34    }
35}
36
37impl Stream for CopyInReceiver {
38    type Item = FrontendMessage;
39
40    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<FrontendMessage>> {
41        if self.done {
42            return Poll::Ready(None);
43        }
44
45        match ready!(self.receiver.poll_next_unpin(cx)) {
46            Some(CopyInMessage::Message(message)) => Poll::Ready(Some(message)),
47            Some(CopyInMessage::Done) => {
48                self.done = true;
49                let mut buf = BytesMut::new();
50                frontend::copy_done(&mut buf);
51                frontend::sync(&mut buf);
52                Poll::Ready(Some(FrontendMessage::Raw(buf.freeze())))
53            }
54            None => {
55                self.done = true;
56                let mut buf = BytesMut::new();
57                frontend::copy_fail("", &mut buf).unwrap();
58                frontend::sync(&mut buf);
59                Poll::Ready(Some(FrontendMessage::Raw(buf.freeze())))
60            }
61        }
62    }
63}
64
65enum SinkState {
66    Active,
67    Closing,
68    Reading,
69}
70
71pin_project! {
72    /// A sink for `COPY ... FROM STDIN` query data.
73    ///
74    /// The copy *must* be explicitly completed via the `Sink::close` or `finish` methods. If it is
75    /// not, the copy will be aborted.
76    pub struct CopyInSink<T> {
77        #[pin]
78        sender: mpsc::Sender<CopyInMessage>,
79        responses: Responses,
80        buf: BytesMut,
81        state: SinkState,
82        #[pin]
83        _p: PhantomPinned,
84        _p2: PhantomData<T>,
85    }
86}
87
88impl<T> CopyInSink<T>
89where
90    T: Buf + 'static + Send,
91{
92    /// A poll-based version of `finish`.
93    pub fn poll_finish(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<u64, Error>> {
94        loop {
95            match self.state {
96                SinkState::Active => {
97                    ready!(self.as_mut().poll_flush(cx))?;
98                    let mut this = self.as_mut().project();
99                    ready!(this.sender.as_mut().poll_ready(cx)).map_err(|_| Error::closed())?;
100                    this.sender
101                        .start_send(CopyInMessage::Done)
102                        .map_err(|_| Error::closed())?;
103                    *this.state = SinkState::Closing;
104                }
105                SinkState::Closing => {
106                    let this = self.as_mut().project();
107                    ready!(this.sender.poll_close(cx)).map_err(|_| Error::closed())?;
108                    *this.state = SinkState::Reading;
109                }
110                SinkState::Reading => {
111                    let this = self.as_mut().project();
112                    match ready!(this.responses.poll_next(cx))? {
113                        Message::CommandComplete(body) => {
114                            let rows = extract_row_affected(&body)?;
115                            return Poll::Ready(Ok(rows));
116                        }
117                        _ => return Poll::Ready(Err(Error::unexpected_message())),
118                    }
119                }
120            }
121        }
122    }
123
124    /// Completes the copy, returning the number of rows inserted.
125    ///
126    /// The `Sink::close` method is equivalent to `finish`, except that it does not return the
127    /// number of rows.
128    pub async fn finish(mut self: Pin<&mut Self>) -> Result<u64, Error> {
129        future::poll_fn(|cx| self.as_mut().poll_finish(cx)).await
130    }
131}
132
133impl<T> Sink<T> for CopyInSink<T>
134where
135    T: Buf + 'static + Send,
136{
137    type Error = Error;
138
139    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
140        self.project()
141            .sender
142            .poll_ready(cx)
143            .map_err(|_| Error::closed())
144    }
145
146    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Error> {
147        let this = self.project();
148
149        let data: Box<dyn Buf + Send> = if item.remaining() > 4096 {
150            if this.buf.is_empty() {
151                Box::new(item)
152            } else {
153                Box::new(this.buf.split().freeze().chain(item))
154            }
155        } else {
156            this.buf.put(item);
157            if this.buf.len() > 4096 {
158                Box::new(this.buf.split().freeze())
159            } else {
160                return Ok(());
161            }
162        };
163
164        let data = CopyData::new(data).map_err(Error::encode)?;
165        this.sender
166            .start_send(CopyInMessage::Message(FrontendMessage::CopyData(data)))
167            .map_err(|_| Error::closed())
168    }
169
170    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
171        let mut this = self.project();
172
173        if !this.buf.is_empty() {
174            ready!(this.sender.as_mut().poll_ready(cx)).map_err(|_| Error::closed())?;
175            let data: Box<dyn Buf + Send> = Box::new(this.buf.split().freeze());
176            let data = CopyData::new(data).map_err(Error::encode)?;
177            this.sender
178                .as_mut()
179                .start_send(CopyInMessage::Message(FrontendMessage::CopyData(data)))
180                .map_err(|_| Error::closed())?;
181        }
182
183        this.sender.poll_flush(cx).map_err(|_| Error::closed())
184    }
185
186    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
187        self.poll_finish(cx).map_ok(|_| ())
188    }
189}
190
191async fn start<T>(client: &InnerClient, buf: Bytes, simple: bool) -> Result<CopyInSink<T>, Error>
192where
193    T: Buf + 'static + Send,
194{
195    let (mut sender, receiver) = mpsc::channel(1);
196    let receiver = CopyInReceiver::new(receiver);
197    let mut responses = client.send(RequestMessages::CopyIn(receiver))?;
198
199    sender
200        .send(CopyInMessage::Message(FrontendMessage::Raw(buf)))
201        .await
202        .map_err(|_| Error::closed())?;
203
204    if !simple {
205        match responses.next().await? {
206            Message::BindComplete => {}
207            _ => return Err(Error::unexpected_message()),
208        }
209    }
210
211    match responses.next().await? {
212        Message::CopyInResponse(_) => {}
213        _ => return Err(Error::unexpected_message()),
214    }
215
216    Ok(CopyInSink {
217        sender,
218        responses,
219        buf: BytesMut::new(),
220        state: SinkState::Active,
221        _p: PhantomPinned,
222        _p2: PhantomData,
223    })
224}
225
226pub async fn copy_in<T>(client: &InnerClient, statement: Statement) -> Result<CopyInSink<T>, Error>
227where
228    T: Buf + 'static + Send,
229{
230    debug!("executing copy in statement {}", statement.name());
231
232    let buf = query::encode(client, &statement, slice_iter(&[]))?;
233    start(client, buf, false).await
234}
235
236pub async fn copy_in_simple<T>(client: &InnerClient, query: &str) -> Result<CopyInSink<T>, Error>
237where
238    T: Buf + 'static + Send,
239{
240    debug!("executing copy in query {}", query);
241
242    let buf = simple_query::encode(client, query)?;
243    start(client, buf, true).await
244}
tokio_postgres/copy_in.rs

tokio_postgres/
copy_in.rs
