Merge pull request #2457 from github/koesie10/auto-model-usages-sarif

Retrieve external API usage snippets using SARIF

Author: koesie10

extensions/ql-vscode/src/codeql-cli/cli.ts

@@ -1073,13 +1073,15 @@ export class CodeQLCliServer implements Disposable {
     resultsPath: string,
     interpretedResultsPath: string,
     sourceInfo?: SourceInfo,
+    args?: string[],
   ): Promise<sarif.Log> {
     const additionalArgs = [
       // TODO: This flag means that we don't group interpreted results
       // by primary location. We may want to revisit whether we call
       // interpretation with and without this flag, or do some
       // grouping client-side.
       "--no-group-results",
+      ...(args ?? []),
     ];
 
     await this.runInterpretCommand(

extensions/ql-vscode/src/data-extensions-editor/auto-model-usages-query.ts

@@ -0,0 +1,136 @@
+import { CancellationTokenSource } from "vscode";
+import { join } from "path";
+import { runQuery } from "./external-api-usage-query";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { QueryRunner } from "../query-server";
+import { DatabaseItem } from "../databases/local-databases";
+import { interpretResultsSarif } from "../query-results";
+import { ProgressCallback } from "../common/vscode/progress";
+
+type Options = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  databaseItem: DatabaseItem;
+  queryStorageDir: string;
+
+  progress: ProgressCallback;
+};
+
+export type UsageSnippetsBySignature = Record<string, string[]>;
+
+export async function getAutoModelUsages({
+  cliServer,
+  queryRunner,
+  databaseItem,
+  queryStorageDir,
+  progress,
+}: Options): Promise<UsageSnippetsBySignature> {
+  const maxStep = 1500;
+
+  const cancellationTokenSource = new CancellationTokenSource();
+
+  // This will re-run the query that was already run when opening the data extensions editor. This
+  // might be unnecessary, but this makes it really easy to get the path to the BQRS file which we
+  // need to interpret the results.
+  const queryResult = await runQuery({
+    cliServer,
+    queryRunner,
+    queryStorageDir,
+    databaseItem,
+    progress: (update) =>
+      progress({
+        maxStep,
+        step: update.step,
+        message: update.message,
+      }),
+    token: cancellationTokenSource.token,
+  });
+  if (!queryResult) {
+    throw new Error("Query failed");
+  }
+
+  progress({
+    maxStep,
+    step: 1100,
+    message: "Retrieving source location prefix",
+  });
+
+  // CodeQL needs to have access to the database to be able to retrieve the
+  // snippets from it. The source location prefix is used to determine the
+  // base path of the database.
+  const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
+    cliServer,
+  );
+  const sourceArchiveUri = databaseItem.sourceArchive;
+  const sourceInfo =
+    sourceArchiveUri === undefined
+      ? undefined
+      : {
+          sourceArchive: sourceArchiveUri.fsPath,
+          sourceLocationPrefix,
+        };
+
+  progress({
+    maxStep,
+    step: 1200,
+    message: "Interpreting results",
+  });
+
+  // Convert the results to SARIF so that Codeql will retrieve the snippets
+  // from the datababe. This means we don't need to do that in the extension
+  // and everything is handled by the CodeQL CLI.
+  const sarif = await interpretResultsSarif(
+    cliServer,
+    {
+      // To interpret the results we need to provide metadata about the query. We could do this using
+      // `resolveMetadata` but that would be an extra call to the CodeQL CLI server and would require
+      // us to know the path to the query on the filesystem. Since we know what the metadata should
+      // look like and the only metadata that the CodeQL CLI requires is an ID and the kind, we can
+      // simply use constants here.
+      kind: "problem",
+      id: "usage",
+    },
+    {
+      resultsPath: queryResult.outputDir.bqrsPath,
+      interpretedResultsPath: join(
+        queryStorageDir,
+        "interpreted-results.sarif",
+      ),
+    },
+    sourceInfo,
+    ["--sarif-add-snippets"],
+  );
+
+  progress({
+    maxStep,
+    step: 1400,
+    message: "Parsing results",
+  });
+
+  const snippets: UsageSnippetsBySignature = {};
+
+  const results = sarif.runs[0]?.results;
+  if (!results) {
+    throw new Error("No results");
+  }
+
+  // This will group the snippets by the method signature.
+  for (const result of results) {
+    const signature = result.message.text;
+
+    const snippet =
+      result.locations?.[0]?.physicalLocation?.contextRegion?.snippet?.text;
+
+    if (!signature || !snippet) {
+      continue;
+    }
+
+    if (!(signature in snippets)) {
+      snippets[signature] = [];
+    }
+
+    snippets[signature].push(snippet);
+  }
+
+  return snippets;
+}

extensions/ql-vscode/src/data-extensions-editor/auto-model.ts

@@ -6,11 +6,13 @@ import {
   Method,
   ModelRequest,
 } from "./auto-model-api";
+import type { UsageSnippetsBySignature } from "./auto-model-usages-query";
 
 export function createAutoModelRequest(
   language: string,
   externalApiUsages: ExternalApiUsage[],
   modeledMethods: Record<string, ModeledMethod>,
+  usages: UsageSnippetsBySignature,
 ): ModelRequest {
   const request: ModelRequest = {
     language,
@@ -29,6 +31,10 @@ export function createAutoModelRequest(
       type: "none",
     };
 
+    const usagesForMethod =
+      usages[externalApiUsage.signature] ??
+      externalApiUsage.usages.map((usage) => usage.label);
+
     const numberOfArguments =
       externalApiUsage.methodParameters === "()"
         ? 0
@@ -48,9 +54,7 @@ export function createAutoModelRequest(
           modeledMethod.type === "none"
             ? undefined
             : toMethodClassification(modeledMethod),
-        usages: externalApiUsage.usages
-          .slice(0, 10)
-          .map((usage) => usage.label),
+        usages: usagesForMethod.slice(0, 10),
         input: `Argument[${argumentIndex}]`,
       };
 

extensions/ql-vscode/src/data-extensions-editor/bqrs.ts

@@ -7,9 +7,9 @@ export function decodeBqrsToExternalApiUsages(
   const methodsByApiName = new Map<string, ExternalApiUsage>();
 
   chunk?.tuples.forEach((tuple) => {
-    const signature = tuple[0] as string;
-    const supported = tuple[1] as boolean;
-    const usage = tuple[2] as Call;
+    const usage = tuple[0] as Call;
+    const signature = tuple[1] as string;
+    const supported = (tuple[2] as string) === "true";
 
     const [packageWithType, methodDeclaration] = signature.split("#");
 

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -45,6 +45,7 @@ import {
   parsePredictedClassifications,
 } from "./auto-model";
 import { showLlmGeneration } from "../config";
+import { getAutoModelUsages } from "./auto-model-usages-query";
 
 function getQlSubmoduleFolder(): WorkspaceFolder | undefined {
   const workspaceFolder = workspace.workspaceFolders?.find(
@@ -385,23 +386,66 @@ export class DataExtensionsEditorView extends AbstractWebview<
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ): Promise<void> {
+    const maxStep = 3000;
+
+    await this.showProgress({
+      step: 0,
+      maxStep,
+      message: "Retrieving usages",
+    });
+
+    const usages = await getAutoModelUsages({
+      cliServer: this.cliServer,
+      queryRunner: this.queryRunner,
+      queryStorageDir: this.queryStorageDir,
+      databaseItem: this.databaseItem,
+      progress: (update) => this.showProgress(update, maxStep),
+    });
+
+    await this.showProgress({
+      step: 1800,
+      maxStep,
+      message: "Creating request",
+    });
+
     const request = createAutoModelRequest(
       this.databaseItem.language,
       externalApiUsages,
       modeledMethods,
+      usages,
     );
 
+    await this.showProgress({
+      step: 2000,
+      maxStep,
+      message: "Sending request",
+    });
+
     const response = await autoModel(this.app.credentials, request);
 
+    await this.showProgress({
+      step: 2500,
+      maxStep,
+      message: "Parsing response",
+    });
+
     const predictedModeledMethods = parsePredictedClassifications(
       response.predicted,
     );
 
+    await this.showProgress({
+      step: 2800,
+      maxStep,
+      message: "Applying results",
+    });
+
     await this.postMessage({
       t: "addModeledMethods",
       modeledMethods: predictedModeledMethods,
       overrideNone: true,
     });
+
+    await this.clearProgress();
   }
 
   /*

extensions/ql-vscode/src/data-extensions-editor/queries/csharp.ts

@@ -5,30 +5,28 @@ export const fetchExternalApisQuery: Query = {
  * @name Usage of APIs coming from external libraries
  * @description A list of 3rd party APIs used in the codebase.
  * @tags telemetry
+ * @kind problem
  * @id cs/telemetry/fetch-external-apis
  */
 
- import csharp
- import semmle.code.csharp.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
- import ExternalApi
- 
- private Call aUsage(ExternalApi api) {
-   result.getTarget().getUnboundDeclaration() = api
- }
- 
- private boolean isSupported(ExternalApi api) {
-   api.isSupported() and result = true
-   or
-   not api.isSupported() and
-   result = false
- }
- 
- from ExternalApi api, string apiName, boolean supported, Call usage
- where
-   apiName = api.getApiName() and
-   supported = isSupported(api) and
-   usage = aUsage(api)
- select apiName, supported, usage
+import csharp
+import ExternalApi
+
+private Call aUsage(ExternalApi api) { result.getTarget().getUnboundDeclaration() = api }
+
+private boolean isSupported(ExternalApi api) {
+  api.isSupported() and result = true
+  or
+  not api.isSupported() and
+  result = false
+}
+
+from ExternalApi api, string apiName, boolean supported, Call usage
+where
+  apiName = api.getApiName() and
+  supported = isSupported(api) and
+  usage = aUsage(api)
+select usage, apiName, supported.toString(), "supported"
 `,
   dependencies: {
     "ExternalApi.qll": `/** Provides classes and predicates related to handling APIs from external libraries. */

extensions/ql-vscode/src/data-extensions-editor/queries/java.ts

@@ -5,11 +5,11 @@ export const fetchExternalApisQuery: Query = {
  * @name Usage of APIs coming from external libraries
  * @description A list of 3rd party APIs used in the codebase. Excludes test and generated code.
  * @tags telemetry
+ * @kind problem
  * @id java/telemetry/fetch-external-apis
  */
 
 import java
-import semmle.code.java.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
 import ExternalApi
 
 private Call aUsage(ExternalApi api) {
@@ -28,7 +28,7 @@ where
   apiName = api.getApiName() and
   supported = isSupported(api) and
   usage = aUsage(api)
-select apiName, supported, usage
+select usage, apiName, supported.toString(), "supported"
 `,
   dependencies: {
     "ExternalApi.qll": `/** Provides classes and predicates related to handling APIs from external libraries. */

extensions/ql-vscode/src/data-extensions-editor/queries/query.ts

@@ -1,4 +1,13 @@
 export type Query = {
+  /**
+   * The main query.
+   *
+   * It should select all usages of external APIs, and return the following result pattern:
+   * - usage: the usage of the external API. This is an entity.
+   * - apiName: the name of the external API. This is a string.
+   * - supported: whether the external API is supported by the extension. This should be a string representation of a boolean to satify the result pattern for a problem query.
+   * - "supported": a string literal. This is required to make the query a valid problem query.
+   */
   mainQuery: string;
   dependencies?: {
     [filename: string]: string;

extensions/ql-vscode/src/pure/bqrs-cli-types.ts

@@ -115,7 +115,7 @@ export type BqrsKind =
   | "Entity";
 
 interface BqrsColumn {
-  name: string;
+  name?: string;
   kind: BqrsKind;
 }
 export interface DecodedBqrsChunk {

extensions/ql-vscode/src/query-results.ts

@@ -139,6 +139,7 @@ export async function interpretResultsSarif(
   metadata: QueryMetadata | undefined,
   resultsPaths: ResultsPaths,
   sourceInfo?: cli.SourceInfo,
+  args?: string[],
 ): Promise<SarifInterpretationData> {
   const { resultsPath, interpretedResultsPath } = resultsPaths;
   let res;
@@ -150,6 +151,7 @@ export async function interpretResultsSarif(
       resultsPath,
       interpretedResultsPath,
       sourceInfo,
+      args,
     );
   }
   return { ...res, t: "SarifInterpretationData" };