Privileges.utils.ts360 lines · main
1import { useQueryClient } from '@tanstack/react-query'
2import { useCallback, useState } from 'react'
3
4import {
5 ALL_PRIVILEGE_TYPES,
6 COLUMN_PRIVILEGE_TYPES,
7 ColumnPrivilegeType,
8} from './Privileges.constants'
9import { grantColumnPrivileges } from '@/data/privileges/column-privileges-grant-mutation'
10import type { ColumnPrivilege } from '@/data/privileges/column-privileges-query'
11import {
12 ColumnPrivilegesRevoke,
13 revokeColumnPrivileges,
14} from '@/data/privileges/column-privileges-revoke-mutation'
15import { privilegeKeys } from '@/data/privileges/keys'
16import {
17 grantTablePrivileges,
18 TablePrivilegesGrant,
19} from '@/data/privileges/table-privileges-grant-mutation'
20import type { PgTablePrivileges } from '@/data/privileges/table-privileges-query'
21import {
22 revokeTablePrivileges,
23 TablePrivilegesRevoke,
24} from '@/data/privileges/table-privileges-revoke-mutation'
25import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
26
27export interface PrivilegeOperation {
28 object: 'table' | 'column'
29 type: 'grant' | 'revoke'
30 id: string | number
31 grantee: string
32 privilege_type: string
33}
34
35export function getDefaultTableCheckedStates(tablePrivilege: PgTablePrivileges) {
36 return Object.fromEntries(
37 ALL_PRIVILEGE_TYPES.map((privilege) => [
38 privilege,
39 tablePrivilege.privileges.find((p) => p.privilege_type === privilege) !== undefined,
40 ])
41 )
42}
43
44export function getDefaultColumnCheckedStates(columnPrivileges: ColumnPrivilege[]) {
45 return Object.fromEntries(
46 columnPrivileges.map((column) => [
47 column.column_id,
48 Object.fromEntries(
49 COLUMN_PRIVILEGE_TYPES.map((privilege) => [
50 privilege,
51 column.privileges.find((p) => p.privilege_type === privilege) !== undefined,
52 ])
53 ),
54 ])
55 )
56}
57
58interface UsePrivilegesStateOptions {
59 tableId: number
60 role: string
61 defaultTableCheckedStates: ReturnType<typeof getDefaultTableCheckedStates>
62 defaultColumnCheckedStates: ReturnType<typeof getDefaultColumnCheckedStates>
63}
64
65function addOrRemoveOperation(
66 operations: PrivilegeOperation[],
67 operation: PrivilegeOperation,
68 /** removes old operations and always adds the new one */
69 forceAdd = false
70): PrivilegeOperation[] {
71 let state = [...operations]
72
73 const oppositeType = operation.type === 'grant' ? 'revoke' : 'grant'
74
75 const existing = state.find((op) => {
76 return (
77 op.object === operation.object &&
78 op.type === oppositeType &&
79 op.id === operation.id &&
80 op.grantee === operation.grantee &&
81 op.privilege_type === operation.privilege_type
82 )
83 })
84
85 if (existing !== undefined) {
86 state = state.filter((op) => op !== existing)
87
88 if (!forceAdd) {
89 return state
90 }
91 }
92
93 state.push(operation)
94
95 return state
96}
97
98export function usePrivilegesState({
99 defaultTableCheckedStates,
100 defaultColumnCheckedStates,
101 tableId,
102 role,
103}: UsePrivilegesStateOptions) {
104 const [operations, setOperations] = useState<PrivilegeOperation[]>([])
105
106 const tableCheckedStates = operations.reduce((acc, op) => {
107 if (op.object === 'table' && op.id === tableId && op.grantee === role) {
108 return {
109 ...acc,
110 [op.privilege_type]: op.type === 'grant',
111 }
112 }
113
114 return acc
115 }, defaultTableCheckedStates)
116
117 const columnCheckedStates = operations.reduce((acc, op) => {
118 let curr = acc
119
120 if (op.object === 'table' && op.grantee === role) {
121 curr = Object.fromEntries(
122 Object.entries(curr).map(([id, column]) => [
123 id,
124 Object.fromEntries(
125 Object.entries(column).map(([privilege, value]) => [
126 privilege,
127 op.privilege_type === privilege ? op.type === 'grant' : value,
128 ])
129 ),
130 ])
131 )
132 }
133
134 if (op.object === 'column' && op.grantee === role) {
135 return {
136 ...curr,
137 [op.id]: {
138 ...curr[op.id],
139 [op.privilege_type]: op.type === 'grant',
140 },
141 }
142 }
143
144 return curr
145 }, defaultColumnCheckedStates)
146
147 function toggleTablePrivilege(privilegeType: string) {
148 const shouldGrant = !tableCheckedStates[privilegeType]
149
150 setOperations((prevState) => {
151 let state = [...prevState]
152
153 if (COLUMN_PRIVILEGE_TYPES.includes(privilegeType as ColumnPrivilegeType)) {
154 if (shouldGrant) {
155 // remove all operations for the columns since
156 // the table privilege will take precedence
157 state = state.filter(
158 (op) =>
159 !(
160 op.object === 'column' &&
161 op.grantee === role &&
162 op.privilege_type === privilegeType
163 )
164 )
165 }
166 }
167
168 state = addOrRemoveOperation(state, {
169 object: 'table',
170 type: shouldGrant ? 'grant' : 'revoke',
171 id: tableId,
172 grantee: role,
173 privilege_type: privilegeType,
174 })
175
176 return state
177 })
178 }
179
180 function toggleColumnPrivilege(columnId: string, privilegeType: string) {
181 const shouldGrant = !columnCheckedStates[columnId][privilegeType]
182
183 setOperations((prevState) => {
184 let state = [...prevState]
185
186 // if the user is revoking a column and the table is enabled
187 if (!shouldGrant && tableCheckedStates[privilegeType]) {
188 // also revoke the table privilege
189 state = addOrRemoveOperation(state, {
190 object: 'table',
191 type: 'revoke',
192 id: tableId,
193 grantee: role,
194 privilege_type: privilegeType,
195 })
196
197 // grant all other enabled columns
198 const operations = Object.entries(columnCheckedStates)
199 .filter(([id]) => id !== columnId)
200 .map(([id, column]) => ({
201 object: 'column' as const,
202 type: column[privilegeType] ? ('grant' as const) : ('revoke' as const),
203 id,
204 grantee: role,
205 privilege_type: privilegeType,
206 }))
207 operations.forEach((op) => {
208 state = addOrRemoveOperation(state, op)
209 })
210 }
211
212 if (shouldGrant) {
213 const areAllOtherColumnsEnabled = Object.entries(columnCheckedStates).every(
214 ([id, column]) => id === columnId || column[privilegeType]
215 )
216
217 if (areAllOtherColumnsEnabled) {
218 // remove all operations for the columns since
219 // the table privilege will take precedence
220 state = state.filter(
221 (op) =>
222 !(
223 op.object === 'column' &&
224 op.grantee === role &&
225 op.privilege_type === privilegeType
226 )
227 )
228
229 // grant the table privilege
230 state = addOrRemoveOperation(state, {
231 object: 'table',
232 type: 'grant',
233 id: tableId,
234 grantee: role,
235 privilege_type: privilegeType,
236 })
237
238 return state
239 }
240 }
241
242 state = addOrRemoveOperation(state, {
243 object: 'column',
244 type: shouldGrant ? 'grant' : 'revoke',
245 id: columnId,
246 grantee: role,
247 privilege_type: privilegeType,
248 })
249
250 return state
251 })
252 }
253
254 const resetOperations = useCallback(() => {
255 setOperations([])
256 }, [])
257
258 return {
259 tableCheckedStates,
260 columnCheckedStates,
261 operations,
262 toggleTablePrivilege,
263 toggleColumnPrivilege,
264 resetOperations,
265 }
266}
267
268export function useApplyPrivilegeOperations(callback?: () => void) {
269 const { data: project } = useSelectedProjectQuery()
270 const queryClient = useQueryClient()
271
272 const [isLoading, setIsLoading] = useState(false)
273
274 const apply = useCallback(
275 async (operations: PrivilegeOperation[]) => {
276 if (!project) return console.error('No project selected')
277
278 setIsLoading(true)
279
280 const tableOperations = operations.filter((op) => op.object === 'table')
281 const columnOperations = operations.filter((op) => op.object === 'column')
282
283 const grantTableOperations = tableOperations
284 .filter((op) => op.type === 'grant')
285 .map((op) => ({
286 relationId: Number(op.id),
287 grantee: op.grantee,
288 privilegeType: op.privilege_type as TablePrivilegesGrant['privilegeType'],
289 }))
290 const revokeTableOperations = tableOperations
291 .filter((op) => op.type === 'revoke')
292 .map((op) => ({
293 relationId: Number(op.id),
294 grantee: op.grantee,
295 privilegeType: op.privilege_type as TablePrivilegesRevoke['privilegeType'],
296 }))
297
298 const grantColumnOperations = columnOperations
299 .filter((op) => op.type === 'grant')
300 .map((op) => ({
301 column_id: String(op.id),
302 grantee: op.grantee,
303 privilege_type: op.privilege_type as ColumnPrivilegesRevoke['privilege_type'],
304 }))
305 const revokeColumnOperations = columnOperations
306 .filter((op) => op.type === 'revoke')
307 .map((op) => ({
308 column_id: String(op.id),
309 grantee: op.grantee,
310 privilege_type: op.privilege_type as ColumnPrivilegesRevoke['privilege_type'],
311 }))
312
313 // annoyingly these can't be run all at once
314 // as postgres can't process them in parallel
315
316 if (revokeTableOperations.length > 0) {
317 await revokeTablePrivileges({
318 projectRef: project.ref,
319 connectionString: project.connectionString,
320 revokes: revokeTableOperations,
321 })
322 }
323 if (grantTableOperations.length > 0) {
324 await grantTablePrivileges({
325 projectRef: project.ref,
326 connectionString: project.connectionString,
327 grants: grantTableOperations,
328 })
329 }
330 if (revokeColumnOperations.length > 0) {
331 await revokeColumnPrivileges({
332 projectRef: project.ref,
333 connectionString: project.connectionString,
334 revokes: revokeColumnOperations,
335 })
336 }
337 if (grantColumnOperations.length > 0) {
338 await grantColumnPrivileges({
339 projectRef: project.ref,
340 connectionString: project.connectionString,
341 grants: grantColumnOperations,
342 })
343 }
344
345 await Promise.all([
346 queryClient.invalidateQueries({ queryKey: privilegeKeys.tablePrivilegesList(project.ref) }),
347 queryClient.invalidateQueries({
348 queryKey: privilegeKeys.columnPrivilegesList(project.ref),
349 }),
350 ])
351
352 setIsLoading(false)
353
354 callback?.()
355 },
356 [callback, project, queryClient]
357 )
358
359 return { apply, isLoading }
360}