2019-07-25 02:18:15 +00:00
|
|
|
use crate::client::{InnerClient, Responses};
|
|
|
|
use crate::codec::FrontendMessage;
|
|
|
|
use crate::connection::RequestMessages;
|
|
|
|
use crate::types::{IsNull, ToSql};
|
2019-08-02 03:43:13 +00:00
|
|
|
use crate::{Error, Portal, Row, Statement};
|
2019-10-13 00:47:47 +00:00
|
|
|
use bytes::{Bytes, BytesMut};
|
2019-10-09 00:22:56 +00:00
|
|
|
use futures::{ready, Stream};
|
2019-11-19 02:12:34 +00:00
|
|
|
use pin_project_lite::pin_project;
|
2019-07-25 02:18:15 +00:00
|
|
|
use postgres_protocol::message::backend::Message;
|
|
|
|
use postgres_protocol::message::frontend;
|
2019-10-16 01:17:10 +00:00
|
|
|
use std::marker::PhantomPinned;
|
2019-07-25 02:18:15 +00:00
|
|
|
use std::pin::Pin;
|
|
|
|
use std::task::{Context, Poll};
|
|
|
|
|
2019-10-09 00:15:41 +00:00
|
|
|
pub async fn query<'a, I>(
|
|
|
|
client: &InnerClient,
|
2019-09-28 18:42:00 +00:00
|
|
|
statement: Statement,
|
|
|
|
params: I,
|
2019-10-09 00:15:41 +00:00
|
|
|
) -> Result<RowStream, Error>
|
2019-09-28 18:42:00 +00:00
|
|
|
where
|
2019-10-09 00:15:41 +00:00
|
|
|
I: IntoIterator<Item = &'a dyn ToSql>,
|
2019-09-28 18:42:00 +00:00
|
|
|
I::IntoIter: ExactSizeIterator,
|
|
|
|
{
|
2019-10-12 23:30:27 +00:00
|
|
|
let buf = encode(client, &statement, params)?;
|
2019-10-09 00:15:41 +00:00
|
|
|
let responses = start(client, buf).await?;
|
|
|
|
Ok(RowStream {
|
|
|
|
statement,
|
|
|
|
responses,
|
2019-10-14 23:42:54 +00:00
|
|
|
_p: PhantomPinned,
|
2019-10-09 00:15:41 +00:00
|
|
|
})
|
2019-07-25 02:18:15 +00:00
|
|
|
}
|
|
|
|
|
2019-10-09 00:22:56 +00:00
|
|
|
pub async fn query_portal(
|
|
|
|
client: &InnerClient,
|
|
|
|
portal: &Portal,
|
2019-08-02 03:43:13 +00:00
|
|
|
max_rows: i32,
|
2019-10-09 00:22:56 +00:00
|
|
|
) -> Result<RowStream, Error> {
|
2019-10-12 23:30:27 +00:00
|
|
|
let buf = client.with_buf(|buf| {
|
|
|
|
frontend::execute(portal.name(), max_rows, buf).map_err(Error::encode)?;
|
|
|
|
frontend::sync(buf);
|
|
|
|
Ok(buf.take().freeze())
|
|
|
|
})?;
|
2019-08-02 03:43:13 +00:00
|
|
|
|
2019-10-09 00:22:56 +00:00
|
|
|
let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
|
2019-08-02 03:43:13 +00:00
|
|
|
|
2019-10-09 00:22:56 +00:00
|
|
|
Ok(RowStream {
|
|
|
|
statement: portal.statement().clone(),
|
|
|
|
responses,
|
2019-10-14 23:42:54 +00:00
|
|
|
_p: PhantomPinned,
|
2019-10-09 00:22:56 +00:00
|
|
|
})
|
2019-08-02 03:43:13 +00:00
|
|
|
}
|
|
|
|
|
2019-09-28 18:42:00 +00:00
|
|
|
pub async fn execute<'a, I>(
|
|
|
|
client: &InnerClient,
|
|
|
|
statement: Statement,
|
|
|
|
params: I,
|
|
|
|
) -> Result<u64, Error>
|
|
|
|
where
|
|
|
|
I: IntoIterator<Item = &'a dyn ToSql>,
|
|
|
|
I::IntoIter: ExactSizeIterator,
|
|
|
|
{
|
2019-10-12 23:30:27 +00:00
|
|
|
let buf = encode(client, &statement, params)?;
|
2019-07-27 03:11:34 +00:00
|
|
|
let mut responses = start(client, buf).await?;
|
2019-07-25 02:18:15 +00:00
|
|
|
|
|
|
|
loop {
|
|
|
|
match responses.next().await? {
|
|
|
|
Message::DataRow(_) => {}
|
|
|
|
Message::CommandComplete(body) => {
|
|
|
|
let rows = body
|
|
|
|
.tag()
|
|
|
|
.map_err(Error::parse)?
|
|
|
|
.rsplit(' ')
|
|
|
|
.next()
|
|
|
|
.unwrap()
|
|
|
|
.parse()
|
|
|
|
.unwrap_or(0);
|
|
|
|
return Ok(rows);
|
|
|
|
}
|
|
|
|
Message::EmptyQueryResponse => return Ok(0),
|
|
|
|
_ => return Err(Error::unexpected_message()),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-10-12 23:30:27 +00:00
|
|
|
async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
|
2019-07-27 03:11:34 +00:00
|
|
|
let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
|
|
|
|
|
|
|
|
match responses.next().await? {
|
|
|
|
Message::BindComplete => {}
|
|
|
|
_ => return Err(Error::unexpected_message()),
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(responses)
|
|
|
|
}
|
|
|
|
|
2019-10-12 23:30:27 +00:00
|
|
|
pub fn encode<'a, I>(client: &InnerClient, statement: &Statement, params: I) -> Result<Bytes, Error>
|
2019-08-02 03:43:13 +00:00
|
|
|
where
|
|
|
|
I: IntoIterator<Item = &'a dyn ToSql>,
|
|
|
|
I::IntoIter: ExactSizeIterator,
|
|
|
|
{
|
2019-10-12 23:30:27 +00:00
|
|
|
client.with_buf(|buf| {
|
|
|
|
encode_bind(statement, params, "", buf)?;
|
|
|
|
frontend::execute("", 0, buf).map_err(Error::encode)?;
|
|
|
|
frontend::sync(buf);
|
|
|
|
Ok(buf.take().freeze())
|
|
|
|
})
|
2019-08-02 03:43:13 +00:00
|
|
|
}
|
|
|
|
|
2019-10-13 00:47:47 +00:00
|
|
|
pub fn encode_bind<'a, I>(
|
|
|
|
statement: &Statement,
|
|
|
|
params: I,
|
|
|
|
portal: &str,
|
|
|
|
buf: &mut BytesMut,
|
|
|
|
) -> Result<(), Error>
|
2019-07-25 02:18:15 +00:00
|
|
|
where
|
|
|
|
I: IntoIterator<Item = &'a dyn ToSql>,
|
|
|
|
I::IntoIter: ExactSizeIterator,
|
|
|
|
{
|
|
|
|
let params = params.into_iter();
|
|
|
|
|
|
|
|
assert!(
|
|
|
|
statement.params().len() == params.len(),
|
|
|
|
"expected {} parameters but got {}",
|
|
|
|
statement.params().len(),
|
|
|
|
params.len()
|
|
|
|
);
|
|
|
|
|
|
|
|
let mut error_idx = 0;
|
|
|
|
let r = frontend::bind(
|
2019-08-02 03:43:13 +00:00
|
|
|
portal,
|
2019-07-25 02:18:15 +00:00
|
|
|
statement.name(),
|
|
|
|
Some(1),
|
|
|
|
params.zip(statement.params()).enumerate(),
|
|
|
|
|(idx, (param, ty)), buf| match param.to_sql_checked(ty, buf) {
|
|
|
|
Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
|
|
|
|
Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
|
|
|
|
Err(e) => {
|
|
|
|
error_idx = idx;
|
|
|
|
Err(e)
|
|
|
|
}
|
|
|
|
},
|
|
|
|
Some(1),
|
2019-10-12 23:30:27 +00:00
|
|
|
buf,
|
2019-07-25 02:18:15 +00:00
|
|
|
);
|
|
|
|
match r {
|
2019-10-12 23:30:27 +00:00
|
|
|
Ok(()) => Ok(()),
|
2019-10-09 22:20:23 +00:00
|
|
|
Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, error_idx)),
|
|
|
|
Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
|
2019-07-25 02:18:15 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-11-19 02:12:34 +00:00
|
|
|
pin_project! {
|
|
|
|
/// A stream of table rows.
|
|
|
|
pub struct RowStream {
|
|
|
|
statement: Statement,
|
|
|
|
responses: Responses,
|
|
|
|
#[pin]
|
|
|
|
_p: PhantomPinned,
|
|
|
|
}
|
2019-07-25 02:18:15 +00:00
|
|
|
}
|
|
|
|
|
2019-10-09 00:15:41 +00:00
|
|
|
impl Stream for RowStream {
|
2019-07-25 02:18:15 +00:00
|
|
|
type Item = Result<Row, Error>;
|
|
|
|
|
2019-10-14 23:42:54 +00:00
|
|
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
|
|
let this = self.project();
|
|
|
|
match ready!(this.responses.poll_next(cx)?) {
|
2019-07-25 02:18:15 +00:00
|
|
|
Message::DataRow(body) => {
|
2019-10-14 23:42:54 +00:00
|
|
|
Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
|
2019-07-25 02:18:15 +00:00
|
|
|
}
|
2019-08-02 03:43:13 +00:00
|
|
|
Message::EmptyQueryResponse
|
|
|
|
| Message::CommandComplete(_)
|
|
|
|
| Message::PortalSuspended => Poll::Ready(None),
|
2019-07-25 02:18:15 +00:00
|
|
|
Message::ErrorResponse(body) => Poll::Ready(Some(Err(Error::db(body)))),
|
|
|
|
_ => Poll::Ready(Some(Err(Error::unexpected_message()))),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|