Fixes to trusting native notebooks (#13292)

Author: DonJayamanne

src/client/datascience/notebook/kernelProvider.ts

@@ -17,8 +17,9 @@ import { KernelSelectionProvider } from '../jupyter/kernels/kernelSelections';
 import { KernelSelector } from '../jupyter/kernels/kernelSelector';
 import { KernelSwitcher } from '../jupyter/kernels/kernelSwitcher';
 import { KernelSelection } from '../jupyter/kernels/types';
+import { INotebookStorageProvider } from '../notebookStorage/notebookStorageProvider';
 import { INotebook, INotebookProvider } from '../types';
-import { getNotebookMetadata, updateKernelInNotebookMetadata } from './helpers/helpers';
+import { getNotebookMetadata, isJupyterNotebook, updateKernelInNotebookMetadata } from './helpers/helpers';
 import { NotebookKernel } from './notebookKernel';
 import { INotebookContentProvider, INotebookExecutionService } from './types';
 @injectable()
@@ -33,6 +34,7 @@ export class VSCodeKernelPickerProvider implements NotebookKernelProvider {
         @inject(KernelSelectionProvider) private readonly kernelSelectionProvider: KernelSelectionProvider,
         @inject(KernelSelector) private readonly kernelSelector: KernelSelector,
         @inject(IVSCodeNotebook) private readonly notebook: IVSCodeNotebook,
+        @inject(INotebookStorageProvider) private readonly storageProvider: INotebookStorageProvider,
         @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider,
         @inject(KernelSwitcher) private readonly kernelSwitcher: KernelSwitcher,
         @inject(INotebookContentProvider) private readonly notebookContentProvider: INotebookContentProvider,
@@ -119,13 +121,21 @@ export class VSCodeKernelPickerProvider implements NotebookKernelProvider {
         }
 
         const document = newKernelInfo.document;
+        if (!isJupyterNotebook(document)) {
+            return;
+        }
         const selection = await newKernelInfo.kernel.validate(document.uri);
         const editor = this.notebook.notebookEditors.find((item) => item.document === document);
         if (!selection || !editor || editor.kernel !== newKernelInfo.kernel) {
             // Possibly closed or different kernel picked.
             return;
         }
-
+        const model = await this.storageProvider.getOrCreateModel(document.uri);
+        if (!model || !model.isTrusted) {
+            // If a model is not trusted, we cannot change the kernel (this results in changes to notebook metadata).
+            // This is because we store selected kernel in the notebook metadata.
+            return;
+        }
         // Change kernel and update metadata.
         const notebook = await this.notebookProvider.getOrCreateNotebook({
             resource: document.uri,

src/client/datascience/notebookStorage/vscNotebookModel.ts

@@ -41,7 +41,9 @@ export class VSCodeNotebookModel extends BaseNotebookModel {
         return this.document?.isDirty === true;
     }
     public get cells(): ICell[] {
-        return this.document
+        // When a notebook is not trusted, return original cells.
+        // This is because the VSCode NotebookDocument object will not have any output in the cells.
+        return this.document && this.isTrusted
             ? this.document.cells.map((cell) => createCellFromVSCNotebookCell(cell, this))
             : this._cells;
     }
@@ -82,7 +84,7 @@ export class VSCodeNotebookModel extends BaseNotebookModel {
     }
     protected generateNotebookJson() {
         const json = super.generateNotebookJson();
-        if (this.document) {
+        if (this.document && this.isTrusted) {
             // The metadata will be in the notebook document.
             const metadata = getNotebookMetadata(this.document);
             if (metadata) {

src/test/datascience/notebook/contentProvider.ds.test.ts

@@ -13,6 +13,7 @@ import { CellErrorOutput, Uri } from 'vscode';
 import { CellDisplayOutput } from '../../../../types/vscode-proposed';
 import { IVSCodeNotebook } from '../../../client/common/application/types';
 import { IDisposable } from '../../../client/common/types';
+import { INotebookStorageProvider } from '../../../client/datascience/notebookStorage/notebookStorageProvider';
 import { INotebookEditorProvider } from '../../../client/datascience/types';
 import { IExtensionTestApi } from '../../common';
 import { EXTENSION_ROOT_DIR_FOR_TESTS, initialize } from '../../initialize';
@@ -49,7 +50,28 @@ suite('DataScience - VSCode Notebook - (Open)', function () {
         testIPynb = Uri.file(await createTemporaryNotebook(templateIPynb, disposables));
     });
     teardown(async () => closeNotebooksAndCleanUpAfterTests(disposables));
+    test('Verify Notebook Json', async () => {
+        const storageProvider = api.serviceContainer.get<INotebookStorageProvider>(INotebookStorageProvider);
+        const file = path.join(
+            EXTENSION_ROOT_DIR_FOR_TESTS,
+            'src',
+            'test',
+            'datascience',
+            'notebook',
+            'testJsonContents.ipynb'
+        );
+        const model = await storageProvider.getOrCreateModel(Uri.file(file));
+        disposables.push(model);
+        model.trust();
+        const jsonStr = fs.readFileSync(file, { encoding: 'utf8' });
+
+        // JSON should be identical, before and after trusting a notebook.
+        assert.deepEqual(JSON.parse(jsonStr), JSON.parse(model.getContent()));
 
+        model.trust();
+
+        assert.deepEqual(JSON.parse(jsonStr), JSON.parse(model.getContent()));
+    });
     test('Verify cells (content, metadata & output)', async () => {
         const vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
         const editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);