diff --git a/postgres/src/config.rs b/postgres/src/config.rs index 7979f0d1..5982e4fc 100644 --- a/postgres/src/config.rs +++ b/postgres/src/config.rs @@ -1,8 +1,11 @@ +use futures::future::Executor; use futures::sync::oneshot; use futures::Future; use log::error; +use std::fmt; use std::path::Path; use std::str::FromStr; +use std::sync::Arc; use std::time::Duration; use tokio_postgres::config::{SslMode, TargetSessionAttrs}; use tokio_postgres::tls::{MakeTlsConnect, TlsConnect}; @@ -10,22 +13,36 @@ use tokio_postgres::{Error, Socket}; use crate::{Client, RUNTIME}; -#[derive(Debug, Clone, PartialEq)] -pub struct Config(tokio_postgres::Config); +#[derive(Clone)] +pub struct Config { + config: tokio_postgres::Config, + executor: Option + Send>>>>, +} + +impl fmt::Debug for Config { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Config") + .field("config", &self.config) + .finish() + } +} impl Default for Config { fn default() -> Config { - Config(tokio_postgres::Config::default()) + Config::new() } } impl Config { pub fn new() -> Config { - Config(tokio_postgres::Config::new()) + Config { + config: tokio_postgres::Config::new(), + executor: None, + } } pub fn user(&mut self, user: &str) -> &mut Config { - self.0.user(user); + self.config.user(user); self } @@ -33,32 +50,32 @@ impl Config { where T: AsRef<[u8]>, { - self.0.password(password); + self.config.password(password); self } pub fn dbname(&mut self, dbname: &str) -> &mut Config { - self.0.dbname(dbname); + self.config.dbname(dbname); self } pub fn options(&mut self, options: &str) -> &mut Config { - self.0.options(options); + self.config.options(options); self } pub fn application_name(&mut self, application_name: &str) -> &mut Config { - self.0.application_name(application_name); + self.config.application_name(application_name); self } pub fn ssl_mode(&mut self, ssl_mode: SslMode) -> &mut Config { - self.0.ssl_mode(ssl_mode); + self.config.ssl_mode(ssl_mode); self } pub fn host(&mut self, host: &str) -> &mut Config { - self.0.host(host); + self.config.host(host); self } @@ -67,27 +84,27 @@ impl Config { where T: AsRef, { - self.0.host_path(host); + self.config.host_path(host); self } pub fn port(&mut self, port: u16) -> &mut Config { - self.0.port(port); + self.config.port(port); self } pub fn connect_timeout(&mut self, connect_timeout: Duration) -> &mut Config { - self.0.connect_timeout(connect_timeout); + self.config.connect_timeout(connect_timeout); self } pub fn keepalives(&mut self, keepalives: bool) -> &mut Config { - self.0.keepalives(keepalives); + self.config.keepalives(keepalives); self } pub fn keepalives_idle(&mut self, keepalives_idle: Duration) -> &mut Config { - self.0.keepalives_idle(keepalives_idle); + self.config.keepalives_idle(keepalives_idle); self } @@ -95,7 +112,15 @@ impl Config { &mut self, target_session_attrs: TargetSessionAttrs, ) -> &mut Config { - self.0.target_session_attrs(target_session_attrs); + self.config.target_session_attrs(target_session_attrs); + self + } + + pub fn executor(&mut self, executor: E) -> &mut Config + where + E: Executor + Send>> + 'static + Sync + Send, + { + self.executor = Some(Arc::new(executor)); self } @@ -106,19 +131,46 @@ impl Config { T::Stream: Send, >::Future: Send, { - let connect = self.0.connect(tls_mode); - let (client, connection) = oneshot::spawn(connect, &RUNTIME.executor()).wait()?; + let (tx, rx) = oneshot::channel(); + let connect = self + .config + .connect(tls_mode) + .then(|r| tx.send(r).map_err(|_| ())); + self.with_executor(|e| e.execute(Box::new(connect))) + .unwrap(); + let (client, connection) = rx.wait().unwrap()?; + let connection = connection.map_err(|e| error!("postgres connection error: {}", e)); - RUNTIME.executor().spawn(connection); + self.with_executor(|e| e.execute(Box::new(connection))) + .unwrap(); Ok(Client::from(client)) } + + fn with_executor(&self, f: F) -> T + where + F: FnOnce(&Executor + Send>>) -> T, + { + match &self.executor { + Some(e) => f(&**e), + None => f(&RUNTIME.executor()), + } + } } impl FromStr for Config { type Err = Error; fn from_str(s: &str) -> Result { - s.parse().map(Config) + s.parse::().map(Config::from) + } +} + +impl From for Config { + fn from(config: tokio_postgres::Config) -> Config { + Config { + config, + executor: None, + } } }