Add image as attachment on copy/paste into cell (#15602)

* add pasted image as attachment

* handle duplicate image logic

* replace with regex

* address PR comments

src/sql/workbench/contrib/notebook/browser/cellViews/markdownToolbar.component.ts

@@ -272,7 +272,11 @@ export class MarkdownToolbarComponent extends AngularDisposable {
 				if (imageCalloutResult.embedImage) {
 					let base64String = await this.getFileContentBase64(URI.file(imageCalloutResult.imagePath));
 					let mimeType = await this.getFileMimeType(URI.file(imageCalloutResult.imagePath));
-					this.cellModel.addAttachment(mimeType, base64String, path.basename(imageCalloutResult.imagePath).replace(' ', ''));
+					const originalImageName: string = path.basename(imageCalloutResult.imagePath).replace(/\s/g, '');
+					let attachmentName = this.cellModel.addAttachment(mimeType, base64String, originalImageName);
+					if (originalImageName !== attachmentName) {
+						imageCalloutResult.insertEscapedMarkdown = `![${attachmentName}](attachment:${attachmentName.replace(/\s/g, '')})`;
+					}
 					await insertFormattedMarkdown(imageCalloutResult.insertEscapedMarkdown, this.getCellEditorControl());
 				}
 				await insertFormattedMarkdown(imageCalloutResult.insertEscapedMarkdown, this.getCellEditorControl());

src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts

@@ -1175,7 +1175,7 @@ suite('Cell Model', function (): void {
 		let imageFilebase64Value = 'data:application/octet-stream;base64,iVBORw0KGgoAAAANSU';
 		let index = imageFilebase64Value.indexOf('base64,');
 		const testImageAttachment: nb.ICellAttachment = { ['image/png']: imageFilebase64Value.substring(index + 7) };
-		const attachments: nb.ICellAttachments = { 'test.png': testImageAttachment };
+		let attachments: nb.ICellAttachments = { 'test.png': testImageAttachment };
 		let notebookModel = new NotebookModelStub({
 			name: '',
 			version: '',
@@ -1189,6 +1189,9 @@ suite('Cell Model', function (): void {
 		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
 		model.addAttachment('image/png', imageFilebase64Value, 'test.png');
 		assert.deepEqual(model.attachments, attachments);
+		attachments = { 'test.png': testImageAttachment, 'test1.png': testImageAttachment };
+		model.addAttachment('image/png', imageFilebase64Value, 'test1.png');
+		assert.deepEqual(model.attachments, attachments, 'addAttachment should add unique images');
 	});
 
 	test('addAttachment should not add an invalid attachment to cell', async function () {
@@ -1207,4 +1210,26 @@ suite('Cell Model', function (): void {
 		cellModel.addAttachment('image/png', imageFilebase64Value, 'test.png');
 		assert.equal(cellModel.attachments, undefined);
 	});
+
+	test('addAttachment should not add a duplicate 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.Code,
+			source: '',
+			metadata: {}
+		};
+		let cellModel = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		cellModel.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		assert.deepEqual(cellModel.attachments, attachments);
+		cellModel.addAttachment('image/png', imageFilebase64Value, 'test.png');
+		assert.deepEqual(cellModel.attachments, attachments, 'addAttachment should not add duplicate images');
+	});
 });

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

@@ -34,7 +34,7 @@ import { IInsightOptions } from 'sql/workbench/common/editor/query/chartState';
 
 let modelId = 0;
 const ads_execute_command = 'ads_execute_command';
-const validBase64OctetStreamRegex = /^data:application\/octet-stream;base64,(?:[A-Za-z0-9+\/]{4})*(?:[A-Za-z0-9+\/]{2}==|[A-Za-z0-9+\/]{3}=|[A-Za-z0-9+\/]{4})/;
+const validBase64OctetStreamRegex = /data:(?:(application\/octet-stream|image\/png));base64,(?:[A-Za-z0-9+\/]{4})*(?:[A-Za-z0-9+\/]{2}==|[A-Za-z0-9+\/]{3}=|[A-Za-z0-9+\/]{4})/;
 export interface QueryResultId {
 	batchId: number;
 	id: number;
@@ -145,11 +145,11 @@ export class CellModel extends Disposable implements ICellModel {
 		return this._metadata;
 	}
 
-	public get attachments() {
+	public get attachments(): nb.ICellAttachments | undefined {
 		return this._attachments;
 	}
 
-	addAttachment(mimeType: string, base64Encoding: string, name: string): void {
+	addAttachment(mimeType: string, base64Encoding: string, name: string): string {
 		// base64Encoded value looks like: data:application/octet-stream;base64,<base64Value>
 		// get the <base64Value> from the string
 		let index = base64Encoding.indexOf('base64,');
@@ -160,11 +160,17 @@ export class CellModel extends Disposable implements ICellModel {
 			if (!this._attachments) {
 				this._attachments = {};
 			}
-			// TO DO: Check if name already exists and message the user?
+			// Check if name already exists and get a unique name
+			if (this._attachments[name] && this._attachments[name][mimeType] !== attachment[mimeType]) {
+				name = this.getUniqueAttachmentName(name.substring(0, name.lastIndexOf('.')), name.substring(name.lastIndexOf('.') + 1));
+			}
+			if (!this._attachments[name]) {
 				this._attachments[name] = attachment;
 				this.sendChangeToNotebook(NotebookChangeType.CellMetadataUpdated);
 			}
 		}
+		return name;
+	}
 
 	private isValidBase64OctetStream(base64Image: string): boolean {
 		return base64Image && validBase64OctetStreamRegex.test(base64Image);
@@ -295,6 +301,7 @@ export class CellModel extends Disposable implements ICellModel {
 	}
 
 	public set source(newSource: string | string[]) {
+		newSource = this.attachImageFromSource(newSource);
 		newSource = this.getMultilineSource(newSource);
 		if (this._source !== newSource) {
 			this._source = newSource;
@@ -305,6 +312,35 @@ export class CellModel extends Disposable implements ICellModel {
 		this._preventNextChartCache = true;
 	}
 
+	private attachImageFromSource(newSource: string | string[]): string | string[] {
+		if (!Array.isArray(newSource) && this.isValidBase64OctetStream(newSource)) {
+			let results;
+			while ((results = validBase64OctetStreamRegex.exec(newSource)) !== null) {
+				let imageName = this.addAttachment(results[1], results[0], 'image.png');
+				newSource = newSource.replace(validBase64OctetStreamRegex, `attachment:${imageName}`);
+			}
+			return newSource;
+		}
+		return newSource;
+	}
+	/**
+	 * Gets unique attachment name to add to cell metadata
+	 * @param imgName a string defining name of the image.
+	 * @param imgExtension extension of the image
+	 * Returns the unique name
+	 */
+	private getUniqueAttachmentName(imgName?: string, imgExtension?: string): string {
+		let nextVal = 0;
+		// Note: this will go forever if it's coded wrong, or you have infinite images in a notebook!
+		while (true) {
+			let imageName = imgName ? `${imgName}${nextVal}.${imgExtension ?? 'png'}` : `image${nextVal}.png`;
+			if (!this._attachments || !this._attachments[imageName]) {
+				return imageName;
+			}
+			nextVal++;
+		}
+	}
+
 	public get modelContentChangedEvent(): IModelContentChangedEvent {
 		return this._modelContentChangedEvent;
 	}

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

@@ -535,7 +535,14 @@ export interface ICellModel {
 	readonly savedConnectionName: string | undefined;
 	readonly attachments: nb.ICellAttachments;
 	readonly currentMode: CellEditModes;
-	addAttachment(mimeType: string, base64Encoding: string, name: string): void;
+	/**
+	 * Adds image as an attachment to cell metadata
+	 * @param mimeType a string defining mimeType of the image. Examples: image/png, image/jpeg
+	 * @param base64Encoding the base64 encoded value of the image
+	 * @param name the name of the image.
+	 * Returns the name of the attachment added to metadata.
+	 */
+	addAttachment(mimeType: string, base64Encoding: string, name: string): string;
 }
 
 export interface IModelFactory {