bug: Dataset Loading Performance

packages/client/src/components/DatasetTable.component.tsx

@@ -1,7 +1,7 @@
 import { DataGrid, GridColDef, GridRowId, GridActionsCellItem } from '@mui/x-data-grid';
 import { useState, useEffect } from 'react';
 import { Dataset, Entry } from '../graphql/graphql';
-import { useEntryForDatasetLazyQuery } from '../graphql/entry/entry';
+import { useEntryForDatasetLazyQuery, useCountEntryForDatasetLazyQuery } from '../graphql/entry/entry';
 import { EntryView } from './EntryView.component';
 import { useTranslation } from 'react-i18next';
 import { useSnackbar } from '../context/Snackbar.context';
@@ -22,6 +22,7 @@ export const DatasetTable: React.FC<DatasetTableProps> = (props) => {
   const [deleteEntryMutation] = useDeleteEntryMutation();
   const confirmation = useConfirmation();
   const [selectedRows, setSelectedRows] = useState<GridRowId[]>([]);
+  const [paginationModel, setPaginationModel] = useState<{ page: number; pageSize: number }>({ page: 0, pageSize: 10 });
 
   const defaultColumns: GridColDef[] = [
     {
@@ -97,19 +98,25 @@ export const DatasetTable: React.FC<DatasetTableProps> = (props) => {
   };
 
   const [entries, setEntries] = useState<Entry[]>([]);
+  const [rowCount, setRowCount] = useState<number>(0);
   const columns = [...defaultColumns, ...(props.additionalColumns ?? [])];
   if (props.supportEntryDelete) {
     columns.push(deleteColumn);
   }
 
   const [entryForDataset, entryForDatasetResult] = useEntryForDatasetLazyQuery();
+  const [entryCount, entryCountResult] = useCountEntryForDatasetLazyQuery();
 
   useEffect(() => {
     reload();
-  }, [props.dataset]);
+  }, [props.dataset, paginationModel]);
 
   const reload = () => {
-    entryForDataset({ variables: { dataset: props.dataset._id }, fetchPolicy: 'network-only' });
+    entryForDataset({
+      variables: { dataset: props.dataset._id, page: paginationModel.page, pageSize: paginationModel.pageSize },
+      fetchPolicy: 'network-only'
+    });
+    entryCount({ variables: { dataset: props.dataset._id } });
   };
 
   // TODO: Add in logic to re-fetch data when the presigned URL expires
@@ -122,11 +129,24 @@ export const DatasetTable: React.FC<DatasetTableProps> = (props) => {
     }
   }, [entryForDatasetResult]);
 
+  useEffect(() => {
+    if (entryCountResult.data) {
+      setRowCount(entryCountResult.data.countEntryForDataset);
+    } else if (entryCountResult.error) {
+      pushSnackbarMessage(t('errors.entryQuery'), 'error');
+      console.error(entryForDatasetResult.error);
+    }
+  }, [entryCountResult]);
+
   return (
     <DataGrid
       getRowHeight={() => 'auto'}
       rows={entries}
+      rowCount={rowCount}
       columns={columns}
+      paginationMode={'server'}
+      paginationModel={paginationModel}
+      onPaginationModelChange={setPaginationModel}
       initialState={{
         pagination: {
           paginationModel: {

packages/client/src/graphql/entry/entry.graphql

@@ -1,5 +1,5 @@
-query entryForDataset($dataset: ID!) {
-  entryForDataset(dataset: $dataset) {
+query entryForDataset($dataset: ID!, $page: Int, $pageSize: Int) {
+  entryForDataset(dataset: $dataset, page: $page, pageSize: $pageSize) {
     _id
     organization
     entryID
@@ -14,6 +14,10 @@ query entryForDataset($dataset: ID!) {
   }
 }
 
+ query countEntryForDataset($dataset: ID!) {
+   countEntryForDataset(dataset: $dataset)
+ }
+
 query entryFromID($entry: ID!) {
   entryFromID(entry: $entry) {
     _id

packages/client/src/graphql/entry/entry.ts

@@ -7,11 +7,20 @@ import * as Apollo from '@apollo/client';
 const defaultOptions = {} as const;
 export type EntryForDatasetQueryVariables = Types.Exact<{
   dataset: Types.Scalars['ID']['input'];
+  page?: Types.InputMaybe<Types.Scalars['Int']['input']>;
+  pageSize?: Types.InputMaybe<Types.Scalars['Int']['input']>;
 }>;
 
 
 export type EntryForDatasetQuery = { __typename?: 'Query', entryForDataset: Array<{ __typename?: 'Entry', _id: string, organization: string, entryID: string, contentType: string, dataset: string, creator: string, dateCreated: any, meta?: any | null, signedUrl: string, signedUrlExpiration: number, isTraining: boolean }> };
 
+export type CountEntryForDatasetQueryVariables = Types.Exact<{
+  dataset: Types.Scalars['ID']['input'];
+}>;
+
+
+export type CountEntryForDatasetQuery = { __typename?: 'Query', countEntryForDataset: number };
+
 export type EntryFromIdQueryVariables = Types.Exact<{
   entry: Types.Scalars['ID']['input'];
 }>;
@@ -28,8 +37,8 @@ export type DeleteEntryMutation = { __typename?: 'Mutation', deleteEntry: boolea
 
 
 export const EntryForDatasetDocument = gql`
-    query entryForDataset($dataset: ID!) {
-  entryForDataset(dataset: $dataset) {
+    query entryForDataset($dataset: ID!, $page: Int, $pageSize: Int) {
+  entryForDataset(dataset: $dataset, page: $page, pageSize: $pageSize) {
     _id
     organization
     entryID
@@ -58,6 +67,8 @@ export const EntryForDatasetDocument = gql`
  * const { data, loading, error } = useEntryForDatasetQuery({
  *   variables: {
  *      dataset: // value for 'dataset'
+ *      page: // value for 'page'
+ *      pageSize: // value for 'pageSize'
  *   },
  * });
  */
@@ -72,6 +83,39 @@ export function useEntryForDatasetLazyQuery(baseOptions?: Apollo.LazyQueryHookOp
 export type EntryForDatasetQueryHookResult = ReturnType<typeof useEntryForDatasetQuery>;
 export type EntryForDatasetLazyQueryHookResult = ReturnType<typeof useEntryForDatasetLazyQuery>;
 export type EntryForDatasetQueryResult = Apollo.QueryResult<EntryForDatasetQuery, EntryForDatasetQueryVariables>;
+export const CountEntryForDatasetDocument = gql`
+    query countEntryForDataset($dataset: ID!) {
+  countEntryForDataset(dataset: $dataset)
+}
+    `;
+
+/**
+ * __useCountEntryForDatasetQuery__
+ *
+ * To run a query within a React component, call `useCountEntryForDatasetQuery` and pass it any options that fit your needs.
+ * When your component renders, `useCountEntryForDatasetQuery` returns an object from Apollo Client that contains loading, error, and data properties
+ * you can use to render your UI.
+ *
+ * @param baseOptions options that will be passed into the query, supported options are listed on: https://www.apollographql.com/docs/react/api/react-hooks/#options;
+ *
+ * @example
+ * const { data, loading, error } = useCountEntryForDatasetQuery({
+ *   variables: {
+ *      dataset: // value for 'dataset'
+ *   },
+ * });
+ */
+export function useCountEntryForDatasetQuery(baseOptions: Apollo.QueryHookOptions<CountEntryForDatasetQuery, CountEntryForDatasetQueryVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useQuery<CountEntryForDatasetQuery, CountEntryForDatasetQueryVariables>(CountEntryForDatasetDocument, options);
+      }
+export function useCountEntryForDatasetLazyQuery(baseOptions?: Apollo.LazyQueryHookOptions<CountEntryForDatasetQuery, CountEntryForDatasetQueryVariables>) {
+          const options = {...defaultOptions, ...baseOptions}
+          return Apollo.useLazyQuery<CountEntryForDatasetQuery, CountEntryForDatasetQueryVariables>(CountEntryForDatasetDocument, options);
+        }
+export type CountEntryForDatasetQueryHookResult = ReturnType<typeof useCountEntryForDatasetQuery>;
+export type CountEntryForDatasetLazyQueryHookResult = ReturnType<typeof useCountEntryForDatasetLazyQuery>;
+export type CountEntryForDatasetQueryResult = Apollo.QueryResult<CountEntryForDatasetQuery, CountEntryForDatasetQueryVariables>;
 export const EntryFromIdDocument = gql`
     query entryFromID($entry: ID!) {
   entryFromID(entry: $entry) {

packages/client/src/graphql/graphql.ts

@@ -468,6 +468,7 @@ export type ProjectPermissionModel = {
 
 export type Query = {
   __typename?: 'Query';
+  countEntryForDataset: Scalars['Int']['output'];
   datasetExists: Scalars['Boolean']['output'];
   entryForDataset: Array<Entry>;
   entryFromID: Entry;
@@ -500,13 +501,20 @@ export type Query = {
 };
 
 
+export type QueryCountEntryForDatasetArgs = {
+  dataset: Scalars['ID']['input'];
+};
+
+
 export type QueryDatasetExistsArgs = {
   name: Scalars['String']['input'];
 };
 
 
 export type QueryEntryForDatasetArgs = {
   dataset: Scalars['ID']['input'];
+  page?: InputMaybe<Scalars['Int']['input']>;
+  pageSize?: InputMaybe<Scalars['Int']['input']>;
 };
 
 

packages/server/src/entry/resolvers/entry.resolver.ts

@@ -1,4 +1,4 @@
-import { Args, ID, Resolver, Query, ResolveField, Parent, Mutation } from '@nestjs/graphql';
+import { Args, ID, Resolver, Query, ResolveField, Parent, Mutation, Int } from '@nestjs/graphql';
 import { Dataset } from '../../dataset/dataset.model';
 import { Entry } from '../models/entry.model';
 import { EntryService } from '../services/entry.service';
@@ -28,13 +28,27 @@ export class EntryResolver {
   @Query(() => [Entry])
   async entryForDataset(
     @Args('dataset', { type: () => ID }, DatasetPipe) dataset: Dataset,
-    @TokenContext() user: TokenPayload
+    @TokenContext() user: TokenPayload,
+    @Args('page', { type: () => Int, nullable: true }) page?: number,
+    @Args('pageSize', { type: () => Int, nullable: true }) pageSize?: number
   ): Promise<Entry[]> {
     if (!(await this.enforcer.enforce(user.user_id, DatasetPermissions.READ, dataset._id.toString()))) {
       throw new UnauthorizedException('User cannot read entries on this dataset');
     }
 
-    return this.entryService.findForDataset(dataset);
+    return this.entryService.findForDataset(dataset, page, pageSize);
+  }
+
+  @Query(() => Int)
+  async countEntryForDataset(
+    @Args('dataset', { type: () => ID }, DatasetPipe) dataset: Dataset,
+    @TokenContext() user: TokenPayload
+  ): Promise<Number> {
+    if (!(await this.enforcer.enforce(user.user_id, DatasetPermissions.READ, dataset._id.toString()))) {
+      throw new UnauthorizedException('User cannot read entries on this dataset');
+    }
+
+    return this.entryService.countForDataset(dataset);
   }
 
   @Query(() => Entry)

packages/server/src/entry/services/entry.service.ts

@@ -48,7 +48,7 @@ export class EntryService {
     await this.entryModel.deleteOne({ _id: entry._id });
   }
 
-  async findForDataset(dataset: Dataset | string): Promise<Entry[]> {
+  async findForDataset(dataset: Dataset | string, page?: number, pageSize?: number): Promise<Entry[]> {
     let id: string = '';
 
     if (typeof dataset === 'string') {
@@ -57,7 +57,26 @@ export class EntryService {
       id = dataset._id.toString();
     }
 
-    return this.entryModel.find({ dataset: id, isTraining: false });
+    const query = this.entryModel.find({ dataset: id, isTraining: false });
+
+    if (page !== undefined && pageSize !== undefined) {
+      const offset = page * pageSize;
+      return await query.skip(offset).limit(pageSize);
+    }
+
+    return query;
+  }
+
+  async countForDataset(dataset: Dataset | string) {
+    let id: string = '';
+
+    if (typeof dataset === 'string') {
+      id = dataset;
+    } else {
+      id = dataset._id.toString();
+    }
+
+    return this.entryModel.count({ dataset: id, isTraining: false });
   }
 
   async exists(entryID: string, dataset: Dataset): Promise<boolean> {