Update dataset column name and role and display row columns in different (#1038)

color based on their role. Also show label badge on dataset grid column

Author: andresgutgon

apps/web/src/actions/datasetsV2/updateColumn.ts

@@ -0,0 +1,40 @@
+'use server'
+
+import { updateDatasetColumn } from '@latitude-data/core/services/datasetsV2/updateColumn'
+import { z } from 'zod'
+
+import { authProcedure } from '../procedures'
+import {
+  DATASET_COLUMN_ROLES,
+  DatasetColumnRole,
+} from '@latitude-data/core/browser'
+import { DatasetsV2Repository } from '@latitude-data/core/repositories'
+
+const datasetColumnRoleSchema = z.enum(
+  Object.values(DATASET_COLUMN_ROLES) as [
+    DatasetColumnRole,
+    ...DatasetColumnRole[],
+  ],
+)
+export const updateDatasetColumnAction = authProcedure
+  .createServerAction()
+  .input(
+    z.object({
+      datasetId: z.string(),
+      identifier: z.string(),
+      name: z.string().min(1, { message: 'Name is required' }),
+      role: datasetColumnRoleSchema,
+    }),
+  )
+  .handler(async ({ input, ctx }) => {
+    const repo = new DatasetsV2Repository(ctx.workspace.id)
+    const dataset = await repo.find(input.datasetId).then((r) => r.unwrap())
+    return updateDatasetColumn({
+      dataset,
+      data: {
+        identifier: input.identifier,
+        name: input.name,
+        role: input.role,
+      },
+    }).then((r) => r.unwrap())
+  })

apps/web/src/app/(private)/datasets/[datasetId]/DatasetDetailTable/DataGrid/UpdateColumnModal/index.tsx

@@ -0,0 +1,79 @@
+import { useFormAction } from '$/hooks/useFormAction'
+import useDatasets from '$/stores/datasetsV2'
+import { DATASET_COLUMN_ROLES, DatasetV2 } from '@latitude-data/core/browser'
+import {
+  Button,
+  FormWrapper,
+  Input,
+  Modal,
+  Select,
+  SelectOption,
+} from '@latitude-data/web-ui'
+import { useMemo } from 'react'
+
+export function UpdateColumnModal({
+  dataset,
+  columnKey,
+  onClose,
+}: {
+  dataset: DatasetV2
+  columnKey: string
+  onClose: () => void
+}) {
+  const column = dataset.columns.find((col) => col.identifier === columnKey)
+  const { updateColumn, isUpdatingColumn } = useDatasets()
+  const options = useMemo<SelectOption<string>[]>(
+    () =>
+      Object.keys(DATASET_COLUMN_ROLES).map((key) => ({
+        label: key,
+        value: key,
+      })),
+    [],
+  )
+  const { action: updateColumnAction } = useFormAction(updateColumn, {
+    onSuccess: () => {
+      onClose()
+    },
+  })
+
+  if (!column) return null
+
+  return (
+    <Modal
+      dismissible
+      open
+      onOpenChange={onClose}
+      title='Edit Column'
+      description={`Edit column ${column.name} in this dataset`}
+      footer={
+        <>
+          <Button variant='outline' fancy onClick={onClose}>
+            Cancel
+          </Button>
+          <Button
+            form='edit-column'
+            type='submit'
+            fancy
+            disabled={isUpdatingColumn}
+          >
+            {isUpdatingColumn ? 'Updating...' : 'Update'}
+          </Button>
+        </>
+      }
+    >
+      <form id='edit-column' action={updateColumnAction}>
+        <FormWrapper>
+          <input type='hidden' name='datasetId' value={dataset.id} />
+          <input type='hidden' name='identifier' value={column.identifier} />
+          <Input required label='Name' name='name' defaultValue={column.name} />
+          <Select<string>
+            options={options}
+            name='role'
+            label='Column role'
+            defaultValue={column.role ?? DATASET_COLUMN_ROLES.parameter}
+          />
+        </FormWrapper>
+      </form>
+    </Modal>
+  )
+}

apps/web/src/app/(private)/datasets/[datasetId]/DatasetDetailTable/DataGrid/index.tsx

@@ -6,16 +6,27 @@ import type {
   RowsChangeData,
   CellClickArgs,
   CellMouseEvent,
+  RenderHeaderCellProps,
 } from '@latitude-data/web-ui/data-grid'
-import { DataGridCellEditor, type EditorCellProps } from '@latitude-data/web-ui'
-import { Text, FloatingPanel, Button } from '@latitude-data/web-ui'
+import {
+  DataGridCellEditor,
+  type EditorCellProps,
+  ReactStateDispatch,
+  Text,
+  FloatingPanel,
+  Button,
+  cn,
+} from '@latitude-data/web-ui'
+import { DatasetHeadText } from '$/app/(private)/datasets/_components/DatasetHeadText'
 import { DatasetRoleStyle } from '$/hooks/useDatasetRoles'
 import { DatasetV2 } from '@latitude-data/core/browser'
 import { ClientPagination } from '@latitude-data/core/lib/pagination/buildPagination'
-import { useCallback, useMemo, useState } from 'react'
+import { Suspense, useCallback, useMemo, useState } from 'react'
 import { LinkableTablePaginationFooter } from '$/components/TablePaginationFooter'
 import { ClientDatasetRow } from '$/stores/datasetRows/rowSerializationHelpers'
 import useDatasetRows from '$/stores/datasetRows'
+import useDatasets from '$/stores/datasetsV2'
+import { UpdateColumnModal } from '$/app/(private)/datasets/[datasetId]/DatasetDetailTable/DataGrid/UpdateColumnModal'
 
 function rowKeyGetter(row: ClientDatasetRow) {
   return row.id
@@ -66,14 +77,41 @@ function renderEditCell(props: RenderEditCellProps<ClientDatasetRow, unknown>) {
   const initialValue = rawValue === undefined ? '' : String(rawValue)
   const valueType = typeof storedRowValue === 'object' ? 'json' : 'text'
   return (
-    <DataGridCellEditor
-      valueType={valueType}
-      value={initialValue}
-      onChange={onChange}
-    />
+    <Suspense fallback={null}>
+      <DataGridCellEditor
+        valueType={valueType}
+        value={initialValue}
+        onChange={onChange}
+      />
+    </Suspense>
   )
 }
 
+const renderHeaderCell =
+  ({
+    setEditColumnKey,
+    column,
+  }: {
+    setEditColumnKey: ReactStateDispatch<string | null>
+    column: DatasetV2['columns'][0]
+  }) =>
+  (props: RenderHeaderCellProps<ClientDatasetRow>) => {
+    const onClickEdit = useCallback(() => {
+      setEditColumnKey(props.column.key)
+    }, [props.column.key, setEditColumnKey])
+    return (
+      <div className='flex items-center gap-x-2'>
+        <DatasetHeadText text={column.name} role={column.role} />
+        <Button
+          className='opacity-30 group-hover/cell-header:opacity-100'
+          variant='nope'
+          iconProps={{ name: 'pencil', color: 'foregroundMuted' }}
+          onClick={onClickEdit}
+        />
+      </div>
+    )
+  }
+
 type Props = DatasetRowsTableProps & {
   updateRows: ReturnType<typeof useDatasetRows>['updateRows']
   deleteRows: ReturnType<typeof useDatasetRows>['deleteRows']
@@ -82,14 +120,21 @@ type Props = DatasetRowsTableProps & {
 
 const countLabel = (count: number) => `${count} rows`
 export default function DataGrid({
-  dataset,
+  dataset: serverDataset,
   rows,
   updateRows,
   deleteRows,
   isDeleting,
   pagination,
+  datasetCellRoleStyles,
 }: Props) {
+  const { data } = useDatasets({}, { fallbackData: [serverDataset] })
+  const dataset = useMemo(() => {
+    return data.find((d) => d.id === serverDataset.id) ?? serverDataset
+  }, [data, serverDataset.id])
+  const { backgroundCssClasses } = datasetCellRoleStyles
   const [selectedRows, setSelectedRows] = useState(() => new Set<number>())
+  const [editColumnKey, setEditColumnKey] = useState<string | null>(null)
   const columns = useMemo<DataGridProps<ClientDatasetRow>['columns']>(() => {
     const dataColumns: DataGridProps<ClientDatasetRow>['columns'] =
       dataset.columns.map((col) => ({
@@ -98,7 +143,16 @@ export default function DataGrid({
         resizable: true,
         selectable: true,
         minWidth: 80,
+        headerCellClass: cn(
+          'group/cell-header',
+          backgroundCssClasses[col.role],
+        ),
         renderEditCell,
+        renderHeaderCell: renderHeaderCell({
+          column: col,
+          setEditColumnKey,
+        }),
+        cellClass: backgroundCssClasses[col.role],
         renderCell,
       }))
 
@@ -136,6 +190,13 @@ export default function DataGrid({
 
   return (
     <>
+      {editColumnKey ? (
+        <UpdateColumnModal
+          dataset={dataset}
+          columnKey={editColumnKey}
+          onClose={() => setEditColumnKey(null)}
+        />
+      ) : null}
       <BaseDataGrid<ClientDatasetRow, unknown, number>
         rowKeyGetter={rowKeyGetter}
         rows={rows}

apps/web/src/stores/datasetsV2.ts

@@ -8,6 +8,7 @@ import useLatitudeAction from '$/hooks/useLatitudeAction'
 import { ROUTES } from '$/services/routes'
 import useSWR, { SWRConfiguration } from 'swr'
 import { compactObject } from '@latitude-data/core/lib/compactObject'
+import { updateDatasetColumnAction } from '$/actions/datasetsV2/updateColumn'
 
 const EMPTY_ARRAY: DatasetV2[] = []
 export default function useDatasets(
@@ -73,6 +74,16 @@ export default function useDatasets(
     },
   })
 
+  const { execute: updateColumn, isPending: isUpdatingColumn } =
+    useLatitudeAction<typeof updateDatasetColumnAction>(
+      updateDatasetColumnAction,
+      {
+        onSuccess: ({ data: dataset }) => {
+          mutate(data.map((ds) => (ds.id === dataset.id ? dataset : ds)))
+        },
+      },
+    )
+
   return {
     data,
     mutate,
@@ -81,6 +92,8 @@ export default function useDatasets(
     createError,
     destroy,
     isDestroying,
+    updateColumn,
+    isUpdatingColumn,
     ...rest,
   }
 }