From b5fdf8c2a72dfacfb527e24e39cbc1903042249d Mon Sep 17 00:00:00 2001
From: Matthew McCune <matthew@matthew.cx>
Date: Tue, 13 Aug 2019 18:31:34 -0700
Subject: [PATCH] Update row count as updates are received (#6642)

---
 src/sql/azdata.d.ts                                |  1 +
 src/sql/platform/query/common/queryModel.ts        |  1 +
 src/sql/platform/query/common/queryModelService.ts | 14 ++++++++++++++
 .../parts/query/browser/statusBarItems.ts          |  8 ++++++++
 4 files changed, 24 insertions(+)

diff --git a/src/sql/azdata.d.ts b/src/sql/azdata.d.ts
index 18989d3315..077ea67812 100644
--- a/src/sql/azdata.d.ts
+++ b/src/sql/azdata.d.ts
@@ -3994,6 +3994,7 @@ declare module 'azdata' {
 	export namespace queryeditor {
 		export type QueryEvent =
 			| 'queryStart'
+			| 'queryUpdate'
 			| 'queryStop'
 			| 'executionPlan'
 			| 'visualize';
diff --git a/src/sql/platform/query/common/queryModel.ts b/src/sql/platform/query/common/queryModel.ts
index 65d7d00f2b..d64bc22758 100644
--- a/src/sql/platform/query/common/queryModel.ts
+++ b/src/sql/platform/query/common/queryModel.ts
@@ -68,6 +68,7 @@ export interface IQueryModelService {
 	showCommitError(error: string): void;
 
 	onRunQueryStart: Event<string>;
+	onRunQueryUpdate: Event<string>;
 	onRunQueryComplete: Event<string>;
 	onQueryEvent: Event<IQueryEvent>;
 
diff --git a/src/sql/platform/query/common/queryModelService.ts b/src/sql/platform/query/common/queryModelService.ts
index 7f58af9877..08d5b9183e 100644
--- a/src/sql/platform/query/common/queryModelService.ts
+++ b/src/sql/platform/query/common/queryModelService.ts
@@ -58,12 +58,14 @@ export class QueryModelService implements IQueryModelService {
 	// MEMBER VARIABLES ////////////////////////////////////////////////////
 	private _queryInfoMap: Map<string, QueryInfo>;
 	private _onRunQueryStart: Emitter<string>;
+	private _onRunQueryUpdate: Emitter<string>;
 	private _onRunQueryComplete: Emitter<string>;
 	private _onQueryEvent: Emitter<IQueryEvent>;
 	private _onEditSessionReady: Emitter<azdata.EditSessionReadyParams>;
 
 	// EVENTS /////////////////////////////////////////////////////////////
 	public get onRunQueryStart(): Event<string> { return this._onRunQueryStart.event; }
+	public get onRunQueryUpdate(): Event<string> { return this._onRunQueryUpdate.event; }
 	public get onRunQueryComplete(): Event<string> { return this._onRunQueryComplete.event; }
 	public get onQueryEvent(): Event<IQueryEvent> { return this._onQueryEvent.event; }
 	public get onEditSessionReady(): Event<azdata.EditSessionReadyParams> { return this._onEditSessionReady.event; }
@@ -75,6 +77,7 @@ export class QueryModelService implements IQueryModelService {
 	) {
 		this._queryInfoMap = new Map<string, QueryInfo>();
 		this._onRunQueryStart = new Emitter<string>();
+		this._onRunQueryUpdate = new Emitter<string>();
 		this._onRunQueryComplete = new Emitter<string>();
 		this._onQueryEvent = new Emitter<IQueryEvent>();
 		this._onEditSessionReady = new Emitter<azdata.EditSessionReadyParams>();
@@ -299,6 +302,17 @@ export class QueryModelService implements IQueryModelService {
 
 			this._fireQueryEvent(uri, 'start');
 		});
+		queryRunner.onResultSetUpdate(() => {
+			this._onRunQueryUpdate.fire(uri);
+
+			let event: IQueryEvent = {
+				type: 'queryUpdate',
+				uri: uri
+			};
+			this._onQueryEvent.fire(event);
+
+			this._fireQueryEvent(uri, 'update');
+		});
 
 		queryRunner.onQueryPlanAvailable(planInfo => {
 			// fire extensibility API event
diff --git a/src/sql/workbench/parts/query/browser/statusBarItems.ts b/src/sql/workbench/parts/query/browser/statusBarItems.ts
index 7e5ea34409..b41ab6042f 100644
--- a/src/sql/workbench/parts/query/browser/statusBarItems.ts
+++ b/src/sql/workbench/parts/query/browser/statusBarItems.ts
@@ -158,6 +158,9 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 				this.disposable.add(runner.onQueryStart(e => {
 					this._displayValue(runner);
 				}));
+				this.disposable.add(runner.onResultSetUpdate(e => {
+					this._displayValue(runner);
+				}));
 				this.disposable.add(runner.onQueryEnd(e => {
 					this._displayValue(runner);
 				}));
@@ -167,6 +170,11 @@ export class RowCountStatusBarContributions extends Disposable implements IWorkb
 						this._displayValue(this.queryModelService.getQueryRunner(uri));
 					}
 				}));
+				this.disposable.add(this.queryModelService.onRunQueryUpdate(e => {
+					if (e === uri) {
+						this._displayValue(this.queryModelService.getQueryRunner(uri));
+					}
+				}));
 				this.disposable.add(this.queryModelService.onRunQueryComplete(e => {
 					if (e === uri) {
 						this._displayValue(this.queryModelService.getQueryRunner(uri));