Skip to content

Commit 9d708ad

Browse files
nenenene
committed
Format LANGUAGE SQL function/procedure bodies (#61)
2 parents e3fe34a + 5e256b3 commit 9d708ad

5 files changed

Lines changed: 252 additions & 2 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"contributors": [
66
"Barna Magyarkuti <bmagyarkuti@gmail.com>",
77
"Bowen Parnell <bparnell@4tel.com.au>",
8+
"Joel Mukuthu <joelmukuthu@gmail.com>",
89
"Rene Saarsoo <nene@triin.net>"
910
],
1011
"license": "GPL-3.0-or-later",

‎src/embed.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@ import { Printer } from "prettier";
22
import { Node } from "sql-parser-cst";
33
import { embedJs } from "./embedJs";
44
import { embedJson } from "./embedJson";
5+
import { embedSql } from "./embedSql";
56

67
export const embed: NonNullable<Printer<Node>["embed"]> = (...args) => {
7-
return embedJson(...args) || embedJs(...args);
8+
return embedJson(...args) || embedJs(...args) || embedSql(...args);
89
};

‎src/embedSql.ts‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import { Printer } from "prettier";
2+
import {
3+
CreateFunctionStmt,
4+
CreateProcedureStmt,
5+
Node,
6+
StringLiteral
7+
} from "sql-parser-cst";
8+
import {
9+
isAsClause,
10+
isCreateFunctionStmt,
11+
isCreateProcedureStmt,
12+
isLanguageClause,
13+
isStringLiteral,
14+
} from "./node_utils";
15+
import { hardline, indent, stripTrailingHardline } from "./print_utils";
16+
17+
export const embedSql: NonNullable<Printer<Node>["embed"]> = (path, options) => {
18+
const node = path.node;
19+
const parent = path.getParentNode(0);
20+
const grandParent = path.getParentNode(1);
21+
22+
if (
23+
isStringLiteral(node) &&
24+
isAsClause(parent) &&
25+
(isCreateFunctionStmt(grandParent) || isCreateProcedureStmt(grandParent)) &&
26+
grandParent.clauses.some(isSqlLanguageClause)
27+
) {
28+
return async (textToDoc) => {
29+
let quote = detectQuote(node);
30+
31+
if (!quote) {
32+
return;
33+
}
34+
35+
if (quote === "'") {
36+
// Convert `'` quotes to `$$` to simplify handling of strings inside the
37+
// function. But bail out if the function contains dollar-quoted strings.
38+
if (node.value.includes("$$")) {
39+
return;
40+
}
41+
quote = "$$";
42+
}
43+
44+
const sql = await textToDoc(node.value, options);
45+
46+
return [
47+
quote,
48+
indent([hardline, stripTrailingHardline(sql)]),
49+
hardline,
50+
quote,
51+
];
52+
};
53+
}
54+
55+
return null;
56+
};
57+
58+
const isSqlLanguageClause = (
59+
clause: CreateFunctionStmt["clauses"][0] | CreateProcedureStmt['clauses'][0],
60+
): boolean => isLanguageClause(clause) && clause.name.name.toLowerCase() === "sql";
61+
62+
const detectQuote = (
63+
node: StringLiteral,
64+
): string | undefined => {
65+
const match = node.text.match(/^('|\$[^$]*\$)/);
66+
return match?.[1];
67+
};

‎test/ddl/function.test.ts‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,101 @@ describe("function", () => {
202202
AS " return /'''|\\"\\"\\"/.test(x) "
203203
`);
204204
});
205+
206+
it(`formats dollar-quoted SQL function`, async () => {
207+
await testPostgresql(dedent`
208+
CREATE FUNCTION my_func()
209+
RETURNS INT64
210+
LANGUAGE sql
211+
AS $$
212+
SELECT 1;
213+
$$
214+
`);
215+
});
216+
217+
it(`reformats SQL in dollar-quoted SQL function`, async () => {
218+
expect(
219+
await pretty(
220+
dedent`
221+
CREATE FUNCTION my_func()
222+
RETURNS INT64
223+
LANGUAGE sql
224+
AS $body$SELECT 1;
225+
select 2$body$
226+
`,
227+
{ dialect: "postgresql" },
228+
),
229+
).toBe(dedent`
230+
CREATE FUNCTION my_func()
231+
RETURNS INT64
232+
LANGUAGE sql
233+
AS $body$
234+
SELECT 1;
235+
SELECT 2;
236+
$body$
237+
`);
238+
});
239+
240+
it(`converts single-quoted SQL functions to dollar-quoted SQL functions`, async () => {
241+
expect(
242+
await pretty(
243+
dedent`
244+
CREATE FUNCTION my_func()
245+
RETURNS TEXT
246+
LANGUAGE sql
247+
AS 'SELECT ''foo'''
248+
`,
249+
{ dialect: "postgresql" },
250+
),
251+
).toBe(dedent`
252+
CREATE FUNCTION my_func()
253+
RETURNS TEXT
254+
LANGUAGE sql
255+
AS $$
256+
SELECT 'foo';
257+
$$
258+
`);
259+
});
260+
261+
it(`does not convert single-quoted SQL functions to dollar-quoted SQL functions when they contain dollar-quoted strings`, async () => {
262+
expect(
263+
await pretty(
264+
dedent`
265+
CREATE FUNCTION my_func()
266+
RETURNS TEXT
267+
LANGUAGE sql
268+
AS 'SELECT $$foo$$'
269+
`,
270+
{ dialect: "postgresql" },
271+
),
272+
).toBe(dedent`
273+
CREATE FUNCTION my_func()
274+
RETURNS TEXT
275+
LANGUAGE sql
276+
AS 'SELECT $$foo$$'
277+
`);
278+
});
279+
280+
it(`handles SQL language identifier case-insensitively`, async () => {
281+
expect(
282+
await pretty(
283+
dedent`
284+
CREATE FUNCTION my_func()
285+
RETURNS INT64
286+
LANGUAGE Sql
287+
AS 'SELECT 1'
288+
`,
289+
{ dialect: "postgresql" },
290+
),
291+
).toBe(dedent`
292+
CREATE FUNCTION my_func()
293+
RETURNS INT64
294+
LANGUAGE Sql
295+
AS $$
296+
SELECT 1;
297+
$$
298+
`);
299+
});
205300
});
206301

207302
describe("drop function", () => {

‎test/ddl/procedure.test.ts‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import dedent from "dedent-js";
2-
import { testBigquery, testPostgresql } from "../test_utils";
2+
import { pretty, testBigquery, testPostgresql } from "../test_utils";
33

44
describe("procedure", () => {
55
describe("create procedure", () => {
@@ -90,6 +90,92 @@ describe("procedure", () => {
9090
`,
9191
);
9292
});
93+
94+
it(`formats dollar-quoted SQL procedure`, async () => {
95+
await testPostgresql(dedent`
96+
CREATE PROCEDURE my_proc()
97+
LANGUAGE sql
98+
AS $$
99+
SELECT 1;
100+
$$
101+
`);
102+
});
103+
104+
it(`reformats SQL in dollar-quoted SQL procedure`, async () => {
105+
expect(
106+
await pretty(
107+
dedent`
108+
CREATE PROCEDURE my_proc()
109+
LANGUAGE sql
110+
AS $body$SELECT 1;
111+
select 2$body$
112+
`,
113+
{ dialect: "postgresql" },
114+
),
115+
).toBe(dedent`
116+
CREATE PROCEDURE my_proc()
117+
LANGUAGE sql
118+
AS $body$
119+
SELECT 1;
120+
SELECT 2;
121+
$body$
122+
`);
123+
});
124+
125+
it(`converts single-quoted SQL procedures to dollar-quoted SQL procedures`, async () => {
126+
expect(
127+
await pretty(
128+
dedent`
129+
CREATE PROCEDURE my_proc()
130+
LANGUAGE sql
131+
AS 'SELECT ''foo'''
132+
`,
133+
{ dialect: "postgresql" },
134+
),
135+
).toBe(dedent`
136+
CREATE PROCEDURE my_proc()
137+
LANGUAGE sql
138+
AS $$
139+
SELECT 'foo';
140+
$$
141+
`);
142+
});
143+
144+
it(`does not convert single-quoted SQL procedures to dollar-quoted SQL procedures when they contain dollar-quoted strings`, async () => {
145+
expect(
146+
await pretty(
147+
dedent`
148+
CREATE PROCEDURE my_proc()
149+
LANGUAGE sql
150+
AS 'SELECT $$foo$$'
151+
`,
152+
{ dialect: "postgresql" },
153+
),
154+
).toBe(dedent`
155+
CREATE PROCEDURE my_proc()
156+
LANGUAGE sql
157+
AS 'SELECT $$foo$$'
158+
`);
159+
});
160+
161+
it(`handles SQL language identifier case-insensitively`, async () => {
162+
expect(
163+
await pretty(
164+
dedent`
165+
CREATE PROCEDURE my_proc()
166+
LANGUAGE Sql
167+
AS 'SELECT 1'
168+
`,
169+
{ dialect: "postgresql" },
170+
),
171+
).toBe(dedent`
172+
CREATE PROCEDURE my_proc()
173+
LANGUAGE Sql
174+
AS $$
175+
SELECT 1;
176+
$$
177+
`);
178+
});
93179
});
94180

95181
describe("drop procedure", () => {

0 commit comments

Comments
 (0)