Make 'Script to notebook' consistent with 'Deploy' when user cancels during password re-acquisition for controller (#13557)

extensions/arc/notebooks/arcDeployment/deploy.postgres.existing.arc.ipynb

@@ -114,6 +114,8 @@
                 "# Login to the data controller.\n",
                 "#\n",
                 "os.environ[\"AZDATA_PASSWORD\"] = os.environ[\"AZDATA_NB_VAR_CONTROLLER_PASSWORD\"]\n",
+                "os.environ[\"KUBECONFIG\"] = controller_kubeconfig\n",
+                "os.environ[\"KUBECTL_CONTEXT\"] = controller_kubectl_context\n",
                 "cmd = f'azdata login -e {controller_endpoint} -u {controller_username}'\n",
                 "out=run_command()"
             ],

extensions/arc/notebooks/arcDeployment/deploy.sql.existing.arc.ipynb

@@ -114,6 +114,8 @@
                 "# Login to the data controller.\n",
                 "#\n",
                 "os.environ[\"AZDATA_PASSWORD\"] = os.environ[\"AZDATA_NB_VAR_CONTROLLER_PASSWORD\"]\n",
+                "os.environ[\"KUBECONFIG\"] = controller_kubeconfig\n",
+                "os.environ[\"KUBECTL_CONTEXT\"] = controller_kubectl_context\n",
                 "cmd = f'azdata login -e {controller_endpoint} -u {controller_username}'\n",
                 "out=run_command()"
             ],

extensions/arc/src/common/api.ts

@@ -4,11 +4,17 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as arc from 'arc';
+import * as rd from 'resource-deployment';
+import * as loc from '../localizedConstants';
 import { PasswordToControllerDialog } from '../ui/dialogs/connectControllerDialog';
 import { AzureArcTreeDataProvider } from '../ui/tree/azureArcTreeDataProvider';
 import { ControllerTreeNode } from '../ui/tree/controllerTreeNode';
-import { UserCancelledError } from './utils';
 
+export class UserCancelledError extends Error implements rd.ErrorWithType {
+	public get type(): rd.ErrorType {
+		return rd.ErrorType.userCancelled;
+	}
+}
 export function arcApi(treeDataProvider: AzureArcTreeDataProvider): arc.IExtension {
 	return {
 		getRegisteredDataControllers: () => getRegisteredDataControllers(treeDataProvider),
@@ -16,12 +22,13 @@ export function arcApi(treeDataProvider: AzureArcTreeDataProvider): arc.IExtensi
 		reacquireControllerPassword: (controllerInfo: arc.ControllerInfo) => reacquireControllerPassword(treeDataProvider, controllerInfo)
 	};
 }
+
 export async function reacquireControllerPassword(treeDataProvider: AzureArcTreeDataProvider, controllerInfo: arc.ControllerInfo): Promise<string> {
 	const dialog = new PasswordToControllerDialog(treeDataProvider);
 	dialog.showDialog(controllerInfo);
 	const model = await dialog.waitForClose();
 	if (!model) {
-		throw new UserCancelledError();
+		throw new UserCancelledError(loc.userCancelledError);
 	}
 	return model.password;
 }

extensions/arc/src/common/utils.ts

@@ -9,8 +9,6 @@ import * as vscode from 'vscode';
 import { ConnectionMode, IconPath, IconPathHelper } from '../constants';
 import * as loc from '../localizedConstants';
 
-export class UserCancelledError extends Error { }
-
 /**
  * Converts the resource type name into the localized Display Name for that type.
  * @param resourceType The resource type name to convert

extensions/arc/src/localizedConstants.ts

@@ -205,3 +205,4 @@ export const noPasswordFound = (controllerName: string) => localize('noPasswordF
 export const noContextFound = (configFile: string) => localize('noContextFound', "No 'contexts' found in the config file: {0}", configFile);
 export const noCurrentContextFound = (configFile: string) => localize('noCurrentContextFound', "No context is marked as 'current-context' in the config file: {0}", configFile);
 export const noNameInContext = (configFile: string) => localize('noNameInContext', "No name field was found in a cluster context in the config file: {0}", configFile);
+export const userCancelledError = localize('userCancelledError', "User cancelled the dialog");

extensions/arc/src/models/controllerModel.ts

@@ -6,7 +6,7 @@
 import { ControllerInfo, ResourceType } from 'arc';
 import * as azdataExt from 'azdata-ext';
 import * as vscode from 'vscode';
-import { UserCancelledError } from '../common/utils';
+import { UserCancelledError } from '../common/api';
 import * as loc from '../localizedConstants';
 import { ConnectToControllerDialog } from '../ui/dialogs/connectControllerDialog';
 import { AzureArcTreeDataProvider } from '../ui/tree/azureArcTreeDataProvider';
@@ -71,7 +71,7 @@ export class ControllerModel {
 					await this.treeDataProvider.addOrUpdateController(model.controllerModel, model.password, false);
 					this._password = model.password;
 				} else {
-					throw new UserCancelledError();
+					throw new UserCancelledError(loc.userCancelledError);
 				}
 			}
 		}

extensions/arc/src/models/miaaModel.ts

@@ -7,8 +7,9 @@ import { MiaaResourceInfo } from 'arc';
 import * as azdata from 'azdata';
 import * as azdataExt from 'azdata-ext';
 import * as vscode from 'vscode';
+import { UserCancelledError } from '../common/api';
 import { Deferred } from '../common/promise';
-import { createCredentialId, parseIpAndPort, UserCancelledError } from '../common/utils';
+import { createCredentialId, parseIpAndPort } from '../common/utils';
 import { credentialNamespace } from '../constants';
 import * as loc from '../localizedConstants';
 import { ConnectToSqlDialog } from '../ui/dialogs/connectSqlDialog';

extensions/arc/src/test/models/controllerModel.test.ts

@@ -11,7 +11,8 @@ import * as sinon from 'sinon';
 import * as TypeMoq from 'typemoq';
 import { v4 as uuid } from 'uuid';
 import * as vscode from 'vscode';
-import { UserCancelledError } from '../../common/utils';
+import * as loc from '../../localizedConstants';
+import { UserCancelledError } from '../../common/api';
 import { ControllerModel } from '../../models/controllerModel';
 import { ConnectToControllerDialog } from '../../ui/dialogs/connectControllerDialog';
 import { AzureArcTreeDataProvider } from '../../ui/tree/azureArcTreeDataProvider';
@@ -39,7 +40,7 @@ describe('ControllerModel', function (): void {
 			// Returning an undefined model here indicates that the dialog closed without clicking "Ok" - usually through the user clicking "Cancel"
 			sinon.stub(ConnectToControllerDialog.prototype, 'waitForClose').returns(Promise.resolve(undefined));
 			const model = new ControllerModel(new AzureArcTreeDataProvider(mockExtensionContext.object), { id: uuid(), url: '127.0.0.1', username: 'admin', name: 'arc', rememberPassword: true, resources: [] });
-			await should(model.azdataLogin()).be.rejectedWith(new UserCancelledError());
+			await should(model.azdataLogin()).be.rejectedWith(new UserCancelledError(loc.userCancelledError));
 		});
 
 		it('Reads password from cred store', async function (): Promise<void> {

extensions/arc/src/ui/tree/controllerTreeNode.ts

@@ -5,7 +5,7 @@
 
 import { MiaaResourceInfo, ResourceInfo, ResourceType } from 'arc';
 import * as vscode from 'vscode';
-import { UserCancelledError } from '../../common/utils';
+import { UserCancelledError } from '../../common/api';
 import * as loc from '../../localizedConstants';
 import { ControllerModel, Registration } from '../../models/controllerModel';
 import { MiaaModel } from '../../models/miaaModel';

extensions/resource-deployment/src/typings/resource-deployment.d.ts

@@ -5,6 +5,14 @@
 declare module 'resource-deployment' {
 	import * as azdata from 'azdata';
 
+	export const enum ErrorType {
+		userCancelled,
+	}
+
+	export interface ErrorWithType extends Error {
+		readonly type: ErrorType;
+	}
+
 	export const enum extension {
 		name = 'Microsoft.resource-deployment'
 	}

extensions/resource-deployment/src/ui/notebookWizard/notebookWizardModel.ts

@@ -12,6 +12,7 @@ import { DeploymentType, NotebookWizardDeploymentProvider, NotebookWizardInfo }
 import { IPlatformService } from '../../services/platformService';
 import { NotebookWizardAutoSummaryPage } from './notebookWizardAutoSummaryPage';
 import { NotebookWizardPage } from './notebookWizardPage';
+import { ErrorType, ErrorWithType } from 'resource-deployment';
 
 export class NotebookWizardModel extends ResourceTypeModel {
 	private _inputComponents: InputComponents = {};
@@ -58,16 +59,27 @@ export class NotebookWizardModel extends ResourceTypeModel {
 	}
 
 	/**
-	 * Generates the notebook and returns true on successful generation
+	 * Generates the notebook and returns true if generation was done and so the wizard should be closed.
 	 **/
 	public async onGenerateScript(): Promise<boolean> {
 		const lastPage = this.wizard.lastPage! as NotebookWizardPage;
 		if (lastPage.validatePage()) {
-			const notebook = await this.prepareNotebookAndEnvironment();
-			await this.openNotebook(notebook);
-			return true;
+			let notebook: Notebook | undefined;
+			try {
+				notebook = await this.prepareNotebookAndEnvironment();
+			} catch (e) {
+				const isUserCancelled = e instanceof Error && 'type' in e && (<ErrorWithType>e).type === ErrorType.userCancelled;
+				// user cancellation is a normal scenario, we just bail out of the wizard without actually opening the notebook, so rethrow for any other case
+				if (!isUserCancelled) {
+					throw e;
+				}
+			}
+			if (notebook) { // open the notebook if it was successfully prepared
+				await this.openNotebook(notebook);
+			}
+			return true; // generation done (or cancelled at user request) so close the wizard
 		} else {
-			return false;
+			return false; // validation failed so do not attempt to generate the notebook and do not close the wizard
 		}
 	}
 
@@ -82,7 +94,7 @@ export class NotebookWizardModel extends ResourceTypeModel {
 		return await this.notebookService.openNotebookWithContent(notebookPath, JSON.stringify(notebook, undefined, 4));
 	}
 
-	private async prepareNotebookAndEnvironment() {
+	private async prepareNotebookAndEnvironment(): Promise<Notebook> {
 		await setModelValues(this.inputComponents, this);
 		const env: NodeJS.ProcessEnv = process.env;
 		this.setEnvironmentVariables(env, (varName) => {