Added pg_binary_protocol attribute to derive send and receive functions for PostgresType

README.md

@@ -53,6 +53,7 @@
 - **Easy Custom Types**
    + `#[derive(PostgresType)]` to use a Rust struct as a Postgres type
       - By default, represented as a CBOR-encoded object in-memory/on-disk, and JSON as human-readable
+      - Supports `#[pg_binary_protocol]` to generate binary protocol send/recv functions
       - Provide custom in-memory/on-disk/human-readable representations
    + `#[derive(PostgresEnum)]` to use a Rust enum as a Postgres enum
    + Composite types supported with the `pgrx::composite_type!("Sample")` macro

articles/forging-sql-from-rust.md

@@ -97,6 +97,7 @@ These `sql` files must be generated from data within the Rust code. The SQL gene
 - Create 'Shell types', in & out functions, and 'Base types' for each `#[derive(PostgresType)]` marked type.
   ```rust
   #[derive(PostgresType, Serialize, Deserialize, Debug, Eq, PartialEq)]
+  #[pg_binary_protocol]
   pub struct Animals {
       names: Vec<String>,
       age_lookup: HashMap<i32, String>,
@@ -140,6 +141,7 @@ These `sql` files must be generated from data within the Rust code. The SQL gene
       Eq, PartialEq, Ord, Hash, PartialOrd,
       PostgresType, Serialize, Deserialize
   )]
+  #[pg_binary_protocol]
   pub struct Thing(String);
   ```
   ```sql

pgrx-macros/src/lib.rs

@@ -698,7 +698,7 @@ fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
     stream.extend(quote! {
         impl ::pgrx::datum::FromDatum for #enum_ident {
             #[inline]
-            unsafe fn from_polymorphic_datum(datum: ::pgrx::pg_sys::Datum, is_null: bool, typeoid: ::pgrx::pg_sys::Oid) -> Option<#enum_ident> {
+            unsafe fn from_polymorphic_datum(datum: ::pgrx::pg_sys::Datum, is_null: bool, _typeoid: ::pgrx::pg_sys::Oid) -> Option<#enum_ident> {
                 if is_null {
                     None
                 } else {
@@ -781,6 +781,7 @@ 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.
+* `pg_binary_protocol`: Use the binary protocol for this type.
 * `pgrx(alignment = "<align>")`: Derive Postgres alignment from Rust type. One of `"on"`, or `"off"`.
 * `sql`: Same arguments as [`#[pgrx(sql = ..)]`](macro@pgrx).
 */
@@ -789,6 +790,7 @@ Optionally accepts the following attributes:
     attributes(
         inoutfuncs,
         pgvarlena_inoutfuncs,
+        pg_binary_protocol,
         bikeshed_postgres_type_manually_impl_from_into_datum,
         requires,
         pgrx
@@ -806,6 +808,9 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
     let has_lifetimes = generics.lifetimes().next();
     let funcname_in = Ident::new(&format!("{name}_in").to_lowercase(), name.span());
     let funcname_out = Ident::new(&format!("{name}_out").to_lowercase(), name.span());
+    let funcname_recv = Ident::new(&format!("{name}_recv").to_lowercase(), name.span());
+    let funcname_send = Ident::new(&format!("{name}_send").to_lowercase(), name.span());
+
     let mut args = parse_postgres_type_args(&ast.attrs);
     let mut stream = proc_macro2::TokenStream::new();
 
@@ -825,7 +830,9 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
         }
     }
 
-    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);
     }
@@ -937,14 +944,14 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
     if args.contains(&PostgresTypeAttribute::Default) {
         stream.extend(quote! {
             #[doc(hidden)]
-            #[::pgrx::pgrx_macros::pg_extern(immutable,parallel_safe)]
+            #[::pgrx::pgrx_macros::pg_extern(immutable, parallel_safe)]
             pub fn #funcname_in #generics(input: Option<&#lifetime ::core::ffi::CStr>) -> Option<#name #generics> {
                 use ::pgrx::inoutfuncs::json_from_slice;
                 input.map(|cstr| json_from_slice(cstr.to_bytes()).ok()).flatten()
             }
 
             #[doc(hidden)]
-            #[::pgrx::pgrx_macros::pg_extern (immutable,parallel_safe)]
+            #[::pgrx::pgrx_macros::pg_extern (immutable, parallel_safe)]
             pub fn #funcname_out #generics(input: #name #generics) -> ::pgrx::ffi::CString {
                 use ::pgrx::inoutfuncs::json_to_vec;
                 let mut bytes = json_to_vec(&input).unwrap();
@@ -1000,7 +1007,57 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
         });
     }
 
-    let sql_graph_entity_item = sql_gen::PostgresTypeDerive::from_derive_input(ast)?;
+    if args.contains(&PostgresTypeAttribute::PgBinaryProtocol) {
+        // At this time, the `PostgresTypeAttribute` does not impact the way we generate
+        // the `recv` and `send` functions.
+        stream.extend(quote! {
+            #[doc(hidden)]
+            #[::pgrx::pgrx_macros::pg_extern(immutable, strict, parallel_safe)]
+            pub fn #funcname_recv #generics(
+                internal: ::pgrx::datum::Internal,
+            ) -> #name #generics {
+                let buf = unsafe { internal.get_mut::<::pgrx::pg_sys::StringInfoData>().unwrap() };
+
+                let mut serialized = ::pgrx::StringInfo::new();
+
+                serialized.push_bytes(&[0u8; ::pgrx::pg_sys::VARHDRSZ]); // reserve space for the header
+                serialized.push_bytes(unsafe {
+                    core::slice::from_raw_parts(
+                        buf.data as *const u8,
+                        buf.len as usize
+                    )
+                });
+
+                let size = serialized.len();
+                let varlena = serialized.into_char_ptr();
+
+                unsafe{
+                    ::pgrx::set_varsize_4b(varlena as *mut ::pgrx::pg_sys::varlena, size as i32);
+                    buf.cursor = buf.len;
+                    ::pgrx::datum::cbor_decode(varlena as *mut ::pgrx::pg_sys::varlena)
+                }
+            }
+            #[doc(hidden)]
+            #[::pgrx::pgrx_macros::pg_extern(immutable, strict, parallel_safe)]
+            pub fn #funcname_send #generics(input: #name #generics) -> Vec<u8> {
+                use ::pgrx::datum::{FromDatum, IntoDatum};
+                let Some(datum): Option<::pgrx::pg_sys::Datum> = input.into_datum() else {
+                    ::pgrx::error!("Datum of type `{}` is unexpectedly NULL.", stringify!(#name));
+                };
+                unsafe {
+                    let Some(serialized): Option<Vec<u8>> = FromDatum::from_datum(datum, false) else {
+                        ::pgrx::error!("Failed to CBOR-serialize Datum to type `{}`.", stringify!(#name));
+                    };
+                    serialized
+                }
+            }
+        });
+    }
+
+    let sql_graph_entity_item = sql_gen::PostgresTypeDerive::from_derive_input(
+        ast,
+        args.contains(&PostgresTypeAttribute::PgBinaryProtocol),
+    )?;
     sql_graph_entity_item.to_tokens(&mut stream);
 
     Ok(stream)
@@ -1097,6 +1154,7 @@ fn impl_guc_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
 #[derive(Debug, Hash, Ord, PartialOrd, Eq, PartialEq)]
 enum PostgresTypeAttribute {
     InOutFuncs,
+    PgBinaryProtocol,
     PgVarlenaInOutFuncs,
     Default,
     ManualFromIntoDatum,
@@ -1112,6 +1170,9 @@ fn parse_postgres_type_args(attributes: &[Attribute]) -> HashSet<PostgresTypeAtt
             "inoutfuncs" => {
                 categorized_attributes.insert(PostgresTypeAttribute::InOutFuncs);
             }
+            "pg_binary_protocol" => {
+                categorized_attributes.insert(PostgresTypeAttribute::PgBinaryProtocol);
+            }
             "pgvarlena_inoutfuncs" => {
                 categorized_attributes.insert(PostgresTypeAttribute::PgVarlenaInOutFuncs);
             }

pgrx-sql-entity-graph/src/lib.rs

@@ -241,6 +241,10 @@ impl ToSql for SqlGraphEntity {
                             in_fn_module_path,
                             out_fn,
                             out_fn_module_path,
+                            receive_fn,
+                            receive_fn_module_path,
+                            send_fn,
+                            send_fn_module_path,
                             ..
                         }) = &context.graph[neighbor]
                         else {
@@ -251,7 +255,16 @@ impl ToSql for SqlGraphEntity {
                             && item.full_path.ends_with(in_fn);
                         let is_out_fn = item.full_path.starts_with(out_fn_module_path)
                             && item.full_path.ends_with(out_fn);
-                        is_in_fn || is_out_fn
+                        let is_receive_fn =
+                            receive_fn_module_path.as_ref().is_some_and(|receive_fn_module_path| {
+                                item.full_path.starts_with(receive_fn_module_path)
+                            }) && receive_fn
+                                .is_some_and(|receive_fn| item.full_path.ends_with(receive_fn));
+                        let is_send_fn =
+                            send_fn_module_path.as_ref().is_some_and(|send_fn_module_path| {
+                                item.full_path.starts_with(send_fn_module_path)
+                            }) && send_fn.is_some_and(|send_fn| item.full_path.ends_with(send_fn));
+                        is_in_fn || is_out_fn || is_receive_fn || is_send_fn
                     })
                 {
                     Ok(String::default())

pgrx-sql-entity-graph/src/postgres_type/entity.rs

@@ -105,6 +105,10 @@ pub struct PostgresTypeEntity {
     pub in_fn_module_path: String,
     pub out_fn: &'static str,
     pub out_fn_module_path: String,
+    pub receive_fn: Option<&'static str>,
+    pub receive_fn_module_path: Option<String>,
+    pub send_fn: Option<&'static str>,
+    pub send_fn_module_path: Option<String>,
     pub to_sql_config: ToSqlConfigEntity,
     pub alignment: Option<usize>,
 }
@@ -152,6 +156,10 @@ impl ToSql for PostgresTypeEntity {
             out_fn,
             out_fn_module_path,
             in_fn,
+            receive_fn,
+            receive_fn_module_path,
+            send_fn,
+            send_fn_module_path,
             alignment,
             ..
         }) = item_node
@@ -217,6 +225,80 @@ impl ToSql for PostgresTypeEntity {
             .ok_or_else(|| eyre!("Could not find out_fn graph entity."))?;
         let out_fn_sql = out_fn_entity.to_sql(context)?;
 
+        let receive_fn_graph_index_and_receive_fn_sql = receive_fn_module_path
+            .as_ref()
+            .zip(*receive_fn)
+            .map(|(receive_fn_module_path, receive_fn)| {
+                let receive_fn_module_path = if !receive_fn_module_path.is_empty() {
+                    receive_fn_module_path.clone()
+                } else {
+                    module_path.to_string() // Presume a local
+                };
+                let receive_fn_path = format!(
+                    "{receive_fn_module_path}{maybe_colons}{receive_fn}",
+                    maybe_colons = if !receive_fn_module_path.is_empty() { "::" } else { "" }
+                );
+
+                // Find the receive function in the context
+                let (_, _index) = context
+                    .externs
+                    .iter()
+                    .find(|(k, _v)| k.full_path == receive_fn_path)
+                    .ok_or_else(|| eyre::eyre!("Did not find `receive_fn`: {receive_fn_path}."))?;
+
+                let (receive_fn_graph_index, receive_fn_entity) = context
+                    .graph
+                    .neighbors_undirected(self_index)
+                    .find_map(|neighbor| match &context.graph[neighbor] {
+                        SqlGraphEntity::Function(func) if func.full_path == receive_fn_path => {
+                            Some((neighbor, func))
+                        }
+                        _ => None,
+                    })
+                    .ok_or_else(|| eyre!("Could not find receive_fn graph entity."))?;
+                let receive_fn_sql = receive_fn_entity.to_sql(context)?;
+
+                Ok::<_, eyre::Report>((receive_fn_graph_index, receive_fn_sql, receive_fn_path))
+            })
+            .transpose()?;
+
+        let send_fn_graph_index_and_send_fn_sql = send_fn_module_path
+            .as_ref()
+            .zip(*send_fn)
+            .map(|(send_fn_module_path, send_fn)| {
+                let send_fn_module_path = if !send_fn_module_path.is_empty() {
+                    send_fn_module_path.clone()
+                } else {
+                    module_path.to_string() // Presume a local
+                };
+                let send_fn_path = format!(
+                    "{send_fn_module_path}{maybe_colons}{send_fn}",
+                    maybe_colons = if !send_fn_module_path.is_empty() { "::" } else { "" }
+                );
+
+                // Find the send function in the context
+                let (_, _index) = context
+                    .externs
+                    .iter()
+                    .find(|(k, _v)| k.full_path == send_fn_path)
+                    .ok_or_else(|| eyre::eyre!("Did not find `send_fn: {}`.", send_fn_path))?;
+
+                let (send_fn_graph_index, send_fn_entity) = context
+                    .graph
+                    .neighbors_undirected(self_index)
+                    .find_map(|neighbor| match &context.graph[neighbor] {
+                        SqlGraphEntity::Function(func) if func.full_path == send_fn_path => {
+                            Some((neighbor, func))
+                        }
+                        _ => None,
+                    })
+                    .ok_or_else(|| eyre!("Could not find send_fn graph entity."))?;
+                let send_fn_sql = send_fn_entity.to_sql(context)?;
+
+                Ok::<_, eyre::Report>((send_fn_graph_index, send_fn_sql, send_fn_path))
+            })
+            .transpose()?;
+
         let shell_type = format!(
             "\n\
                 -- {file}:{line}\n\
@@ -244,6 +326,29 @@ impl ToSql for PostgresTypeEntity {
             })
             .unwrap_or_default();
 
+        let (receive_send_attributes, receive_send_sql) = receive_fn_graph_index_and_receive_fn_sql
+            .zip(send_fn_graph_index_and_send_fn_sql)
+            .map(|((receive_fn_graph_index, receive_fn_sql, receive_fn_path), (send_fn_graph_index, send_fn_sql, send_fn_path))| {
+                let receive_fn = receive_fn.unwrap();
+                let send_fn = send_fn.unwrap();
+                (
+                    format! {
+                        "\
+                        \tRECEIVE = {schema_prefix_receive_fn}{receive_fn}, /* {receive_fn_path} */\n\
+                        \tSEND = {schema_prefix_send_fn}{send_fn}, /* {send_fn_path} */\n\
+                        ",
+                        schema_prefix_receive_fn = context.schema_prefix_for(&receive_fn_graph_index),
+                        schema_prefix_send_fn = context.schema_prefix_for(&send_fn_graph_index),
+                    },
+                    format! {
+                        "\n\
+                        {receive_fn_sql}\n\
+                        {send_fn_sql}\n\
+                        "
+                    }
+                )
+            }).unwrap_or_default();
+
         let materialized_type = format! {
             "\n\
                 -- {file}:{line}\n\
@@ -252,14 +357,24 @@ impl ToSql for PostgresTypeEntity {
                     \tINTERNALLENGTH = variable,\n\
                     \tINPUT = {schema_prefix_in_fn}{in_fn}, /* {in_fn_path} */\n\
                     \tOUTPUT = {schema_prefix_out_fn}{out_fn}, /* {out_fn_path} */\n\
+                    {receive_send_attributes}\
                     \tSTORAGE = extended{alignment}\n\
                 );\
             ",
             schema = context.schema_prefix_for(&self_index),
             schema_prefix_in_fn = context.schema_prefix_for(&in_fn_graph_index),
-            schema_prefix_out_fn = context.schema_prefix_for(&out_fn_graph_index),
+            schema_prefix_out_fn = context.schema_prefix_for(&out_fn_graph_index)
         };
 
-        Ok(shell_type + "\n" + &in_fn_sql + "\n" + &out_fn_sql + "\n" + &materialized_type)
+        let result = shell_type
+            + "\n"
+            + &in_fn_sql
+            + "\n"
+            + &out_fn_sql
+            + &receive_send_sql
+            + "\n"
+            + &materialized_type;
+
+        Ok(result)
     }
 }