Support composites

sfackler · sfackler · commit d0f50fea3139 · 2016-10-07T22:03:05.000-07:00
diff --git a/postgres-derive-internals/src/accepts.rs b/postgres-derive-internals/src/accepts.rs
@@ -1,6 +1,8 @@
 use std::fmt::Write;
+use quote::{Tokens, ToTokens};
 
 use enums::Variant;
+use composites::Field;
 
 pub fn enum_body(name: &str, variants: &[Variant]) -> String {
     let mut body = String::new();
@@ -34,3 +36,39 @@ pub fn enum_body(name: &str, variants: &[Variant]) -> String {
 
     body
 }
+
+pub fn composite_body(name: &str, trait_: &str, fields: &[Field]) -> String {
+    let mut body = String::new();
+
+    write!(body, "
+        if type_.name() != \"{}\" {{
+            return false;
+        }}
+
+        match *type_.kind() {{
+            ::postgres::types::Kind::Composite(ref fields) => {{
+                if fields.len() != {} {{
+                    return false;
+                }}
+
+                fields.iter().all(|f| {{
+                    match f.name() {{", name, fields.len()).unwrap();
+
+    for field in fields {
+        let mut tokens = Tokens::new();
+        field.type_.to_tokens(&mut tokens);
+        write!(body, "
+                        \"{}\" => <{} as ::postgres::types::{}>::accepts(f.type_()),",
+               field.name, tokens, trait_).unwrap();
+    }
+
+    write!(body, "\
+                        _ => false,\
+                    }}\
+                }})\
+            }}\
+            _ => false,\
+        }}").unwrap();
+
+    body
+}
diff --git a/postgres-derive-internals/src/composites.rs b/postgres-derive-internals/src/composites.rs
@@ -0,0 +1,22 @@
+use syn::{self, Ident, Ty};
+
+use overrides::Overrides;
+
+pub struct Field {
+    pub name: String,
+    pub ident: Ident,
+    pub type_: Ty,
+}
+
+impl Field {
+    pub fn parse(raw: &syn::Field) -> Result<Field, String> {
+        let overrides = try!(Overrides::extract(&raw.attrs));
+
+        let ident = raw.ident.as_ref().unwrap().clone();
+        Ok(Field {
+            name: overrides.name.unwrap_or_else(|| ident.to_string()),
+            ident: ident,
+            type_: raw.ty.clone(),
+        })
+    }
+}
diff --git a/postgres-derive-internals/src/fromsql.rs b/postgres-derive-internals/src/fromsql.rs
@@ -1,8 +1,9 @@
 use std::fmt::Write;
-use syn::{Body, Ident, MacroInput, VariantData, Field};
+use syn::{self, Body, Ident, MacroInput, VariantData};
 use quote::{Tokens, ToTokens};
 
 use accepts;
+use composites::Field;
 use enums::Variant;
 use overrides::Overrides;
 
@@ -20,6 +21,11 @@ pub fn expand_derive_fromsql(input: &MacroInput) -> Result<String, String> {
             let field = &fields[0];
             (domain_accepts_body(field), domain_body(&input.ident, field))
         }
+        Body::Struct(VariantData::Struct(ref fields)) => {
+            let fields: Vec<Field> = try!(fields.iter().map(Field::parse).collect());
+            (accepts::composite_body(&name, "FromSql", &fields),
+             composite_body(&input.ident, &fields))
+        }
         _ => {
             return Err("#[derive(ToSql)] may only be applied to structs, single field tuple \
                         structs, and enums".to_owned())
@@ -64,16 +70,108 @@ fn enum_body(ident: &Ident, variants: &[Variant]) -> String {
     out
 }
 
-fn domain_accepts_body(field: &Field) -> String {
+fn domain_accepts_body(field: &syn::Field) -> String {
     let mut tokens = Tokens::new();
     field.ty.to_tokens(&mut tokens);
     format!("
         <{} as ::postgres::types::FromSql>::accepts(type_)", tokens)
 }
 
-fn domain_body(ident: &Ident, field: &Field) -> String {
+fn domain_body(ident: &Ident, field: &syn::Field) -> String {
     let mut tokens = Tokens::new();
     field.ty.to_tokens(&mut tokens);
     format!("\
         <{} as ::postgres::types::FromSql>::from_sql(_type, buf, _info).map({})", tokens, ident)
 }
+
+fn composite_body(ident: &Ident, fields: &[Field]) -> String {
+    let mut out = "
+        fn read_be_i32(buf: &mut &[u8]) -> ::std::io::Result<i32> {
+            let mut bytes = [0; 4];
+            try!(::std::io::Read::read_exact(buf, &mut bytes));
+            let num = ((bytes[0] as i32) << 24) |
+                      ((bytes[1] as i32) << 16) |
+                      ((bytes[2] as i32) << 8) |
+                      (bytes[3] as i32);
+            ::std::result::Result::Ok(num)
+        }
+
+        fn read_value<T>(type_: &::postgres::types::Type,
+                         buf: &mut &[u8],
+                         info: &::postgres::types::SessionInfo)
+                         -> ::std::result::Result<T,
+                                                  ::std::boxed::Box<::std::error::Error +
+                                                                    ::std::marker::Sync +
+                                                                    ::std::marker::Send>>
+            where T: ::postgres::types::FromSql
+        {
+            let len = try!(read_be_i32(buf));
+            let value = if len < 0 {
+                ::std::option::Option::None
+            } else {
+                if len as usize > buf.len() {
+                    return ::std::result::Result::Err(
+                        ::std::convert::Into::into(\"invalid buffer size\"));
+                }
+                let (head, tail) = buf.split_at(len as usize);
+                *buf = tail;
+                ::std::option::Option::Some(&head[..])
+            };
+            ::postgres::types::FromSql::from_sql_nullable(type_, value, info)
+        }
+
+        let fields = match *_type.kind() {
+            ::postgres::types::Kind::Composite(ref fields) => fields,
+            _ => unreachable!(),
+        };
+
+        let mut buf = buf;
+        let num_fields = try!(read_be_i32(&mut buf));
+        if num_fields as usize != fields.len() {
+            return ::std::result::Result::Err(
+                ::std::convert::Into::into(format!(\"invalid field count: {} vs {}\", num_fields,
+                                                   fields.len())));
+        }
+        ".to_owned();
+
+    for field in fields {
+        write!(out, "
+        let mut __{} = ::std::option::Option::None;", field.ident).unwrap();
+    }
+
+    write!(out, "
+        for field in fields {{
+            let oid = try!(read_be_i32(&mut buf)) as u32;
+            if oid != field.type_().oid() {{
+                return ::std::result::Result::Err(
+                    ::std::convert::Into::into(\"unexpected OID\"));
+            }}\
+            \
+            match field.name() {{").unwrap();
+
+    for field in fields {
+        write!(out, "
+                \"{}\" => {{
+                    __{} = ::std::option::Option::Some(
+                        try!(read_value(field.type_(), &mut buf, _info)));
+                }}",
+               field.name, field.ident).unwrap();
+    }
+
+    write!(out, "
+                _ => unreachable!(),
+            }}
+        }}
+
+        ::std::result::Result::Ok({} {{", ident).unwrap();
+
+    for field in fields {
+        write!(out, "
+            {}: __{}.unwrap(),", field.ident, field.ident).unwrap();
+    }
+
+    write!(out, "
+        }})").unwrap();
+
+    out
+}
diff --git a/postgres-derive-internals/src/lib.rs b/postgres-derive-internals/src/lib.rs
@@ -9,6 +9,7 @@ use fromsql::expand_derive_fromsql;
 use tosql::expand_derive_tosql;
 
 mod accepts;
+mod composites;
 mod enums;
 mod fromsql;
 mod overrides;
diff --git a/postgres-derive-internals/src/tosql.rs b/postgres-derive-internals/src/tosql.rs
@@ -1,8 +1,9 @@
 use std::fmt::Write;
-use syn::{Body, Ident, MacroInput, VariantData, Field};
+use syn::{self, Body, Ident, MacroInput, VariantData};
 use quote::{Tokens, ToTokens};
 
 use accepts;
+use composites::Field;
 use enums::Variant;
 use overrides::Overrides;
 
@@ -20,6 +21,11 @@ pub fn expand_derive_tosql(input: &MacroInput) -> Result<String, String> {
             let field = &fields[0];
             (domain_accepts_body(&name, &field), domain_body())
         }
+        Body::Struct(VariantData::Struct(ref fields)) => {
+            let fields: Vec<Field> = try!(fields.iter().map(Field::parse).collect());
+            (accepts::composite_body(&name, "ToSql", &fields),
+             composite_body(&fields))
+        }
         _ => {
             return Err("#[derive(ToSql)] may only be applied to structs, single field tuple \
                         structs, and enums".to_owned())
@@ -60,12 +66,12 @@ fn enum_body(ident: &Ident, variants: &[Variant]) -> String {
         };
 
         buf.extend_from_slice(s.as_bytes());
-        Ok(::postgres::types::IsNull::No)");
+        ::std::result::Result::Ok(::postgres::types::IsNull::No)");
 
     out
 }
 
-fn domain_accepts_body(name: &str, field: &Field) -> String {
+fn domain_accepts_body(name: &str, field: &syn::Field) -> String {
     let mut tokens = Tokens::new();
     field.ty.to_tokens(&mut tokens);
 
@@ -91,3 +97,56 @@ fn domain_body() -> String {
 
         ::postgres::types::ToSql::to_sql(&self.0, type_, buf, _info)".to_owned()
 }
+
+fn composite_body(fields: &[Field]) -> String {
+    let mut out = "
+        fn write_be_i32<W>(buf: &mut W, n: i32) -> ::std::io::Result<()>
+            where W: ::std::io::Write
+        {
+            let be = [(n >> 24) as u8, (n >> 16) as u8, (n >> 8) as u8, n as u8];
+            buf.write_all(&be)
+        }
+
+        let fields = match *_type.kind() {
+            ::postgres::types::Kind::Composite(ref fields) => fields,
+            _ => unreachable!(),
+        };
+
+        try!(write_be_i32(buf, fields.len() as i32));
+
+        for field in fields {
+            try!(write_be_i32(buf, field.type_().oid() as i32));
+
+            let base = buf.len();
+            try!(write_be_i32(buf, 0));
+            let r = match field.name() {".to_owned();
+
+    for field in fields {
+        write!(out, "
+                \"{}\" => ::postgres::types::ToSql::to_sql(&self.{}, field.type_(), buf, _info),",
+               field.name, field.ident).unwrap();
+    }
+
+    write!(out, "\
+                _ => unreachable!(),\
+            }};\
+            \
+            let count = match try!(r) {{\
+                ::postgres::types::IsNull::Yes => -1,\
+                ::postgres::types::IsNull::No => {{\
+                    let len = buf.len() - base - 4;\
+                    if len > i32::max_value() as usize {{\
+                        return ::std::result::Result::Err(\
+                            ::std::convert::Into::into(\"value too large to transmit\"));\
+                    }}\
+                    len as i32\
+                }}\
+            }};\
+            \
+            try!(write_be_i32(&mut &mut buf[base..base + 4], count));\
+        }}\
+        \
+        ::std::result::Result::Ok(::postgres::types::IsNull::No)").unwrap();
+
+    out
+}
diff --git a/postgres-derive/tests/composites.rs b/postgres-derive/tests/composites.rs
