Merge pull request #1683 from github/koesie10/fix-duplicate-downloads

Fix duplicate variant analysis results downloads

extensions/ql-vscode/src/remote-queries/variant-analysis-manager.ts

@@ -26,6 +26,8 @@ import { createTimestampFile, showAndLogErrorMessage } from '../helpers';
 import * as fs from 'fs-extra';
 
 export class VariantAnalysisManager extends DisposableObject implements VariantAnalysisViewManager<VariantAnalysisView> {
+  private static readonly REPO_STATES_FILENAME = 'repo_states.json';
+
   private readonly _onVariantAnalysisAdded = this.push(new EventEmitter<VariantAnalysis>());
   public readonly onVariantAnalysisAdded = this._onVariantAnalysisAdded.event;
   private readonly _onVariantAnalysisStatusUpdated = this.push(new EventEmitter<VariantAnalysis>());
@@ -40,6 +42,8 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
   private static readonly maxConcurrentDownloads = 3;
   private readonly queue = new PQueue({ concurrency: VariantAnalysisManager.maxConcurrentDownloads });
 
+  private readonly repoStates = new Map<number, Record<number, VariantAnalysisScannedRepositoryState>>();
+
   constructor(
     private readonly ctx: ExtensionContext,
     private readonly storagePath: string,
@@ -60,6 +64,15 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
       this._onVariantAnalysisRemoved.fire(variantAnalysis);
     } else {
       await this.setVariantAnalysis(variantAnalysis);
+
+      try {
+        const repoStates = await fs.readJson(this.getRepoStatesStoragePath(variantAnalysis.id));
+        this.repoStates.set(variantAnalysis.id, repoStates);
+      } catch (e) {
+        // Ignore this error, we simply might not have downloaded anything yet
+        this.repoStates.set(variantAnalysis.id, {});
+      }
+
       if (!await isVariantAnalysisComplete(variantAnalysis, this.makeResultDownloadChecker(variantAnalysis))) {
         void commands.executeCommand('codeQL.monitorVariantAnalysis', variantAnalysis);
       }
@@ -120,6 +133,10 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
     return this.variantAnalyses.get(variantAnalysisId);
   }
 
+  public async getRepoStates(variantAnalysisId: number): Promise<VariantAnalysisScannedRepositoryState[]> {
+    return Object.values(this.repoStates.get(variantAnalysisId) ?? {});
+  }
+
   public get variantAnalysesSize(): number {
     return this.variantAnalyses.size;
   }
@@ -152,6 +169,8 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
 
     await this.prepareStorageDirectory(variantAnalysis.id);
 
+    this.repoStates.set(variantAnalysis.id, {});
+
     this._onVariantAnalysisAdded.fire(variantAnalysis);
   }
 
@@ -166,6 +185,14 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
 
   private async onRepoStateUpdated(variantAnalysisId: number, repoState: VariantAnalysisScannedRepositoryState): Promise<void> {
     await this.getView(variantAnalysisId)?.updateRepoState(repoState);
+
+    let repoStates = this.repoStates.get(variantAnalysisId);
+    if (!repoStates) {
+      repoStates = {};
+      this.repoStates.set(variantAnalysisId, repoStates);
+    }
+
+    repoStates[repoState.repositoryId] = repoState;
   }
 
   public async monitorVariantAnalysis(
@@ -180,6 +207,10 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
     variantAnalysis: VariantAnalysis,
     cancellationToken: CancellationToken
   ): Promise<void> {
+    if (this.repoStates.get(variantAnalysis.id)?.[scannedRepo.repository.id]?.downloadStatus === VariantAnalysisScannedRepositoryDownloadStatus.Succeeded) {
+      return;
+    }
+
     const repoState = {
       repositoryId: scannedRepo.repository.id,
       downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Pending,
@@ -216,11 +247,19 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
       repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.InProgress;
       await this.onRepoStateUpdated(variantAnalysis.id, repoState);
 
+      try {
         await this.variantAnalysisResultsManager.download(credentials, variantAnalysis.id, repoTask, this.getVariantAnalysisStorageLocation(variantAnalysis.id));
+      } catch (e) {
+        repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.Failed;
+        await this.onRepoStateUpdated(variantAnalysis.id, repoState);
+        throw new Error(`Could not download the results for variant analysis with id: ${variantAnalysis.id}. Error: ${getErrorMessage(e)}`);
+      }
     }
 
     repoState.downloadStatus = VariantAnalysisScannedRepositoryDownloadStatus.Succeeded;
     await this.onRepoStateUpdated(variantAnalysis.id, repoState);
+
+    await fs.outputJson(this.getRepoStatesStoragePath(variantAnalysis.id), this.repoStates.get(variantAnalysis.id));
   }
 
   public async enqueueDownload(
@@ -242,6 +281,13 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
     );
   }
 
+  private getRepoStatesStoragePath(variantAnalysisId: number): string {
+    return path.join(
+      this.getVariantAnalysisStorageLocation(variantAnalysisId),
+      VariantAnalysisManager.REPO_STATES_FILENAME
+    );
+  }
+
   /**
    * Prepares a directory for storing results for a variant analysis.
    * This directory contains a timestamp file, which will be

extensions/ql-vscode/src/remote-queries/variant-analysis-view-manager.ts

@@ -1,4 +1,4 @@
-import { VariantAnalysis } from './shared/variant-analysis';
+import { VariantAnalysis, VariantAnalysisScannedRepositoryState } from './shared/variant-analysis';
 
 export interface VariantAnalysisViewInterface {
   variantAnalysisId: number;
@@ -10,4 +10,5 @@ export interface VariantAnalysisViewManager<T extends VariantAnalysisViewInterfa
   unregisterView(view: T): void;
 
   getVariantAnalysis(variantAnalysisId: number): Promise<VariantAnalysis | undefined>;
+  getRepoStates(variantAnalysisId: number): Promise<VariantAnalysisScannedRepositoryState[]>;
 }

extensions/ql-vscode/src/remote-queries/variant-analysis-view.ts

@@ -127,6 +127,16 @@ export class VariantAnalysisView extends AbstractWebview<ToVariantAnalysisMessag
       t: 'setVariantAnalysis',
       variantAnalysis,
     });
+
+    const repoStates = await this.manager.getRepoStates(this.variantAnalysisId);
+    if (repoStates.length === 0) {
+      return;
+    }
+
+    await this.postMessage({
+      t: 'setRepoStates',
+      repoStates,
+    });
   }
 
   private async openQueryFile(): Promise<void> {

extensions/ql-vscode/src/view/variant-analysis/RepoRow.tsx

@@ -119,7 +119,7 @@ export const RepoRow = ({
     }
   }, [resultsLoaded, resultsLoading]);
 
-  const disabled = !status || !isCompletedAnalysisRepoStatus(status);
+  const disabled = !status || !isCompletedAnalysisRepoStatus(status) || (status === VariantAnalysisRepoStatus.Succeeded && downloadStatus !== VariantAnalysisScannedRepositoryDownloadStatus.Succeeded);
   const expandableContentLoaded = status && (status !== VariantAnalysisRepoStatus.Succeeded || resultsLoaded);
 
   return (

extensions/ql-vscode/src/view/variant-analysis/__tests__/RepoRow.spec.tsx

@@ -1,6 +1,9 @@
 import * as React from 'react';
 import { render as reactRender, screen } from '@testing-library/react';
-import { VariantAnalysisRepoStatus } from '../../../remote-queries/shared/variant-analysis';
+import {
+  VariantAnalysisRepoStatus,
+  VariantAnalysisScannedRepositoryDownloadStatus
+} from '../../../remote-queries/shared/variant-analysis';
 import userEvent from '@testing-library/user-event';
 import { RepoRow, RepoRowProps } from '../RepoRow';
 import { createMockRepositoryWithMetadata } from '../../../vscode-tests/factories/remote-queries/shared/repository';
@@ -50,7 +53,7 @@ describe(RepoRow.name, () => {
     })).toBeDisabled();
   });
 
-  it('renders the succeeded state', () => {
+  it('renders the succeeded state without download status', () => {
     render({
       status: VariantAnalysisRepoStatus.Succeeded,
       resultCount: 178,
@@ -60,6 +63,42 @@ describe(RepoRow.name, () => {
       name: 'Success',
     })).toBeInTheDocument();
     expect(screen.getByText('178')).toBeInTheDocument();
+    expect(screen.getByRole<HTMLButtonElement>('button', {
+      expanded: false
+    })).toBeDisabled();
+  });
+
+  it('renders the succeeded state with pending download status', () => {
+    render({
+      status: VariantAnalysisRepoStatus.Succeeded,
+      resultCount: 178,
+      downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Pending,
+    });
+
+    expect(screen.getByRole<HTMLButtonElement>('button', {
+      expanded: false
+    })).toBeDisabled();
+  });
+
+  it('renders the succeeded state with in progress download status', () => {
+    render({
+      status: VariantAnalysisRepoStatus.Succeeded,
+      resultCount: 178,
+      downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+    });
+
+    expect(screen.getByRole<HTMLButtonElement>('button', {
+      expanded: false
+    })).toBeDisabled();
+  });
+
+  it('renders the succeeded state with succeeded download status', () => {
+    render({
+      status: VariantAnalysisRepoStatus.Succeeded,
+      resultCount: 178,
+      downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+    });
+
     expect(screen.getByRole<HTMLButtonElement>('button', {
       expanded: false
     })).toBeEnabled();
@@ -217,6 +256,7 @@ describe(RepoRow.name, () => {
   it('can expand the repo item when succeeded and loaded', async () => {
     render({
       status: VariantAnalysisRepoStatus.Succeeded,
+      downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
       interpretedResults: [],
     });
 
@@ -232,6 +272,7 @@ describe(RepoRow.name, () => {
   it('can expand the repo item when succeeded and not loaded', async () => {
     const { rerender } = render({
       status: VariantAnalysisRepoStatus.Succeeded,
+      downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
     });
 
     await userEvent.click(screen.getByRole('button', {

extensions/ql-vscode/src/view/variant-analysis/__tests__/VariantAnalysisAnalyzedRepos.spec.tsx

@@ -3,6 +3,7 @@ import { render as reactRender, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import {
   VariantAnalysisRepoStatus,
+  VariantAnalysisScannedRepositoryDownloadStatus,
   VariantAnalysisStatus
 } from '../../../remote-queries/shared/variant-analysis';
 import { VariantAnalysisAnalyzedRepos, VariantAnalysisAnalyzedReposProps } from '../VariantAnalysisAnalyzedRepos';
@@ -110,6 +111,12 @@ describe(VariantAnalysisAnalyzedRepos.name, () => {
 
   it('renders the interpreted result for a succeeded repo', async () => {
     render({
+      repositoryStates: [
+        {
+          repositoryId: 2,
+          downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+        }
+      ],
       repositoryResults: [
         {
           variantAnalysisId: 1,

extensions/ql-vscode/src/vscode-tests/cli-integration/remote-queries/variant-analysis-manager.test.ts

@@ -19,6 +19,7 @@ import { createMockScannedRepos } from '../../factories/remote-queries/shared/sc
 import {
   VariantAnalysis,
   VariantAnalysisScannedRepository,
+  VariantAnalysisScannedRepositoryDownloadStatus,
   VariantAnalysisStatus,
 } from '../../../remote-queries/shared/variant-analysis';
 import { createTimestampFile } from '../../../helpers';
@@ -27,6 +28,9 @@ import { VariantAnalysisRepoTask } from '../../../remote-queries/gh-api/variant-
 
 describe('Variant Analysis Manager', async function() {
   let sandbox: sinon.SinonSandbox;
+  let pathExistsStub: sinon.SinonStub;
+  let readJsonStub: sinon.SinonStub;
+  let outputJsonStub: sinon.SinonStub;
   let cli: CodeQLCliServer;
   let cancellationTokenSource: CancellationTokenSource;
   let variantAnalysisManager: VariantAnalysisManager;
@@ -42,6 +46,9 @@ describe('Variant Analysis Manager', async function() {
     sandbox.stub(config, 'isVariantAnalysisLiveResultsEnabled').returns(false);
     sandbox.stub(fs, 'mkdirSync');
     sandbox.stub(fs, 'writeFile');
+    pathExistsStub = sandbox.stub(fs, 'pathExists').callThrough();
+    readJsonStub = sandbox.stub(fs, 'readJson').callThrough();
+    outputJsonStub = sandbox.stub(fs, 'outputJson');
 
     cancellationTokenSource = new CancellationTokenSource();
 
@@ -65,6 +72,74 @@ describe('Variant Analysis Manager', async function() {
     sandbox.restore();
   });
 
+  describe('rehydrateVariantAnalysis', () => {
+    const variantAnalysis = createMockVariantAnalysis({});
+
+    describe('when the directory does not exist', () => {
+      beforeEach(() => {
+        pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(false);
+      });
+
+      it('should fire the removed event if the file does not exist', async () => {
+        const stub = sandbox.stub();
+        variantAnalysisManager.onVariantAnalysisRemoved(stub);
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis);
+
+        expect(stub).to.have.been.calledOnce;
+        sinon.assert.calledWith(pathExistsStub, path.join(storagePath, variantAnalysis.id.toString()));
+      });
+    });
+
+    describe('when the directory exists', () => {
+      beforeEach(() => {
+        pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(true);
+      });
+
+      it('should store the variant analysis', async () => {
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis);
+
+        expect(await variantAnalysisManager.getVariantAnalysis(variantAnalysis.id)).to.deep.equal(variantAnalysis);
+      });
+
+      it('should not error if the repo states file does not exist', async () => {
+        readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).rejects(new Error('File does not exist'));
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis);
+
+        sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+        expect(await variantAnalysisManager.getRepoStates(variantAnalysis.id)).to.deep.equal([]);
+      });
+
+      it('should read in the repo states if it exists', async () => {
+        readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).resolves({
+          [scannedRepos[0].repository.id]: {
+            repositoryId: scannedRepos[0].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+          },
+          [scannedRepos[1].repository.id]: {
+            repositoryId: scannedRepos[1].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+          },
+        });
+
+        await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis);
+
+        sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+        expect(await variantAnalysisManager.getRepoStates(variantAnalysis.id)).to.have.same.deep.members([
+          {
+            repositoryId: scannedRepos[0].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+          },
+          {
+            repositoryId: scannedRepos[1].repository.id,
+            downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+          },
+        ]);
+      });
+    });
+  });
+
   describe('when credentials are invalid', async () => {
     beforeEach(async () => { sandbox.stub(Credentials, 'initialize').resolves(undefined); });
 
@@ -159,6 +234,168 @@ describe('Variant Analysis Manager', async function() {
 
           expect(getVariantAnalysisRepoResultStub.calledOnce).to.be.true;
         });
+
+        it('should skip the download if the repository has already been downloaded', async () => {
+          // First, do a download so it is downloaded. This avoids having to mock the repo states.
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          getVariantAnalysisRepoStub.resetHistory();
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          expect(getVariantAnalysisRepoStub.notCalled).to.be.true;
+        });
+
+        it('should write the repo state when the download is successful', async () => {
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should not write the repo state when the download fails', async () => {
+          getVariantAnalysisRepoResultStub.rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e: any) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+        });
+
+        it('should have a failed repo state when the repo task API fails', async () => {
+          getVariantAnalysisRepoStub.onFirstCall().rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[1],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Failed,
+            },
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should have a failed repo state when the download fails', async () => {
+          getVariantAnalysisRepoResultStub.onFirstCall().rejects(new Error('Failed to download'));
+
+          try {
+            await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+              scannedRepos[0],
+              variantAnalysis,
+              cancellationTokenSource.token
+            );
+            fail('Expected an error to be thrown');
+          } catch (e) {
+            // we can ignore this error, we expect this
+          }
+
+          sinon.assert.notCalled(outputJsonStub);
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[1],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Failed,
+            },
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+          });
+        });
+
+        it('should update the repo state correctly', async () => {
+          // To set some initial repo states, we need to mock the correct methods so that the repo states are read in.
+          // The actual tests for these are in rehydrateVariantAnalysis, so we can just mock them here and test that
+          // the methods are called.
+
+          pathExistsStub.withArgs(path.join(storagePath, variantAnalysis.id.toString())).resolves(true);
+          // This will read in the correct repo states
+          readJsonStub.withArgs(path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json')).resolves({
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+            [scannedRepos[2].repository.id]: {
+              repositoryId: scannedRepos[2].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+            },
+          });
+
+          await variantAnalysisManager.rehydrateVariantAnalysis(variantAnalysis);
+          sinon.assert.calledWith(readJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'));
+
+          await variantAnalysisManager.autoDownloadVariantAnalysisResult(
+            scannedRepos[0],
+            variantAnalysis,
+            cancellationTokenSource.token
+          );
+
+          sinon.assert.calledWith(outputJsonStub, path.join(storagePath, variantAnalysis.id.toString(), 'repo_states.json'), {
+            [scannedRepos[1].repository.id]: {
+              repositoryId: scannedRepos[1].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            },
+            [scannedRepos[2].repository.id]: {
+              repositoryId: scannedRepos[2].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.InProgress,
+            },
+            [scannedRepos[0].repository.id]: {
+              repositoryId: scannedRepos[0].repository.id,
+              downloadStatus: VariantAnalysisScannedRepositoryDownloadStatus.Succeeded,
+            }
+          });
+        });
       });
 
       describe('enqueueDownload', async () => {