Adds a execute and return result message (#383)

* inital request

* refactored query execution failure callback to take exception

* added failure callback to execute and return

* added test for query execute and return

* updated params

* removed dead code

* addressed feedback; added multiple active result set support; updated tests

* addessed feedback and added testing and errors and verification

* change <= to ==

* changed name of trashQ to removedQuery

src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs

@@ -315,7 +315,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 Query = query,
                 OwnerUri = ownerUri
             };
-            await queryExecutionService.InterServiceExecuteQuery(executeParams, eventSender,
+            await queryExecutionService.InterServiceExecuteQuery(executeParams, null, eventSender,
                 queryCreateSuccessCallback, queryCreateFailureCallback,
                 queryCompleteSuccessCallback, queryCompleteFailureCallback);
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.cs

@@ -245,6 +245,22 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string QueryServiceResultSetHasNoResults
+        {
+            get
+            {
+                return Keys.GetString(Keys.QueryServiceResultSetHasNoResults);
+            }
+        }
+
+        public static string QueryServiceResultSetTooLarge
+        {
+            get
+            {
+                return Keys.GetString(Keys.QueryServiceResultSetTooLarge);
+            }
+        }
+
         public static string QueryServiceSaveAsResultSetNotComplete
         {
             get
@@ -3468,6 +3484,12 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string QueryServiceResultSetAddNoRows = "QueryServiceResultSetAddNoRows";
 
 
+            public const string QueryServiceResultSetHasNoResults = "QueryServiceResultSetHasNoResults";
+
+
+            public const string QueryServiceResultSetTooLarge = "QueryServiceResultSetTooLarge";
+
+
             public const string QueryServiceSaveAsResultSetNotComplete = "QueryServiceSaveAsResultSetNotComplete";
 
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.resx

@@ -265,6 +265,14 @@
     <value>Cannot add row to result buffer, data reader does not contain rows</value>
     <comment></comment>
   </data>
+  <data name="QueryServiceResultSetHasNoResults" xml:space="preserve">
+    <value>Query has no results to return</value>
+    <comment></comment>
+  </data>
+  <data name="QueryServiceResultSetTooLarge" xml:space="preserve">
+    <value>Result set has too many rows to be safely loaded</value>
+    <comment></comment>
+  </data>
   <data name="QueryServiceSaveAsResultSetNotComplete" xml:space="preserve">
     <value>Result cannot be saved until query execution has completed</value>
     <comment></comment>

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -108,6 +108,10 @@ QueryServiceMessageSenderNotSql = Sender for OnInfoMessage event must be a SqlCo
 
 QueryServiceResultSetAddNoRows = Cannot add row to result buffer, data reader does not contain rows
 
+QueryServiceResultSetHasNoResults = Query has no results to return
+
+QueryServiceResultSetTooLarge = Result set has too many rows to be safely loaded
+
 ### Save As Requests
 
 QueryServiceSaveAsResultSetNotComplete = Result cannot be saved until query execution has completed

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.xlf

@@ -2101,6 +2101,15 @@
         <target state="new">The database {0} is not accessible.</target>
         <note></note>
       </trans-unit>
+      <trans-unit id="QueryServiceResultSetHasNoResults">
+        <source>Query has no results to return</source>
+        <target state="new">Query has no results to return</target>
+        <note></note>
+      </trans-unit>
+      <trans-unit id="QueryServiceResultSetTooLarge">
+        <source>Result set has too many rows to be safely loaded</source>
+        <target state="new">Result set has too many rows to be safely loaded</target>
+      </trans-unit>
       <trans-unit id="Backup_TaskName">
         <source>Backup Database</source>
         <target state="new">Backup Database</target>

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/SimpleExecuteRequest.cs

@@ -0,0 +1,54 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteRequests
+{
+    /// <summary>
+    /// Parameters for executing a query from a provided string
+    /// </summary>
+    public class SimpleExecuteParams
+    {
+        /// <summary>
+        /// The string to execute
+        /// </summary>
+        public string QueryString { get; set; }
+
+        /// <summary>
+        /// The owneruri to get connection from
+        /// </summary>
+        public string OwnerUri { get; set; }
+    }
+
+    /// <summary>
+    /// Result
+    /// </summary>
+    public class SimpleExecuteResult
+    {
+
+        /// <summary>
+        /// The number of rows that was returned with the resultset
+        /// </summary>
+        public long RowCount { get; set; }
+
+        /// <summary>
+        /// Details about the columns that are provided as solutions
+        /// </summary>
+        public DbColumnWrapper[] ColumnInfo { get; set; }
+
+        /// <summary>
+        /// 2D array of the cell values requested from result set
+        /// </summary>
+        public DbCellValue[][] Rows { get; set; }
+    }
+
+    public class SimpleExecuteRequest
+    {
+        public static readonly
+            RequestType<SimpleExecuteParams, SimpleExecuteResult> Type =
+            RequestType<SimpleExecuteParams, SimpleExecuteResult>.Create("query/simpleexecute");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs

@@ -137,6 +137,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(SaveResultsAsExcelRequest.Type, HandleSaveResultsAsExcelRequest);
             serviceHost.SetRequestHandler(SaveResultsAsJsonRequest.Type, HandleSaveResultsAsJsonRequest);
             serviceHost.SetRequestHandler(QueryExecutionPlanRequest.Type, HandleExecutionPlanRequest);
+            serviceHost.SetRequestHandler(SimpleExecuteRequest.Type, HandleSimpleExecuteRequest);
 
             // Register handler for shutdown event
             serviceHost.RegisterShutdownTask((shutdownParams, requestContext) =>
@@ -165,7 +166,88 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             Func<string, Task> queryCreateFailureAction = message => requestContext.SendError(message);
 
             // Use the internal handler to launch the query
-            return InterServiceExecuteQuery(executeParams, requestContext, queryCreateSuccessAction, queryCreateFailureAction, null, null);
+            return InterServiceExecuteQuery(executeParams, null, requestContext, queryCreateSuccessAction, queryCreateFailureAction, null, null);
+        }
+
+        /// <summary>
+        /// Handles a request to execute a string and return the result
+        /// </summary>
+        internal Task HandleSimpleExecuteRequest(SimpleExecuteParams executeParams,
+            RequestContext<SimpleExecuteResult> requestContext)
+        {
+             ExecuteStringParams executeStringParams = new ExecuteStringParams
+            {
+                Query = executeParams.QueryString,
+                // generate guid as the owner uri to make sure every query is unique
+                OwnerUri = Guid.NewGuid().ToString()
+            };
+
+            // get connection
+            ConnectionInfo connInfo;
+            if (!ConnectionService.TryFindConnection(executeParams.OwnerUri, out connInfo))
+            {
+                return requestContext.SendError(SR.QueryServiceQueryInvalidOwnerUri);
+            }
+
+            if (connInfo.ConnectionDetails.MultipleActiveResultSets == null || connInfo.ConnectionDetails.MultipleActiveResultSets == false) {
+                // if multipleActive result sets is not allowed, don't specific a connection and make the ownerURI the true owneruri
+                connInfo = null;
+                executeStringParams.OwnerUri = executeParams.OwnerUri;
+            }
+
+            Func<string, Task> queryCreateFailureAction = message => requestContext.SendError(message);
+
+            ResultOnlyContext<SimpleExecuteResult> newContext = new ResultOnlyContext<SimpleExecuteResult>(requestContext);
+
+            // handle sending event back when the query completes
+            Query.QueryAsyncEventHandler queryComplete = async q =>
+            {
+                Query removedQuery;
+                // check to make sure any results were recieved
+                if (q.Batches.Length == 0 || q.Batches[0].ResultSets.Count == 0) 
+                {
+                    await requestContext.SendError(SR.QueryServiceResultSetHasNoResults);
+                    ActiveQueries.TryRemove(executeStringParams.OwnerUri, out removedQuery);
+                    return;
+                } 
+
+                var rowCount = q.Batches[0].ResultSets[0].RowCount;
+                // check to make sure there is a safe amount of rows to load into memory
+                if (rowCount > Int32.MaxValue) 
+                {
+                    await requestContext.SendError(SR.QueryServiceResultSetTooLarge);
+                    ActiveQueries.TryRemove(executeStringParams.OwnerUri, out removedQuery);
+                    return;
+                }
+                
+                SubsetParams subsetRequestParams = new SubsetParams
+                {
+                    OwnerUri = executeStringParams.OwnerUri,
+                    BatchIndex = 0,
+                    ResultSetIndex = 0,
+                    RowsStartIndex = 0,
+                    RowsCount = Convert.ToInt32(rowCount)
+                };
+                // get the data to send back
+                ResultSetSubset subset = await InterServiceResultSubset(subsetRequestParams);
+                SimpleExecuteResult result = new SimpleExecuteResult
+                {
+                    RowCount = q.Batches[0].ResultSets[0].RowCount,
+                    ColumnInfo = q.Batches[0].ResultSets[0].Columns,
+                    Rows = subset.Rows
+                };
+                await requestContext.SendResult(result);
+                // remove the active query since we are done with it
+                ActiveQueries.TryRemove(executeStringParams.OwnerUri, out removedQuery);
+            };
+
+            // handle sending error back when query fails
+            Query.QueryAsyncErrorEventHandler queryFail = async (q, e) =>
+            {
+                await requestContext.SendError(e);
+            };
+
+            return InterServiceExecuteQuery(executeStringParams, connInfo, newContext, null, queryCreateFailureAction, queryComplete, queryFail);
         }
 
         /// <summary>
@@ -325,6 +407,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// custom actions to be taken upon creation of query and failure to create query.
         /// </summary>
         /// <param name="executeParams">Parameters for execution</param>
+        /// <param name="connInfo">Connection Info to use; will try and get the connection from owneruri if not provided</param>
         /// <param name="queryEventSender">Event sender that will send progressive events during execution of the query</param>
         /// <param name="queryCreateSuccessFunc">
         /// Callback for when query has been created successfully. If result is <c>true</c>, query
@@ -342,6 +425,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// Callback to call when query has completed execution with errors. May be <c>null</c>.
         /// </param>
         public async Task InterServiceExecuteQuery(ExecuteRequestParamsBase executeParams, 
+            ConnectionInfo connInfo,
             IEventSender queryEventSender,
             Func<Query, Task<bool>> queryCreateSuccessFunc,
             Func<string, Task> queryCreateFailFunc,
@@ -355,7 +439,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             try
             {
                 // Get a new active query
-                newQuery = CreateQuery(executeParams);
+                newQuery = CreateQuery(executeParams, connInfo);
                 if (queryCreateSuccessFunc != null && !await queryCreateSuccessFunc(newQuery))
                 {
                     // The callback doesn't want us to continue, for some reason
@@ -441,11 +525,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Private Helpers
 
-        private Query CreateQuery(ExecuteRequestParamsBase executeParams)
+        private Query CreateQuery(ExecuteRequestParamsBase executeParams, ConnectionInfo connInfo)
         {
             // Attempt to get the connection for the editor
             ConnectionInfo connectionInfo;
-            if (!ConnectionService.TryFindConnection(executeParams.OwnerUri, out connectionInfo))
+            if (connInfo != null) {
+                connectionInfo = connInfo;
+            } else if (!ConnectionService.TryFindConnection(executeParams.OwnerUri, out connectionInfo))
             {
                 throw new ArgumentOutOfRangeException(nameof(executeParams.OwnerUri), SR.QueryServiceQueryInvalidOwnerUri);
             }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultOnlyContext.cs

@@ -0,0 +1,30 @@
+// 
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.SqlTools.Hosting.Protocol;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    /// <summary>
+    /// Implementation of IEventSender that swallows events without doing anything with them.
+    /// In the future this class could be used to roll up all the events and send
+    /// them all at once
+    /// </summary>
+    public class ResultOnlyContext<TResult> : IEventSender
+    {
+        private readonly RequestContext<TResult> origContext;
+
+        public ResultOnlyContext(RequestContext<TResult> context) {
+            origContext = context;
+        }
+
+        public virtual async Task SendEvent<TParams>(EventType<TParams> eventType, TParams eventParams)
+        {
+            // no op to swallow events
+            // in the future this could be used to roll up events and send them back in the result
+        }
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs

@@ -103,7 +103,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
             // If: I call the inter-service API to execute with a null execute params
             // Then: It should throw
             await Assert.ThrowsAsync<ArgumentNullException>(
-                () => qes.InterServiceExecuteQuery(null, eventSender, null, null, null, null));
+                () => qes.InterServiceExecuteQuery(null, null, eventSender, null, null, null, null));
         }
 
         [Fact]
@@ -116,7 +116,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
             // If: I call the inter-service API to execute a query with a a null event sender
             // Then: It should throw
             await Assert.ThrowsAsync<ArgumentNullException>(
-                () => qes.InterServiceExecuteQuery(executeParams, null, null, null, null, null));
+                () => qes.InterServiceExecuteQuery(executeParams, null, null, null, null, null, null));
         }
 
         [Fact]
@@ -432,6 +432,48 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
 
+        [Fact]
+        public async Task SimpleExecuteErrorWithNoResultsTest()
+        {
+            var queryService = Common.GetPrimedExecutionService(null, true, false, null);
+            var queryParams = new SimpleExecuteParams { OwnerUri = Constants.OwnerUri, QueryString = Constants.StandardQuery };
+            var efv = new EventFlowValidator<SimpleExecuteResult>()
+                .AddSimpleExecuteErrorValidator(SR.QueryServiceResultSetHasNoResults)
+                .Complete();
+            await queryService.HandleSimpleExecuteRequest(queryParams, efv.Object);
+
+            Query q;
+            queryService.ActiveQueries.TryGetValue(Constants.OwnerUri, out q);
+
+            // wait on the task to finish
+            q.ExecutionTask.Wait();
+
+            efv.Validate();
+
+            Assert.Equal(0, queryService.ActiveQueries.Count);
+        }
+        
+        [Fact]
+        public async Task SimpleExecuteVerifyResultsTest()
+        {
+            var queryService = Common.GetPrimedExecutionService(Common.StandardTestDataSet, true, false, null);
+            var queryParams = new SimpleExecuteParams { OwnerUri = Constants.OwnerUri, QueryString = Constants.StandardQuery };
+            var efv = new EventFlowValidator<SimpleExecuteResult>()
+                .AddSimpleExecuteQueryResultValidator(Common.StandardTestDataSet)
+                .Complete();
+            await queryService.HandleSimpleExecuteRequest(queryParams, efv.Object);
+
+            Query q;
+            queryService.ActiveQueries.TryGetValue(Constants.OwnerUri, out q);
+
+            // wait on the task to finish
+            q.ExecutionTask.Wait();
+            
+            efv.Validate();
+
+            Assert.Equal(0, queryService.ActiveQueries.Count);
+        }
+
         #endregion
 
         private static WorkspaceService<SqlToolsSettings> GetDefaultWorkspaceService(string query)
@@ -444,6 +486,25 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.Execution
 
     public static class QueryExecutionEventFlowValidatorExtensions
     {
+
+        public static EventFlowValidator<SimpleExecuteResult> AddSimpleExecuteQueryResultValidator(
+            this EventFlowValidator<SimpleExecuteResult> efv, TestResultSet[] testData)
+        {
+            return efv.AddResultValidation(p => 
+            {
+                Assert.Equal(p.RowCount, testData[0].Rows.Count);
+            });
+        }
+
+        public static EventFlowValidator<SimpleExecuteResult> AddSimpleExecuteErrorValidator(
+            this EventFlowValidator<SimpleExecuteResult> efv, string expectedMessage)
+        {
+            return efv.AddSimpleErrorValidation((m, e) => 
+            {
+                Assert.Equal(m, expectedMessage);
+            });
+        }
+
         public static EventFlowValidator<ExecuteRequestResult> AddStandardQueryResultValidator(
             this EventFlowValidator<ExecuteRequestResult> efv)
         {