From de43d9992aa166010085a27cfba0780f6d734837 Mon Sep 17 00:00:00 2001 From: Greg Richardson Date: Sun, 16 Jun 2024 20:39:59 -0600 Subject: [PATCH] feat(db-playground): seeds tab --- .../database-playground/app/api/chat/route.ts | 5 +- apps/database-playground/app/page.tsx | 97 ++++- apps/database-playground/lib/sql-util.ts | 193 ++++++++++ apps/database-playground/package.json | 4 +- package-lock.json | 7 +- packages/ai-commands/package.json | 3 +- packages/ai-commands/src/sql/rls.test.ts | 7 +- packages/ai-commands/test/sql-util.ts | 307 +-------------- packages/common/package.json | 5 +- packages/common/sql-util.ts | 352 ++++++++++++++++++ 10 files changed, 651 insertions(+), 329 deletions(-) create mode 100644 apps/database-playground/lib/sql-util.ts create mode 100644 packages/common/sql-util.ts diff --git a/apps/database-playground/app/api/chat/route.ts b/apps/database-playground/app/api/chat/route.ts index d739263299d..444d1331eb3 100644 --- a/apps/database-playground/app/api/chat/route.ts +++ b/apps/database-playground/app/api/chat/route.ts @@ -9,6 +9,9 @@ export const maxDuration = 30 const reportSchema = z.object({ name: z.string(), description: z.string() }) export type Report = typeof reportSchema._type +export const tabsSchema = z.enum(['diagram', 'migrations', 'seeds']) +export type TabValue = typeof tabsSchema._type + export async function POST(req: Request) { const { messages } = await req.json() @@ -80,7 +83,7 @@ export async function POST(req: Request) { Switches to a different tab. `, parameters: z.object({ - tab: z.enum(['diagram', 'migrations']), + tab: tabsSchema, }), }, }, diff --git a/apps/database-playground/app/page.tsx b/apps/database-playground/app/page.tsx index d2b92373512..f0ab1b3f234 100644 --- a/apps/database-playground/app/page.tsx +++ b/apps/database-playground/app/page.tsx @@ -8,7 +8,9 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from '@ui/components/shadcn/ import { nanoid } from 'ai' import { useChat } from 'ai/react' import Chart from 'chart.js/auto' +import { assertDefined } from 'common/sql-util' import { AnimatePresence, LazyMotion, m } from 'framer-motion' +import { parseQuery } from 'libpg-query/wasm' import { throttle } from 'lodash' import { ArrowUp, Square } from 'lucide-react' import { useCallback, useMemo, useRef, useState } from 'react' @@ -24,6 +26,8 @@ import { Button } from 'ui/src/components/shadcn/ui/button' import SchemaGraph from '~/components/schema/graph' import { useTablesQuery } from '~/data/tables/tables-query' import { useReportSuggestions } from '~/lib/hooks' +import { groupStatements } from '~/lib/sql-util' +import { TabValue, tabsSchema } from './api/chat/route' const loadFramerFeatures = () => import('./framer-features').then((res) => res.default) @@ -35,13 +39,14 @@ export default function Page() { const [migrationSql, setMigrationSql] = useState( '-- Migrations will appear here as you chat with Supabase AI' ) - const [tab, setTab] = useState('diagram') + const [seedSql, setSeedSql] = useState('-- Seeds will appear here as you chat with Supabase AI') + const [tab, setTab] = useState('diagram') const [brainstormIdeas] = useState(false) const { reports } = useReportSuggestions(db, { enabled: brainstormIdeas }) const { data: tables, refetch } = useTablesQuery({ schemas: ['public'], includeColumns: true }) - const { messages, input, setInput, handleInputChange, append, stop, isLoading } = useChat({ + const { messages, input, setInput, handleInputChange, append, stop, isLoading, error } = useChat({ api: 'api/chat', maxToolRoundtrips: 10, // Provide the LLM with the current schema before the chat starts @@ -82,19 +87,43 @@ export default function Page() { case 'executeSql': { try { const { sql } = toolCall.args as any + console.log(sql) + + const parseResult = await parseQuery(sql) + + assertDefined(parseResult.stmts, 'Expected parse result to contain statements') + + const { seeds, migrations } = groupStatements(parseResult.stmts) + const results = await db.exec(sql) - setMigrationSql((s) => { - const newSql = (s + '\n' + sql).trim() - return format(newSql, { - language: 'postgresql', - keywordCase: 'lower', - identifierCase: 'lower', - dataTypeCase: 'lower', - functionCase: 'lower', + // TODO: use libpg-query de-parser once released + // This assumes every statement is a seed or migration, + // which might not be true + if (seeds.length > 0) { + setSeedSql((s) => { + const newSql = (s + '\n' + sql).trim() + return format(newSql, { + language: 'postgresql', + keywordCase: 'lower', + identifierCase: 'lower', + dataTypeCase: 'lower', + functionCase: 'lower', + }) }) - }) + } else if (migrations.length > 0) { + setMigrationSql((s) => { + const newSql = (s + '\n' + sql).trim() + return format(newSql, { + language: 'postgresql', + keywordCase: 'lower', + identifierCase: 'lower', + dataTypeCase: 'lower', + functionCase: 'lower', + }) + }) + } const { data: tables, error } = await refetch() @@ -204,11 +233,12 @@ export default function Page() { setTab(tab)} + onValueChange={(tab) => setTab(tabsSchema.parse(tab))} > - + Diagram Migrations + Seeds @@ -249,6 +279,47 @@ export default function Page() { await editor.getAction('editor.action.formatDocument').run() }) + // Run format on the initial value + await editor.getAction('editor.action.formatDocument').run() + }} + /> + + + { + // Register pgsql formatter + monaco.languages.registerDocumentFormattingEditProvider('pgsql', { + async provideDocumentFormattingEdits(model) { + const currentCode = editor.getValue() + const formattedCode = format(currentCode, { + language: 'postgresql', + keywordCase: 'lower', + }) + return [ + { + range: model.getFullModelRange(), + text: formattedCode, + }, + ] + }, + }) + + // Format on cmd+s + editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyS, async () => { + await editor.getAction('editor.action.formatDocument').run() + }) + // Run format on the initial value await editor.getAction('editor.action.formatDocument').run() }} diff --git a/apps/database-playground/lib/sql-util.ts b/apps/database-playground/lib/sql-util.ts new file mode 100644 index 00000000000..32ba4d83d15 --- /dev/null +++ b/apps/database-playground/lib/sql-util.ts @@ -0,0 +1,193 @@ +import { Statement, assertDefined, unwrapNode, unwrapOneOfNode } from 'common/sql-util' +import { Node, RawStmt } from 'libpg-query/wasm' + +/** + * Groups SQL statements into queries, seeds, and migrations + * based on heuristics. + */ +export function groupStatements(rawStatements: RawStmt[]) { + return rawStatements.reduce<{ + queries: Statement[] + seeds: Statement[] + migrations: Statement[] + }>( + (merged, stmt) => { + assertDefined(stmt.stmt, 'Expected statement to exist') + + const queryStatement = unwrapQueryStatement(stmt.stmt) + + if (queryStatement) { + return { + ...merged, + queries: [...merged.queries, queryStatement], + } + } + + const seedStatement = unwrapSeedStatement(stmt.stmt) + + if (seedStatement) { + return { + ...merged, + seeds: [...merged.seeds, seedStatement], + } + } + + const migrationStatement = unwrapMigrationStatement(stmt.stmt) + + if (migrationStatement) { + return { + ...merged, + migrations: [...merged.migrations, stmt.stmt], + } + } + + throw new Error('Unsupported SQL statement') + }, + { queries: [], seeds: [], migrations: [] } + ) +} + +export function unwrapSeedStatement(node: Node) { + // Insert, update, and delete statements are always considered seeds + const statement = unwrapOneOfNode(node, ['InsertStmt', 'UpdateStmt', 'DeleteStmt', 'MergeStmt']) + + if (statement) { + return statement + } + + // Select statements are considered seeds if they insert into another table + const selectStatement = unwrapNode(node, 'SelectStmt') + + if (selectStatement?.intoClause) { + return selectStatement + } +} + +export function unwrapQueryStatement(node: Node) { + // Select statements are considered queries if they don't insert into another table + const selectStatement = unwrapNode(node, 'SelectStmt') + + if (selectStatement && !selectStatement.intoClause) { + return selectStatement + } +} + +export function unwrapMigrationStatement(node: Node) { + const statement = unwrapOneOfNode(node, [ + 'ReturnStmt', + 'PLAssignStmt', + 'AlterTableStmt', + 'AlterDomainStmt', + 'SetOperationStmt', + 'GrantStmt', + 'GrantRoleStmt', + 'AlterDefaultPrivilegesStmt', + 'ClosePortalStmt', + 'ClusterStmt', + 'CopyStmt', + 'CreateStmt', + 'DefineStmt', + 'DropStmt', + 'TruncateStmt', + 'CommentStmt', + 'FetchStmt', + 'IndexStmt', + 'CreateFunctionStmt', + 'AlterFunctionStmt', + 'DoStmt', + 'RenameStmt', + 'RuleStmt', + 'NotifyStmt', + 'ListenStmt', + 'UnlistenStmt', + 'TransactionStmt', + 'ViewStmt', + 'LoadStmt', + 'CreateDomainStmt', + 'CreatedbStmt', + 'DropdbStmt', + 'VacuumStmt', + 'ExplainStmt', + 'CreateTableAsStmt', + 'CreateSeqStmt', + 'AlterSeqStmt', + 'VariableSetStmt', + 'VariableShowStmt', + 'DiscardStmt', + 'CreateTrigStmt', + 'CreatePLangStmt', + 'CreateRoleStmt', + 'AlterRoleStmt', + 'DropRoleStmt', + 'LockStmt', + 'ConstraintsSetStmt', + 'ReindexStmt', + 'CheckPointStmt', + 'CreateSchemaStmt', + 'AlterDatabaseStmt', + 'AlterDatabaseRefreshCollStmt', + 'AlterDatabaseSetStmt', + 'AlterRoleSetStmt', + 'CreateConversionStmt', + 'CreateCastStmt', + 'CreateOpClassStmt', + 'CreateOpFamilyStmt', + 'AlterOpFamilyStmt', + 'PrepareStmt', + 'ExecuteStmt', + 'DeallocateStmt', + 'DeclareCursorStmt', + 'CreateTableSpaceStmt', + 'DropTableSpaceStmt', + 'AlterObjectDependsStmt', + 'AlterObjectSchemaStmt', + 'AlterOwnerStmt', + 'AlterOperatorStmt', + 'AlterTypeStmt', + 'DropOwnedStmt', + 'ReassignOwnedStmt', + 'CompositeTypeStmt', + 'CreateEnumStmt', + 'CreateRangeStmt', + 'AlterEnumStmt', + 'AlterTSDictionaryStmt', + 'AlterTSConfigurationStmt', + 'CreateFdwStmt', + 'AlterFdwStmt', + 'CreateForeignServerStmt', + 'AlterForeignServerStmt', + 'CreateUserMappingStmt', + 'AlterUserMappingStmt', + 'DropUserMappingStmt', + 'AlterTableSpaceOptionsStmt', + 'AlterTableMoveAllStmt', + 'SecLabelStmt', + 'CreateForeignTableStmt', + 'ImportForeignSchemaStmt', + 'CreateExtensionStmt', + 'AlterExtensionStmt', + 'AlterExtensionContentsStmt', + 'CreateEventTrigStmt', + 'AlterEventTrigStmt', + 'RefreshMatViewStmt', + 'ReplicaIdentityStmt', + 'AlterSystemStmt', + 'CreatePolicyStmt', + 'AlterPolicyStmt', + 'CreateTransformStmt', + 'CreateAmStmt', + 'CreatePublicationStmt', + 'AlterPublicationStmt', + 'CreateSubscriptionStmt', + 'AlterSubscriptionStmt', + 'DropSubscriptionStmt', + 'CreateStatsStmt', + 'AlterCollationStmt', + 'CallStmt', + 'AlterStatsStmt', + ]) + + if (statement) { + return statement + } +} diff --git a/apps/database-playground/package.json b/apps/database-playground/package.json index dd3a627ed63..c33f9e8a890 100644 --- a/apps/database-playground/package.json +++ b/apps/database-playground/package.json @@ -19,10 +19,12 @@ "ai-sandbox": "*", "chart.js": "^4.4.3", "chartjs-adapter-date-fns": "^3.0.0", + "common": "*", "config": "*", "date-fns": "^3.6.0", "framer-motion": "^11.2.10", "katex": "^0.16.10", + "libpg-query": "^15.2.0", "monaco-editor": "^0.49.0", "next": "14.2.3", "react": "^18", @@ -49,4 +51,4 @@ "tailwindcss": "^3.4.1", "typescript": "^5" } -} +} \ No newline at end of file diff --git a/package-lock.json b/package-lock.json index c4343a4e46d..142453ef82c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -210,10 +210,12 @@ "ai-sandbox": "*", "chart.js": "^4.4.3", "chartjs-adapter-date-fns": "^3.0.0", + "common": "*", "config": "*", "date-fns": "^3.6.0", "framer-motion": "^11.2.10", "katex": "^0.16.10", + "libpg-query": "^15.2.0", "monaco-editor": "^0.49.0", "next": "14.2.3", "react": "^18", @@ -42470,7 +42472,6 @@ "version": "1.0.2", "resolved": "https://registry.npmjs.org/minipass-collect/-/minipass-collect-1.0.2.tgz", "integrity": "sha512-6T6lH0H8OG9kITm/Jm6tdooIbogG9e0tLgpY6mphXSm/A9u8Nq1ryBG+Qspiub9LjWlBPsPS3tWQ/Botq4FdxA==", - "dev": true, "dependencies": { "minipass": "^3.0.0" }, @@ -53749,6 +53750,7 @@ "@serafin/schema-builder": "^0.18.5", "@supabase/supabase-js": "*", "ai": "^2.2.29", + "common": "*", "common-tags": "^1.8.2", "config": "*", "js-tiktoken": "^1.0.10", @@ -54138,6 +54140,7 @@ "devDependencies": { "@types/react": "^18.2.24", "@types/react-dom": "^18.2.8", + "libpg-query": "^15.2.0", "tsconfig": "*", "typescript": "^5.4.3" }, @@ -55534,4 +55537,4 @@ "dev": true } } -} \ No newline at end of file +} diff --git a/packages/ai-commands/package.json b/packages/ai-commands/package.json index c6fd2214ea2..3efbbeabebb 100644 --- a/packages/ai-commands/package.json +++ b/packages/ai-commands/package.json @@ -13,6 +13,7 @@ "@serafin/schema-builder": "^0.18.5", "@supabase/supabase-js": "*", "ai": "^2.2.29", + "common": "*", "common-tags": "^1.8.2", "config": "*", "js-tiktoken": "^1.0.10", @@ -36,4 +37,4 @@ "tsconfig": "*", "typescript": "^5.4.3" } -} +} \ No newline at end of file diff --git a/packages/ai-commands/src/sql/rls.test.ts b/packages/ai-commands/src/sql/rls.test.ts index 93dda53c743..86ab6517309 100644 --- a/packages/ai-commands/src/sql/rls.test.ts +++ b/packages/ai-commands/src/sql/rls.test.ts @@ -1,6 +1,5 @@ import { describe, expect, test } from '@jest/globals' import { codeBlock } from 'common-tags' -import OpenAI from 'openai' import { assertAndRenderColumn, assertAndUnwrapNode, @@ -8,14 +7,14 @@ import { assertEachSideOfExpression, assertEitherSideOfExpression, assertNodeType, - getPolicies, - getPolicyInfo, parseConstant, renderFields, renderJsonExpression, renderTargets, unwrapNode, -} from '../../test/sql-util' +} from 'common/sql-util' +import OpenAI from 'openai' +import { getPolicies, getPolicyInfo } from '../../test/sql-util' import { collectStream, extractMarkdownSql, withMetadata } from '../../test/util' import { chatRlsPolicy } from './rls' diff --git a/packages/ai-commands/test/sql-util.ts b/packages/ai-commands/test/sql-util.ts index 79fae6a4795..c715f9acff5 100644 --- a/packages/ai-commands/test/sql-util.ts +++ b/packages/ai-commands/test/sql-util.ts @@ -1,188 +1,5 @@ -import { A_Const, A_Expr, ColumnRef, CreatePolicyStmt, Node, parseQuery } from 'libpg-query' - -export type PolicyInfo = { - name: string - relation: string - command?: string - roles: string[] - usingNode?: Node - withCheckNode?: Node -} - -/** - * Extracts keys from a union type. - */ -type ExtractKeys = T extends T ? keyof T : never - -/** - * Unwraps a Node to get its underlying value. - */ -type NodeValue> = T extends Record - ? V - : never - -export class AssertionError extends Error { - constructor(message: string) { - super(message) - - // Pop the top line from the stack trace so that - // debug tools will reference the code calling the - // assertion function and not this error within it - if (this.stack) { - // Capture the current stack trace and split it into lines - const stackLines = this.stack.split('\n') - - // Remove the second line which is the current constructor - stackLines.splice(1, 1) - - // Reassign the modified stack trace back to the error object - this.stack = stackLines.join('\n') - } - } -} - -/** - * Asserts that a value is defined. - * - * Useful for type narrowing. - */ -export function assertDefined(value: T | undefined, errorMessage: string): asserts value is T { - if (value === undefined) { - throw new AssertionError(errorMessage) - } -} - -/** - * Asserts that a `Node` is a specific type. - * - * Useful for type narrowing. - */ -export function assertNodeType( - node: Node, - type: ExtractKeys, - errorMessage: string -): asserts node is T { - if (!(type in node)) { - throw new AssertionError(errorMessage) - } -} - -/** - * Asserts that a `Node` is a specific type and - * unwraps its underlying value. - * - * @returns The unwrapped `Node` value. - * @throws If `node` is not of type `type`. - */ -export function assertAndUnwrapNode>( - node: Node, - type: U, - errorMessage: string -): NodeValue { - if (!(type in node)) { - throw new AssertionError(errorMessage) - } - return (node as any)[type] -} - -/** - * Unwraps a `Node`'s underlying value. - * - * @returns The unwrapped `Node` value or `undefined` if - * the node is not of type `type`. - */ -export function unwrapNode>( - node: Node, - type: U -): NodeValue | undefined { - if (!(type in node)) { - return undefined - } - return (node as any)[type] -} - -/** - * Asserts that either the left or right side of the - * expression is processed through `fn` without throwing - * any errors. - * - * If both sides throw errors, the assertion fails and the - * error from the left side will be thrown. - */ -export function assertEitherSideOfExpression( - expression: A_Expr, - fn: (node: Node, side: 'left' | 'right') => U -): U { - assertDefined(expression.lexpr, 'Expected left side of expression to exist') - assertDefined(expression.rexpr, 'Expected right side of expression to exist') - - try { - return fn(expression.lexpr, 'left') - } catch (leftError) { - try { - return fn(expression.rexpr, 'right') - } catch (rightError) { - throw leftError - } - } -} -/** - * Asserts that both sides of the expression are processed - * without throwing any errors. - * - * Order doesn't matter. As long as `firstFn` and `secondFn` - * pass separately on either side of the expression, the assertion - * will pass. Otherwise if `firstFn` and `secondFn` both fail - * after trying on both sides separately, the assertion will fail. - */ -export function assertEachSideOfExpression( - expression: A_Expr, - firstFn: (node: Node) => void, - secondFn: (node: Node) => void -): void { - assertDefined(expression.lexpr, 'Expected left side of expression to exist') - assertDefined(expression.rexpr, 'Expected right side of expression to exist') - - let firstSide: Node - let secondSide: Node - - try { - // Try `firstFn` on the left first - firstSide = expression.lexpr - secondSide = expression.rexpr - firstFn(firstSide) - } catch (firstError) { - // Otherwise try `firstFn` on the right - firstSide = expression.rexpr - secondSide = expression.lexpr - try { - firstFn(firstSide) - } catch (secondError) { - // `firstFn` failed on both sides, so we - // need to throw an error not matter what - - // Perform one more test using `secondFn` - // to help determine which error to show - // for `firstFn` - try { - secondFn(secondSide) - } catch (_) { - throw firstError - } - throw secondError - } - } - - try { - // `firstFn` passed on one of the sides, so - // try `secondFn` on the opposite side - secondFn(secondSide) - } catch (err) { - // `secondFn` failed, so we need to throw an error - throw err - } -} - +import { PolicyInfo, assertAndUnwrapNode, assertDefined, unwrapNode } from 'common/sql-util' +import { CreatePolicyStmt, parseQuery } from 'libpg-query' /** * Extracts all the `CREATE POLICY` statements * from a SQL string as parsed ASTs. @@ -244,123 +61,3 @@ export async function getPolicyInfo(createPolicyStatement: CreatePolicyStmt) { return policyInfo } - -export function renderTargets(targets: Node[], renderTarget: (node: Node) => T): T[] { - return targets.map((node) => { - const target = assertAndUnwrapNode( - node, - 'ResTarget', - 'Expected target list to contain ResTargets' - ) - - assertDefined(target.val, 'Expected ResTarget to have a val') - return renderTarget(target.val) - }) -} - -export function renderColumn(column: ColumnRef) { - assertDefined(column.fields, 'Expected column to have fields') - return renderFields(column.fields) -} - -export function assertAndRenderColumn(node: Node, errorMessage: string) { - const column = assertAndUnwrapNode(node, 'ColumnRef', errorMessage) - return renderColumn(column) -} - -export function renderJsonExpression(expression: A_Expr): string { - assertDefined(expression.name, 'Expected expression to have an operator') - - if (expression.name.length > 1) { - throw new AssertionError('Only one JSON operator supported per expression') - } - - const [name] = expression.name - const operatorString = assertAndUnwrapNode( - name, - 'String', - 'Expected JSON operator to be a string' - ) - assertDefined(operatorString.sval, 'PG string expected to have an sval') - const operator = operatorString.sval - - if (!['->', '->>'].includes(operator)) { - throw new AssertionError(`Invalid JSON operator ${operator}`) - } - - assertDefined(expression.lexpr, 'Expected JSON expression to have a left-side component') - assertDefined(expression.rexpr, 'Expected JSON expression to have a right-side component') - - let left: string | number - let right: string | number - - const leftConstant = unwrapNode(expression.lexpr, 'A_Const') - const leftColumn = unwrapNode(expression.lexpr, 'ColumnRef') - const leftFuncCall = unwrapNode(expression.lexpr, 'FuncCall') - const leftExpression = unwrapNode(expression.lexpr, 'A_Expr') - - if (leftConstant) { - // JSON path cannot contain a float - if ('fval' in leftConstant) { - throw new AssertionError('Invalid JSON path: Expression cannot contain a float') - } - left = `'${parseConstant(leftConstant)}'` - } else if (leftColumn) { - left = renderColumn(leftColumn) - } else if (leftFuncCall) { - assertDefined(leftFuncCall.funcname, 'Expected function call to have a name') - const functionName = renderFields(leftFuncCall.funcname) - left = `${functionName}()` - } else if (leftExpression) { - left = renderJsonExpression(leftExpression) - } else { - throw new AssertionError('Invalid JSON path') - } - - const rightConstant = unwrapNode(expression.rexpr, 'A_Const') - - if (rightConstant) { - // JSON path cannot contain a float - if ('fval' in rightConstant) { - throw new AssertionError('Invalid JSON path: Expression cannot contain a float') - } - right = `'${parseConstant(rightConstant)}'` - } else { - throw new AssertionError('Invalid JSON path') - } - - return `${left}${operator}${right}` -} - -export function renderFields(fields: Node[]) { - const nameSegments = fields - .map((field) => { - const stringField = unwrapNode(field, 'String') - const starField = unwrapNode(field, 'A_Star') - - if (stringField !== undefined) { - return stringField.sval - } else if (starField !== undefined) { - return '*' - } else { - const [internalType] = Object.keys(field) - throw new Error(`Unsupported internal type '${internalType}' for fields`) - } - }) - .filter((name): name is string => name !== undefined) - - return nameSegments.join('.') -} - -export function parseConstant(constant: A_Const) { - if ('sval' in constant) { - return constant.sval?.sval ?? '' - } else if ('ival' in constant) { - // The PG parser turns 0 into undefined, so convert it back here - return constant.ival?.ival ?? 0 - } else if ('fval' in constant) { - return constant.fval?.fval ? parseFloat(constant.fval.fval) : 0 - } else { - throw new AssertionError(`Constant values must be a string, integer, or float`) - } -} diff --git a/packages/common/package.json b/packages/common/package.json index 3765018aff6..1071cc0c29d 100644 --- a/packages/common/package.json +++ b/packages/common/package.json @@ -19,14 +19,15 @@ "devDependencies": { "@types/react": "^18.2.24", "@types/react-dom": "^18.2.8", + "libpg-query": "^15.2.0", "tsconfig": "*", "typescript": "^5.4.3" }, "peerDependencies": { - "@supabase/supabase-js": "*", "@supabase/auth-js": "^2.63.0", + "@supabase/supabase-js": "*", "next": "*", "react": "*", "react-dom": "*" } -} +} \ No newline at end of file diff --git a/packages/common/sql-util.ts b/packages/common/sql-util.ts new file mode 100644 index 00000000000..f5698cc1615 --- /dev/null +++ b/packages/common/sql-util.ts @@ -0,0 +1,352 @@ +import { A_Const, A_Expr, ColumnRef, Node } from 'libpg-query' + +export type PolicyInfo = { + name: string + relation: string + command?: string + roles: string[] + usingNode?: Node + withCheckNode?: Node +} + +/** + * Extracts keys from a union type. + */ +export type ExtractKeys = T extends T ? keyof T : never + +/** + * Filters strings ending in the specified suffix. + */ +export type StringWithSuffix = T extends `${infer _Prefix}${S}` + ? T + : never + +/** + * Extracts all nodes who's keys end in the specified suffix. + */ +export type NodeWithSuffix = Node extends infer U + ? U extends { + [K in keyof U]: K extends `${infer _Prefix}${S}` ? U : never + }[keyof U] + ? U[keyof U] + : never + : never + +/** + * Every Node type. + */ +export type NodeType = ExtractKeys + +/** + * Every statement Node type. + */ +export type StatementType = StringWithSuffix + +/** + * Every statement Node. + */ +export type Statement = NodeWithSuffix<'Stmt'> + +/** + * Unwraps a Node to get its underlying value. + */ +export type NodeValue = T extends Record ? V : never + +export class AssertionError extends Error { + constructor(message: string) { + super(message) + + // Pop the top line from the stack trace so that + // debug tools will reference the code calling the + // assertion function and not this error within it + if (this.stack) { + // Capture the current stack trace and split it into lines + const stackLines = this.stack.split('\n') + + // Remove the second line which is the current constructor + stackLines.splice(1, 1) + + // Reassign the modified stack trace back to the error object + this.stack = stackLines.join('\n') + } + } +} + +/** + * Asserts that a value is defined. + * + * Useful for type narrowing. + */ +export function assertDefined(value: T | undefined, errorMessage: string): asserts value is T { + if (value === undefined) { + throw new AssertionError(errorMessage) + } +} + +/** + * Asserts that a `Node` is a specific type. + * + * Useful for type narrowing. + */ +export function assertNodeType( + node: Node, + type: ExtractKeys, + errorMessage: string +): asserts node is T { + if (!(type in node)) { + throw new AssertionError(errorMessage) + } +} + +/** + * Asserts that a `Node` is a specific type and + * unwraps its underlying value. + * + * @returns The unwrapped `Node` value. + * @throws If `node` is not of type `type`. + */ +export function assertAndUnwrapNode( + node: Node, + type: U, + errorMessage: string +): NodeValue { + if (!(type in node)) { + throw new AssertionError(errorMessage) + } + return (node as any)[type] +} + +/** + * Unwraps a `Node`'s underlying value. + * + * @returns The unwrapped `Node` value or `undefined` if + * the node is not of type `type`. + */ +export function unwrapNode( + node: Node, + type: U +): NodeValue | undefined { + if (!(type in node)) { + return undefined + } + return (node as any)[type] +} + +/** + * Unwraps a `Node`'s underlying value. + * + * @returns The unwrapped `Node` value or `undefined` if + * the node is not one of the node types in `types`. + */ +export function unwrapOneOfNode( + node: Node, + types: U +): { [K in keyof U]: NodeValue }[number] | undefined { + const nodeType = types.find((type) => unwrapNode(node, type)) + + if (nodeType) { + return unwrapNode(node, nodeType) + } +} + +/** + * Asserts that either the left or right side of the + * expression is processed through `fn` without throwing + * any errors. + * + * If both sides throw errors, the assertion fails and the + * error from the left side will be thrown. + */ +export function assertEitherSideOfExpression( + expression: A_Expr, + fn: (node: Node, side: 'left' | 'right') => U +): U { + assertDefined(expression.lexpr, 'Expected left side of expression to exist') + assertDefined(expression.rexpr, 'Expected right side of expression to exist') + + try { + return fn(expression.lexpr, 'left') + } catch (leftError) { + try { + return fn(expression.rexpr, 'right') + } catch (rightError) { + throw leftError + } + } +} +/** + * Asserts that both sides of the expression are processed + * without throwing any errors. + * + * Order doesn't matter. As long as `firstFn` and `secondFn` + * pass separately on either side of the expression, the assertion + * will pass. Otherwise if `firstFn` and `secondFn` both fail + * after trying on both sides separately, the assertion will fail. + */ +export function assertEachSideOfExpression( + expression: A_Expr, + firstFn: (node: Node) => void, + secondFn: (node: Node) => void +): void { + assertDefined(expression.lexpr, 'Expected left side of expression to exist') + assertDefined(expression.rexpr, 'Expected right side of expression to exist') + + let firstSide: Node + let secondSide: Node + + try { + // Try `firstFn` on the left first + firstSide = expression.lexpr + secondSide = expression.rexpr + firstFn(firstSide) + } catch (firstError) { + // Otherwise try `firstFn` on the right + firstSide = expression.rexpr + secondSide = expression.lexpr + try { + firstFn(firstSide) + } catch (secondError) { + // `firstFn` failed on both sides, so we + // need to throw an error not matter what + + // Perform one more test using `secondFn` + // to help determine which error to show + // for `firstFn` + try { + secondFn(secondSide) + } catch (_) { + throw firstError + } + throw secondError + } + } + + try { + // `firstFn` passed on one of the sides, so + // try `secondFn` on the opposite side + secondFn(secondSide) + } catch (err) { + // `secondFn` failed, so we need to throw an error + throw err + } +} + +export function renderTargets(targets: Node[], renderTarget: (node: Node) => T): T[] { + return targets.map((node) => { + const target = assertAndUnwrapNode( + node, + 'ResTarget', + 'Expected target list to contain ResTargets' + ) + + assertDefined(target.val, 'Expected ResTarget to have a val') + return renderTarget(target.val) + }) +} + +export function renderColumn(column: ColumnRef) { + assertDefined(column.fields, 'Expected column to have fields') + return renderFields(column.fields) +} + +export function assertAndRenderColumn(node: Node, errorMessage: string) { + const column = assertAndUnwrapNode(node, 'ColumnRef', errorMessage) + return renderColumn(column) +} + +export function renderJsonExpression(expression: A_Expr): string { + assertDefined(expression.name, 'Expected expression to have an operator') + + if (expression.name.length > 1) { + throw new AssertionError('Only one JSON operator supported per expression') + } + + const [name] = expression.name + const operatorString = assertAndUnwrapNode( + name, + 'String', + 'Expected JSON operator to be a string' + ) + assertDefined(operatorString.sval, 'PG string expected to have an sval') + const operator = operatorString.sval + + if (!['->', '->>'].includes(operator)) { + throw new AssertionError(`Invalid JSON operator ${operator}`) + } + + assertDefined(expression.lexpr, 'Expected JSON expression to have a left-side component') + assertDefined(expression.rexpr, 'Expected JSON expression to have a right-side component') + + let left: string | number + let right: string | number + + const leftConstant = unwrapNode(expression.lexpr, 'A_Const') + const leftColumn = unwrapNode(expression.lexpr, 'ColumnRef') + const leftFuncCall = unwrapNode(expression.lexpr, 'FuncCall') + const leftExpression = unwrapNode(expression.lexpr, 'A_Expr') + + if (leftConstant) { + // JSON path cannot contain a float + if ('fval' in leftConstant) { + throw new AssertionError('Invalid JSON path: Expression cannot contain a float') + } + left = `'${parseConstant(leftConstant)}'` + } else if (leftColumn) { + left = renderColumn(leftColumn) + } else if (leftFuncCall) { + assertDefined(leftFuncCall.funcname, 'Expected function call to have a name') + const functionName = renderFields(leftFuncCall.funcname) + left = `${functionName}()` + } else if (leftExpression) { + left = renderJsonExpression(leftExpression) + } else { + throw new AssertionError('Invalid JSON path') + } + + const rightConstant = unwrapNode(expression.rexpr, 'A_Const') + + if (rightConstant) { + // JSON path cannot contain a float + if ('fval' in rightConstant) { + throw new AssertionError('Invalid JSON path: Expression cannot contain a float') + } + right = `'${parseConstant(rightConstant)}'` + } else { + throw new AssertionError('Invalid JSON path') + } + + return `${left}${operator}${right}` +} + +export function renderFields(fields: Node[]) { + const nameSegments = fields + .map((field) => { + const stringField = unwrapNode(field, 'String') + const starField = unwrapNode(field, 'A_Star') + + if (stringField !== undefined) { + return stringField.sval + } else if (starField !== undefined) { + return '*' + } else { + const [internalType] = Object.keys(field) + throw new Error(`Unsupported internal type '${internalType}' for fields`) + } + }) + .filter((name): name is string => name !== undefined) + + return nameSegments.join('.') +} + +export function parseConstant(constant: A_Const) { + if ('sval' in constant) { + return constant.sval?.sval ?? '' + } else if ('ival' in constant) { + // The PG parser turns 0 into undefined, so convert it back here + return constant.ival?.ival ?? 0 + } else if ('fval' in constant) { + return constant.fval?.fval ? parseFloat(constant.fval.fval) : 0 + } else { + throw new AssertionError(`Constant values must be a string, integer, or float`) + } +}