rust-postgres/tokio-postgres/src/client.rs

189 lines
5.0 KiB
Rust
Raw Normal View History

2019-07-24 02:54:22 +00:00
use crate::codec::BackendMessages;
use crate::connection::{Request, RequestMessages};
use crate::prepare::prepare;
2019-07-25 02:18:15 +00:00
use crate::query::{execute, query, Query};
use crate::types::{Oid, ToSql, Type};
2019-07-24 02:54:22 +00:00
use crate::{Error, Statement};
use fallible_iterator::FallibleIterator;
use futures::channel::mpsc;
2019-07-25 02:18:15 +00:00
use futures::future;
use futures::{ready, StreamExt};
2019-07-24 03:16:31 +00:00
use parking_lot::Mutex;
2019-07-24 02:54:22 +00:00
use postgres_protocol::message::backend::Message;
2019-07-24 03:16:31 +00:00
use std::collections::HashMap;
2019-07-24 02:54:22 +00:00
use std::future::Future;
use std::sync::Arc;
use std::task::{Context, Poll};
2019-07-24 02:54:22 +00:00
pub struct Responses {
receiver: mpsc::Receiver<BackendMessages>,
cur: BackendMessages,
}
impl Responses {
2019-07-25 02:18:15 +00:00
pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Result<Message, Error>> {
2019-07-24 02:54:22 +00:00
loop {
match self.cur.next().map_err(Error::parse)? {
2019-07-25 02:18:15 +00:00
Some(Message::ErrorResponse(body)) => return Poll::Ready(Err(Error::db(body))),
Some(message) => return Poll::Ready(Ok(message)),
2019-07-24 02:54:22 +00:00
None => {}
}
2019-07-25 02:18:15 +00:00
match ready!(self.receiver.poll_next_unpin(cx)) {
2019-07-24 02:54:22 +00:00
Some(messages) => self.cur = messages,
2019-07-25 02:18:15 +00:00
None => return Poll::Ready(Err(Error::closed())),
2019-07-24 02:54:22 +00:00
}
}
}
2019-07-25 02:18:15 +00:00
pub async fn next(&mut self) -> Result<Message, Error> {
future::poll_fn(|cx| self.poll_next(cx)).await
}
2019-07-24 02:54:22 +00:00
}
2019-07-24 03:16:31 +00:00
struct State {
has_typeinfo: bool,
has_typeinfo_composite: bool,
has_typeinfo_enum: bool,
types: HashMap<Oid, Type>,
}
2019-07-24 02:54:22 +00:00
pub struct InnerClient {
sender: mpsc::UnboundedSender<Request>,
2019-07-24 03:16:31 +00:00
state: Mutex<State>,
2019-07-24 02:54:22 +00:00
}
impl InnerClient {
pub fn send(&self, messages: RequestMessages) -> Result<Responses, Error> {
let (sender, receiver) = mpsc::channel(1);
let request = Request { messages, sender };
self.sender
.unbounded_send(request)
.map_err(|_| Error::closed())?;
Ok(Responses {
receiver,
cur: BackendMessages::empty(),
})
}
2019-07-24 03:16:31 +00:00
pub fn has_typeinfo(&self) -> bool {
self.state.lock().has_typeinfo
}
pub fn set_has_typeinfo(&self) {
self.state.lock().has_typeinfo = true;
}
pub fn has_typeinfo_composite(&self) -> bool {
self.state.lock().has_typeinfo_composite
}
pub fn set_has_typeinfo_composite(&self) {
self.state.lock().has_typeinfo_composite = true;
}
pub fn has_typeinfo_enum(&self) -> bool {
self.state.lock().has_typeinfo_enum
}
pub fn set_has_typeinfo_enum(&self) {
self.state.lock().has_typeinfo_enum = true;
}
pub fn type_(&self, oid: Oid) -> Option<Type> {
self.state.lock().types.get(&oid).cloned()
}
pub fn set_type(&self, oid: Oid, type_: Type) {
self.state.lock().types.insert(oid, type_);
}
2019-07-24 02:54:22 +00:00
}
pub struct Client {
inner: Arc<InnerClient>,
process_id: i32,
secret_key: i32,
}
impl Client {
pub(crate) fn new(
sender: mpsc::UnboundedSender<Request>,
process_id: i32,
secret_key: i32,
) -> Client {
Client {
2019-07-24 03:16:31 +00:00
inner: Arc::new(InnerClient {
sender,
state: Mutex::new(State {
has_typeinfo: false,
has_typeinfo_composite: false,
has_typeinfo_enum: false,
types: HashMap::new(),
}),
}),
process_id,
secret_key,
}
}
2019-07-24 02:54:22 +00:00
pub(crate) fn inner(&self) -> Arc<InnerClient> {
self.inner.clone()
}
pub fn prepare<'a>(
&mut self,
query: &'a str,
) -> impl Future<Output = Result<Statement, Error>> + 'a {
self.prepare_typed(query, &[])
}
pub fn prepare_typed<'a>(
&mut self,
query: &'a str,
parameter_types: &'a [Type],
) -> impl Future<Output = Result<Statement, Error>> + 'a {
prepare(self.inner(), query, parameter_types)
}
2019-07-25 02:18:15 +00:00
pub fn query<'a>(
&mut self,
statement: &'a Statement,
params: &'a [&dyn ToSql],
) -> impl Future<Output = Result<Query, Error>> + 'a {
self.query_iter(statement, params.iter().cloned())
}
pub fn query_iter<'a, I>(
&mut self,
statement: &'a Statement,
params: I,
) -> impl Future<Output = Result<Query, Error>> + 'a
where
I: IntoIterator<Item = &'a dyn ToSql> + 'a,
I::IntoIter: ExactSizeIterator,
{
query(self.inner(), statement, params)
}
pub fn execute<'a>(
&mut self,
statement: &'a Statement,
params: &'a [&dyn ToSql],
) -> impl Future<Output = Result<u64, Error>> + 'a {
self.execute_iter(statement, params.iter().cloned())
}
pub fn execute_iter<'a, I>(
&mut self,
statement: &'a Statement,
params: I,
) -> impl Future<Output = Result<u64, Error>> + 'a
where
I: IntoIterator<Item = &'a dyn ToSql> + 'a,
I::IntoIter: ExactSizeIterator,
{
execute(self.inner(), statement, params)
}
}