Some prep for custom type lookup

Author: sfackler

tokio-postgres/Cargo.toml

@@ -31,11 +31,11 @@ with-serde_json-1 = ["serde-1", "serde_json-1"]
 "with-uuid-0_7" = ["uuid-07"]
 
 [dependencies]
-antidote = "1.0"
 bytes = "0.4"
 fallible-iterator = "0.2"
 futures-preview = { version = "0.3.0-alpha.17", features = ["nightly", "async-await"] }
 log = "0.4"
+parking_lot = "0.9"
 percent-encoding = "1.0"
 phf = "0.7.23"
 postgres-protocol = { version = "0.4.1", path = "../postgres-protocol" }

tokio-postgres/src/client.rs

@@ -1,12 +1,14 @@
 use crate::codec::BackendMessages;
 use crate::connection::{Request, RequestMessages};
 use crate::prepare::prepare;
-use crate::types::Type;
+use crate::types::{Oid, Type};
 use crate::{Error, Statement};
 use fallible_iterator::FallibleIterator;
 use futures::channel::mpsc;
 use futures::{Stream, StreamExt};
+use parking_lot::Mutex;
 use postgres_protocol::message::backend::Message;
+use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -34,8 +36,16 @@ impl Responses {
     }
 }
 
+struct State {
+    has_typeinfo: bool,
+    has_typeinfo_composite: bool,
+    has_typeinfo_enum: bool,
+    types: HashMap<Oid, Type>,
+}
+
 pub struct InnerClient {
     sender: mpsc::UnboundedSender<Request>,
+    state: Mutex<State>,
 }
 
 impl InnerClient {
@@ -51,6 +61,38 @@ impl InnerClient {
             cur: BackendMessages::empty(),
         })
     }
+
+    pub fn has_typeinfo(&self) -> bool {
+        self.state.lock().has_typeinfo
+    }
+
+    pub fn set_has_typeinfo(&self) {
+        self.state.lock().has_typeinfo = true;
+    }
+
+    pub fn has_typeinfo_composite(&self) -> bool {
+        self.state.lock().has_typeinfo_composite
+    }
+
+    pub fn set_has_typeinfo_composite(&self) {
+        self.state.lock().has_typeinfo_composite = true;
+    }
+
+    pub fn has_typeinfo_enum(&self) -> bool {
+        self.state.lock().has_typeinfo_enum
+    }
+
+    pub fn set_has_typeinfo_enum(&self) {
+        self.state.lock().has_typeinfo_enum = true;
+    }
+
+    pub fn type_(&self, oid: Oid) -> Option<Type> {
+        self.state.lock().types.get(&oid).cloned()
+    }
+
+    pub fn set_type(&self, oid: Oid, type_: Type) {
+        self.state.lock().types.insert(oid, type_);
+    }
 }
 
 pub struct Client {
@@ -66,7 +108,15 @@ impl Client {
         secret_key: i32,
     ) -> Client {
         Client {
-            inner: Arc::new(InnerClient { sender }),
+            inner: Arc::new(InnerClient {
+                sender,
+                state: Mutex::new(State {
+                    has_typeinfo: false,
+                    has_typeinfo_composite: false,
+                    has_typeinfo_enum: false,
+                    types: HashMap::new(),
+                }),
+            }),
             process_id,
             secret_key,
         }

tokio-postgres/src/prepare.rs

@@ -62,10 +62,14 @@ pub async fn prepare(
     Ok(Statement::new(&client, name, parameters, columns))
 }
 
-async fn get_type(client: &InnerClient, oid: Oid) -> Result<Type, Error> {
+async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
     if let Some(type_) = Type::from_oid(oid) {
         return Ok(type_);
     }
 
+    if let Some(type_) = client.type_(oid) {
+        return Ok(type_);
+    }
+
     unimplemented!()
 }

tokio-postgres/tests/test/main.rs

@@ -5,8 +5,8 @@ use futures::{try_join, FutureExt};
 use tokio::net::TcpStream;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::tls::{NoTls, NoTlsStream};
-use tokio_postgres::{Client, Config, Connection, Error};
 use tokio_postgres::types::Type;
+use tokio_postgres::{Client, Config, Connection, Error};
 
 mod parse;
 #[cfg(feature = "runtime")]
@@ -113,24 +113,6 @@ async fn pipelined_prepare() {
 }
 
 /*
-#[test]
-fn pipelined_prepare() {
-    let _ = env_logger::try_init();
-    let mut runtime = Runtime::new().unwrap();
-
-    let (mut client, connection) = runtime.block_on(connect("user=postgres")).unwrap();
-    let connection = connection.map_err(|e| panic!("{}", e));
-    runtime.handle().spawn(connection).unwrap();
-
-    let prepare1 = client.prepare("SELECT $1::HSTORE[]");
-    let prepare2 = client.prepare("SELECT $1::HSTORE[]");
-    let prepare = prepare1.join(prepare2);
-    runtime.block_on(prepare).unwrap();
-
-    drop(client);
-    runtime.run().unwrap();
-}
-
 #[test]
 fn insert_select() {
     let _ = env_logger::try_init();