feat(typescript): add functions setof type introspection

src/lib/PostgresMetaTypes.ts

@@ -33,7 +33,7 @@ export default class PostgresMetaTypes {
         t.typrelid = 0
         or (
           select
-            c.relkind ${includeTableTypes ? `in ('c', 'r')` : `= 'c'`}
+            c.relkind ${includeTableTypes ? `in ('c', 'r', 'v')` : `= 'c'`}
           from
             pg_class c
           where

src/lib/sql/functions.sql

@@ -44,6 +44,25 @@ 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 and exists (
+      select 1 from pg_class c 
+      where c.oid = rt.typrelid 
+      -- exclude custom types relation from what is considered a set of table
+      and c.relkind in ('r', 'p', 'v', 'm', 'f')
+    ) 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 +95,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

@@ -27,6 +27,14 @@ export const apply = async ({
   const columnsByTableId = Object.fromEntries<PostgresColumn[]>(
     [...tables, ...foreignTables, ...views, ...materializedViews].map((t) => [t.id, []])
   )
+  // group types by id for quicker lookup
+  const typesById = types.reduce(
+    (acc, type) => {
+      acc[type.id] = type
+      return acc
+    },
+    {} as Record<string, (typeof types)[number]>
+  )
   columns
     .filter((c) => c.table_id in columnsByTableId)
     .sort(({ name: a }, { name: b }) => a.localeCompare(b))
@@ -45,6 +53,7 @@ export type Database = {
       const schemaViews = [...views, ...materializedViews]
         .filter((view) => view.schema === schema.name)
         .sort(({ name: a }, { name: b }) => a.localeCompare(b))
+
       const schemaFunctions = functions
         .filter((func) => {
           if (func.schema !== schema.name) {
@@ -94,7 +103,7 @@ export type Database = {
                       ...schemaFunctions
                         .filter((fn) => fn.argument_types === table.name)
                         .map((fn) => {
-                          const type = types.find(({ id }) => id === fn.return_type_id)
+                          const type = typesById[fn.return_type_id]
                           let tsType = 'unknown'
                           if (type) {
                             tsType = pgTypeToTsType(type.name, { types, schemas, tables, views })
@@ -285,9 +294,8 @@ export type Database = {
                           if (inArgs.length === 0) {
                             return 'Record<PropertyKey, never>'
                           }
-
                           const argsNameAndType = inArgs.map(({ name, type_id, has_default }) => {
-                            const type = types.find(({ id }) => id === type_id)
+                            const type = typesById[type_id]
                             let tsType = 'unknown'
                             if (type) {
                               tsType = pgTypeToTsType(type.name, { types, schemas, tables, views })
@@ -343,7 +351,31 @@ 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
+                          ? `SetofOptions: {
+                        from: ${fns
+                          .map((fnd) => {
+                            if (fnd.args.length > 0 && fnd.args[0].table_name) {
+                              const tableType = typesById[fnd.args[0].type_id]
+                              return JSON.stringify(tableType.format)
+                            } else {
+                              // If the function can be called with scalars or without any arguments, then add a * matching everything
+                              return '"*"'
+                            }
+                          })
+                          // Dedup before join
+                          .filter((value, index, self) => self.indexOf(value) === index)
+                          .toSorted()
+                          .join(' | ')}
+                        to: ${JSON.stringify(fns[0].return_table_name)}
+                        isOneToOne: ${fns[0].returns_multiple_rows ? false : true}
+                      }
+                      `
+                          : ''
+                      }
                     }`
               )
             })()}

test/db/00-init.sql

@@ -55,6 +55,17 @@ $$ language plpgsql;
 CREATE VIEW todos_view AS SELECT * FROM public.todos;
 -- For testing typegen on view-to-view relationships
 create view users_view as select * from public.users;
+-- Create a more complex view for testing
+CREATE VIEW user_todos_summary_view AS
+SELECT 
+    u.id as user_id,
+    u.name as user_name,
+    u.status as user_status,
+    COUNT(t.id) as todo_count,
+    array_agg(t.details) FILTER (WHERE t.details IS NOT NULL) as todo_details
+FROM public.users u
+LEFT JOIN public.todos t ON t."user-id" = u.id
+GROUP BY u.id, u.name, u.status;
 
 create materialized view todos_matview as select * from public.todos;
 
@@ -181,3 +192,81 @@ 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;
+$$;
+
+
+-- Function returning view using scalar as input
+CREATE OR REPLACE FUNCTION public.get_single_user_summary_from_view(search_user_id bigint)
+RETURNS SETOF user_todos_summary_view
+LANGUAGE SQL STABLE
+ROWS 1
+AS $$
+  SELECT * FROM user_todos_summary_view WHERE user_id = search_user_id;
+$$;
+-- Function returning view using table row as input
+CREATE OR REPLACE FUNCTION public.get_single_user_summary_from_view(user_row users)
+RETURNS SETOF user_todos_summary_view
+LANGUAGE SQL STABLE
+ROWS 1 
+AS $$
+  SELECT * FROM user_todos_summary_view WHERE user_id = user_row.id;
+$$;
+-- Function returning view using another view row as input
+CREATE OR REPLACE FUNCTION public.get_single_user_summary_from_view(userview_row users_view)
+RETURNS SETOF user_todos_summary_view
+LANGUAGE SQL STABLE
+ROWS 1 
+AS $$
+  SELECT * FROM user_todos_summary_view WHERE user_id = userview_row.id;
+$$;
+
+
+-- Function returning view using scalar as input
+CREATE OR REPLACE FUNCTION public.get_todos_from_user(search_user_id bigint)
+RETURNS SETOF todos
+LANGUAGE SQL STABLE
+AS $$
+  SELECT * FROM todos WHERE "user-id" = search_user_id;
+$$;
+-- Function returning view using table row as input
+CREATE OR REPLACE FUNCTION public.get_todos_from_user(user_row users)
+RETURNS SETOF todos
+LANGUAGE SQL STABLE
+AS $$
+  SELECT * FROM todos WHERE "user-id" = user_row.id;
+$$;
+-- Function returning view using another view row as input
+CREATE OR REPLACE FUNCTION public.get_todos_from_user(userview_row users_view)
+RETURNS SETOF todos
+LANGUAGE SQL STABLE
+AS $$
+  SELECT * FROM todos WHERE "user-id" = userview_row.id;
+$$;