Joshenlim/fe 4175 add row limit in explorerquerysourcemenu (#49098)

## Context

Related to Explorer/Notebook - builds on the ExplorerQuerySourceMenu by
adding an option for row limit in both Notebook Query cell + Query Tab

## Side note RE persistence of the selected row limit value
Note that for QueryTab - its intentional that for now that the row limit
isn't scoped to the query draft atm as I wanna avoid making changes to
`explorer-query` atm as there's a couple of PRs in flux that touches
that file. So will handle that separately

^ This means that switching between query tabs will not change nor
persist the row limit

<img width="931" height="335" alt="image"
src="https://github.com/user-attachments/assets/d1d52ee7-7c1d-42aa-a6ae-1d7d99ab95c9"
/>

<img width="1381" height="486" alt="image"
src="https://github.com/user-attachments/assets/689368c9-c0fc-4000-a09e-f59bfa7afa97"
/>

## To test
- [ ] Verify that row limit behaviour works in notebooks
- [ ] Verify that row limit behaviour works in explorer query tab

apps/studio/components/interfaces/Explorer/ExplorerQuerySourceMenu.test.tsx

@@ -44,7 +44,12 @@ describe('ExplorerQuerySourceMenu', () => {
   ) =>
     customRender(
       <FeatureFlagContext.Provider value={{ configcat: flags, posthog: {}, hasLoaded: true }}>
-        <ExplorerQuerySourceMenu source={source} onSourceChange={vi.fn()} />
+        <ExplorerQuerySourceMenu
+          rowLimit={100}
+          onRowLimitChange={vi.fn()}
+          source={source}
+          onSourceChange={vi.fn()}
+        />
       </FeatureFlagContext.Provider>
     )
 
@@ -53,6 +58,8 @@ describe('ExplorerQuerySourceMenu', () => {
 
     customRender(
       <ExplorerQuerySourceMenu
+        rowLimit={100}
+        onRowLimitChange={vi.fn()}
         source={{
           _tag: 'logs',
           time_range: { _tag: 'relative_time_range', amount: 1, unit: 'hour' },
@@ -72,6 +79,8 @@ describe('ExplorerQuerySourceMenu', () => {
 
     customRender(
       <ExplorerQuerySourceMenu
+        rowLimit={100}
+        onRowLimitChange={vi.fn()}
         source={{
           _tag: 'logs',
           time_range: { _tag: 'relative_time_range', amount: 1, unit: 'hour' },

apps/studio/components/interfaces/Explorer/ExplorerQuerySourceMenu.tsx

@@ -9,6 +9,7 @@ import {
   DropdownMenuTrigger,
 } from 'ui'
 
+import { RowLimitSubMenu } from '../SQLEditor/UtilityPanel/QuerySourceMenu/RowLimitSubMenu'
 import { DatabaseParametersSubMenu } from '@/components/interfaces/QuerySources/DatabaseParametersSubMenu'
 import { LogsCustomRangeDialog } from '@/components/interfaces/QuerySources/LogsCustomRangeDialog'
 import { LogsTimeRangeSubMenu } from '@/components/interfaces/QuerySources/LogsTimeRangeSubMenu'
@@ -23,6 +24,8 @@ import {
 } from '@/data/query-sources/query-source-registry'
 
 export type ExplorerQuerySourceMenuProps = {
+  rowLimit?: number
+  onRowLimitChange?: (val: number) => void
   source: QuerySourceBinding
   onSourceChange: (source: QuerySourceBinding) => void
 }
@@ -37,6 +40,8 @@ export type ExplorerQuerySourceMenuProps = {
  * has SQL to preserve or discard and a fresh draft does not.
  */
 export const ExplorerQuerySourceMenu = ({
+  rowLimit = 100,
+  onRowLimitChange,
   source,
   onSourceChange,
 }: ExplorerQuerySourceMenuProps) => {
@@ -97,12 +102,20 @@ export const ExplorerQuerySourceMenu = ({
           <DropdownMenuSeparator />
 
           {source._tag === 'database' ? (
-            <DatabaseParametersSubMenu
-              identifier={source.database_identifier ?? ref}
-              onIdentifierChange={(database_identifier) =>
-                onSourceChange({ _tag: 'database', database_identifier })
-              }
-            />
+            <>
+              <DatabaseParametersSubMenu
+                identifier={source.database_identifier ?? ref}
+                onIdentifierChange={(database_identifier) =>
+                  onSourceChange({ _tag: 'database', database_identifier })
+                }
+              />
+              {onRowLimitChange !== undefined && (
+                <RowLimitSubMenu
+                  value={rowLimit}
+                  onValueChange={(val) => onRowLimitChange(Number(val))}
+                />
+              )}
+            </>
           ) : (
             <LogsTimeRangeSubMenu
               range={source.time_range}

apps/studio/components/interfaces/Explorer/QueryCell/QueryCell.utils.test.ts

@@ -6,6 +6,7 @@ import {
   cloneQueryCell,
   DEFAULT_CELL_ROW_LIMIT,
   getCellDisplay,
+  setCellRowLimit,
   setCellSql,
   toQueryModel,
 } from './QueryCell.utils'
@@ -122,6 +123,16 @@ describe('setCellSql', () => {
   })
 })
 
+describe('setCellRowLimit', () => {
+  it('writes the row limit onto a database cell without touching its query', () => {
+    expect(setCellRowLimit(DATABASE_CELL, 500)).toEqual({ ...DATABASE_CELL, row_limit: 500 })
+  })
+
+  it('leaves a log cell unchanged, since it has no row limit concept', () => {
+    expect(setCellRowLimit(LOG_CELL, 500)).toEqual(LOG_CELL)
+  })
+})
+
 describe('cloneQueryCell', () => {
   it('copies the chart series array rather than aliasing it', () => {
     const clone = cloneQueryCell(DATABASE_CELL)

apps/studio/components/interfaces/Explorer/QueryCell/QueryCell.utils.ts

@@ -117,6 +117,22 @@ export function setCellSql(cell: Snapshot<QueryCell>, sql: string): QueryCell {
   }
 }
 
+/**
+ * Writes a new row limit onto a database cell. A log cell has no row limit concept, so it
+ * passes through unchanged.
+ */
+export function setCellRowLimit(cell: Snapshot<QueryCell>, rowLimit: number): QueryCell {
+  if (cell._tag === 'log_cell') return cloneQueryCell(cell)
+
+  return {
+    ...copyQueryCellBase(cell),
+    _tag: 'database_cell',
+    unchecked_sql: cell.unchecked_sql,
+    row_limit: rowLimit,
+    database_identifier: cell.database_identifier,
+  }
+}
+
 /**
  * Builds the editor's query model from a cell and the editor's live text buffer. Branding
  * the buffer is the editor boundary the safe-SQL model expects; which brand applies is

apps/studio/components/interfaces/Explorer/QueryCell/index.tsx

@@ -10,6 +10,7 @@ import {
   cloneChartConfig,
   cloneQueryCell,
   getCellDisplay,
+  setCellRowLimit,
   setCellSql,
   toQueryModel,
 } from './QueryCell.utils'
@@ -78,6 +79,9 @@ export const QueryCell = ({ cell }: QueryCellProps) => {
       chart: cloneChartConfig(display.chart),
     }))
 
+  const handleRowLimitChange = (rowLimit: number) =>
+    updateQueryCell((candidate) => setCellRowLimit(candidate, rowLimit))
+
   return (
     <SortableSection
       id={cell.id}
@@ -97,6 +101,7 @@ export const QueryCell = ({ cell }: QueryCellProps) => {
         onSqlCommit={handleSqlCommit}
         onSourceChange={handleSourceChange}
         onResultChange={setResult}
+        onRowLimitChange={handleRowLimitChange}
         onDisplayChange={handleDisplayChange}
       />
     </SortableSection>

apps/studio/components/interfaces/Explorer/QueryEditor.tsx

@@ -76,6 +76,7 @@ export type QueryEditorProps = {
   onSqlCommit?: (sql: string) => void
   onSourceChange?: (source: QuerySourceBinding) => void
   onResultChange: (result: QueryResult) => void
+  onRowLimitChange?: (val: number) => void
   onDisplayChange?: (display: QueryDisplay) => void
 }
 
@@ -97,6 +98,7 @@ export const QueryEditor = ({
   onSqlCommit,
   onSourceChange,
   onResultChange,
+  onRowLimitChange,
   onDisplayChange,
 }: QueryEditorProps) => {
   const sql = query.uncheckedSql
@@ -205,6 +207,8 @@ export const QueryEditor = ({
             <ExplorerQuerySourceMenu
               source={toQuerySourceBinding(query)}
               onSourceChange={onSourceChange}
+              rowLimit={rowLimit}
+              onRowLimitChange={onRowLimitChange}
             />
           )}
           {display && onDisplayChange && (
@@ -267,7 +271,7 @@ export const QueryEditor = ({
         {rowLimit && (
           <>
             <p>·</p>
-            <p>Limit {rowLimit} rows</p>
+            <p>{rowLimit < 0 ? 'No row limit' : `Limit ${rowLimit} rows`}</p>
           </>
         )}
       </ExplorerQueryFooter>

apps/studio/components/interfaces/Explorer/QueryTab.tsx

@@ -10,15 +10,16 @@ import { toQuerySourceBinding } from '@/data/query-sources/query-source-registry
 import { explorerQueryState, useExplorerQueryStateSnapshot } from '@/state/explorer-query'
 import { createTabId, TabsStateContext } from '@/state/tabs'
 
-const QUERY_ROW_LIMIT = 100
-
 /** Query-tab lifecycle adapter around the shared QueryEditor. */
 export const QueryTab = () => {
   const { id, ref } = useParams()
   const router = useRouter()
   const tabs = useContext(TabsStateContext)
   const querySnap = useExplorerQueryStateSnapshot()
+
+  const [rowLimit, setRowLimit] = useState<number>(100)
   const [restoredQueryKey, setRestoredQueryKey] = useState<string>()
+
   const stateDraft = id ? querySnap.drafts[id] : undefined
   const draft = stateDraft?.projectRef === ref ? stateDraft : undefined
   const result = draft && id ? querySnap.results[id] : undefined
@@ -81,7 +82,7 @@ export const QueryTab = () => {
       : {
           ...toQuerySourceBinding(draft),
           uncheckedSql: draft.uncheckedSql,
-          rowLimit: QUERY_ROW_LIMIT,
+          rowLimit,
         }
 
   return (
@@ -99,6 +100,7 @@ export const QueryTab = () => {
       onSqlChange={(sql) => explorerQueryState.updateDraft({ id, sql })}
       onSourceChange={(source) => explorerQueryState.updateDraft({ id, source })}
       onResultChange={handleResultChange}
+      onRowLimitChange={setRowLimit}
     />
   )
 }

apps/studio/components/interfaces/SQLEditor/UtilityPanel/QuerySourceMenu/QuerySourceMenu.tsx

@@ -149,7 +149,10 @@ export const QuerySourceMenu = ({ id, runSource, canCreateLogsSnippet }: QuerySo
                 />
               )}
               <RunAsSubMenu />
-              <RowLimitSubMenu />
+              <RowLimitSubMenu
+                value={sessionSnap.limit}
+                onValueChange={(val) => sessionSnap.setLimit(Number(val))}
+              />
             </>
           )}
         </DropdownMenuContent>

apps/studio/components/interfaces/SQLEditor/UtilityPanel/QuerySourceMenu/RowLimitSubMenu.tsx

@@ -7,11 +7,14 @@ import {
 } from 'ui'
 
 import { ROWS_PER_PAGE_OPTIONS } from '../../SQLEditor.constants'
-import { useSqlEditorSessionSnapshot } from '@/state/sql-editor/sql-editor-session-state'
 
-export const RowLimitSubMenu = () => {
-  const sessionSnap = useSqlEditorSessionSnapshot()
-  const currentLabel = ROWS_PER_PAGE_OPTIONS.find((opt) => opt.value === sessionSnap.limit)?.label
+interface RowLimitSubMenuProps {
+  value: number
+  onValueChange: (value: string) => void
+}
+
+export const RowLimitSubMenu = ({ value, onValueChange }: RowLimitSubMenuProps) => {
+  const currentLabel = ROWS_PER_PAGE_OPTIONS.find((opt) => opt.value === value)?.label
 
   return (
     <DropdownMenuSub>
@@ -22,10 +25,7 @@ export const RowLimitSubMenu = () => {
         </div>
       </DropdownMenuSubTrigger>
       <DropdownMenuSubContent className="w-40">
-        <DropdownMenuRadioGroup
-          value={sessionSnap.limit.toString()}
-          onValueChange={(val) => sessionSnap.setLimit(Number(val))}
-        >
+        <DropdownMenuRadioGroup value={value.toString()} onValueChange={onValueChange}>
           {ROWS_PER_PAGE_OPTIONS.map((option) => (
             <DropdownMenuRadioItem key={option.label} value={option.value.toString()}>
               {option.label}