From 73f00b63ceea105cd8a2585ec8a6327dbb5d70c2 Mon Sep 17 00:00:00 2001
From: Cory Rivera <corivera@microsoft.com>
Date: Mon, 27 Feb 2023 11:54:35 -0800
Subject: [PATCH] Open books in VS Code notebooks if the setting is enabled.
 (#22031)

---
 extensions/notebook/package.json             |  3 ++
 extensions/notebook/src/book/bookTreeView.ts | 32 +++++++++++++++-----
 extensions/notebook/src/typings/refs.d.ts    |  1 +
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/extensions/notebook/package.json b/extensions/notebook/package.json
index 1dea07d6f2..bf70bc7c16 100644
--- a/extensions/notebook/package.json
+++ b/extensions/notebook/package.json
@@ -19,6 +19,9 @@
       "supported": true
     }
   },
+  "enabledApiProposals": [
+    "notebookEditor"
+  ],
   "contributes": {
     "configuration": {
       "type": "object",
diff --git a/extensions/notebook/src/book/bookTreeView.ts b/extensions/notebook/src/book/bookTreeView.ts
index 1f3729a380..ddf6dc082b 100644
--- a/extensions/notebook/src/book/bookTreeView.ts
+++ b/extensions/notebook/src/book/bookTreeView.ts
@@ -395,13 +395,23 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 		}
 	}
 
+	private get useVSCodeNotebooks(): boolean {
+		let workbench = vscode.workspace.getConfiguration('workbench');
+		return workbench?.get<boolean>('useVSCodeNotebooks') && workbench?.get<boolean>('enablePreviewFeatures');
+	}
+
 	async openNotebook(resource: string): Promise<void> {
 		try {
 			await vscode.commands.executeCommand(constants.BuiltInCommands.SetContext, constants.unsavedBooksContextKey, false);
 			if (this._openAsUntitled) {
 				await this.openNotebookAsUntitled(resource);
 			} else {
-				await azdata.nb.showNotebookDocument(vscode.Uri.file(resource));
+				if (this.useVSCodeNotebooks) {
+					let doc = await vscode.workspace.openNotebookDocument(vscode.Uri.file(resource));
+					await vscode.window.showNotebookDocument(doc);
+				} else {
+					await azdata.nb.showNotebookDocument(vscode.Uri.file(resource));
+				}
 				// let us keep a list of already visited notebooks so that we do not trust them again, potentially
 				// overriding user changes
 				let normalizedResource = path.normalize(resource);
@@ -532,13 +542,19 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 	async openNotebookAsUntitled(resource: string): Promise<void> {
 		try {
 			await vscode.commands.executeCommand(constants.BuiltInCommands.SetContext, constants.unsavedBooksContextKey, true);
-			let untitledFileName: vscode.Uri = this.getUntitledNotebookUri(resource);
-			let document: vscode.TextDocument = await vscode.workspace.openTextDocument(resource);
-			await azdata.nb.showNotebookDocument(untitledFileName, {
-				connectionProfile: null,
-				initialContent: document.getText(),
-				initialDirtyState: false
-			});
+			let untitledFileUri: vscode.Uri = this.getUntitledNotebookUri(resource);
+
+			if (this.useVSCodeNotebooks) {
+				let doc = await vscode.workspace.openNotebookDocument(untitledFileUri);
+				await vscode.window.showNotebookDocument(doc);
+			} else {
+				let document: vscode.TextDocument = await vscode.workspace.openTextDocument(resource);
+				await azdata.nb.showNotebookDocument(untitledFileUri, {
+					connectionProfile: null,
+					initialContent: document.getText(),
+					initialDirtyState: false
+				});
+			}
 		} catch (e) {
 			void vscode.window.showErrorMessage(loc.openUntitledNotebookError(resource, e instanceof Error ? e.message : e));
 		}
diff --git a/extensions/notebook/src/typings/refs.d.ts b/extensions/notebook/src/typings/refs.d.ts
index 420c12b6ad..487f690eef 100644
--- a/extensions/notebook/src/typings/refs.d.ts
+++ b/extensions/notebook/src/typings/refs.d.ts
@@ -6,4 +6,5 @@
 /// <reference path='../../../../src/sql/azdata.d.ts'/>
 /// <reference path='../../../../src/sql/azdata.proposed.d.ts'/>
 /// <reference path='../../../../src/vscode-dts/vscode.d.ts'/>
+/// <reference path='../../../../src/vscode-dts/vscode.proposed.notebookEditor.d.ts'/>
 /// <reference types='@types/node'/>