feat(db-playground): seeds tab

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,
         }),
       },
     },

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<TabValue>('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() {
         <Tabs
           className="flex-1 h-full flex flex-col items-stretch"
           value={tab}
-          onValueChange={(tab) => setTab(tab)}
+          onValueChange={(tab) => setTab(tabsSchema.parse(tab))}
         >
-          <TabsList className="grid w-full grid-cols-2">
+          <TabsList className="grid w-full grid-cols-3">
             <TabsTrigger value="diagram">Diagram</TabsTrigger>
             <TabsTrigger value="migrations">Migrations</TabsTrigger>
+            <TabsTrigger value="seeds">Seeds</TabsTrigger>
           </TabsList>
           <TabsContent value="diagram" className="h-full">
             <SchemaGraph schema="public" />
@@ -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()
+              }}
+            />
+          </TabsContent>
+          <TabsContent value="seeds" className="h-full py-4 rounded-md bg-[#1e1e1e]">
+            <Editor
+              language="pgsql"
+              value={seedSql}
+              theme="vs-dark"
+              options={{
+                tabSize: 2,
+                minimap: {
+                  enabled: false,
+                },
+                fontSize: 13,
+                readOnly: true,
+              }}
+              onMount={async (editor, monaco) => {
+                // 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()
               }}

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
+  }
+}

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"
   }
-}
+}

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
     }
   }
-}
+}

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"
   }
-}
+}

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'
 

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> = T extends T ? keyof T : never
-
-/**
- * Unwraps a Node to get its underlying value.
- */
-type NodeValue<T extends Node, U extends ExtractKeys<Node>> = T extends Record<U, infer V>
-  ? 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<T>(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<T extends Node>(
-  node: Node,
-  type: ExtractKeys<T>,
-  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<T extends Node, U extends ExtractKeys<Node>>(
-  node: Node,
-  type: U,
-  errorMessage: string
-): NodeValue<T, U> {
-  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<T extends Node, U extends ExtractKeys<Node>>(
-  node: Node,
-  type: U
-): NodeValue<T, U> | 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<U>(
-  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<U>(
-  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<T>(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`)
-  }
-}

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": "*"
   }
-}
+}

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> = T extends T ? keyof T : never
+
+/**
+ * Filters strings ending in the specified suffix.
+ */
+export type StringWithSuffix<T extends string, S extends string> = T extends `${infer _Prefix}${S}`
+  ? T
+  : never
+
+/**
+ * Extracts all nodes who's keys end in the specified suffix.
+ */
+export type NodeWithSuffix<S extends string> = 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<Node>
+
+/**
+ * Every statement Node type.
+ */
+export type StatementType = StringWithSuffix<NodeType, 'Stmt'>
+
+/**
+ * Every statement Node.
+ */
+export type Statement = NodeWithSuffix<'Stmt'>
+
+/**
+ * Unwraps a Node to get its underlying value.
+ */
+export type NodeValue<T extends Node, U extends NodeType> = T extends Record<U, infer V> ? 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<T>(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<T extends Node>(
+  node: Node,
+  type: ExtractKeys<T>,
+  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<T extends Node, U extends NodeType>(
+  node: Node,
+  type: U,
+  errorMessage: string
+): NodeValue<T, U> {
+  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<T extends Node, U extends NodeType>(
+  node: Node,
+  type: U
+): NodeValue<T, U> | 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<T extends Node, U extends readonly [NodeType, ...NodeType[]]>(
+  node: Node,
+  types: U
+): { [K in keyof U]: NodeValue<T, U[K]> }[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<U>(
+  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<U>(
+  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<T>(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`)
+  }
+}