Support custom types

Author: sfackler

tokio-postgres/Cargo.toml

@@ -32,6 +32,7 @@ circle-ci = { repository = "sfackler/rust-postgres" }
 "with-uuid-0.6" = ["postgres-shared/with-uuid-0.6"]
 
 [dependencies]
+antidote = "1.0"
 bytes = "0.4"
 fallible-iterator = "0.1.3"
 futures = "0.1.7"

tokio-postgres/src/lib.rs

@@ -1,3 +1,4 @@
+extern crate antidote;
 extern crate bytes;
 extern crate fallible_iterator;
 extern crate futures_cpupool;

tokio-postgres/src/proto/client.rs

@@ -1,7 +1,10 @@
+use antidote::Mutex;
 use futures::sync::mpsc;
 use postgres_protocol;
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
+use std::collections::HashMap;
+use std::sync::Arc;
 
 use disconnected;
 use error::{self, Error};
@@ -10,7 +13,7 @@ use proto::execute::ExecuteFuture;
 use proto::prepare::PrepareFuture;
 use proto::query::QueryStream;
 use proto::statement::Statement;
-use types::{IsNull, ToSql, Type};
+use types::{IsNull, Oid, ToSql, Type};
 
 pub struct PendingRequest {
     sender: mpsc::UnboundedSender<Request>,
@@ -28,13 +31,30 @@ impl PendingRequest {
     }
 }
 
+pub struct State {
+    pub types: HashMap<Oid, Type>,
+    pub typeinfo_query: Option<Statement>,
+    pub typeinfo_enum_query: Option<Statement>,
+    pub typeinfo_composite_query: Option<Statement>,
+}
+
+#[derive(Clone)]
 pub struct Client {
+    pub state: Arc<Mutex<State>>,
     sender: mpsc::UnboundedSender<Request>,
 }
 
 impl Client {
     pub fn new(sender: mpsc::UnboundedSender<Request>) -> Client {
-        Client { sender }
+        Client {
+            state: Arc::new(Mutex::new(State {
+                types: HashMap::new(),
+                typeinfo_query: None,
+                typeinfo_enum_query: None,
+                typeinfo_composite_query: None,
+            })),
+            sender,
+        }
     }
 
     pub fn prepare(&mut self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
@@ -45,7 +65,7 @@ impl Client {
             Ok(())
         });
 
-        PrepareFuture::new(pending, self.sender.clone(), name)
+        PrepareFuture::new(pending, self.sender.clone(), name, self.clone())
     }
 
     pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {

tokio-postgres/src/proto/mod.rs

@@ -20,6 +20,9 @@ mod query;
 mod row;
 mod socket;
 mod statement;
+mod typeinfo;
+mod typeinfo_composite;
+mod typeinfo_enum;
 
 pub use proto::cancel::CancelFuture;
 pub use proto::client::Client;

tokio-postgres/src/proto/prepare.rs

@@ -1,14 +1,17 @@
 use fallible_iterator::FallibleIterator;
 use futures::sync::mpsc;
-use futures::{Poll, Stream};
-use postgres_protocol::message::backend::{Message, ParameterDescriptionBody, RowDescriptionBody};
+use futures::{Future, Poll, Stream};
+use postgres_protocol::message::backend::Message;
 use state_machine_future::RentToOwn;
+use std::mem;
+use std::vec;
 
 use error::{self, Error};
-use proto::client::PendingRequest;
+use proto::client::{Client, PendingRequest};
 use proto::connection::Request;
 use proto::statement::Statement;
-use types::Type;
+use proto::typeinfo::TypeinfoFuture;
+use types::{Oid, Type};
 use Column;
 use {bad_response, disconnected};
 
@@ -19,33 +22,57 @@ pub enum Prepare {
         request: PendingRequest,
         sender: mpsc::UnboundedSender<Request>,
         name: String,
+        client: Client,
     },
     #[state_machine_future(transitions(ReadParameterDescription))]
     ReadParseComplete {
         sender: mpsc::UnboundedSender<Request>,
         receiver: mpsc::Receiver<Message>,
         name: String,
+        client: Client,
     },
     #[state_machine_future(transitions(ReadRowDescription))]
     ReadParameterDescription {
         sender: mpsc::UnboundedSender<Request>,
         receiver: mpsc::Receiver<Message>,
         name: String,
+        client: Client,
     },
     #[state_machine_future(transitions(ReadReadyForQuery))]
     ReadRowDescription {
         sender: mpsc::UnboundedSender<Request>,
         receiver: mpsc::Receiver<Message>,
         name: String,
-        parameters: ParameterDescriptionBody,
+        parameters: Vec<Oid>,
+        client: Client,
     },
-    #[state_machine_future(transitions(Finished))]
+    #[state_machine_future(transitions(GetParameterTypes, GetColumnTypes, Finished))]
     ReadReadyForQuery {
         sender: mpsc::UnboundedSender<Request>,
         receiver: mpsc::Receiver<Message>,
         name: String,
-        parameters: ParameterDescriptionBody,
-        columns: Option<RowDescriptionBody>,
+        parameters: Vec<Oid>,
+        columns: Vec<(String, Oid)>,
+        client: Client,
+    },
+    #[state_machine_future(transitions(GetColumnTypes, Finished))]
+    GetParameterTypes {
+        future: TypeinfoFuture,
+        remaining_parameters: vec::IntoIter<Oid>,
+        sender: mpsc::UnboundedSender<Request>,
+        name: String,
+        parameters: Vec<Type>,
+        columns: Vec<(String, Oid)>,
+    },
+    #[state_machine_future(transitions(Finished))]
+    GetColumnTypes {
+        future: TypeinfoFuture,
+        cur_column_name: String,
+        remaining_columns: vec::IntoIter<(String, Oid)>,
+        sender: mpsc::UnboundedSender<Request>,
+        name: String,
+        parameters: Vec<Type>,
+        columns: Vec<Column>,
     },
     #[state_machine_future(ready)]
     Finished(Statement),
@@ -62,6 +89,7 @@ impl PollPrepare for Prepare {
             sender: state.sender,
             receiver,
             name: state.name,
+            client: state.client,
         })
     }
 
@@ -76,6 +104,7 @@ impl PollPrepare for Prepare {
                 sender: state.sender,
                 receiver: state.receiver,
                 name: state.name,
+                client: state.client,
             }),
             Some(Message::ErrorResponse(body)) => Err(error::__db(body)),
             Some(_) => Err(bad_response()),
@@ -94,7 +123,8 @@ impl PollPrepare for Prepare {
                 sender: state.sender,
                 receiver: state.receiver,
                 name: state.name,
-                parameters: body,
+                parameters: body.parameters().collect()?,
+                client: state.client,
             }),
             Some(_) => Err(bad_response()),
             None => Err(disconnected()),
@@ -107,9 +137,12 @@ impl PollPrepare for Prepare {
         let message = try_receive!(state.receiver.poll());
         let state = state.take();
 
-        let body = match message {
-            Some(Message::RowDescription(body)) => Some(body),
-            Some(Message::NoData) => None,
+        let columns = match message {
+            Some(Message::RowDescription(body)) => body
+                .fields()
+                .map(|f| (f.name().to_string(), f.type_oid()))
+                .collect()?,
+            Some(Message::NoData) => vec![],
             Some(_) => return Err(bad_response()),
             None => return Err(disconnected()),
         };
@@ -119,7 +152,8 @@ impl PollPrepare for Prepare {
             receiver: state.receiver,
             name: state.name,
             parameters: state.parameters,
-            columns: body,
+            columns,
+            client: state.client,
         })
     }
 
@@ -130,33 +164,103 @@ impl PollPrepare for Prepare {
         let state = state.take();
 
         match message {
-            Some(Message::ReadyForQuery(_)) => {
-                // FIXME handle custom types
-                let parameters = state
-                    .parameters
-                    .parameters()
-                    .map(|oid| Type::from_oid(oid).unwrap())
-                    .collect()?;
-                let columns = match state.columns {
-                    Some(body) => body
-                        .fields()
-                        .map(|f| {
-                            Column::new(f.name().to_string(), Type::from_oid(f.type_oid()).unwrap())
-                        })
-                        .collect()?,
-                    None => vec![],
-                };
-
-                transition!(Finished(Statement::new(
-                    state.sender,
-                    state.name,
-                    parameters,
-                    columns
-                )))
+            Some(Message::ReadyForQuery(_)) => {}
+            Some(_) => return Err(bad_response()),
+            None => return Err(disconnected()),
+        }
+
+        let mut parameters = state.parameters.into_iter();
+        if let Some(oid) = parameters.next() {
+            transition!(GetParameterTypes {
+                future: TypeinfoFuture::new(oid, state.client),
+                remaining_parameters: parameters,
+                sender: state.sender,
+                name: state.name,
+                parameters: vec![],
+                columns: state.columns,
+            });
+        }
+
+        let mut columns = state.columns.into_iter();
+        if let Some((name, oid)) = columns.next() {
+            transition!(GetColumnTypes {
+                future: TypeinfoFuture::new(oid, state.client),
+                cur_column_name: name,
+                remaining_columns: columns,
+                sender: state.sender,
+                name: state.name,
+                parameters: vec![],
+                columns: vec![],
+            });
+        }
+
+        transition!(Finished(Statement::new(
+            state.sender,
+            state.name,
+            vec![],
+            vec![]
+        )))
+    }
+
+    fn poll_get_parameter_types<'a>(
+        state: &'a mut RentToOwn<'a, GetParameterTypes>,
+    ) -> Poll<AfterGetParameterTypes, Error> {
+        let client = loop {
+            let (ty, client) = try_ready!(state.future.poll());
+            state.parameters.push(ty);
+
+            match state.remaining_parameters.next() {
+                Some(oid) => state.future = TypeinfoFuture::new(oid, client),
+                None => break client,
             }
-            Some(_) => Err(bad_response()),
-            None => Err(disconnected()),
+        };
+        let state = state.take();
+
+        let mut columns = state.columns.into_iter();
+        if let Some((name, oid)) = columns.next() {
+            transition!(GetColumnTypes {
+                future: TypeinfoFuture::new(oid, client),
+                cur_column_name: name,
+                remaining_columns: columns,
+                sender: state.sender,
+                name: state.name,
+                parameters: state.parameters,
+                columns: vec![],
+            })
         }
+
+        transition!(Finished(Statement::new(
+            state.sender,
+            state.name,
+            state.parameters,
+            vec![],
+        )))
+    }
+
+    fn poll_get_column_types<'a>(
+        state: &'a mut RentToOwn<'a, GetColumnTypes>,
+    ) -> Poll<AfterGetColumnTypes, Error> {
+        loop {
+            let (ty, client) = try_ready!(state.future.poll());
+            let name = mem::replace(&mut state.cur_column_name, String::new());
+            state.columns.push(Column::new(name, ty));
+
+            match state.remaining_columns.next() {
+                Some((name, oid)) => {
+                    state.cur_column_name = name;
+                    state.future = TypeinfoFuture::new(oid, client);
+                }
+                None => break,
+            }
+        }
+        let state = state.take();
+
+        transition!(Finished(Statement::new(
+            state.sender,
+            state.name,
+            state.parameters,
+            state.columns,
+        )))
     }
 }
 
@@ -165,7 +269,8 @@ impl PrepareFuture {
         request: PendingRequest,
         sender: mpsc::UnboundedSender<Request>,
         name: String,
+        client: Client,
     ) -> PrepareFuture {
-        Prepare::start(request, sender, name)
+        Prepare::start(request, sender, name, client)
     }
 }