From ae178efe3e090a6d822017729463c64d2996e5ba Mon Sep 17 00:00:00 2001
From: Benjamin Russell <benrr101@outlook.com>
Date: Fri, 31 Mar 2017 13:59:44 -0700
Subject: [PATCH] Fixing misc edit bugs (async/await, table not found) (#300)

* Async/await issue when requesting a edit/subset
* Returning good error messages when table/view doesn't exist
---
 .../EditData/EditDataService.cs                | 18 +++++++++++++-----
 .../EditData/SmoEditMetadataFactory.cs         |  4 ++++
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs
index 464b68ae..23700cdb 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs
@@ -198,18 +198,26 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             });
         }
 
-        internal Task HandleSubsetRequest(EditSubsetParams subsetParams,
+        internal async Task HandleSubsetRequest(EditSubsetParams subsetParams,
             RequestContext<EditSubsetResult> requestContext)
         {
-            return HandleSessionRequest(subsetParams, requestContext, session =>
+            try
             {
-                EditRow[] rows = session.GetRows(subsetParams.RowStartIndex, subsetParams.RowCount).Result;
-                return new EditSubsetResult
+                EditSession session = GetActiveSessionOrThrow(subsetParams.OwnerUri);
+
+                EditRow[] rows = await session.GetRows(subsetParams.RowStartIndex, subsetParams.RowCount);
+                EditSubsetResult result = new EditSubsetResult
                 {
                     RowCount = rows.Length,
                     Subset = rows
                 };
-            });
+
+                await requestContext.SendResult(result);
+            }
+            catch(Exception e)
+            {
+                await requestContext.SendError(e.Message);
+            }
         }
 
         internal Task HandleUpdateCellRequest(EditUpdateCellParams updateParams,
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/SmoEditMetadataFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/SmoEditMetadataFactory.cs
index f51a8320..d9834cb6 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/EditData/SmoEditMetadataFactory.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/EditData/SmoEditMetadataFactory.cs
@@ -77,6 +77,10 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
             // A bug in SMO makes it necessary to call refresh to attain certain properties (such as IsMemoryOptimized)
             smoResult.Refresh();
+            if (smoResult.State != SqlSmoState.Existing)
+            {
+                throw new ArgumentOutOfRangeException(nameof(objectNamedParts), SR.EditDataObjectNotFound);
+            }
             
             // Generate the edit column metadata
             List<EditColumnMetadata> editColumns = new List<EditColumnMetadata>();