feat(ui): add expandable row support to DataTable (#9940)

ui/app/(prowler)/demo-expandable-table/page.tsx

@@ -0,0 +1,214 @@
+"use client";
+
+import { ColumnDef } from "@tanstack/react-table";
+import { CloudIcon, FolderIcon, ServerIcon } from "lucide-react";
+import { notFound } from "next/navigation";
+
+import { DataTable } from "@/components/ui/table/data-table";
+import { DataTableExpandAllToggle } from "@/components/ui/table/data-table-expand-all-toggle";
+import { DataTableExpandableCell } from "@/components/ui/table/data-table-expandable-cell";
+
+/**
+ * Demo page for the Expandable DataTable component.
+ * Only accessible in development mode.
+ *
+ * Showcases:
+ * 1. Hierarchical rows with expand/collapse
+ * 2. Expand all / collapse all toggle
+ * 3. Row selection with child auto-selection
+ */
+
+const IS_DEV = process.env.NODE_ENV === "development";
+
+// Type constants following project conventions
+const PROVIDER_TYPES = {
+  ORGANIZATION: "organization",
+  OU: "ou",
+  ACCOUNT: "account",
+} as const;
+type ProviderType = (typeof PROVIDER_TYPES)[keyof typeof PROVIDER_TYPES];
+
+const PROVIDER_STATUSES = {
+  CONNECTED: "connected",
+  DISCONNECTED: "disconnected",
+  PENDING: "pending",
+} as const;
+type ProviderStatus =
+  (typeof PROVIDER_STATUSES)[keyof typeof PROVIDER_STATUSES];
+
+interface HierarchicalProvider {
+  id: string;
+  name: string;
+  type: ProviderType;
+  status: ProviderStatus;
+  resourceCount: number;
+  children?: HierarchicalProvider[];
+}
+
+const tableData: HierarchicalProvider[] = [
+  {
+    id: "org-1",
+    name: "AWS Organization",
+    type: PROVIDER_TYPES.ORGANIZATION,
+    status: PROVIDER_STATUSES.CONNECTED,
+    resourceCount: 1250,
+    children: [
+      {
+        id: "ou-prod",
+        name: "Production OU",
+        type: PROVIDER_TYPES.OU,
+        status: PROVIDER_STATUSES.CONNECTED,
+        resourceCount: 800,
+        children: [
+          {
+            id: "acc-prod-1",
+            name: "prod-web-services",
+            type: PROVIDER_TYPES.ACCOUNT,
+            status: PROVIDER_STATUSES.CONNECTED,
+            resourceCount: 450,
+          },
+          {
+            id: "acc-prod-2",
+            name: "prod-databases",
+            type: PROVIDER_TYPES.ACCOUNT,
+            status: PROVIDER_STATUSES.CONNECTED,
+            resourceCount: 350,
+          },
+        ],
+      },
+      {
+        id: "ou-dev",
+        name: "Development OU",
+        type: PROVIDER_TYPES.OU,
+        status: PROVIDER_STATUSES.CONNECTED,
+        resourceCount: 450,
+        children: [
+          {
+            id: "acc-dev-1",
+            name: "dev-sandbox",
+            type: PROVIDER_TYPES.ACCOUNT,
+            status: PROVIDER_STATUSES.PENDING,
+            resourceCount: 200,
+          },
+          {
+            id: "acc-dev-2",
+            name: "dev-testing",
+            type: PROVIDER_TYPES.ACCOUNT,
+            status: PROVIDER_STATUSES.DISCONNECTED,
+            resourceCount: 250,
+          },
+        ],
+      },
+    ],
+  },
+];
+
+const STATUS_COLORS = {
+  connected: "text-green-500",
+  disconnected: "text-red-500",
+  pending: "text-yellow-500",
+} as const;
+
+const TYPE_ICONS = {
+  organization: ServerIcon,
+  ou: FolderIcon,
+  account: CloudIcon,
+} as const;
+
+const columns: ColumnDef<HierarchicalProvider>[] = [
+  {
+    accessorKey: "name",
+    size: 300,
+    header: ({ table }) => (
+      <div className="flex items-center gap-2">
+        <DataTableExpandAllToggle table={table} />
+        <span>Name</span>
+      </div>
+    ),
+    cell: ({ row }) => {
+      const Icon = TYPE_ICONS[row.original.type];
+      return (
+        <DataTableExpandableCell row={row}>
+          <div className="flex items-center gap-2">
+            <Icon className="h-4 w-4 shrink-0" />
+            <span>{row.original.name}</span>
+          </div>
+        </DataTableExpandableCell>
+      );
+    },
+  },
+  {
+    accessorKey: "type",
+    size: 120,
+    header: "Type",
+    cell: ({ row }) => <span className="capitalize">{row.original.type}</span>,
+  },
+  {
+    accessorKey: "status",
+    size: 120,
+    header: "Status",
+    cell: ({ row }) => (
+      <span className={STATUS_COLORS[row.original.status]}>
+        {row.original.status}
+      </span>
+    ),
+  },
+  {
+    accessorKey: "resourceCount",
+    size: 100,
+    header: "Resources",
+    cell: ({ row }) => row.original.resourceCount.toLocaleString(),
+  },
+];
+
+export default function DemoExpandableTablePage() {
+  if (!IS_DEV) {
+    notFound();
+  }
+
+  return (
+    <div className="container mx-auto space-y-12 p-8">
+      <h1 className="text-3xl font-bold">Expandable DataTable Demo</h1>
+
+      {/* Expandable DataTable */}
+      <section className="space-y-4">
+        <div>
+          <h2 className="text-xl font-semibold">Expandable DataTable</h2>
+          <p className="text-text-neutral-secondary text-sm">
+            Table with hierarchical rows. Click chevron to expand/collapse, or
+            use the header icon to expand/collapse all.
+          </p>
+        </div>
+
+        <DataTable
+          columns={columns}
+          data={tableData}
+          getSubRows={(row) => row.children}
+          defaultExpanded={true}
+        />
+      </section>
+
+      {/* Expandable DataTable with Row Selection */}
+      <section className="space-y-4">
+        <div>
+          <h2 className="text-xl font-semibold">
+            Expandable DataTable with Row Selection
+          </h2>
+          <p className="text-text-neutral-secondary text-sm">
+            Selecting a parent auto-selects all children
+            (enableSubRowSelection).
+          </p>
+        </div>
+
+        <DataTable
+          columns={columns}
+          data={tableData}
+          getSubRows={(row) => row.children}
+          enableRowSelection
+          enableSubRowSelection
+          defaultExpanded={{ "org-1": true, "ou-prod": true }}
+        />
+      </section>
+    </div>
+  );
+}

ui/components/ui/table/data-table-animated-row.tsx

@@ -0,0 +1,91 @@
+"use client";
+
+import { Cell, flexRender, Row } from "@tanstack/react-table";
+import { motion } from "framer-motion";
+
+import { cn } from "@/lib/utils";
+
+interface DataTableAnimatedRowProps<TData> {
+  row: Row<TData>;
+}
+
+/**
+ * DataTableAnimatedRow renders a table row with smooth expand/collapse animations.
+ *
+ * The trick: You cannot animate <tr> height directly (tables ignore it).
+ * Instead, we wrap each cell's content in a motion.div and animate THAT.
+ *
+ * How it works:
+ * 1. The <tr> itself is not animated
+ * 2. Each <td> contains a motion.div wrapper
+ * 3. The wrapper animates height from 0 to "auto"
+ * 4. overflow-hidden clips content during animation
+ * 5. Padding is on the inner content, not the td
+ */
+export function DataTableAnimatedRow<TData>({
+  row,
+}: DataTableAnimatedRowProps<TData>) {
+  return (
+    <motion.tr
+      initial="collapsed"
+      animate="open"
+      exit="collapsed"
+      variants={{
+        open: { opacity: 1 },
+        collapsed: { opacity: 0 },
+      }}
+      transition={{ duration: 0.2 }}
+      data-state={row.getIsSelected() ? "selected" : undefined}
+      className={cn(
+        "transition-colors",
+        "[&>td:first-child]:rounded-l-full [&>td:last-child]:rounded-r-full",
+        "hover:bg-bg-neutral-tertiary",
+        "data-[state=selected]:bg-bg-neutral-tertiary",
+      )}
+    >
+      {row.getVisibleCells().map((cell: Cell<TData, unknown>, index, cells) => {
+        const isFirst = index === 0;
+        const isLast = index === cells.length - 1;
+
+        return (
+          <td key={cell.id} className="overflow-hidden p-0">
+            <motion.div
+              initial="collapsed"
+              animate="open"
+              exit="collapsed"
+              variants={{
+                open: {
+                  height: "auto",
+                  opacity: 1,
+                  transition: {
+                    height: { duration: 0.2, ease: "easeOut" },
+                    opacity: { duration: 0.15, delay: 0.05 },
+                  },
+                },
+                collapsed: {
+                  height: 0,
+                  opacity: 0,
+                  transition: {
+                    height: { duration: 0.2, ease: "easeIn" },
+                    opacity: { duration: 0.1 },
+                  },
+                },
+              }}
+              className="overflow-hidden"
+            >
+              <div
+                className={cn(
+                  "px-1.5 py-2",
+                  isFirst && "pl-3",
+                  isLast && "pr-3",
+                )}
+              >
+                {flexRender(cell.column.columnDef.cell, cell.getContext())}
+              </div>
+            </motion.div>
+          </td>
+        );
+      })}
+    </motion.tr>
+  );
+}

ui/components/ui/table/data-table-expand-all-toggle.tsx

@@ -0,0 +1,69 @@
+"use client";
+
+import { Table } from "@tanstack/react-table";
+import { Maximize2Icon, Minimize2Icon } from "lucide-react";
+
+import { cn } from "@/lib/utils";
+
+interface DataTableExpandAllToggleProps<TData> {
+  table: Table<TData>;
+}
+
+/**
+ * DataTableExpandAllToggle provides a button in the table header to expand
+ * or collapse all rows at once.
+ *
+ * Features:
+ * - Shows Maximize2 icon when rows are collapsed (click to expand all)
+ * - Shows Minimize2 icon when rows are expanded (click to collapse all)
+ * - Accessible with proper aria-label
+ * - Only renders when the table has expandable rows
+ *
+ * @example
+ * ```tsx
+ * // In column definition header:
+ * {
+ *   id: "name",
+ *   header: ({ table }) => (
+ *     <div className="flex items-center gap-2">
+ *       <DataTableExpandAllToggle table={table} />
+ *       <span>Name</span>
+ *     </div>
+ *   ),
+ *   cell: ({ row }) => (
+ *     <DataTableExpandableCell row={row}>
+ *       <span>{row.original.name}</span>
+ *     </DataTableExpandableCell>
+ *   ),
+ * }
+ * ```
+ */
+export function DataTableExpandAllToggle<TData>({
+  table,
+}: DataTableExpandAllToggleProps<TData>) {
+  const isAllExpanded = table.getIsAllRowsExpanded();
+  const canExpand = table.getCanSomeRowsExpand();
+
+  if (!canExpand) {
+    return null;
+  }
+
+  return (
+    <button
+      onClick={() => table.toggleAllRowsExpanded(!isAllExpanded)}
+      className={cn(
+        "rounded p-1 transition-colors",
+        "hover:bg-prowler-white/10",
+        "focus-visible:ring-border-input-primary-press focus-visible:ring-2 focus-visible:outline-none",
+      )}
+      aria-label={isAllExpanded ? "Collapse all rows" : "Expand all rows"}
+      aria-expanded={isAllExpanded}
+    >
+      {isAllExpanded ? (
+        <Minimize2Icon className="h-4 w-4" />
+      ) : (
+        <Maximize2Icon className="h-4 w-4" />
+      )}
+    </button>
+  );
+}

ui/components/ui/table/data-table-expand-toggle.tsx

@@ -0,0 +1,58 @@
+"use client";
+
+import { Row } from "@tanstack/react-table";
+import { ChevronRightIcon } from "lucide-react";
+
+import { cn } from "@/lib/utils";
+
+interface DataTableExpandToggleProps<TData> {
+  row: Row<TData>;
+}
+
+/**
+ * DataTableExpandToggle provides a clickable chevron button for expanding/collapsing
+ * table rows that have sub-rows.
+ *
+ * Features:
+ * - Only shows for rows that can expand (have sub-rows)
+ * - Provides consistent spacing for rows without sub-rows
+ * - Animates chevron rotation on expand/collapse
+ * - Accessible with proper aria-label
+ *
+ * @example
+ * ```tsx
+ * // In column definition:
+ * {
+ *   id: "expand",
+ *   cell: ({ row }) => <DataTableExpandToggle row={row} />,
+ * }
+ * ```
+ */
+export function DataTableExpandToggle<TData>({
+  row,
+}: DataTableExpandToggleProps<TData>) {
+  if (!row.getCanExpand()) {
+    // Return a spacer div for alignment when row has no sub-rows
+    return <div className="w-4" />;
+  }
+
+  return (
+    <button
+      onClick={row.getToggleExpandedHandler()}
+      className={cn(
+        "rounded p-1 transition-colors",
+        "hover:bg-prowler-white/10",
+        "focus-visible:ring-border-input-primary-press focus-visible:ring-2 focus-visible:outline-none",
+      )}
+      aria-label={row.getIsExpanded() ? "Collapse row" : "Expand row"}
+      aria-expanded={row.getIsExpanded()}
+    >
+      <ChevronRightIcon
+        className={cn(
+          "h-4 w-4 shrink-0 transition-transform duration-200",
+          row.getIsExpanded() && "rotate-90",
+        )}
+      />
+    </button>
+  );
+}

ui/components/ui/table/data-table-expandable-cell.tsx

@@ -0,0 +1,69 @@
+"use client";
+
+import { Row } from "@tanstack/react-table";
+import { CornerDownRightIcon } from "lucide-react";
+
+import { DataTableExpandToggle } from "./data-table-expand-toggle";
+
+/** Indentation per nesting level in rem units */
+const INDENT_PER_LEVEL_REM = 1.5;
+
+interface DataTableExpandableCellProps<TData> {
+  row: Row<TData>;
+  children: React.ReactNode;
+  /** Whether to show the expand/collapse toggle (default: true) */
+  showToggle?: boolean;
+}
+
+/**
+ * DataTableExpandableCell is a wrapper component for table cells that need
+ * to display content with proper indentation for nested rows.
+ *
+ * Features:
+ * - Automatically indents content based on row depth
+ * - Shows CornerDownRight icon for child rows (depth > 0)
+ * - Optionally includes the expand/collapse toggle for parent rows
+ * - Maintains proper alignment for all nesting levels
+ *
+ * @example
+ * ```tsx
+ * // In column definition:
+ * {
+ *   accessorKey: "name",
+ *   header: "Name",
+ *   cell: ({ row }) => (
+ *     <DataTableExpandableCell row={row}>
+ *       <span>{row.original.name}</span>
+ *     </DataTableExpandableCell>
+ *   ),
+ * }
+ * ```
+ */
+export function DataTableExpandableCell<TData>({
+  row,
+  children,
+  showToggle = true,
+}: DataTableExpandableCellProps<TData>) {
+  const isChildRow = row.depth > 0;
+  const canExpand = row.getCanExpand();
+
+  return (
+    <div
+      className="flex items-center gap-2"
+      style={{ paddingLeft: `${row.depth * INDENT_PER_LEVEL_REM}rem` }}
+    >
+      {showToggle && (
+        <>
+          {canExpand ? (
+            <DataTableExpandToggle row={row} />
+          ) : isChildRow ? (
+            <CornerDownRightIcon className="h-4 w-4 shrink-0" />
+          ) : (
+            <div className="w-4" />
+          )}
+        </>
+      )}
+      {children}
+    </div>
+  );
+}

ui/components/ui/table/data-table.tsx

@@ -3,8 +3,10 @@
 import {
   ColumnDef,
   ColumnFiltersState,
+  ExpandedState,
   flexRender,
   getCoreRowModel,
+  getExpandedRowModel,
   getFilteredRowModel,
   getSortedRowModel,
   OnChangeFn,
@@ -13,7 +15,8 @@ import {
   SortingState,
   useReactTable,
 } from "@tanstack/react-table";
-import { useState } from "react";
+import { AnimatePresence } from "framer-motion";
+import { useEffect, useRef, useState } from "react";
 
 import {
   Table,
@@ -23,12 +26,20 @@ import {
   TableHeader,
   TableRow,
 } from "@/components/ui/table";
+import { DataTableAnimatedRow } from "@/components/ui/table/data-table-animated-row";
 import { DataTablePagination } from "@/components/ui/table/data-table-pagination";
 import { DataTableSearch } from "@/components/ui/table/data-table-search";
 import { useFilterTransitionOptional } from "@/contexts";
 import { cn } from "@/lib";
 import { FilterOption, MetaDataProps } from "@/types";
 
+/**
+ * Default column size used by TanStack Table when no explicit size is set.
+ * We skip applying inline width styles for columns with this default value
+ * to allow them to flex naturally within the table layout.
+ */
+const DEFAULT_COLUMN_SIZE = 150;
+
 interface DataTableProviderProps<TData, TValue> {
   columns: ColumnDef<TData, TValue>[];
   data: TData[];
@@ -42,6 +53,16 @@ interface DataTableProviderProps<TData, TValue> {
   getRowCanSelect?: (row: Row<TData>) => boolean;
   /** Show search bar in the table toolbar */
   showSearch?: boolean;
+  /** Function to extract sub-rows from a row for hierarchical data */
+  getSubRows?: (row: TData) => TData[] | undefined;
+  /** Controlled expanded state */
+  expanded?: ExpandedState;
+  /** Callback when expanded state changes */
+  onExpandedChange?: OnChangeFn<ExpandedState>;
+  /** Auto-select children when parent selected (default: true) */
+  enableSubRowSelection?: boolean;
+  /** Expand all rows by default, or provide specific expanded state */
+  defaultExpanded?: boolean | ExpandedState;
   /** Prefix for URL params to avoid conflicts (e.g., "findings" -> "findingsPage") */
   paramPrefix?: string;
 
@@ -87,6 +108,11 @@ export function DataTable<TData, TValue>({
   onRowSelectionChange,
   getRowCanSelect,
   showSearch = false,
+  getSubRows,
+  expanded: controlledExpanded,
+  onExpandedChange,
+  enableSubRowSelection = true,
+  defaultExpanded,
   paramPrefix = "",
   controlledSearch,
   onSearchChange,
@@ -98,6 +124,12 @@ export function DataTable<TData, TValue>({
 }: DataTableProviderProps<TData, TValue>) {
   const [sorting, setSorting] = useState<SortingState>([]);
   const [columnFilters, setColumnFilters] = useState<ColumnFiltersState>([]);
+  // ExpandedState should be a Record<string, boolean> for individual row control
+  // Note: We don't use `true` (boolean) as it makes rows permanently expanded
+  const [expanded, setExpanded] = useState<ExpandedState>(() => {
+    if (typeof defaultExpanded === "object") return defaultExpanded;
+    return {};
+  });
 
   // Get transition state from context for loading indicator
   const filterTransition = useFilterTransitionOptional();
@@ -116,18 +148,48 @@ export function DataTable<TData, TValue>({
     getFilteredRowModel: getFilteredRowModel(),
     onRowSelectionChange,
     manualPagination: true,
+    // Expansion support for hierarchical data
+    getSubRows,
+    getExpandedRowModel: getSubRows ? getExpandedRowModel() : undefined,
+    enableSubRowSelection,
+    onExpandedChange: onExpandedChange ?? setExpanded,
     state: {
       sorting,
       columnFilters,
       rowSelection: rowSelection ?? {},
+      expanded: controlledExpanded ?? expanded,
     },
   });
 
+  // Track whether initial expansion has been applied
+  const hasInitiallyExpanded = useRef(false);
+
+  // Expand all rows on mount when defaultExpanded={true}
+  useEffect(() => {
+    if (
+      !hasInitiallyExpanded.current &&
+      defaultExpanded === true &&
+      getSubRows
+    ) {
+      table.toggleAllRowsExpanded(true);
+      hasInitiallyExpanded.current = true;
+    }
+  }, [defaultExpanded, getSubRows, table]);
+
   // Calculate selection key to force header re-render when selection changes
   const selectionKey = rowSelection
     ? Object.keys(rowSelection).filter((k) => rowSelection[k]).length
     : 0;
 
+  // Calculate expansion key to force header re-render when expansion changes
+  const currentExpanded = controlledExpanded ?? expanded;
+  const expansionKey =
+    currentExpanded === true
+      ? "all"
+      : typeof currentExpanded === "object"
+        ? Object.keys(currentExpanded).filter((k) => currentExpanded[k]).length
+        : 0;
+
   // Format total entries count
   const totalEntries = metadata?.pagination?.count ?? 0;
   const formattedTotal = totalEntries.toLocaleString();
@@ -161,13 +223,21 @@ export function DataTable<TData, TValue>({
           )}
         </div>
       )}
-      <Table>
+      <Table className={getSubRows ? "table-fixed" : undefined}>
         <TableHeader>
           {table.getHeaderGroups().map((headerGroup) => (
-            <TableRow key={`${headerGroup.id}-${selectionKey}`}>
+            <TableRow key={`${headerGroup.id}-${selectionKey}-${expansionKey}`}>
               {headerGroup.headers.map((header) => {
+                const size = header.getSize();
                 return (
-                  <TableHead key={header.id}>
+                  <TableHead
+                    key={header.id}
+                    style={
+                      getSubRows && size !== DEFAULT_COLUMN_SIZE
+                        ? { width: `${size}px` }
+                        : undefined
+                    }
+                  >
                     {header.isPlaceholder
                       ? null
                       : flexRender(
@@ -181,26 +251,38 @@ export function DataTable<TData, TValue>({
           ))}
         </TableHeader>
         <TableBody>
-          {rows?.length ? (
-            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>
-                ))}
+          <AnimatePresence initial={false}>
+            {rows?.length ? (
+              rows.map((row) =>
+                getSubRows && row.depth > 0 ? (
+                  <DataTableAnimatedRow key={row.id} row={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>
-            ))
-          ) : (
-            <TableRow>
-              <TableCell colSpan={columns.length} className="h-24 text-center">
-                No results.
-              </TableCell>
-            </TableRow>
-          )}
+            )}
+          </AnimatePresence>
         </TableBody>
       </Table>
       {metadata && (