Fix ModelView tree components to keep custom state (#11239)

* Fix ModelView tree components to keep custom state

* override postGetChildren

* deepClone

src/sql/workbench/browser/modelComponents/treeViewDataProvider.ts

@@ -11,6 +11,24 @@ import { IModelViewTreeViewDataProvider, ITreeComponentItem } from 'sql/workbenc
 import { INotificationService } from 'vs/platform/notification/common/notification';
 // eslint-disable-next-line code-import-patterns
 import * as vsTreeView from 'vs/workbench/api/browser/mainThreadTreeViews';
+import { ResolvableTreeItem } from 'vs/workbench/common/views';
+import { deepClone } from 'vs/base/common/objects';
+
+export class ResolvableTreeComponentItem extends ResolvableTreeItem implements ITreeComponentItem {
+
+	checked?: boolean;
+	enabled?: boolean;
+	onCheckedChanged?: (checked: boolean) => void;
+	children?: ITreeComponentItem[];
+
+	constructor(treeItem: ITreeComponentItem, resolve?: (() => Promise<ITreeComponentItem | undefined>)) {
+		super(treeItem, resolve);
+		this.checked = treeItem.checked;
+		this.enabled = treeItem.enabled;
+		this.onCheckedChanged = treeItem.onCheckedChanged;
+		this.children = deepClone(treeItem.children);
+	}
+}
 
 export class TreeViewDataProvider extends vsTreeView.TreeViewDataProvider implements IModelViewTreeViewDataProvider {
 	constructor(handle: number, treeViewId: string,
@@ -32,4 +50,24 @@ export class TreeViewDataProvider extends vsTreeView.TreeViewDataProvider implem
 
 	refresh(itemsToRefreshByHandle: { [treeItemHandle: string]: ITreeComponentItem }) {
 	}
+
+	/**
+	 * Returns the set of mapped ResolvableTreeComponentItems
+	 * @override
+	 * @param elements The elements to map
+	 */
+	protected async postGetChildren(elements: ITreeComponentItem[]): Promise<ResolvableTreeComponentItem[]> {
+		const result: ResolvableTreeComponentItem[] = [];
+		const hasResolve = await this.hasResolve;
+		if (elements) {
+			for (const element of elements) {
+				const resolvable = new ResolvableTreeComponentItem(element, hasResolve ? () => {
+					return this._proxy.$resolve(this.treeViewId, element.handle);
+				} : undefined);
+				this.itemsMap.set(element.handle, resolvable);
+				result.push(resolvable);
+			}
+		}
+		return result;
+	}
 }

src/vs/workbench/api/browser/mainThreadTreeViews.ts

@@ -164,8 +164,8 @@ export type TreeItemHandle = string;
 // {{SQL CARBON EDIT}}
 export class TreeViewDataProvider implements ITreeViewDataProvider {
 
-	private readonly itemsMap: Map<TreeItemHandle, ITreeItem> = new Map<TreeItemHandle, ITreeItem>();
+	protected readonly itemsMap: Map<TreeItemHandle, ITreeItem> = new Map<TreeItemHandle, ITreeItem>(); // {{SQL CARBON EDIT}} For use by Component Tree View
-	private hasResolve: Promise<boolean>;
+	protected hasResolve: Promise<boolean>; // {{SQL CARBON EDIT}} For use by Component Tree View
 
 	// {{SQL CARBON EDIT}}
 	constructor(protected readonly treeViewId: string,
@@ -219,7 +219,7 @@ export class TreeViewDataProvider implements ITreeViewDataProvider {
 		return this.itemsMap.size === 0;
 	}
 
-	private async postGetChildren(elements: ITreeItem[]): Promise<ResolvableTreeItem[]> {
+	protected async postGetChildren(elements: ITreeItem[]): Promise<ResolvableTreeItem[]> { // {{SQL CARBON EDIT}} For use by Component Tree View
 		const result: ResolvableTreeItem[] = [];
 		const hasResolve = await this.hasResolve;
 		if (elements) {