From d7fa45bd3cb43a72a5fbe3cb9a4435f3295099ae Mon Sep 17 00:00:00 2001
From: Keven Leone <kevenleone.me@gmail.com>
Date: Sat, 2 May 2026 23:12:46 -0300
Subject: [PATCH 1/2] fix: render checkboxes in row selection column fixes:
 #1504

The __ps_select column rendered empty cells with no visual
state indicator. Render CheckboxTable inside header and cell
with pointer-events-none so selection stays handled by
DataGrid mousedown handlers.
---
 ui/studio/grid/test-utils.tsx             | 23 +++++++++++++++++----
 ui/studio/views/sql/SqlView.tsx           | 25 +++++++++++++++++++----
 ui/studio/views/table/ActiveTableView.tsx | 25 +++++++++++++++++++----
 3 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/ui/studio/grid/test-utils.tsx b/ui/studio/grid/test-utils.tsx
index a071bd51..cbad90d5 100644
--- a/ui/studio/grid/test-utils.tsx
+++ b/ui/studio/grid/test-utils.tsx
@@ -2,6 +2,7 @@ import type { ColumnDef } from "@tanstack/react-table";
 import { act } from "react";
 import { type Mock, vi } from "vitest";
 
+import { CheckboxTable } from "../../components/ui/checkbox-table";
 import { TableHead } from "../../components/ui/table";
 import { Cell, type CellProps } from "../cell/Cell";
 
@@ -58,14 +59,28 @@ 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">
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={table.getIsAllRowsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </TableHead>
       );
     },
-    cell() {
+    cell({ row }) {
       return (props: Omit<CellProps, "children" | "ref">) => (
-        <Cell data-select="true" {...props} />
+        <Cell data-select="true" {...props}>
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={row.getIsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </Cell>
       );
     },
   };
diff --git a/ui/studio/views/sql/SqlView.tsx b/ui/studio/views/sql/SqlView.tsx
index 904ec5ac..c1d9fb61 100644
--- a/ui/studio/views/sql/SqlView.tsx
+++ b/ui/studio/views/sql/SqlView.tsx
@@ -24,6 +24,7 @@ import { createSqlEditorSchemaFromIntrospection } from "../../../../data/sql-edi
 import { getTopLevelSqlStatementAtCursor } from "../../../../data/sql-statements";
 import { Button } from "../../../components/ui/button";
 import { Input } from "../../../components/ui/input";
+import { CheckboxTable } from "../../../components/ui/checkbox-table";
 import { TableHead, TableRow } from "../../../components/ui/table";
 import { useColumnPinning } from "../../../hooks/use-column-pinning";
 import { useIntrospection } from "../../../hooks/use-introspection";
@@ -103,15 +104,31 @@ 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">
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={table.getIsAllRowsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </TableHead>
+      );
     };
   },
   cell({ row }) {
-    void row;
     return (props: Omit<CellProps, "children" | "ref">) => {
-      return <Cell data-select="true" {...props} />;
+      return (
+        <Cell data-select="true" {...props}>
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={row.getIsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </Cell>
+      );
     };
   },
 } satisfies ColumnDef<Record<string, unknown>>;
diff --git a/ui/studio/views/table/ActiveTableView.tsx b/ui/studio/views/table/ActiveTableView.tsx
index cdf7f11b..8388e0af 100644
--- a/ui/studio/views/table/ActiveTableView.tsx
+++ b/ui/studio/views/table/ActiveTableView.tsx
@@ -39,6 +39,7 @@ import {
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from "../../../components/ui/dropdown-menu";
+import { CheckboxTable } from "../../../components/ui/checkbox-table";
 import { TableHead } from "../../../components/ui/table";
 import { useActiveTableInsert } from "../../../hooks/use-active-table-insert";
 import { useActiveTableQuery } from "../../../hooks/use-active-table-query";
@@ -1516,15 +1517,31 @@ 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">
+              <div className="flex items-center justify-center h-full w-full">
+                <CheckboxTable
+                  checked={table.getIsAllRowsSelected()}
+                  className="pointer-events-none h-4 w-4"
+                />
+              </div>
+            </TableHead>
+          );
         };
       },
       cell({ row }) {
-        void row;
         return (props: Omit<CellProps, "children" | "ref">) => {
-          return <Cell data-select="true" {...props} />;
+          return (
+            <Cell data-select="true" {...props}>
+              <div className="flex items-center justify-center h-full w-full">
+                <CheckboxTable
+                  checked={row.getIsSelected()}
+                  className="pointer-events-none h-4 w-4"
+                />
+              </div>
+            </Cell>
+          );
         };
       },
     },

From 06a3fd420f4a554edbacec53f7c8ee29ed95a1d6 Mon Sep 17 00:00:00 2001
From: Keven Leone <kevenleone.me@gmail.com>
Date: Sat, 2 May 2026 23:29:02 -0300
Subject: [PATCH 2/2] fix: improve row selection handling and clean up code,
 fixes: #1504

---
 ui/studio/grid/DataGrid.tsx | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/ui/studio/grid/DataGrid.tsx b/ui/studio/grid/DataGrid.tsx
index 6a611298..38e86cdf 100644
--- a/ui/studio/grid/DataGrid.tsx
+++ b/ui/studio/grid/DataGrid.tsx
@@ -927,7 +927,7 @@ export function DataGrid(props: DataGridProps) {
       });
       const shouldEnablePreview = Boolean(
         dragDropTarget.compatibleOverId &&
-        dragDropTarget.compatibleOverId !== activeId,
+          dragDropTarget.compatibleOverId !== activeId,
       );
       setIsColumnReorderPreviewEnabled((current) =>
         current === shouldEnablePreview ? current : shouldEnablePreview,
@@ -1470,8 +1470,8 @@ export function DataGrid(props: DataGridProps) {
       const text = hasRowSelection
         ? getSelectedRowClipboardText()
         : hasCellSelection
-          ? getSelectedClipboardText()
-          : getFocusedCellClipboardText();
+        ? getSelectedClipboardText()
+        : getFocusedCellClipboardText();
 
       if (!text) {
         return;
@@ -2064,12 +2064,26 @@ export function DataGrid(props: DataGridProps) {
         return;
       }
 
-      selectSingleRowMode({ rowId, rowIndex, drag: isPrimaryButton });
+      clearNativeTextSelection();
+      clearCellSelectionState();
+
+      const nextSelection = { ...rowSelectionState };
+
+      if (nextSelection[rowId]) {
+        delete nextSelection[rowId];
+      } else {
+        nextSelection[rowId] = true;
+      }
+
+      setRowSelection(nextSelection);
+
+      rowSelectionDragRef.current = isPrimaryButton;
+      rowSelectionAnchorRef.current = rowIndex;
     },
     [
       clearCellSelectionState,
+      clearNativeTextSelection,
       rowSelectionState,
-      selectSingleRowMode,
       setRowSelection,
     ],
   );