konnosDashboard
pg-meta-functions.ts390 lines · main
1import { z } from 'zod'
2
3import { DEFAULT_SYSTEM_SCHEMAS } from './constants'
4import { filterByList } from './helpers'
5import {
6 ident,
7 joinSqlFragments,
8 keyword,
9 literal,
10 safeSql,
11 type SafeSqlFragment,
12} from './pg-format'
13import { FUNCTIONS_SQL } from './sql/functions'
14
15export const pgFunctionZod = z.object({
16 id: z.number(),
17 schema: z.string(),
18 name: z.string(),
19 language: z.string(),
20 definition: z.string(),
21 complete_statement: z.string(),
22 args: z.array(
23 z.object({
24 mode: z.union([
25 z.literal('in'),
26 z.literal('out'),
27 z.literal('inout'),
28 z.literal('variadic'),
29 z.literal('table'),
30 ]),
31 name: z.string(),
32 type_id: z.number(),
33 has_default: z.boolean(),
34 })
35 ),
36 argument_types: z.string(),
37 identity_argument_types: z.string(),
38 return_type_id: z.number(),
39 return_type: z.string(),
40 return_type_relation_id: z.union([z.number(), z.null()]),
41 is_set_returning_function: z.boolean(),
42 behavior: z.union([z.literal('IMMUTABLE'), z.literal('STABLE'), z.literal('VOLATILE')]),
43 security_definer: z.boolean(),
44 config_params: z.union([z.record(z.string(), z.string()), z.null()]),
45})
46
47export type PGFunction = z.infer<typeof pgFunctionZod>
48
49export const pgFunctionArrayZod = z.array(pgFunctionZod)
50export const pgFunctionOptionalZod = z.optional(pgFunctionZod)
51
52export function list({
53 includeSystemSchemas = false,
54 includedSchemas,
55 excludedSchemas,
56 limit,
57 offset,
58}: {
59 includeSystemSchemas?: boolean
60 includedSchemas?: string[]
61 excludedSchemas?: string[]
62 limit?: number
63 offset?: number
64} = {}): {
65 sql: SafeSqlFragment
66 zod: typeof pgFunctionArrayZod
67} {
68 let sql = safeSql`
69 with f as (
70 ${FUNCTIONS_SQL}
71 )
72 select
73 f.*
74 from f
75 `
76 const filter = filterByList(
77 includedSchemas,
78 excludedSchemas,
79 !includeSystemSchemas ? DEFAULT_SYSTEM_SCHEMAS : undefined
80 )
81 if (filter) {
82 sql = safeSql`${sql} where schema ${filter}`
83 }
84 if (limit) {
85 sql = safeSql`${sql} limit ${literal(limit)}`
86 }
87 if (offset) {
88 sql = safeSql`${sql} offset ${literal(offset)}`
89 }
90
91 return {
92 sql,
93 zod: pgFunctionArrayZod,
94 }
95}
96
97type FunctionsRetrieveReturn = {
98 sql: SafeSqlFragment
99 zod: typeof pgFunctionOptionalZod
100}
101
102export function retrieve({ id }: { id: number }): FunctionsRetrieveReturn
103export function retrieve({
104 name,
105 schema,
106 args,
107}: {
108 name: string
109 schema: string
110 args: string[]
111}): FunctionsRetrieveReturn
112export function retrieve({
113 id,
114 name,
115 schema = 'public',
116 args = [],
117}: {
118 id?: number
119 name?: string
120 schema?: string
121 args?: string[]
122}): FunctionsRetrieveReturn {
123 if (id) {
124 const sql = safeSql`
125 with f as (
126 ${FUNCTIONS_SQL}
127 )
128 select
129 f.*
130 from f where id = ${literal(id)};`
131
132 return {
133 sql,
134 zod: pgFunctionOptionalZod,
135 }
136 } else if (name && schema && args) {
137 const argsFragment = args.length
138 ? safeSql`(
139 select string_agg(type_oid::text, ' ') from (
140 select (
141 split_args.arr[
142 array_length(
143 split_args.arr,
144 1
145 )
146 ]::regtype::oid
147 ) as type_oid from (
148 select string_to_array(
149 unnest(
150 array[${joinSqlFragments(args.map(literal), ',')}]
151 ),
152 ' '
153 ) as arr
154 ) as split_args
155 ) args
156 )`
157 : literal('')
158 const sql = safeSql`with f as (
159 ${FUNCTIONS_SQL}
160 )
161 select
162 f.*
163 from f join pg_proc as p on id = p.oid where schema = ${literal(schema)} and name = ${literal(name)} and p.proargtypes::text = ${argsFragment}`
164
165 return {
166 sql,
167 zod: pgFunctionOptionalZod,
168 }
169 } else {
170 throw new Error('Must provide either id or name and schema')
171 }
172}
173
174export const pgFunctionCreateZod = z.object({
175 name: z.string(),
176 definition: z.string(),
177 args: z.array(z.string()).optional(),
178 behavior: z.enum(['IMMUTABLE', 'STABLE', 'VOLATILE']).optional(),
179 config_params: z.record(z.string(), z.string()).optional(),
180 schema: z.string().optional(),
181 language: z.string().optional(),
182 return_type: z.string().optional(),
183 security_definer: z.boolean().optional(),
184})
185
186// Zod validates `args`, `return_type`, and `config_params` values as runtime
187// strings; the SafeSqlFragment brand is a separate compile-time trust check.
188// `args` items are `"name type"` fragments, `return_type` is a type
189// expression (e.g. "integer", "SETOF users"), and each `config_params` value
190// is interpolated raw after `SET <param> TO` — all three drop into the
191// CREATE FUNCTION statement uninspected, so callers must promote untrusted
192// input via acceptUntrustedSql (and ensure the promotion satisfies safety
193// requirements) before calling create(). `'FROM CURRENT'` is a sentinel for
194// the `SET <param> FROM CURRENT` grammar and is detected by string equality.
195export type PGFunctionCreate = Omit<
196 z.infer<typeof pgFunctionCreateZod>,
197 'args' | 'return_type' | 'config_params'
198> & {
199 args?: Array<SafeSqlFragment>
200 return_type?: SafeSqlFragment
201 config_params?: Record<string, SafeSqlFragment | 'FROM CURRENT'>
202}
203
204// `update()` and `remove()` reuse signature pieces from a previously-fetched
205// function. Callers must pass a value whose raw-SQL fields (`argument_types`,
206// `identity_argument_types`, `return_type`, and `config_params` values)
207// have been branded at the API/database boundary.
208export type PGSavedFunction = Omit<
209 PGFunction,
210 'argument_types' | 'identity_argument_types' | 'return_type' | 'config_params'
211> & {
212 argument_types: SafeSqlFragment
213 identity_argument_types: SafeSqlFragment
214 return_type: SafeSqlFragment
215 config_params: Record<string, SafeSqlFragment> | null
216}
217
218// PostgreSQL configuration parameters can be namespaced custom GUCs
219// (e.g. `app.jwt_secret`, `pgaudit.log`). Neither `keyword` nor `ident`
220// fits: `keyword`'s regex rejects `.`, and `ident` would quote the whole
221// value as a single identifier (`"app.jwt_secret"`), which Postgres reads
222// as one literal name rather than a qualified reference. Split on `.`,
223// `ident` each segment, and rejoin with a literal `.`.
224function qualifiedIdent(value: string): SafeSqlFragment {
225 return value
226 .split('.')
227 .map(ident)
228 .reduce<SafeSqlFragment>(
229 (acc, part, i) => (i === 0 ? part : safeSql`${acc}.${part}`),
230 safeSql``
231 )
232}
233
234function _generateCreateFunctionSql(
235 {
236 name,
237 schema,
238 args,
239 definition,
240 return_type,
241 language,
242 behavior,
243 security_definer,
244 config_params,
245 }: PGFunctionCreate,
246 { replace = false } = {}
247): SafeSqlFragment {
248 const argsFragment = args && args.length > 0 ? joinSqlFragments(args, ', ') : safeSql``
249 const configParamsFragment =
250 config_params && Object.keys(config_params).length > 0
251 ? joinSqlFragments(
252 Object.entries(config_params).map(([param, value]) =>
253 value === 'FROM CURRENT'
254 ? safeSql`SET ${qualifiedIdent(param)} FROM CURRENT`
255 : safeSql`SET ${qualifiedIdent(param)} TO ${value}`
256 ),
257 '\n'
258 )
259 : safeSql``
260
261 return safeSql`
262 CREATE ${replace ? safeSql`OR REPLACE` : safeSql``} FUNCTION ${ident(schema!)}.${ident(name!)}(${argsFragment})
263 RETURNS ${return_type!}
264 AS ${literal(definition)}
265 LANGUAGE ${keyword(language!)}
266 ${keyword(behavior!)}
267 CALLED ON NULL INPUT
268 ${security_definer ? safeSql`SECURITY DEFINER` : safeSql`SECURITY INVOKER`}
269 ${configParamsFragment};
270 `
271}
272
273export function create({
274 name,
275 schema = 'public',
276 args = [],
277 definition,
278 return_type = safeSql`void`,
279 language = 'sql',
280 behavior = 'VOLATILE',
281 security_definer = false,
282 config_params = {},
283}: PGFunctionCreate) {
284 const sql = _generateCreateFunctionSql({
285 name,
286 schema,
287 args,
288 definition,
289 return_type,
290 language,
291 behavior,
292 security_definer,
293 config_params,
294 })
295
296 return {
297 sql,
298 zod: z.void(),
299 }
300}
301
302export const pgFunctionUpdateZod = z.object({
303 name: z.string().optional(),
304 schema: z.string().optional(),
305 definition: z.string().optional(),
306})
307
308export type PGFunctionUpdate = z.infer<typeof pgFunctionUpdateZod>
309
310function splitArgumentTypes(argumentTypes: SafeSqlFragment): Array<SafeSqlFragment> {
311 return argumentTypes.split(', ') as Array<SafeSqlFragment>
312}
313
314export function update(
315 currentFunc: PGSavedFunction,
316 { name, schema, definition }: PGFunctionUpdate
317): { sql: SafeSqlFragment; zod: z.ZodType<void> } {
318 const args = splitArgumentTypes(currentFunc.argument_types)
319 const identityArgs = currentFunc.identity_argument_types
320
321 const updateDefinitionSql =
322 typeof definition === 'string'
323 ? _generateCreateFunctionSql(
324 {
325 ...currentFunc,
326 definition,
327 args,
328 config_params: currentFunc.config_params ?? {},
329 },
330 { replace: true }
331 )
332 : safeSql``
333
334 const updateNameSql =
335 name && name !== currentFunc.name
336 ? safeSql`ALTER FUNCTION ${ident(currentFunc.schema)}.${ident(currentFunc.name)}(${identityArgs}) RENAME TO ${ident(name)};`
337 : safeSql``
338
339 const updateSchemaSql =
340 schema && schema !== currentFunc.schema
341 ? safeSql`ALTER FUNCTION ${ident(currentFunc.schema)}.${ident(name || currentFunc.name)}(${identityArgs}) SET SCHEMA ${ident(schema)};`
342 : safeSql``
343
344 const sql = safeSql`
345 DO LANGUAGE plpgsql $$
346 BEGIN
347 IF ${typeof definition === 'string' ? safeSql`TRUE` : safeSql`FALSE`} THEN
348 ${updateDefinitionSql}
349
350 IF (
351 SELECT id
352 FROM (${FUNCTIONS_SQL}) AS f
353 WHERE f.schema = ${literal(currentFunc.schema)}
354 AND f.name = ${literal(currentFunc.name)}
355 AND f.identity_argument_types = ${literal(identityArgs)}
356 ) != ${literal(currentFunc.id)} THEN
357 RAISE EXCEPTION ${literal(`Cannot find function "${currentFunc.schema}"."${currentFunc.name}"(${identityArgs})`)};
358 END IF;
359 END IF;
360
361 ${updateNameSql}
362
363 ${updateSchemaSql}
364 END;
365 $$;
366 `
367
368 return {
369 sql,
370 zod: z.void(),
371 }
372}
373
374export const pgFunctionDeleteZod = z.object({
375 cascade: z.boolean().default(false).optional(),
376})
377
378export type PGFunctionDelete = z.infer<typeof pgFunctionDeleteZod>
379
380export function remove(
381 func: PGSavedFunction,
382 { cascade = false }: PGFunctionDelete = {}
383): { sql: SafeSqlFragment; zod: z.ZodType<void> } {
384 const sql = safeSql`DROP FUNCTION ${ident(func.schema)}.${ident(func.name)}(${func.identity_argument_types}) ${cascade ? safeSql`CASCADE` : safeSql`RESTRICT`};`
385
386 return {
387 sql,
388 zod: z.void(),
389 }
390}