Merge pull request #454 from supabase/add_array_support

add support for array arguments and return types in UDFs

docs/changelog.md

@@ -57,3 +57,4 @@
 ## master
 - bugfix: make non-default args non-null in UDFs
 - bugfix: default value of a string type argument in a UDF was wrapped in single quotes
+- feature: add support for array types in UDFs

docs/functions.md

@@ -278,6 +278,97 @@ Functions returning multiple rows of a table or view are exposed as [collections
 
     A set returning function with any of its argument names clashing with argument names of a collection (`first`, `last`, `before`, `after`, `filter`, or `orderBy`) will not be exposed.
 
+Functions accepting or returning arrays of non-composite types are also supported. In the following example, the `ids` array is used to filter rows from the `Account` table:
+
+=== "Function"
+
+    ```sql
+    create table "Account"(
+      id serial primary key,
+      email varchar(255) not null
+    );
+
+    insert into "Account"(email)
+    values
+      ('[email protected]'),
+      ('[email protected]'),
+      ('[email protected]');
+
+    create function "accountsByIds"("ids" int[])
+      returns setof "Account"
+      stable
+      language sql
+    as $$ select id, email from "Account" where id = any(ids); $$;
+    ```
+
+=== "QueryType"
+
+    ```graphql
+    type Query {
+      accountsByIds(
+        ids: Int[]!
+
+        """Query the first `n` records in the collection"""
+        first: Int
+
+        """Query the last `n` records in the collection"""
+        last: Int
+
+        """Query values in the collection before the provided cursor"""
+        before: Cursor
+
+        """Query values in the collection after the provided cursor"""
+        after: Cursor
+
+        """Filters to apply to the results set when querying from the collection"""
+        filter: AccountFilter
+
+        """Sort order to apply to the collection"""
+        orderBy: [AccountOrderBy!]
+      ): AccountConnection
+    }
+    ```
+
+=== "Query"
+
+    ```graphql
+    query {
+      accountsByIds(ids: [1, 2]) {
+        edges {
+          node {
+            id
+            email
+          }
+        }
+      }
+    }
+    ```
+
+=== "Response"
+
+    ```json
+    {
+      "data": {
+        "accountsByIds": {
+          "edges": [
+            {
+              "node": {
+                "id": 1,
+                "email": "[email protected]"
+              }
+            },
+            {
+              "node": {
+                "id": 2,
+                "email": "[email protected]"
+              }
+            }
+          ]
+        }
+      }
+    }
+    ```
+
 ## Default Arguments
 
 Functions with default arguments can have their default arguments omitted.
@@ -328,4 +419,5 @@ The following features are not yet supported. Any function using these features
 * Functions with a nameless argument
 * Functions returning void
 * Variadic functions
-* Function that accept or return an array type
+* Functions that accept or return an array of composite type
+* Functions that accept or return an enum type or an array of enum type

src/builder.rs

@@ -564,6 +564,7 @@ pub struct FunctionCallBuilder {
 
 pub enum FuncCallReturnTypeBuilder {
     Scalar,
+    List,
     Node(NodeBuilder),
     Connection(ConnectionBuilder),
 }
@@ -600,6 +601,7 @@ where
 
             let return_type_builder = match func_call_resp_type.return_type.deref() {
                 __Type::Scalar(_) => FuncCallReturnTypeBuilder::Scalar,
+                __Type::List(_) => FuncCallReturnTypeBuilder::List,
                 __Type::Node(_) => {
                     let node_builder = to_node_builder(
                         field,
@@ -628,7 +630,7 @@ where
                             .unmodified_type()
                             .name()
                             .ok_or("Encountered type without name in function call builder")?
-                    ))
+                    ));
                 }
             };
 
@@ -2193,7 +2195,12 @@ impl __Schema {
                             None => __TypeField::PossibleTypes(None),
                         },
                         "ofType" => {
-                            let unwrapped_type_builder = match type_ {
+                            let field_type = if let __Type::FuncCallResponse(func_call_resp_type) = type_ {
+                                func_call_resp_type.return_type.deref()
+                            } else {
+                                type_
+                            };
+                            let unwrapped_type_builder = match field_type {
                                 __Type::List(list_type) => {
                                     let inner_type: __Type = (*(list_type.type_)).clone();
                                     Some(self.to_type_builder_from_type(

src/sql_types.rs

@@ -123,8 +123,14 @@ impl Function {
 
     fn arg_types_are_supported(&self, types: &HashMap<u32, Arc<Type>>) -> bool {
         self.args().all(|(arg_type, _, _, _)| {
-            if let Some(return_type) = types.get(&arg_type) {
-                return_type.category == TypeCategory::Other
+            if let Some(arg_type) = types.get(&arg_type) {
+                let array_element_type_is_supported = self.array_element_type_is_supported(
+                    &arg_type.category,
+                    arg_type.array_element_type_oid,
+                    types,
+                );
+                arg_type.category == TypeCategory::Other
+                    || (arg_type.category == TypeCategory::Array && array_element_type_is_supported)
             } else {
                 false
             }
@@ -133,16 +139,43 @@ impl Function {
 
     fn return_type_is_supported(&self, types: &HashMap<u32, Arc<Type>>) -> bool {
         if let Some(return_type) = types.get(&self.type_oid) {
+            let array_element_type_is_supported = self.array_element_type_is_supported(
+                &return_type.category,
+                return_type.array_element_type_oid,
+                types,
+            );
             return_type.category != TypeCategory::Pseudo
+                && return_type.category != TypeCategory::Enum
                 && return_type.name != "record"
                 && return_type.name != "trigger"
                 && return_type.name != "event_trigger"
-                && !self.type_name.ends_with("[]")
+                && array_element_type_is_supported
         } else {
             false
         }
     }
 
+    fn array_element_type_is_supported(
+        &self,
+        type_category: &TypeCategory,
+        array_element_type_oid: Option<u32>,
+        types: &HashMap<u32, Arc<Type>>,
+    ) -> bool {
+        if *type_category == TypeCategory::Array {
+            if let Some(array_element_type_oid) = array_element_type_oid {
+                if let Some(array_element_type) = types.get(&array_element_type_oid) {
+                    array_element_type.category == TypeCategory::Other
+                } else {
+                    false
+                }
+            } else {
+                false
+            }
+        } else {
+            true
+        }
+    }
+
     fn is_function_overloaded(&self, function_name_to_count: &HashMap<&String, u32>) -> bool {
         if let Some(&count) = function_name_to_count.get(&self.name) {
             count > 1

src/transpile.rs

@@ -560,7 +560,7 @@ impl FunctionCallBuilder {
         let func_name = quote_ident(&self.function.name);
 
         let query = match &self.return_type_builder {
-            FuncCallReturnTypeBuilder::Scalar => {
+            FuncCallReturnTypeBuilder::Scalar | FuncCallReturnTypeBuilder::List => {
                 let type_adjustment_clause = apply_suffix_casts(self.function.type_oid);
                 format!("select to_jsonb({func_schema}.{func_name}{args_clause}{type_adjustment_clause}) {block_name};")
             }
@@ -1353,6 +1353,9 @@ fn apply_suffix_casts(type_oid: u32) -> String {
         20 => "::text",           // bigints as text
         114 | 3802 => "#>> '{}'", // json/b as stringified
         1700 => "::text",         // numeric as text
+        1016 => "::text[]",       // bigint arrays as array of text
+        199 | 3807 => "::text[]", // json/b array as array of text
+        1231 => "::text[]",       // numeric array as array of text
         _ => "",
     }
     .to_string()