From 7c1710d396ec36a0bbca7f19957ef1a7c1ebe64f Mon Sep 17 00:00:00 2001
From: Karl Burtram <karlb@microsoft.com>
Date: Mon, 8 Oct 2018 12:41:04 -0700
Subject: [PATCH] Reopen connections prior to creating query execution data
 readers (#704)

* Reopen connections prior to creating query execution data readers

* Reopen connection if an exception was rasied executing the query

* Fix unit tests
---
 .../Connection/ConnectionService.cs           | 23 +++++++++++++++++++
 .../QueryExecution/Batch.cs                   |  5 ++++
 .../QueryExecution/Query.cs                   |  9 ++++++++
 .../Execution/ServiceIntegrationTests.cs      |  2 ++
 4 files changed, 39 insertions(+)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
index 215b83b4..f4115a73 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
@@ -1528,5 +1528,28 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             
             return null;
         }
+
+        public static void EnsureConnectionIsOpen(DbConnection conn, bool forceReopen = false)
+        {
+            // verify that the connection is open
+            if (conn.State != ConnectionState.Open || forceReopen)
+            {
+                try
+                {
+                    // close it in case it's in some non-Closed state
+                    conn.Close();
+                }
+                catch
+                {
+                    // ignore any exceptions thrown from .Close
+                    // if the connection is really broken the .Open call will throw
+                }
+                finally
+                {
+                    // try to reopen the connection
+                    conn.Open();
+                }
+            }
+        }
     }
 }
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
index fe2e1f36..83dd650f 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
@@ -17,6 +17,7 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.Utility;
 using System.Globalization;
 using System.Collections.ObjectModel;
+using Microsoft.SqlTools.ServiceLayer.Connection;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
@@ -364,6 +365,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     // so add a newline to the end of the query. See https://github.com/Microsoft/sqlopsstudio/issues/1424
                     dbCommand.CommandText += Environment.NewLine;
 
+                    ConnectionService.EnsureConnectionIsOpen(conn);
+
                     // Fetch schema info separately, since CommandBehavior.KeyInfo will include primary
                     // key columns in the result set, even if they weren't part of the select statement.
                     // Extra key columns get added to the end, so just correlate via Column Ordinal.
@@ -380,6 +383,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     }
                 }
 
+                ConnectionService.EnsureConnectionIsOpen(conn);
+
                 // Execute the command to get back a reader
                 using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(cancellationToken))
                 {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
index 5c49f560..51556b66 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
@@ -436,6 +436,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 {
                     ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, newDatabaseName);
                 }
+
+                foreach (Batch b in Batches)
+                {
+                    if (b.HasError)
+                    {
+                        ConnectionService.EnsureConnectionIsOpen(sqlConn, forceReopen: true);
+                        break;
+                    }
+                }
             }
         }
 
diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs
index a58fc6ca..1501ef52 100644
--- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs
@@ -416,6 +416,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
                 .AddStandardBatchStartValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(1)
+                .AddStandardQueryCompleteValidator(1)
                 .Complete();
 
             await Common.AwaitExecution(queryService, queryParams, efv.Object);
@@ -442,6 +443,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
                 .AddStandardBatchStartValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(1)
+                .AddStandardQueryCompleteValidator(1)
                 .Complete();
             await Common.AwaitExecution(queryService, queryParams, efv.Object);