|
| 1 | +import * as fetch from "node-fetch"; |
| 2 | +import * as unzipper from "unzipper"; |
| 3 | +import { ExtensionContext, Uri, ProgressOptions, ProgressLocation, commands, window } from "vscode"; |
| 4 | +import * as fs from "fs-extra"; |
| 5 | +import * as path from "path"; |
| 6 | +import { DatabaseManager } from "./databases"; |
| 7 | +import { ProgressCallback, showAndLogErrorMessage, withProgress } from "./helpers"; |
| 8 | + |
| 9 | +export default async function promptFetchDatabase(dbm: DatabaseManager, ctx: ExtensionContext) { |
| 10 | + try { |
| 11 | + const databaseUrl = await window.showInputBox({ |
| 12 | + prompt: 'Enter URL of zipfile of database to download' |
| 13 | + }); |
| 14 | + |
| 15 | + if (databaseUrl) { |
| 16 | + validateUrl(databaseUrl); |
| 17 | + |
| 18 | + const progressOptions: ProgressOptions = { |
| 19 | + location: ProgressLocation.Notification, |
| 20 | + title: 'Adding database from URL', |
| 21 | + cancellable: false, |
| 22 | + }; |
| 23 | + await withProgress(progressOptions, async progress => await databaseFetcher(databaseUrl, dbm, ctx, progress)); |
| 24 | + commands.executeCommand('codeQLDatabases.focus'); |
| 25 | + } |
| 26 | + } catch (e) { |
| 27 | + showAndLogErrorMessage(e.message); |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +async function databaseFetcher( |
| 32 | + databaseUrl: string, |
| 33 | + databasesManager: DatabaseManager, |
| 34 | + ctx: ExtensionContext, |
| 35 | + progressCallback: ProgressCallback |
| 36 | +): Promise<void> { |
| 37 | + progressCallback({ |
| 38 | + maxStep: 3, |
| 39 | + message: 'Downloading database', |
| 40 | + step: 1 |
| 41 | + }); |
| 42 | + const storagePath = ctx.storagePath || ctx.globalStoragePath; |
| 43 | + if (!storagePath) { |
| 44 | + throw new Error("No storage path specified."); |
| 45 | + } |
| 46 | + const unzipPath = await getStorageFolder(storagePath, databaseUrl); |
| 47 | + |
| 48 | + const response = await fetch.default(databaseUrl); |
| 49 | + const unzipStream = unzipper.Extract({ |
| 50 | + path: unzipPath |
| 51 | + }); |
| 52 | + progressCallback({ |
| 53 | + maxStep: 3, |
| 54 | + message: 'Unzipping database', |
| 55 | + step: 2 |
| 56 | + }); |
| 57 | + await new Promise((resolve, reject) => { |
| 58 | + response.body.on('error', reject); |
| 59 | + unzipStream.on('error', reject); |
| 60 | + unzipStream.on('close', resolve); |
| 61 | + response.body.pipe(unzipStream); |
| 62 | + }); |
| 63 | + progressCallback({ |
| 64 | + maxStep: 3, |
| 65 | + message: 'Opening database', |
| 66 | + step: 3 |
| 67 | + }); |
| 68 | + |
| 69 | + // if there is a single directory inside, then assume that's what we want to import |
| 70 | + const dirs = await fs.readdir(unzipPath); |
| 71 | + const dbPath = dirs?.length === 1 && (await fs.stat(path.join(unzipPath, dirs[0]))).isDirectory |
| 72 | + ? path.join(unzipPath, dirs[0]) |
| 73 | + : unzipPath; |
| 74 | + |
| 75 | + // might need to upgrade before importing... |
| 76 | + const item = await databasesManager.openDatabase(Uri.parse(dbPath)); |
| 77 | + databasesManager.setCurrentDatabaseItem(item); |
| 78 | +} |
| 79 | + |
| 80 | +async function getStorageFolder(storagePath: string, urlStr: string) { |
| 81 | + const url = Uri.parse(urlStr); |
| 82 | + let lastName = path.basename(url.path).substring(0, 255); |
| 83 | + if (lastName.endsWith(".zip")) { |
| 84 | + lastName = lastName.substring(0, lastName.length - 4); |
| 85 | + } |
| 86 | + |
| 87 | + const realpath = await fs.realpath(storagePath); |
| 88 | + let folderName = path.join(realpath, lastName); |
| 89 | + let counter = 0; |
| 90 | + while (await fs.pathExists(folderName)) { |
| 91 | + counter++; |
| 92 | + folderName = path.join(realpath, `${lastName}-${counter}`); |
| 93 | + if (counter > 100) { |
| 94 | + throw new Error("Could not find a unique name for downloaded database."); |
| 95 | + } |
| 96 | + } |
| 97 | + return folderName; |
| 98 | +} |
| 99 | + |
| 100 | + |
| 101 | +function validateUrl(databaseUrl: string) { |
| 102 | + let uri; |
| 103 | + try { |
| 104 | + uri = Uri.parse(databaseUrl, true); |
| 105 | + } catch (e) { |
| 106 | + throw new Error(`Invalid url: ${databaseUrl}`); |
| 107 | + } |
| 108 | + |
| 109 | + if (uri.scheme !== 'https') { |
| 110 | + throw new Error('Must use https for downloading a database.'); |
| 111 | + } |
| 112 | +} |
0 commit comments