From f06fd7f6f7874327562014e611c2e6dfc71d315e Mon Sep 17 00:00:00 2001 From: Barbara Valdez <34872381+barbaravaldez@users.noreply.github.com> Date: Wed, 1 Jun 2022 16:23:46 -0700 Subject: [PATCH] Fix split cell not working with attachments (#19587) --- .../test/electron-browser/cell.test.ts | 90 +++++++++++++++++++ .../services/notebook/browser/models/cell.ts | 10 +-- .../browser/models/modelInterfaces.ts | 7 ++ .../notebook/browser/models/notebookModel.ts | 9 +- 4 files changed, 108 insertions(+), 8 deletions(-) diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts index f8941d3f6f..2f74df6ba6 100644 --- a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts +++ b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts @@ -1079,6 +1079,96 @@ suite('Cell Model', function (): void { assert.deepStrictEqual(serializedCell.attachments, undefined, 'JSON should not include attachments if attachments do not exist'); }); + test('Should remove unused attachments name when updating cell source', async function () { + const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}'); + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '', + attachments: cellAttachment + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + assert.deepStrictEqual(model.attachments, contents.attachments, 'Attachments do not match in cellModel'); + + model.source = 'Test'; + + assert.notDeepStrictEqual(model.attachments, contents.attachments, 'Unused attachments are not removed after updating cell source'); + }); + + test('Should not remove attachments that are still referenced in cell after updating the cell source', async function () { + const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}'); + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![ads.png](attachment:ads.png)', + attachments: cellAttachment + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + + model.source = '![ads.png](attachment:ads.png) \n Test'; + + assert.deepStrictEqual(model.attachments, contents.attachments, 'Attachments referenced in cell were removed'); + }); + + test('Should update cell attachments', async function () { + const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}'); + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![ads.png](attachment:ads.png)' + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.updateAttachmentsFromSource('![ads.png](attachment:ads.png)', cellAttachment); + + assert.deepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are not updated correctly'); + }); + + test('Should not update cell attachments if they are not referenced in cell source', async function () { + const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}'); + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![ads.png](attachment:ads.png)' + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.updateAttachmentsFromSource('test', cellAttachment); + + assert.notDeepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are updated when they are not referenced in cell source'); + }); + + test('Should not update cell attachments if the attachment reference is not in the correct format', async function () { + const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}'); + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![ads.png](attachment:ads.png)' + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.updateAttachmentsFromSource('ads.png', cellAttachment); + + assert.notDeepStrictEqual(model.attachments, cellAttachment, 'Cell attachments are updated when the reference is not in the correct format'); + }); + test('Should not have cache chart data after new cell created', async function () { let notebookModel = new NotebookModelStub({ name: '', diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index 68be80d6e1..eff3c4a362 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/src/sql/workbench/services/notebook/browser/models/cell.ts @@ -334,7 +334,7 @@ export class CellModel extends Disposable implements ICellModel { } public set source(newSource: string | string[]) { - this.cleanUnusedAttachments(Array.isArray(newSource) ? newSource.join() : newSource); + this.updateAttachmentsFromSource(Array.isArray(newSource) ? newSource.join() : newSource); newSource = this.attachImageFromSource(newSource); newSource = this.getMultilineSource(newSource); if (this._source !== newSource) { @@ -358,12 +358,8 @@ export class CellModel extends Disposable implements ICellModel { return newSource; } - /** - * Cleans up the attachments, removing any ones that aren't being currently used in the specified source string. - * @param source The new source string to check for attachments being used - */ - private cleanUnusedAttachments(source: string): void { - const originalAttachments = this._attachments; + public updateAttachmentsFromSource(source: string, attachments?: nb.ICellAttachments): void { + const originalAttachments = attachments ? attachments : this._attachments; this._attachments = {}; // Find existing attachments in the form ![...](attachment:...) so that we can make sure we keep those attachments const attachmentRegex = /!\[.*?\]\(attachment:(.*?)\)/g; diff --git a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts index 4c2ca7739a..99fc54a990 100644 --- a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts +++ b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts @@ -554,6 +554,13 @@ export interface ICellModel { * Returns the name of the attachment added to metadata. */ addAttachment(mimeType: string, base64Encoding: string, name: string): string; + /** + * Updates the current cell attachments with the attachments provided. + * If no attachments are passed in then it cleans up the current cell attachments and removes any ones that aren't being currently used in the specified source string. + * @param source The new source string to check for attachments being used + * @param attachments (Optional) The new attachments for the cell + */ + updateAttachmentsFromSource(source: string, attachments?: nb.ICellAttachments): void; richTextCursorPosition: ICaretPosition; markdownCursorPosition: IPosition; /** diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index 67ae36e435..b176470587 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts @@ -570,10 +570,11 @@ export class NotebookModel extends Disposable implements INotebookModel { let range = model.getFullModelRange(); let selection = editorControl.getSelection(); let source = this.cells[index].source; - let newCell = undefined, tailCell = undefined, partialSource = undefined; + let newCell: ICellModel = undefined, tailCell: ICellModel = undefined, partialSource = undefined; let newCellIndex = index; let tailCellIndex = index; let splitCells: SplitCell[] = []; + let attachments = {}; // Save UI state let showMarkdown = this.cells[index].showMarkdown; @@ -607,6 +608,9 @@ export class NotebookModel extends Disposable implements INotebookModel { partialSource = source.slice(selection.startLineNumber - 1, selection.startLineNumber)[0].slice(0, selection.startColumn - 1); headsource = headsource.concat(partialSource.toString()); } + // Save attachments before updating cell contents + attachments = this.cells[index].attachments; + // No need to update attachments, since unused attachments are removed when updating the cell source this.cells[index].source = headsource; splitCells.push({ cell: this.cells[index], prefix: undefined }); } @@ -628,6 +632,7 @@ export class NotebookModel extends Disposable implements INotebookModel { //If the selection is not from the start of the cell, create a new cell. if (headContent.length) { newCell = this.createCell(cellType, language); + newCell.updateAttachmentsFromSource(newSource.join(), attachments); newCell.source = newSource; newCellIndex++; this.insertCell(newCell, newCellIndex, false); @@ -651,6 +656,7 @@ export class NotebookModel extends Disposable implements INotebookModel { if (tailSource[0] === '\r\n' || tailSource[0] === '\n') { newlinesBeforeTailCellContent = tailSource.splice(0, 1)[0]; } + tailCell.updateAttachmentsFromSource(tailSource.join(), attachments); tailCell.source = tailSource; tailCellIndex = newCellIndex + 1; this.insertCell(tailCell, tailCellIndex, false); @@ -684,6 +690,7 @@ export class NotebookModel extends Disposable implements INotebookModel { let firstCell = cells[0].cell; // Append the other cell sources to the first cell for (let i = 1; i < cells.length; i++) { + firstCell.attachments = { ...firstCell.attachments, ...cells[i].cell.attachments }; firstCell.source = cells[i].prefix ? [...firstCell.source, ...cells[i].prefix, ...cells[i].cell.source] : [...firstCell.source, ...cells[i].cell.source]; } // Set newly created cell as active cell