tokio_postgres/
copy_out.rs

1use crate::client::{InnerClient, Responses};
2use crate::codec::FrontendMessage;
3use crate::connection::RequestMessages;
4use crate::{query, slice_iter, Error, Statement};
5use bytes::Bytes;
6use futures_util::{ready, Stream};
7use log::debug;
8use pin_project_lite::pin_project;
9use postgres_protocol::message::backend::Message;
10use std::marker::PhantomPinned;
11use std::pin::Pin;
12use std::task::{Context, Poll};
13
14pub async fn copy_out(client: &InnerClient, statement: Statement) -> Result<CopyOutStream, Error> {
15    debug!("executing copy out statement {}", statement.name());
16
17    let buf = query::encode(client, &statement, slice_iter(&[]))?;
18    let responses = start(client, buf).await?;
19    Ok(CopyOutStream {
20        responses,
21        _p: PhantomPinned,
22    })
23}
24
25async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
26    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
27
28    match responses.next().await? {
29        Message::BindComplete => {}
30        _ => return Err(Error::unexpected_message()),
31    }
32
33    match responses.next().await? {
34        Message::CopyOutResponse(_) => {}
35        _ => return Err(Error::unexpected_message()),
36    }
37
38    Ok(responses)
39}
40
41pin_project! {
42    /// A stream of `COPY ... TO STDOUT` query data.
43    pub struct CopyOutStream {
44        responses: Responses,
45        #[pin]
46        _p: PhantomPinned,
47    }
48}
49
50impl Stream for CopyOutStream {
51    type Item = Result<Bytes, Error>;
52
53    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
54        let this = self.project();
55
56        match ready!(this.responses.poll_next(cx)?) {
57            Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
58            Message::CopyDone => Poll::Ready(None),
59            _ => Poll::Ready(Some(Err(Error::unexpected_message()))),
60        }
61    }
62}
tokio_postgres/copy_out.rs

tokio_postgres/
copy_out.rs
