Retrieve external API usage snippets using SARIF

Author: koesie10

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

@@ -1050,13 +1050,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,129 @@
+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: Pick<
+    CodeQLCliServer,
+    "resolveDatabase" | "resolveQlpacks" | "interpretBqrsSarif"
+  >;
+  queryRunner: Pick<QueryRunner, "createQueryRun" | "logger">;
+  databaseItem: Pick<
+    DatabaseItem,
+    | "contents"
+    | "databaseUri"
+    | "language"
+    | "sourceArchive"
+    | "getSourceLocationPrefix"
+  >;
+  queryStorageDir: string;
+
+  progress: ProgressCallback;
+};
+
+export async function getAutoModelUsages({
+  cliServer,
+  queryRunner,
+  databaseItem,
+  queryStorageDir,
+  progress,
+}: Options): Promise<Record<string, string[]>> {
+  const maxStep = 1500;
+
+  const cancellationTokenSource = new CancellationTokenSource();
+
+  const queryResult = await runQuery("usagesQuery", {
+    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 locatin prefix",
+  });
+
+  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",
+  });
+
+  const sarif = await interpretResultsSarif(
+    cliServer,
+    {
+      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: Record<string, string[]> = {};
+
+  const results = sarif.runs[0]?.results;
+  if (!results) {
+    throw new Error("No results");
+  }
+
+  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

@@ -11,6 +11,7 @@ export function createAutoModelRequest(
   language: string,
   externalApiUsages: ExternalApiUsage[],
   modeledMethods: Record<string, ModeledMethod>,
+  usages: Record<string, string[]>,
 ): ModelRequest {
   const request: ModelRequest = {
     language,
@@ -29,6 +30,10 @@ export function createAutoModelRequest(
       type: "none",
     };
 
+    const usagesForMethod =
+      usages[externalApiUsage.signature] ??
+      externalApiUsage.usages.map((usage) => usage.label);
+
     const numberOfArguments =
       externalApiUsage.methodParameters === "()"
         ? 0
@@ -48,9 +53,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/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(
@@ -242,7 +243,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
     const cancellationTokenSource = new CancellationTokenSource();
 
     try {
-      const queryResult = await runQuery({
+      const queryResult = await runQuery("mainQuery", {
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
         databaseItem: this.databaseItem,
@@ -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/external-api-usage-query.ts

@@ -16,6 +16,7 @@ import { QueryResultType } from "../pure/new-messages";
 import { join } from "path";
 import { redactableError } from "../pure/errors";
 import { QueryLanguage } from "../common/query-language";
+import { Query } from "./queries/query";
 
 export type RunQueryOptions = {
   cliServer: Pick<CodeQLCliServer, "resolveQlpacks">;
@@ -27,14 +28,17 @@ export type RunQueryOptions = {
   token: CancellationToken;
 };
 
-export async function runQuery({
-  cliServer,
-  queryRunner,
-  databaseItem,
-  queryStorageDir,
-  progress,
-  token,
-}: RunQueryOptions): Promise<CoreCompletedQuery | undefined> {
+export async function runQuery(
+  queryName: keyof Omit<Query, "dependencies">,
+  {
+    cliServer,
+    queryRunner,
+    databaseItem,
+    queryStorageDir,
+    progress,
+    token,
+  }: RunQueryOptions,
+): Promise<CoreCompletedQuery | undefined> {
   // The below code is temporary to allow for rapid prototyping of the queries. Once the queries are stabilized, we will
   // move these queries into the `github/codeql` repository and use them like any other contextual (e.g. AST) queries.
   // This is intentionally not pretty code, as it will be removed soon.
@@ -51,7 +55,7 @@ export async function runQuery({
 
   const queryDir = (await dir({ unsafeCleanup: true })).path;
   const queryFile = join(queryDir, "FetchExternalApis.ql");
-  await writeFile(queryFile, query.mainQuery, "utf8");
+  await writeFile(queryFile, query[queryName], "utf8");
 
   if (query.dependencies) {
     for (const [filename, contents] of Object.entries(query.dependencies)) {

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

@@ -2,33 +2,52 @@ import { Query } from "./query";
 
 export const fetchExternalApisQuery: Query = {
   mainQuery: `/**
- * @name Usage of APIs coming from external libraries
- * @description A list of 3rd party APIs used in the codebase.
- * @tags telemetry
- * @id cs/telemetry/fetch-external-apis
- */
+* @name Usage of APIs coming from external libraries
+* @description A list of 3rd party APIs used in the codebase.
+* @tags telemetry
+* @id cs/telemetry/fetch-external-apis
+*/
+
+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 apiName, supported, usage
+`,
+  usagesQuery: `/**
+* @name Usage of APIs coming from external libraries
+* @description A list of 3rd party APIs used in the codebase.
+* @kind problem
+* @id cs/telemetry/fetch-external-api-usages
+*/
+
+import csharp
+import ExternalApi
+
+private Call aUsage(ExternalApi api) {
+ result.getTarget().getUnboundDeclaration() = api
+}
 
- 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
+from ExternalApi api, string apiName, Call usage
+where
+ apiName = api.getApiName() and
+ usage = aUsage(api)
+select usage, apiName
 `,
   dependencies: {
     "ExternalApi.qll": `/** Provides classes and predicates related to handling APIs from external libraries. */