Add queryInfo to query events (#19777)

* Add queryInfo to query events

* docs

src/sql/azdata.d.ts

@@ -4992,14 +4992,18 @@ declare module 'azdata' {
 			| 'executionPlan'
 			| 'visualize';
 
-		/**
-		 * args for each event type
-		 * queryStart: undefined
-		 * queryStop: undefined
-		 * executionPlan: string
-		 * visualize: ResultSetSummary
-		 */
 		export interface QueryEventListener {
+			/**
+			 * A callback that is called whenever a query event occurs
+			 * @param type The type of query event
+			 * @param document The document this event was sent by
+			 * @param args The extra information for the event, if any
+			 * The args sent depend on the type of event :
+			 * queryStart: undefined
+			 * queryStop: undefined
+			 * executionPlan: string (the plan itself)
+			 * visualize: ResultSetSummary (the result set to be visualized)
+			 */
 			onQueryEvent(type: QueryEventType, document: QueryDocument, args: ResultSetSummary | string | undefined): void;
 		}
 

src/sql/azdata.proposed.d.ts

@@ -1525,4 +1525,52 @@ declare module 'azdata' {
 		 */
 		link: LinkArea;
 	}
+
+	export namespace queryeditor {
+
+		export interface IQueryMessage {
+			/**
+			 * The message string
+			 */
+			message: string;
+			/**
+			 * Whether this message is an error message or not
+			 */
+			isError: boolean;
+			/**
+			 * The timestamp for when this message was sent
+			 */
+			time?: string;
+		}
+
+		/**
+		 * Information about a query that was executed
+		 */
+		export interface IQueryInfo {
+			/**
+			 * Any messages that have been received from the query provider
+			 */
+			messages: IQueryMessage[];
+			/**
+			 * The text of the query statement
+			 */
+			queryText?: string;
+		}
+
+		export interface QueryEventListener {
+			/**
+			 * An event that is fired for query events
+			 * @param type The type of query event
+			 * @param document The document this event was sent by
+			 * @param args The extra information for the event, if any
+			 * The args sent depend on the type of event :
+			 * queryStart: undefined
+			 * queryStop: undefined
+			 * executionPlan: string (the plan itself)
+			 * 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;
+		}
+	}
 }

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

@@ -18,6 +18,9 @@ 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 {
@@ -32,7 +35,8 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 		@IEditorService private _editorService: IEditorService,
 		@IQueryManagementService private _queryManagementService: IQueryManagementService,
 		@ILogService private _logService: ILogService,
-		@IQueryEditorService private _queryEditorService: IQueryEditorService
+		@IQueryEditorService private _queryEditorService: IQueryEditorService,
+		@IModelService private _modelService: IModelService
 	) {
 		super();
 		if (extHostContext) {
@@ -118,7 +122,34 @@ 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);
-			this._proxy.$onQueryEvent(connectionProfile?.providerName, handle, event.uri, event);
+			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);
 		});
 		this._queryEventListenerDisposables.set(handle, disposable);
 	}

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

@@ -6,10 +6,10 @@
 import { IMainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostQueryEditorShape, SqlMainContext, MainThreadQueryEditorShape } from 'sql/workbench/api/common/sqlExtHost.protocol';
 import * as azdata from 'azdata';
-import { IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
 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';
 
 class ExtHostQueryDocument implements azdata.queryeditor.QueryDocument {
 	constructor(
@@ -64,11 +64,11 @@ export class ExtHostQueryEditor implements ExtHostQueryEditorShape {
 		});
 	}
 
-	public $onQueryEvent(providerId: string, handle: number, fileUri: string, event: IQueryEvent): void {
+	public $onQueryEvent(providerId: string, handle: number, fileUri: string, event: IExtHostQueryEvent): 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);
+			listener.onQueryEvent(event.type, new ExtHostQueryDocument(providerId, fileUri, this._proxy), params, event.queryInfo);
 		}
 	}
 

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

@@ -23,11 +23,11 @@ import {
 	IModelViewWizardDetails, IModelViewWizardPageDetails, IExecuteManagerDetails, INotebookSessionDetails,
 	INotebookKernelDetails, INotebookFutureDetails, FutureMessageType, INotebookFutureDone, INotebookEditOperation,
 	NotebookChangeKind,
-	ISerializationManagerDetails
+	ISerializationManagerDetails,
+	IExtHostQueryEvent
 } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { IUndoStopOptions } from 'vs/workbench/api/common/extHost.protocol';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
-import { IQueryEvent } from 'sql/workbench/services/query/common/queryModel';
 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';
@@ -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: IQueryEvent): void;
+	$onQueryEvent(providerId: string, handle: number, fileUri: string, event: IExtHostQueryEvent): void;
 }
 
 export interface MainThreadQueryEditorShape extends IDisposable {

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

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { nb, IConnectionProfile } from 'azdata';
+import * as azdata from 'azdata';
 import * as vsExtTypes from 'vs/workbench/api/common/extHostTypes';
 import { URI } from 'vs/base/common/uri';
 
@@ -448,7 +448,7 @@ export enum AzureResource {
 }
 
 export class TreeItem extends vsExtTypes.TreeItem {
-	payload?: IConnectionProfile;
+	payload?: azdata.IConnectionProfile;
 	providerHandle?: string;
 }
 
@@ -613,7 +613,7 @@ export enum FutureMessageType {
 export interface INotebookFutureDone {
 	succeeded: boolean;
 	rejectReason: string;
-	message: nb.IShellMessage;
+	message: azdata.nb.IShellMessage;
 }
 
 export interface ICellRange {
@@ -686,7 +686,7 @@ export interface INotebookEditOperation {
 	/**
 	 * The cell metadata to use for the edit operation (only for some edit operations)
 	 */
-	cell: Partial<nb.ICellContents>;
+	cell: Partial<azdata.nb.ICellContents>;
 	/**
 	 * Whether to append the content to the existing content or replace it.
 	 */
@@ -1057,3 +1057,14 @@ 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;
+}