Add restore from S3 url

extensions/mssql/package.json

@@ -1714,6 +1714,7 @@
     }
   },
   "dependencies": {
+    "@aws-sdk/client-s3": "^3.490.0",
     "@microsoft/ads-extension-telemetry": "^3.0.1",
     "@microsoft/ads-service-downloader": "^1.2.1",
     "dataprotocol-client": "github:Microsoft/sqlops-dataprotocolclient#2.0.0",

extensions/mssql/src/contracts.ts

@@ -10,6 +10,7 @@ import { ConnectParams } from 'dataprotocol-client/lib/protocol';
 import * as mssql from 'mssql';
 import { DatabaseFileData } from 'mssql';
 import { BackupResponse } from 'azdata';
+import { CredentialInfo } from 'azdata';
 
 // ------------------------------- < Telemetry Sent Event > ------------------------------------
 
@@ -1725,6 +1726,23 @@ export interface PurgeQueryStoreDataRequestParams {
 	database: string;
 }
 
+export namespace CreateCredentialRequest {
+	export const type = new RequestType<CreateCredentialRequestParams, void, void, void>('objectManagement/createCredentialRequest');
+}
+
+export interface CreateCredentialRequestParams {
+	credentialInfo: CredentialInfo;
+	connectionUri: string;
+}
+
+export namespace GetCredentialNamesRequest {
+	export const type = new RequestType<GetCredentialNamesRequestParams, string[], void, void>('objectManagement/getCredentialNamesRequest');
+}
+
+export interface GetCredentialNamesRequestParams {
+	connectionUri: string;
+}
+
 // ------------------------------- < Object Management > ------------------------------------
 
 // ------------------------------- < Encryption IV/KEY updation Event > ------------------------------------

extensions/mssql/src/mssql.d.ts

@@ -1038,6 +1038,17 @@ declare module 'mssql' {
 		 * @param database The target database.
 		 */
 		purgeQueryStoreData(connectionUri: string, database: string): Thenable<void>;
+		/**
+		 * Create a new credential
+		 * @param connectionUri The URI of the server connection.
+		 * @param credentialInfo
+		 */
+		createCredential(connectionUri: string, credentialInfo: azdata.CredentialInfo): Thenable<void>;
+		/**
+		 * Gets all the credentials that exist in the current server
+		 * @param connectionUri The URI of the server connection.
+		 */
+		getCredentialNames(connectionUri: string): Thenable<string[]>;
 	}
 
 	/**

extensions/mssql/src/objectManagement/localizedConstants.ts

@@ -344,6 +344,7 @@ export const RestorePlanSectionText = localize('objectManagement.restoreDatabase
 export const RestoreFromBackupFileOptionText = localize('objectManagement.restoreDatabase.restoreFromBackupFileOptionText', "Backup file");
 export const RestoreFromDatabaseOptionText = localize('objectManagement.restoreDatabase.restoreFromDatabaseOptionText', "Database");
 export const RestoreFromUrlText = localize('objectManagement.restoreDatabase.restoreFromUrlText', "URL");
+export const RestoreFromS3UrlText = localize('objectManagement.restoreDatabase.restoreFromS3UrlText', "S3 URL");
 export const BackupFolderPathTitle = localize('objectManagement.restoreDatabase.backupFolderPathTitle', "Please enter one or more file paths separated by commas");
 export const RestoreDatabaseFilesAsText = localize('objectManagement.restoreDatabase.restoreDatabaseFilesAsText', "Restore database files as");
 export const RestoreDatabaseFileDetailsText = localize('objectManagement.restoreDatabase.restoreDatabaseFileDetailsText', "Restore database file Details");
@@ -600,6 +601,17 @@ export const NotAvailableText = localize('objectManagement.databaseProperties.no
 export const PurgeQueryStoreDataMessage = (databaseName: string) => localize('objectManagement.databaseProperties.purgeQueryStoreDataMessage', "Are you sure you want to purge the Query Store data from '{0}'?", databaseName);
 export const fileGroupsNameInput = localize('objectManagement.filegroupsNameInput', "Filegroup Name");
 
+// S3 credentials
+export const SelectS3BackupFileDialogTitle = localize('objectManagement.selectS3BackupFileDialogTitle', "Select S3 Storage Backup File");
+export const RegionSpecificEndpointText = localize('objectManagement.regionSpecificEndpointLabel', "Region-specific endpoint");
+export const SecretKeyText = localize('objectManagement.secretKeyLabel', "Secret Key");
+export const AccessKeyText = localize('objectManagement.accessKeyLabel', "Access Key");
+export const RegionText = localize('objectManagement.regionLabel', "Region");
+export const AddCredentialsText = localize('objectManagement.addCredentialsLabel', "Add Credentials");
+export const SelectS3BucketText = localize('objectManagement.SelectS3BucketLabel', "Select S3 Bucket");
+export const SelectBackupFileText = localize('objectManagement.SelectBackupFileLabel', "Select Backup File");
+export const InvalidS3UrlError = localize('objectManagement.InvalidS3UrlError', "Invalid S3 endpoint");
+
 // Util functions
 export function getNodeTypeDisplayName(type: string, inTitle: boolean = false): string {
 	switch (type) {

extensions/mssql/src/objectManagement/objectManagementService.ts

@@ -12,6 +12,7 @@ import { ObjectManagement, IObjectManagementService, DatabaseFileData, BackupInf
 import { ClientCapabilities } from 'vscode-languageclient';
 import { AppContext } from '../appContext';
 import { BackupResponse } from 'azdata';
+import { CredentialInfo } from 'azdata';
 
 export class ObjectManagementService extends BaseService implements IObjectManagementService {
 	public static asFeature(context: AppContext): ISqlOpsFeature {
@@ -106,6 +107,16 @@ export class ObjectManagementService extends BaseService implements IObjectManag
 		const params: contracts.PurgeQueryStoreDataRequestParams = { connectionUri, database };
 		return this.runWithErrorHandling(contracts.PurgeQueryStoreDataRequest.type, params);
 	}
+
+	async createCredential(connectionUri: string, credentialInfo: azdata.CredentialInfo): Promise<void> {
+		const params: contracts.CreateCredentialRequestParams = { connectionUri, credentialInfo };
+		return this.runWithErrorHandling(contracts.CreateCredentialRequest.type, params);
+	}
+
+	async getCredentialNames(connectionUri: string): Promise<string[]> {
+		const params: contracts.GetCredentialNamesRequestParams = { connectionUri };
+		return this.runWithErrorHandling(contracts.GetCredentialNamesRequest.type, params);
+	}
 }
 
 const ServerLevelSecurableTypes: SecurableTypeMetadata[] = [
@@ -305,6 +316,14 @@ export class TestObjectManagementService implements IObjectManagementService {
 		return this.delayAndResolve([]);
 	}
 
+	async createCredential(connectionUri: string, credentialInfo: CredentialInfo): Promise<void> {
+		return this.delayAndResolve();
+	}
+
+	async getCredentialNames(connectionUri: string): Promise<string[]> {
+		return this.delayAndResolve([]);
+	}
+
 	private generateSearchResult(objectType: ObjectManagement.NodeType, schema: string | undefined, count: number): ObjectManagement.SearchResultItem[] {
 		let items: ObjectManagement.SearchResultItem[] = [];
 		for (let i = 0; i < count; i++) {

extensions/mssql/src/objectManagement/ui/S3AddBackupFileDialog.ts

@@ -0,0 +1,177 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as vscode from 'vscode';
+import * as localizedConstants from '../localizedConstants';
+import { DefaultInputWidth, DialogBase } from '../../ui/dialogBase';
+import * as awsClient from '@aws-sdk/client-s3';
+import { IObjectManagementService } from 'mssql';
+
+export interface S3AddBackupFileDialogResult {
+	s3Url: vscode.Uri;
+	secretKey: string;
+	accessKey: string;
+	backupFilePath: string;
+}
+
+export class S3AddBackupFileDialog extends DialogBase<S3AddBackupFileDialogResult> {
+	private s3UrlInputBox: azdata.InputBoxComponent;
+	private secretKeyInputBox: azdata.InputBoxComponent;
+	private accessKeyInputBox: azdata.InputBoxComponent;
+	private credentialButton: azdata.ButtonComponent;
+	private regionInputBox: azdata.InputBoxComponent;
+	private bucketDropdown: azdata.DropDownComponent;
+	private backupFilesDropdown: azdata.DropDownComponent;
+	private result: S3AddBackupFileDialogResult;
+	private objectManagementService: IObjectManagementService;
+	private credentialInfo: azdata.CredentialInfo;
+	private connectionUri: string;
+	private s3Client: awsClient.S3Client;
+
+	constructor(objectManagementService: IObjectManagementService, connectionUri: string) {
+		super(localizedConstants.SelectS3BackupFileDialogTitle, localizedConstants.SelectS3BackupFileDialogTitle);
+		this.result = {
+			s3Url: undefined,
+			secretKey: undefined,
+			accessKey: undefined,
+			backupFilePath: undefined
+		};
+
+		// Relabel Cancel button to Back, since clicking cancel on an inner dialog makes it seem like it would close the whole dialog overall
+		this.dialogObject.cancelButton.label = localizedConstants.BackButtonLabel;
+		this.dialogObject.okButton.label = localizedConstants.AddButton;
+		this.dialogObject.okButton.enabled = false;
+
+		this.objectManagementService = objectManagementService;
+		this.connectionUri = connectionUri;
+
+		this.dialogObject.okButton.onClick(async () => {
+			this.result.backupFilePath = `s3://${this.bucketDropdown.value}.s3.${this.regionInputBox.value}.amazonaws.com/${this.backupFilesDropdown.value}`;
+			this.credentialInfo = {
+				secret: `${this.result.accessKey}:${this.result.secretKey}`,
+				identity: 'S3 Access Key',
+				name: this.result.backupFilePath,
+				createDate: undefined,
+				dateLastModified: undefined,
+				providerName: 'MSSQL',
+				id: undefined
+			}
+			await this.objectManagementService.createCredential(this.connectionUri, this.credentialInfo);
+		});
+	}
+
+	protected async initialize(): Promise<void> {
+		this.s3UrlInputBox = this.createInputBox(async (value) => {
+			this.result.s3Url = vscode.Uri.parse(value);
+			this.regionInputBox.value = this.result.s3Url.toString().split(".")[1];
+			this.enableCredentialButton();
+		}, {
+			ariaLabel: localizedConstants.RegionSpecificEndpointText,
+			inputType: 'text',
+			placeHolder: 'https://s3.{region}.amazonaws.com'
+		});
+		const s3UrlContainer = this.createLabelInputContainer(localizedConstants.RegionSpecificEndpointText, this.s3UrlInputBox, true);
+
+		this.secretKeyInputBox = this.createInputBox(async (value) => {
+			this.result.secretKey = value;
+			this.enableCredentialButton();
+		}, {
+			ariaLabel: localizedConstants.SecretKeyText,
+			inputType: 'password',
+		});
+		const secretKeyContainer = this.createLabelInputContainer(localizedConstants.SecretKeyText, this.secretKeyInputBox, true);
+
+		this.accessKeyInputBox = this.createInputBox(async (value) => {
+			this.result.accessKey = value;
+			this.enableCredentialButton();
+		}, {
+			ariaLabel: localizedConstants.AccessKeyText,
+			inputType: 'password',
+		});
+		const accessKeyContainer = this.createLabelInputContainer(localizedConstants.AccessKeyText, this.accessKeyInputBox, true);
+
+		this.credentialButton = this.createButton(localizedConstants.AddCredentialsText, localizedConstants.AddCredentialsText, async () => {
+			this.createS3Client();
+			await this.setBucketDropdown();
+		}, false, DefaultInputWidth);
+		const credentialButtonContainer = this.createLabelInputContainer(' ', this.credentialButton);
+
+		this.regionInputBox = this.createInputBox(async () => { }, {
+			ariaLabel: localizedConstants.RegionText,
+			inputType: 'text',
+			enabled: false
+		});
+		const regionInputContainer = this.createLabelInputContainer(localizedConstants.RegionText, this.regionInputBox, true);
+
+		this.bucketDropdown = this.createDropdown(localizedConstants.SelectS3BucketText, async (newValue) => {
+			await this.setBackupFilesDropdown(newValue);
+		}, [], '', false);
+		const bucketContainer = this.createLabelInputContainer(localizedConstants.SelectS3BucketText, this.bucketDropdown, true);
+
+		this.backupFilesDropdown = this.createDropdown(localizedConstants.SelectBackupFileText, async () => {
+			// enable ok button once we have a backup file to restore
+			this.dialogObject.okButton.enabled = true;
+		}, [], '', false);
+		const backupFilesContainer = this.createLabelInputContainer(localizedConstants.SelectBackupFileText, this.backupFilesDropdown, true);
+
+		this.formContainer.addItems([s3UrlContainer, secretKeyContainer, accessKeyContainer, credentialButtonContainer, regionInputContainer, bucketContainer, backupFilesContainer]);
+	}
+
+	protected override async validateInput(): Promise<string[]> {
+		const errors = await super.validateInput();
+		if (!this.result.s3Url.toString().includes('s3') && (this.result.s3Url.scheme !== 'https' || this.result.s3Url.scheme !== 'https')) {
+			errors.push(localizedConstants.InvalidS3UrlError);
+		}
+		return errors;
+	}
+
+	private createS3Client(): void {
+		this.s3Client = new awsClient.S3Client({
+			forcePathStyle: true,
+			region: this.regionInputBox.value,
+			endpoint: this.result.s3Url.toString(),
+			credentials: {
+				accessKeyId: this.dialogResult.accessKey,
+				secretAccessKey: this.dialogResult.secretKey
+			}
+		});
+	}
+
+	private async getBucketList(): Promise<string[]> {
+		let command = new awsClient.ListBucketsCommand({});
+		let response = await this.s3Client.send(command);
+		return response.Buckets.map(r => r.Name);
+	}
+	/**
+	 * Gets a list of all the backup files in S3 storage bucket
+	 */
+	private async getBackupFiles(bucket: string): Promise<string[]> {
+		const input: awsClient.ListObjectVersionsCommandInput = {
+			Bucket: bucket
+		};
+		let command = new awsClient.ListObjectsV2Command(input);
+		let response = await this.s3Client.send(command);
+		return response.Contents.filter(r => r.Key.endsWith('.bak')).map(r => r.Key);
+	}
+
+	private async setBackupFilesDropdown(bucket: string): Promise<void> {
+		this.backupFilesDropdown.values = await this.getBackupFiles(bucket);
+		this.backupFilesDropdown.enabled = true;
+	}
+
+	private async setBucketDropdown(): Promise<void> {
+		this.bucketDropdown.values = await this.getBucketList();
+		this.bucketDropdown.enabled = true;
+	}
+
+	private enableCredentialButton(): void {
+		this.credentialButton.enabled = (this.result.s3Url && this.result.accessKey && this.result.secretKey) !== undefined;
+	}
+
+	public override get dialogResult(): S3AddBackupFileDialogResult | undefined {
+		return this.result;
+	}
+}