Format LANGUAGE SQL function bodies

Author: joelmukuthu

src/embed.ts

@@ -2,7 +2,8 @@ import { Printer } from "prettier";
 import { Node } from "sql-parser-cst";
 import { embedJs } from "./embedJs";
 import { embedJson } from "./embedJson";
+import { embedSql } from "./embedSql";
 
 export const embed: NonNullable<Printer<Node>["embed"]> = (...args) => {
-  return embedJson(...args) || embedJs(...args);
+  return embedJson(...args) || embedJs(...args) || embedSql(...args);
 };

src/embedSql.ts

@@ -0,0 +1,56 @@
+import { Printer } from "prettier";
+import { CreateFunctionStmt, Node, StringLiteral } from "sql-parser-cst";
+import {
+  isAsClause,
+  isCreateFunctionStmt,
+  isLanguageClause,
+  isStringLiteral,
+} from "./node_utils";
+import { hardline, indent, stripTrailingHardline } from "./print_utils";
+
+export const embedSql: NonNullable<Printer<Node>["embed"]> = (path, options) => {
+  const node = path.node;
+  const parent = path.getParentNode(0);
+  const grandParent = path.getParentNode(1);
+
+  if (
+    isStringLiteral(node) &&
+    isAsClause(parent) &&
+    isCreateFunctionStmt(grandParent) &&
+    grandParent.clauses.some(isSqlLanguageClause)
+  ) {
+    return async (textToDoc) => {
+      const quote = detectQuote(node);
+
+      if (quote) {
+        const sql = await textToDoc(node.value, options);
+
+        return [
+          quote,
+          indent([hardline, stripTrailingHardline(sql)]),
+          hardline,
+          quote,
+        ];
+      }
+    };
+  }
+
+  return null;
+};
+
+const isSqlLanguageClause = (
+  clause: CreateFunctionStmt["clauses"][0],
+): boolean => isLanguageClause(clause) && clause.name.name === "sql";
+
+const detectQuote = (
+  node: StringLiteral,
+): string | undefined => {
+  const match = node.text.match(/^('|\$[^$]*\$)/);
+  const quote = match?.[1];
+
+  if (quote && node.text.endsWith(quote)) {
+    return quote;
+  }
+
+  return undefined;
+};

test/ddl/function.test.ts

@@ -202,6 +202,74 @@ describe("function", () => {
         AS " return /'''|\\"\\"\\"/.test(x) "
       `);
     });
+
+    it(`formats dollar-quoted SQL function`, async () => {
+      await testPostgresql(dedent`
+        CREATE FUNCTION my_func()
+        RETURNS INT64
+        LANGUAGE sql
+        AS $$
+          SELECT 1;
+        $$
+      `);
+    });
+
+    it(`reformats SQL in dollar-quoted SQL function`, async () => {
+      expect(
+        await pretty(
+          dedent`
+            CREATE FUNCTION my_func()
+            RETURNS INT64
+            LANGUAGE sql
+            AS $body$SELECT 1;
+            select 2$body$
+          `,
+          { dialect: "postgresql" },
+        ),
+      ).toBe(dedent`
+        CREATE FUNCTION my_func()
+        RETURNS INT64
+        LANGUAGE sql
+        AS $body$
+          SELECT 1;
+          SELECT 2;
+        $body$
+      `);
+    });
+
+    it(`formats single-quoted SQL function`, async () => {
+      await testPostgresql(dedent`
+        CREATE FUNCTION my_func()
+        RETURNS TEXT
+        LANGUAGE sql
+        AS '
+          SELECT ''foo'';
+        '
+      `);
+    });
+
+    it(`reformats SQL in single-quoted SQL function`, async () => {
+      expect(
+        await pretty(
+          dedent`
+            CREATE FUNCTION my_func()
+            RETURNS TEXT
+            LANGUAGE sql
+            AS 'SELECT ''foo'';
+            select ''bar'''
+          `,
+          { dialect: "postgresql" },
+        ),
+      ).toBe(dedent`
+        CREATE FUNCTION my_func()
+        RETURNS TEXT
+        LANGUAGE sql
+        AS '
+          SELECT ''foo'';
+          SELECT ''bar'';
+        '
+      `);
+    });
   });
 
   describe("drop function", () => {