Allow custom executors

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<Arc<Executor<Box<Future<Item = (), Error = ()> + 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<Path>,
     {
-        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<E>(&mut self, executor: E) -> &mut Config
+    where
+        E: Executor<Box<Future<Item = (), Error = ()> + Send>> + 'static + Sync + Send,
+    {
+        self.executor = Some(Arc::new(executor));
         self
     }
 
@@ -106,19 +131,46 @@ impl Config {
         T::Stream: Send,
         <T::TlsConnect as TlsConnect<Socket>>::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<F, T>(&self, f: F) -> T
+    where
+        F: FnOnce(&Executor<Box<Future<Item = (), Error = ()> + 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<Config, Error> {
-        s.parse().map(Config)
+        s.parse::<tokio_postgres::Config>().map(Config::from)
+    }
+}
+
+impl From<tokio_postgres::Config> for Config {
+    fn from(config: tokio_postgres::Config) -> Config {
+        Config {
+            config,
+            executor: None,
+        }
     }
 }