Skip to content

Commit

Permalink
chore(dev): dataset editor mvp
Browse files Browse the repository at this point in the history
  • Loading branch information
bcsherma committed Dec 14, 2024
1 parent cb612d4 commit f3b5f35
Show file tree
Hide file tree
Showing 11 changed files with 833 additions and 280 deletions.
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import {GridPaginationModel} from '@mui/x-data-grid-pro';

const MAX_PAGE_SIZE = 100;
export const DEFAULT_PAGE_SIZE = 100;
export const DEFAULT_PAGE_SIZE = 25;

export const getValidPaginationModel = (
queryPage: string | undefined,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ import {TABLE_ID_EDGE_NAME} from '../wfReactInterface/constants';
import {useWFHooks} from '../wfReactInterface/context';
import {SortBy} from '../wfReactInterface/traceServerClientTypes';

const RowId = styled.span`
export const RowId = styled.span`
font-family: 'Inconsolata', monospace;
`;
RowId.displayName = 'S.RowId';
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,215 @@
import {Box} from '@mui/material';
import {GridPaginationModel, GridRenderCellParams} from '@mui/x-data-grid-pro';
import React, {FC, useCallback, useMemo, useState} from 'react';
import styled from 'styled-components';

import {isWeaveObjectRef, parseRefMaybe} from '../../../../../../react';
import {Tooltip} from '../../../../../Tooltip';
import {flattenObjectPreservingWeaveTypes} from '../../../Browse2/browse2Util';
import {DEFAULT_PAGE_SIZE} from '../../grid/pagination';
import {StyledDataGrid} from '../../StyledDataGrid';
import {A} from '../common/Links';
import {useEditContext} from '../EditContext'; // Import the context
import {useWFHooks} from '../wfReactInterface/context';
import {SortBy} from '../wfReactInterface/traceServerClientTypes';
import {RowId} from './DataTableView'; // Import shared components

type DatasetObjectVal = {
_type: 'Dataset';
name: string | null;
description: string | null;
rows: string;
_class_name: 'Dataset';
_bases: ['Object', 'BaseModel'];
};

interface EditableDataTableViewProps {
datasetObjectId: string;
datasetObject: DatasetObjectVal;
fullHeight?: boolean;
}

// Add styled component for edited cells
const EditedCellHighlightWrapper = styled(Box)`
.edited-cell {
background-color: rgba(25, 118, 210, 0.1);
border-radius: 8px;
transition: background-color 0.5s ease;
}
`;

export const EditableDataTableView: FC<EditableDataTableViewProps> = props => {
const {useTableRowsQuery, useTableQueryStats} = useWFHooks();
const [sortBy] = useState<SortBy[]>([]);
const {editedCellsMap, processRowUpdate, setRowIndices} = useEditContext(); // Use the context

const [paginationModel, setPaginationModel] = useState<GridPaginationModel>({
page: 0,
pageSize: DEFAULT_PAGE_SIZE,
});

// Parse table ref
const parsedRef = useMemo(
() => parseRefMaybe(props.datasetObject.rows),
[props.datasetObject.rows]
);
const lookupKey = useMemo(() => {
if (
parsedRef == null ||
!isWeaveObjectRef(parsedRef) ||
parsedRef.weaveKind !== 'table'
) {
return null;
}
return {
entity: parsedRef.entityName,
project: parsedRef.projectName,
digest: parsedRef.artifactVersion,
};
}, [parsedRef]);

// Fetch row count
const numRowsQuery = useTableQueryStats(
lookupKey?.entity ?? '',
lookupKey?.project ?? '',
lookupKey?.digest ?? '',
{skip: lookupKey == null}
);

// Fetch rows
const fetchQuery = useTableRowsQuery(
lookupKey?.entity ?? '',
lookupKey?.project ?? '',
lookupKey?.digest ?? '',
undefined,
paginationModel.pageSize,
paginationModel.page * paginationModel.pageSize,
sortBy,
{skip: lookupKey == null}
);

// Convert data to list of dictionaries and flatten nested objects
const dataAsListOfDict = useMemo(() => {
return (fetchQuery.result?.rows ?? []).map(row => {
let val = row;
if (val == null) {
return {};
} else if (typeof val === 'object' && !Array.isArray(val)) {
if ('val' in val) {
val = val.val; // Extract val field
}
return flattenObjectPreservingWeaveTypes(val);
}
return {'': val};
});
}, [fetchQuery.result?.rows]);

// Reapply edits when rows are fetched
const rows = useMemo(() => {
if (!fetchQuery.loading && fetchQuery.result?.rows) {
return dataAsListOfDict.map((row, i) => {
const digest = fetchQuery.result!.rows[i].digest;
const rowKey = `${digest}`;
const editedRow = editedCellsMap.get(rowKey);
const baseRow = editedRow ? {...row, ...editedRow} : row;
setRowIndices(prev => {
const updatedMap = new Map(prev);
updatedMap.set(rowKey, i);
return updatedMap;
});
return {
id: digest,
...baseRow,
};
});
}
return [];
}, [
fetchQuery.loading,
fetchQuery.result,
dataAsListOfDict,
editedCellsMap,
setRowIndices,
]);

// Generate columns with cell class names for edited cells
const columns = useMemo(() => {
const firstRow = rows[0] ?? {};
return Object.keys(firstRow).map(field => {
if (field === 'id') {
return {
field,
headerName: 'id',
width: 50,
editable: false,
filterable: false,
sortable: false,
renderCell: (params: GridRenderCellParams) => {
const id = params.value;
const rowLabel = id ? id.slice(-4) : params.id;
const rowSpan = (
<Tooltip trigger={<RowId>{rowLabel}</RowId>} content={id} />
);
return <A onClick={() => {}}>{rowSpan}</A>;
},
};
}
return {
field,
headerName: field,
flex: 1,
editable: true,
sortable: false,
filterable: false,
cellClassName: (params: any) => {
const rowKey = `${params.row.id}`;
const editedRow = editedCellsMap.get(rowKey);
return editedRow && editedRow[field] !== undefined
? 'edited-cell'
: '';
},
};
});
}, [rows, editedCellsMap]);

// Handle pagination model change
const handlePaginationModelChange = useCallback(
(newModel: GridPaginationModel) => {
setPaginationModel(newModel);
},
[]
);

return (
<div
style={{
display: 'flex',
flexDirection: 'column',
width: '100%',
height: props.fullHeight ? '100%' : 'inherit',
}}>
<EditedCellHighlightWrapper>
<StyledDataGrid
density="compact"
rows={rows}
columns={columns}
editMode="cell"
pagination={true}
paginationMode="server"
paginationModel={paginationModel}
onPaginationModelChange={handlePaginationModelChange}
rowCount={numRowsQuery.result?.count ?? 0}
pageSizeOptions={[5, 10, 20, 50, 100]}
disableMultipleColumnsSorting
loading={fetchQuery.loading}
disableRowSelectionOnClick
keepBorders={false}
sx={{
border: 'none',
}}
processRowUpdate={processRowUpdate}
/>
</EditedCellHighlightWrapper>
</div>
);
};
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher';
import {OBJECT_ATTR_EDGE_NAME} from '../wfReactInterface/constants';
import {WeaveCHTable, WeaveCHTableSourceRefContext} from './DataTableView';
import {EditableDataTableView} from './EditableDataTableView';
import {ObjectViewer} from './ObjectViewer';
import {getValueType, traverse} from './traverse';
import {ValueView} from './ValueView';
Expand All @@ -28,10 +29,12 @@ const EXPANDED_IDS_LENGTH = 200;
type Data = Record<string, any>;

type ObjectViewerSectionProps = {
objectId: string;
title: string;
data: Data;
noHide?: boolean;
isExpanded?: boolean;
isEditing?: boolean;
};

const TitleRow = styled.div`
Expand Down Expand Up @@ -90,6 +93,7 @@ const ObjectViewerSectionNonEmptyMemoed = React.memo(
);

const ObjectViewerSectionNonEmpty = ({
objectId,
title,
data,
noHide,
Expand Down Expand Up @@ -229,10 +233,12 @@ const ObjectViewerSectionNonEmpty = ({
};

export const ObjectViewerSection = ({
objectId,
title,
data,
noHide,
isExpanded,
isEditing,
}: ObjectViewerSectionProps) => {
const currentRef = useContext(WeaveCHTableSourceRefContext);

Expand Down Expand Up @@ -261,14 +267,13 @@ export const ObjectViewerSection = ({
if (numKeys === 1 && '_result' in data) {
let value = data._result;
if (isWeaveRef(value)) {
// Little hack to make sure that we render refs
// inside the expansion table view
value = {' ': value};
}
const valueType = getValueType(value);
if (valueType === 'object' || (valueType === 'array' && value.length > 0)) {
return (
<ObjectViewerSectionNonEmptyMemoed
objectId={objectId}
title={title}
data={value}
noHide={noHide}
Expand All @@ -291,9 +296,6 @@ export const ObjectViewerSection = ({
);
}

// Here we have a very special case for when the section is viewing a dataset.
// Instead of rending the generic renderer, we directly render a full-screen
// data table.
if (
data._type === 'Dataset' &&
data._class_name === 'Dataset' &&
Expand All @@ -307,9 +309,22 @@ export const ObjectViewerSection = ({
height: '100%',
overflow: 'hidden',
}}>
<WeaveCHTable tableRefUri={data.rows} fullHeight />
<TitleRow>
<Title>{title}</Title>
</TitleRow>
{isEditing ? (
<EditableDataTableView
datasetObjectId={objectId}
// @ts-expect-error
datasetObject={data}
fullHeight
/>
) : (
<WeaveCHTable tableRefUri={data.rows} fullHeight />
)}
</Box>
);

if (currentRef != null) {
return (
<WeaveCHTableSourceRefContext.Provider
Expand All @@ -327,6 +342,7 @@ export const ObjectViewerSection = ({
data={data}
noHide={noHide}
isExpanded={isExpanded}
objectId={objectId}
/>
);
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
import React, {createContext, useContext} from 'react';

interface EditContextType {
editedCellsMap: Map<string, any>;
setEditedCellsMap: React.Dispatch<React.SetStateAction<Map<string, any>>>;
rowIndices: Map<string, number>;
setRowIndices: React.Dispatch<React.SetStateAction<Map<string, number>>>;
processRowUpdate: (newRow: any, oldRow: any) => any;
}

export const EditContext = createContext<EditContextType | undefined>(
undefined
);

export const useEditContext = () => {
const context = useContext(EditContext);
if (!context) {
throw new Error('useEditContext must be used within an EditProvider');
}
return context;
};
Loading

0 comments on commit f3b5f35

Please sign in to comment.