From 4271e93aa3e9b885697c1cb5b7f2be3b0d129bf7 Mon Sep 17 00:00:00 2001
From: Manuel Maute <manuel.maute@bradbit.com>
Date: Tue, 17 Jan 2023 18:47:03 +0100
Subject: [PATCH] feat: add typings for functions returning table

---
 src/server/templates/typescript.ts | 12 +++++++++++-
 test/db/00-init.sql                | 12 ++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/server/templates/typescript.ts b/src/server/templates/typescript.ts
index bb3213c2..84e79e60 100644
--- a/src/server/templates/typescript.ts
+++ b/src/server/templates/typescript.ts
@@ -222,7 +222,7 @@ export interface Database {
                 ([fnName, fns]) =>
                   `${JSON.stringify(fnName)}: ${fns
                     .map(
-                      ({ args, return_type }) => `{
+                      ({ args, return_type, complete_statement }) => `{
                   Args: ${(() => {
                     const inArgs = args.filter(({ mode }) => mode === 'in')
 
@@ -277,6 +277,16 @@ export interface Database {
                       )} }[]`
                     }
 
+                    const setOfTable = schemaTables.find(
+                      ({ schema: _schema, name }) => _schema === schema.name && name === return_type
+                    )
+
+                    if (setOfTable) {
+                      return `Database[${JSON.stringify(schema.name)}]['Tables'][${JSON.stringify(
+                        return_type
+                      )}]['Row']${complete_statement.includes('SETOF') ? '[]' : ''}`
+                    }
+
                     return pgTypeToTsType(return_type, types, schemas)
                   })()}
                 }`
diff --git a/test/db/00-init.sql b/test/db/00-init.sql
index 17c65bc2..644fae0a 100644
--- a/test/db/00-init.sql
+++ b/test/db/00-init.sql
@@ -65,3 +65,15 @@ create foreign table foreign_table (
   name text,
   status user_status
 ) server foreign_server options (schema_name 'public', table_name 'users');
+
+create or replace function public.get_user_by_id(integer) returns public.users as $$
+select * from public.users where id = $1;
+$$ language sql stable;
+
+create or replace function public.get_users() returns setof public.users as $$
+select * from public.users;
+$$ language sql stable;
+
+create or replace function public.get_users_typed() returns table (id int, name text) as $$
+select id, name from public.users;
+$$ language sql stable;
\ No newline at end of file