Merge pull request #401 from jcreedcmu/jcreed/pagination

Implement pagination for BQRS results.

Author: jcreedcmu

extensions/ql-vscode/src/adapt.ts

@@ -101,3 +101,34 @@ export function adaptBqrs(schema: AdaptedSchema, page: DecodedBqrsChunk): RawRes
     rows: page.tuples.map(adaptRow),
   };
 }
+
+/**
+ * This type has two branches; we are in the process of changing from
+ * one to the other. The old way is to parse them inside the webview,
+ * the new way is to parse them in the extension. The main motivation
+ * for this transition is to make pagination possible in such a way
+ * that only one page needs to be sent from the extension to the webview.
+ */
+export type ParsedResultSets = ExtensionParsedResultSets | WebviewParsedResultSets;
+
+/**
+ * The old method doesn't require any nontrivial information to be included here,
+ * just a tag to indicate that it is being used.
+ */
+export interface WebviewParsedResultSets {
+  t: 'WebviewParsed';
+  selectedTable?: string; // when undefined, means 'show default table'
+}
+
+/**
+ * The new method includes which bqrs page is being sent, and the
+ * actual results parsed on the extension side.
+ */
+export interface ExtensionParsedResultSets {
+  t: 'ExtensionParsed';
+  pageNumber: number;
+  numPages: number;
+  selectedTable?: string; // when undefined, means 'show default table'
+  resultSetNames: string[];
+  resultSet: RawResultSet;
+}

extensions/ql-vscode/src/interface-types.ts

@@ -1,6 +1,6 @@
 import * as sarif from 'sarif';
 import { ResolvableLocationValue } from 'semmle-bqrs';
-import { RawResultSet } from './adapt';
+import { ParsedResultSets } from './adapt';
 
 /**
  * Only ever show this many results per run in interpreted results.
@@ -12,6 +12,11 @@ export const INTERPRETED_RESULTS_PER_RUN_LIMIT = 100;
  */
 export const RAW_RESULTS_LIMIT = 10000;
 
+/**
+ * Show this many rows in a raw result table at a time.
+ */
+export const RAW_RESULTS_PAGE_SIZE = 100;
+
 export interface DatabaseInfo {
   name: string;
   databaseUri: string;
@@ -81,9 +86,10 @@ export interface SetStateMsg {
 
   /**
    * An experimental way of providing results from the extension.
-   * Should be undefined unless config.EXPERIMENTAL_BQRS_SETTING is set to true.
+   * Should be in the WebviewParsedResultSets branch of the type
+   * unless config.EXPERIMENTAL_BQRS_SETTING is set to true.
    */
-  resultSets?: RawResultSet[];
+  parsedResultSets: ParsedResultSets;
 }
 
 /** Advance to the next or previous path no in the path viewer */
@@ -101,7 +107,8 @@ export type FromResultsViewMsg =
   | ToggleDiagnostics
   | ChangeRawResultsSortMsg
   | ChangeInterpretedResultsSortMsg
-  | ResultViewLoaded;
+  | ResultViewLoaded
+  | ChangePage;
 
 interface ViewSourceFileMsg {
   t: 'viewSourceFile';
@@ -122,6 +129,12 @@ interface ResultViewLoaded {
   t: 'resultViewLoaded';
 }
 
+interface ChangePage {
+  t: 'changePage';
+  pageNumber: number; // 0-indexed, displayed to the user as 1-indexed
+  selectedTable: string;
+}
+
 export enum SortDirection {
   asc, desc
 }

extensions/ql-vscode/src/interface-utils.ts

@@ -0,0 +1,22 @@
+import { RawResultSet } from "./adapt";
+import { ResultSetSchema } from "semmle-bqrs";
+import { Interpretation } from "./interface-types";
+
+export const SELECT_TABLE_NAME = '#select';
+export const ALERTS_TABLE_NAME = 'alerts';
+
+export type RawTableResultSet = { t: 'RawResultSet' } & RawResultSet;
+export type PathTableResultSet = { t: 'SarifResultSet'; readonly schema: ResultSetSchema; name: string } & Interpretation;
+
+export type ResultSet =
+  | RawTableResultSet
+  | PathTableResultSet;
+
+export function getDefaultResultSet(resultSets: readonly ResultSet[]): string {
+  return getDefaultResultSetName(resultSets.map(resultSet => resultSet.schema.name));
+}
+
+export function getDefaultResultSetName(resultSetNames: readonly string[]): string {
+  // Choose first available result set from the array
+  return [ALERTS_TABLE_NAME, SELECT_TABLE_NAME, resultSetNames[0]].filter(resultSetName => resultSetNames.includes(resultSetName))[0];
+}

extensions/ql-vscode/src/interface.ts

@@ -10,14 +10,15 @@ import { CodeQLCliServer } from './cli';
 import { DatabaseItem, DatabaseManager } from './databases';
 import { showAndLogErrorMessage } from './helpers';
 import { assertNever } from './helpers-pure';
-import { FromResultsViewMsg, Interpretation, INTERPRETED_RESULTS_PER_RUN_LIMIT, IntoResultsViewMsg, QueryMetadata, ResultsPaths, SortedResultSetInfo, SortedResultsMap, InterpretedResultsSortState, SortDirection } from './interface-types';
+import { FromResultsViewMsg, Interpretation, INTERPRETED_RESULTS_PER_RUN_LIMIT, IntoResultsViewMsg, QueryMetadata, ResultsPaths, SortedResultSetInfo, SortedResultsMap, InterpretedResultsSortState, SortDirection, RAW_RESULTS_PAGE_SIZE } from './interface-types';
 import { Logger } from './logging';
 import * as messages from './messages';
 import { CompletedQuery, interpretResults } from './query-results';
 import { QueryInfo, tmpDir } from './run-queries';
 import { parseSarifLocation, parseSarifPlainTextMessage } from './sarif-utils';
-import { adaptSchema, adaptBqrs, RawResultSet } from './adapt';
+import { adaptSchema, adaptBqrs, RawResultSet, ParsedResultSets } from './adapt';
 import { EXPERIMENTAL_BQRS_SETTING } from './config';
+import { getDefaultResultSetName } from './interface-utils';
 
 /**
  * interface.ts
@@ -115,8 +116,13 @@ function sortInterpretedResults(results: Sarif.Result[], sortState: InterpretedR
   }
 }
 
+function numPagesOfResultSet(resultSet: RawResultSet): number {
+  return Math.ceil(resultSet.schema.tupleCount / RAW_RESULTS_PAGE_SIZE);
+}
+
 export class InterfaceManager extends DisposableObject {
   private _displayedQuery?: CompletedQuery;
+  private _interpretation?: Interpretation;
   private _panel: vscode.WebviewPanel | undefined;
   private _panelLoaded = false;
   private _panelLoadedCallBacks: (() => void)[] = [];
@@ -288,6 +294,9 @@ export class InterfaceManager extends DisposableObject {
           query.updateInterpretedSortState(this.cliServer, msg.sortState)
         );
         break;
+      case "changePage":
+        await this.showPageOfResults(msg.selectedTable, msg.pageNumber);
+        break;
       default:
         assertNever(msg);
     }
@@ -339,6 +348,7 @@ export class InterfaceManager extends DisposableObject {
     );
 
     this._displayedQuery = results;
+    this._interpretation = interpretation;
 
     const panel = this.getPanel();
     await this.waitForPanelLoaded();
@@ -364,18 +374,37 @@ export class InterfaceManager extends DisposableObject {
       });
     }
 
-    let resultSets: RawResultSet[] | undefined;
+    const getParsedResultSets = async (): Promise<ParsedResultSets> => {
+      if (EXPERIMENTAL_BQRS_SETTING.getValue()) {
+        const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath, RAW_RESULTS_PAGE_SIZE);
+
+        const resultSetNames = schemas["result-sets"].map(resultSet => resultSet.name);
+
+        // This may not wind up being the page we actually show, if there are interpreted results,
+        // but speculatively send it anyway.
+        const selectedTable = getDefaultResultSetName(resultSetNames);
+        const schema = schemas["result-sets"].find(resultSet => resultSet.name == selectedTable)!;
+        if (schema === undefined) {
+          return { t: 'WebviewParsed' };
+        }
 
-    if (EXPERIMENTAL_BQRS_SETTING.getValue()) {
-      resultSets = [];
-      const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath);
-      for (const schema of schemas["result-sets"]) {
-        const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name);
+        const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name, RAW_RESULTS_PAGE_SIZE, schema.pagination?.offsets[0]);
         const adaptedSchema = adaptSchema(schema);
         const resultSet = adaptBqrs(adaptedSchema, chunk);
-        resultSets.push(resultSet);
+
+        return {
+          t: 'ExtensionParsed',
+          pageNumber: 0,
+          numPages: numPagesOfResultSet(resultSet),
+          resultSet,
+          selectedTable: undefined,
+          resultSetNames
+        };
       }
-    }
+      else {
+        return { t: 'WebviewParsed' };
+      }
+    };
 
     await this.postMessage({
       t: "setState",
@@ -384,14 +413,67 @@ export class InterfaceManager extends DisposableObject {
       resultsPath: this.convertPathToWebviewUri(
         results.query.resultsPaths.resultsPath
       ),
-      resultSets,
+      parsedResultSets: await getParsedResultSets(),
       sortedResultsMap,
       database: results.database,
       shouldKeepOldResultsWhileRendering,
       metadata: results.query.metadata
     });
   }
 
+  /**
+   * Show a page of raw results from the chosen table.
+   */
+  public async showPageOfResults(selectedTable: string, pageNumber: number): Promise<void> {
+    const results = this._displayedQuery;
+    if (results === undefined) {
+      throw new Error('trying to view a page of a query that is not loaded');
+    }
+
+    const sortedResultsMap: SortedResultsMap = {};
+    results.sortedResultsInfo.forEach(
+      (v, k) =>
+        (sortedResultsMap[k] = this.convertPathPropertiesToWebviewUris(
+          v
+        ))
+    );
+
+    const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath, RAW_RESULTS_PAGE_SIZE);
+
+    const resultSetNames = schemas["result-sets"].map(resultSet => resultSet.name);
+
+    const schema = schemas["result-sets"].find(resultSet => resultSet.name == selectedTable)!;
+    if (schema === undefined)
+      throw new Error(`Query result set '${selectedTable}' not found.`);
+
+    const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name, RAW_RESULTS_PAGE_SIZE, schema.pagination?.offsets[pageNumber]);
+    const adaptedSchema = adaptSchema(schema);
+    const resultSet = adaptBqrs(adaptedSchema, chunk);
+
+    const parsedResultSets: ParsedResultSets = {
+      t: 'ExtensionParsed',
+      pageNumber,
+      resultSet,
+      numPages: numPagesOfResultSet(resultSet),
+      selectedTable: selectedTable,
+      resultSetNames
+    };
+
+    await this.postMessage({
+      t: "setState",
+      interpretation: this._interpretation,
+      origResultsPaths: results.query.resultsPaths,
+      resultsPath: this.convertPathToWebviewUri(
+        results.query.resultsPaths.resultsPath
+      ),
+      parsedResultSets,
+      sortedResultsMap,
+      database: results.database,
+      shouldKeepOldResultsWhileRendering: false,
+      metadata: results.query.metadata
+    });
+  }
+
   private async getTruncatedResults(
     metadata: QueryMetadata | undefined,
     resultsPaths: ResultsPaths,

extensions/ql-vscode/src/view/alert-table.tsx

@@ -5,9 +5,10 @@ import * as Keys from '../result-keys';
 import { LocationStyle } from 'semmle-bqrs';
 import * as octicons from './octicons';
 import { className, renderLocation, ResultTableProps, zebraStripe, selectableZebraStripe, jumpToLocation, nextSortDirection } from './result-table-utils';
-import { PathTableResultSet, onNavigation, NavigationEvent, vscode } from './results';
+import { onNavigation, NavigationEvent, vscode } from './results';
 import { parseSarifPlainTextMessage, parseSarifLocation } from '../sarif-utils';
 import { InterpretedResultsSortColumn, SortDirection, InterpretedResultsSortState } from '../interface-types';
+import { PathTableResultSet } from '../interface-utils';
 
 export type PathTableProps = ResultTableProps & { resultSet: PathTableResultSet };
 export interface PathTableState {

extensions/ql-vscode/src/view/raw-results-table.tsx

@@ -1,12 +1,14 @@
 import * as React from "react";
 import { renderLocation, ResultTableProps, zebraStripe, className, nextSortDirection } from "./result-table-utils";
-import { RawTableResultSet, vscode } from "./results";
+import { vscode } from "./results";
 import { ResultValue } from "../adapt";
 import { SortDirection, RAW_RESULTS_LIMIT, RawResultsSortState } from "../interface-types";
+import { RawTableResultSet } from "../interface-utils";
 
 export type RawTableProps = ResultTableProps & {
   resultSet: RawTableResultSet;
   sortState?: RawResultsSortState;
+  offset: number;
 };
 
 export class RawTable extends React.Component<RawTableProps, {}> {
@@ -28,7 +30,7 @@ export class RawTable extends React.Component<RawTableProps, {}> {
       <tr key={rowIndex} {...zebraStripe(rowIndex)}>
         {
           [
-            <td key={-1}>{rowIndex + 1}</td>,
+            <td key={-1}>{rowIndex + 1 + this.props.offset}</td>,
             ...row.map((value, columnIndex) =>
               <td key={columnIndex}>
                 {

extensions/ql-vscode/src/view/result-table-utils.tsx

@@ -1,15 +1,17 @@
 import * as React from 'react';
 import { LocationValue, ResolvableLocationValue, tryGetResolvableLocation } from 'semmle-bqrs';
 import { RawResultsSortState, QueryMetadata, SortDirection } from '../interface-types';
-import { ResultSet, vscode } from './results';
+import { vscode } from './results';
 import { assertNever } from '../helpers-pure';
+import { ResultSet } from '../interface-utils';
 
 export interface ResultTableProps {
   resultSet: ResultSet;
   databaseUri: string;
   metadata?: QueryMetadata;
   resultsPath: string | undefined;
   sortState?: RawResultsSortState;
+  offset: number;
 
   /**
    * Holds if there are any raw results. When that is the case, we