Skip to content

Commit

Permalink
Problem: text type representation is not always efficient
Browse files Browse the repository at this point in the history
For this reason, Postgres allows types to have an external binary
representation. Also, some clients insist on using binary
representation.

Solution: introduce SendRecvFuncs trait and `sendrecvfuncs` attribute

These are used to specify how external binary representation encoding is
accomplished.
  • Loading branch information
yrashk committed Nov 24, 2022
1 parent c0ce2a8 commit 0564fc4
Show file tree
Hide file tree
Showing 13 changed files with 426 additions and 29 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

23 changes: 23 additions & 0 deletions pgx-examples/custom_types/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
43 changes: 40 additions & 3 deletions pgx-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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();

Expand All @@ -710,7 +716,7 @@ 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);
}
Expand Down Expand Up @@ -803,7 +809,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
Expand Down Expand Up @@ -895,6 +928,7 @@ fn impl_guc_enum(ast: DeriveInput) -> proc_macro2::TokenStream {
enum PostgresTypeAttribute {
InOutFuncs,
PgVarlenaInOutFuncs,
SendReceiveFuncs,
Default,
}

Expand All @@ -912,6 +946,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
Expand Down
3 changes: 2 additions & 1 deletion pgx-tests/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ pg12 = [ "pgx/pg12" ]
pg13 = [ "pgx/pg13" ]
pg14 = [ "pgx/pg14" ]
pg15 = [ "pgx/pg15" ]
pg_test = [ ]
pg_test = [ "bytes" ]

[package.metadata.docs.rs]
features = ["pg14"]
Expand All @@ -44,6 +44,7 @@ serde_json = "1.0.88"
time = "0.3.17"
eyre = "0.6.8"
thiserror = "1.0"
bytes = { version = "1.2.1", optional = true }

[dependencies.pgx]
path = "../pgx"
Expand Down
1 change: 1 addition & 0 deletions pgx-tests/src/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
61 changes: 56 additions & 5 deletions pgx-tests/src/tests/postgres_type_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -152,16 +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 {
use std::error::Error;
use postgres::types::{FromSql, IsNull, ToSql, Type};
use postgres::types::private::BytesMut;
#[allow(unused_imports)]
use crate as pgx_tests;

use crate::tests::postgres_type_tests::{
CustomTextFormatSerializedEnumType, CustomTextFormatSerializedType, JsonEnumType, JsonType,
VarlenaEnumType, VarlenaType,
};
use crate::tests::postgres_type_tests::{BinaryEncodedType, CustomTextFormatSerializedEnumType,
CustomTextFormatSerializedType, JsonEnumType, JsonType,
VarlenaEnumType, VarlenaType};
use pgx::prelude::*;
use pgx::PgVarlena;

Expand Down Expand Up @@ -246,4 +262,39 @@ 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);
}
}
41 changes: 41 additions & 0 deletions pgx-tests/src/tests/stringinfo_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
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][..]));
}

}
13 changes: 11 additions & 2 deletions pgx-utils/src/sql_entity_graph/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand All @@ -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,
}
Expand Down
Loading

0 comments on commit 0564fc4

Please sign in to comment.