Switch query events to just send ranges instead of full query text (#19823)

* Switch query events to just send ranges instead of full query text

* undo azdata changes

* fix type

* comment + remove unneeded ?

extensions/query-history/src/queryHistoryProvider.ts

@@ -5,6 +5,7 @@
 
 import * as vscode from 'vscode';
 import * as azdata from 'azdata';
+import { EOL } from 'os';
 import { QueryHistoryNode } from './queryHistoryNode';
 
 const QUERY_HISTORY_CONFIG_SECTION = 'queryHistory';
@@ -23,9 +24,19 @@ export class QueryHistoryProvider implements vscode.TreeDataProvider<QueryHistor
 
 	constructor() {
 		this._disposables.push(azdata.queryeditor.registerQueryEventListener({
-			onQueryEvent: async (type: azdata.queryeditor.QueryEventType, document: azdata.queryeditor.QueryDocument, args: azdata.ResultSetSummary | string | undefined, queryInfo?: azdata.queryeditor.IQueryInfo) => {
+			onQueryEvent: async (type: azdata.queryeditor.QueryEventType, document: azdata.queryeditor.QueryDocument, args: azdata.ResultSetSummary | string | undefined, queryInfo?: azdata.queryeditor.QueryInfo) => {
 				if (this._captureEnabled && queryInfo && type === 'queryStop') {
-					const queryText = queryInfo.queryText ?? '';
+					const textDocuments = vscode.workspace.textDocuments;
+					// We need to compare URIs, but the event Uri comes in as string so while it should be in the same format as
+					// the textDocument uri.toString() we parse it into a vscode.Uri first to be absolutely sure.
+					const textDocument = textDocuments.find(e => e.uri.toString() === vscode.Uri.parse(document.uri).toString());
+					if (!textDocument) {
+						// If we couldn't find the document then we can't get the text so just log the error and move on
+						console.error(`Couldn't find text document with URI ${document.uri} for query event`);
+						return;
+					}
+					// Combine all the text from the batches back together
+					const queryText = queryInfo.range.map(r => textDocument.getText(r) ?? '').join(EOL);
 					const connProfile = await azdata.connection.getConnection(document.uri);
 					const isError = queryInfo.messages.find(m => m.isError) ? false : true;
 					// Add to the front of the list so the new item appears at the top

src/sql/azdata.proposed.d.ts

@@ -1542,7 +1542,7 @@ declare module 'azdata' {
 	}
 	export namespace queryeditor {
 
-		export interface IQueryMessage {
+		export interface QueryMessage {
 			/**
 			 * The message string
 			 */
@@ -1560,15 +1560,15 @@ declare module 'azdata' {
 		/**
 		 * Information about a query that was executed
 		 */
-		export interface IQueryInfo {
+		export interface QueryInfo {
 			/**
 			 * Any messages that have been received from the query provider
 			 */
-			messages: IQueryMessage[];
+			messages: QueryMessage[];
 			/**
-			 * The text of the query statement
+			 * The ranges for each batch that has executed so far
 			 */
-			queryText?: string;
+			range: vscode.Range[];
 		}
 
 		export interface QueryEventListener {
@@ -1584,7 +1584,7 @@ declare module 'azdata' {
 			 * visualize: ResultSetSummary (the result set to be visualized)
 			 * @param queryInfo The information about the query that triggered this event
 			 */
-			onQueryEvent(type: QueryEventType, document: QueryDocument, args: ResultSetSummary | string | undefined, queryInfo: IQueryInfo): void;
+			onQueryEvent(type: QueryEventType, document: QueryDocument, args: ResultSetSummary | string | undefined, queryInfo: QueryInfo): void;
 		}
 	}
 }

src/sql/workbench/api/browser/mainThreadQueryEditor.ts

@@ -18,9 +18,6 @@ import { ConnectionProfile } from 'sql/platform/connection/common/connectionProf
 import { ILogService } from 'vs/platform/log/common/log';
 import { URI } from 'vs/base/common/uri';
 import { IQueryEditorService } from 'sql/workbench/services/queryEditor/common/queryEditorService';
-import { IModelService } from 'vs/editor/common/services/modelService';
-import { Range } from 'vs/editor/common/core/range';
-import { IExtHostQueryEvent } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 @extHostNamedCustomer(SqlMainContext.MainThreadQueryEditor)
 export class MainThreadQueryEditor extends Disposable implements MainThreadQueryEditorShape {
@@ -35,8 +32,7 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 		@IEditorService private _editorService: IEditorService,
 		@IQueryManagementService private _queryManagementService: IQueryManagementService,
 		@ILogService private _logService: ILogService,
-		@IQueryEditorService private _queryEditorService: IQueryEditorService,
-		@IModelService private _modelService: IModelService
+		@IQueryEditorService private _queryEditorService: IQueryEditorService
 	) {
 		super();
 		if (extHostContext) {
@@ -121,35 +117,8 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 
 	public $registerQueryInfoListener(handle: number): void {
 		const disposable = this._queryModelService.onQueryEvent(event => {
-			let connectionProfile = this._connectionManagementService.getConnectionProfile(event.uri);
-			const uri: URI = URI.parse(event.uri);
-			const model = this._modelService.getModel(uri);
-			// Get the query text from the model - we do it here so we can send the query text for all events
-			// to the extension host from one place
-			let queryText: string | undefined = undefined;
-			if (model) {
-				// VS Range is 1 based so offset values by 1. The endLine we get back from SqlToolsService is incremented
-				// by 1 from the original input range sent in as well so take that into account and don't modify
-				queryText = event.queryInfo.range.length > 0 ?
-					model.getValueInRange(new Range(
-						event.queryInfo.range[0].startLineNumber,
-						event.queryInfo.range[0].startColumn,
-						event.queryInfo.range[0].endLineNumber,
-						event.queryInfo.range[0].endColumn)) :
-					// If no specific selection get the entire text
-					model.getValue();
-			}
-			// Convert into an IExtHostQueryEvent with the properties it expects
-			const extHostEvent: IExtHostQueryEvent = {
-				type: event.type,
-				uri: event.uri,
-				params: event.params,
-				queryInfo: {
-					messages: event.queryInfo.messages,
-					queryText
-				}
-			};
-			this._proxy.$onQueryEvent(connectionProfile?.providerName, handle, event.uri, extHostEvent);
+			const connectionProfile = this._connectionManagementService.getConnectionProfile(event.uri);
+			this._proxy.$onQueryEvent(connectionProfile?.providerName, handle, event.uri, event);
 		});
 		this._queryEventListenerDisposables.set(handle, disposable);
 	}

src/sql/workbench/api/common/extHostQueryEditor.ts

@@ -9,7 +9,8 @@ import * as azdata from 'azdata';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { Disposable } from 'vs/workbench/api/common/extHostTypes';
 import { URI } from 'vs/base/common/uri';
-import { IExtHostQueryEvent } from 'sql/workbench/api/common/sqlExtHostTypes';
+import { IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
+import * as sqlTypeConverters from 'sql/workbench/api/common/sqlExtHostTypeConverters';
 
 class ExtHostQueryDocument implements azdata.queryeditor.QueryDocument {
 	constructor(
@@ -64,11 +65,11 @@ export class ExtHostQueryEditor implements ExtHostQueryEditorShape {
 		});
 	}
 
-	public $onQueryEvent(providerId: string, handle: number, fileUri: string, event: IExtHostQueryEvent): void {
+	public $onQueryEvent(providerId: string, handle: number, fileUri: string, event: IQueryEvent): void {
 		let listener: azdata.queryeditor.QueryEventListener = this._queryListeners[handle];
 		if (listener) {
 			let params = event.params && event.params.planXml ? event.params.planXml : event.params;
-			listener.onQueryEvent(event.type, new ExtHostQueryDocument(providerId, fileUri, this._proxy), params, event.queryInfo);
+			listener.onQueryEvent(event.type, new ExtHostQueryDocument(providerId, fileUri, this._proxy), params, sqlTypeConverters.QueryInfo.to(event.queryInfo));
 		}
 	}
 

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -23,14 +23,14 @@ import {
 	IModelViewWizardDetails, IModelViewWizardPageDetails, IExecuteManagerDetails, INotebookSessionDetails,
 	INotebookKernelDetails, INotebookFutureDetails, FutureMessageType, INotebookFutureDone, INotebookEditOperation,
 	NotebookChangeKind,
-	ISerializationManagerDetails,
-	IExtHostQueryEvent
+	ISerializationManagerDetails
 } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { IUndoStopOptions } from 'vs/workbench/api/common/extHost.protocol';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { EditorViewColumn } from 'vs/workbench/api/common/shared/editor';
 import { TreeDataTransferDTO } from 'vs/workbench/api/common/shared/treeDataTransfer';
 import { ITelemetryEventProperties } from 'sql/platform/telemetry/common/telemetry';
+import { IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
 
 export abstract class ExtHostAzureBlobShape {
 	public $createSas(connectionUri: string, blobContainerUri: string, blobStorageKey: string, storageAccountName: string, expirationDate: string): Thenable<mssql.CreateSasResponse> { throw ni(); }
@@ -921,7 +921,7 @@ export interface MainThreadModelViewDialogShape extends IDisposable {
 	$setDirty(handle: number, isDirty: boolean): void;
 }
 export interface ExtHostQueryEditorShape {
-	$onQueryEvent(providerId: string, handle: number, fileUri: string, event: IExtHostQueryEvent): void;
+	$onQueryEvent(providerId: string, handle: number, fileUri: string, event: IQueryEvent): void;
 }
 
 export interface MainThreadQueryEditorShape extends IDisposable {

src/sql/workbench/api/common/sqlExtHostTypeConverters.ts

@@ -0,0 +1,21 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type * as azdata from 'azdata';
+import { IQueryInfo } from 'sql/workbench/services/query/common/queryModel';
+import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
+
+export namespace QueryInfo {
+
+	export function to(queryInfo: IQueryInfo | undefined): azdata.queryeditor.QueryInfo | undefined {
+		if (!queryInfo) {
+			return undefined;
+		}
+		return {
+			messages: queryInfo.messages,
+			range: queryInfo.range.map(r => typeConverters.Range.to(r))
+		};
+	}
+}

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -1057,14 +1057,3 @@ export namespace executionPlan {
 		None = 4
 	}
 }
-
-/**
- * Query event info to send to the extension host. This is a separate type from IQueryEvent
- * since this has the query text ranges converted into the actual query text.
- */
-export interface IExtHostQueryEvent {
-	type: azdata.queryeditor.QueryEventType;
-	uri: string;
-	queryInfo: azdata.queryeditor.IQueryInfo;
-	params?: any;
-}

src/sql/workbench/services/query/common/queryModelService.ts

@@ -35,23 +35,23 @@ export interface QueryEvent {
 export class QueryInfo {
 	public queryRunner?: EditQueryRunner;
 	public dataService?: DataService;
-	public queryEventQueue?: QueryEvent[];
-	public range?: Array<IRange>;
+	public queryEventQueue: QueryEvent[] = [];
+	public range: Array<IRange> = [];
 	public selectionSnippet?: string;
 
 	// Notes if the angular components have obtained the DataService. If not, all messages sent
 	// via the data service will be lost.
-	public dataServiceReady?: boolean;
+	public dataServiceReady: boolean = false;
 
-	constructor() {
-		this.dataServiceReady = false;
-		this.queryEventQueue = [];
-		this.range = [];
-	}
+	constructor() { }
 
 	public set uri(newUri: string) {
-		this.queryRunner.uri = newUri;
-		this.dataService.uri = newUri;
+		if (this.queryRunner) {
+			this.queryRunner.uri = newUri;
+		}
+		if (this.dataService) {
+			this.dataService.uri = newUri;
+		}
 	}
 }
 
@@ -271,7 +271,7 @@ export class QueryModelService implements IQueryModelService {
 						text: strings.format(nls.localize('runQueryBatchStartLine', "Line {0}"), b.range.startLineNumber)
 					};
 				}
-				info.range!.push(b.range);
+				info.range.push(b.range);
 			}
 			let message = {
 				message: messageText,
@@ -293,7 +293,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: queryRunner.uri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -311,7 +311,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: queryRunner.uri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -327,7 +327,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: queryRunner.uri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				}
 			};
@@ -343,7 +343,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: planInfo.fileUri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				},
 				params: planInfo
@@ -358,7 +358,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: qp2Info.fileUri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				},
 				params: qp2Info.planGraphs
@@ -372,7 +372,7 @@ export class QueryModelService implements IQueryModelService {
 				uri: queryRunner.uri,
 				queryInfo:
 				{
-					range: info.range!,
+					range: info.range,
 					messages: info.queryRunner!.messages
 				},
 				params: resultSetInfo
@@ -514,7 +514,7 @@ export class QueryModelService implements IQueryModelService {
 					uri: ownerUri,
 					queryInfo:
 					{
-						range: info.range!,
+						range: info.range,
 						messages: info.queryRunner!.messages
 					},
 				};
@@ -531,7 +531,7 @@ export class QueryModelService implements IQueryModelService {
 					uri: ownerUri,
 					queryInfo:
 					{
-						range: info.range!,
+						range: info.range,
 						messages: info.queryRunner!.messages
 					},
 				};