feat: capture migrations during Live Share session

apps/postgres-new/components/app-provider.tsx

@@ -8,6 +8,8 @@ import { User } from '@supabase/supabase-js'
 import { useQueryClient } from '@tanstack/react-query'
 import { Mutex } from 'async-mutex'
 import { debounce } from 'lodash'
+import { isQuery, parseQuery as parseQueryMessage } from '@supabase-labs/pg-protocol/frontend'
+import { isErrorResponse } from '@supabase-labs/pg-protocol/backend'
 import {
   createContext,
   PropsWithChildren,
@@ -32,105 +34,31 @@ import {
 import { legacyDomainHostname } from '~/lib/util'
 import { parse, serialize } from '~/lib/websocket-protocol'
 import { createClient } from '~/utils/supabase/client'
+import { assertDefined, isMigrationStatement } from '~/lib/sql-util'
+import type { ParseResult } from 'libpg-query/wasm'
+import { generateId, Message } from 'ai'
+import { getMessagesQueryKey } from '~/data/messages/messages-query'
 
 export type AppProps = PropsWithChildren
 
 // Create a singleton DbManager that isn't exposed to double mounting
 const dbManager = typeof window !== 'undefined' ? new DbManager() : undefined
 
-export default function AppProvider({ children }: AppProps) {
-  const [isLoadingUser, setIsLoadingUser] = useState(true)
-  const [user, setUser] = useState<User>()
-  const [isSignInDialogOpen, setIsSignInDialogOpen] = useState(false)
-  const [isRenameDialogOpen, setIsRenameDialogOpen] = useState(false)
-  const [isRateLimited, setIsRateLimited] = useState(false)
-
-  const focusRef = useRef<FocusHandle>(null)
-
+function useLiveShare() {
   const supabase = createClient()
-
-  useEffect(() => {
-    const {
-      data: { subscription },
-    } = supabase.auth.onAuthStateChange((e) => {
-      focusRef.current?.focus()
-    })
-
-    return () => subscription.unsubscribe()
-  }, [supabase])
-
-  const loadUser = useCallback(async () => {
-    setIsLoadingUser(true)
-    try {
-      const { data, error } = await supabase.auth.getUser()
-
-      if (error) {
-        // TODO: handle error
-        setUser(undefined)
-        return
-      }
-
-      const { user } = data
-
-      setUser(user)
-
-      return user
-    } finally {
-      setIsLoadingUser(false)
-    }
-  }, [supabase])
-
-  useEffect(() => {
-    loadUser()
-  }, [loadUser])
-
-  const signIn = useCallback(async () => {
-    const { error } = await supabase.auth.signInWithOAuth({
-      provider: 'github',
-      options: {
-        redirectTo: window.location.toString(),
-      },
-    })
-
-    if (error) {
-      // TODO: handle sign in error
-    }
-
-    const user = await loadUser()
-    return user
-  }, [supabase, loadUser])
-
-  const signOut = useCallback(async () => {
-    const { error } = await supabase.auth.signOut()
-
-    if (error) {
-      // TODO: handle sign out error
-    }
-
-    setUser(undefined)
-  }, [supabase])
-
-  const pgliteVersion = process.env.NEXT_PUBLIC_PGLITE_VERSION
-  const { value: pgVersion } = useAsyncMemo(async () => {
-    if (!dbManager) {
-      throw new Error('dbManager is not available')
-    }
-
-    return await dbManager.getRuntimePgVersion()
-  }, [dbManager])
-
   const queryClient = useQueryClient()
-
   const [liveSharedDatabaseId, setLiveSharedDatabaseId] = useState<string | null>(null)
   const [connectedClientIp, setConnectedClientIp] = useState<string | null>(null)
   const [liveShareWebsocket, setLiveShareWebsocket] = useState<WebSocket | null>(null)
+
   const cleanUp = useCallback(() => {
     setLiveShareWebsocket(null)
     setLiveSharedDatabaseId(null)
     setConnectedClientIp(null)
   }, [setLiveShareWebsocket, setLiveSharedDatabaseId, setConnectedClientIp])
+
   const startLiveShare = useCallback(
-    async (databaseId: string) => {
+    async (databaseId: string, options?: { captureMigrations?: boolean }) => {
       if (!dbManager) {
         throw new Error('dbManager is not available')
       }
@@ -205,6 +133,43 @@ export default function AppProvider({ children }: AppProps) {
 
           ws.send(serialize(connectionId, response))
 
+          // Capture migrations if enabled
+          if (options?.captureMigrations && !isErrorResponse(response)) {
+            const { deparse, parseQuery } = await import('libpg-query/wasm')
+            if (isQuery(message)) {
+              const parsedMessage = parseQueryMessage(message)
+              const parseResult = await parseQuery(parsedMessage.query)
+              assertDefined(parseResult.stmts, 'Expected stmts to exist in parse result')
+              const migrationStmts = parseResult.stmts.filter(isMigrationStatement)
+              if (migrationStmts.length > 0) {
+                const filteredAst: ParseResult = {
+                  version: parseResult.version,
+                  stmts: migrationStmts,
+                }
+                const migrationSql = await deparse(filteredAst)
+                const chatMessage: Message = {
+                  id: generateId(),
+                  role: 'assistant',
+                  content: '',
+                  toolInvocations: [
+                    {
+                      state: 'result',
+                      toolCallId: generateId(),
+                      toolName: 'executeSql',
+                      args: { sql: migrationSql },
+                      result: { success: true },
+                    },
+                  ],
+                }
+                await dbManager.createMessage(databaseId, chatMessage)
+                // invalidate messages query to refresh the migrations tab
+                await queryClient.invalidateQueries({
+                  queryKey: getMessagesQueryKey(databaseId),
+                })
+              }
+            }
+          }
+
           // Refresh table UI when safe to do so
           // A backend response can have multiple wire messages
           const backendMessages = Array.from(getMessages(response))
@@ -232,19 +197,105 @@ export default function AppProvider({ children }: AppProps) {
 
       setLiveShareWebsocket(ws)
     },
-    [cleanUp, supabase.auth]
+    [cleanUp, supabase.auth, queryClient]
   )
+
   const stopLiveShare = useCallback(() => {
     liveShareWebsocket?.close()
     cleanUp()
   }, [cleanUp, liveShareWebsocket])
-  const liveShare = {
+
+  return {
     start: startLiveShare,
     stop: stopLiveShare,
     databaseId: liveSharedDatabaseId,
     clientIp: connectedClientIp,
     isLiveSharing: Boolean(liveSharedDatabaseId),
   }
+}
+type LiveShare = ReturnType<typeof useLiveShare>
+
+export default function AppProvider({ children }: AppProps) {
+  const [isLoadingUser, setIsLoadingUser] = useState(true)
+  const [user, setUser] = useState<User>()
+  const [isSignInDialogOpen, setIsSignInDialogOpen] = useState(false)
+  const [isRenameDialogOpen, setIsRenameDialogOpen] = useState(false)
+  const [isRateLimited, setIsRateLimited] = useState(false)
+
+  const focusRef = useRef<FocusHandle>(null)
+
+  const supabase = createClient()
+
+  useEffect(() => {
+    const {
+      data: { subscription },
+    } = supabase.auth.onAuthStateChange((e) => {
+      focusRef.current?.focus()
+    })
+
+    return () => subscription.unsubscribe()
+  }, [supabase])
+
+  const loadUser = useCallback(async () => {
+    setIsLoadingUser(true)
+    try {
+      const { data, error } = await supabase.auth.getUser()
+
+      if (error) {
+        // TODO: handle error
+        setUser(undefined)
+        return
+      }
+
+      const { user } = data
+
+      setUser(user)
+
+      return user
+    } finally {
+      setIsLoadingUser(false)
+    }
+  }, [supabase])
+
+  useEffect(() => {
+    loadUser()
+  }, [loadUser])
+
+  const signIn = useCallback(async () => {
+    const { error } = await supabase.auth.signInWithOAuth({
+      provider: 'github',
+      options: {
+        redirectTo: window.location.toString(),
+      },
+    })
+
+    if (error) {
+      // TODO: handle sign in error
+    }
+
+    const user = await loadUser()
+    return user
+  }, [supabase, loadUser])
+
+  const signOut = useCallback(async () => {
+    const { error } = await supabase.auth.signOut()
+
+    if (error) {
+      // TODO: handle sign out error
+    }
+
+    setUser(undefined)
+  }, [supabase])
+
+  const pgliteVersion = process.env.NEXT_PUBLIC_PGLITE_VERSION
+  const { value: pgVersion } = useAsyncMemo(async () => {
+    if (!dbManager) {
+      throw new Error('dbManager is not available')
+    }
+
+    return await dbManager.getRuntimePgVersion()
+  }, [dbManager])
+
   const [isLegacyDomain, setIsLegacyDomain] = useState(false)
   const [isLegacyDomainRedirect, setIsLegacyDomainRedirect] = useState(false)
 
@@ -259,6 +310,8 @@ export default function AppProvider({ children }: AppProps) {
     setIsRenameDialogOpen(isLegacyDomain || isLegacyDomainRedirect)
   }, [])
 
+  const liveShare = useLiveShare()
+
   return (
     <AppContext.Provider
       value={{
@@ -305,13 +358,7 @@ export type AppContextValues = {
   dbManager?: DbManager
   pgliteVersion?: string
   pgVersion?: string
-  liveShare: {
-    start: (databaseId: string) => Promise<void>
-    stop: () => void
-    databaseId: string | null
-    clientIp: string | null
-    isLiveSharing: boolean
-  }
+  liveShare: LiveShare
   isLegacyDomain: boolean
   isLegacyDomainRedirect: boolean
 }

apps/postgres-new/components/sidebar.tsx

@@ -579,7 +579,7 @@ function LiveShareMenuItem(props: ConnectMenuItemProps) {
         if (liveShare.isLiveSharing) {
           liveShare.stop()
         }
-        liveShare.start(props.databaseId)
+        liveShare.start(props.databaseId, { captureMigrations: true })
         router.push(`/db/${props.databaseId}`)
         props.setIsPopoverOpen(false)
       }}

apps/postgres-new/package.json

@@ -25,6 +25,7 @@
     "@radix-ui/react-tabs": "^1.1.0",
     "@radix-ui/react-tooltip": "^1.1.2",
     "@std/tar": "npm:@jsr/std__tar@^0.1.1",
+    "@supabase-labs/pg-protocol": "https://pkg.pr.new/supabase-community/pg-protocol/@supabase-labs/pg-protocol@86c902a",
     "@supabase/postgres-meta": "^0.81.2",
     "@supabase/ssr": "^0.4.0",
     "@supabase/supabase-js": "^2.45.0",