Object Explorer: Refresh OE node if database restored (#607)

* refresh OE node if restored

* refresh OE when task is completed

* check task id for restore

src/sql/parts/disasterRecovery/restore/common/restoreServiceImpl.ts

@@ -22,6 +22,9 @@ import { MssqlRestoreInfo } from 'sql/parts/disasterRecovery/restore/mssqlRestor
 import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { ProviderConnectionInfo } from 'sql/parts/connection/common/providerConnectionInfo';
 import * as Utils from 'sql/parts/connection/common/utils';
+import { IObjectExplorerService } from 'sql/parts/registeredServer/common/objectExplorerService';
+import { ITaskService } from 'sql/parts/taskHistory/common/taskService';
+import { TaskStatus, TaskNode } from 'sql/parts/taskHistory/common/taskNode';
 
 export class RestoreService implements IRestoreService {
 
@@ -135,13 +138,17 @@ export class RestoreDialogController implements IRestoreDialogController {
 	private _ownerUri: string;
 	private _sessionId: string;
 	private readonly _restoreFeature = 'Restore';
+	private readonly _restoreTaskName: string = 'Restore Database';
+	private readonly _restoreCompleted : string = 'Completed';
 	private _optionValues: { [optionName: string]: any } = {};
 
 	constructor(
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IRestoreService private _restoreService: IRestoreService,
 		@IConnectionManagementService private _connectionService: IConnectionManagementService,
-		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
+		@IObjectExplorerService private _objectExplorerService: IObjectExplorerService,
+		@ITaskService private _taskService: ITaskService
 	) {
 	}
 
@@ -153,9 +160,29 @@ export class RestoreDialogController implements IRestoreDialogController {
 			restoreOption.taskExecutionMode = TaskExecutionMode.executeAndScript;
 		}
 
-		this._restoreService.restore(this._ownerUri, restoreOption);
-		let restoreDialog = this._restoreDialogs[this._currentProvider];
-		restoreDialog.close();
+		this._restoreService.restore(this._ownerUri, restoreOption).then(result => {
+			const self = this;
+			let connectionProfile = self._connectionService.getConnectionProfile(self._ownerUri);
+			let activeNode = self._objectExplorerService.getObjectExplorerNode(connectionProfile);
+			this._taskService.onTaskComplete(response => {
+				if (result.taskId === response.id && this.isSuccessfulRestore(response) && activeNode) {
+					self._objectExplorerService.refreshTreeNode(activeNode.getSession(), activeNode).then(result => {
+						self._objectExplorerService.getServerTreeView().refreshTree();
+					});
+				}
+			});
+			let restoreDialog = this._restoreDialogs[this._currentProvider];
+			restoreDialog.close();
+		});
+	}
+
+	private isSuccessfulRestore(response: TaskNode): boolean {
+		return (response.taskName === this._restoreTaskName &&
+				response.message === this._restoreCompleted &&
+				(response.status === TaskStatus.succeeded ||
+					response.status === TaskStatus.succeededWithWarning) &&
+				(response.taskExecutionMode === TaskExecutionMode.execute ||
+					response.taskExecutionMode === TaskExecutionMode.executeAndScript));
 	}
 
 	private handleMssqlOnValidateFile(overwriteTargetDatabase: boolean = false): void {
@@ -309,7 +336,6 @@ export class RestoreDialogController implements IRestoreDialogController {
 						let restoreDialog = this._restoreDialogs[this._currentProvider] as OptionsDialog;
 						restoreDialog.open(this.getRestoreOption(), this._optionValues);
 					}
-
 					resolve(result);
 				}, error => {
 					reject(error);

src/sql/parts/registeredServer/common/objectExplorerService.ts

@@ -63,6 +63,8 @@ export interface IObjectExplorerService {
 	isFocused(): boolean;
 
 	onSelectionOrFocusChange: Event<void>;
+
+	getServerTreeView(): ServerTreeView;
 }
 
 interface SessionStatus {
@@ -417,4 +419,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	public isFocused(): boolean {
 		return this._serverTreeView.isFocused();
 	}
+
+	public getServerTreeView() {
+		return this._serverTreeView;
+	}
 }

src/sql/parts/taskHistory/common/taskService.ts

@@ -9,13 +9,13 @@ import { TaskNode, TaskStatus, TaskExecutionMode } from 'sql/parts/taskHistory/c
 import { IQueryEditorService } from 'sql/parts/query/common/queryEditorService';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import Event, { Emitter } from 'vs/base/common/event';
-export const SERVICE_ID = 'taskHistoryService';
 import { ILifecycleService } from 'vs/platform/lifecycle/common/lifecycle';
 import { IChoiceService } from 'vs/platform/message/common/message';
 import { localize } from 'vs/nls';
 import Severity from 'vs/base/common/severity';
 import { TPromise } from 'vs/base/common/winjs.base';
 
+export const SERVICE_ID = 'taskHistoryService';
 export const ITaskService = createDecorator<ITaskService>(SERVICE_ID);
 
 export interface ITaskService {