diff --git a/apps/studio/components/interfaces/Database/Tests/TestRow.tsx b/apps/studio/components/interfaces/Database/Tests/TestRow.tsx new file mode 100644 index 00000000000..b0d2c11c7ac --- /dev/null +++ b/apps/studio/components/interfaces/Database/Tests/TestRow.tsx @@ -0,0 +1,193 @@ +import dayjs from 'dayjs' +import { MoreVertical, Play, Trash } from 'lucide-react' +import { forwardRef, useImperativeHandle, useState } from 'react' +import { toast } from 'sonner' +import { + Badge, + Button, + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, + TableCell, + TableRow, +} from 'ui' +import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' +import { ButtonTooltip } from 'components/ui/ButtonTooltip' + +import { DatabaseTest } from 'data/database-tests/database-tests-query' +import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' +import { useSelectedProject } from 'hooks/misc/useSelectedProject' +import { useDatabaseTestDeleteMutation } from 'data/database-tests/database-test-delete-mutation' + +export interface TestRowHandle { + runTest: () => Promise<{ error: any | null }> + setStatus: (status: 'queued' | 'passed' | 'failed' | undefined) => void + setIsRunning: (running: boolean) => void +} + +interface TestRowProps { + test: DatabaseTest + prependQuery?: DatabaseTest | null + canRun?: boolean + onSelectTest: (test: DatabaseTest) => void +} + +// Utility function to validate test query format +const isValidTestQuery = (query: string): boolean => { + const trimmedQuery = query.trim().toLowerCase() + return trimmedQuery.startsWith('begin;') && trimmedQuery.endsWith('rollback;') +} + +const TestRow = forwardRef( + ({ test, prependQuery, canRun = false, onSelectTest }, ref) => { + const project = useSelectedProject() + const [isRunning, setIsRunning] = useState(false) + const [status, setStatus] = useState<'queued' | 'passed' | 'failed' | undefined>() + const [lastRun, setLastRun] = useState() + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) + + const { mutateAsync: runTestMutation } = useExecuteSqlMutation() + const { mutate: deleteTest, isLoading: isDeleting } = useDatabaseTestDeleteMutation({ + onSuccess: () => { + setIsDeleteModalOpen(false) + }, + }) + + const isTestQueryValid = isValidTestQuery(test.query) + + const handleRunTest = async () => { + if (!isTestQueryValid) { + toast.error( + `Test "${test.name}" has invalid format. Must start with BEGIN; and end with ROLLBACK;` + ) + return { error: new Error('Invalid test format') } + } + + setIsRunning(true) + setStatus(undefined) + try { + // Determine the SQL to run + let sqlToRun = test.query + + // If a setup test exists, prepend it to regular tests + if (prependQuery && canRun) { + sqlToRun = `${prependQuery.query}\n\n${test.query}` + } + + const res = await runTestMutation({ + projectRef: project?.ref, + connectionString: project?.connectionString, + sql: sqlToRun, + }) + const { result } = res + console.log('result', res) + const hasFailedTests = result.some((r: any) => r.finish?.includes('failed')) + if (hasFailedTests) { + toast.error(`Test "${test.name}" failed.`) + setStatus('failed') + } else { + toast.success(`Test "${test.name}" passed!`) + setStatus('passed') + } + setLastRun(new Date().toISOString()) + return { error: null } + } catch (error: any) { + toast.error(`Failed to run test "${test.name}": ${error.message}`) + setStatus('failed') + setLastRun(new Date().toISOString()) + return { error } + } finally { + setIsRunning(false) + } + } + + useImperativeHandle(ref, () => ({ + runTest: handleRunTest, + setStatus, + setIsRunning, + })) + + const onDelete = () => { + if (!project) return + deleteTest({ projectRef: project.ref, id: test.id }) + } + + const getStatusBadge = () => { + if (isRunning) { + return Running + } + if (status === 'queued') { + return Queued + } + if (status === 'passed') { + return Passed + } + if (status === 'failed') { + return Failed + } + return Not run + } + + return ( + <> + + +

onSelectTest(test)}> + {test.name} +

+ + {getStatusBadge()} + {lastRun ? dayjs(lastRun).fromNow() : 'Never'} + +
+ {canRun && ( + } + loading={isRunning} + disabled={isRunning || status === 'queued' || !isTestQueryValid} + onClick={handleRunTest} + tooltip={{ + content: { + side: 'bottom', + text: !isTestQueryValid + ? 'Test query must start with BEGIN; and end with ROLLBACK;' + : undefined, + }, + }} + /> + )} + + +
+ + + setIsDeleteModalOpen(false)} + onConfirm={onDelete} + loading={isDeleting} + > +

+ Are you sure you want to delete the test "{test.name}"? This action cannot be undone. +

+ + + ) + } +) +TestRow.displayName = 'TestRow' + +export default TestRow diff --git a/apps/studio/components/interfaces/Database/Tests/TestsList.tsx b/apps/studio/components/interfaces/Database/Tests/TestsList.tsx new file mode 100644 index 00000000000..d3b1de7d272 --- /dev/null +++ b/apps/studio/components/interfaces/Database/Tests/TestsList.tsx @@ -0,0 +1,311 @@ +import { partition } from 'lodash' +import { useRouter } from 'next/router' +import { forwardRef, useImperativeHandle, useMemo, useRef, useState } from 'react' +import { Search, MoreVertical, Trash } from 'lucide-react' +import { + Alert, + Button, + Card, + CardHeader, + CardTitle, + Input, + Table, + TableBody, + TableHead, + TableHeader, + TableRow, + TableCell, + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from 'ui' +import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' + +import ShimmeringLoader from 'components/ui/ShimmeringLoader' +import { useDatabaseTestsQuery } from 'data/database-tests/database-tests-query' +import { DatabaseTest } from 'data/database-tests/database-tests-query' +import { useDatabaseTestDeleteMutation } from 'data/database-tests/database-test-delete-mutation' +import { useSelectedProject } from 'hooks/misc/useSelectedProject' +import { toast } from 'sonner' +import TestRow, { TestRowHandle } from './TestRow' + +// Utility function to validate test query format +const isValidTestQuery = (query: string): boolean => { + const trimmedQuery = query.trim().toLowerCase() + return trimmedQuery.startsWith('begin;') && trimmedQuery.endsWith('rollback;') +} + +// Simple component for setup test rows - just name and actions +interface SetupTestRowProps { + test: DatabaseTest + onSelectTest: (test: DatabaseTest) => void +} + +const SetupTestRow = ({ test, onSelectTest }: SetupTestRowProps) => { + const project = useSelectedProject() + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) + + const { mutate: deleteTest, isLoading: isDeleting } = useDatabaseTestDeleteMutation({ + onSuccess: () => { + setIsDeleteModalOpen(false) + }, + }) + + const onDelete = () => { + if (!project) return + deleteTest({ projectRef: project.ref, id: test.id }) + } + + return ( +
+

onSelectTest(test)}> + {test.name} +

+ + +
+ ) +} + +export interface TestsListHandle { + runAllTests: () => void +} + +interface DatabaseTestsListProps { + onSelectTest: (test: DatabaseTest) => void +} + +const DatabaseTestsList = forwardRef( + ({ onSelectTest }, ref) => { + const router = useRouter() + const [search, setSearch] = useState('') + const testRowRefs = useRef<{ [key: string]: TestRowHandle | null }>({}) + + const { + data: tests, + isLoading, + isSuccess, + isError, + error, + } = useDatabaseTestsQuery({ + projectRef: router.query.ref as string, + }) + + // Split tests into setup and regular groups + const { setupTests, regularTests, combinedSetupQuery } = useMemo(() => { + if (!tests) return { setupTests: [], regularTests: [], combinedSetupQuery: null } + + const setupTests = tests.filter((test) => test.name.startsWith('000')) + const regularTests = tests.filter((test) => !test.name.startsWith('000')) + + // Sort setup tests alphabetically to ensure consistent order + const sortedSetupTests = setupTests.sort((a, b) => a.name.localeCompare(b.name)) + + // Combine all setup test queries into a single query + const combinedSetupQuery = + sortedSetupTests.length > 0 + ? ({ + id: 'combined-setup', + name: 'Combined Setup Tests', + query: sortedSetupTests.map((test) => test.query).join('\n\n'), + } as DatabaseTest) + : null + + return { + setupTests: sortedSetupTests, + regularTests, + combinedSetupQuery, + } + }, [tests]) + + // Filter both groups based on search + const filteredSetupTests = useMemo(() => { + if (search.length === 0) return setupTests + return setupTests.filter((t) => t.name.toLowerCase().includes(search.toLowerCase())) + }, [setupTests, search]) + + const filteredRegularTests = useMemo(() => { + if (search.length === 0) return regularTests + return regularTests.filter((t) => t.name.toLowerCase().includes(search.toLowerCase())) + }, [regularTests, search]) + + const runAllTests = async () => { + const allTests = [...filteredSetupTests, ...filteredRegularTests] + if (allTests.length === 0) { + toast.error('No tests to run.') + return + } + + // Filter out invalid tests + const validTests = allTests.filter((test) => isValidTestQuery(test.query)) + const invalidTests = allTests.filter((test) => !isValidTestQuery(test.query)) + + if (invalidTests.length > 0) { + toast.error( + `${invalidTests.length} test(s) skipped - must start with BEGIN; and end with ROLLBACK;` + ) + } + + if (validTests.length === 0) { + toast.error('No valid tests to run.') + return + } + + // Always run setup test first, then regular tests alphabetically + const validSetupTests = validTests.filter((test) => test.name.startsWith('000')) + const validRegularTests = validTests.filter((test) => !test.name.startsWith('000')) + + const sortedRegularTests = [...validRegularTests].sort((a, b) => a.name.localeCompare(b.name)) + const sortedTests = [...validSetupTests, ...sortedRegularTests] + + // Set all tests to queued status initially + sortedTests.forEach((test) => { + testRowRefs.current[test.id]?.setStatus('queued') + }) + + let passedCount = 0 + let failedCount = 0 + + // Run tests sequentially + for (const test of sortedTests) { + const testRef = testRowRefs.current[test.id] + if (testRef) { + // Clear queued status and run the test + testRef.setStatus(undefined) + const result = await testRef.runTest() + + if (result?.error) { + failedCount++ + } else { + passedCount++ + } + } + } + + // Show final results + if (failedCount > 0) { + toast.error(`${failedCount} out of ${validTests.length} tests failed.`) + } else { + toast.success(`All ${validTests.length} tests passed!`) + } + } + + useImperativeHandle(ref, () => ({ + runAllTests, + })) + + // Pass combined setup query to child components for prepending + const setupTest = combinedSetupQuery + + const renderTestTable = (tests: DatabaseTest[], isSetupSection = false) => { + if (isSetupSection) { + // Simplified view for setup tests - just name and actions + return ( +
+ {tests.length === 0 ? ( +

+ No setup tests found. Create tests with names starting with "000" to set up test + environment. +

+ ) : ( + tests.map((test) => ( + + )) + )} +
+ ) + } + + // Regular tests view with full TestRow component + return ( + + + + Name + Status + Last Run + + + + + {tests.map((test) => ( + (testRowRefs.current[test.id] = el)} + key={test.id} + test={test} + prependQuery={setupTest} + canRun={true} + onSelectTest={onSelectTest} + /> + ))} + +
+ ) + } + + return ( +
+ {isLoading && } + {isError && ( +
+ + {(error as any)?.message} + +
+ )} + {isSuccess && ( + <> + {filteredSetupTests.length === 0 && + filteredRegularTests.length === 0 && + setupTests.length === 0 && ( +
No tests found.
+ )} + + {/* Setup Tests Group - Always show */} + + + Before Tests + + {renderTestTable(filteredSetupTests, true)} + + + {/* Regular Tests Group */} + {filteredRegularTests.length > 0 && ( + + + Tests + + {renderTestTable(filteredRegularTests, false)} + + )} + + )} +
+ ) + } +) + +DatabaseTestsList.displayName = 'DatabaseTestsList' +export default DatabaseTestsList diff --git a/apps/studio/components/layouts/DatabaseLayout/Database.Commands.tsx b/apps/studio/components/layouts/DatabaseLayout/Database.Commands.tsx index ba42375641b..c1c3a05bc14 100644 --- a/apps/studio/components/layouts/DatabaseLayout/Database.Commands.tsx +++ b/apps/studio/components/layouts/DatabaseLayout/Database.Commands.tsx @@ -52,6 +52,13 @@ export function useDatabaseGotoCommands(options?: CommandOptions) { route: `/project/${ref}/database/functions`, defaultHidden: true, }, + { + id: 'nav-database-tests', + name: 'Tests', + value: 'Database: Tests', + route: `/project/${ref}/database/tests`, + defaultHidden: true, + }, { id: 'nav-database-extensions', name: 'Extensions', diff --git a/apps/studio/components/layouts/DatabaseLayout/DatabaseMenu.utils.tsx b/apps/studio/components/layouts/DatabaseLayout/DatabaseMenu.utils.tsx index 48388f21f59..efeb98ed3b3 100644 --- a/apps/studio/components/layouts/DatabaseLayout/DatabaseMenu.utils.tsx +++ b/apps/studio/components/layouts/DatabaseLayout/DatabaseMenu.utils.tsx @@ -39,6 +39,12 @@ export const generateDatabaseMenu = ( url: `/project/${ref}/database/triggers`, items: [], }, + { + name: 'Tests', + key: 'tests', + url: `/project/${ref}/database/tests`, + items: [], + }, { name: 'Enumerated Types', key: 'types', diff --git a/apps/studio/components/ui/EditorPanel/EditorPanel.tsx b/apps/studio/components/ui/EditorPanel/EditorPanel.tsx index ee70fab8c8f..4acab158d0f 100644 --- a/apps/studio/components/ui/EditorPanel/EditorPanel.tsx +++ b/apps/studio/components/ui/EditorPanel/EditorPanel.tsx @@ -101,8 +101,8 @@ export const EditorPanel = ({ onChange }: EditorPanelProps) => { } } - if (editorPanel.onSave) { - editorPanel.onSave(currentValue) + if (editorPanel.onRun) { + editorPanel.onRun(currentValue) } executeSql({ @@ -371,6 +371,11 @@ export const EditorPanel = ({ onChange }: EditorPanelProps) => { )}
+ {editorPanel.onSave && ( + + )}
diff --git a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx index b434044f48c..59fe121cd45 100644 --- a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx +++ b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx @@ -1,4 +1,4 @@ -import { Code, Play } from 'lucide-react' +import { Code, Play, Edit, MoreVertical } from 'lucide-react' import { DragEvent, ReactNode, useEffect, useMemo, useState } from 'react' import { Bar, BarChart, CartesianGrid, Cell, Tooltip, XAxis, YAxis } from 'recharts' import { toast } from 'sonner' @@ -7,20 +7,40 @@ import { useParams } from 'common' import { ReportBlockContainer } from 'components/interfaces/Reports/ReportBlock/ReportBlockContainer' import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig' import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results' +import { useIsInlineEditorEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext' +import useNewQuery from 'components/interfaces/SQLEditor/hooks' import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' +import { useDatabaseTestCreateMutation } from 'data/database-tests/database-test-create-mutation' +import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import dayjs from 'dayjs' import { Parameter, parseParameters } from 'lib/sql-parameters' import { Dashboards } from 'types' -import { ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui' +import { + ChartContainer, + ChartTooltipContent, + cn, + CodeBlock, + SQL_ICON, + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from 'ui' import ShimmeringLoader from 'ui-patterns/ShimmeringLoader' import { ButtonTooltip } from '../ButtonTooltip' import { CHART_COLORS } from '../Charts/Charts.constants' import SqlWarningAdmonition from '../SqlWarningAdmonition' import { BlockViewConfiguration } from './BlockViewConfiguration' -import { EditQueryButton } from './EditQueryButton' import { ParametersPopover } from './ParametersPopover' import { getCumulativeResults } from './QueryBlock.utils' +import { useRouter } from 'next/router' +import Link from 'next/link' +import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization' +import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state' +import { useAppStateSnapshot } from 'state/app-state' +import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2' +import { DiffType } from 'components/interfaces/SQLEditor/SQLEditor.types' export const DEFAULT_CHART_CONFIG: ChartConfig = { type: 'bar', @@ -122,7 +142,17 @@ export const QueryBlock = ({ onDragStart, onResults, }: QueryBlockProps) => { + const router = useRouter() const { ref } = useParams() + const { newQuery } = useNewQuery() + const sqlEditorSnap = useSqlEditorV2StateSnapshot() + const { setEditorPanel } = useAppStateSnapshot() + const snap = useAiAssistantStateSnapshot() + const org = useSelectedOrganization() + const { mutate: sendEvent } = useSendEventMutation() + const isInlineEditorEnabled = useIsInlineEditorEnabled() + const isInSQLEditor = router.pathname.includes('/sql') + const isInNewSnippet = router.pathname.endsWith('/sql') const [chartSettings, setChartSettings] = useState(chartConfig) const { xKey, yKey, view = 'table' } = chartSettings @@ -159,6 +189,15 @@ export const QueryBlock = ({ const postgresConnectionString = primaryDatabase?.connectionString const readOnlyConnectionString = primaryDatabase?.connection_string_read_only + const { mutate: createTest, isLoading: isCreatingTest } = useDatabaseTestCreateMutation({ + onSuccess: () => { + toast.success('Test created successfully') + }, + onError: (error) => { + toast.error(`Failed to create test: ${error.message}`) + }, + }) + const chartData = chartSettings.cumulative ? getCumulativeResults({ rows: formattedQueryResult ?? [] }, chartSettings) : formattedQueryResult @@ -211,6 +250,139 @@ export const QueryBlock = ({ } } + const handleSaveAsSnippet = () => { + if (!sql) return + + if (id !== undefined) { + // Navigate to existing snippet + router.push(`/project/${ref}/sql/${id}`) + return + } + + if (!isInSQLEditor || isInNewSnippet) { + if (isInlineEditorEnabled) { + setEditorPanel({ + open: true, + initialValue: sql, + }) + snap.closeAssistant() + } else { + if (sql) newQuery(sql, label) + } + sendEvent({ + action: 'assistant_edit_in_sql_editor_clicked', + properties: { + isInSQLEditor, + isInNewSnippet, + }, + groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' }, + }) + } else { + // Show dropdown options for SQL editor context + sqlEditorSnap.setDiffContent(sql, DiffType.Addition) + } + } + + const handleSaveAsTest = () => { + if (!sql) return + + createTest({ + projectRef: ref!, + connectionString: postgresConnectionString, + query: sql, + }) + } + + const renderQueryActionsDropdown = () => { + if (id !== undefined) { + // If this is an existing snippet, just show the edit button + return ( + } + tooltip={{ + content: { side: 'bottom', text: 'Edit in SQL Editor' }, + }} + > + + + ) + } + + if (!isInSQLEditor || isInNewSnippet) { + // Show dropdown for save options + return ( + + + } + tooltip={{ + content: { side: 'bottom', text: 'Save options' }, + }} + /> + + + + + Save as snippet + + + + Save as test + + + + ) + } else { + // In SQL editor, show the original dropdown behavior + return ( + + + } + tooltip={ + !!sql + ? { content: { side: 'bottom', text: 'Edit in SQL Editor' } } + : { content: { side: 'bottom', text: undefined } } + } + /> + + {!!sql && ( + + sqlEditorSnap.setDiffContent(sql, DiffType.Addition)} + > + Insert code + + sqlEditorSnap.setDiffContent(sql, DiffType.Modification)} + > + Replace code + + newQuery(sql, label)}> + Create new snippet + + + + Save as test + + + )} + + ) + } + } + useEffect(() => { setChartSettings(chartConfig) }, [chartConfig]) @@ -295,7 +467,7 @@ export const QueryBlock = ({ )} - + {renderQueryActionsDropdown()} {(showRunButtonIfNotReadOnly || !readOnlyError) && ( void +} = {}) => { + const queryClient = useQueryClient() + const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation() + + return useMutation({ + mutationFn: async ({ projectRef, query }: DatabaseTestCreateVariables) => { + // AI title generation + const { title } = await generateSqlTitle({ sql: query }) + + // This is a mock implementation. + // In a real scenario, you would save the test to a file or a database table. + const key = databaseTestsKeys.list(projectRef) + const currentTests = queryClient.getQueryData(key) ?? [] + const newTest: DatabaseTest = { + id: uuidv4(), + name: title, + query, + } + queryClient.setQueryData(key, [...currentTests, newTest]) + return newTest + }, + onSuccess: (newTest) => { + toast.success(`Test "${newTest.name}" created.`) + onSuccess?.() + }, + onError: (error: any) => { + toast.error(`Failed to create test: ${error.message}`) + }, + }) +} diff --git a/apps/studio/data/database-tests/database-test-delete-mutation.ts b/apps/studio/data/database-tests/database-test-delete-mutation.ts new file mode 100644 index 00000000000..b50d41a2ebc --- /dev/null +++ b/apps/studio/data/database-tests/database-test-delete-mutation.ts @@ -0,0 +1,37 @@ +import { useMutation, useQueryClient } from '@tanstack/react-query' +import { toast } from 'sonner' + +import { databaseTestsKeys } from './database-tests-key' +import { DatabaseTest } from './database-tests-query' + +type DatabaseTestDeleteVariables = { + projectRef: string + id: string +} + +export const useDatabaseTestDeleteMutation = ({ + onSuccess, +}: { + onSuccess?: () => void +} = {}) => { + const queryClient = useQueryClient() + + return useMutation({ + mutationFn: async ({ projectRef, id }: DatabaseTestDeleteVariables) => { + // This is a mock implementation. + // In a real scenario, you would delete the test file or remove it from a database table. + const key = databaseTestsKeys.list(projectRef) + const currentTests = queryClient.getQueryData(key) ?? [] + const updatedTests = currentTests.filter((test) => test.id !== id) + queryClient.setQueryData(key, updatedTests) + return id + }, + onSuccess: (id) => { + toast.success(`Test successfully deleted.`) + onSuccess?.() + }, + onError: (error: any) => { + toast.error(`Failed to delete test: ${error.message}`) + }, + }) +} diff --git a/apps/studio/data/database-tests/database-test-update-mutation.ts b/apps/studio/data/database-tests/database-test-update-mutation.ts new file mode 100644 index 00000000000..eada746ba4d --- /dev/null +++ b/apps/studio/data/database-tests/database-test-update-mutation.ts @@ -0,0 +1,48 @@ +import { useMutation, useQueryClient } from '@tanstack/react-query' +import { toast } from 'sonner' + +import { useSqlTitleGenerateMutation } from 'data/ai/sql-title-mutation' +import { databaseTestsKeys } from './database-tests-key' +import { DatabaseTest } from './database-tests-query' + +type DatabaseTestUpdateVariables = { + projectRef: string + id: string + query: string + name?: string +} + +export const useDatabaseTestUpdateMutation = ({ + onSuccess, +}: { + onSuccess?: () => void +} = {}) => { + const queryClient = useQueryClient() + const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation() + + return useMutation({ + mutationFn: async ({ projectRef, id, query, name }: DatabaseTestUpdateVariables) => { + // Generate AI title if name is not provided + const finalName = name ?? (await generateSqlTitle({ sql: query })).title + + // This is a mock implementation. + // In a real scenario, you would update the test in a file or a database table. + const key = databaseTestsKeys.list(projectRef) + const currentTests = queryClient.getQueryData(key) ?? [] + const updatedTests = currentTests.map((test) => + test.id === id ? { ...test, name: finalName, query } : test + ) + queryClient.setQueryData(key, updatedTests) + + const updatedTest = updatedTests.find((test) => test.id === id) + return updatedTest! + }, + onSuccess: (updatedTest) => { + toast.success(`Test "${updatedTest.name}" updated.`) + onSuccess?.() + }, + onError: (error: any) => { + toast.error(`Failed to update test: ${error.message}`) + }, + }) +} diff --git a/apps/studio/data/database-tests/database-tests-key.ts b/apps/studio/data/database-tests/database-tests-key.ts new file mode 100644 index 00000000000..36d8b2af575 --- /dev/null +++ b/apps/studio/data/database-tests/database-tests-key.ts @@ -0,0 +1,3 @@ +export const databaseTestsKeys = { + list: (projectRef: string | undefined) => [projectRef, 'database-tests'] as const, +} diff --git a/apps/studio/data/database-tests/database-tests-query.ts b/apps/studio/data/database-tests/database-tests-query.ts new file mode 100644 index 00000000000..759ab4e1648 --- /dev/null +++ b/apps/studio/data/database-tests/database-tests-query.ts @@ -0,0 +1,137 @@ +import { useQuery, UseQueryOptions } from '@tanstack/react-query' +import { databaseTestsKeys } from './database-tests-key' + +export type DatabaseTest = { + id: string + name: string + query: string +} + +export type DatabaseTestsVariables = { + projectRef?: string + connectionString?: string +} + +// This is a placeholder. Replace with your actual data fetching logic. +const MOCK_TESTS: DatabaseTest[] = [ + { + id: '0', + name: '000-setup-tests-hooks', + query: ` +-- install tests utilities +-- install pgtap extension for testing +create extension if not exists pgtap with schema extensions; +/* +--------------------- +---- install dbdev ---- +---------------------- +Requires: +- pg_tle: https://github.com/aws/pg_tle +- pgsql-http: https://github.com/pramsey/pgsql-http +*/ +create extension if not exists http with schema extensions; +create extension if not exists pg_tle; +drop extension if exists "supabase-dbdev"; +select pgtle.uninstall_extension_if_exists('supabase-dbdev'); +select + pgtle.install_extension( + 'supabase-dbdev', + resp.contents ->> 'version', + 'PostgreSQL package manager', + resp.contents ->> 'sql' + ) +from http( + ( + 'GET', + 'https://api.database.dev/rest/v1/' + || 'package_versions?select=sql,version' + || '&package_name=eq.supabase-dbdev' + || '&order=version.desc' + || '&limit=1', + array[ + ('apiKey', 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InhtdXB0cHBsZnZpaWZyYndtbXR2Iiwicm9sZSI6ImFub24iLCJpYXQiOjE2ODAxMDczNzIsImV4cCI6MTk5NTY4MzM3Mn0.z2CN0mvO2No8wSi46Gw59DFGCTJrzM0AQKsu_5k134s')::http_header + ], + null, + null + ) +) x, +lateral ( + select + ((row_to_json(x) -> 'content') #>> '{}')::json -> 0 +) resp(contents); +create extension "supabase-dbdev"; +select dbdev.install('supabase-dbdev'); +drop extension if exists "supabase-dbdev"; +create extension "supabase-dbdev"; +-- Install test helpers +select dbdev.install('basejump-supabase_test_helpers'); +create extension if not exists "basejump-supabase_test_helpers" version '0.0.6'; + `, + }, + { + id: '1', + name: '01_test_users_table', + query: ` +BEGIN; +SELECT plan(1); +SELECT has_table('users'); +SELECT * FROM finish(); +ROLLBACK; + `.trim(), + }, + { + id: '2', + name: '02_test_rls_policy', + query: ` +BEGIN; +SELECT plan(1); +SELECT policies_are('public', 'profiles', ARRAY['Profiles are public']); +SELECT * FROM finish(); +ROLLBACK; + `.trim(), + }, + { + id: '3', + name: '03_test_rls_public', + query: ` +begin; +select plan(1); +-- Verify RLS is enabled on all tables in the public schema +select tests.rls_enabled('public'); +select * from finish(); +rollback; + `, + }, +] + +export async function getDatabaseTests( + { projectRef, connectionString }: DatabaseTestsVariables, + signal?: AbortSignal +) { + if (!projectRef) throw new Error('projectRef is required') + + // In a real implementation, you would fetch this from a database or file system. + // For now, we return mock data. + // The timeout is to simulate network latency. + await new Promise((resolve) => setTimeout(resolve, 500)) + + return MOCK_TESTS +} + +export type DatabaseTestsData = Awaited> +export type DatabaseTestsError = unknown + +export const useDatabaseTestsQuery = ( + { projectRef, connectionString }: DatabaseTestsVariables, + { enabled = true, ...options }: UseQueryOptions = {} +) => + useQuery( + databaseTestsKeys.list(projectRef), + ({ signal }) => getDatabaseTests({ projectRef, connectionString }, signal), + { + staleTime: 0, + refetchOnWindowFocus: false, + refetchOnMount: true, + ...options, + } + ) diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index ac07b32bb10..50198d39d00 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -27,6 +27,7 @@ export const toolSetValidationSchema = z.record( 'get_advisors', // Local tools + 'get_policies', 'display_query', 'display_edge_function', 'rename_chat', @@ -89,11 +90,12 @@ export const TOOL_CATEGORY_MAP: Record = { rename_chat: TOOL_CATEGORIES.UI, search_docs: TOOL_CATEGORIES.UI, - // Schema tools - MCP + // Schema tools - MCP and local list_tables: TOOL_CATEGORIES.SCHEMA, list_extensions: TOOL_CATEGORIES.SCHEMA, list_edge_functions: TOOL_CATEGORIES.SCHEMA, list_branches: TOOL_CATEGORIES.SCHEMA, + get_policies: TOOL_CATEGORIES.SCHEMA, // Log tools - MCP and local get_advisors: TOOL_CATEGORIES.LOG, diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index fe57a9ed10d..b102795d0a7 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -8,6 +8,7 @@ import { IS_PLATFORM } from 'common' import { getOrganizations } from 'data/organizations/organizations-query' import { getProjects } from 'data/projects/projects-query' import { executeSql } from 'data/sql/execute-sql-query' +import { getDatabasePolicies } from 'data/database-policies/database-policies-query' import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi' import { getModel } from 'lib/ai/model' import apiWrapper from 'lib/api/apiWrapper' @@ -170,6 +171,56 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { return { status: 'Chat request sent to client' } }, }), + get_policies: tool({ + description: 'Get existing RLS policies for one or more schemas.', + parameters: z.object({ + schemas: z.array(z.string()).describe('The schema names to get the policies for.'), + }), + execute: async ({ schemas }) => { + const includeSchemaMetadata = aiOptInLevel !== 'disabled' + if (!includeSchemaMetadata) { + return 'You have disabled AI features that can read your schema. Please re-enable them to use this tool.' + } + + const headers = { + 'Content-Type': 'application/json', + ...(authorization && { Authorization: authorization }), + } + + try { + const data = await getDatabasePolicies( + { + projectRef, + connectionString, + schema: schemas?.join(','), + }, + undefined, + headers + ) + + if (data.length === 0) { + return `No policies found for the specified schemas: ${schemas.join(', ')}.` + } + + const formattedPolicies = data + .map( + (policy) => ` +- Policy Name: "${policy.name}" +Table: ${policy.schema}.${policy.table} +Action: ${policy.action} +Roles: ${policy.roles.join(', ')} +Command: ${policy.command} +Definition: ${policy.definition} +${policy.check ? `Check: ${policy.check}` : ''}` + ) + .join('\n') + + return `Here are the existing policies for schemas ${schemas.join(', ')}:\n${formattedPolicies}` + } catch (error: any) { + return `Failed to fetch policies: ${error.message}` + } + }, + }), } // Get a list of all schemas to add to context @@ -753,6 +804,84 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { CREATE INDEX idx_products_tenant ON products(tenant_id); \`\`\` + # Database testing with pgTAP + - You can also help the user write tests for their database. + - Supabase uses pgTAP to run database tests. + - Assume that the pgTAP extension is already installed, as well as the \`database.dev\` and \`basejump-supabase_test_helpers\` extensions. + - Tests you create can be saved for running periodically + - Always get schema information using \`list_tables\` + - If writing tests for RLS policies, first get RLS information using \`list_policies\` + + ## Writing tests + - Always wrap tests in a transaction (\`BEGIN;\` and \`ROLLBACK;\`) to isolate them from other tests. + - Use \`SELECT plan(N);\` where N is the number of tests you are running. + - Use test helpers to simplify the tests: + - \`tests.create_supabase_user('user_id')\`: Creates a test user. + - \`tests.authenticate_as('user_id')\`: Switches the context to the created user. + - \`tests.rls_enabled('schema_name')\`: Verifies that RLS is enabled on all tables in a schema. + - Here is an example of a test: + \`\`\`sql + begin; + -- Assuming 000-setup-tests-hooks.sql file is present to use tests helpers + select plan(4); + -- Set up test data + -- Create test supabase users + select tests.create_supabase_user('user1@test.com'); + select tests.create_supabase_user('user2@test.com'); + -- Create test data + insert into public.todos (task, user_id) values + ('User 1 Task 1', tests.get_supabase_uid('user1@test.com')), + ('User 1 Task 2', tests.get_supabase_uid('user1@test.com')), + ('User 2 Task 1', tests.get_supabase_uid('user2@test.com')); + -- Test as User 1 + select tests.authenticate_as('user1@test.com'); + -- Test 1: User 1 should only see their own todos + select results_eq( + 'select count(*) from todos', + ARRAY[2::bigint], + 'User 1 should only see their 2 todos' + ); + -- Test 2: User 1 can create their own todo + select lives_ok( + $$insert into todos (task, user_id) values ('New Task', tests.get_supabase_uid('user1@test.com'))$$, + 'User 1 can create their own todo' + ); + -- Test as User 2 + select tests.authenticate_as('user2@test.com'); + -- Test 3: User 2 should only see their own todos + select results_eq( + 'select count(*) from todos', + ARRAY[1::bigint], + 'User 2 should only see their 1 todo' + ); + -- Test 4: User 2 cannot modify User 1's todo + SELECT results_ne( + $$ update todos set task = 'Hacked!' where user_id = tests.get_supabase_uid('user1@test.com') returning 1 $$, + $$ values(1) $$, + 'User 2 cannot modify User 1 todos' + ); + select * from finish(); + rollback; + \`\`\` + + For checking row updates, first return the updated row and then check if the result is not equal or equal to the expected result. + e.g. SELECT results_ne( + $$ update todos set task = 'Hacked!' where user_id = tests.get_supabase_uid('user1@test.com') returning 1 $$, + $$ values(1) $$, + 'User 2 cannot modify User 1 todos' + + SELECT results_eq( + $$ + UPDATE public.messages + SET message = 'Hacked by user 1' + WHERE user_id = tests.get_supabase_uid('user1@test.com') + RETURNING 1 + $$, + $$ VALUES(1) $$, + 'User 1 can update their own messages' + ); + ); + # Edge Functions - Use the \`display_edge_function\` tool to generate complete, high-quality Edge Functions in TypeScript for the Deno runtime. - **Dependencies**: diff --git a/apps/studio/pages/project/[ref]/database/tests.tsx b/apps/studio/pages/project/[ref]/database/tests.tsx new file mode 100644 index 00000000000..4e1f1528064 --- /dev/null +++ b/apps/studio/pages/project/[ref]/database/tests.tsx @@ -0,0 +1,193 @@ +import { PermissionAction } from '@supabase/shared-types/out/constants' +import Link from 'next/link' +import { useRef, useState } from 'react' +import { Alert, Button, AiIconAnimation } from 'ui' +import { ButtonTooltip } from 'components/ui/ButtonTooltip' + +import { DatabaseTest } from 'data/database-tests/database-tests-query' +import { useDatabaseExtensionsQuery } from 'data/database-extensions/database-extensions-query' +import { Play } from 'lucide-react' +import TestsList, { TestsListHandle } from 'components/interfaces/Database/Tests/TestsList' +import DatabaseLayout from 'components/layouts/DatabaseLayout/DatabaseLayout' +import DefaultLayout from 'components/layouts/DefaultLayout' +import { PageLayout } from 'components/layouts/PageLayout/PageLayout' +import { ScaffoldContainer, ScaffoldSection } from 'components/layouts/Scaffold' +import NoPermission from 'components/ui/NoPermission' +import { useCheckPermissions, usePermissionsLoaded } from 'hooks/misc/useCheckPermissions' +import { useSelectedProject } from 'hooks/misc/useSelectedProject' +import { useAppStateSnapshot } from 'state/app-state' +import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state' +import type { NextPageWithLayout } from 'types' +import { toast } from 'sonner' +import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' +import { useDatabaseTestCreateMutation } from 'data/database-tests/database-test-create-mutation' +import { useDatabaseTestUpdateMutation } from 'data/database-tests/database-test-update-mutation' + +const DatabaseTestsPage: NextPageWithLayout = () => { + const { setEditorPanel } = useAppStateSnapshot() + const aiSnap = useAiAssistantStateSnapshot() + const project = useSelectedProject() + const [showRunAllTestsModal, setShowRunAllTestsModal] = useState(false) + const testsListRef = useRef(null) + + const canReadTests = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_READ, 'tables') + const canCreateTests = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'tables') + const isPermissionsLoaded = usePermissionsLoaded() + + const { data: extensions, isLoading: isLoadingExtensions } = useDatabaseExtensionsQuery({ + projectRef: project?.ref, + connectionString: project?.connectionString, + }) + const pgtap = extensions?.find((ext) => ext.name === 'pgtap') + const isPgtapInstalled = pgtap !== undefined && pgtap.installed_version !== null + + const { mutate: createTest, isLoading: isCreatingTest } = useDatabaseTestCreateMutation({ + onSuccess: () => { + setEditorPanel({ open: false }) + }, + }) + + const { mutate: updateTest, isLoading: isUpdatingTest } = useDatabaseTestUpdateMutation({ + onSuccess: () => { + setEditorPanel({ open: false }) + }, + }) + + const onSelectTest = (test: DatabaseTest) => { + setEditorPanel({ + open: true, + initialValue: test.query, + label: `Edit test: ${test.name}`, + saveLabel: 'Save test', + onSave: (query) => { + updateTest({ + projectRef: project?.ref!, + id: test.id, + query, + }) + }, + }) + } + + const createNewTest = () => { + setEditorPanel({ + open: true, + initialValue: ` +BEGIN; +SELECT plan(1); +-- Your test here +SELECT * FROM finish(); +ROLLBACK; + `.trim(), + label: 'Create new test', + saveLabel: 'Save test', + onSave: (query) => { + createTest({ + projectRef: project?.ref!, + connectionString: project?.connectionString, + query, + }) + }, + }) + } + + const onRunAllTests = () => { + setShowRunAllTestsModal(true) + } + + const primaryActions = ( +
+ + } + onClick={() => + aiSnap.newChat({ + name: 'Create a new test', + open: true, + initialInput: `Write a pgTAP test to...`, + }) + } + tooltip={{ + content: { + side: 'bottom', + text: !canCreateTests + ? 'You need additional permissions to create tests' + : 'Create with Supabase Assistant', + }, + }} + /> +
+ ) + + const secondaryActions = ( + + ) + + return ( + + + +
+ {!isLoadingExtensions && !isPgtapInstalled && ( + + + Enable pgTAP + + + } + > + The pgTAP extension is required to run database tests. Please enable it from the + database extensions page. + + )} + +
+ + + setShowRunAllTestsModal(false)} + onConfirm={() => { + testsListRef.current?.runAllTests() + setShowRunAllTestsModal(false) + }} + > +

+ Are you sure you want to run all the tests? This may take a while. +

+ + + ) +} + +DatabaseTestsPage.getLayout = (page) => ( + + {page} + +) + +export default DatabaseTestsPage diff --git a/apps/studio/state/app-state.ts b/apps/studio/state/app-state.ts index 9303cc4ed12..fccb1e0581d 100644 --- a/apps/studio/state/app-state.ts +++ b/apps/studio/state/app-state.ts @@ -15,6 +15,7 @@ type EditorPanelType = { label?: string saveLabel?: string onSave?: (value: string) => void + onRun?: (value: string) => void functionName?: string templates?: Template[] initialPrompt?: string diff --git a/packages/ui/src/components/shadcn/ui/table.tsx b/packages/ui/src/components/shadcn/ui/table.tsx index a812d5fbfd6..9445387ad3d 100644 --- a/packages/ui/src/components/shadcn/ui/table.tsx +++ b/packages/ui/src/components/shadcn/ui/table.tsx @@ -64,7 +64,7 @@ const TableHead = React.forwardRef<