Improve SELECT/PERFORM/RETURN formatting

src/syntax/procedural_language.ts

@@ -182,7 +182,10 @@ export const proceduralLanguageMap: CstToDocMap<AllProceduralNodes> = {
   // CALL
   call_stmt: (print) => group(print.spaced(["callKw", "func"])),
   // RETURN
-  return_stmt: (print) => group(print.spaced(["returnKw", "expr"])),
+  return_stmt: (print, node) =>
+    node.expr
+      ? group([print("returnKw"), indent([line, print("expr")])])
+      : print("returnKw"),
   return_next_stmt: (print) => group(print.spaced(["returnNextKw", "expr"])),
   return_query_stmt: (print, node) => {
     if (node.expr.type === "execute_expr") {

src/syntax/select.ts

@@ -1,8 +1,9 @@
 import { Doc } from "prettier";
-import { AllSelectNodes, LimitClause } from "sql-parser-cst";
+import { AllSelectNodes, LimitClause, SelectClause } from "sql-parser-cst";
 import { PrintFn } from "../PrintFn";
 import { isDefined } from "../utils";
 import { CstToDocMap } from "../CstToDocMap";
+import { isEmpty } from "../node_utils";
 import {
   join,
   line,
@@ -62,12 +63,16 @@ export const selectMap: CstToDocMap<AllSelectNodes> = {
     ]),
 
   // SELECT clause
-  select_clause: (print, node, path, opts) =>
-    group([
+  select_clause: (print, node, path, opts) => {
+    if (hasSingleColumn(node)) {
+      return group(print.spaced(["selectKw", "modifiers", "columns"]));
+    }
+    return group([
       print.spaced(["selectKw", "modifiers"]),
       node.columns ? indent([line, print("columns")]) : [],
       containsNewline(node, opts) ? breakParent : [],
-    ]),
+    ]);
+  },
   select_all: (print) => print.spaced(["allKw"]),
   select_distinct: (print) => print.spaced(["distinctKw"]),
   select_distinct_on: (print) =>
@@ -305,3 +310,10 @@ const printLimitValues = (
     return print("count");
   }
 };
+
+const hasSingleColumn = (node: SelectClause): boolean => {
+  if (!node.columns) {
+    return false;
+  }
+  return node.columns.items.filter((item) => !isEmpty(item)).length === 1;
+};

test/expr/expr.test.ts

@@ -139,15 +139,14 @@ describe("expr", () => {
 
   it(`preserves parenthesis around compound-SELECT inside function arguments`, async () => {
     await test(dedent`
-      SELECT
-        coalesce(
-          '',
-          (
-            SELECT x FROM xs
-            UNION
-            SELECT y FROM ys
-          )
+      SELECT coalesce(
+        '',
+        (
+          SELECT x FROM xs
+          UNION
+          SELECT y FROM ys
         )
+      )
       FROM tbl
     `);
   });
@@ -163,74 +162,68 @@ describe("expr", () => {
   describe("case", () => {
     it(`formats CASE expression always on multiple lines`, async () => {
       await test(dedent`
-        SELECT
-          CASE x
-            WHEN 1 THEN 'A'
-            ELSE 'B'
-          END
+        SELECT CASE x
+          WHEN 1 THEN 'A'
+          ELSE 'B'
+        END
       `);
     });
 
     it(`formats CASE expression with base expression`, async () => {
       await test(dedent`
-        SELECT
-          CASE status
-            WHEN 1 THEN 'good'
-            WHEN 2 THEN 'bad'
-            ELSE 'unknown'
-          END
+        SELECT CASE status
+          WHEN 1 THEN 'good'
+          WHEN 2 THEN 'bad'
+          ELSE 'unknown'
+        END
       `);
     });
 
     it(`formats CASE expression without base expression`, async () => {
       await test(dedent`
-        SELECT
-          CASE
-            WHEN status = 1 THEN 'good'
-            WHEN status = 2 THEN 'bad'
-            ELSE 'unknown'
-          END
+        SELECT CASE
+          WHEN status = 1 THEN 'good'
+          WHEN status = 2 THEN 'bad'
+          ELSE 'unknown'
+        END
       `);
     });
 
     it(`breaks long WHEN/THEN into separate lines`, async () => {
       await test(
         dedent`
-          SELECT
-            CASE
-              WHEN column_name = 1 THEN
-                result_name
-            END
+          SELECT CASE
+            WHEN column_name = 1 THEN
+              result_name
+          END
         `,
-        { printWidth: 40 },
+        { printWidth: 35 },
       );
     });
 
     it(`breaks multiple long WHEN/THEN clauses without blank lines between them`, async () => {
       await test(
         dedent`
-          SELECT
-            CASE
-              WHEN column_name = 1 THEN
-                result_name
-              WHEN column_name = 2 THEN
-                other_result
-              ELSE foo
-            END
+          SELECT CASE
+            WHEN column_name = 1 THEN
+              result_name
+            WHEN column_name = 2 THEN
+              other_result
+            ELSE foo
+          END
         `,
-        { printWidth: 40 },
+        { printWidth: 35 },
       );
     });
 
     it(`breaks long ELSE into separate lines`, async () => {
       await test(
         dedent`
-          SELECT
-            CASE
-              WHEN x = 1 THEN 'A'
-              ELSE
-                long_fallback_expression
-            END
+          SELECT CASE
+            WHEN x = 1 THEN 'A'
+            ELSE
+              long_fallback_expression
+          END
         `,
         { printWidth: 25 },
       );
@@ -239,14 +232,13 @@ describe("expr", () => {
     it(`indents long multi-condition WHEN clauses`, async () => {
       await test(
         dedent`
-          SELECT
-            CASE
-              WHEN
-                column_name = 1
-                AND (other_name = 2 OR other_name = 3)
-              THEN
-                result_name
-            END
+          SELECT CASE
+            WHEN
+              column_name = 1
+              AND (other_name = 2 OR other_name = 3)
+            THEN
+              result_name
+          END
         `,
         { printWidth: 50 },
       );
@@ -255,12 +247,11 @@ describe("expr", () => {
     it(`indents long multi-expression THEN clauses`, async () => {
       await test(
         dedent`
-          SELECT
-            CASE
-              WHEN column_name = 1 THEN
-                result_name = 1
-                AND other_name = 2
-            END
+          SELECT CASE
+            WHEN column_name = 1 THEN
+              result_name = 1
+              AND other_name = 2
+          END
         `,
         { printWidth: 45 },
       );
@@ -290,10 +281,9 @@ describe("expr", () => {
 
     it(`formats BigQuery array field access to multiple lines`, async () => {
       await testBigquery(dedent`
-        SELECT
-          ["Coffee Cup", "Tea Kettle", "Milk Glass"][
-            SAFE_OFFSET(some_really_long_index_number)
-          ]
+        SELECT ["Coffee Cup", "Tea Kettle", "Milk Glass"][
+          SAFE_OFFSET(some_really_long_index_number)
+        ]
       `);
     });
 

test/expr/func.test.ts

@@ -4,37 +4,33 @@ import { pretty, test, testBigquery, testPostgresql } from "../test_utils";
 // Functions and function-like language constructs
 describe("functions", () => {
   it(`formats function call to single line`, async () => {
-    expect(await pretty(`SELECT sqrt(1, 2, 3)`, { printWidth: 16 }))
+    expect(await pretty(`SELECT sqrt(1, 2, 3)`, { printWidth: 25 }))
       .toBe(dedent`
-      SELECT
-        sqrt(1, 2, 3)
+      SELECT sqrt(1, 2, 3)
     `);
   });
 
   it(`formats function call to multiple lines`, async () => {
     expect(await pretty(`SELECT sqrt(1, 2, 3)`, { printWidth: 10 }))
       .toBe(dedent`
-      SELECT
-        sqrt(
-          1,
-          2,
-          3
-        )
+      SELECT sqrt(
+        1,
+        2,
+        3
+      )
     `);
   });
 
   it(`keeps empty function args on one line`, async () => {
     expect(await pretty(`SELECT my_func()`, { printWidth: 10 })).toBe(dedent`
-      SELECT
-        my_func()
+      SELECT my_func()
     `);
   });
 
   it(`does not treat count(DISTINCT) as empty function args`, async () => {
-    expect(await pretty(`SELECT count(DISTINCT id)`, { printWidth: 25 }))
+    expect(await pretty(`SELECT count(DISTINCT id)`, { printWidth: 30 }))
       .toBe(dedent`
-      SELECT
-        count(DISTINCT id)
+      SELECT count(DISTINCT id)
     `);
   });
 
@@ -47,8 +43,7 @@ describe("functions", () => {
     ).toBe(dedent`
       CREATE FUNCTION my_func() AS
         (
-          SELECT
-            1
+          SELECT 1
         )
     `);
   });
@@ -57,10 +52,9 @@ describe("functions", () => {
     expect(
       await pretty(`SELECT my_func(/* comment */)`, { printWidth: 25 }),
     ).toBe(dedent`
-      SELECT
-        my_func(
-          /* comment */
-        )
+      SELECT my_func(
+        /* comment */
+      )
     `);
   });
 
@@ -75,11 +69,10 @@ describe("functions", () => {
         { printWidth: 25 },
       ),
     ).toBe(dedent`
-      SELECT
-        my_func(
-          -- comment
+      SELECT my_func(
+        -- comment
 
-        )
+      )
     `);
   });
 
@@ -90,13 +83,12 @@ describe("functions", () => {
   });
   it(`formats long named function arguments`, async () => {
     await testBigquery(dedent`
-      SELECT
-        concat_lower_or_upper(
-          first_parameter =>
-            another_function_call(another_function_param => 'Hohoho Hello'),
-          second_parameter => 'World',
-          uppercase => TRUE
-        )
+      SELECT concat_lower_or_upper(
+        first_parameter =>
+          another_function_call(another_function_param => 'Hohoho Hello'),
+        second_parameter => 'World',
+        uppercase => TRUE
+      )
     `);
   });
   // Behavior when sqlCanonicalSyntax: false
@@ -134,15 +126,14 @@ describe("functions", () => {
 
     it(`formats long combo of DISTINCT, IGNORE NULLS, ORDER BY, LIMIT to multiple lines`, async () => {
       await testBigquery(dedent`
-        SELECT
-          my_function_name(
-            DISTINCT
-            first_argument,
-            second_argument
-            IGNORE NULLS
-            ORDER BY some_field_name, other_field_name
-            LIMIT 10000, 200
-          )
+        SELECT my_function_name(
+          DISTINCT
+          first_argument,
+          second_argument
+          IGNORE NULLS
+          ORDER BY some_field_name, other_field_name
+          LIMIT 10000, 200
+        )
       `);
     });
   });