rust-postgres/tokio-postgres/src/proto/simple_query.rs
2018-12-08 17:52:20 -08:00

57 lines
1.8 KiB
Rust

use futures::sync::mpsc;
use futures::{Poll, Stream};
use postgres_protocol::message::backend::Message;
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
use crate::proto::client::{Client, PendingRequest};
use crate::Error;
#[derive(StateMachineFuture)]
pub enum SimpleQuery {
#[state_machine_future(start, transitions(ReadResponse))]
Start {
client: Client,
request: PendingRequest,
},
#[state_machine_future(transitions(Finished))]
ReadResponse { receiver: mpsc::Receiver<Message> },
#[state_machine_future(ready)]
Finished(()),
#[state_machine_future(error)]
Failed(Error),
}
impl PollSimpleQuery for SimpleQuery {
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
let state = state.take();
let receiver = state.client.send(state.request)?;
transition!(ReadResponse { receiver })
}
fn poll_read_response<'a>(
state: &'a mut RentToOwn<'a, ReadResponse>,
) -> Poll<AfterReadResponse, Error> {
loop {
let message = try_ready_receive!(state.receiver.poll());
match message {
Some(Message::CommandComplete(_))
| Some(Message::RowDescription(_))
| Some(Message::DataRow(_))
| Some(Message::EmptyQueryResponse) => {}
Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
Some(Message::ReadyForQuery(_)) => transition!(Finished(())),
Some(_) => return Err(Error::unexpected_message()),
None => return Err(Error::closed()),
}
}
}
}
impl SimpleQueryFuture {
pub fn new(client: Client, request: PendingRequest) -> SimpleQueryFuture {
SimpleQuery::start(client, request)
}
}