feat(typescript): add functions setof type introspection

src/lib/sql/functions.sql

@@ -44,6 +44,20 @@ select
   pg_get_function_result(f.oid) as return_type,
   nullif(rt.typrelid::int8, 0) as return_type_relation_id,
   f.proretset as is_set_returning_function,
+  case
+    when f.proretset and rt.typrelid != 0 then true
+    else false
+  end as returns_set_of_table,
+  case 
+    when f.proretset and rt.typrelid != 0 then
+      (select relname from pg_class where oid = rt.typrelid)
+    else null
+  end as return_table_name,
+  case
+    when f.proretset then
+      coalesce(f.prorows, 0) > 1
+    else false
+  end as returns_multiple_rows,
   case
     when f.provolatile = 'i' then 'IMMUTABLE'
     when f.provolatile = 's' then 'STABLE'
@@ -76,32 +90,48 @@ from
     select
       oid,
       jsonb_agg(jsonb_build_object(
-        'mode', t2.mode,
+        'mode', mode,
         'name', name,
         'type_id', type_id,
-        'has_default', has_default
+        'has_default', has_default,
+        'table_name', table_name
       )) as args
     from
       (
         select
-          oid,
-          unnest(arg_modes) as mode,
-          unnest(arg_names) as name,
-          unnest(arg_types)::int8 as type_id,
-          unnest(arg_has_defaults) as has_default
+          t1.oid,
+          t2.mode,
+          t1.name,
+          t1.type_id,
+          t1.has_default,
+          case 
+            when pt.typrelid != 0 then pc.relname 
+            else null 
+          end as table_name
         from
-          functions
-      ) as t1,
-      lateral (
-        select
-          case
-            when t1.mode = 'i' then 'in'
-            when t1.mode = 'o' then 'out'
-            when t1.mode = 'b' then 'inout'
-            when t1.mode = 'v' then 'variadic'
-            else 'table'
-          end as mode
-      ) as t2
+          (
+            select
+              oid,
+              unnest(arg_modes) as mode,
+              unnest(arg_names) as name,
+              unnest(arg_types)::int8 as type_id,
+              unnest(arg_has_defaults) as has_default
+            from
+              functions
+          ) as t1
+          cross join lateral (
+            select
+              case
+                when t1.mode = 'i' then 'in'
+                when t1.mode = 'o' then 'out'
+                when t1.mode = 'b' then 'inout'
+                when t1.mode = 'v' then 'variadic'
+                else 'table'
+              end as mode
+          ) as t2
+          left join pg_type pt on pt.oid = t1.type_id
+          left join pg_class pc on pc.oid = pt.typrelid
+      ) sub
     group by
-      t1.oid
+      oid
   ) f_args on f_args.oid = f.oid

src/lib/types.ts

@@ -148,6 +148,7 @@ const postgresFunctionSchema = Type.Object({
       name: Type.String(),
       type_id: Type.Number(),
       has_default: Type.Boolean(),
+      table_name: Type.Union([Type.String(), Type.Null()]),
     })
   ),
   argument_types: Type.String(),
@@ -156,6 +157,9 @@ const postgresFunctionSchema = Type.Object({
   return_type: Type.String(),
   return_type_relation_id: Type.Union([Type.Integer(), Type.Null()]),
   is_set_returning_function: Type.Boolean(),
+  returns_set_of_table: Type.Boolean(),
+  return_table_name: Type.Union([Type.String(), Type.Null()]),
+  returns_multiple_rows: Type.Boolean(),
   behavior: Type.Union([
     Type.Literal('IMMUTABLE'),
     Type.Literal('STABLE'),

src/server/templates/typescript.ts

@@ -343,7 +343,26 @@ export type Database = {
                         }
 
                         return 'unknown'
-                      })()}${fns[0].is_set_returning_function ? '[]' : ''}
+                      })()}${fns[0].is_set_returning_function && fns[0].returns_multiple_rows ? '[]' : ''}
+                      ${
+                        // if the function return a set of a table and some definition take in parameter another table
+                        fns[0].returns_set_of_table &&
+                        fns.some((fnd) => fnd.args.length === 1 && fnd.args[0].table_name)
+                          ? `SetofOptions: {
+                        from: ${fns
+                          // if the function take a row as first parameter
+                          .filter((fnd) => fnd.args.length === 1 && fnd.args[0].table_name)
+                          .map((fnd) => {
+                            const arg_type = types.find((t) => t.id === fnd.args[0].type_id)
+                            return JSON.stringify(arg_type?.format)
+                          })
+                          .join(' | ')}
+                        to: ${JSON.stringify(fns[0].return_table_name)}
+                        isOneToOne: ${fns[0].returns_multiple_rows ? false : true}
+                      }
+                      `
+                          : ''
+                      }
                     }`
               )
             })()}

test/db/00-init.sql

@@ -181,3 +181,32 @@ LANGUAGE SQL STABLE
 AS $$
   SELECT * FROM public.todos WHERE "user-id" = todo_row."user-id";
 $$;
+
+-- SETOF composite_type - Returns multiple rows of a custom composite type
+CREATE OR REPLACE FUNCTION public.get_composite_type_data()
+RETURNS SETOF composite_type_with_array_attribute
+LANGUAGE SQL STABLE
+AS $$
+  SELECT ROW(ARRAY['hello', 'world']::text[])::composite_type_with_array_attribute
+  UNION ALL
+  SELECT ROW(ARRAY['foo', 'bar']::text[])::composite_type_with_array_attribute;
+$$;
+
+-- SETOF record - Returns multiple rows with structure defined in the function
+CREATE OR REPLACE FUNCTION public.get_user_summary()
+RETURNS SETOF record
+LANGUAGE SQL STABLE
+AS $$
+  SELECT u.id, name, count(t.id) as todo_count
+  FROM public.users u
+  LEFT JOIN public.todos t ON t."user-id" = u.id
+  GROUP BY u.id, u.name;
+$$;
+
+-- SETOF scalar_type - Returns multiple values of a basic type
+CREATE OR REPLACE FUNCTION public.get_user_ids()
+RETURNS SETOF bigint
+LANGUAGE SQL STABLE
+AS $$
+  SELECT id FROM public.users;
+$$;