fix: data table fixes

Signed-off-by: Varun Raj <varun@skcript.com>

package.json

@@ -19,6 +19,7 @@
     "@radix-ui/react-alert-dialog": "^1.1.4",
     "@radix-ui/react-avatar": "^1.1.0",
     "@radix-ui/react-checkbox": "^1.1.2",
+    "@radix-ui/react-context-menu": "^2.2.16",
     "@radix-ui/react-dialog": "^1.1.1",
     "@radix-ui/react-dropdown-menu": "^2.1.1",
     "@radix-ui/react-hover-card": "^1.1.1",

src/components/ui/context-menu.tsx

@@ -0,0 +1,197 @@
+import * as React from "react"
+import * as ContextMenuPrimitive from "@radix-ui/react-context-menu"
+import { cn } from "@/lib/utils"
+import { CheckIcon, ChevronRightIcon, DotFilledIcon } from "@radix-ui/react-icons"
+
+const ContextMenu = ContextMenuPrimitive.Root
+
+const ContextMenuTrigger = ContextMenuPrimitive.Trigger
+
+const ContextMenuGroup = ContextMenuPrimitive.Group
+
+const ContextMenuPortal = ContextMenuPrimitive.Portal
+
+const ContextMenuSub = ContextMenuPrimitive.Sub
+
+const ContextMenuRadioGroup = ContextMenuPrimitive.RadioGroup
+
+const ContextMenuSubTrigger = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.SubTrigger>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.SubTrigger> & {
+    inset?: boolean
+  }
+>(({ className, inset, children, ...props }, ref) => (
+  <ContextMenuPrimitive.SubTrigger
+    ref={ref}
+    className={cn(
+      "flex cursor-default select-none items-center rounded-sm px-2 py-1.5 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[state=open]:bg-accent data-[state=open]:text-accent-foreground",
+      inset && "pl-8",
+      className
+    )}
+    {...props}
+  >
+    {children}
+    <ChevronRightIcon className="ml-auto h-4 w-4" />
+  </ContextMenuPrimitive.SubTrigger>
+))
+ContextMenuSubTrigger.displayName = ContextMenuPrimitive.SubTrigger.displayName
+
+const ContextMenuSubContent = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.SubContent>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.SubContent>
+>(({ className, ...props }, ref) => (
+  <ContextMenuPrimitive.SubContent
+    ref={ref}
+    className={cn(
+      "z-50 min-w-[8rem] overflow-hidden rounded-md border bg-popover p-1 text-popover-foreground shadow-lg data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 origin-[--radix-context-menu-content-transform-origin]",
+      className
+    )}
+    {...props}
+  />
+))
+ContextMenuSubContent.displayName = ContextMenuPrimitive.SubContent.displayName
+
+const ContextMenuContent = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.Content>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.Content>
+>(({ className, ...props }, ref) => (
+  <ContextMenuPrimitive.Portal>
+    <ContextMenuPrimitive.Content
+      ref={ref}
+      className={cn(
+        "z-50 max-h-[--radix-context-menu-content-available-height] min-w-[8rem] overflow-y-auto overflow-x-hidden rounded-md border bg-popover p-1 text-popover-foreground shadow-md data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 origin-[--radix-context-menu-content-transform-origin]",
+        className
+      )}
+      {...props}
+    />
+  </ContextMenuPrimitive.Portal>
+))
+ContextMenuContent.displayName = ContextMenuPrimitive.Content.displayName
+
+const ContextMenuItem = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.Item>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.Item> & {
+    inset?: boolean
+  }
+>(({ className, inset, ...props }, ref) => (
+  <ContextMenuPrimitive.Item
+    ref={ref}
+    className={cn(
+      "relative flex cursor-default select-none items-center rounded-sm px-2 py-1.5 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50",
+      inset && "pl-8",
+      className
+    )}
+    {...props}
+  />
+))
+ContextMenuItem.displayName = ContextMenuPrimitive.Item.displayName
+
+const ContextMenuCheckboxItem = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.CheckboxItem>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.CheckboxItem>
+>(({ className, children, checked, ...props }, ref) => (
+  <ContextMenuPrimitive.CheckboxItem
+    ref={ref}
+    className={cn(
+      "relative flex cursor-default select-none items-center rounded-sm py-1.5 pl-8 pr-2 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50",
+      className
+    )}
+    checked={checked}
+    {...props}
+  >
+    <span className="absolute left-2 flex h-3.5 w-3.5 items-center justify-center">
+      <ContextMenuPrimitive.ItemIndicator>
+        <CheckIcon className="h-4 w-4" />
+      </ContextMenuPrimitive.ItemIndicator>
+    </span>
+    {children}
+  </ContextMenuPrimitive.CheckboxItem>
+))
+ContextMenuCheckboxItem.displayName =
+  ContextMenuPrimitive.CheckboxItem.displayName
+
+const ContextMenuRadioItem = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.RadioItem>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.RadioItem>
+>(({ className, children, ...props }, ref) => (
+  <ContextMenuPrimitive.RadioItem
+    ref={ref}
+    className={cn(
+      "relative flex cursor-default select-none items-center rounded-sm py-1.5 pl-8 pr-2 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50",
+      className
+    )}
+    {...props}
+  >
+    <span className="absolute left-2 flex h-3.5 w-3.5 items-center justify-center">
+      <ContextMenuPrimitive.ItemIndicator>
+        <DotFilledIcon className="h-4 w-4 fill-current" />
+      </ContextMenuPrimitive.ItemIndicator>
+    </span>
+    {children}
+  </ContextMenuPrimitive.RadioItem>
+))
+ContextMenuRadioItem.displayName = ContextMenuPrimitive.RadioItem.displayName
+
+const ContextMenuLabel = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.Label>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.Label> & {
+    inset?: boolean
+  }
+>(({ className, inset, ...props }, ref) => (
+  <ContextMenuPrimitive.Label
+    ref={ref}
+    className={cn(
+      "px-2 py-1.5 text-sm font-semibold text-foreground",
+      inset && "pl-8",
+      className
+    )}
+    {...props}
+  />
+))
+ContextMenuLabel.displayName = ContextMenuPrimitive.Label.displayName
+
+const ContextMenuSeparator = React.forwardRef<
+  React.ElementRef<typeof ContextMenuPrimitive.Separator>,
+  React.ComponentPropsWithoutRef<typeof ContextMenuPrimitive.Separator>
+>(({ className, ...props }, ref) => (
+  <ContextMenuPrimitive.Separator
+    ref={ref}
+    className={cn("-mx-1 my-1 h-px bg-border", className)}
+    {...props}
+  />
+))
+ContextMenuSeparator.displayName = ContextMenuPrimitive.Separator.displayName
+
+const ContextMenuShortcut = ({
+  className,
+  ...props
+}: React.HTMLAttributes<HTMLSpanElement>) => {
+  return (
+    <span
+      className={cn(
+        "ml-auto text-xs tracking-widest text-muted-foreground",
+        className
+      )}
+      {...props}
+    />
+  )
+}
+ContextMenuShortcut.displayName = "ContextMenuShortcut"
+
+export {
+  ContextMenu,
+  ContextMenuTrigger,
+  ContextMenuContent,
+  ContextMenuItem,
+  ContextMenuCheckboxItem,
+  ContextMenuRadioItem,
+  ContextMenuLabel,
+  ContextMenuSeparator,
+  ContextMenuShortcut,
+  ContextMenuGroup,
+  ContextMenuPortal,
+  ContextMenuSub,
+  ContextMenuSubContent,
+  ContextMenuSubTrigger,
+  ContextMenuRadioGroup,
+}

src/components/ui/data-table.tsx

@@ -22,6 +22,13 @@ import {
   DropdownMenuContent,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu"
+import {
+  ContextMenu,
+  ContextMenuContent,
+  ContextMenuItem,
+  ContextMenuSeparator,
+  ContextMenuTrigger,
+} from "@/components/ui/context-menu"
 import {
   TableRoot,
   TableBody,
@@ -35,190 +42,221 @@ interface DataTableProps<TData, TValue> {
   columns: ColumnDef<TData, TValue>[]
   data: TData[]
   onRowSelectionChange?: (selectedIds: string[]) => void
-  selectedIds?: string[]
   getRowId?: (row: TData) => string
   searchValue?: string
   onSearchChange?: (value: string) => void
 }
 
-export function DataTable<TData, TValue>({
-  columns,
-  data,
-  onRowSelectionChange,
-  selectedIds = [],
-  getRowId,
-  searchValue,
-  onSearchChange,
-}: DataTableProps<TData, TValue>) {
-  const [sorting, setSorting] = React.useState<SortingState>([])
-  const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>([])
-  const [columnVisibility, setColumnVisibility] = React.useState<VisibilityState>({})
-  
-  // Initialize row selection from selectedIds
-  const initialRowSelection = React.useMemo(() => {
-    const selection: Record<string, boolean> = {}
-    selectedIds.forEach(id => {
-      selection[id] = true
+export interface DataTableRef<TData> {
+  selectRows: (rowIds: string[]) => void
+  selectAllRows: () => void
+  deselectAllRows: () => void
+  getSelectedRows: () => string[]
+  clearSelection: () => void
+}
+
+const DataTable = React.forwardRef<DataTableRef<any>, DataTableProps<any, any>>(
+  ({ columns, data, onRowSelectionChange, getRowId, searchValue, onSearchChange }, ref) => {
+    const [sorting, setSorting] = React.useState<SortingState>([])
+    const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>([])
+    const [columnVisibility, setColumnVisibility] = React.useState<VisibilityState>({})
+    const [rowSelection, setRowSelection] = React.useState<Record<string, boolean>>({})
+
+    const table = useReactTable({
+      data,
+      columns,
+      getRowId: getRowId || ((row: any) => row.id),
+      onSortingChange: setSorting,
+      onColumnFiltersChange: setColumnFilters,
+      getCoreRowModel: getCoreRowModel(),
+      getPaginationRowModel: getPaginationRowModel(),
+      getSortedRowModel: getSortedRowModel(),
+      getFilteredRowModel: getFilteredRowModel(),
+      onColumnVisibilityChange: setColumnVisibility,
+      onRowSelectionChange: setRowSelection,
+      state: {
+        sorting,
+        columnFilters,
+        columnVisibility,
+        rowSelection,
+      },
+      // Add these options to prevent potential issues
+      enableRowSelection: true,
+      enableMultiRowSelection: true,
+      enableSorting: true,
+      enableFilters: true,
+      enableColumnFilters: true,
     })
-    return selection
-  }, [selectedIds])
-  
-  const [rowSelection, setRowSelection] = React.useState<Record<string, boolean>>(initialRowSelection)
 
-  const table = useReactTable({
-    data,
-    columns,
-    getRowId: getRowId || ((row: any) => row.id),
-    onSortingChange: setSorting,
-    onColumnFiltersChange: setColumnFilters,
-    getCoreRowModel: getCoreRowModel(),
-    getPaginationRowModel: getPaginationRowModel(),
-    getSortedRowModel: getSortedRowModel(),
-    getFilteredRowModel: getFilteredRowModel(),
-    onColumnVisibilityChange: setColumnVisibility,
-    onRowSelectionChange: setRowSelection,
-    state: {
-      sorting,
-      columnFilters,
-      columnVisibility,
-      rowSelection,
-    },
-    // Add these options to prevent potential issues
-    enableRowSelection: true,
-    enableMultiRowSelection: true,
-    enableSorting: true,
-    enableFilters: true,
-    enableColumnFilters: true,
-  })
+    // Expose methods via ref
+    React.useImperativeHandle(ref, () => ({
+      selectRows: (rowIds: string[]) => {
+        const newSelection: Record<string, boolean> = {}
+        rowIds.forEach(id => {
+          newSelection[id] = true
+        })
+        setRowSelection(newSelection)
+      },
+      selectAllRows: () => {
+        const allRowIds = table.getRowModel().rows.map(row => row.id)
+        const newSelection: Record<string, boolean> = {}
+        allRowIds.forEach(id => {
+          newSelection[id] = true
+        })
+        setRowSelection(newSelection)
+      },
+      deselectAllRows: () => {
+        setRowSelection({})
+      },
+      getSelectedRows: () => {
+        return Object.keys(rowSelection).filter(id => rowSelection[id])
+      },
+      clearSelection: () => {
+        setRowSelection({})
+      },
+    }), [table])
 
-  // Sync external selection with table selection
-  React.useEffect(() => {
-    const newRowSelection: Record<string, boolean> = {}
-    selectedIds.forEach((id) => {
-      newRowSelection[id] = true
-    })
-    setRowSelection(newRowSelection)
-  }, [selectedIds])
+    // Sync external search with table filter
+    React.useEffect(() => {
+      if (searchValue !== undefined) {
+        table.getColumn("albumName")?.setFilterValue(searchValue)
+      }
+    }, [searchValue, table])
 
-  // Sync external search with table filter
-  React.useEffect(() => {
-    if (searchValue !== undefined) {
-      table.getColumn("albumName")?.setFilterValue(searchValue)
+    // Notify parent of selection changes
+    React.useEffect(() => {
+      const selectedRowIds = Object.keys(rowSelection).filter(id => rowSelection[id])
+      onRowSelectionChange?.(selectedRowIds)
+    }, [rowSelection, onRowSelectionChange])
+
+    const handleShowAllColumns = () => {
+      table.toggleAllColumnsVisible(true)
     }
-  }, [searchValue, table])
 
-  // Notify parent of selection changes
-  React.useEffect(() => {
-    const selectedRowIds = Object.keys(rowSelection).filter(id => rowSelection[id])
-    onRowSelectionChange?.(selectedRowIds)
-  }, [rowSelection, onRowSelectionChange])
+    const handleHideAllColumns = () => {
+      table.toggleAllColumnsVisible(false)
+    }
 
-  return (
-    <div className="w-full">
-      <div className="flex items-center py-4">
-        <DropdownMenu>
-          <DropdownMenuTrigger asChild>
-            <Button variant="outline" className="ml-auto">
-              Columns <ChevronDown />
-            </Button>
-          </DropdownMenuTrigger>
-          <DropdownMenuContent align="end">
-            {table
-              .getAllColumns()
-              .filter((column) => column.getCanHide())
-              .map((column) => {
-                return (
-                  <DropdownMenuCheckboxItem
-                    key={column.id}
-                    className="capitalize"
-                    checked={column.getIsVisible()}
-                    onCheckedChange={(value) =>
-                      column.toggleVisibility(!!value)
-                    }
-                  >
-                    {column.id}
-                  </DropdownMenuCheckboxItem>
-                )
-              })}
-          </DropdownMenuContent>
-        </DropdownMenu>
-      </div>
-      <div className="overflow-hidden rounded-md border">
-        <TableRoot>
-          <TableHeader>
-            {table.getHeaderGroups().map((headerGroup) => (
-              <TableRow key={headerGroup.id}>
-                {headerGroup.headers.map((header) => {
-                  return (
-                    <TableHead key={header.id}>
-                      {header.isPlaceholder
-                        ? null
-                        : flexRender(
-                            header.column.columnDef.header,
-                            header.getContext()
-                          )}
-                    </TableHead>
-                  )
-                })}
-              </TableRow>
-            ))}
-          </TableHeader>
-          <TableBody>
-            {table.getRowModel().rows?.length ? (
-              table.getRowModel().rows.map((row) => (
-                <TableRow
-                  key={row.id}
-                  data-state={row.getIsSelected() && "selected"}
-                >
-                  {row.getVisibleCells().map((cell) => (
-                    <TableCell 
-                      key={cell.id}
-                    >
-                      {flexRender(
-                        cell.column.columnDef.cell,
-                        cell.getContext()
-                      )}
-                    </TableCell>
-                  ))}
+    return (
+      <div className="w-full">
+        <div className="overflow-hidden rounded-md border">
+          <TableRoot>
+            <TableHeader>
+              {table.getHeaderGroups().map((headerGroup) => (
+                <TableRow key={headerGroup.id}>
+                  {headerGroup.headers.map((header) => {
+                    return (
+                      <ContextMenu key={header.id}>
+                        <ContextMenuTrigger asChild>
+                          <TableHead className="cursor-context-menu">
+                            {header.isPlaceholder
+                              ? null
+                              : flexRender(
+                                  header.column.columnDef.header,
+                                  header.getContext()
+                                )}
+                          </TableHead>
+                        </ContextMenuTrigger>
+                        <ContextMenuContent>
+                          <ContextMenuItem
+                            onClick={() => header.column.toggleVisibility()}
+                          >
+                            {header.column.getIsVisible() ? "Hide Column" : "Show Column"}
+                          </ContextMenuItem>
+                          <ContextMenuSeparator />
+                          <ContextMenuItem onClick={handleShowAllColumns}>
+                            Show All Columns
+                          </ContextMenuItem>
+                          <ContextMenuItem onClick={handleHideAllColumns}>
+                            Hide All Columns
+                          </ContextMenuItem>
+                          <ContextMenuSeparator />
+                          <div className="px-2 py-1.5 text-sm font-medium text-muted-foreground">
+                            Column Visibility
+                          </div>
+                          {table
+                            .getAllColumns()
+                            .filter((column) => column.getCanHide())
+                            .map((column) => {
+                              return (
+                                <ContextMenuItem
+                                  key={column.id}
+                                  onClick={() => column.toggleVisibility()}
+                                  className="flex items-center justify-between"
+                                >
+                                  <span className="capitalize">{column.id}</span>
+                                  <div className={`w-2 h-2 rounded-full ${column.getIsVisible() ? 'bg-green-500' : 'bg-gray-300'}`} />
+                                </ContextMenuItem>
+                              )
+                            })}
+                        </ContextMenuContent>
+                      </ContextMenu>
+                    )
+                  })}
                 </TableRow>
-              ))
-            ) : (
-              <TableRow>
-                <TableCell
-                  colSpan={columns.length}
-                  className="h-24 text-center"
-                >
-                  No results.
-                </TableCell>
-              </TableRow>
-            )}
-          </TableBody>
-        </TableRoot>
-      </div>
-      <div className="flex items-center justify-end space-x-2 py-4">
-        <div className="flex-1 text-sm text-muted-foreground">
-          {table.getFilteredSelectedRowModel().rows.length} of{" "}
-          {table.getFilteredRowModel().rows.length} row(s) selected.
+              ))}
+            </TableHeader>
+            <TableBody>
+              {table.getRowModel().rows?.length ? (
+                table.getRowModel().rows.map((row) => (
+                  <TableRow
+                    key={row.id}
+                    data-state={row.getIsSelected() && "selected"}
+                  >
+                    {row.getVisibleCells().map((cell) => (
+                      <TableCell 
+                        key={cell.id}
+                      >
+                        {flexRender(
+                          cell.column.columnDef.cell,
+                          cell.getContext()
+                        )}
+                      </TableCell>
+                    ))}
+                  </TableRow>
+                ))
+              ) : (
+                <TableRow>
+                  <TableCell
+                    colSpan={columns.length}
+                    className="h-24 text-center"
+                  >
+                    No results.
+                  </TableCell>
+                </TableRow>
+              )}
+            </TableBody>
+          </TableRoot>
         </div>
-        <div className="space-x-2">
-          <Button
-            variant="outline"
-            size="sm"
-            onClick={() => table.previousPage()}
-            disabled={!table.getCanPreviousPage()}
-          >
-            Previous
-          </Button>
-          <Button
-            variant="outline"
-            size="sm"
-            onClick={() => table.nextPage()}
-            disabled={!table.getCanNextPage()}
-          >
-            Next
-          </Button>
+        <div className="flex items-center justify-end space-x-2 py-4">
+          <div className="flex-1 text-sm text-muted-foreground">
+            {table.getFilteredSelectedRowModel().rows.length} of{" "}
+            {table.getFilteredRowModel().rows.length} row(s) selected.
+          </div>
+          <div className="space-x-2">
+            <Button
+              variant="outline"
+              size="sm"
+              onClick={() => table.previousPage()}
+              disabled={!table.getCanPreviousPage()}
+            >
+              Previous
+            </Button>
+            <Button
+              variant="outline"
+              size="sm"
+              onClick={() => table.nextPage()}
+              disabled={!table.getCanNextPage()}
+            >
+              Next
+            </Button>
+          </div>
         </div>
       </div>
-    </div>
-  )
-} 
+    )
+  }
+)
+
+DataTable.displayName = "DataTable"
+
+export default DataTable

src/pages/albums/index.tsx

@@ -16,7 +16,7 @@ import { AlertDialog } from '@/components/ui/alert-dialog'
 import toast from 'react-hot-toast'
 import { useAlbumSelection } from '@/hooks'
 import FloatingBar from '@/components/shared/FloatingBar'
-import { DataTable } from '@/components/ui/data-table'
+import DataTable, { DataTableRef } from '@/components/ui/data-table'
 import { albumColumns } from '@/components/albums/AlbumTableColumns'
 
 const SORT_BY_OPTIONS = [
@@ -41,6 +41,8 @@ export default function AlbumListPage() {
   const [errorMessage, setErrorMessage] = useState('')
   const albumShareDialogRef = useRef<IAlbumShareDialogRef>(null);
   const [deleting, setDeleting] = useState(false)
+
+  const dataTableRef = useRef<DataTableRef<IAlbum>>(null)
   
   // Separate selection state for table view
   const [tableSelectedIds, setTableSelectedIds] = useState<string[]>([])
@@ -103,9 +105,9 @@ export default function AlbumListPage() {
     })
   }
 
-  const handleTableSelectionChange = (selectedRowIds: string[]) => {
+  const handleTableSelectionChange = useCallback((selectedRowIds: string[]) => {
     setTableSelectedIds(selectedRowIds)
-  }
+  }, [])
 
   const handleSelectAll = () => {
     if (viewMode === 'table') {
@@ -140,11 +142,11 @@ export default function AlbumListPage() {
           <DataTable
             columns={albumColumns}
             data={searchedAlbums}
-            selectedIds={currentSelectedIds}
             onRowSelectionChange={handleTableSelectionChange}
             getRowId={(row) => row.id}
             searchValue={search}
             onSearchChange={setSearch}
+            ref={dataTableRef}
           />
         </div>
       )