fix: improve row selection handling and clean up code, fixes: #1504

.changeset/checkbox-row-selection.md

@@ -0,0 +1,5 @@
+---
+"@prisma/studio-core": patch
+---
+
+Fix row-selection checkboxes rendering as empty cells with no visual state indicator, and fix drag-selecting a range of rows silently dropping previously selected rows outside that range.

ui/studio/cell/SelectHeaderCell.tsx

@@ -4,18 +4,23 @@ import { CheckboxTable } from "@/ui/components/ui/checkbox-table";
 
 export interface SelectHeaderCellProps {
   table: Table<Record<string, unknown>>;
-  readonly: boolean;
 }
 
 export function SelectHeaderCell(props: SelectHeaderCellProps) {
-  const { table, readonly } = props;
+  const { table } = props;
 
   return (
-    <CheckboxTable
-      checked={table.getIsAllRowsSelected()}
-      onCheckedChange={(value) => table.toggleAllRowsSelected(!!value)}
-      disabled={readonly}
-      className="w-full h-full rounded-none border-none"
-    />
+    <div className="flex items-center justify-center h-full w-full">
+      <CheckboxTable
+        checked={
+          table.getIsAllRowsSelected()
+            ? true
+            : table.getIsSomeRowsSelected()
+              ? "indeterminate"
+              : false
+        }
+        className="pointer-events-none h-4 w-4"
+      />
+    </div>
   );
 }

ui/studio/cell/SelectRowCell.tsx

@@ -4,18 +4,17 @@ import { CheckboxTable } from "@/ui/components/ui/checkbox-table";
 
 export interface SelectRowCellProps {
   row: Row<Record<string, unknown>>;
-  readonly: boolean;
 }
 
 export function SelectRowCell(props: SelectRowCellProps) {
-  const { row, readonly } = props;
+  const { row } = props;
 
   return (
-    <CheckboxTable
-      checked={row.getIsSelected()}
-      disabled={!row.getCanSelect() || readonly}
-      onCheckedChange={row.getToggleSelectedHandler()}
-      className="w-full h-full rounded-none border-none"
-    />
+    <div className="flex items-center justify-center h-full w-full">
+      <CheckboxTable
+        checked={row.getIsSelected()}
+        className="pointer-events-none h-4 w-4"
+      />
+    </div>
   );
 }

ui/studio/grid/DataGrid.interactions.test.tsx

@@ -1095,6 +1095,105 @@ describe("DataGrid interactions", () => {
     cleanup();
   });
 
+  it("renders a checkbox in the row selection column that reflects selection state", () => {
+    createSelection({ isCollapsed: true });
+
+    const { cleanup, getCell } = renderGrid({
+      columnDefs: createReadOnlyColumns({ includeRowSelector: true }),
+    });
+
+    const checkbox = getCell(0, "__ps_select").querySelector(
+      '[role="checkbox"]',
+    );
+
+    expect(checkbox).not.toBeNull();
+    expect(checkbox?.getAttribute("data-state")).toBe("unchecked");
+
+    dispatchMouse(getCell(0, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    expect(
+      getCell(0, "__ps_select")
+        .querySelector('[role="checkbox"]')
+        ?.getAttribute("data-state"),
+    ).toBe("checked");
+
+    cleanup();
+  });
+
+  it("shows an indeterminate header checkbox when only some rows are selected", () => {
+    createSelection({ isCollapsed: true });
+
+    const { cleanup, container, getCell } = renderGrid({
+      columnDefs: createReadOnlyColumns({ includeRowSelector: true }),
+    });
+
+    dispatchMouse(getCell(0, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    const headerCheckbox = container.querySelector(
+      'th[aria-label="Row selection spacer"] [role="checkbox"]',
+    );
+
+    expect(headerCheckbox?.getAttribute("data-state")).toBe("indeterminate");
+
+    cleanup();
+  });
+
+  it("preserves existing row selection when drag-selecting a different range", () => {
+    createSelection({ isCollapsed: true });
+
+    const { cleanup, getCell, getSelectedRowCount } = renderGrid({
+      columnDefs: createReadOnlyColumns({ includeRowSelector: true }),
+    });
+
+    dispatchMouse(getCell(0, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    expect(getSelectedRowCount()).toBe(1);
+
+    dispatchMouse(getCell(2, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(getCell(1, "__ps_select"), "mouseover", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    expect(getSelectedRowCount()).toBe(3);
+    expect(getCell(0, "__ps_select").closest("tr")?.dataset.rowSelected).toBe(
+      "true",
+    );
+    expect(getCell(1, "__ps_select").closest("tr")?.dataset.rowSelected).toBe(
+      "true",
+    );
+    expect(getCell(2, "__ps_select").closest("tr")?.dataset.rowSelected).toBe(
+      "true",
+    );
+
+    cleanup();
+  });
+
+  it("does not re-select a row from a stray mouse movement after a deselecting click", () => {
+    createSelection({ isCollapsed: true });
+
+    const { cleanup, getCell, getSelectedRowCount } = renderGrid({
+      columnDefs: createReadOnlyColumns({ includeRowSelector: true }),
+    });
+
+    dispatchMouse(getCell(0, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    expect(getSelectedRowCount()).toBe(1);
+
+    dispatchMouse(getCell(0, "__ps_select"), "mousedown", { button: 0 });
+    dispatchMouse(getCell(1, "__ps_select"), "mouseover", { button: 0 });
+    dispatchMouse(window, "mouseup");
+
+    expect(getSelectedRowCount()).toBe(0);
+    expect(
+      getCell(1, "__ps_select").closest("tr")?.dataset.rowSelected,
+    ).toBeUndefined();
+
+    cleanup();
+  });
+
   it("selects all rows when clicking the top-left spacer header cell", () => {
     createSelection({ isCollapsed: true });
 

ui/studio/grid/DataGrid.tsx

@@ -687,6 +687,7 @@ export function DataGrid(props: DataGridProps) {
   } | null>(null);
   const rowSelectionAnchorRef = useRef<number | null>(null);
   const rowSelectionDragRef = useRef(false);
+  const rowSelectionDragBaseRef = useRef<RowSelectionState>({});
   const previousSelectionScopeKeyRef = useRef<string | undefined>(
     selectionScopeKey,
   );
@@ -1470,6 +1471,7 @@ export function DataGrid(props: DataGridProps) {
     const handleMouseUp = (event: MouseEvent) => {
       rowSelectionDragRef.current = false;
       rowSelectionAnchorRef.current = null;
+      rowSelectionDragBaseRef.current = {};
 
       const pointerSelection = pointerSelectionRef.current;
 
@@ -1660,6 +1662,7 @@ export function DataGrid(props: DataGridProps) {
   const resetRowSelectionInteractionState = useCallback(() => {
     rowSelectionAnchorRef.current = null;
     rowSelectionDragRef.current = false;
+    rowSelectionDragBaseRef.current = {};
   }, []);
 
   const resetSelectionInteractionState = useCallback(() => {
@@ -2029,7 +2032,9 @@ export function DataGrid(props: DataGridProps) {
     (fromRowIndex: number, toRowIndex: number) => {
       const start = Math.min(fromRowIndex, toRowIndex);
       const end = Math.max(fromRowIndex, toRowIndex);
-      const nextSelection: RowSelectionState = {};
+      const nextSelection: RowSelectionState = {
+        ...rowSelectionDragBaseRef.current,
+      };
       const rowModel = table.getRowModel().rows;
 
       for (let index = start; index <= end; index++) {
@@ -2118,29 +2123,33 @@ export function DataGrid(props: DataGridProps) {
         return;
       }
 
-      if (isPrimaryButton) {
-        event.preventDefault();
-      }
+      event.preventDefault();
       event.stopPropagation();
+      clearNativeTextSelection();
+      clearCellSelectionState();
 
-      if (event.shiftKey) {
-        clearNativeTextSelection();
-        clearCellSelectionState();
-        const nextSelection = { ...rowSelectionState };
+      const previousSelection = rowSelectionState;
+      const isSelecting = previousSelection[rowId] !== true;
+      const nextSelection = { ...previousSelection };
 
-        if (nextSelection[rowId] === true) {
-          delete nextSelection[rowId];
-        } else {
-          nextSelection[rowId] = true;
-        }
+      if (isSelecting) {
+        nextSelection[rowId] = true;
+      } else {
+        delete nextSelection[rowId];
+      }
+
+      setRowSelection(nextSelection);
+      rowSelectionAnchorRef.current = rowIndex;
 
-        setRowSelection(nextSelection);
+      // Shift+click toggles a single row without starting a drag-range
+      // selection; deselecting a row also skips arming the drag so a tiny
+      // mouse movement afterward doesn't re-select it via mouseenter.
+      if (isSelecting && !event.shiftKey) {
+        rowSelectionDragBaseRef.current = previousSelection;
+        rowSelectionDragRef.current = true;
+      } else {
         rowSelectionDragRef.current = false;
-        rowSelectionAnchorRef.current = rowIndex;
-        return;
       }
-
-      selectSingleRowMode({ rowId, rowIndex, drag: isPrimaryButton });
     },
     [
       clearCellSelectionState,

ui/studio/grid/test-utils.tsx

@@ -4,6 +4,8 @@ import { type Mock, vi } from "vitest";
 
 import { TableHead } from "../../components/ui/table";
 import { Cell, type CellProps } from "../cell/Cell";
+import { SelectHeaderCell } from "../cell/SelectHeaderCell";
+import { SelectRowCell } from "../cell/SelectRowCell";
 
 export type GridRow = Record<string, unknown>;
 export type GridColumnDef = ColumnDef<GridRow>;
@@ -58,14 +60,18 @@ export function createReadOnlyColumns(args?: {
     enableSorting: false,
     size: 35,
     minSize: 35,
-    header() {
+    header({ table }) {
       return (props: Omit<CellProps, "children" | "ref">) => (
-        <TableHead {...props} aria-label="Row selection spacer" />
+        <TableHead {...props} aria-label="Row selection spacer">
+          <SelectHeaderCell table={table} />
+        </TableHead>
       );
     },
-    cell() {
+    cell({ row }) {
       return (props: Omit<CellProps, "children" | "ref">) => (
-        <Cell data-select="true" {...props} />
+        <Cell data-select="true" {...props}>
+          <SelectRowCell row={row} />
+        </Cell>
       );
     },
   };

ui/studio/views/sql/SqlView.tsx

@@ -32,6 +32,8 @@ import { useNavigation } from "../../../hooks/use-navigation";
 import type { CellProps } from "../../cell/Cell";
 import { Cell } from "../../cell/Cell";
 import { getCell } from "../../cell/get-cell";
+import { SelectHeaderCell } from "../../cell/SelectHeaderCell";
+import { SelectRowCell } from "../../cell/SelectRowCell";
 import { useStudio } from "../../context";
 import { DataGrid, type DataGridProps } from "../../grid/DataGrid";
 import { DataGridDraggableHeaderCell } from "../../grid/DataGridDraggableHeaderCell";
@@ -103,15 +105,21 @@ const SQL_ROW_SELECTION_COLUMN_DEF = {
   size: 35,
   minSize: 35,
   header({ table }) {
-    void table;
     return (props: Omit<CellProps, "children" | "ref">) => {
-      return <TableHead {...props} aria-label="Row selection spacer" />;
+      return (
+        <TableHead {...props} aria-label="Row selection spacer">
+          <SelectHeaderCell table={table} />
+        </TableHead>
+      );
     };
   },
   cell({ row }) {
-    void row;
     return (props: Omit<CellProps, "children" | "ref">) => {
-      return <Cell data-select="true" {...props} />;
+      return (
+        <Cell data-select="true" {...props}>
+          <SelectRowCell row={row} />
+        </Cell>
+      );
     };
   },
 } satisfies ColumnDef<Record<string, unknown>>;

ui/studio/views/table/ActiveTableView.tsx

@@ -65,6 +65,8 @@ import {
 } from "../../cell/Cell";
 import { getCell } from "../../cell/get-cell";
 import { Link, RelationLink } from "../../cell/Link";
+import { SelectHeaderCell } from "../../cell/SelectHeaderCell";
+import { SelectRowCell } from "../../cell/SelectRowCell";
 import { WriteableCell } from "../../cell/WriteableCell";
 import { useRegisterCommandPaletteActions } from "../../CommandPalette";
 import { type TableUiState, useStudio } from "../../context";
@@ -1546,15 +1548,21 @@ export function ActiveTableView(_props: ViewProps) {
       size: 35,
       minSize: 35,
       header({ table }) {
-        void table;
         return (props: Omit<CellProps, "children" | "ref">) => {
-          return <TableHead {...props} aria-label="Row selection spacer" />;
+          return (
+            <TableHead {...props} aria-label="Row selection spacer">
+              <SelectHeaderCell table={table} />
+            </TableHead>
+          );
         };
       },
       cell({ row }) {
-        void row;
         return (props: Omit<CellProps, "children" | "ref">) => {
-          return <Cell data-select="true" {...props} />;
+          return (
+            <Cell data-select="true" {...props}>
+              <SelectRowCell row={row} />
+            </Cell>
+          );
         };
       },
     },