konnosDashboard
useTestQueryRLS.ts200 lines · main
1import { type SafeSqlFragment, type UntrustedSqlFragment } from '@supabase/pg-meta'
2import { useState } from 'react'
3import { toast } from 'sonner'
4
5import { checkIfAppendLimitRequired, suffixWithLimit } from '../../SQLEditor/SQLEditor.utils'
6import { type ParseQueryResults } from './RLSTester.types'
7import { filterTablePolicies } from './useTestQueryRLS.utils'
8import { useParseClientCodeMutation } from '@/data/ai/parse-client-code-mutation'
9import { useDatabasePoliciesQuery } from '@/data/database-policies/database-policies-query'
10import { useCheckTableRLSStatusMutation } from '@/data/database/table-check-rls-mutation'
11import { useParseSQLQueryMutation } from '@/data/misc/parse-query-mutation'
12import { useExecuteSqlMutation } from '@/data/sql/execute-sql-mutation'
13import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
14import { wrapWithRoleImpersonation } from '@/lib/role-impersonation'
15import { usePostgresSandbox } from '@/state/postgres-sandbox/sandbox'
16import {
17 isRoleImpersonationEnabled,
18 useGetImpersonatedRoleState,
19 useImpersonatedUser,
20 useRoleImpersonationStateSnapshot,
21} from '@/state/role-impersonation-state'
22
23const limit = 100
24
25/**
26 * [Joshen] Testing a SQL query for its RLS access involves 3 async steps
27 * 0. (Optional) Inferring client library code to SQL query via the AI Assistant
28 * 1. Parsing the provided SQL query to retrieve its operation type + tables involved
29 * 2. Checking for tables involved if they've got RLS enabled
30 * 3. Actually running the query to retrieve the results
31 *
32 * Errors should all be handled as part of the UI instead of toasts, hence the empty onError
33 * handlers to mute the default error handlers within the react query mutationhooks
34 */
35export const useTestQueryRLS = () => {
36 const { data: project } = useSelectedProjectQuery()
37 const { role } = useRoleImpersonationStateSnapshot()
38
39 const { sandbox } = usePostgresSandbox()
40 const getImpersonatedRoleState = useGetImpersonatedRoleState()
41 const impersonatedRoleState = getImpersonatedRoleState()
42 const user = useImpersonatedUser()
43
44 const [isLoading, setIsLoading] = useState(false)
45 const [sandboxError, setSandboxError] = useState<Error>()
46
47 const { data: policies = [] } = useDatabasePoliciesQuery({
48 projectRef: project?.ref,
49 connectionString: project?.connectionString,
50 })
51
52 const { mutateAsync: executeSql, error: executeSqlMutationError } = useExecuteSqlMutation({
53 onError: () => {},
54 })
55 const executeSqlError = sandbox ? sandboxError : executeSqlMutationError
56
57 const {
58 mutateAsync: parseClientCode,
59 isPending: isInferring,
60 error: parseClientCodeError,
61 } = useParseClientCodeMutation({
62 onError: () => {},
63 })
64
65 const inferSQLFromLib = async (
66 value: string,
67 onInferSQL: (unchecked_sql: UntrustedSqlFragment) => void
68 ) => {
69 const { unchecked_sql, valid } = await parseClientCode({ code: value })
70 if (valid && unchecked_sql != null) {
71 onInferSQL(unchecked_sql)
72 } else {
73 toast.error('Client library code provided is not valid')
74 }
75 }
76
77 const { mutateAsync: parseQuery, error: parseQueryError } = useParseSQLQueryMutation({
78 onError: () => {},
79 })
80
81 const { mutateAsync: getTableRLSStatus, error: getTableRLSStatusError } =
82 useCheckTableRLSStatusMutation({
83 onError: () => {},
84 })
85
86 const testQuery = async ({
87 value,
88 option,
89 onExecuteSQL,
90 onParseQuery,
91 }: {
92 value: SafeSqlFragment
93 option: 'anon' | 'authenticated'
94 onExecuteSQL: ({
95 result,
96 isAutoLimit,
97 }: {
98 result: Object[] | null
99 isAutoLimit: boolean
100 }) => void
101 onParseQuery: (results?: ParseQueryResults) => void
102 }) => {
103 if (!project) return console.error('Project is required')
104
105 if (option === 'authenticated' && !user) {
106 return toast('Select which user to test as before running the query')
107 }
108
109 try {
110 setIsLoading(true)
111 setSandboxError(undefined)
112
113 const { appendAutoLimit } = checkIfAppendLimitRequired(value, limit)
114 const formattedSql = suffixWithLimit(value, limit)
115 const data = await parseQuery({ sql: formattedSql })
116
117 if (data.operation !== 'SELECT') {
118 return toast('Only SELECT statements are supported with the RLS Tester at the moment')
119 }
120
121 const formattedTables = data.tables.map((x) => {
122 const [schema, table] = x.includes('.') ? x.split('.') : ['public', x]
123 return { schema, table }
124 })
125 const response = await getTableRLSStatus({
126 projectRef: project?.ref,
127 connectionString: project?.connectionString,
128 tables: formattedTables,
129 })
130
131 const tables = response
132 .map(({ table, schema, rls_enabled }) => {
133 const tablePolicies = filterTablePolicies({
134 policies,
135 schema,
136 table,
137 role: role?.role,
138 operation: data.operation,
139 })
140 return {
141 table,
142 schema,
143 isRLSEnabled: rls_enabled,
144 tablePolicies,
145 }
146 })
147 .sort((a, b) => {
148 const aFirst = a.isRLSEnabled && a.tablePolicies.length === 0
149 const bFirst = b.isRLSEnabled && b.tablePolicies.length === 0
150 return Number(bFirst) - Number(aFirst)
151 })
152
153 const autoLimit = appendAutoLimit ? limit : undefined
154 const sql = wrapWithRoleImpersonation(formattedSql, impersonatedRoleState)
155
156 const { result } = sandbox
157 ? await sandbox.run({ sql }).catch((e) => {
158 setSandboxError(e instanceof Error ? e : new Error(String(e)))
159 throw e
160 })
161 : await executeSql({
162 sql,
163 autoLimit,
164 projectRef: project.ref,
165 connectionString: project.connectionString,
166 isRoleImpersonationEnabled: isRoleImpersonationEnabled(impersonatedRoleState.role),
167 isStatementTimeoutDisabled: true,
168 handleError: (e) => {
169 throw e
170 },
171 queryKey: ['rls-tester'],
172 })
173 onExecuteSQL({ result, isAutoLimit: !!autoLimit })
174
175 onParseQuery({
176 tables,
177 operation: data.operation,
178 role: role?.role,
179 user,
180 })
181 } catch (error) {
182 onExecuteSQL({ result: null, isAutoLimit: false })
183 onParseQuery(undefined)
184 } finally {
185 setIsLoading(false)
186 }
187 }
188
189 return {
190 limit,
191 testQuery,
192 inferSQLFromLib,
193 isLoading,
194 isInferring,
195 executeSqlError,
196 parseQueryError,
197 parseClientCodeError,
198 getTableRLSStatusError,
199 }
200}