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(
18        client,
19        &statement,
20        slice_iter(&[]),
21        query::ResultFormat::Binary,
22    )?;
23    let responses = start(client, buf).await?;
24    Ok(CopyOutStream {
25        responses,
26        _p: PhantomPinned,
27    })
28}
29
30async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
31    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
32
33    match responses.next().await? {
34        Message::BindComplete => {}
35        _ => return Err(Error::unexpected_message()),
36    }
37
38    match responses.next().await? {
39        Message::CopyOutResponse(_) => {}
40        _ => return Err(Error::unexpected_message()),
41    }
42
43    Ok(responses)
44}
45
46pin_project! {
47    /// A stream of `COPY ... TO STDOUT` query data.
48    pub struct CopyOutStream {
49        responses: Responses,
50        #[pin]
51        _p: PhantomPinned,
52    }
53}
54
55impl Stream for CopyOutStream {
56    type Item = Result<Bytes, Error>;
57
58    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
59        let this = self.project();
60
61        match ready!(this.responses.poll_next(cx)?) {
62            Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
63            Message::CopyDone => Poll::Ready(None),
64            _ => Poll::Ready(Some(Err(Error::unexpected_message()))),
65        }
66    }
67}
tokio_postgres/copy_out.rs

tokio_postgres/
copy_out.rs
