Skip to content

Commit 0b5a248

Browse files
committed
Problem: text type representation is not always efficient
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.
1 parent c0ce2a8 commit 0b5a248

File tree

13 files changed

+449
-27
lines changed

13 files changed

+449
-27
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pgx-examples/custom_types/README.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,29 @@ fn do_a_thing(mut input: PgVarlena<MyType>) -> PgVarlena<MyType> {
110110
}
111111
```
112112

113+
## External Binary Representation
114+
115+
PostgreSQL allows types to have an external binary representation for more efficient communication with
116+
clients (as a matter of fact, Rust's [https://crates.io/crates/postgres](postgres) crate uses binary types
117+
exclusively). By default, `PostgresType` do not have any external binary representation, however, this can
118+
be done by specifying `#[sendrecvfuncs]` attribute on the type and implementing `SendRecvFuncs` trait:
119+
120+
```rust
121+
#[derive(PostgresType, Serialize, Deserialize, Debug, PartialEq)]
122+
#[sendrecvfuncs]
123+
pub struct BinaryEncodedType(Vec<u8>);
124+
125+
impl SendRecvFuncs for BinaryEncodedType {
126+
fn send(&self) -> Vec<u8> {
127+
self.0.clone()
128+
}
129+
130+
fn recv(buffer: &[u8]) -> Self {
131+
Self(buffer.to_vec())
132+
}
133+
}
134+
```
135+
113136
## Notes
114137

115138
- For serde-compatible types, you can use the `#[inoutfuncs]` annotation (instead of `#[pgvarlena_inoutfuncs]`) if you'd

pgx-macros/src/lib.rs

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -681,9 +681,13 @@ Optionally accepts the following attributes:
681681
682682
* `inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the type.
683683
* `pgvarlena_inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the `PgVarlena` of this type.
684+
* `sendrecvfuncs`: Define binary send/receive functions for the type.
684685
* `sql`: Same arguments as [`#[pgx(sql = ..)]`](macro@pgx).
685686
*/
686-
#[proc_macro_derive(PostgresType, attributes(inoutfuncs, pgvarlena_inoutfuncs, requires, pgx))]
687+
#[proc_macro_derive(
688+
PostgresType,
689+
attributes(inoutfuncs, pgvarlena_inoutfuncs, sendrecvfuncs, requires, pgx)
690+
)]
687691
pub fn postgres_type(input: TokenStream) -> TokenStream {
688692
let ast = parse_macro_input!(input as syn::DeriveInput);
689693

@@ -696,6 +700,8 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
696700
let has_lifetimes = generics.lifetimes().next();
697701
let funcname_in = Ident::new(&format!("{}_in", name).to_lowercase(), name.span());
698702
let funcname_out = Ident::new(&format!("{}_out", name).to_lowercase(), name.span());
703+
let funcname_send = Ident::new(&format!("{}_send", name).to_lowercase(), name.span());
704+
let funcname_recv = Ident::new(&format!("{}_recv", name).to_lowercase(), name.span());
699705
let mut args = parse_postgres_type_args(&ast.attrs);
700706
let mut stream = proc_macro2::TokenStream::new();
701707

@@ -710,7 +716,9 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
710716
_ => panic!("#[derive(PostgresType)] can only be applied to structs or enums"),
711717
}
712718

713-
if args.is_empty() {
719+
if !args.contains(&PostgresTypeAttribute::InOutFuncs)
720+
&& !args.contains(&PostgresTypeAttribute::PgVarlenaInOutFuncs)
721+
{
714722
// assume the user wants us to implement the InOutFuncs
715723
args.insert(PostgresTypeAttribute::Default);
716724
}
@@ -803,7 +811,34 @@ fn impl_postgres_type(ast: DeriveInput) -> proc_macro2::TokenStream {
803811
});
804812
}
805813

806-
let sql_graph_entity_item = PostgresType::from_derive_input(ast).unwrap();
814+
if args.contains(&PostgresTypeAttribute::SendReceiveFuncs) {
815+
stream.extend(quote! {
816+
#[doc(hidden)]
817+
#[pg_extern(immutable,parallel_safe,strict)]
818+
pub fn #funcname_recv #generics(input: ::pgx::Internal) -> #name #generics {
819+
let mut buffer0 = unsafe {
820+
input
821+
.get_mut::<::pgx::pg_sys::StringInfoData>()
822+
.expect("Can't retrieve StringInfo pointer")
823+
};
824+
let mut buffer = StringInfo::from_pg(buffer0 as *mut _).expect("failed to construct StringInfo");
825+
let slice = buffer.read(..).expect("failure reading StringInfo");
826+
::pgx::SendRecvFuncs::recv(slice)
827+
}
828+
829+
#[doc(hidden)]
830+
#[pg_extern(immutable,parallel_safe,strict)]
831+
pub fn #funcname_send #generics(input: #name #generics) -> Vec<u8> {
832+
::pgx::SendRecvFuncs::send(&input)
833+
}
834+
});
835+
}
836+
837+
let sql_graph_entity_item = PostgresType::from_derive_input(
838+
ast,
839+
args.contains(&PostgresTypeAttribute::SendReceiveFuncs),
840+
)
841+
.unwrap();
807842
sql_graph_entity_item.to_tokens(&mut stream);
808843

809844
stream
@@ -895,6 +930,7 @@ fn impl_guc_enum(ast: DeriveInput) -> proc_macro2::TokenStream {
895930
enum PostgresTypeAttribute {
896931
InOutFuncs,
897932
PgVarlenaInOutFuncs,
933+
SendReceiveFuncs,
898934
Default,
899935
}
900936

@@ -912,6 +948,9 @@ fn parse_postgres_type_args(attributes: &[Attribute]) -> HashSet<PostgresTypeAtt
912948
"pgvarlena_inoutfuncs" => {
913949
categorized_attributes.insert(PostgresTypeAttribute::PgVarlenaInOutFuncs);
914950
}
951+
"sendrecvfuncs" => {
952+
categorized_attributes.insert(PostgresTypeAttribute::SendReceiveFuncs);
953+
}
915954

916955
_ => {
917956
// we can just ignore attributes we don't understand

pgx-tests/Cargo.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ pg12 = [ "pgx/pg12" ]
2020
pg13 = [ "pgx/pg13" ]
2121
pg14 = [ "pgx/pg14" ]
2222
pg15 = [ "pgx/pg15" ]
23-
pg_test = [ ]
23+
pg_test = [ "bytes" ]
2424

2525
[package.metadata.docs.rs]
2626
features = ["pg14"]
@@ -44,6 +44,7 @@ serde_json = "1.0.88"
4444
time = "0.3.17"
4545
eyre = "0.6.8"
4646
thiserror = "1.0"
47+
bytes = { version = "1.2.1", optional = true }
4748

4849
[dependencies.pgx]
4950
path = "../pgx"

pgx-tests/src/tests/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ mod schema_tests;
3939
mod shmem_tests;
4040
mod spi_tests;
4141
mod srf_tests;
42+
mod stringinfo_tests;
4243
mod struct_type_tests;
4344
mod trigger_tests;
4445
mod uuid_tests;

pgx-tests/src/tests/postgres_type_tests.rs

Lines changed: 64 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ Use of this source code is governed by the MIT license that can be found in the
88
*/
99
use pgx::cstr_core::CStr;
1010
use pgx::prelude::*;
11-
use pgx::{InOutFuncs, PgVarlena, PgVarlenaInOutFuncs, StringInfo};
11+
use pgx::{InOutFuncs, PgVarlena, PgVarlenaInOutFuncs, SendRecvFuncs, StringInfo};
1212
use serde::{Deserialize, Serialize};
1313
use std::str::FromStr;
1414

@@ -152,15 +152,32 @@ pub enum JsonEnumType {
152152
E2 { b: f32 },
153153
}
154154

155+
#[derive(PostgresType, Serialize, Deserialize, Debug, PartialEq)]
156+
#[sendrecvfuncs]
157+
pub struct BinaryEncodedType(Vec<u8>);
158+
159+
impl SendRecvFuncs for BinaryEncodedType {
160+
fn send(&self) -> Vec<u8> {
161+
self.0.clone()
162+
}
163+
164+
fn recv(buffer: &[u8]) -> Self {
165+
Self(buffer.to_vec())
166+
}
167+
}
168+
155169
#[cfg(any(test, feature = "pg_test"))]
156170
#[pgx::pg_schema]
157171
mod tests {
158172
#[allow(unused_imports)]
159173
use crate as pgx_tests;
174+
use postgres::types::private::BytesMut;
175+
use postgres::types::{FromSql, IsNull, ToSql, Type};
176+
use std::error::Error;
160177

161178
use crate::tests::postgres_type_tests::{
162-
CustomTextFormatSerializedEnumType, CustomTextFormatSerializedType, JsonEnumType, JsonType,
163-
VarlenaEnumType, VarlenaType,
179+
BinaryEncodedType, CustomTextFormatSerializedEnumType, CustomTextFormatSerializedType,
180+
JsonEnumType, JsonType, VarlenaEnumType, VarlenaType,
164181
};
165182
use pgx::prelude::*;
166183
use pgx::PgVarlena;
@@ -246,4 +263,48 @@ mod tests {
246263
.expect("SPI returned NULL");
247264
assert!(matches!(result, JsonEnumType::E1 { a } if a == 1.0));
248265
}
266+
267+
#[pg_test]
268+
fn test_binary_encoded_type() {
269+
impl ToSql for BinaryEncodedType {
270+
fn to_sql(
271+
&self,
272+
_ty: &Type,
273+
out: &mut BytesMut,
274+
) -> Result<IsNull, Box<dyn Error + Sync + Send>>
275+
where
276+
Self: Sized,
277+
{
278+
use bytes::BufMut;
279+
out.put_slice(self.0.as_slice());
280+
Ok(IsNull::No)
281+
}
282+
283+
fn accepts(_ty: &Type) -> bool
284+
where
285+
Self: Sized,
286+
{
287+
true
288+
}
289+
290+
postgres::types::to_sql_checked!();
291+
}
292+
293+
impl<'a> FromSql<'a> for BinaryEncodedType {
294+
fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
295+
Ok(Self(raw.to_vec()))
296+
}
297+
298+
fn accepts(_ty: &Type) -> bool {
299+
true
300+
}
301+
}
302+
303+
// postgres client uses binary types so we can use it to test this functionality
304+
let (mut client, _) = pgx_tests::client().unwrap();
305+
let val = BinaryEncodedType(vec![0, 1, 2]);
306+
let result = client.query("SELECT $1::BinaryEncodedType", &[&val]).unwrap();
307+
let val1: BinaryEncodedType = result[0].get(0);
308+
assert_eq!(val, val1);
309+
}
249310
}
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
Portions Copyright 2019-2021 ZomboDB, LLC.
3+
Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <[email protected]>
4+
5+
All rights reserved.
6+
7+
Use of this source code is governed by the MIT license that can be found in the LICENSE file.
8+
*/
9+
10+
#[cfg(any(test, feature = "pg_test"))]
11+
#[pgx::pg_schema]
12+
mod tests {
13+
#[allow(unused_imports)]
14+
use crate as pgx_tests;
15+
16+
use pgx::*;
17+
18+
#[pg_test]
19+
fn test_string_info_read_full() {
20+
let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
21+
assert_eq!(string_info.read(..), Some(&[1, 2, 3, 4, 5][..]));
22+
assert_eq!(string_info.read(..), Some(&[][..]));
23+
assert_eq!(string_info.read(..=1), None);
24+
}
25+
26+
#[pg_test]
27+
fn test_string_info_read_offset() {
28+
let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
29+
assert_eq!(string_info.read(1..), Some(&[2, 3, 4, 5][..]));
30+
assert_eq!(string_info.read(..), Some(&[][..]));
31+
}
32+
33+
#[pg_test]
34+
fn test_string_info_read_cap() {
35+
let mut string_info = StringInfo::from(vec![1, 2, 3, 4, 5]);
36+
assert_eq!(string_info.read(..=1), Some(&[1][..]));
37+
assert_eq!(string_info.read(1..=2), Some(&[3][..]));
38+
assert_eq!(string_info.read(..), Some(&[4, 5][..]));
39+
}
40+
}

pgx-utils/src/sql_entity_graph/mod.rs

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,8 @@ impl ToSql for SqlGraphEntity {
205205
if context.graph.neighbors_undirected(context.externs.get(item).unwrap().clone()).any(|neighbor| {
206206
let neighbor_item = &context.graph[neighbor];
207207
match neighbor_item {
208-
SqlGraphEntity::Type(PostgresTypeEntity { in_fn, in_fn_module_path, out_fn, out_fn_module_path, .. }) => {
208+
SqlGraphEntity::Type(PostgresTypeEntity { in_fn, in_fn_module_path, out_fn, out_fn_module_path, send_fn, recv_fn,
209+
send_fn_module_path, recv_fn_module_path, .. }) => {
209210
let is_in_fn = item.full_path.starts_with(in_fn_module_path) && item.full_path.ends_with(in_fn);
210211
if is_in_fn {
211212
tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an in_fn.");
@@ -214,7 +215,15 @@ impl ToSql for SqlGraphEntity {
214215
if is_out_fn {
215216
tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an out_fn.");
216217
}
217-
is_in_fn || is_out_fn
218+
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());
219+
if is_send_fn {
220+
tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an send_fn.");
221+
}
222+
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());
223+
if is_recv_fn {
224+
tracing::trace!(r#type = %neighbor_item.dot_identifier(), "Skipping, is an recv_fn.");
225+
}
226+
is_in_fn || is_out_fn || is_send_fn || is_recv_fn
218227
},
219228
_ => false,
220229
}

0 commit comments

Comments
 (0)