External binary representation support (SEND/RECEIVE for custom types)

pgx-examples/custom_types/README.md

@@ -110,6 +110,29 @@ fn do_a_thing(mut input: PgVarlena<MyType>) -> PgVarlena<MyType> {
 }
 ```
 
+## External Binary Representation
+
+PostgreSQL allows types to have an external binary representation for more efficient communication with
+clients (as a matter of fact, Rust's [https://crates.io/crates/postgres](postgres) crate uses binary types
+exclusively). By default, `PostgresType` do not have any external binary representation, however, this can
+be done by specifying `#[sendrecvfuncs]` attribute on the type and implementing `SendRecvFuncs` trait:
+
+```rust
+#[derive(PostgresType, Serialize, Deserialize, Debug, PartialEq)]
+#[sendrecvfuncs]
+pub struct BinaryEncodedType(Vec<u8>);
+
+impl SendRecvFuncs for BinaryEncodedType {
+    fn send(&self) -> Vec<u8> {
+        self.0.clone()
+    }
+
+    fn recv(buffer: &[u8]) -> Self {
+        Self(buffer.to_vec())
+    }
+}
+```
+
 ## Notes
 
 - For serde-compatible types, you can use the `#[inoutfuncs]` annotation (instead of `#[pgvarlena_inoutfuncs]`) if you'd 

pgx-macros/src/lib.rs

@@ -681,9 +681,13 @@ Optionally accepts the following attributes:
 
 * `inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the type.
 * `pgvarlena_inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the `PgVarlena` of this type.
+* `sendrecvfuncs`: Define binary send/receive functions for the type.
 * `sql`: Same arguments as [`#[pgx(sql = ..)]`](macro@pgx).
 */
-#[proc_macro_derive(PostgresType, attributes(inoutfuncs, pgvarlena_inoutfuncs, requires, pgx))]
+#[proc_macro_derive(
+    PostgresType,
+    attributes(inoutfuncs, pgvarlena_inoutfuncs, sendrecvfuncs, requires, pgx)
+)]
 pub fn postgres_type(input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as syn::DeriveInput);
 
@@ -696,6 +700,8 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
     let has_lifetimes = generics.lifetimes().next();
     let funcname_in = Ident::new(&format!("{}_in", name).to_lowercase(), name.span());
     let funcname_out = Ident::new(&format!("{}_out", name).to_lowercase(), name.span());
+    let funcname_send = Ident::new(&format!("{}_send", name).to_lowercase(), name.span());
+    let funcname_recv = Ident::new(&format!("{}_recv", name).to_lowercase(), name.span());
     let mut args = parse_postgres_type_args(&ast.attrs);
     let mut stream = proc_macro2::TokenStream::new();
 
@@ -710,7 +716,9 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
         _ => panic!("#[derive(PostgresType)] can only be applied to structs or enums"),
     }
 
-    if args.is_empty() {
+    if !args.contains(&PostgresTypeAttribute::InOutFuncs)
+        && !args.contains(&PostgresTypeAttribute::PgVarlenaInOutFuncs)
+    {
         // assume the user wants us to implement the InOutFuncs
         args.insert(PostgresTypeAttribute::Default);
     }
@@ -803,7 +811,34 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
         });
     }
 
-    let sql_graph_entity_item = PostgresType::from_derive_input(ast).unwrap();
+    if args.contains(&PostgresTypeAttribute::SendReceiveFuncs) {
+        stream.extend(quote! {
+            #[doc(hidden)]
+            #[pg_extern(immutable,parallel_safe,strict)]
+            pub fn #funcname_recv #generics(input: ::pgx::Internal) -> #name #generics {
+                let mut buffer0 = unsafe {
+                    input
+                    .get_mut::<::pgx::pg_sys::StringInfoData>()
+                    .expect("Can't retrieve StringInfo pointer")
+                };
+                let mut buffer = StringInfo::from_pg(buffer0 as *mut _).expect("failed to construct StringInfo");
+                let slice = buffer.read(..).expect("failure reading StringInfo");
+                ::pgx::SendRecvFuncs::recv(slice)
+            }
+
+            #[doc(hidden)]
+            #[pg_extern(immutable,parallel_safe,strict)]
+            pub fn #funcname_send #generics(input: #name #generics) -> Vec<u8> {
+                ::pgx::SendRecvFuncs::send(&input)
+            }
+        });
+    }
+
+    let sql_graph_entity_item = PostgresType::from_derive_input(
+        ast,
+        args.contains(&PostgresTypeAttribute::SendReceiveFuncs),
+    )
+    .unwrap();
     sql_graph_entity_item.to_tokens(&mut stream);
 
     stream
@@ -895,6 +930,7 @@ fn impl_guc_enum(ast: DeriveInput) -> proc_macro2::TokenStream {
 enum PostgresTypeAttribute {
     InOutFuncs,
     PgVarlenaInOutFuncs,
+    SendReceiveFuncs,
     Default,
 }
 
@@ -912,6 +948,9 @@ fn parse_postgres_type_args(attributes: &[Attribute]) -> HashSet<PostgresTypeAtt
             "pgvarlena_inoutfuncs" => {
                 categorized_attributes.insert(PostgresTypeAttribute::PgVarlenaInOutFuncs);
             }
+            "sendrecvfuncs" => {
+                categorized_attributes.insert(PostgresTypeAttribute::SendReceiveFuncs);
+            }
 
             _ => {
                 // we can just ignore attributes we don't understand

pgx-tests/Cargo.toml

@@ -44,6 +44,7 @@ serde_json = "1.0.88"
 time = "0.3.17"
 eyre = "0.6.8"
 thiserror = "1.0"
+bytes = "1.2.1"
 
 [dependencies.pgx]
 path = "../pgx"

pgx-tests/src/tests/mod.rs

@@ -39,6 +39,7 @@ mod schema_tests;
 mod shmem_tests;
 mod spi_tests;
 mod srf_tests;
+mod stringinfo_tests;
 mod struct_type_tests;
 mod trigger_tests;
 mod uuid_tests;

pgx-tests/src/tests/postgres_type_tests.rs

@@ -8,7 +8,7 @@ Use of this source code is governed by the MIT license that can be found in the
 */
 use pgx::cstr_core::CStr;
 use pgx::prelude::*;
-use pgx::{InOutFuncs, PgVarlena, PgVarlenaInOutFuncs, StringInfo};
+use pgx::{InOutFuncs, PgVarlena, PgVarlenaInOutFuncs, SendRecvFuncs, StringInfo};
 use serde::{Deserialize, Serialize};
 use std::str::FromStr;
 
@@ -152,15 +152,32 @@ pub enum JsonEnumType {
     E2 { b: f32 },
 }
 
+#[derive(PostgresType, Serialize, Deserialize, Debug, PartialEq)]
+#[sendrecvfuncs]
+pub struct BinaryEncodedType(Vec<u8>);
+
+impl SendRecvFuncs for BinaryEncodedType {
+    fn send(&self) -> Vec<u8> {
+        self.0.clone()
+    }
+
+    fn recv(buffer: &[u8]) -> Self {
+        Self(buffer.to_vec())
+    }
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgx::pg_schema]
 mod tests {
     #[allow(unused_imports)]
     use crate as pgx_tests;
+    use postgres::types::private::BytesMut;
+    use postgres::types::{FromSql, IsNull, ToSql, Type};
+    use std::error::Error;
 
     use crate::tests::postgres_type_tests::{
-        CustomTextFormatSerializedEnumType, CustomTextFormatSerializedType, JsonEnumType, JsonType,
-        VarlenaEnumType, VarlenaType,
+        BinaryEncodedType, CustomTextFormatSerializedEnumType, CustomTextFormatSerializedType,
+        JsonEnumType, JsonType, VarlenaEnumType, VarlenaType,
     };
     use pgx::prelude::*;
     use pgx::PgVarlena;
@@ -246,4 +263,48 @@ mod tests {
                 .expect("SPI returned NULL");
         assert!(matches!(result, JsonEnumType::E1 { a } if a == 1.0));
     }
+
+    #[pg_test]
+    fn test_binary_encoded_type() {
+        impl ToSql for BinaryEncodedType {
+            fn to_sql(
+                &self,
+                _ty: &Type,
+                out: &mut BytesMut,
+            ) -> Result<IsNull, Box<dyn Error + Sync + Send>>
+            where
+                Self: Sized,
+            {
+                use bytes::BufMut;
+                out.put_slice(self.0.as_slice());
+                Ok(IsNull::No)
+            }
+
+            fn accepts(_ty: &Type) -> bool
+            where
+                Self: Sized,
+            {
+                true
+            }
+
+            postgres::types::to_sql_checked!();
+        }
+
+        impl<'a> FromSql<'a> for BinaryEncodedType {
+            fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
+                Ok(Self(raw.to_vec()))
+            }
+
+            fn accepts(_ty: &Type) -> bool {
+                true
+            }
+        }
+
+        // postgres client uses binary types so we can use it to test this functionality
+        let (mut client, _) = pgx_tests::client().unwrap();
+        let val = BinaryEncodedType(vec![0, 1, 2]);
+        let result = client.query("SELECT $1::BinaryEncodedType", &[&val]).unwrap();
+        let val1: BinaryEncodedType = result[0].get(0);
+        assert_eq!(val, val1);
+    }
 }

pgx-tests/src/tests/stringinfo_tests.rs

@@ -0,0 +1,40 @@
+/*
+Portions Copyright 2019-2021 ZomboDB, LLC.
+Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <[email protected]>
+
+All rights reserved.
+
+Use of this source code is governed by the MIT license that can be found in the LICENSE file.
+*/
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pgx::pg_schema]
+mod tests {
+    #[allow(unused_imports)]
+    use crate as pgx_tests;
+
+    use pgx::*;
+
+    #[pg_test]
+    fn test_string_info_read_full() {
+        let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
+        assert_eq!(string_info.read(..), Some(&[1, 2, 3, 4, 5][..]));
+        assert_eq!(string_info.read(..), Some(&[][..]));
+        assert_eq!(string_info.read(..=1), None);
+    }
+
+    #[pg_test]
+    fn test_string_info_read_offset() {
+        let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
+        assert_eq!(string_info.read(1..), Some(&[2, 3, 4, 5][..]));
+        assert_eq!(string_info.read(..), Some(&[][..]));
+    }
+
+    #[pg_test]
+    fn test_string_info_read_cap() {
+        let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
+        assert_eq!(string_info.read(..=1), Some(&[1][..]));
+        assert_eq!(string_info.read(1..=2), Some(&[3][..]));
+        assert_eq!(string_info.read(..), Some(&[4, 5][..]));
+    }
+}

pgx-utils/src/sql_entity_graph/mod.rs

@@ -205,7 +205,8 @@ impl ToSql for SqlGraphEntity {
                 if context.graph.neighbors_undirected(context.externs.get(item).unwrap().clone()).any(|neighbor| {
                     let neighbor_item = &context.graph[neighbor];
                     match neighbor_item {
-                        SqlGraphEntity::Type(PostgresTypeEntity { in_fn, in_fn_module_path, out_fn, out_fn_module_path, .. }) => {
+                        SqlGraphEntity::Type(PostgresTypeEntity { in_fn, in_fn_module_path, out_fn, out_fn_module_path, send_fn, recv_fn,
+                                                 send_fn_module_path, recv_fn_module_path, .. }) => {
                             let is_in_fn = item.full_path.starts_with(in_fn_module_path) && item.full_path.ends_with(in_fn);
                             if is_in_fn {
                                 tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an in_fn.");
@@ -214,7 +215,15 @@ impl ToSql for SqlGraphEntity {
                             if is_out_fn {
                                 tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an out_fn.");
                             }
-                            is_in_fn || is_out_fn
+                            let is_send_fn = send_fn.is_some() && item.full_path.starts_with(send_fn_module_path) && item.full_path.ends_with(send_fn.unwrap_or_default());
+                            if is_send_fn {
+                                tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an send_fn.");
+                            }
+                            let is_recv_fn = recv_fn.is_some() && item.full_path.starts_with(recv_fn_module_path) && item.full_path.ends_with(recv_fn.unwrap_or_default());
+                            if is_recv_fn {
+                                tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an recv_fn.");
+                            }
+                            is_in_fn || is_out_fn || is_send_fn || is_recv_fn
                         },
                         _ => false,
                     }