diff --git a/src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts b/src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts index 3612ea3532..e28b7bcdfe 100644 --- a/src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts +++ b/src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts @@ -366,8 +366,6 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges { private updateCellSource(): void { let textOutputElement = this.output.nativeElement; let newCellSource: string = this._htmlMarkdownConverter.convert(textOutputElement.innerHTML); - // reset cell attachments to remove unused image data since we're going to go through each of them again - this.cellModel.attachments = {}; this.cellModel.source = newCellSource; this._changeRef.detectChanges(); } 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 8ab3f89e20..332c115afb 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 @@ -1234,6 +1234,97 @@ suite('Cell Model', function (): void { assert.deepStrictEqual(cellModel.attachments, attachments, 'addAttachment should not add duplicate images'); }); + test('deleting attachment from cell source should remove it from attachments should add a valid attachment to cell', async function () { + let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU'; + let index = imageFilebase64Value.indexOf('base64,'); + const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) }; + let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment }; + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '', + metadata: {} + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.addAttachment('image/png', imageFilebase64Value, 'test.png'); + assert.deepStrictEqual(model.attachments, attachments, 'attachment was not added initially'); + model.source = ''; + assert.deepStrictEqual(model.attachments, {}, 'attachments should be empty after clearing attachment from cell'); + }); + + test('modifying cell source with existing attachment keeps attachment', async function () { + let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU'; + let index = imageFilebase64Value.indexOf('base64,'); + const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) }; + let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment }; + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![image](attachment:test.png)', + metadata: {} + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.addAttachment('image/png', imageFilebase64Value, 'test.png'); + assert.deepStrictEqual(model.attachments, attachments, 'attachment was not added initially'); + model.source = 'Some new text ' + model.source; + assert.deepStrictEqual(model.attachments, attachments, 'attachment should still exist after modifying cell source'); + }); + + test('modifying cell source with multiple attachments keeps all attachments', async function () { + let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU'; + let index = imageFilebase64Value.indexOf('base64,'); + const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) }; + let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment, 'test2.png': testImageAttachment }; + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![image](attachment:test.png)\n![image2](attachment:test2.png)', + metadata: {} + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.addAttachment('image/png', imageFilebase64Value, 'test.png'); + model.addAttachment('image/png', imageFilebase64Value, 'test2.png'); + assert.deepStrictEqual(model.attachments, attachments, 'attachments were not added initially'); + model.source = 'Some new text ' + model.source; + assert.deepStrictEqual(model.attachments, attachments, 'attachments should still exist after modifying cell source'); + }); + + test('removing attachment keeps any others present in cell', async function () { + let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU'; + let index = imageFilebase64Value.indexOf('base64,'); + const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) }; + let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment, 'test2.png': testImageAttachment }; + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Markdown, + source: '![image](attachment:test.png)\n![image2](attachment:test2.png)', + metadata: {} + }; + let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }); + model.addAttachment('image/png', imageFilebase64Value, 'test.png'); + model.addAttachment('image/png', imageFilebase64Value, 'test2.png'); + assert.deepStrictEqual(model.attachments, attachments, 'attachments were not added initially'); + model.source = '![image2](attachment:test2.png)'; + delete attachments['test.png']; + assert.deepStrictEqual(model.attachments, attachments, 'we should still have one attachment after removing the other'); + }); + test('cell should fire onCurrentEditModeChanged on edit', async function () { let notebookModel = new NotebookModelStub({ diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index 1e2321d939..8d85c1bfd5 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/src/sql/workbench/services/notebook/browser/models/cell.ts @@ -307,6 +307,7 @@ export class CellModel extends Disposable implements ICellModel { } public set source(newSource: string | string[]) { + this.cleanUnusedAttachments(Array.isArray(newSource) ? newSource.join() : newSource); newSource = this.attachImageFromSource(newSource); newSource = this.getMultilineSource(newSource); if (this._source !== newSource) { @@ -329,6 +330,22 @@ 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; + this._attachments = {}; + // Find existing attachments in the form ![...](attachment:...) so that we can make sure we keep those attachments + const attachmentRegex = /!\[.*?\]\(attachment:(.*?)\)/g; + let match; + while (match = attachmentRegex.exec(source)) { // eslint-disable-line no-cond-assign + this._attachments[match[1]] = originalAttachments[match[1]]; + } + } + /** * Gets unique attachment name to add to cell metadata * @param imgName a string defining name of the image.