Use totalRows if available from params to validate MAX_EXPORT_ROW_COUNT (#41384)

* Use totalRows if available from params to validate MAX_EXPORT_ROW_COUNT

* Clean

apps/studio/components/grid/components/header/Header.tsx

@@ -1,6 +1,5 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { ArrowUp, ChevronDown, FileText, Trash } from 'lucide-react'
-import Link from 'next/link'
 import { ReactNode, useState } from 'react'
 import { toast } from 'sonner'
 
@@ -22,7 +21,6 @@ import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { useAsyncCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject'
-import { DOCS_URL } from 'lib/constants'
 import { RoleImpersonationState } from 'lib/role-impersonation'
 import {
   useRoleImpersonationStateSnapshot,
@@ -45,22 +43,6 @@ import { FilterPopover } from './filter/FilterPopover'
 import { formatRowsForCSV } from './Header.utils'
 import { SortPopover } from './sort/SortPopover'
 
-// [Joshen] CSV exports require this guard as a fail-safe if the table is
-// just too large for a browser to keep all the rows in memory before
-// exporting. Either that or export as multiple CSV sheets with max n rows each
-export const MAX_EXPORT_ROW_COUNT = 500000
-export const MAX_EXPORT_ROW_COUNT_MESSAGE = (
-  <>
-    Sorry! We're unable to support exporting row counts larger than{' '}
-    {MAX_EXPORT_ROW_COUNT.toLocaleString('en-US')} at the moment. Alternatively, you may consider
-    using{' '}
-    <Link href={`${DOCS_URL}/reference/cli/supabase-db-dump`} target="_blank">
-      pg_dump
-    </Link>{' '}
-    via our CLI instead.
-  </>
-)
-
 export type HeaderProps = {
   customHeader: ReactNode
   isRefetching: boolean

apps/studio/components/layouts/TableEditorLayout/ExportAllRows.tsx

@@ -1,16 +1,13 @@
+import { useQueryClient, type QueryClient } from '@tanstack/react-query'
 import saveAs from 'file-saver'
 import Papa from 'papaparse'
 import { useCallback, useState, type ReactNode } from 'react'
 
-import { useQueryClient, type QueryClient } from '@tanstack/react-query'
 import { IS_PLATFORM } from 'common'
-import {
-  MAX_EXPORT_ROW_COUNT,
-  MAX_EXPORT_ROW_COUNT_MESSAGE,
-} from 'components/grid/components/header/Header'
 import { parseSupaTable } from 'components/grid/SupabaseGrid.utils'
 import type { Filter, Sort, SupaTable } from 'components/grid/types'
 import { formatTableRowsToSQL } from 'components/interfaces/TableGridEditor/TableEntity.utils'
+import { InlineLink } from 'components/ui/InlineLink'
 import { ENTITY_TYPE } from 'data/entity-types/entity-type-constants'
 import type { Entity } from 'data/entity-types/entity-types-infinite-query'
 import { tableEditorKeys } from 'data/table-editor/keys'
@@ -18,6 +15,7 @@ import { getTableEditor, type TableEditorData } from 'data/table-editor/table-ed
 import { isTableLike } from 'data/table-editor/table-editor-types'
 import { fetchAllTableRows } from 'data/table-rows/table-rows-query'
 import { useStaticEffectEvent } from 'hooks/useStaticEffectEvent'
+import { DOCS_URL } from 'lib/constants'
 import type { RoleImpersonationState } from 'lib/role-impersonation'
 import { ConfirmationModal } from 'ui-patterns/Dialogs/ConfirmationModal'
 import {
@@ -34,6 +32,19 @@ import {
 } from './ExportAllRows.errors'
 import { useProgressToasts } from './ExportAllRows.progress'
 
+// [Joshen] CSV exports require this guard as a fail-safe if the table is
+// just too large for a browser to keep all the rows in memory before
+// exporting. Either that or export as multiple CSV sheets with max n rows each
+const MAX_EXPORT_ROW_COUNT = 500000
+const MAX_EXPORT_ROW_COUNT_MESSAGE = (
+  <p>
+    Sorry! We're unable to support exporting row counts larger than{' '}
+    {MAX_EXPORT_ROW_COUNT.toLocaleString('en-US')} at the moment. Alternatively, you may consider
+    using <InlineLink href={`${DOCS_URL}/reference/cli/supabase-db-dump`}>pg_dump</InlineLink> via
+    our CLI instead.
+  </p>
+)
+
 type OutputCallbacks = {
   convertToOutputFormat: (formattedRows: Record<string, unknown>[], table: SupaTable) => string
   convertToBlob: (str: string) => Blob
@@ -49,6 +60,7 @@ type FetchAllRowsParams = {
   filters?: Filter[]
   sorts?: Sort[]
   roleImpersonationState?: RoleImpersonationState
+  totalRows?: number
   startCallback?: () => void
   progressCallback?: (progress: number) => void
 } & OutputCallbacks
@@ -67,6 +79,7 @@ const fetchAllRows = async ({
   filters,
   sorts,
   roleImpersonationState,
+  totalRows,
   startCallback,
   progressCallback,
   convertToOutputFormat,
@@ -112,7 +125,14 @@ const fetchAllRows = async ({
     }
   }
 
-  if (isTableLike(table) && table.live_rows_estimate > MAX_EXPORT_ROW_COUNT) {
+  if (totalRows !== undefined) {
+    if (totalRows > MAX_EXPORT_ROW_COUNT) {
+      return {
+        status: 'error',
+        error: new TableTooLargeError(table.name, totalRows, MAX_EXPORT_ROW_COUNT),
+      }
+    }
+  } else if (isTableLike(table) && table.live_rows_estimate > MAX_EXPORT_ROW_COUNT) {
     return {
       status: 'error',
       error: new TableTooLargeError(table.name, table.live_rows_estimate, MAX_EXPORT_ROW_COUNT),
@@ -283,6 +303,7 @@ export const useExportAllRowsGeneric = (
               filters: params.filters,
               sorts: params.sorts,
               roleImpersonationState: params.roleImpersonationState,
+              totalRows: params.totalRows,
               startCallback: () => {
                 startProgressTracker({
                   id: entity.id,