#[cfg(feature = "runtime")] use crate::cancel_query; use crate::codec::BackendMessages; use crate::config::{Host, SslMode}; use crate::connection::{Request, RequestMessages}; use crate::simple_query::SimpleQueryStream; use crate::copy_out::CopyStream; use crate::query::RowStream; use crate::slice_iter; #[cfg(feature = "runtime")] use crate::tls::MakeTlsConnect; use crate::tls::TlsConnect; use crate::to_statement::ToStatement; use crate::types::{Oid, ToSql, Type}; #[cfg(feature = "runtime")] use crate::Socket; use crate::{cancel_query_raw, copy_in, copy_out, query, Transaction}; use crate::{prepare, SimpleQueryMessage}; use crate::{simple_query, Row}; use crate::{Error, Statement}; use bytes::IntoBuf; use fallible_iterator::FallibleIterator; use futures::channel::mpsc; use futures::{future, TryStream, TryStreamExt}; use futures::{ready, StreamExt}; use parking_lot::Mutex; use postgres_protocol::message::backend::Message; use std::collections::HashMap; use std::error; use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Duration; use tokio::io::{AsyncRead, AsyncWrite}; pub struct Responses { receiver: mpsc::Receiver, cur: BackendMessages, } impl Responses { pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { loop { match self.cur.next().map_err(Error::parse)? { Some(Message::ErrorResponse(body)) => return Poll::Ready(Err(Error::db(body))), Some(message) => return Poll::Ready(Ok(message)), None => {} } match ready!(self.receiver.poll_next_unpin(cx)) { Some(messages) => self.cur = messages, None => return Poll::Ready(Err(Error::closed())), } } } pub async fn next(&mut self) -> Result { future::poll_fn(|cx| self.poll_next(cx)).await } } struct State { typeinfo: Option, typeinfo_composite: Option, typeinfo_enum: Option, types: HashMap, } pub struct InnerClient { sender: mpsc::UnboundedSender, state: Mutex, } impl InnerClient { pub fn send(&self, messages: RequestMessages) -> Result { 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(), }) } pub fn typeinfo(&self) -> Option { self.state.lock().typeinfo.clone() } pub fn set_typeinfo(&self, statement: &Statement) { self.state.lock().typeinfo = Some(statement.clone()); } pub fn typeinfo_composite(&self) -> Option { self.state.lock().typeinfo_composite.clone() } pub fn set_typeinfo_composite(&self, statement: &Statement) { self.state.lock().typeinfo_composite = Some(statement.clone()); } pub fn typeinfo_enum(&self) -> Option { self.state.lock().typeinfo_enum.clone() } pub fn set_typeinfo_enum(&self, statement: &Statement) { self.state.lock().typeinfo_enum = Some(statement.clone()); } pub fn type_(&self, oid: Oid) -> Option { self.state.lock().types.get(&oid).cloned() } pub fn set_type(&self, oid: Oid, type_: &Type) { self.state.lock().types.insert(oid, type_.clone()); } } #[derive(Clone)] pub(crate) struct SocketConfig { pub host: Host, pub port: u16, pub connect_timeout: Option, pub keepalives: bool, pub keepalives_idle: Duration, } /// An asynchronous PostgreSQL client. /// /// The client is one half of what is returned when a connection is established. Users interact with the database /// through this client object. pub struct Client { inner: Arc, #[cfg(feature = "runtime")] socket_config: Option, ssl_mode: SslMode, process_id: i32, secret_key: i32, } impl Client { pub(crate) fn new( sender: mpsc::UnboundedSender, ssl_mode: SslMode, process_id: i32, secret_key: i32, ) -> Client { Client { inner: Arc::new(InnerClient { sender, state: Mutex::new(State { typeinfo: None, typeinfo_composite: None, typeinfo_enum: None, types: HashMap::new(), }), }), #[cfg(feature = "runtime")] socket_config: None, ssl_mode, process_id, secret_key, } } pub(crate) fn inner(&self) -> &Arc { &self.inner } #[cfg(feature = "runtime")] pub(crate) fn set_socket_config(&mut self, socket_config: SocketConfig) { self.socket_config = Some(socket_config); } /// Creates a new prepared statement. /// /// Prepared statements can be executed repeatedly, and may contain query parameters (indicated by `$1`, `$2`, etc), /// which are set when executed. Prepared statements can only be used with the connection that created them. pub async fn prepare(&self, query: &str) -> Result { self.prepare_typed(query, &[]).await } /// Like `prepare`, but allows the types of query parameters to be explicitly specified. /// /// The list of types may be smaller than the number of parameters - the types of the remaining parameters will be /// inferred. For example, `client.prepare_typed(query, &[])` is equivalent to `client.prepare(query)`. pub async fn prepare_typed( &self, query: &str, parameter_types: &[Type], ) -> Result { prepare::prepare(&self.inner, query, parameter_types).await } /// Executes a statement, returning a vector of the resulting rows. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. pub async fn query( &self, statement: &T, params: &[&(dyn ToSql + Sync)], ) -> Result, Error> where T: ?Sized + ToStatement, { self.query_raw(statement, slice_iter(params)) .await? .try_collect() .await } /// The maximally flexible version of [`query`]. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. /// /// [`query`]: #method.query pub async fn query_raw<'a, T, I>(&self, statement: &T, params: I) -> Result where T: ?Sized + ToStatement, I: IntoIterator, I::IntoIter: ExactSizeIterator, { let statement = statement.__convert().into_statement(self).await?; query::query(&self.inner, statement, params).await } /// Executes a statement, returning the number of rows modified. /// /// If the statement does not modify any rows (e.g. `SELECT`), 0 is returned. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. pub async fn execute( &self, statement: &T, params: &[&(dyn ToSql + Sync)], ) -> Result where T: ?Sized + ToStatement, { self.execute_raw(statement, slice_iter(params)).await } /// The maximally flexible version of [`execute`]. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. /// /// [`execute`]: #method.execute pub async fn execute_raw<'a, T, I>(&self, statement: &T, params: I) -> Result where T: ?Sized + ToStatement, I: IntoIterator, I::IntoIter: ExactSizeIterator, { let statement = statement.__convert().into_statement(self).await?; query::execute(self.inner(), statement, params).await } /// Executes a `COPY FROM STDIN` statement, returning the number of rows created. /// /// The data in the provided stream is passed along to the server verbatim; it is the caller's responsibility to /// ensure it uses the proper format. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. pub async fn copy_in( &self, statement: &T, params: &[&(dyn ToSql + Sync)], stream: S, ) -> Result where T: ?Sized + ToStatement, S: TryStream, S::Ok: IntoBuf, ::Buf: 'static + Send, S::Error: Into>, { let statement = statement.__convert().into_statement(self).await?; let params = slice_iter(params); copy_in::copy_in(self.inner(), statement, params, stream).await } /// Executes a `COPY TO STDOUT` statement, returning a stream of the resulting data. /// /// # Panics /// /// Panics if the number of parameters provided does not match the number expected. pub async fn copy_out( &self, statement: &T, params: &[&(dyn ToSql + Sync)], ) -> Result where T: ?Sized + ToStatement, { let statement = statement.__convert().into_statement(self).await?; let params = slice_iter(params); copy_out::copy_out(self.inner(), statement, params).await } /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows. /// /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that /// point. The simple query protocol returns the values in rows as strings rather than in their binary encodings, /// so the associated row type doesn't work with the `FromSql` trait. Rather than simply returning a stream over the /// rows, this method returns a stream over an enum which indicates either the completion of one of the commands, /// or a row of data. This preserves the framing between the separate statements in the request. /// /// # Warning /// /// Prepared statements should be use for any query which contains user-specified data, as they provided the /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass /// them to this method! pub async fn simple_query( &self, query: &str, ) -> Result, Error> { self.simple_query_raw(query).await?.try_collect().await } pub(crate) async fn simple_query_raw(&self, query: &str) -> Result { simple_query::simple_query(self.inner(), query).await } /// Executes a sequence of SQL statements using the simple query protocol. /// /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that /// point. This is intended for use when, for example, initializing a database schema. /// /// # Warning /// /// Prepared statements should be use for any query which contains user-specified data, as they provided the /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass /// them to this method! pub async fn batch_execute(&self, query: &str) -> Result<(), Error> { simple_query::batch_execute(self.inner(), query).await } /// Begins a new database transaction. /// /// The transaction will roll back by default - use the `commit` method to commit it. pub async fn transaction(&mut self) -> Result, Error> { self.batch_execute("BEGIN").await?; Ok(Transaction::new(self)) } /// Attempts to cancel an in-progress query. /// /// The server provides no information about whether a cancellation attempt was successful or not. An error will /// only be returned if the client was unable to connect to the database. /// /// Requires the `runtime` Cargo feature (enabled by default). #[cfg(feature = "runtime")] pub async fn cancel_query(&self, tls: T) -> Result<(), Error> where T: MakeTlsConnect, { cancel_query::cancel_query( self.socket_config.clone(), self.ssl_mode, tls, self.process_id, self.secret_key, ) .await } /// Like `cancel_query`, but uses a stream which is already connected to the server rather than opening a new /// connection itself. pub async fn cancel_query_raw(&self, stream: S, tls: T) -> Result<(), Error> where S: AsyncRead + AsyncWrite + Unpin, T: TlsConnect, { cancel_query_raw::cancel_query_raw( stream, self.ssl_mode, tls, self.process_id, self.secret_key, ) .await } /// Determines if the connection to the server has already closed. /// /// In that case, all future queries will fail. pub fn is_closed(&self) -> bool { self.inner.sender.is_closed() } }