Notebook Views initialization fix (#17109)

Separate the Views load from the initialization. This way we can load previously created views, and only add the new views data to the document when needed. For now, this happens only when a view is created.

src/sql/workbench/services/notebook/browser/models/notebookExtension.ts

@@ -9,27 +9,36 @@ import { deepClone } from 'vs/base/common/objects';
 
 export class NotebookExtension<TNotebookMeta, TCellMeta> {
 	readonly version = 1;
-	readonly extensionName = 'azuredatastudio';
 	readonly extensionNamespace = 'extensions';
 
-	public getNotebookMetadata(notebook: INotebookModel): TNotebookMeta {
+	private _extensionName: string;
+
+	public constructor(extensionName: string) {
+		this._extensionName = extensionName;
+	}
+
+	public get extensionName(): string {
+		return this._extensionName;
+	}
+
+	public getExtensionMetadata(notebook: INotebookModel): TNotebookMeta {
 		const metadata = notebook.getMetaValue(this.extensionNamespace) || {};
 		return metadata[this.extensionName] as TNotebookMeta;
 	}
 
-	public setNotebookMetadata(notebook: INotebookModel, metadata: TNotebookMeta) {
+	public setExtensionMetadata(notebook: INotebookModel, metadata: TNotebookMeta) {
 		const meta = {};
 		meta[this.extensionName] = metadata;
 		notebook.setMetaValue(this.extensionNamespace, meta);
 		notebook.serializationStateChanged(NotebookChangeType.MetadataChanged);
 	}
 
-	public getCellMetadata(cell: ICellModel): TCellMeta {
+	public getExtensionCellMetadata(cell: ICellModel): TCellMeta {
 		const namespaceMeta = cell.metadata[this.extensionNamespace] || {};
 		return namespaceMeta[this.extensionName] as TCellMeta;
 	}
 
-	public setCellMetadata(cell: ICellModel, metadata: TCellMeta) {
+	public setExtensionCellMetadata(cell: ICellModel, metadata: TCellMeta) {
 		const meta = {};
 		meta[this.extensionName] = metadata;
 		cell.metadata[this.extensionNamespace] = meta;

src/sql/workbench/services/notebook/browser/notebookViews/notebookViewModel.ts

@@ -51,11 +51,11 @@ export class NotebookViewModel implements INotebookView {
 	}
 
 	protected initializeCell(cell: ICellModel, idx: number) {
-		let meta = this._notebookViews.getCellMetadata(cell);
+		let meta = this._notebookViews.getExtensionCellMetadata(cell);
 
 		if (!meta) {
 			this._notebookViews.initializeCell(cell);
-			meta = this._notebookViews.getCellMetadata(cell);
+			meta = this._notebookViews.getExtensionCellMetadata(cell);
 		}
 
 		// Ensure that we are not duplicting view entries in cell metadata
@@ -91,7 +91,7 @@ export class NotebookViewModel implements INotebookView {
 	}
 
 	public getCellMetadata(cell: ICellModel): INotebookViewCell {
-		const meta = this._notebookViews.getCellMetadata(cell);
+		const meta = this._notebookViews.getExtensionCellMetadata(cell);
 		return meta?.views?.find(view => view.guid === this.guid);
 	}
 

src/sql/workbench/services/notebook/browser/notebookViews/notebookViewsExtension.ts

@@ -13,30 +13,40 @@ import { INotebookView, INotebookViewCell, INotebookViewCellMetadata, INotebookV
 
 export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetadata, INotebookViewCellMetadata> implements INotebookViews {
 	static readonly defaultViewName = localize('notebookView.untitledView', "Untitled View");
+	static readonly extension = 'notebookviews';
 
 	readonly maxNameIterationAttempts = 100;
-	readonly extension = 'azuredatastudio';
 	override readonly version = 1;
 
-	protected _metadata: INotebookViewMetadata;
+	protected _metadata: INotebookViewMetadata | undefined;
+	private _initialized: boolean = false;
 	private _onViewDeleted = new Emitter<void>();
 	private _onActiveViewChanged = new Emitter<void>();
 
 	constructor(protected _notebook: INotebookModel) {
-		super();
-		this.loadOrInitialize();
+		super(NotebookViewsExtension.extension);
+		this.load();
 	}
 
-	public loadOrInitialize() {
-		this._metadata = this.getNotebookMetadata(this._notebook);
+	public load(): void {
+		this._metadata = this.getExtensionMetadata();
+
+		if (this._metadata) {
+			this._metadata.views = this._metadata.views.map(view => NotebookViewModel.load(view.guid, this));
+			this._initialized = true;
+		}
+	}
+
+	public initialize() {
+		this._metadata = this.getExtensionMetadata();
 
 		if (!this._metadata) {
 			this.initializeNotebook();
 			this.initializeCells();
 			this.commit();
-		} else {
-			this._metadata.views = this._metadata.views.map(view => NotebookViewModel.load(view.guid, this));
 		}
+
+		this._initialized = true;
 	}
 
 	protected initializeNotebook() {
@@ -59,12 +69,17 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 			views: []
 		};
 
-		this.setCellMetadata(cell, meta);
+		this.setExtensionCellMetadata(cell, meta);
 	}
 
 	public createNewView(name?: string): INotebookView {
 		const viewName = name || this.generateDefaultViewName();
 
+		// If the notebook has not been initialized, do it now
+		if (!this.initialized) {
+			this.initialize();
+		}
+
 		const view = new NotebookViewModel(viewName, this);
 		view.initialize(true);
 
@@ -77,21 +92,21 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 	}
 
 	public removeView(guid: string) {
-		let viewToRemove = this._metadata.views.findIndex(view => view.guid === guid);
-		if (viewToRemove !== -1) {
-			let removedView = this._metadata.views.splice(viewToRemove, 1);
+		let viewToRemove = this._metadata?.views.findIndex(view => view.guid === guid);
+		if (viewToRemove >= 0) {
+			let removedView = this._metadata?.views.splice(viewToRemove, 1);
 
 			// Remove view data for each cell
 			if (removedView.length === 1) {
 				this._notebook?.cells.forEach((cell) => {
-					let meta = this.getCellMetadata(cell);
+					let meta = this.getExtensionCellMetadata(cell);
 					meta.views = meta.views.filter(x => x.guid !== removedView[0].guid);
-					this.setCellMetadata(cell, meta);
+					this.setExtensionCellMetadata(cell, meta);
 				});
 			}
 		}
 
-		if (guid === this._metadata.activeView) {
+		if (guid === this._metadata?.activeView) {
 			this._metadata.activeView = undefined;
 		}
 
@@ -113,13 +128,13 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 	}
 
 	public updateCell(cell: ICellModel, currentView: INotebookView, cellData: INotebookViewCell, override: boolean = false) {
-		const cellMetadata = this.getCellMetadata(cell);
+		const cellMetadata = this.getExtensionCellMetadata(cell);
 		if (cellMetadata) {
 			const viewToUpdate = cellMetadata.views.findIndex(view => view.guid === currentView.guid);
 
 			if (viewToUpdate >= 0) {
 				cellMetadata.views[viewToUpdate] = override ? cellData : { ...cellMetadata.views[viewToUpdate], ...cellData };
-				this.setCellMetadata(cell, cellMetadata);
+				this.setExtensionCellMetadata(cell, cellMetadata);
 			}
 		}
 	}
@@ -129,7 +144,7 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 	}
 
 	public getViews(): INotebookView[] {
-		return this._metadata.views;
+		return this._metadata?.views ?? [];
 	}
 
 	public get metadata(): INotebookViewMetadata {
@@ -137,21 +152,27 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 	}
 
 	public getCells(): INotebookViewCellMetadata[] {
-		return this._notebook.cells.map(cell => this.getCellMetadata(cell));
+		return this._notebook.cells.map(cell => this.getExtensionCellMetadata(cell));
+	}
+
+	public override getExtensionMetadata(): INotebookViewMetadata {
+		return super.getExtensionMetadata(this._notebook);
 	}
 
 	public getActiveView(): INotebookView {
-		return this.getViews().find(view => view.guid === this._metadata.activeView);
+		return this.getViews().find(view => view.guid === this._metadata?.activeView);
 	}
 
 	public setActiveView(view: INotebookView) {
-		this._metadata.activeView = view.guid;
-		this._onActiveViewChanged.fire();
+		if (this._metadata) {
+			this._metadata.activeView = view.guid;
+			this._onActiveViewChanged.fire();
+		}
 	}
 
 	public commit() {
 		this._metadata = Object.assign({}, this._metadata);
-		this.setNotebookMetadata(this._notebook, this._metadata);
+		this.setExtensionMetadata(this._notebook, this._metadata);
 	}
 
 	public viewNameIsTaken(name: string): boolean {
@@ -165,4 +186,8 @@ export class NotebookViewsExtension extends NotebookExtension<INotebookViewMetad
 	public get onActiveViewChanged(): Event<void> {
 		return this._onActiveViewChanged.event;
 	}
+
+	public get initialized(): boolean {
+		return this._initialized;
+	}
 }