Merge pull request #3455 from github/koesie10/extension-managed-location

Fix incomplete removal of databases

Author: koesie10

extensions/ql-vscode/src/databases/database-fetcher.ts

@@ -399,6 +399,7 @@ async function databaseArchiveFetcher(
       nameOverride,
       {
         addSourceArchiveFolder,
+        extensionManagedLocation: unzipPath,
       },
     );
     return item;

extensions/ql-vscode/src/databases/local-databases/database-item-impl.ts

@@ -66,6 +66,10 @@ export class DatabaseItemImpl implements DatabaseItem {
     return this.options.origin;
   }
 
+  public get extensionManagedLocation(): string | undefined {
+    return this.options.extensionManagedLocation;
+  }
+
   public resolveSourceFile(uriStr: string | undefined): Uri {
     const sourceArchive = this.sourceArchive;
     const uri = uriStr ? Uri.parse(uriStr, true) : undefined;

extensions/ql-vscode/src/databases/local-databases/database-item.ts

@@ -31,6 +31,12 @@ export interface DatabaseItem {
    */
   readonly origin: DatabaseOrigin | undefined;
 
+  /**
+   * The location of the base storage location as managed by the extension, or undefined
+   * if unknown or not managed by the extension.
+   */
+  readonly extensionManagedLocation: string | undefined;
+
   /** If the database is invalid, describes why. */
   readonly error: Error | undefined;
 

extensions/ql-vscode/src/databases/local-databases/database-manager.ts

@@ -82,6 +82,10 @@ function eventFired<T>(
 }
 
 type OpenDatabaseOptions = {
+  /**
+   * A location that is managed by the extension.
+   */
+  extensionManagedLocation?: string;
   isTutorialDatabase?: boolean;
   /**
    * Whether to add a workspace folder containing the source archive to the workspace. Default is true.
@@ -141,6 +145,7 @@ export class DatabaseManager extends DisposableObject {
     makeSelected = true,
     displayName?: string,
     {
+      extensionManagedLocation,
       isTutorialDatabase = false,
       addSourceArchiveFolder = addDatabaseSourceToWorkspace(),
     }: OpenDatabaseOptions = {},
@@ -149,6 +154,7 @@ export class DatabaseManager extends DisposableObject {
       uri,
       origin,
       displayName,
+      extensionManagedLocation,
     );
 
     return await this.addExistingDatabaseItem(
@@ -202,6 +208,7 @@ export class DatabaseManager extends DisposableObject {
     uri: vscode.Uri,
     origin: DatabaseOrigin | undefined,
     displayName: string | undefined,
+    extensionManagedLocation?: string,
   ): Promise<DatabaseItemImpl> {
     const contents = await DatabaseResolver.resolveDatabaseContents(uri);
     const fullOptions: FullDatabaseOptions = {
@@ -210,6 +217,7 @@ export class DatabaseManager extends DisposableObject {
       dateAdded: Date.now(),
       language: await this.getPrimaryLanguage(uri.fsPath),
       origin,
+      extensionManagedLocation,
     };
     const databaseItem = new DatabaseItemImpl(uri, contents, fullOptions);
 
@@ -370,6 +378,7 @@ export class DatabaseManager extends DisposableObject {
     let dateAdded = undefined;
     let language = undefined;
     let origin = undefined;
+    let extensionManagedLocation = undefined;
     if (state.options) {
       if (typeof state.options.displayName === "string") {
         displayName = state.options.displayName;
@@ -379,6 +388,7 @@ export class DatabaseManager extends DisposableObject {
       }
       language = state.options.language;
       origin = state.options.origin;
+      extensionManagedLocation = state.options.extensionManagedLocation;
     }
 
     const dbBaseUri = vscode.Uri.parse(state.uri, true);
@@ -392,6 +402,7 @@ export class DatabaseManager extends DisposableObject {
       dateAdded,
       language,
       origin,
+      extensionManagedLocation,
     };
     const item = new DatabaseItemImpl(dbBaseUri, undefined, fullOptions);
 
@@ -583,15 +594,20 @@ export class DatabaseManager extends DisposableObject {
     // Remove this database item from the allow-list
     await this.deregisterDatabase(item);
 
+    // Find whether we know directly which directory we should remove
+    const directoryToRemove = item.extensionManagedLocation
+      ? vscode.Uri.file(item.extensionManagedLocation)
+      : item.databaseUri;
+
     // Delete folder from file system only if it is controlled by the extension
-    if (this.isExtensionControlledLocation(item.databaseUri)) {
+    if (this.isExtensionControlledLocation(directoryToRemove)) {
       void extLogger.log("Deleting database from filesystem.");
-      await remove(item.databaseUri.fsPath).then(
-        () => void extLogger.log(`Deleted '${item.databaseUri.fsPath}'`),
+      await remove(directoryToRemove.fsPath).then(
+        () => void extLogger.log(`Deleted '${directoryToRemove.fsPath}'`),
         (e: unknown) =>
           void extLogger.log(
             `Failed to delete '${
-              item.databaseUri.fsPath
+              directoryToRemove.fsPath
             }'. Reason: ${getErrorMessage(e)}`,
           ),
       );

extensions/ql-vscode/src/databases/local-databases/database-options.ts

@@ -5,10 +5,12 @@ export interface DatabaseOptions {
   dateAdded?: number | undefined;
   language?: string;
   origin?: DatabaseOrigin;
+  extensionManagedLocation?: string;
 }
 
 export interface FullDatabaseOptions extends DatabaseOptions {
   dateAdded: number | undefined;
   language: string | undefined;
   origin: DatabaseOrigin | undefined;
+  extensionManagedLocation: string | undefined;
 }

extensions/ql-vscode/test/factories/databases/databases.ts

@@ -14,11 +14,12 @@ export function mockDbOptions(): FullDatabaseOptions {
     origin: {
       type: "folder",
     },
+    extensionManagedLocation: undefined,
   };
 }
 
 export function createMockDB(
-  dir: DirResult,
+  dir: DirResult | string,
   dbOptions = mockDbOptions(),
   // source archive location must be a real(-ish) location since
   // tests will add this to the workspace location
@@ -38,10 +39,18 @@ export function createMockDB(
   );
 }
 
-export function sourceLocationUri(dir: DirResult) {
+export function sourceLocationUri(dir: DirResult | string) {
+  if (typeof dir === "string") {
+    return Uri.file(join(dir, "src.zip"));
+  }
+
   return Uri.file(join(dir.name, "src.zip"));
 }
 
-export function dbLocationUri(dir: DirResult) {
+export function dbLocationUri(dir: DirResult | string) {
+  if (typeof dir === "string") {
+    return Uri.file(join(dir, "db"));
+  }
+
   return Uri.file(join(dir.name, "db"));
 }

extensions/ql-vscode/test/vscode-tests/minimal-workspace/local-queries/local-databases.test.ts

@@ -242,6 +242,37 @@ describe("local databases", () => {
       await expect(pathExists(mockDbItem.databaseUri.fsPath)).resolves.toBe(
         false,
       );
+      await expect(pathExists(dir.name)).resolves.toBe(true);
+    });
+
+    it("should remove a database item with an extension managed location", async () => {
+      const dbLocation = join(dir.name, "org-repo-12");
+      await ensureDir(dbLocation);
+
+      const mockDbItem = createMockDB(dbLocation, {
+        ...mockDbOptions(),
+        extensionManagedLocation: dbLocation,
+      });
+      await ensureDir(mockDbItem.databaseUri.fsPath);
+
+      // pretend that this item is the first workspace folder in the list
+      jest
+        .spyOn(mockDbItem, "belongsToSourceArchiveExplorerUri")
+        .mockReturnValue(true);
+
+      await (databaseManager as any).addDatabaseItem(mockDbItem);
+
+      updateSpy.mockClear();
+
+      await databaseManager.removeDatabaseItem(mockDbItem);
+
+      expect(databaseManager.databaseItems).toEqual([]);
+      expect(updateSpy).toHaveBeenCalledWith("databaseList", []);
+      // should remove the folder
+      expect(workspace.updateWorkspaceFolders).toHaveBeenCalledWith(0, 1);
+
+      // should delete the complete extension managed location
+      await expect(pathExists(dbLocation)).resolves.toBe(false);
     });
 
     it("should remove a database item outside of the extension controlled area", async () => {
@@ -604,6 +635,7 @@ describe("local databases", () => {
         origin: {
           type: "folder",
         },
+        extensionManagedLocation: undefined,
       };
       mockDbItem = createMockDB(dir, options);