Parameterize Connection over the stream

Author: sfackler

tokio-postgres/src/lib.rs

@@ -28,6 +28,7 @@ use postgres_shared::rows::RowIndex;
 use std::error::Error as StdError;
 use std::fmt;
 use std::sync::atomic::{AtomicUsize, Ordering};
+use tokio_io::{AsyncRead, AsyncWrite};
 
 #[doc(inline)]
 pub use postgres_shared::stmt::Column;
@@ -38,7 +39,7 @@ pub use postgres_shared::{CancelData, Notification};
 
 use error::{DbError, Error};
 use params::ConnectParams;
-use tls::TlsConnect;
+use tls::{TlsConnect, TlsStream};
 use types::{FromSql, ToSql, Type};
 
 pub mod error;
@@ -125,9 +126,12 @@ impl Client {
 }
 
 #[must_use = "futures do nothing unless polled"]
-pub struct Connection(proto::Connection);
+pub struct Connection<S>(proto::Connection<S>);
 
-impl Connection {
+impl<S> Connection<S>
+where
+    S: AsyncRead + AsyncWrite,
+{
     pub fn cancel_data(&self) -> CancelData {
         self.0.cancel_data()
     }
@@ -141,7 +145,10 @@ impl Connection {
     }
 }
 
-impl Future for Connection {
+impl<S> Future for Connection<S>
+where
+    S: AsyncRead + AsyncWrite,
+{
     type Item = ();
     type Error = Error;
 
@@ -173,10 +180,10 @@ impl Future for CancelQuery {
 pub struct Handshake(proto::HandshakeFuture);
 
 impl Future for Handshake {
-    type Item = (Client, Connection);
+    type Item = (Client, Connection<Box<TlsStream>>);
     type Error = Error;
 
-    fn poll(&mut self) -> Poll<(Client, Connection), Error> {
+    fn poll(&mut self) -> Poll<(Client, Connection<Box<TlsStream>>), Error> {
         let (client, connection) = try_ready!(self.0.poll());
 
         Ok(Async::Ready((Client(client), Connection(connection))))

tokio-postgres/src/proto/connection.rs

@@ -5,10 +5,10 @@ use postgres_protocol::message::frontend;
 use std::collections::{HashMap, VecDeque};
 use std::io;
 use tokio_codec::Framed;
+use tokio_io::{AsyncRead, AsyncWrite};
 
 use proto::codec::PostgresCodec;
 use proto::copy_in::CopyInReceiver;
-use tls::TlsStream;
 use {AsyncMessage, CancelData, Notification};
 use {DbError, Error};
 
@@ -32,8 +32,8 @@ enum State {
     Closing,
 }
 
-pub struct Connection {
-    stream: Framed<Box<TlsStream>, PostgresCodec>,
+pub struct Connection<S> {
+    stream: Framed<S, PostgresCodec>,
     cancel_data: CancelData,
     parameters: HashMap<String, String>,
     receiver: mpsc::UnboundedReceiver<Request>,
@@ -43,13 +43,16 @@ pub struct Connection {
     state: State,
 }
 
-impl Connection {
+impl<S> Connection<S>
+where
+    S: AsyncRead + AsyncWrite,
+{
     pub fn new(
-        stream: Framed<Box<TlsStream>, PostgresCodec>,
+        stream: Framed<S, PostgresCodec>,
         cancel_data: CancelData,
         parameters: HashMap<String, String>,
         receiver: mpsc::UnboundedReceiver<Request>,
-    ) -> Connection {
+    ) -> Connection<S> {
         Connection {
             stream,
             cancel_data,
@@ -295,7 +298,10 @@ impl Connection {
     }
 }
 
-impl Future for Connection {
+impl<S> Future for Connection<S>
+where
+    S: AsyncRead + AsyncWrite,
+{
     type Item = ();
     type Error = Error;
 

tokio-postgres/src/proto/handshake.rs

@@ -61,7 +61,7 @@ pub enum Handshake {
         parameters: HashMap<String, String>,
     },
     #[state_machine_future(ready)]
-    Finished((Client, Connection)),
+    Finished((Client, Connection<Box<TlsStream>>)),
     #[state_machine_future(error)]
     Failed(Error),
 }