Add option to use notebook json contents instead of file string when …

…creating a notebook. (microsoft#18972)
GazimbeIsaac · Apr 8, 2022 · c2cc32a · c2cc32a
1 parent 828c676
commit c2cc32a
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 19 deletions.
diff --git a/src/sql/workbench/api/common/sqlExtHost.protocol.ts b/src/sql/workbench/api/common/sqlExtHost.protocol.ts
@@ -1009,7 +1009,7 @@ export interface INotebookShowOptions {
 	providerId?: string;
 	connectionProfile?: azdata.IConnectionProfile;
 	defaultKernel?: azdata.nb.IKernelSpec;
-	initialContent?: string;
+	initialContent?: string | azdata.nb.INotebookContents;
 	initialDirtyState?: boolean;
 }
 

diff --git a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts
@@ -222,6 +222,7 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 	private _providers: string[];
 	private _standardKernels: IStandardKernelWithProvider[];
 	private _connectionProfile: IConnectionProfile;
+	private _notebookContents: azdata.nb.INotebookContents;
 	private _defaultKernel: azdata.nb.IKernelSpec;
 	public hasBootstrapped = false;
 	// Holds the HTML content for the editor when the editor discards this input and loads another
@@ -283,7 +284,7 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 	public get contentLoader(): IContentLoader {
 		if (!this._contentLoader) {
 			let contentManager = this.instantiationService.createInstance(LocalContentManager);
-			this._contentLoader = this.instantiationService.createInstance(NotebookEditorContentLoader, this, contentManager);
+			this._contentLoader = this.instantiationService.createInstance(NotebookEditorContentLoader, this, contentManager, this._notebookContents);
 		}
 		return this._contentLoader;
 	}
@@ -319,6 +320,11 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 		return this._connectionProfile;
 	}
 
+	public setNotebookContents(value: azdata.nb.INotebookContents) {
+		this._notebookContents = value;
+		(this.contentLoader as NotebookEditorContentLoader).notebookContents = value;
+	}
+
 	public get standardKernels(): IStandardKernelWithProvider[] {
 		return this._standardKernels;
 	}
@@ -461,7 +467,7 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 				this._standardKernels.push(...standardKernels);
 			}
 			let serializationProvider = await this.notebookService.getOrCreateSerializationManager(this._providerId, this._resource);
-			this._contentLoader = this.instantiationService.createInstance(NotebookEditorContentLoader, this, serializationProvider.contentManager);
+			this._contentLoader = this.instantiationService.createInstance(NotebookEditorContentLoader, this, serializationProvider.contentManager, this._notebookContents);
 		}
 	}
 
@@ -541,12 +547,18 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 export class NotebookEditorContentLoader implements IContentLoader {
 	constructor(
 		private notebookInput: NotebookInput,
-		private contentManager: azdata.nb.ContentManager) {
+		private contentManager: azdata.nb.ContentManager,
+		public notebookContents: azdata.nb.INotebookContents | undefined) {
 	}
 
 	async loadContent(): Promise<azdata.nb.INotebookContents> {
-		let notebookEditorModel = await this.notebookInput.resolve();
-		let notebookContents = await this.contentManager.deserializeNotebook(notebookEditorModel.contentString);
+		let notebookContents: azdata.nb.INotebookContents;
+		if (this.notebookContents) {
+			notebookContents = this.notebookContents;
+		} else {
+			let notebookEditorModel = await this.notebookInput.resolve();
+			notebookContents = await this.contentManager.deserializeNotebook(notebookEditorModel.contentString);
+		}
 
 		// Special case .NET Interactive kernel spec to handle inconsistencies between notebook providers and jupyter kernel specs
 		if (notebookContents.metadata?.kernelspec?.display_name?.startsWith(DotnetInteractiveJupyterLabelPrefix)) {

diff --git a/src/sql/workbench/services/notebook/browser/interface.ts b/src/sql/workbench/services/notebook/browser/interface.ts
@@ -10,8 +10,9 @@ import { IStandardKernelWithProvider } from 'sql/workbench/services/notebook/bro
 import { IEditorInput } from 'vs/workbench/common/editor';
 
 export interface INotebookInput extends IEditorInput {
-	defaultKernel?: azdata.nb.IKernelSpec,
-	connectionProfile?: azdata.IConnectionProfile,
+	defaultKernel?: azdata.nb.IKernelSpec;
+	connectionProfile?: azdata.IConnectionProfile;
+	setNotebookContents(contents: azdata.nb.INotebookContents): void;
 	isDirty(): boolean;
 	setDirty(boolean);
 	readonly notebookUri: URI;

diff --git a/src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts b/src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts
@@ -185,6 +185,7 @@ export class NotebookService extends Disposable implements INotebookService {
 	private _trustedCacheQueue: URI[] = [];
 	private _unTrustedCacheQueue: URI[] = [];
 	private _onCodeCellExecutionStart: Emitter<void> = new Emitter<void>();
+	private _notebookInputsMap: Map<string, IEditorInput> = new Map();
 
 	constructor(
 		@ILifecycleService lifecycleService: ILifecycleService,
@@ -255,7 +256,7 @@ export class NotebookService extends Disposable implements INotebookService {
 		do {
 			uri = URI.from({ scheme: Schemas.untitled, path: `Notebook-${counter}` });
 			counter++;
-		} while (this._untitledEditorService.get(uri));
+		} while (this._untitledEditorService.get(uri) || this._notebookInputsMap.has(uri.toString())); // Also have to check stored inputs, since those might not be opened in an editor yet.
 		return uri;
 	}
 
@@ -268,16 +269,9 @@ export class NotebookService extends Disposable implements INotebookService {
 			resource = uri;
 		}
 
-		let serializedContent: string;
-		if (contents) {
-			// Have to serialize contents again first, since our notebook code assumes input is based on the raw file contents
-			let manager = await this.getOrCreateSerializationManager(providerId, uri);
-			serializedContent = await manager.contentManager.serializeNotebook(contents);
-		}
-
 		let options: INotebookShowOptions = {
 			providerId: providerId,
-			initialContent: serializedContent
+			initialContent: contents
 		};
 		return this.createNotebookInput(options, resource);
 	}
@@ -286,19 +280,31 @@ export class NotebookService extends Disposable implements INotebookService {
 		let uri: URI;
 		if (resource) {
 			uri = URI.revive(resource);
+			if (this._notebookInputsMap.has(uri.toString())) {
+				return this._notebookInputsMap.get(uri.toString());
+			}
 		} else {
 			uri = this.getUntitledFileUri();
 		}
 		let isUntitled: boolean = uri.scheme === Schemas.untitled;
 
 		let fileInput: IEditorInput;
 		let languageMode = options.providerId === INTERACTIVE_PROVIDER_ID ? INTERACTIVE_LANGUAGE_MODE : DEFAULT_NB_LANGUAGE_MODE;
+		let initialStringContents: string;
+		if (options.initialContent) {
+			if (typeof options.initialContent === 'string') {
+				initialStringContents = options.initialContent;
+			} else {
+				let manager = await this.getOrCreateSerializationManager(options.providerId, uri);
+				initialStringContents = await manager.contentManager.serializeNotebook(options.initialContent);
+			}
+		}
 		if (isUntitled && path.isAbsolute(uri.fsPath)) {
-			const model = this._untitledEditorService.create({ associatedResource: uri, mode: languageMode, initialValue: options.initialContent });
+			const model = this._untitledEditorService.create({ associatedResource: uri, mode: languageMode, initialValue: initialStringContents });
 			fileInput = this._instantiationService.createInstance(UntitledTextEditorInput, model);
 		} else {
 			if (isUntitled) {
-				const model = this._untitledEditorService.create({ untitledResource: uri, mode: languageMode, initialValue: options.initialContent });
+				const model = this._untitledEditorService.create({ untitledResource: uri, mode: languageMode, initialValue: initialStringContents });
 				fileInput = this._instantiationService.createInstance(UntitledTextEditorInput, model);
 			} else {
 				fileInput = this._editorService.createEditorInput({ forceFile: true, resource: uri, mode: languageMode });
@@ -312,6 +318,9 @@ export class NotebookService extends Disposable implements INotebookService {
 			if (isINotebookInput(fileInput)) {
 				fileInput.defaultKernel = options.defaultKernel;
 				fileInput.connectionProfile = options.connectionProfile;
+				if (typeof options.initialContent !== 'string') {
+					fileInput.setNotebookContents(options.initialContent);
+				}
 
 				if (isUntitled) {
 					let untitledModel = await fileInput.resolve();
@@ -327,6 +336,7 @@ export class NotebookService extends Disposable implements INotebookService {
 			throw new Error(localize('failedToCreateNotebookInput', "Failed to create notebook input for provider '{0}'", options.providerId));
 		}
 
+		this._notebookInputsMap.set(uri.toString(), fileInput);
 		return fileInput;
 	}
 
@@ -616,6 +626,8 @@ export class NotebookService extends Disposable implements INotebookService {
 		if (this._editors.delete(editor.id)) {
 			this._onNotebookEditorRemove.fire(editor);
 		}
+		this._notebookInputsMap.delete(editor.notebookParams.notebookUri.toString());
+
 		// Remove the manager from the tracked list, and let the notebook provider know that it should update its mappings
 		this.sendNotebookCloseToProvider(editor);
 	}