Fix intellisense for .NET Interactive SQL kernel (#19254)

src/sql/workbench/api/common/notebooks/notebookUtils.ts

@@ -17,6 +17,10 @@ const DotnetInteractiveLanguagePrefix = 'dotnet-interactive.';
 export const DotnetInteractiveDisplayName = '.NET Interactive';
 
 export function convertToVSCodeNotebookCell(cellKind: azdata.nb.CellType, cellIndex: number, cellUri: URI, docUri: URI, cellLanguage: string, cellSource?: string | string[]): vscode.NotebookCell {
+	// We only use this notebook field for .NET Interactive's intellisense, which only uses the notebook's URI
+	let notebook = <vscode.NotebookDocument>{
+		uri: docUri
+	};
 	return <vscode.NotebookCell>{
 		kind: cellKind === CellTypes.Code ? NotebookCellKind.Code : NotebookCellKind.Markup,
 		index: cellIndex,
@@ -24,10 +28,9 @@ export function convertToVSCodeNotebookCell(cellKind: azdata.nb.CellType, cellIn
 			uri: cellUri,
 			languageId: cellLanguage,
 			getText: () => Array.isArray(cellSource) ? cellSource.join('') : (cellSource ?? ''),
+			notebook: notebook
 		},
-		notebook: <vscode.NotebookDocument>{
+		notebook: notebook,
-			uri: docUri
-		},
 		outputs: [],
 		metadata: {},
 		mime: undefined

src/sql/workbench/common/constants.ts

@@ -41,6 +41,7 @@ export const INTERACTIVE_PROVIDER_ID = 'dotnet-interactive';
 export const INTERACTIVE_LANGUAGE_MODE = 'dib';
 export const DEFAULT_NB_LANGUAGE_MODE = 'notebook';
 export const TSGOPS_WEB_QUALITY = 'tsgops-image';
+export const CELL_URI_PATH_PREFIX = 'notebook-editor-';
 
 // The version of the notebook file format that we support
 export const NBFORMAT = 4;

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -235,6 +235,9 @@ export class ServerManagerStub implements nb.ServerManager {
 }
 
 export class NotebookServiceStub implements INotebookService {
+	getNotebookURIForCell(cellUri: URI): URI {
+		throw new Error('Method not implemented.');
+	}
 	getSupportedLanguagesForProvider(provider: string, kernelDisplayName?: string): Promise<string[]> {
 		throw new Error('Method not implemented.');
 	}

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

@@ -35,6 +35,7 @@ import { CellOutputEdit, CellOutputDataEdit } from 'sql/workbench/services/noteb
 import { ILogService } from 'vs/platform/log/common/log';
 import { IModeService } from 'vs/editor/common/services/modeService';
 import { ICellMetadata } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { CELL_URI_PATH_PREFIX } from 'sql/workbench/common/constants';
 
 let modelId = 0;
 const ads_execute_command = 'ads_execute_command';
@@ -1121,7 +1122,7 @@ export class CellModel extends Disposable implements ICellModel {
 	}
 
 	private createUri(): void {
-		let uri = URI.from({ scheme: Schemas.untitled, path: `notebook-editor-${this.id}` });
+		let uri = URI.from({ scheme: Schemas.untitled, path: `${CELL_URI_PATH_PREFIX}${this.id}` });
 		// Use this to set the internal (immutable) and public (shared with extension) uri properties
 		this.cellUri = uri;
 	}

src/sql/workbench/services/notebook/browser/notebookService.ts

@@ -144,6 +144,8 @@ export interface INotebookService {
 	openNotebook(resource: UriComponents, options: INotebookShowOptions): Promise<IEditorPane | undefined>;
 
 	getUntitledUriPath(originalTitle: string): string;
+
+	getNotebookURIForCell(cellUri: URI): URI | undefined;
 }
 
 export interface IExecuteProvider {

src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts

@@ -371,6 +371,19 @@ export class NotebookService extends Disposable implements INotebookService {
 		return title;
 	}
 
+	public getNotebookURIForCell(cellUri: URI): URI | undefined {
+		for (let editor of this.listNotebookEditors()) {
+			if (editor.cells) {
+				for (let cell of editor.cells) {
+					if (cell.cellUri === cellUri) {
+						return editor.notebookParams.notebookUri;
+					}
+				}
+			}
+		}
+		return undefined;
+	}
+
 	private updateSQLRegistrationWithConnectionProviders() {
 		// Update the SQL extension
 		let sqlNotebookKernels = this._providerToStandardKernels.get(notebookConstants.SQL);

src/sql/workbench/test/electron-browser/api/vscodeNotebookApi.test.ts

@@ -425,6 +425,9 @@ suite('Notebook Serializer', () => {
 		assert.deepStrictEqual(actual.document.uri, expected.document.uri);
 		assert.strictEqual(actual.document.languageId, expected.document.languageId);
 		assert.deepStrictEqual(actual.notebook.uri, expected.notebook.uri);
+		assert.deepStrictEqual(actual.document.notebook.uri, expected.document.notebook.uri);
+		assert.deepStrictEqual(actual.document.notebook.uri, expected.notebook.uri);
+		assert.deepStrictEqual(actual.notebook.uri, expected.document.notebook.uri);
 	}
 	function validateCellsMatch(actual: vscode.NotebookCell[], expected: vscode.NotebookCell[]): void {
 		assert.strictEqual(actual.length, expected.length, 'Cell arrays did not have equal lengths.');

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

@@ -33,6 +33,8 @@ import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService
 import { IPathService } from 'vs/workbench/services/path/common/pathService';
 import { diffSets, diffMaps } from 'vs/base/common/collections';
 import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
+import { Schemas } from 'vs/base/common/network';
+import { CELL_URI_PATH_PREFIX } from 'sql/workbench/common/constants';
 
 
 class TextEditorSnapshot {
@@ -390,13 +392,20 @@ export class MainThreadDocumentsAndEditors {
 	}
 
 	private _toModelAddData(model: ITextModel): IModelAddedData {
+		// {{SQL CARBON EDIT}}
+		// Check if this TextModel is part of a notebook cell
+		let notebookUri: URI;
+		if (model.uri.scheme === Schemas.untitled && model.uri.path.startsWith(CELL_URI_PATH_PREFIX)) {
+			notebookUri = this._notebookService.getNotebookURIForCell(model.uri);
+		}
 		return {
 			uri: model.uri,
 			versionId: model.getVersionId(),
 			lines: model.getLinesContent(),
 			EOL: model.getEOL(),
 			modeId: model.getLanguageIdentifier().language,
-			isDirty: this._textFileService.isDirty(model.uri)
+			isDirty: this._textFileService.isDirty(model.uri),
+			notebookUri: notebookUri
 		};
 	}
 

src/vs/workbench/api/common/extHost.protocol.ts

@@ -1212,6 +1212,8 @@ export interface IModelAddedData {
 	EOL: string;
 	modeId: string;
 	isDirty: boolean;
+	// {{SQL CARBON EDIT}}
+	notebookUri?: URI;
 }
 export interface ExtHostDocumentsShape {
 	$acceptModelModeChanged(strURL: UriComponents, newModeId: string): void;

src/vs/workbench/api/common/extHostDocumentsAndEditors.ts

@@ -97,6 +97,14 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 					}
 				}
 				if (!ref) {
+					// {{SQL CARBON EDIT}}
+					// Add URI of the notebook that is using this document for a cell's editor.
+					if (!data.notebook && data.notebookUri) {
+						// We only use this notebook field for .NET Interactive's intellisense, which only uses the notebook's URI
+						data.notebook = <vscode.NotebookDocument>{
+							uri: URI.revive(data.notebookUri)
+						};
+					}
 					ref = new Reference(new ExtHostDocumentData(
 						this._extHostRpc.getProxy(MainContext.MainThreadDocuments),
 						resource,

src/vs/workbench/services/workingCopy/common/workingCopyService.ts

@@ -11,6 +11,7 @@ import { Disposable, IDisposable, toDisposable, DisposableStore, dispose } from
 import { ResourceMap } from 'vs/base/common/map';
 import { IWorkingCopy, IWorkingCopyIdentifier } from 'vs/workbench/services/workingCopy/common/workingCopy';
 import { Schemas } from 'vs/base/common/network'; // {{SQL CARBON EDIT}} @chlafreniere need to block working copies of notebook editors from being tracked
+import { CELL_URI_PATH_PREFIX } from 'sql/workbench/common/constants';
 
 export const IWorkingCopyService = createDecorator<IWorkingCopyService>('workingCopyService');
 
@@ -142,7 +143,7 @@ export class WorkingCopyService extends Disposable implements IWorkingCopyServic
 		}
 
 		// {{SQL CARBON EDIT}} @chlafreniere need to block working copies of notebook editors from being tracked
-		if (workingCopy.resource.path.includes('notebook-editor-') && workingCopy.resource.scheme === Schemas.untitled) {
+		if (workingCopy.resource.path.includes(CELL_URI_PATH_PREFIX) && workingCopy.resource.scheme === Schemas.untitled) {
 			return new DisposableStore();
 		}