refactor: Make reusable AsyncTable

src/common/sortingHelpers.ts

@@ -1,7 +1,7 @@
 export type SortingDirection = 'asc' | 'desc'
 
 export interface SortingState {
-    rowId: string;
+    columnId: string;
     direction: SortingDirection;
 }
 

src/components/AsyncTable/index.tsx

@@ -0,0 +1,265 @@
+import {
+    sceneGraph, 
+    SceneObject,
+    SceneObjectState,
+    SceneObjectBase,
+    SceneComponentProps,
+    SceneVariables,
+    SceneVariableSet,
+} from '@grafana/scenes';
+import React, { useEffect, useMemo } from 'react';
+import { DataFrameView } from '@grafana/data';
+import { InteractiveTable } from '../InteractiveTable/InterativeTable';
+import { LinkCell } from 'pages/Workloads/components/LinkCell';
+import { asyncQueryRunner } from 'common/queryHelpers';
+import { resolveVariable } from 'common/variableHelpers';
+import { CellContext, ColumnDef, ColumnSort, Row } from '@tanstack/react-table';
+import { FormattedCell } from 'pages/Workloads/components/FormattedCell';
+import { SortingState } from 'common/sortingHelpers';
+
+export type CellType = 'link' | 'formatted'
+
+type LinkCellProps = {
+    urlBuilder?: (value: any) => string;
+}
+
+export type FormattedCellProps = {
+    decimals?: number;
+    format?: string;
+}
+
+export type CellProps = LinkCellProps | FormattedCellProps;
+
+export interface ColumnSortingConfig<TableRow> {
+    enabled: boolean;
+    type?: 'label' | 'value';
+    local?: boolean;
+    compare?: (a: TableRow, b: TableRow, direction: 'asc' | 'desc') => number;
+}
+
+export interface Column<TableRow> {
+    id: string;
+    header: string;
+    accessor?: (row: TableRow) => any;
+    cellType?: CellType;
+    cellProps?: CellProps;
+    sortingConfig: ColumnSortingConfig<TableRow>;
+    columns?: Array<Column<TableRow>>;
+}
+
+interface TableState<TableRow> extends SceneObjectState {
+    expandedRows?: SceneObject[];
+    asyncRowData?: Map<string, number[]>;
+    visibleRowIds?: string;
+    sorting?: SortingState;
+    columns: Array<Column<TableRow>>;
+    createRowId: (row: TableRow) => string;
+    asyncDataRowMapper: (row: TableRow, asyncRowData: any) => void;
+    rootQueryBuilder: (variables: SceneVariables | SceneVariableSet, sorting: SortingState, sortingConfig: ColumnSortingConfig<TableRow>) => any;
+    rowQueryBuilder: (rows: TableRow[], variables: SceneVariables | SceneVariableSet) => any;
+    expandedRowBuilder?: (row: TableRow) => SceneObject;
+}
+
+interface ExpandedRowProps<TableRow> {
+    table: AsyncTable<TableRow>;
+    row: TableRow;
+}
+
+function ExpandedRow<TableRow>({ table, row }: ExpandedRowProps<TableRow>) {
+    const { expandedRows } = table.useState();
+
+    const rowScene = expandedRows?.find((scene) => scene.state.key === table.state.createRowId(row));
+
+    useEffect(() => {
+      if (!rowScene && table.state.expandedRowBuilder) {
+        const newRowScene = table.state.expandedRowBuilder(row);
+        table.setState({ expandedRows: [...(table.state.expandedRows ?? []), newRowScene] });
+      }
+    }, [row, table, rowScene]);
+
+    return rowScene ? <rowScene.Component model={rowScene} /> : null;
+}
+
+function mapColumn<TableRow>(column: Column<TableRow>): ColumnDef<TableRow> {
+
+    let cell = undefined;
+    switch (column.cellType) {
+        case 'link':
+            cell = (props: CellContext<TableRow, any>) => LinkCell('namespaces', props.row.getValue(column.id))
+            break;
+        case 'formatted':
+            const formattedCellProps = column.cellProps as FormattedCellProps;
+            cell = (props: CellContext<TableRow, any>) => FormattedCell({
+                value: props.row.getValue(column.id),
+                decimals: formattedCellProps.decimals,
+                format: formattedCellProps.format,
+            })
+            break;
+    }
+
+    return {
+        id: column.id,
+        header: column.header,
+        enableSorting: column.sortingConfig.enabled,
+        accessorFn: column.accessor,
+        cell: cell,
+        columns: column.columns?.map((column) => mapColumn(column)),
+    }
+}
+
+export class AsyncTable<TableRow> extends SceneObjectBase<TableState<TableRow>> {
+
+    private onSortFn  = this.onSort.bind(this);;
+    private onRowsChangedFn = this.onRowsChanged.bind(this);
+    private createRowIdFn = this.createRowId.bind(this);
+
+    constructor(state: TableState<TableRow>) {
+        super({ ...state, asyncRowData: new Map<string, number[]>() });
+    }
+
+    private setAsyncRowData(data: any) {
+        this.setState({ ...this.state, asyncRowData: data });
+    }
+
+    private setVisibleRowIds(ids: string) {
+        this.setState({ ...this.state, visibleRowIds: ids });
+    }
+
+    private createRowId(row: TableRow) {
+        return this.state.createRowId(row);
+    }
+
+    private getColumnById(id: string) {
+
+        const findColumn = (columns: Array<Column<TableRow>>): Column<TableRow> | undefined => {
+            for (const column of columns) {
+                if (column.id === id) {
+                    return column;
+                }
+
+                if (column.columns) {
+                    const found = findColumn(column.columns);
+                    if (found) {
+                        return found;
+                    }
+                }
+            }
+
+            return undefined;
+        }
+
+        return findColumn(this.state.columns);
+    }
+
+    private onRowsChanged(rows: Array<Row<TableRow>>) {
+
+        const ids = rows.map((row: Row<TableRow>) => this.createRowId(row.original)).join(',');
+
+        if (!ids || ids.length === 0 || this.state.visibleRowIds === ids) {
+            return;
+        }
+
+        const sceneVariables = sceneGraph.getVariables(this)
+        const timeRange = sceneGraph.getTimeRange(this)
+        const datasourceVariable = resolveVariable(sceneVariables, 'datasource')
+
+        asyncQueryRunner({
+            datasource: {
+                uid: datasourceVariable?.toString(),
+                type: 'prometheus',
+            },
+
+            queries: [
+                ...this.state.rowQueryBuilder(rows.map(row => row.original), sceneVariables),
+            ],
+            $timeRange: timeRange.clone(),
+        }).then((data) => {
+            this.setVisibleRowIds(ids);
+            this.setAsyncRowData(data);
+        });
+    }
+
+    private onSort(newSorting: ColumnSort[]) {
+        if (newSorting && newSorting.length > 0) {
+
+            const newSortingState: SortingState = {
+                columnId: newSorting[0].id,
+                direction: newSorting[0].desc ? 'desc' : 'asc'
+            }
+
+            const newState: TableState<TableRow> = {
+                ...this.state,
+                sorting: newSortingState
+            }
+
+            const sortingConfig = this.getColumnById(newSortingState.columnId)?.sortingConfig;
+
+            if (sortingConfig && sortingConfig.local === false) {
+                newState.$data = this.state.rootQueryBuilder(sceneGraph.getVariables(this), newSortingState, sortingConfig)
+            }
+
+            this.setState(newState)
+        }
+    }
+
+    static Component = (props: SceneComponentProps<AsyncTable<any>>) => {
+        const { data } = sceneGraph.getData(props.model).useState();
+        const { asyncRowData, sorting, columns, asyncDataRowMapper } = props.model.useState();
+        const sortingConfig = sorting ? props.model.getColumnById(sorting.columnId)?.sortingConfig : undefined;
+
+        const columnDefs = useMemo(
+            () => {
+                return columns.map((column) => {
+                    return mapColumn(column);
+                })
+            },
+            [columns]
+        );
+
+        const tableData = useMemo(() => {
+            if (!data || data.series.length === 0) {
+                return [];
+            }
+
+            const frame = data.series[0];
+            const view = new DataFrameView<any>(frame);
+            const rows = view.toArray();
+
+            for (const row of rows) {
+                asyncDataRowMapper(row, asyncRowData);
+            }
+
+            if (sorting && sortingConfig && sortingConfig.compare && sortingConfig.enabled) {
+                return rows.sort((a, b) => {
+                    return sortingConfig.compare!(a, b, sorting.direction)
+                })
+            }
+
+            return rows;
+        }, [data, asyncRowData, sorting, sortingConfig, asyncDataRowMapper]);
+
+        const currentSorting = useMemo(() => {
+            if (sorting) {
+                return [{
+                    id: sorting.columnId,
+                    desc: sorting.direction === 'desc'
+                }];
+            }
+
+            return [];
+        }, [sorting]);
+
+        return (
+            <InteractiveTable
+                columns={columnDefs}
+                currentSorting={currentSorting}
+                getRowId={props.model.createRowIdFn}
+                data={tableData}
+                renderExpandedRow={ props.model.state.expandedRowBuilder ? (row) => <ExpandedRow table={props.model} row={row} /> : undefined}
+                pageSize={10}
+                onRowsChanged={props.model.onRowsChangedFn}
+                onSort={props.model.onSortFn}
+            />
+        );
+    };
+}