Fix a problem with fns returning tuples with Attributes (#440)

Signed-off-by: Ana Hobden <[email protected]>
pgcentralfoundation · Feb 16, 2022 · d55c5d1 · d55c5d1
1 parent ed9f061
commit d55c5d1
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 6 deletions.
diff --git a/pgx-tests/src/tests/pg_extern_args_tests.rs b/pgx-tests/src/tests/pg_extern_args_tests.rs
@@ -1,6 +1,16 @@
 // Copyright 2020 ZomboDB, LLC <[email protected]>. All rights reserved. Use of this source code is
 // governed by the MIT license that can be found in the LICENSE file.
 
+use pgx::*;
+
+#[pg_extern(immutable)]
+fn returns_tuple_with_attributes() -> (
+    name!(arg, String),
+    name!(arg2, String),
+) {
+   ("hi".to_string(), "bye".to_string())
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgx::pg_schema]
 mod tests {

diff --git a/pgx-utils/src/sql_entity_graph/pg_extern/attribute.rs b/pgx-utils/src/sql_entity_graph/pg_extern/attribute.rs
@@ -26,9 +26,9 @@ pub enum Attribute {
     Sql(ToSqlConfig),
 }
 
-impl ToTokens for Attribute {
-    fn to_tokens(&self, tokens: &mut TokenStream2) {
-        let quoted = match self {
+impl Attribute {
+    pub(crate) fn to_sql_entity_graph_tokens(&self) -> TokenStream2 {
+        match self {
             Attribute::Immutable => quote! { pgx::datum::sql_entity_graph::ExternArgs::Immutable },
             Attribute::Strict => quote! { pgx::datum::sql_entity_graph::ExternArgs::Strict },
             Attribute::Stable => quote! { pgx::datum::sql_entity_graph::ExternArgs::Stable },
@@ -65,7 +65,52 @@ impl ToTokens for Attribute {
             }
             // This attribute is handled separately
             Attribute::Sql(_) => {
-                return;
+                quote! { }
+            }
+        }
+    }
+}
+
+impl ToTokens for Attribute {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        let quoted = match self {
+            Attribute::Immutable => quote! { immutable },
+            Attribute::Strict => quote! { strict },
+            Attribute::Stable => quote! { stable },
+            Attribute::Volatile => quote! { volatile },
+            Attribute::Raw => quote! { raw },
+            Attribute::NoGuard => quote! { no_guard },
+            Attribute::ParallelSafe => {
+                quote! { parallel_safe }
+            }
+            Attribute::ParallelUnsafe => {
+                quote! { parallel_unsafe }
+            }
+            Attribute::ParallelRestricted => {
+                quote! { parallel_restricted }
+            }
+            Attribute::Error(s) => {
+                quote! { error = #s }
+            }
+            Attribute::Schema(s) => {
+                quote! { schema = #s }
+            }
+            Attribute::Name(s) => {
+                quote! { name = #s }
+            }
+            Attribute::Cost(s) => {
+                quote! { cost = #s }
+            }
+            Attribute::Requires(items) => {
+                let items_iter = items
+                    .iter()
+                    .map(|x| x.to_token_stream())
+                    .collect::<Vec<_>>();
+                quote! { requires = [#(#items_iter),*] }
+            }
+            // This attribute is handled separately
+            Attribute::Sql(to_sql_config) => {
+                quote! { sql = #to_sql_config }
             }
         };
         tokens.append_all(quoted);
@@ -128,7 +173,7 @@ impl Parse for Attribute {
                     }
                 }
             }
-            _ => return Err(syn::Error::new(Span::call_site(), "Invalid option")),
+            e => return Err(syn::Error::new(Span::call_site(), format!("Invalid option `{}` inside `{} {}`", e, ident.to_string(), input.to_string()))),
         };
         Ok(found)
     }

diff --git a/pgx-utils/src/sql_entity_graph/pg_extern/mod.rs b/pgx-utils/src/sql_entity_graph/pg_extern/mod.rs
@@ -219,7 +219,9 @@ impl ToTokens for PgExtern {
         let name = self.name();
         let schema = self.schema();
         let schema_iter = schema.iter();
-        let extern_attrs = self.attrs.iter().collect::<Punctuated<_, Token![,]>>();
+        let extern_attrs = self.attrs.iter()
+            .map(|attr| attr.to_sql_entity_graph_tokens())
+            .collect::<Punctuated<_, Token![,]>>();
         let search_path = self.search_path().into_iter();
         let inputs = self.inputs().unwrap();
         let returns = match self.returns() {