konnosDashboard
sql-util.ts365 lines · main
1import { A_Const, A_Expr, ColumnRef, CreatePolicyStmt, Node, parseQuery } from 'libpg-query/wasm'
2
3export type PolicyInfo = {
4 name: string
5 relation: string
6 command?: string
7 roles: string[]
8 usingNode?: Node
9 withCheckNode?: Node
10}
11
12/**
13 * Extracts keys from a union type.
14 */
15type ExtractKeys<T> = T extends T ? keyof T : never
16
17/**
18 * Unwraps a Node to get its underlying value.
19 */
20type NodeValue<T extends Node, U extends ExtractKeys<Node>> =
21 T extends Record<U, infer V> ? V : never
22
23export class AssertionError extends Error {
24 constructor(message: string) {
25 super(message)
26
27 // Pop the top line from the stack trace so that
28 // debug tools will reference the code calling the
29 // assertion function and not this error within it
30 if (this.stack) {
31 // Capture the current stack trace and split it into lines
32 const stackLines = this.stack.split('\n')
33
34 // Remove the second line which is the current constructor
35 stackLines.splice(1, 1)
36
37 // Reassign the modified stack trace back to the error object
38 this.stack = stackLines.join('\n')
39 }
40 }
41}
42
43/**
44 * Asserts that a value is defined.
45 *
46 * Useful for type narrowing.
47 */
48export function assertDefined<T>(value: T | undefined, errorMessage: string): asserts value is T {
49 if (value === undefined) {
50 throw new AssertionError(errorMessage)
51 }
52}
53
54/**
55 * Asserts that a `Node` is a specific type.
56 *
57 * Useful for type narrowing.
58 */
59export function assertNodeType<T extends Node>(
60 node: Node,
61 type: ExtractKeys<T>,
62 errorMessage: string
63): asserts node is T {
64 if (!(type in node)) {
65 throw new AssertionError(errorMessage)
66 }
67}
68
69/**
70 * Asserts that a `Node` is a specific type and
71 * unwraps its underlying value.
72 *
73 * @returns The unwrapped `Node` value.
74 * @throws If `node` is not of type `type`.
75 */
76export function assertAndUnwrapNode<T extends Node, U extends ExtractKeys<Node>>(
77 node: Node,
78 type: U,
79 errorMessage: string
80): NodeValue<T, U> {
81 if (!(type in node)) {
82 throw new AssertionError(errorMessage)
83 }
84 return (node as any)[type]
85}
86
87/**
88 * Unwraps a `Node`'s underlying value.
89 *
90 * @returns The unwrapped `Node` value or `undefined` if
91 * the node is not of type `type`.
92 */
93export function unwrapNode<T extends Node, U extends ExtractKeys<Node>>(
94 node: Node,
95 type: U
96): NodeValue<T, U> | undefined {
97 if (!(type in node)) {
98 return undefined
99 }
100 return (node as any)[type]
101}
102
103/**
104 * Asserts that either the left or right side of the
105 * expression is processed through `fn` without throwing
106 * any errors.
107 *
108 * If both sides throw errors, the assertion fails and the
109 * error from the left side will be thrown.
110 */
111export function assertEitherSideOfExpression<U>(
112 expression: A_Expr,
113 fn: (node: Node, side: 'left' | 'right') => U
114): U {
115 assertDefined(expression.lexpr, 'Expected left side of expression to exist')
116 assertDefined(expression.rexpr, 'Expected right side of expression to exist')
117
118 try {
119 return fn(expression.lexpr, 'left')
120 } catch (leftError) {
121 try {
122 return fn(expression.rexpr, 'right')
123 } catch (rightError) {
124 throw leftError
125 }
126 }
127}
128/**
129 * Asserts that both sides of the expression are processed
130 * without throwing any errors.
131 *
132 * Order doesn't matter. As long as `firstFn` and `secondFn`
133 * pass separately on either side of the expression, the assertion
134 * will pass. Otherwise if `firstFn` and `secondFn` both fail
135 * after trying on both sides separately, the assertion will fail.
136 */
137export function assertEachSideOfExpression<U>(
138 expression: A_Expr,
139 firstFn: (node: Node) => void,
140 secondFn: (node: Node) => void
141): void {
142 assertDefined(expression.lexpr, 'Expected left side of expression to exist')
143 assertDefined(expression.rexpr, 'Expected right side of expression to exist')
144
145 let firstSide: Node
146 let secondSide: Node
147
148 try {
149 // Try `firstFn` on the left first
150 firstSide = expression.lexpr
151 secondSide = expression.rexpr
152 firstFn(firstSide)
153 } catch (firstError) {
154 // Otherwise try `firstFn` on the right
155 firstSide = expression.rexpr
156 secondSide = expression.lexpr
157 try {
158 firstFn(firstSide)
159 } catch (secondError) {
160 // `firstFn` failed on both sides, so we
161 // need to throw an error not matter what
162
163 // Perform one more test using `secondFn`
164 // to help determine which error to show
165 // for `firstFn`
166 try {
167 secondFn(secondSide)
168 } catch (_) {
169 throw firstError
170 }
171 throw secondError
172 }
173 }
174
175 try {
176 // `firstFn` passed on one of the sides, so
177 // try `secondFn` on the opposite side
178 secondFn(secondSide)
179 } catch (err) {
180 // `secondFn` failed, so we need to throw an error
181 throw err
182 }
183}
184
185/**
186 * Extracts all the `CREATE POLICY` statements
187 * from a SQL string as parsed ASTs.
188 */
189export async function getPolicies(sql: string) {
190 const result = await parseQuery(sql)
191
192 assertDefined(result.stmts, 'Expected parse result to contain statements')
193
194 return result.stmts.reduce<CreatePolicyStmt[]>((filtered, stmt) => {
195 assertDefined(stmt.stmt, 'Expected statement to exist')
196
197 const createPolicyStatement = unwrapNode(stmt.stmt, 'CreatePolicyStmt')
198
199 if (createPolicyStatement) {
200 return [...filtered, createPolicyStatement]
201 }
202
203 return filtered
204 }, [])
205}
206
207/**
208 * Parses a Postgres SQL policy.
209 *
210 * @returns Information about the policy, including its name, table, command, and expressions.
211 */
212export async function getPolicyInfo(createPolicyStatement: CreatePolicyStmt) {
213 assertDefined(createPolicyStatement.policy_name, 'Expected policy to have a name')
214 assertDefined(createPolicyStatement.table?.relname, 'Expected policy to have a relation')
215
216 const name = createPolicyStatement.policy_name
217 const relation = createPolicyStatement.table.relname
218 const command = createPolicyStatement.cmd_name
219
220 const roles =
221 createPolicyStatement.roles?.map((node) => {
222 const roleSpec = assertAndUnwrapNode(
223 node,
224 'RoleSpec',
225 'Expected roles to contain a list of RoleSpec'
226 )
227 assertDefined(roleSpec.rolename, 'Expected RoleSpec to have a rolename')
228
229 return roleSpec.rolename
230 }) ?? []
231
232 const usingExpression = createPolicyStatement.qual
233 const checkExpression = createPolicyStatement.with_check
234
235 const policyInfo: PolicyInfo = {
236 name,
237 relation,
238 command,
239 roles,
240 usingNode: usingExpression,
241 withCheckNode: checkExpression,
242 }
243
244 return policyInfo
245}
246
247export function renderTargets<T>(targets: Node[], renderTarget: (node: Node) => T): T[] {
248 return targets.map((node) => {
249 const target = assertAndUnwrapNode(
250 node,
251 'ResTarget',
252 'Expected target list to contain ResTargets'
253 )
254
255 assertDefined(target.val, 'Expected ResTarget to have a val')
256 return renderTarget(target.val)
257 })
258}
259
260export function renderColumn(column: ColumnRef) {
261 assertDefined(column.fields, 'Expected column to have fields')
262 return renderFields(column.fields)
263}
264
265export function assertAndRenderColumn(node: Node, errorMessage: string) {
266 const column = assertAndUnwrapNode(node, 'ColumnRef', errorMessage)
267 return renderColumn(column)
268}
269
270export function renderJsonExpression(expression: A_Expr): string {
271 assertDefined(expression.name, 'Expected expression to have an operator')
272
273 if (expression.name.length > 1) {
274 throw new AssertionError('Only one JSON operator supported per expression')
275 }
276
277 const [name] = expression.name
278 const operatorString = assertAndUnwrapNode(
279 name,
280 'String',
281 'Expected JSON operator to be a string'
282 )
283 assertDefined(operatorString.sval, 'PG string expected to have an sval')
284 const operator = operatorString.sval
285
286 if (!['->', '->>'].includes(operator)) {
287 throw new AssertionError(`Invalid JSON operator ${operator}`)
288 }
289
290 assertDefined(expression.lexpr, 'Expected JSON expression to have a left-side component')
291 assertDefined(expression.rexpr, 'Expected JSON expression to have a right-side component')
292
293 let left: string | number
294 let right: string | number
295
296 const leftConstant = unwrapNode(expression.lexpr, 'A_Const')
297 const leftColumn = unwrapNode(expression.lexpr, 'ColumnRef')
298 const leftFuncCall = unwrapNode(expression.lexpr, 'FuncCall')
299 const leftExpression = unwrapNode(expression.lexpr, 'A_Expr')
300
301 if (leftConstant) {
302 // JSON path cannot contain a float
303 if ('fval' in leftConstant) {
304 throw new AssertionError('Invalid JSON path: Expression cannot contain a float')
305 }
306 left = `'${parseConstant(leftConstant)}'`
307 } else if (leftColumn) {
308 left = renderColumn(leftColumn)
309 } else if (leftFuncCall) {
310 assertDefined(leftFuncCall.funcname, 'Expected function call to have a name')
311 const functionName = renderFields(leftFuncCall.funcname)
312 left = `${functionName}()`
313 } else if (leftExpression) {
314 left = renderJsonExpression(leftExpression)
315 } else {
316 throw new AssertionError('Invalid JSON path')
317 }
318
319 const rightConstant = unwrapNode(expression.rexpr, 'A_Const')
320
321 if (rightConstant) {
322 // JSON path cannot contain a float
323 if ('fval' in rightConstant) {
324 throw new AssertionError('Invalid JSON path: Expression cannot contain a float')
325 }
326 right = `'${parseConstant(rightConstant)}'`
327 } else {
328 throw new AssertionError('Invalid JSON path')
329 }
330
331 return `${left}${operator}${right}`
332}
333
334export function renderFields(fields: Node[]) {
335 const nameSegments = fields
336 .map((field) => {
337 const stringField = unwrapNode(field, 'String')
338 const starField = unwrapNode(field, 'A_Star')
339
340 if (stringField !== undefined) {
341 return stringField.sval
342 } else if (starField !== undefined) {
343 return '*'
344 } else {
345 const [internalType] = Object.keys(field)
346 throw new Error(`Unsupported internal type '${internalType}' for fields`)
347 }
348 })
349 .filter((name): name is string => name !== undefined)
350
351 return nameSegments.join('.')
352}
353
354export function parseConstant(constant: A_Const) {
355 if ('sval' in constant) {
356 return constant.sval?.sval ?? ''
357 } else if ('ival' in constant) {
358 // The PG parser turns 0 into undefined, so convert it back here
359 return constant.ival?.ival ?? 0
360 } else if ('fval' in constant) {
361 return constant.fval?.fval ? parseFloat(constant.fval.fval) : 0
362 } else {
363 throw new AssertionError(`Constant values must be a string, integer, or float`)
364 }
365}