Bring back old editor resolving logic to fix saving in notebooks. (#21186)

src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts

@@ -11,6 +11,7 @@ import { IExtensionService } from 'vs/workbench/services/extensions/common/exten
 import { NotebookInput } from 'sql/workbench/contrib/notebook/browser/models/notebookInput';
 import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
 import { IResourceEditorInput } from 'vs/platform/editor/common/editor';
+import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService';
 
 export class FileNotebookInput extends NotebookInput {
 	public static ID: string = 'workbench.editorinputs.fileNotebookInput';
@@ -23,9 +24,10 @@ export class FileNotebookInput extends NotebookInput {
 		@ITextModelService textModelService: ITextModelService,
 		@IInstantiationService instantiationService: IInstantiationService,
 		@INotebookService notebookService: INotebookService,
-		@IExtensionService extensionService: IExtensionService
+		@IExtensionService extensionService: IExtensionService,
+		@IEditorResolverService editorResolverService: IEditorResolverService,
 	) {
-		super(title, resource, textInput, showActions, textModelService, instantiationService, notebookService, extensionService);
+		super(title, resource, textInput, showActions, textModelService, instantiationService, notebookService, extensionService, editorResolverService);
 	}
 
 	public override get textInput(): FileEditorInput {

src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { IRevertOptions, GroupIdentifier, EditorInputCapabilities, IUntypedEditorInput } from 'vs/workbench/common/editor';
+import { IRevertOptions, GroupIdentifier, EditorInputCapabilities, IUntypedEditorInput, isEditorInputWithOptionsAndGroup, DEFAULT_EDITOR_ASSOCIATION } from 'vs/workbench/common/editor';
 import { Emitter, Event } from 'vs/base/common/event';
 import { URI } from 'vs/base/common/uri';
 import * as resources from 'vs/base/common/resources';
@@ -40,6 +40,8 @@ import { Extensions as LanguageAssociationExtensions, ILanguageAssociationRegist
 import { NotebookLanguage } from 'sql/workbench/common/constants';
 import { convertToInternalInteractiveKernelMetadata } from 'sql/workbench/api/common/notebooks/notebookUtils';
 import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfiguration';
+import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService';
+import { isEqual } from 'vs/base/common/resources';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 const languageAssociationRegistry = Registry.as<ILanguageAssociationRegistry>(LanguageAssociationExtensions.LanguageAssociations);
@@ -247,7 +249,8 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 		@ITextModelService private textModelService: ITextModelService,
 		@IInstantiationService private instantiationService: IInstantiationService,
 		@INotebookService private notebookService: INotebookService,
-		@IExtensionService private extensionService: IExtensionService
+		@IExtensionService private extensionService: IExtensionService,
+		@IEditorResolverService private readonly editorResolverService: IEditorResolverService,
 	) {
 		super();
 		this._standardKernels = [];
@@ -335,18 +338,42 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 
 	override async save(groupId: number, options?: ITextFileSaveOptions): Promise<EditorInput | undefined> {
 		await this.updateModel();
-		let input: any = await this.textInput.save(groupId, options);
+		let untypedInput = await this.textInput.save(groupId, options);
-		await this.setTrustForNewEditor(input);
+		let editorInput = await this.createEditorInput(untypedInput, groupId, false);
+		await this.setTrustForNewEditor(editorInput);
 		const langAssociation = languageAssociationRegistry.getAssociationForLanguage(NotebookLanguage.Ipynb);
-		return langAssociation.convertInput(input);
+		return langAssociation.convertInput(editorInput);
 	}
 
 	override async saveAs(group: number, options?: ITextFileSaveOptions): Promise<EditorInput | undefined> {
 		await this.updateModel();
-		let input: any = await this.textInput.saveAs(group, options);
+		let untypedInput = await this.textInput.saveAs(group, options);
-		await this.setTrustForNewEditor(input);
+		let editorInput = await this.createEditorInput(untypedInput, group, true);
+		await this.setTrustForNewEditor(editorInput);
 		const langAssociation = languageAssociationRegistry.getAssociationForLanguage(NotebookLanguage.Ipynb);
-		return langAssociation.convertInput(input);
+		return langAssociation.convertInput(editorInput);
+	}
+
+	private async createEditorInput(untypedEditor: IUntypedEditorInput | undefined, group: GroupIdentifier, saveAs: boolean): Promise<EditorInput | undefined> {
+		// If this save operation results in a new editor, either
+		// because it was saved to disk (e.g. from untitled) or
+		// through an explicit "Save As", make sure to replace it.
+		if (!untypedEditor) {
+			return undefined; // if we have an undefined input, then the save was cancelled, so do nothing here
+		}
+
+		if ('resource' in untypedEditor) {
+			let target = untypedEditor.resource;
+			if (target.scheme !== this.textInput.resource.scheme || (saveAs && !isEqual(target, this.textInput.preferredResource))
+			) {
+				const editor = await this.editorResolverService.resolveEditor({ resource: target, options: { override: DEFAULT_EDITOR_ASSOCIATION.id } }, group);
+				if (isEditorInputWithOptionsAndGroup(editor)) {
+					return editor.editor;
+				}
+			}
+		}
+
+		return this.textInput;
 	}
 
 	private async setTrustForNewEditor(newInput: EditorInput | undefined): Promise<void> {

src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts

@@ -12,6 +12,7 @@ import { INotebookService } from 'sql/workbench/services/notebook/browser/notebo
 import { UntitledTextEditorInput } from 'vs/workbench/services/untitled/common/untitledTextEditorInput';
 import { EditorInputCapabilities } from 'vs/workbench/common/editor';
 import { UNTITLED_NOTEBOOK_TYPEID } from 'sql/workbench/common/constants';
+import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService';
 
 export class UntitledNotebookInput extends NotebookInput {
 	public static ID: string = UNTITLED_NOTEBOOK_TYPEID;
@@ -23,9 +24,10 @@ export class UntitledNotebookInput extends NotebookInput {
 		@ITextModelService textModelService: ITextModelService,
 		@IInstantiationService instantiationService: IInstantiationService,
 		@INotebookService notebookService: INotebookService,
-		@IExtensionService extensionService: IExtensionService
+		@IExtensionService extensionService: IExtensionService,
+		@IEditorResolverService editorResolverService: IEditorResolverService,
 	) {
-		super(title, resource, textInput, true, textModelService, instantiationService, notebookService, extensionService);
+		super(title, resource, textInput, true, textModelService, instantiationService, notebookService, extensionService, editorResolverService);
 		// Set the mode explicitly so that the auto language detection doesn't run and mark the model as being JSON
 		this.textInput.resolve().then(() => this.setMode(textInput.model.getLanguageId()));
 	}

src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts

@@ -58,6 +58,7 @@ import { CellModel } from 'sql/workbench/services/notebook/browser/models/cell';
 import { IEditorOptions } from 'vs/platform/editor/common/editor';
 import { ITextResourceConfigurationService } from 'vs/editor/common/services/textResourceConfiguration';
 import { FindReplaceStateChangedEvent, INewFindReplaceState } from 'vs/editor/contrib/find/browser/findState';
+import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService';
 
 class NotebookModelStub extends stubs.NotebookModelStub {
 	public contentChangedEmitter = new Emitter<NotebookContentChange>();
@@ -103,10 +104,11 @@ suite('Test class NotebookEditor:', () => {
 	let queryTextEditor: QueryTextEditor;
 	let untitledNotebookInput: UntitledNotebookInput;
 	let notebookEditorStub: NotebookEditorStub;
+	let editorResolverService: IEditorResolverService;
 
 	setup(async () => {
 		// setup services
-		({ instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub } = setupServices({ instantiationService, workbenchThemeService }));
+		({ instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, editorResolverService } = setupServices({ instantiationService, workbenchThemeService }));
 		// Create notebookEditor
 		notebookEditor = createNotebookEditor(instantiationService, workbenchThemeService, notebookService);
 	});
@@ -127,7 +129,7 @@ suite('Test class NotebookEditor:', () => {
 			const untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, untitledTextEditorService.create({ associatedResource: untitledUri }));
 			const untitledNotebookInput = new UntitledNotebookInput(
 				testTitle, untitledUri, untitledTextInput,
-				undefined, instantiationService, notebookService, extensionService
+				undefined, instantiationService, notebookService, extensionService, editorResolverService
 			);
 			const testNotebookEditor = new NotebookEditorStub({ cellGuid: cellTextEditorGuid, editor: queryTextEditor, model: notebookModel, notebookParams: <INotebookParams>{ notebookUri: untitledNotebookInput.notebookUri } });
 			notebookService.addNotebookEditor(testNotebookEditor);
@@ -714,9 +716,10 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins
 	const untitledUri = URI.from({ scheme: Schemas.untitled, path: 'NotebookEditor.Test-TestPath' });
 	const untitledTextEditorService = instantiationService.get(IUntitledTextEditorService);
 	const untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, untitledTextEditorService.create({ associatedResource: untitledUri }));
+	let editorResolverService = instantiationService.get(IEditorResolverService);
 	const untitledNotebookInput = new UntitledNotebookInput(
 		testTitle, untitledUri, untitledTextInput,
-		undefined, instantiationService, notebookService, extensionService
+		undefined, instantiationService, notebookService, extensionService, editorResolverService
 	);
 
 	const cellTextEditorGuid = generateUuid();
@@ -731,7 +734,7 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins
 	);
 	const notebookEditorStub = new NotebookEditorStub({ cellGuid: cellTextEditorGuid, editor: queryTextEditor, model: new NotebookModelStub(), notebookParams: <INotebookParams>{ notebookUri: untitledNotebookInput.notebookUri } });
 	notebookService.addNotebookEditor(notebookEditorStub);
-	return { instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub };
+	return { instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, editorResolverService };
 }
 
 function createNotebookEditor(instantiationService: TestInstantiationService, workbenchThemeService: WorkbenchThemeService, notebookService: NotebookService) {

src/sql/workbench/contrib/notebook/test/browser/notebookInput.test.ts

@@ -22,6 +22,7 @@ import { TestInstantiationService } from 'vs/platform/instantiation/test/common/
 import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService';
 import { EditorInputCapabilities } from 'vs/workbench/common/editor';
 import { LocalContentManager } from 'sql/workbench/services/notebook/common/localContentManager';
+import { IEditorResolverService } from 'vs/workbench/services/editor/common/editorResolverService';
 
 suite('Notebook Input', function (): void {
 	const instantiationService = workbenchInstantiationService();
@@ -55,20 +56,22 @@ suite('Notebook Input', function (): void {
 	let untitledTextInput: UntitledTextEditorInput;
 	let untitledNotebookInput: UntitledNotebookInput;
 
+	const editorResolverService = instantiationService.get(IEditorResolverService);
+
 	setup(() => {
 		const accessor = instantiationService.createInstance(ServiceAccessor);
 		const service = accessor.untitledTextEditorService;
 		untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, service.create({ associatedResource: untitledUri }));
 		untitledNotebookInput = new UntitledNotebookInput(
 			testTitle, untitledUri, untitledTextInput,
-			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, editorResolverService);
 	});
 
 	test('File Notebook Input', async function (): Promise<void> {
 		let fileUri = URI.from({ scheme: Schemas.file, path: 'TestPath' });
 		let fileNotebookInput = new FileNotebookInput(
 			testTitle, fileUri, undefined, true,
-			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, editorResolverService);
 
 		let inputId = fileNotebookInput.typeId;
 		assert.strictEqual(inputId, FileNotebookInput.ID);