fix(typegen): set returning function

Author: soedirgo

src/lib/sql/functions.sql

@@ -25,7 +25,7 @@ with functions as (
     p.prokind = 'f'
 )
 select
-  f.oid :: int8 as id,
+  f.oid::int8 as id,
   n.nspname as schema,
   f.proname as name,
   l.lanname as language,
@@ -40,7 +40,9 @@ select
   coalesce(f_args.args, '[]') as args,
   pg_get_function_arguments(f.oid) as argument_types,
   pg_get_function_identity_arguments(f.oid) as identity_argument_types,
-  t.typname as return_type,
+  rt.typname as return_type,
+  nullif(rt.typrelid::int8, 0) as return_type_relation_id,
+  f.proretset as is_set_returning_function,
   case
     when f.provolatile = 'i' then 'IMMUTABLE'
     when f.provolatile = 's' then 'STABLE'
@@ -52,7 +54,7 @@ from
   functions f
   left join pg_namespace n on f.pronamespace = n.oid
   left join pg_language l on f.prolang = l.oid
-  left join pg_type t on t.oid = f.prorettype
+  left join pg_type rt on rt.oid = f.prorettype
   left join (
     select
       oid,

src/lib/types.ts

@@ -108,6 +108,8 @@ const postgresFunctionSchema = Type.Object({
   argument_types: Type.String(),
   identity_argument_types: Type.String(),
   return_type: Type.String(),
+  return_type_relation_id: Type.Union([Type.Integer(), Type.Null()]),
+  is_set_returning_function: Type.Boolean(),
   behavior: Type.Union([
     Type.Literal('IMMUTABLE'),
     Type.Literal('STABLE'),

src/server/templates/typescript.ts

@@ -222,7 +222,12 @@ export interface Database {
                 ([fnName, fns]) =>
                   `${JSON.stringify(fnName)}: ${fns
                     .map(
-                      ({ args, return_type }) => `{
+                      ({
+                        args,
+                        return_type,
+                        return_type_relation_id,
+                        is_set_returning_function,
+                      }) => `{
                   Args: ${(() => {
                     const inArgs = args.filter(({ mode }) => mode === 'in')
 
@@ -252,14 +257,16 @@ export interface Database {
                       return { name, type: 'unknown', has_default }
                     })
 
-                    return `{ ${argsNameAndType.map(
-                      ({ name, type, has_default }) =>
-                        `${JSON.stringify(name)}${has_default ? '?' : ''}: ${type}`
-                    )} }`
+                    return `{
+                      ${argsNameAndType.map(
+                        ({ name, type, has_default }) =>
+                          `${JSON.stringify(name)}${has_default ? '?' : ''}: ${type}`
+                      )}
+                    }`
                   })()}
-                  Returns: ${(() => {
+                  Returns: (${(() => {
+                    // Case 1: `returns table`.
                     const tableArgs = args.filter(({ mode }) => mode === 'table')
-
                     if (tableArgs.length > 0) {
                       const argsNameAndType = tableArgs.map(({ name, type_id }) => {
                         let type = arrayTypes.find(({ id }) => id === type_id)
@@ -278,13 +285,35 @@ export interface Database {
                         return { name, type: 'unknown' }
                       })
 
-                      return `{ ${argsNameAndType.map(
-                        ({ name, type }) => `${JSON.stringify(name)}: ${type}`
-                      )} }[]`
+                      return `{
+                        ${argsNameAndType.map(
+                          ({ name, type }) => `${JSON.stringify(name)}: ${type}`
+                        )}
+                      }`
+                    }
+
+                    // Case 2: returns a relation's row type.
+                    const relation = [...tables, ...views].find(
+                      ({ id }) => id === return_type_relation_id
+                    )
+                    if (relation) {
+                      return `{
+                        ${relation.columns
+                          .sort(({ name: a }, { name: b }) => a.localeCompare(b))
+                          .map(
+                            (column) =>
+                              `${JSON.stringify(column.name)}: ${pgTypeToTsType(
+                                column.format,
+                                types,
+                                schemas
+                              )} ${column.is_nullable ? '| null' : ''}`
+                          )}
+                      }`
                     }
 
+                    // Case 3: returns base/composite/enum type.
                     return pgTypeToTsType(return_type, types, schemas)
-                  })()}
+                  })()})${is_set_returning_function ? '[]' : ''}
                 }`
                     )
                     .join('|')}`
@@ -367,7 +396,7 @@ const pgTypeToTsType = (
   } else if (pgType === 'void') {
     return 'undefined'
   } else if (pgType === 'record') {
-    return 'Record<string, unknown>[]'
+    return 'Record<string, unknown>'
   } else if (pgType.startsWith('_')) {
     return `(${pgTypeToTsType(pgType.substring(1), types, schemas)})[]`
   } else {

test/lib/functions.ts

@@ -32,9 +32,11 @@ test('list', async () => {
       "definition": "select $1 + $2;",
       "id": Any<Number>,
       "identity_argument_types": "integer, integer",
+      "is_set_returning_function": false,
       "language": "sql",
       "name": "add",
       "return_type": "int4",
+      "return_type_relation_id": null,
       "schema": "public",
       "security_definer": false,
     }
@@ -129,9 +131,11 @@ test('retrieve, create, update, delete', async () => {
         "definition": "select a + b",
         "id": Any<Number>,
         "identity_argument_types": "a smallint, b smallint",
+        "is_set_returning_function": false,
         "language": "sql",
         "name": "test_func",
         "return_type": "int4",
+        "return_type_relation_id": null,
         "schema": "public",
         "security_definer": true,
       },
@@ -176,9 +180,11 @@ test('retrieve, create, update, delete', async () => {
         "definition": "select a + b",
         "id": Any<Number>,
         "identity_argument_types": "a smallint, b smallint",
+        "is_set_returning_function": false,
         "language": "sql",
         "name": "test_func",
         "return_type": "int4",
+        "return_type_relation_id": null,
         "schema": "public",
         "security_definer": true,
       },
@@ -227,9 +233,11 @@ test('retrieve, create, update, delete', async () => {
         "definition": "select b - a",
         "id": Any<Number>,
         "identity_argument_types": "a smallint, b smallint",
+        "is_set_returning_function": false,
         "language": "sql",
         "name": "test_func_renamed",
         "return_type": "int4",
+        "return_type_relation_id": null,
         "schema": "test_schema",
         "security_definer": true,
       },
@@ -274,9 +282,11 @@ test('retrieve, create, update, delete', async () => {
         "definition": "select b - a",
         "id": Any<Number>,
         "identity_argument_types": "a smallint, b smallint",
+        "is_set_returning_function": false,
         "language": "sql",
         "name": "test_func_renamed",
         "return_type": "int4",
+        "return_type_relation_id": null,
         "schema": "test_schema",
         "security_definer": true,
       },