Fix/restore connection issue (#416)

* fixed the connection issue used for restore

* fixed a test

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);
             }
         }
 

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;

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);