Use column names as keys for table data in SQL cell outputs (#22688)

src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts

@@ -11,7 +11,7 @@ import { IContextMenuService, IContextViewService } from 'vs/platform/contextvie
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-import { IDataResource } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
+import { IDataResource, rowHasColumnNameKeys } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
 import { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner';
 import { ResultSetSummary, ResultSetSubset, ICellValue } from 'sql/workbench/services/query/common/query';
 import { INotificationService } from 'vs/platform/notification/common/notification';
@@ -37,7 +37,6 @@ import { IInsightOptions } from 'sql/workbench/common/editor/query/chartState';
 import { NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
 import { IQueryModelService } from 'sql/workbench/services/query/common/queryModel';
 import { ActionsOrientation } from 'vs/base/browser/ui/actionbar/actionbar';
-import { values } from 'vs/base/common/collections';
 import { URI } from 'vs/base/common/uri';
 import { QueryResultId } from 'sql/workbench/services/notebook/browser/models/cell';
 import { equals } from 'vs/base/common/arrays';
@@ -195,7 +194,7 @@ function reorderGridData(source: IDataResource): void {
 	if (source.data.length > 0) {
 		let rowKeys = Object.keys(source.data[0]);
 		if (!equals(columnNames, rowKeys)) {
-			// SQL notebooks do not use columnName as key (instead use indices)
+			// Older SQL notebooks use indices as keys instead of the column name.
 			// Indicies indicate the row is ordered properly
 			// We must check the data to know if it is in index form
 			let notIndexOrderKeys = false;
@@ -371,19 +370,27 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	private transformSource(source: IDataResource): void {
-		this._rows = source.data.map(row => {
-			let rowData: azdata.DbCellValue[] = [];
-			Object.keys(row).forEach((val, index) => {
-				let displayValue = String(values(row)[index]);
-				// Since the columns[0] represents the row number, start at 1
-				rowData.push({
-					displayValue: displayValue,
-					isNull: false,
-					invariantCultureDisplayValue: displayValue
-				});
+		if (source.data.length > 0) {
+			let columns = source.schema.fields;
+			// Rows are either indexed by column name or ordinal number, so check for one column name to see if it uses that format
+			let useColumnNameKey = rowHasColumnNameKeys(source.data[0], source.schema.fields.map(field => field.name));
+			this._rows = source.data.map(row => {
+				let rowData: azdata.DbCellValue[] = [];
+				for (let index = 0; index < Object.keys(row).length; index++) {
+					let key = useColumnNameKey ? columns[index].name : index;
+					let displayValue = String(row[key]);
+					// Since the columns[0] represents the row number, start at 1
+					rowData.push({
+						displayValue: displayValue,
+						isNull: false,
+						invariantCultureDisplayValue: displayValue
+					});
+				}
+				return rowData;
 			});
-			return rowData;
-		});
+		} else {
+			this._rows = [];
+		}
 	}
 
 	public updateResultSet(resultSet: ResultSetSummary, rows: ICellValue[][]): void {

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

@@ -57,19 +57,6 @@ suite('Data Resource Data Provider', function () {
 	let cellModel = TypeMoq.Mock.ofType(CellModel);
 	let configurationService = new TestConfigurationService();
 
-	// Create test data with two rows and two columns
-	let source: IDataResource = {
-		data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
-		schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
-	};
-	let resultSet: ResultSetSummary = {
-		batchId: 0,
-		columnInfo: [{ columnName: 'col1' }, { columnName: 'col2' }],
-		complete: true,
-		id: 0,
-		rowCount: 2
-	};
-
 	suiteSetup(async () => {
 		let notebookModel = await createandLoadNotebookModel();
 		cellModel.setup(x => x.notebookModel).returns(() => notebookModel);
@@ -98,6 +85,33 @@ suite('Data Resource Data Provider', function () {
 	});
 
 	test('serializeResults call is successful', async function (): Promise<void> {
+		// Create test data with two rows and two columns
+		let source: IDataResource = {
+			data: [{ 'col1': '1', 'col2': '2' }, { 'col1': '3', 'col2': '4' }],
+			schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
+		};
+
+		await runSerializeTest(source);
+	});
+
+	test('serializeResults call is successful with ordinal keys for source data', async function (): Promise<void> {
+		// Create test data with two rows and two columns
+		let source: IDataResource = {
+			data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
+			schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
+		};
+
+		await runSerializeTest(source);
+	});
+
+	async function runSerializeTest(source: IDataResource): Promise<void> {
+		let resultSet: ResultSetSummary = {
+			batchId: 0,
+			columnInfo: [{ columnName: 'col1' }, { columnName: 'col2' }],
+			complete: true,
+			id: 0,
+			rowCount: 2
+		};
 		let tempFolderPath = path.join(os.tmpdir(), `TestDataResourceDataProvider_${uuid.v4()}`);
 		await fs.mkdir(tempFolderPath);
 		let dataResourceDataProvider = new DataResourceDataProvider(
@@ -140,5 +154,5 @@ suite('Data Resource Data Provider', function () {
 
 		const withHeadersResult = await fs.readFile(withHeadersFile.fsPath);
 		assert.strictEqual(withHeadersResult.toString(), 'col1 col2\n1 2\n3 4\n', 'result data should include headers');
-	});
+	}
 });

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

@@ -37,7 +37,7 @@ suite('SQL Future', function () {
 		};
 		const expectedData = {
 			'application/vnd.dataresource+json': {
-				data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
+				data: [{ 'col1': '1', 'col2': '2' }, { 'col1': '3', 'col2': '4' }],
 				schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
 			},
 			'text/html': ['<table>', '<tr><th>col1</th><th>col2</th></tr>', '<tr><td>1</td><td>2</td></tr>', '<tr><td>3</td><td>4</td></tr>', '</table>']

src/sql/workbench/services/notebook/browser/outputs/tableRenderers.ts

@@ -8,14 +8,13 @@ import { Table } from 'sql/base/browser/ui/table/table';
 import { textFormatter } from 'sql/base/browser/ui/table/formatters';
 import { RowNumberColumn } from 'sql/base/browser/ui/table/plugins/rowNumberColumn.plugin';
 import { escape } from 'sql/base/common/strings';
-import { IDataResource } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
+import { IDataResource, IDataResourceRow, rowHasColumnNameKeys } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
 import { attachTableStyler } from 'sql/platform/theme/common/styler';
 import { IThemeService } from 'vs/platform/theme/common/themeService';
 import { MouseWheelSupport } from 'sql/base/browser/ui/table/plugins/mousewheelTableScroll.plugin';
 import { AutoColumnSize } from 'sql/base/browser/ui/table/plugins/autoSizeColumns.plugin';
 import { AdditionalKeyBindings } from 'sql/base/browser/ui/table/plugins/additionalKeyBindings.plugin';
 import { RESULTS_GRID_DEFAULTS } from 'sql/workbench/common/constants';
-import { values } from 'vs/base/common/collections';
 import { IAccessibilityService } from 'vs/platform/accessibility/common/accessibility';
 import { IQuickInputService } from 'vs/platform/quickinput/common/quickInput';
 
@@ -88,13 +87,17 @@ export function renderDataResource(
 }
 
 // SlickGrid requires columns and data to be in a very specific format; this code was adapted from tableInsight.component.ts
-export function transformData(rows: any[], columns: Slick.Column<any>[]): { [key: string]: string }[] {
+function transformData(rows: IDataResourceRow[], columns: Slick.Column<any>[]): IDataResourceRow[] {
+	// Rows are either indexed by column name or ordinal number, so check for one column name to see if it uses that format
+	let useColumnNameKey = rowHasColumnNameKeys(rows[0], columns.map(column => column.name));
 	return rows.map(row => {
 		let dataWithSchema = {};
 		Object.keys(row).forEach((val, index) => {
-			let displayValue = String(values(row)[index]);
 			// Since the columns[0] represents the row number, start at 1
-			dataWithSchema[columns[index + 1].field] = {
+			let columnName = columns[index + 1].field;
+			let sourceKey = useColumnNameKey ? columnName : index + 1;
+			let displayValue = String(row[sourceKey]);
+			dataWithSchema[columnName] = {
 				displayValue: displayValue,
 				ariaLabel: escape(displayValue),
 				isNull: false
@@ -104,7 +107,7 @@ export function transformData(rows: any[], columns: Slick.Column<any>[]): { [key
 	});
 }
 
-export function transformColumns(columns: string[]): Slick.Column<any>[] {
+function transformColumns(columns: string[]): Slick.Column<any>[] {
 	return columns.map((col, index) => {
 		return <Slick.Column<any>>{
 			name: col,

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -608,7 +608,7 @@ export class SQLFuture extends Disposable implements FutureInternal {
 			this._rowsMap.set(key, rows.concat(queryResult.rows));
 
 			// Convert rows to data resource and html and send to cell model to be saved
-			let dataResourceRows = this.convertRowsToDataResource(queryResult.rows);
+			let dataResourceRows = this.convertRowsToDataResource(resultSet.columnInfo, queryResult.rows);
 			let saveData = this._dataToSaveMap.get(key);
 			saveData['application/vnd.dataresource+json'].data = saveData['application/vnd.dataresource+json'].data.concat(dataResourceRows);
 			let htmlRows = this.convertRowsToHtml(queryResult.rows, key);
@@ -702,11 +702,12 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		return htmlTable;
 	}
 
-	private convertRowsToDataResource(rows: ICellValue[][]): any[] {
+	private convertRowsToDataResource(columns: IColumn[], rows: ICellValue[][]): IDataResourceRow[] {
 		return rows.map(row => {
-			let rowObject: { [key: string]: any; } = {};
+			let rowObject = {};
 			row.forEach((val, index) => {
-				rowObject[index] = val.displayValue;
+				let columnName = columns[index].columnName;
+				rowObject[columnName] = val.displayValue;
 			});
 			return rowObject;
 		});
@@ -775,7 +776,7 @@ export class SQLFuture extends Disposable implements FutureInternal {
 
 export interface IDataResource {
 	schema: IDataResourceFields;
-	data: any[];
+	data: IDataResourceRow[];
 }
 
 export interface IDataResourceFields {
@@ -787,6 +788,20 @@ export interface IDataResourceSchema {
 	type?: string;
 }
 
+export interface IDataResourceRow {
+	[key: string]: any;
+}
+
+/**
+ * Determines whether a row from a query result set uses column name keys to access its cell data, rather than ordinal number.
+ * @param row The data row to inspect.
+ * @param columnNames The array of column names from the result set's column schema. Column names can be in any order.
+ */
+export function rowHasColumnNameKeys(row: IDataResourceRow, columnNames: string[]): boolean {
+	let columnNameSet = new Set(columnNames);
+	return Object.keys(row).every(rowKey => columnNameSet.has(rowKey));
+}
+
 class ExternalScriptMagic {
 
 	constructor(private language: string) {