diff --git a/apps/studio/components/interfaces/Explorer/NotebookEditor.tsx b/apps/studio/components/interfaces/Explorer/NotebookEditor.tsx index 0949b02db74..3b95637b049 100644 --- a/apps/studio/components/interfaces/Explorer/NotebookEditor.tsx +++ b/apps/studio/components/interfaces/Explorer/NotebookEditor.tsx @@ -116,13 +116,9 @@ export const NotebookEditor = () => { switch (cell._tag) { case 'markdown_cell': return - case 'database_cell': - return - case 'log_cell': - // [Joshen] Will eventually hook it up - return null + return } })} diff --git a/apps/studio/components/interfaces/Explorer/QueryCell/index.tsx b/apps/studio/components/interfaces/Explorer/QueryCell/index.tsx index 00f63c1c9d1..57c22e307aa 100644 --- a/apps/studio/components/interfaces/Explorer/QueryCell/index.tsx +++ b/apps/studio/components/interfaces/Explorer/QueryCell/index.tsx @@ -7,11 +7,19 @@ import { MoveCellDropdownContent } from '../MoveCellDropdownContent' import { QueryEditor } from '../QueryEditor' import { type QueryDisplay, type QueryResult } from '../types' import { SortableSection } from '@/components/ui/SortableSection' -import { type DatabaseCell as DatabaseCellSchema } from '@/data/content/notebooks/notebook-schema' +import { + type DatabaseCell as DatabaseCellSchema, + type LogCell as LogCellSchema, +} from '@/data/content/notebooks/notebook-schema' +import { untrustedLogSql } from '@/data/logs/safe-analytics-sql' +import { + createDefaultCellSource, + type CellSource, +} from '@/data/query-sources/query-source-registry' import { useCurrentNotebook, useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state' interface QueryCellProps { - cell: Snapshot + cell: Snapshot } /** @@ -20,9 +28,6 @@ interface QueryCellProps { * - Database selection logic * * QueryCell atm minimally supports running queries and rendering results - * - * [Joshen] TODO: handleUpdateCell might be able to shift into notebook-state, so component - * doesn't need to have context of the other cells */ type QueryCellUpdate = { sql: string } | { title: string } | { display: QueryDisplay } @@ -31,41 +36,88 @@ type QueryCellUpdate = { sql: string } | { title: string } | { display: QueryDis export const QueryCell = ({ cell }: QueryCellProps) => { const snap = useNotebooksStateSnapshot() const currentNotebook = useCurrentNotebook() - const cells = currentNotebook?.notebook.content?.cells ?? [] - const [sql, setSql] = useState(cell.unchecked_sql) + const { id, title: cellTitle, view, chart, unchecked_sql } = cell + const rowLimit = 'row_limit' in cell ? cell.row_limit : undefined + const source = + cell._tag === 'database_cell' + ? createDefaultCellSource('database') + : createDefaultCellSource('logs') + + const [sql, setSql] = useState(unchecked_sql) const [result, setResult] = useState() - const title = cell.title ?? 'Untitled snippet' + const title = cellTitle ?? 'Untitled snippet' const display: QueryDisplay = { - view: cell.view ?? 'table', - chart: cell.chart ? { ...cell.chart, y_columns: [...cell.chart.y_columns] } : undefined, + view: view ?? 'table', + chart: chart ? { ...chart, y_columns: [...chart.y_columns] } : undefined, + } + + const handleSourceChange = (source: CellSource) => { + const notebookId = currentNotebook?.notebook.id + if (!notebookId) return + + snap.updateCell({ + id: notebookId, + cellId: id, + updater: (candidate) => { + if (source.type === 'database' && candidate._tag === 'log_cell') { + const { _tag, time_range, unchecked_sql, ...rest } = candidate + return { + ...rest, + _tag: 'database_cell' as const, + row_limit: 100, + unchecked_sql: untrustedSql(unchecked_sql), + } + } + + if (source.type === 'logs' && candidate._tag === 'database_cell') { + const { _tag, row_limit, unchecked_sql, ...rest } = candidate + return { + ...rest, + _tag: 'log_cell' as const, + time_range: { + _tag: 'relative_time_range' as const, + unit: 'hour' as const, + amount: 1, + }, + unchecked_sql: untrustedLogSql(unchecked_sql), + } + } + + return candidate + }, + }) } const handleUpdateCell = (payload: QueryCellUpdate) => { const notebookId = currentNotebook?.notebook.id if (!notebookId) return - const nextCells = cells.map((candidate) => { - if (candidate.id !== cell.id || candidate._tag !== 'database_cell') return candidate + snap.updateCell({ + id: notebookId, + cellId: id, + updater: (candidate) => { + if (candidate._tag !== 'database_cell' && candidate._tag !== 'log_cell') return candidate - if ('sql' in payload) { - return { ...candidate, unchecked_sql: untrustedSql(payload.sql) } - } + if ('sql' in payload) { + return candidate._tag === 'database_cell' + ? { ...candidate, unchecked_sql: untrustedSql(payload.sql) } + : { ...candidate, unchecked_sql: untrustedLogSql(payload.sql) } + } - if ('title' in payload) { - const nextTitle = payload.title.trim() - return nextTitle ? { ...candidate, title: nextTitle } : candidate - } + if ('title' in payload) { + const nextTitle = payload.title.trim() + return nextTitle ? { ...candidate, title: nextTitle } : candidate + } - return { - ...candidate, - view: payload.display.view, - chart: payload.display.chart, - } + return { + ...candidate, + view: payload.display.view, + chart: payload.display.chart, + } + }, }) - - snap.updateCells({ id: notebookId, cells: nextCells }) } return ( @@ -76,18 +128,22 @@ export const QueryCell = ({ cell }: QueryCellProps) => { gripClassName="mt-2 opacity-0 group-hover:opacity-100 has-[[data-state=open]]:opacity-100 transition" > handleUpdateCell({ title })} onSqlChange={setSql} onSqlCommit={(sql) => handleUpdateCell({ sql })} + onSourceChange={handleSourceChange} onResultChange={setResult} - onDisplayChange={(display) => handleUpdateCell({ display })} + onDisplayChange={ + cell._tag === 'database_cell' ? (display) => handleUpdateCell({ display }) : undefined + } /> ) diff --git a/apps/studio/components/interfaces/Explorer/QueryEditor.tsx b/apps/studio/components/interfaces/Explorer/QueryEditor.tsx index 537e858d23d..8c79ef00777 100644 --- a/apps/studio/components/interfaces/Explorer/QueryEditor.tsx +++ b/apps/studio/components/interfaces/Explorer/QueryEditor.tsx @@ -46,7 +46,7 @@ export type QueryEditorProps = { sql: string source?: CellSource result?: QueryResult - rowLimit: number + rowLimit?: number display?: QueryDisplay toolbarActions?: ReactNode onTitleChange: (title: string) => void @@ -238,8 +238,12 @@ export const QueryEditor = ({

{(result?.rows ?? []).length.toLocaleString()} rows

-

·

-

Limit {rowLimit} rows

+ {rowLimit && ( + <> +

·

+

Limit {rowLimit} rows

+ + )} ) diff --git a/apps/studio/data/content/content-remap.test.ts b/apps/studio/data/content/content-remap.test.ts index 1c288ea7810..36f46481571 100644 --- a/apps/studio/data/content/content-remap.test.ts +++ b/apps/studio/data/content/content-remap.test.ts @@ -118,6 +118,7 @@ describe('remapSqlContentField', () => { _tag: 'database_cell', id: DATABASE_CELL.id, row_limit: 100, + view: 'table', unchecked_sql: untrustedSql(DATABASE_CELL.sql), }) expect(databaseCell).not.toHaveProperty('sql') @@ -125,6 +126,7 @@ describe('remapSqlContentField', () => { _tag: 'log_cell', id: LOG_CELL.id, time_range: LOG_CELL.time_range, + view: 'table', unchecked_sql: untrustedLogSql(LOG_CELL.sql), }) expect(logCell).not.toHaveProperty('sql') diff --git a/apps/studio/data/content/notebooks/notebook-schema.test.ts b/apps/studio/data/content/notebooks/notebook-schema.test.ts index f50bd20c9a6..c6fc30abbaa 100644 --- a/apps/studio/data/content/notebooks/notebook-schema.test.ts +++ b/apps/studio/data/content/notebooks/notebook-schema.test.ts @@ -225,6 +225,7 @@ describe('notebookDomainSchema', () => { _tag: 'database_cell', id: FULL_NOTEBOOK.cells[1].id, row_limit: 100, + view: 'table', unchecked_sql: untrustedSql('select * from auth.users limit 100'), }) expect(databaseCell).not.toHaveProperty('sql') diff --git a/apps/studio/data/content/notebooks/notebook-schema.ts b/apps/studio/data/content/notebooks/notebook-schema.ts index 545267b75e9..c326011f3b7 100644 --- a/apps/studio/data/content/notebooks/notebook-schema.ts +++ b/apps/studio/data/content/notebooks/notebook-schema.ts @@ -51,7 +51,7 @@ const databaseCellSchema = z.object({ title: z.string().optional(), sql: z.string(), row_limit: z.number(), - view: z.enum(['table', 'chart']).default('table').optional(), + view: z.enum(['table', 'chart']).optional(), chart: chartConfigSchema.optional(), }) @@ -61,6 +61,7 @@ const logCellSchema = z.object({ title: z.string().optional(), sql: z.string(), time_range: timeRangeSchema, + view: z.enum(['table', 'chart']).optional(), chart: chartConfigSchema.optional(), }) @@ -126,19 +127,19 @@ export const agentNotebookSchema = z.object({ export type AgentNotebook = z.infer export type AgentCell = z.infer -// The domain shape: parses the same wire cell (`cellSchema`) and transforms `sql` into a -// branded `unchecked_sql`. +// The domain shape: parses the same wire cell (`cellSchema`), transforms `sql` into a +// branded `unchecked_sql`, and defaults `view` to 'table' const cellDomainSchema = cellSchema.transform((cell) => { switch (cell._tag) { case 'markdown_cell': return cell case 'database_cell': { - const { sql, ...rest } = cell - return { ...rest, unchecked_sql: untrustedSql(sql) } + const { sql, view, ...rest } = cell + return { ...rest, view: view ?? 'table', unchecked_sql: untrustedSql(sql) } } case 'log_cell': { - const { sql, ...rest } = cell - return { ...rest, unchecked_sql: untrustedLogSql(sql) } + const { sql, view, ...rest } = cell + return { ...rest, view: view ?? 'table', unchecked_sql: untrustedLogSql(sql) } } } }) diff --git a/apps/studio/lib/ai/tools/notebook-tools.test.ts b/apps/studio/lib/ai/tools/notebook-tools.test.ts index 1231cbd509e..c6889a74a06 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.test.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.test.ts @@ -208,12 +208,14 @@ describe('ai/tools/notebook-tools', () => { _tag: 'database_cell', id: 'cell-2', row_limit: 100, + view: 'table', sql: 'select * from auth.users limit 100', }, { _tag: 'log_cell', id: 'cell-3', time_range: { _tag: 'relative_time_range', unit: 'hour', amount: 1 }, + view: 'table', sql: 'select timestamp, event_message from edge_logs limit 10', }, ], diff --git a/apps/studio/state/notebooks/notebooks-state.ts b/apps/studio/state/notebooks/notebooks-state.ts index cd0c2829696..14449cf9279 100644 --- a/apps/studio/state/notebooks/notebooks-state.ts +++ b/apps/studio/state/notebooks/notebooks-state.ts @@ -121,6 +121,30 @@ export const notebooksState = proxy({ notebooksState.updateCells({ id, cells: nextCells }) }, + /** + * Update a single cell in a notebook's cell array via an updater callback. + * The caller decides how the cell's content should change (e.g. field + * defaults, tag conversion) since that's a UI concern, not a state one — + * this only finds the cell by id and re-saves the array. + */ + updateCell: ({ + id, + cellId, + updater, + }: { + id: string + cellId: string + updater: (cell: Notebooks.Cell) => Notebooks.Cell + }) => { + const stateNotebook = notebooksState.notebooks[id] + if (!stateNotebook?.notebook.content) return + + const nextCells = stateNotebook.notebook.content.cells.map((cell) => + cell.id === cellId ? updater(cell) : cell + ) + notebooksState.updateCells({ id, cells: nextCells }) + }, + /** * Remove a single cell from a notebook's cell array. */