From d5b2bcdcb797c2a6753db2e34e149c081a0648d7 Mon Sep 17 00:00:00 2001
From: Leila Lali <llali@microsoft.com>
Date: Mon, 17 Jul 2017 15:57:30 -0700
Subject: [PATCH] Fix/restore connection issue (#416)

* fixed the connection issue used for restore

* fixed a test
---
 .../DisasterRecoveryService.cs                | 14 +++++++----
 .../RestoreOperation/RestoreDatabaseHelper.cs | 23 ++++++++++++++++++-
 .../RestoreDatabaseServiceTests.cs            |  3 ++-
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
index 0f25faa4..d2006a61 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/DisasterRecoveryService.cs
@@ -123,9 +123,10 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             RestoreParams restoreParams,
             RequestContext<RestorePlanResponse> requestContext)
         {
+            RestorePlanResponse response = new RestorePlanResponse();
+
             try
             {
-                RestorePlanResponse response = new RestorePlanResponse();
                 ConnectionInfo connInfo;
                 bool supported = IsBackupRestoreOperationSupported(restoreParams, out connInfo);
 
@@ -143,7 +144,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             }
             catch (Exception ex)
             {
-                await requestContext.SendError(ex.ToString());
+                response.CanRestore = false;
+                response.ErrorMessage = ex.Message;
+                await requestContext.SendResult(response);
             }
         }
 
@@ -154,9 +157,10 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             RestoreParams restoreParams,
             RequestContext<RestoreResponse> requestContext)
         {
+            RestoreResponse response = new RestoreResponse();
+
             try
             {
-                RestoreResponse response = new RestoreResponse();
                 ConnectionInfo connInfo;
                 bool supported = IsBackupRestoreOperationSupported(restoreParams, out connInfo);
 
@@ -199,7 +203,9 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery
             }
             catch (Exception ex)
             {
-                await requestContext.SendError(ex.ToString());
+                response.Result = false;
+                response.ErrorMessage = ex.Message;
+                await requestContext.SendResult(response);
             }
         }
 
diff --git a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs
index 2c633cb4..e4322c5d 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/DisasterRecovery/RestoreOperation/RestoreDatabaseHelper.cs
@@ -5,12 +5,16 @@
 
 using System;
 using System.Linq;
+using System.Data.Common;
+using System.Data.SqlClient;
 using System.Threading.Tasks;
 using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.Smo;
 using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 using Microsoft.SqlTools.ServiceLayer.DisasterRecovery.Contracts;
 using Microsoft.SqlTools.ServiceLayer.TaskServices;
+using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 {
@@ -219,7 +223,24 @@ namespace Microsoft.SqlTools.ServiceLayer.DisasterRecovery.RestoreOperation
 
             if (connInfo != null)
             {
-                Server server = new Server(new ServerConnection(connInfo.ConnectionDetails.ServerName));
+                SqlConnection connection;
+                DbConnection dbConnection = connInfo.AllConnections.First();
+                ReliableSqlConnection reliableSqlConnection = dbConnection as ReliableSqlConnection;
+                SqlConnection sqlConnection = dbConnection as SqlConnection;
+                if (reliableSqlConnection != null)
+                {
+                    connection = reliableSqlConnection.GetUnderlyingConnection();
+                }
+                else if (sqlConnection != null)
+                {
+                    connection = sqlConnection;
+                }
+                else
+                {
+                    Logger.Write(LogLevel.Warning, "Cannot find any sql connection for restore operation");
+                    return null;
+                }
+                Server server = new Server(new ServerConnection(connection));
 
                 RestoreDatabaseTaskDataObject restoreDataObject = new RestoreDatabaseTaskDataObject(server, restoreParams.DatabaseName);
                 restoreDataObject.RestoreParams = restoreParams;
diff --git a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs
index b9483d4a..4c44be41 100644
--- a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs
@@ -185,11 +185,12 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 if (canRestore)
                 {
                     Assert.True(response.DbFiles.Any());
-                    Assert.Equal(response.DatabaseName, "BackupTestDb");
                     if (string.IsNullOrEmpty(targetDatabase))
                     {
                         targetDatabase = response.DatabaseName;
                     }
+                    Assert.Equal(response.DatabaseName, targetDatabase);
+                   
                     if(execute)
                     {
                         await DropDatabase(targetDatabase);