konnosDashboard
policy.ts166 lines · main
1import { generateText, Output, stepCountIs } from 'ai'
2import { IS_PLATFORM } from 'common'
3import { source } from 'common-tags'
4import { NextApiRequest, NextApiResponse } from 'next'
5import { z } from 'zod'
6
7import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
8import { getOrgAIDetails } from '@/lib/ai/ai-details'
9import { getModel } from '@/lib/ai/model'
10import { DEFAULT_COMPLETION_MODEL } from '@/lib/ai/model.utils'
11import { RLS_PROMPT } from '@/lib/ai/prompts'
12import { getTools } from '@/lib/ai/tools'
13import apiWrapper from '@/lib/api/apiWrapper'
14
15const policySchema = z.object({
16 sql: z.string().describe('The generated Postgres CREATE POLICY statement.'),
17 name: z.string().describe('The name of the policy.'),
18 command: z
19 .enum(['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'ALL'])
20 .describe('The SQL command this policy applies to.'),
21 definition: z
22 .string()
23 .optional()
24 .describe('The USING clause expression (for SELECT, UPDATE, DELETE).'),
25 check: z.string().optional().describe('The WITH CHECK clause expression (for INSERT, UPDATE).'),
26 action: z
27 .enum(['PERMISSIVE', 'RESTRICTIVE'])
28 .default('PERMISSIVE')
29 .describe('Whether the policy is PERMISSIVE or RESTRICTIVE.'),
30 roles: z.array(z.string()).default(['public']).describe('The roles this policy applies to.'),
31})
32
33const requestBodySchema = z.object({
34 tableName: z.string().min(1),
35 schema: z.string().default('public'),
36 columns: z.array(z.string()).optional(),
37 projectRef: z.string().min(1),
38 connectionString: z.string().min(1),
39 orgSlug: z.string().optional(),
40 message: z.string().optional(),
41})
42
43async function handler(req: NextApiRequest, res: NextApiResponse) {
44 const { method } = req
45
46 switch (method) {
47 case 'POST':
48 return handlePost(req, res)
49 default:
50 res.setHeader('Allow', ['POST'])
51 res.status(405).json({ data: null, error: { message: `Method ${method} Not Allowed` } })
52 }
53}
54
55export async function handlePost(req: NextApiRequest, res: NextApiResponse) {
56 const authorization = req.headers.authorization
57 const accessToken = authorization?.replace('Bearer ', '')
58
59 if (IS_PLATFORM && !accessToken) {
60 return res.status(401).json({ error: 'Authorization token is required' })
61 }
62
63 const body = typeof req.body === 'string' ? JSON.parse(req.body) : req.body
64 const { data, error: parseError } = requestBodySchema.safeParse(body)
65
66 if (parseError) {
67 return res.status(400).json({ error: 'Invalid request body', issues: parseError.issues })
68 }
69
70 const { tableName, schema, columns = [], projectRef, connectionString, orgSlug, message } = data
71
72 let aiOptInLevel: AiOptInLevel = 'disabled'
73
74 if (!IS_PLATFORM) {
75 aiOptInLevel = 'schema'
76 }
77
78 if (IS_PLATFORM && orgSlug && authorization) {
79 try {
80 const { aiOptInLevel: orgAIOptInLevel } = await getOrgAIDetails({
81 orgSlug,
82 authorization,
83 })
84
85 aiOptInLevel = orgAIOptInLevel
86 } catch (error) {
87 return res.status(400).json({
88 error: 'There was an error fetching your organization details',
89 })
90 }
91 }
92
93 try {
94 const { modelParams, error: modelError } = await getModel({
95 provider: 'openai',
96 modelEntry: DEFAULT_COMPLETION_MODEL,
97 })
98
99 if (modelError) {
100 return res.status(500).json({ error: modelError.message })
101 }
102
103 const tools = await getTools({
104 projectRef,
105 connectionString,
106 authorization,
107 aiOptInLevel,
108 accessToken,
109 })
110
111 const { experimental_output } = await generateText({
112 ...modelParams,
113 stopWhen: stepCountIs(5),
114 prompt: source`
115 You are a Postgres RLS (Row Level Security) expert.
116 Determine the most appropriate policies for the "${schema}"."${tableName}" table within a Briven project.
117
118 ${columns.length > 0 ? `Table columns: ${columns.join(', ')}` : 'No column metadata provided.'}
119
120 ${message ? `User request: ${message}` : ''}
121
122 RLS Guide: ${RLS_PROMPT}
123
124 Requirements:
125 - Use the available planning and schema tools (like "list_policies" or "list_tables") to inspect the "${schema}" schema and existing policies before generating new ones.
126 - Ensure policies strictly adhere to the existing schema
127 - Return a curated list of recommended CREATE POLICY statements as JSON.
128 - Each policy must include: name, sql, command (SELECT/INSERT/UPDATE/DELETE/ALL), action (PERMISSIVE/RESTRICTIVE), roles (array of role names).
129 - Include "definition" (USING clause expression without the USING keyword) for SELECT, UPDATE, DELETE policies.
130 - Include "check" (WITH CHECK clause expression without the WITH CHECK keywords) for INSERT, UPDATE policies.
131 - Avoid duplicating existing policies and reference the public schema and typical Briven best practices when deciding the coverage.
132 - Prefer PERMISSIVE policies unless a RESTRICTIVE policy is explicitly required
133 `,
134 tools,
135 experimental_output: Output.object({
136 schema: z.object({
137 policies: z.array(policySchema),
138 }),
139 }),
140 })
141
142 // Add table and schema to each policy from the request
143 const policies = (experimental_output?.policies ?? []).map((policy) => ({
144 ...policy,
145 table: tableName,
146 schema,
147 }))
148
149 return res.json(policies)
150 } catch (error) {
151 if (error instanceof Error) {
152 console.error(`AI policy generation failed: ${error.message}`)
153 return res.status(500).json({
154 error: 'Failed to generate policy. Please try again.',
155 })
156 }
157 return res.status(500).json({
158 error: 'An unknown error occurred.',
159 })
160 }
161}
162
163const wrapper = (req: NextApiRequest, res: NextApiResponse) =>
164 apiWrapper(req, res, handler, { withAuth: true })
165
166export default wrapper