diff --git a/apps/postgres-new/components/app-provider.tsx b/apps/postgres-new/components/app-provider.tsx index 88663bbc..14c2136a 100644 --- a/apps/postgres-new/components/app-provider.tsx +++ b/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() - const [isSignInDialogOpen, setIsSignInDialogOpen] = useState(false) - const [isRenameDialogOpen, setIsRenameDialogOpen] = useState(false) - const [isRateLimited, setIsRateLimited] = useState(false) - - const focusRef = useRef(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(null) const [connectedClientIp, setConnectedClientIp] = useState(null) const [liveShareWebsocket, setLiveShareWebsocket] = useState(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 + +export default function AppProvider({ children }: AppProps) { + const [isLoadingUser, setIsLoadingUser] = useState(true) + const [user, setUser] = useState() + const [isSignInDialogOpen, setIsSignInDialogOpen] = useState(false) + const [isRenameDialogOpen, setIsRenameDialogOpen] = useState(false) + const [isRateLimited, setIsRateLimited] = useState(false) + + const focusRef = useRef(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 ( Promise - stop: () => void - databaseId: string | null - clientIp: string | null - isLiveSharing: boolean - } + liveShare: LiveShare isLegacyDomain: boolean isLegacyDomainRedirect: boolean } diff --git a/apps/postgres-new/components/sidebar.tsx b/apps/postgres-new/components/sidebar.tsx index 9db7d255..f69181af 100644 --- a/apps/postgres-new/components/sidebar.tsx +++ b/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) }} diff --git a/apps/postgres-new/package.json b/apps/postgres-new/package.json index 15b733c3..8d400652 100644 --- a/apps/postgres-new/package.json +++ b/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", diff --git a/package-lock.json b/package-lock.json index 85591eb1..53c5af93 100644 --- a/package-lock.json +++ b/package-lock.json @@ -93,6 +93,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", @@ -3478,6 +3479,12 @@ "@jsr/std__streams": "^1.0.5" } }, + "node_modules/@supabase-labs/pg-protocol": { + "version": "0.1.0", + "resolved": "https://pkg.pr.new/supabase-community/pg-protocol/@supabase-labs/pg-protocol@86c902a", + "integrity": "sha512-EtSMB7z+iaHfqnodTf78sR8vz4G6JxLUNkMSxWEFEr2NiJjoHDYUCKdHLpodTrJPpMy0lbqy4RO88KwEJo3h6w==", + "license": "MIT" + }, "node_modules/@supabase/auth-js": { "version": "2.65.0", "license": "MIT",