Format RETURNS TABLE (..)

nene · Jan 27, 2024 · dbda44b · dbda44b
1 parent 46be069
commit dbda44b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/src/syntax/data_type.ts b/src/syntax/data_type.ts
@@ -3,7 +3,7 @@ import { isArray } from "../utils";
 import { CstToDocMap } from "../CstToDocMap";
 import { group, indent, softline } from "../print_utils";
 
-export const dataTypeMap: Partial<CstToDocMap<AllDataTypeNodes>> = {
+export const dataTypeMap: CstToDocMap<AllDataTypeNodes> = {
   // print single-word types as `TYPE(10)` and multi-word types as `MY TYPE (10)`
   named_data_type: (print, node) =>
     (isArray(node.nameKw) ? print.spaced : print)(["nameKw", "params"]),
@@ -16,4 +16,5 @@ export const dataTypeMap: Partial<CstToDocMap<AllDataTypeNodes>> = {
   array_type_param: (print) => print.spaced(["dataType", "constraints"]),
   struct_type_param: (print) =>
     print.spaced(["name", "dataType", "constraints"]),
+  table_data_type: (print) => print.spaced(["tableKw", "columns"]),
 };
diff --git a/test/ddl/function.test.ts b/test/ddl/function.test.ts
@@ -45,16 +45,24 @@ describe("function", () => {
 
     it(`formats RETURNS clause`, async () => {
       await testBigquery(dedent`
-        CREATE FUNCTION IF NOT EXISTS my_func()
+        CREATE FUNCTION my_func()
         RETURNS INT64
         AS
           (SELECT 1)
       `);
     });
 
+    it(`formats RETURNS TABLE`, async () => {
+      await testPostgresql(dedent`
+        CREATE FUNCTION foo()
+        RETURNS TABLE (id INT, name TEXT)
+        AS 'SELECT 1'
+      `);
+    });
+
     it(`formats OPTIONS (...)`, async () => {
       await testBigquery(dedent`
-        CREATE FUNCTION IF NOT EXISTS my_func()
+        CREATE FUNCTION my_func()
         AS
           (SELECT 1)
         OPTIONS (description = 'constant-value function')