Factor in filters and sorts for export via CLI option (#37399)

* Factor in filters and sorts for export via CLI option

* Add align

* Update

* Smol fix

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

@@ -1,8 +1,13 @@
+import { useState } from 'react'
+
 import { useParams } from 'common'
+import { Filter, Sort, SupaTable } from 'components/grid/types'
 import { getConnectionStrings } from 'components/interfaces/Connect/DatabaseSettings.utils'
 import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
+import { getAllTableRowsSql } from 'data/table-rows/table-rows-query'
 import { pluckObjectFields } from 'lib/helpers'
-import { useState } from 'react'
+import { RoleImpersonationState, wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { useRoleImpersonationStateSnapshot } from 'state/role-impersonation-state'
 import {
   Button,
   cn,
@@ -22,13 +27,25 @@ import {
 import { Admonition } from 'ui-patterns'
 
 interface ExportDialogProps {
-  table?: { name: string; schema: string }
+  table?: SupaTable
+  filters?: Filter[]
+  sorts?: Sort[]
+  ignoreRoleImpersonation?: boolean
+
   open: boolean
   onOpenChange: (open: boolean) => void
 }
 
-export const ExportDialog = ({ table, open, onOpenChange }: ExportDialogProps) => {
+export const ExportDialog = ({
+  table,
+  filters = [],
+  sorts = [],
+  ignoreRoleImpersonation = false,
+  open,
+  onOpenChange,
+}: ExportDialogProps) => {
   const { ref: projectRef } = useParams()
+  const roleImpersonationState = useRoleImpersonationStateSnapshot()
 
   const [selectedTab, setSelectedTab] = useState<string>('csv')
 
@@ -48,10 +65,18 @@ export const ExportDialog = ({ table, open, onOpenChange }: ExportDialogProps) =
   })
 
   const outputName = `${table?.name}_rows`
+  const queryChains = !table ? undefined : getAllTableRowsSql({ table, sorts, filters })
+  const query = !!queryChains
+    ? ignoreRoleImpersonation
+      ? queryChains.toSql()
+      : wrapWithRoleImpersonation(
+          queryChains.toSql(),
+          roleImpersonationState as RoleImpersonationState
+        )
+    : ''
 
   const csvExportCommand = `
-${connectionStrings.direct.psql} -c "COPY (SELECT * FROM "${table?.schema}"."${table?.name}") TO STDOUT WITH CSV HEADER DELIMITER ',';" > ${outputName}.csv
-`.trim()
+${connectionStrings.direct.psql} -c "COPY (${query}) TO STDOUT WITH CSV HEADER DELIMITER ',';" > ${outputName}.csv`.trim()
 
   const sqlExportCommand = `
 pg_dump -h ${db_host} -p ${db_port} -d ${db_name} -U ${db_user} --table="${table?.schema}.${table?.name}" --data-only --column-inserts > ${outputName}.sql
@@ -95,6 +120,12 @@ pg_dump -h ${db_host} -p ${db_port} -d ${db_name} -U ${db_user} --table="${table
                 value={sqlExportCommand}
                 className="[&_code]:text-[12px] [&_code]:text-foreground"
               />
+              <Admonition
+                type="note"
+                className="mt-2"
+                title="Filters are not supported when exporting as SQL via pg_dump"
+                description="If you'd like to export as SQL, we recommend creating a view first then exporting the data from there via pg_dump instead"
+              />
             </TabsContent_Shadcn_>
           </Tabs_Shadcn_>
 
@@ -107,15 +138,11 @@ pg_dump -h ${db_host} -p ${db_port} -d ${db_name} -U ${db_user} --table="${table
           </p>
 
           {selectedTab === 'sql' && (
-            <Admonition
-              type="note"
-              title="The pg_dump version needs to match your Postgres version"
-            >
-              <p className="!leading-normal">
-                If you run into a server version mismatch error, you will need to update{' '}
-                <code>pg_dump</code> before running the command.
-              </p>
-            </Admonition>
+            <p className="text-sm text-foreground-light">
+              Note: <code>pg_dump</code> needs to match your project's Postgres version. If you run
+              into a server version mismatch error, you will need to update <code>pg_dump</code>{' '}
+              before running the command.
+            </p>
           )}
         </DialogSection>
         <DialogFooter>

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

@@ -504,11 +504,9 @@ const RowHeader = () => {
               Export
             </Button>
           </DropdownMenuTrigger>
-          <DropdownMenuContent className={snap.allRowsSelected ? 'w-52' : 'w-40'}>
+          <DropdownMenuContent align="start" className={snap.allRowsSelected ? 'w-52' : 'w-40'}>
             <DropdownMenuItem onClick={onRowsExportCSV}>Export as CSV</DropdownMenuItem>
             <DropdownMenuItem onClick={onRowsExportSQL}>Export as SQL</DropdownMenuItem>
-            {/* [Joshen] Should make this available for all cases, but that'll involve updating
-            the Dialog's SQL output to be dynamic based on any filters applied */}
             {snap.allRowsSelected ? (
               <DropdownMenuItem className="group" onClick={() => setShowExportModal(true)}>
                 <div>
@@ -535,7 +533,9 @@ const RowHeader = () => {
       </div>
 
       <ExportDialog
-        table={{ name: snap.table.name, schema: snap.table.schema ?? '' }}
+        table={snap.table}
+        filters={filters}
+        sorts={sorts}
         open={showExportModal}
         onOpenChange={() => setShowExportModal(false)}
       />

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

@@ -1,5 +1,5 @@
 import saveAs from 'file-saver'
-import { Clipboard, Copy, Download, Edit, Lock, MoreHorizontal, Trash } from 'lucide-react'
+import { Clipboard, Copy, Download, Edit, Lock, MoreVertical, Trash } from 'lucide-react'
 import Link from 'next/link'
 import Papa from 'papaparse'
 import { toast } from 'sonner'
@@ -30,6 +30,7 @@ import { useTableEditorStateSnapshot } from 'state/table-editor'
 import { createTabId, useTabsStateSnapshot } from 'state/tabs'
 import {
   Badge,
+  Button,
   cn,
   copyToClipboard,
   DropdownMenu,
@@ -242,7 +243,7 @@ const EntityListItem: ItemRenderer<Entity, EntityListItemProps> = ({
           isOpened: isOpened && !isPreview,
           isPreview,
         }),
-        'px-4'
+        'pl-4 pr-1'
       )}
       onDoubleClick={(e) => {
         e.preventDefault()
@@ -285,8 +286,16 @@ const EntityListItem: ItemRenderer<Entity, EntityListItemProps> = ({
 
         {canEdit && (
           <DropdownMenu>
-            <DropdownMenuTrigger className="text-foreground-lighter transition-all text-transparent group-hover:text-foreground data-[state=open]:text-foreground">
-              <MoreHorizontal size={14} strokeWidth={2} />
+            <DropdownMenuTrigger
+              asChild
+              className="text-foreground-lighter transition-all text-transparent group-hover:text-foreground data-[state=open]:text-foreground"
+            >
+              <Button
+                type="text"
+                className="w-6 h-6"
+                icon={<MoreVertical size={14} strokeWidth={2} />}
+                onClick={(e) => e.preventDefault()}
+              />
             </DropdownMenuTrigger>
             <DropdownMenuContent side="bottom" align="start" className="w-44">
               <DropdownMenuItem

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

@@ -5,6 +5,8 @@ import { useEffect, useMemo, useState } from 'react'
 import { useParams } from 'common'
 import { useBreakpoint } from 'common/hooks/useBreakpoint'
 import { ExportDialog } from 'components/grid/components/header/ExportDialog'
+import { parseSupaTable } from 'components/grid/SupabaseGrid.utils'
+import { SupaTable } from 'components/grid/types'
 import { ProtectedSchemaWarning } from 'components/interfaces/Database/ProtectedSchemaWarning'
 import EditorMenuListSkeleton from 'components/layouts/TableEditorLayout/EditorMenuListSkeleton'
 import AlertError from 'components/ui/AlertError'
@@ -13,7 +15,7 @@ import InfiniteList from 'components/ui/InfiniteList'
 import SchemaSelector from 'components/ui/SchemaSelector'
 import { ENTITY_TYPE } from 'data/entity-types/entity-type-constants'
 import { useEntityTypesQuery } from 'data/entity-types/entity-types-infinite-query'
-import { useTableEditorQuery } from 'data/table-editor/table-editor-query'
+import { getTableEditor, useTableEditorQuery } from 'data/table-editor/table-editor-query'
 import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useLocalStorage } from 'hooks/misc/useLocalStorage'
 import { useQuerySchemaState } from 'hooks/misc/useSchemaQueryState'
@@ -40,14 +42,14 @@ import EntityListItem from './EntityListItem'
 import { TableMenuEmptyState } from './TableMenuEmptyState'
 
 export const TableEditorMenu = () => {
-  const { id: _id } = useParams()
+  const { id: _id, ref: projectRef } = useParams()
   const id = _id ? Number(_id) : undefined
   const snap = useTableEditorStateSnapshot()
   const { selectedSchema, setSelectedSchema } = useQuerySchemaState()
   const isMobile = useBreakpoint()
 
   const [searchText, setSearchText] = useState<string>('')
-  const [tableToExport, setTableToExport] = useState<{ name: string; schema: string }>()
+  const [tableToExport, setTableToExport] = useState<SupaTable>()
   const [visibleTypes, setVisibleTypes] = useState<string[]>(Object.values(ENTITY_TYPE))
   const [sort, setSort] = useLocalStorage<'alphabetical' | 'grouped-alphabetical'>(
     'table-editor-sort',
@@ -93,13 +95,24 @@ export const TableEditorMenu = () => {
     id,
   })
 
+  const tableEditorTabsCleanUp = useTableEditorTabsCleanUp()
+
+  const onSelectExportCLI = async (id: number) => {
+    const table = await getTableEditor({
+      id: id,
+      projectRef,
+      connectionString: project?.connectionString,
+    })
+    const supaTable = table && parseSupaTable(table)
+    setTableToExport(supaTable)
+  }
+
   useEffect(() => {
     if (selectedTable?.schema) {
       setSelectedSchema(selectedTable.schema)
     }
   }, [selectedTable?.schema])
 
-  const tableEditorTabsCleanUp = useTableEditorTabsCleanUp()
   useEffect(() => {
     // Clean up tabs + recent items for any tables that might have been removed outside of the dashboard session
     if (entityTypes && !searchText) {
@@ -255,11 +268,7 @@ export const TableEditorMenu = () => {
                       projectRef: project?.ref!,
                       id: Number(id),
                       isSchemaLocked,
-                      onExportCLI: () => {
-                        const entity = entityTypes?.find((x) => x.id === id)
-                        if (!entity) return
-                        setTableToExport({ name: entity.name, schema: entity.schema })
-                      },
+                      onExportCLI: () => onSelectExportCLI(Number(id)),
                     }}
                     getItemSize={() => 28}
                     hasNextPage={hasNextPage}
@@ -274,6 +283,7 @@ export const TableEditorMenu = () => {
       </div>
 
       <ExportDialog
+        ignoreRoleImpersonation
         table={tableToExport}
         open={!!tableToExport}
         onOpenChange={(open) => {
@@ -283,5 +293,3 @@ export const TableEditorMenu = () => {
     </>
   )
 }
-
-export default TableEditorMenu

apps/studio/data/table-rows/table-rows-query.ts

@@ -67,35 +67,15 @@ export async function executeWithRetry<T>(
   throw new Error('Max retries reached without success')
 }
 
-// TODO: fetchAllTableRows is used for CSV export, but since it doesn't actually truncate anything, (compare to getTableRows)
-// this is not suitable and will cause crashes on the pg-meta side given big tables
-// (either when the number of rows exceeds Blob size or if the columns in the rows are too large).
-// We should handle those errors gracefully, maybe adding a hint to the user about how to extract
-// the CSV to their machine via a direct command line connection (e.g., pg_dump), which will be much more
-// reliable for large data extraction.
-export const fetchAllTableRows = async ({
-  projectRef,
-  connectionString,
+export const getAllTableRowsSql = ({
   table,
   filters = [],
   sorts = [],
-  roleImpersonationState,
-  progressCallback,
 }: {
-  projectRef: string
-  connectionString?: string | null
   table: SupaTable
   filters?: Filter[]
   sorts?: Sort[]
-  roleImpersonationState?: RoleImpersonationState
-  progressCallback?: (value: number) => void
 }) => {
-  if (IS_PLATFORM && !connectionString) {
-    console.error('Connection string is required')
-    return []
-  }
-
-  const rows: any[] = []
   const query = new Query()
 
   const arrayBasedColumns = table.columns
@@ -129,6 +109,40 @@ export const fetchAllTableRows = async ({
     })
   }
 
+  return queryChains
+}
+
+// TODO: fetchAllTableRows is used for CSV export, but since it doesn't actually truncate anything, (compare to getTableRows)
+// this is not suitable and will cause crashes on the pg-meta side given big tables
+// (either when the number of rows exceeds Blob size or if the columns in the rows are too large).
+// We should handle those errors gracefully, maybe adding a hint to the user about how to extract
+// the CSV to their machine via a direct command line connection (e.g., pg_dump), which will be much more
+// reliable for large data extraction.
+export const fetchAllTableRows = async ({
+  projectRef,
+  connectionString,
+  table,
+  filters = [],
+  sorts = [],
+  roleImpersonationState,
+  progressCallback,
+}: {
+  projectRef: string
+  connectionString?: string | null
+  table: SupaTable
+  filters?: Filter[]
+  sorts?: Sort[]
+  roleImpersonationState?: RoleImpersonationState
+  progressCallback?: (value: number) => void
+}) => {
+  if (IS_PLATFORM && !connectionString) {
+    console.error('Connection string is required')
+    return []
+  }
+
+  const rows: any[] = []
+  const queryChains = getAllTableRowsSql({ table, sorts, filters })
+
   const rowsPerPage = 500
   const THROTTLE_DELAY = 500
 

apps/studio/pages/project/[ref]/editor/[id].tsx

@@ -6,7 +6,7 @@ import DefaultLayout from 'components/layouts/DefaultLayout'
 import { EditorBaseLayout } from 'components/layouts/editors/EditorBaseLayout'
 import { useProjectContext } from 'components/layouts/ProjectLayout/ProjectContext'
 import TableEditorLayout from 'components/layouts/TableEditorLayout/TableEditorLayout'
-import TableEditorMenu from 'components/layouts/TableEditorLayout/TableEditorMenu'
+import { TableEditorMenu } from 'components/layouts/TableEditorLayout/TableEditorMenu'
 import { useTableEditorQuery } from 'data/table-editor/table-editor-query'
 import { createTabId, useTabsStateSnapshot } from 'state/tabs'
 import type { NextPageWithLayout } from 'types'

apps/studio/pages/project/[ref]/editor/index.tsx

@@ -6,7 +6,7 @@ import SidePanelEditor from 'components/interfaces/TableGridEditor/SidePanelEdit
 import DefaultLayout from 'components/layouts/DefaultLayout'
 import { EditorBaseLayout } from 'components/layouts/editors/EditorBaseLayout'
 import TableEditorLayout from 'components/layouts/TableEditorLayout/TableEditorLayout'
-import TableEditorMenu from 'components/layouts/TableEditorLayout/TableEditorMenu'
+import { TableEditorMenu } from 'components/layouts/TableEditorLayout/TableEditorMenu'
 import { NewTab } from 'components/layouts/Tabs/NewTab'
 import { useAppStateSnapshot } from 'state/app-state'
 import { editorEntityTypes, useTabsStateSnapshot } from 'state/tabs'

apps/studio/pages/project/[ref]/editor/new.tsx

@@ -5,7 +5,7 @@ import { SidePanelEditor } from 'components/interfaces/TableGridEditor'
 import DefaultLayout from 'components/layouts/DefaultLayout'
 import { EditorBaseLayout } from 'components/layouts/editors/EditorBaseLayout'
 import TableEditorLayout from 'components/layouts/TableEditorLayout/TableEditorLayout'
-import TableEditorMenu from 'components/layouts/TableEditorLayout/TableEditorMenu'
+import { TableEditorMenu } from 'components/layouts/TableEditorLayout/TableEditorMenu'
 import { NewTab } from 'components/layouts/Tabs/NewTab'
 import type { NextPageWithLayout } from 'types'