diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx new file mode 100644 index 000000000000..1c014ff512c3 --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx @@ -0,0 +1,147 @@ +import {Box, Popover} from '@mui/material'; +import { + GridEditInputCell, + GridRenderCellParams, + GridRenderEditCellParams, +} from '@mui/x-data-grid-pro'; +import {Icon} from '@wandb/weave/components/Icon'; +import React, {useState} from 'react'; +import styled from 'styled-components'; + +const commonCellStyles = { + height: '100%', + width: '100%', + fontFamily: '"Source Sans Pro", sans-serif', + fontSize: '14px', + lineHeight: '1.5', + padding: '8px 12px', + display: 'flex', + alignItems: 'center', + transition: 'background-color 0.2s ease', +}; + +interface CellViewingRendererProps extends GridRenderCellParams { + isEdited?: boolean; + isDeleted?: boolean; + isNew?: boolean; +} + +const StyledEditCell = styled(GridEditInputCell)` + textarea { + height: 100% !important; + padding: 8px 12px; + font-family: 'Source Sans Pro', sans-serif; + font-size: 14px; + line-height: 1.5; + } + .MuiInputBase-root { + height: 100%; + padding: 0; + } + .MuiInputBase-input { + height: 100% !important; + } +`; + +export const CellViewingRenderer: React.FC = ({ + value, + isEdited = false, + isDeleted = false, + isNew = false, + api, + id, + field, +}) => { + const [isHovered, setIsHovered] = useState(false); + + const handleEditClick = (event: React.MouseEvent) => { + event.stopPropagation(); + api.startCellEditMode({id, field}); + }; + + const getBackgroundColor = () => { + if (isDeleted) { + return 'rgba(255, 0, 0, 0.1)'; + } + if (isEdited) { + return 'rgba(0, 128, 128, 0.1)'; + } + if (isNew) { + return 'rgba(0, 255, 0, 0.1)'; + } + return 'transparent'; + }; + + return ( + setIsHovered(true)} + onMouseLeave={() => setIsHovered(false)} + sx={{ + ...commonCellStyles, + position: 'relative', + cursor: 'pointer', + backgroundColor: getBackgroundColor(), + opacity: isDeleted ? 0.5 : 1, + textDecoration: isDeleted ? 'line-through' : 'none', + '&:hover': { + backgroundColor: 'rgba(0, 0, 0, 0.04)', + }, + }}> + {value} + {isHovered && ( + + + + )} + + ); +}; + +export const CellEditingRenderer: React.FC< + GridRenderEditCellParams +> = params => { + return ( + <> + + + params.api.stopCellEditMode({id: params.id, field: params.field}) + } + anchorOrigin={{ + vertical: 'bottom', + horizontal: 'left', + }} + transformOrigin={{ + vertical: 'top', + horizontal: 'left', + }} + sx={{ + '& .MuiPopover-paper': { + minWidth: '200px', + minHeight: '100px', + padding: '8px', + }, + }}> + + + + ); +}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx new file mode 100644 index 000000000000..5a0bd561e75c --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx @@ -0,0 +1,132 @@ +import React, {createContext, useCallback, useContext, useState} from 'react'; + +interface DatasetRow { + // id: string | number; + [key: string]: any; + ___weave: { + id: string; + // The index must be set for rows of an existing dataset. + index?: number; + isNew?: boolean; + }; +} + +interface EditedCell { + [fieldName: string]: unknown; +} + +interface DatasetEditContextType { + /** Map of edited cells, keyed by row absolute index */ + editedCellsMap: Map; + setEditedCellsMap: React.Dispatch< + React.SetStateAction> + >; + /** Map of complete edited rows, keyed by row absolute index */ + editedRows: Map; + setEditedRows: React.Dispatch>>; + /** Callback to process row updates from the data grid */ + processRowUpdate: (newRow: DatasetRow, oldRow: DatasetRow) => DatasetRow; + /** Array of row indices that have been marked for deletion */ + deletedRows: number[]; + setDeletedRows: React.Dispatch>; + /** Map of newly added rows, keyed by temporary row ID */ + addedRows: Map; + setAddedRows: React.Dispatch>>; + /** Reset the context to its initial state */ + reset: () => void; +} + +export const DatasetEditContext = createContext< + DatasetEditContextType | undefined +>(undefined); + +export const useDatasetEditContext = () => { + const context = useContext(DatasetEditContext); + if (!context) { + throw new Error( + 'useDatasetEditContext must be used within a DatasetEditProvider' + ); + } + return context; +}; + +interface DatasetEditProviderProps { + children: React.ReactNode; +} + +export const DatasetEditProvider: React.FC = ({ + children, +}) => { + const [editedCellsMap, setEditedCellsMap] = useState>( + new Map() + ); + const [editedRows, setEditedRows] = useState>( + new Map() + ); + const [deletedRows, setDeletedRows] = useState([]); + const [addedRows, setAddedRows] = useState>( + new Map() + ); + + const processRowUpdate = useCallback( + (newRow: DatasetRow, oldRow: DatasetRow): DatasetRow => { + const changedField = Object.keys(newRow).find( + key => newRow[key] !== oldRow[key] && key !== 'id' + ); + + if (changedField) { + const rowKey = String(oldRow.___weave.id); + const rowIndex = oldRow.___weave.index; + if (oldRow.___weave.isNew) { + setAddedRows(prev => { + const updatedMap = new Map(prev); + updatedMap.set(rowKey, newRow); + return updatedMap; + }); + } else { + setEditedCellsMap(prev => { + const existingEdits = prev.get(rowIndex!) || {}; + const updatedMap = new Map(prev); + updatedMap.set(rowIndex!, { + ...existingEdits, + [changedField]: newRow[changedField], + }); + return updatedMap; + }); + setEditedRows(prev => { + const updatedMap = new Map(prev); + updatedMap.set(rowIndex!, newRow); + return updatedMap; + }); + } + } + return newRow; + }, + [] + ); + + const reset = useCallback(() => { + setEditedCellsMap(new Map()); + setEditedRows(new Map()); + setDeletedRows([]); + setAddedRows(new Map()); + }, []); + + return ( + + {children} + + ); +}; diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx index f3aed14e9efd..3e80255554f9 100644 --- a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx @@ -1,13 +1,15 @@ -import Box from '@mui/material/Box'; -import React, {useMemo} from 'react'; +import {Box, Popover, Typography} from '@mui/material'; +import React, {useCallback, useMemo, useState} from 'react'; +import {Link, useHistory} from 'react-router-dom'; +import {toast} from 'react-toastify'; +import {Button} from '../../../../Button'; import {Icon} from '../../../../Icon'; import {LoadingDots} from '../../../../LoadingDots'; import {Tailwind} from '../../../../Tailwind'; +import {useWeaveflowCurrentRouteContext} from '../context'; import {WeaveCHTableSourceRefContext} from '../pages/CallPage/DataTableView'; -import {ObjectViewerSection} from '../pages/CallPage/ObjectViewerSection'; -import {objectVersionText} from '../pages/common/Links'; -import {ObjectVersionsLink} from '../pages/common/Links'; +import {ObjectVersionsLink, objectVersionText} from '../pages/common/Links'; import {CenteredAnimatedLoader} from '../pages/common/Loader'; import { ScrollableTabContent, @@ -15,25 +17,60 @@ import { } from '../pages/common/SimplePageLayout'; import {TabUseDataset} from '../pages/TabUseDataset'; import {useWFHooks} from '../pages/wfReactInterface/context'; +import {TableUpdateSpec} from '../pages/wfReactInterface/traceServerClientTypes'; import {objectVersionKeyToRefUri} from '../pages/wfReactInterface/utilities'; import {ObjectVersionSchema} from '../pages/wfReactInterface/wfDataModelHooksInterface'; import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher'; +import {useDatasetEditContext} from './DatasetEditorContext'; +import {EditableDatasetView} from './EditableDatasetView'; + +const PUBLISHED_LINK_STYLES = { + color: 'rgb(94, 234, 212)', + textDecoration: 'none', + fontFamily: 'Inconsolata', + fontWeight: 600, +} as const; + +const POPOVER_STYLES = { + '& .MuiPopover-paper': { + marginTop: '8px', + marginRight: '8px', + }, +} as const; + +const CODE_STYLES = { + fontFamily: 'Inconsolata', + fontWeight: 600, + backgroundColor: 'rgba(0, 0, 0, 0.04)', + padding: '2px 4px', + borderRadius: '4px', +} as const; export const DatasetVersionPage: React.FC<{ objectVersion: ObjectVersionSchema; -}> = ({objectVersion}) => { - const {useRootObjectVersions, useRefsData} = useWFHooks(); + refExtra?: string; +}> = ({objectVersion, refExtra}) => { + const {editedCellsMap, editedRows, deletedRows, addedRows} = + useDatasetEditContext(); + const router = useWeaveflowCurrentRouteContext(); + const {useRootObjectVersions, useRefsData, useTableUpdate, useObjCreate} = + useWFHooks(); + + const [isEditing, setIsEditing] = useState(false); + const [publishAnchorEl, setPublishAnchorEl] = useState( + null + ); + const entityName = objectVersion.entity; const projectName = objectVersion.project; const objectName = objectVersion.objectId; const objectVersionIndex = objectVersion.versionIndex; + const projectId = `${entityName}/${projectName}`; const objectVersions = useRootObjectVersions( entityName, projectName, - { - objectIds: [objectName], - }, + {objectIds: [objectName]}, undefined, true ); @@ -53,12 +90,203 @@ export const DatasetVersionPage: React.FC<{ typeof viewerData !== 'object' || viewerData === null || Array.isArray(viewerData); - if (dataIsPrimitive) { - return {_result: viewerData}; - } - return viewerData; + return dataIsPrimitive ? {_result: viewerData} : viewerData; }, [viewerData]); + const originalTableDigest = viewerDataAsObject?.rows?.split('/').pop() ?? ''; + + const handleEditClick = useCallback(() => setIsEditing(true), []); + const handleCancelClick = useCallback(() => setIsEditing(false), []); + const handlePublishClose = () => setPublishAnchorEl(null); + + const handlePublishClick = useCallback( + (event: React.MouseEvent) => { + if ( + editedCellsMap.size > 0 || + deletedRows.length > 0 || + addedRows.size > 0 + ) { + setPublishAnchorEl(event.currentTarget); + } + }, + [editedCellsMap, deletedRows, addedRows] + ); + + const cleanRow = (row: any) => { + return Object.fromEntries( + Object.entries(row).filter(([key]) => !['___weave'].includes(key)) + ); + }; + + const convertEditsToTableUpdateSpec = useCallback(() => { + const updates: TableUpdateSpec[] = []; + + editedRows.forEach((editedRow, rowIndex) => { + if (rowIndex !== undefined) { + updates.push({pop: {index: rowIndex}}); + updates.push({ + insert: { + index: rowIndex, + row: cleanRow(editedRow), + }, + }); + } + }); + + deletedRows + .sort((a, b) => b - a) + .forEach(rowIndex => { + updates.push({pop: {index: rowIndex}}); + }); + + Array.from(addedRows.values()) + .reverse() + .forEach(row => { + updates.push({ + insert: { + index: 0, + row: cleanRow(row), + }, + }); + }); + + return updates; + }, [editedRows, deletedRows, addedRows]); + + const tableUpdate = useTableUpdate(); + const objCreate = useObjCreate(); + + const history = useHistory(); + + const handlePublish = useCallback(async () => { + setIsEditing(false); + setPublishAnchorEl(null); + + const tableUpdateSpecs = convertEditsToTableUpdateSpec(); + const tableUpdateResp = await tableUpdate( + projectId, + originalTableDigest, + tableUpdateSpecs + ); + const tableRef = `weave:///${projectId}/table/${tableUpdateResp.digest}`; + + const newObjVersion = await objCreate(projectId, objectName, { + ...objectVersion.val, + rows: tableRef, + }); + + const url = router.objectVersionUIUrl( + entityName, + projectName, + objectName, + newObjVersion, + undefined, + refExtra + ); + + toast( +
+ + Published{' '} + + {objectName}:v{objectVersionCount} + +
+ ); + history.push(url); + }, [ + objectVersionCount, + history, + refExtra, + router, + objectName, + objectVersion.val, + convertEditsToTableUpdateSpec, + projectId, + objCreate, + tableUpdate, + originalTableDigest, + entityName, + projectName, + ]); + + const renderEditingControls = () => ( +
+ + + Editing dataset + + + +
+ ); + + const renderPublishPopover = () => ( + + + + Publish changes to a new version of{' '} + {objectName}? + + + + + + + + ); + return ( -
-
-

Name

- -
- {objectName} - {objectVersions.loading ? ( - - ) : ( - - ({objectVersionCount} version - {objectVersionCount !== 1 ? 's' : ''}) - - )} - -
-
+
+
+
+

Name

+ +
+ {objectName} + {objectVersions.loading ? ( + + ) : ( + + ({objectVersionCount} version + {objectVersionCount !== 1 ? 's' : ''}) + + )} + +
+
+
+
+

Version

+

{objectVersionIndex}

+
-
-

Version

-

{objectVersionIndex}

+
+ {isEditing ? ( + renderEditingControls() + ) : ( + + )}
+ {renderPublishPopover()}
} @@ -114,22 +359,16 @@ export const DatasetVersionPage: React.FC<{ label: 'Rows', content: ( - + {data.loading ? ( ) : ( - diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx new file mode 100644 index 000000000000..94a811c3ca58 --- /dev/null +++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx @@ -0,0 +1,495 @@ +import {Box} from '@mui/material'; +import { + GridColDef, + GridPaginationModel, + GridRenderCellParams, + GridRowModel, + GridSortModel, +} from '@mui/x-data-grid-pro'; +import {A} from '@wandb/weave/common/util/links'; +import {Button} from '@wandb/weave/components/Button'; +import {RowId} from '@wandb/weave/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView'; +import {Tooltip} from '@wandb/weave/components/Tooltip'; +import React, { + FC, + useCallback, + useContext, + useEffect, + useMemo, + useState, +} from 'react'; +import {useHistory} from 'react-router-dom'; +import {v4 as uuidv4} from 'uuid'; + +import {isWeaveObjectRef, parseRef, parseRefMaybe} from '../../../../../react'; +import {flattenObjectPreservingWeaveTypes} from '../../Browse2/browse2Util'; +import {CellValue} from '../../Browse2/CellValue'; +import {useWeaveflowCurrentRouteContext} from '../context'; +import {WeaveCHTableSourceRefContext} from '../pages/CallPage/DataTableView'; +import {TABLE_ID_EDGE_NAME} from '../pages/wfReactInterface/constants'; +import {useWFHooks} from '../pages/wfReactInterface/context'; +import {SortBy} from '../pages/wfReactInterface/traceServerClientTypes'; +import {StyledDataGrid} from '../StyledDataGrid'; +import {CellEditingRenderer, CellViewingRenderer} from './CellRenderers'; +import {useDatasetEditContext} from './DatasetEditorContext'; + +const ADDED_ROW_ID_PREFIX = 'new-'; + +// Dataset object schema as it is stored in the database. +interface DatasetObjectVal { + _type: 'Dataset'; + name: string | null; + description: string | null; + rows: string; + _class_name: 'Dataset'; + _bases: ['Object', 'BaseModel']; +} + +interface EditableDataTableViewProps { + datasetObject: DatasetObjectVal; + isEditing: boolean; +} + +interface ControlCellProps { + params: GridRenderCellParams; + deleteRow: (absoluteIndex: number) => void; + restoreRow: (absoluteIndex: number) => void; + deleteAddedRow: (rowId: string) => void; + isDeleted: boolean; + isNew: boolean; +} + +interface ControlsColumnHeaderProps { + onAddRow: () => void; +} + +const ControlCell: FC = ({ + params, + deleteRow, + restoreRow, + deleteAddedRow, + isDeleted, + isNew, +}) => ( + +