Make save result async (#107)

* Make save results asynchronous

* Prevent write share of file

* Lock objects in stages

* Create Save result objects

* refactor and write rows in batches

* CHange batchSize from test value

* Remove await in handler

* Removing the file reader as a member of the resultset

* Change Dispose to wait for save

* Change concurrentBag

* PascalCase variables

* Modify function signature and tests

* Safe file methods

* refactor ResultSets to Ilist and remove ToList

* Change dictionary key and prevent add to saveTasks during dispose

* Simplify row concatenation

* Fix prevent add

* Fix prevent add

* Add methods to expose saveTasks and isBeingDisposed

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

@@ -125,7 +125,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <summary>
         /// The result sets of the batch execution
         /// </summary>
-        public IEnumerable<ResultSet> ResultSets
+        public IList<ResultSet> ResultSets
         {
             get { return resultSets; }
         }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs

@@ -60,17 +60,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
     /// Parameters to save results as CSV
     /// </summary>
     public class SaveResultsAsCsvRequestParams: SaveResultsRequestParams{
-        
-        /// <summary>
-        /// CSV - Write values in quotes 
-        /// </summary>
-        public Boolean ValueInQuotes { get; set; }
-
-        /// <summary>
-        /// The encoding of the file to save results in
-        /// </summary>
-        public string FileEncoding { get; set; }
-
         /// <summary>
         /// Include headers of columns in CSV
         /// </summary>

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

@@ -0,0 +1,46 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+using System;
+using System.IO;
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    internal static class FileUtils
+    {
+        /// <summary>
+        /// Checks if file exists and swallows exceptions, if any
+        /// </summary>
+        /// <param name="path"> path of the file</param>
+        /// <returns></returns>
+        internal static bool SafeFileExists(string path)
+        {
+            try
+            {
+                return File.Exists(path);
+            }
+            catch (Exception)
+            {
+                // Swallow exception
+                return false;
+            }
+        }
+
+        /// <summary>
+        /// Deletes a file and swallows exceptions, if any
+        /// </summary>
+        /// <param name="path"></param>
+        internal static void SafeFileDelete(string path)
+        {
+            try
+            {
+                File.Delete(path);
+            }
+            catch (Exception)
+            {
+                // Swallow exception, do nothing
+            }
+        }
+
+    }
+}

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

@@ -4,7 +4,6 @@
 //
 using System;
 using System.Collections.Concurrent;
-using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
@@ -16,7 +15,6 @@ using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
-using Newtonsoft.Json;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
@@ -252,7 +250,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <summary>
         /// Process request to save a resultSet to a file in CSV format
         /// </summary>
-        public async Task HandleSaveResultsAsCsvRequest(SaveResultsAsCsvRequestParams saveParams,
+        internal async Task HandleSaveResultsAsCsvRequest(SaveResultsAsCsvRequestParams saveParams,
             RequestContext<SaveResultRequestResult> requestContext)
         {
             // retrieve query for OwnerUri
@@ -265,67 +263,39 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 });
                 return;
             }
-            try
-            {
-                using (StreamWriter csvFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create)))
-                {
-                    // get the requested resultSet from query
-                    Batch selectedBatch = result.Batches[saveParams.BatchIndex];
-                    ResultSet selectedResultSet = (selectedBatch.ResultSets.ToList())[saveParams.ResultSetIndex];
-                    int columnCount = 0;
-                    int rowCount = 0;
-                    int columnStartIndex = 0;
-                    int rowStartIndex = 0;
 
-                    // set column, row counts depending on whether save request is for entire result set or a subset
-                    if (SaveResults.isSaveSelection(saveParams))
-                    {   
-                        columnCount = saveParams.ColumnEndIndex.Value - saveParams.ColumnStartIndex.Value + 1;
-                        rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
-                        columnStartIndex = saveParams.ColumnStartIndex.Value;
-                        rowStartIndex =saveParams.RowStartIndex.Value;
-                    }
-                    else
-                    {                
-                        columnCount = selectedResultSet.Columns.Length;
-                        rowCount = (int)selectedResultSet.RowCount;
-                    }
 
-                    // write column names if include headers option is chosen
+            ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
-                    if (saveParams.IncludeHeaders)
+            if (!selectedResultSet.IsBeingDisposed)
             {
-                        await csvFile.WriteLineAsync( string.Join( ",", selectedResultSet.Columns.Skip(columnStartIndex).Take(columnCount).Select( column =>
+                // Create SaveResults object and add success and error handlers to respective events
-                                            SaveResults.EncodeCsvField(column.ColumnName) ?? string.Empty)));
+                SaveResults saveAsCsv = new SaveResults();
-                    }
 
-                    // retrieve rows and write as csv
+                SaveResults.AsyncSaveEventHandler successHandler = async message =>
-                    ResultSetSubset resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex, rowCount);
-                    foreach (var row in resultSubset.Rows)
                 {
-                        await csvFile.WriteLineAsync( string.Join( ",", row.Skip(columnStartIndex).Take(columnCount).Select( field => 
+                    selectedResultSet.RemoveSaveTask(saveParams.FilePath);
-                                            SaveResults.EncodeCsvField((field != null) ? field.ToString(): "NULL"))));
+                    await requestContext.SendResult(new SaveResultRequestResult { Messages = message });
-                    }
+                };
+                saveAsCsv.SaveCompleted += successHandler;
+                SaveResults.AsyncSaveEventHandler errorHandler = async message =>
+                {
+                    selectedResultSet.RemoveSaveTask(saveParams.FilePath);
+                    await requestContext.SendError(message);
+                };
+                saveAsCsv.SaveFailed += errorHandler;
 
-                }
+                saveAsCsv.SaveResultSetAsCsv(saveParams, requestContext, result);
+
+                // Associate the ResultSet with the save task
+                selectedResultSet.AddSaveTask(saveParams.FilePath, saveAsCsv.SaveTask);
 
-                // Successfully wrote file, send success result
-                await requestContext.SendResult(new SaveResultRequestResult { Messages = null });
-            }
-            catch(Exception ex)
-            {
-                // Delete file when exception occurs
-                if (File.Exists(saveParams.FilePath))
-                {
-                    File.Delete(saveParams.FilePath);
-                }
-                await requestContext.SendError(ex.Message);
             }
         }
 
         /// <summary>
         /// Process request to save a resultSet to a file in JSON format
         /// </summary>
-        public async Task HandleSaveResultsAsJsonRequest(SaveResultsAsJsonRequestParams saveParams,
+        internal async Task HandleSaveResultsAsJsonRequest(SaveResultsAsJsonRequestParams saveParams,
             RequestContext<SaveResultRequestResult> requestContext)
         {
             // retrieve query for OwnerUri
@@ -338,73 +308,31 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 });
                 return;
             }
-            try
-            {
-                using (StreamWriter jsonFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create)))
-                using (JsonWriter jsonWriter = new JsonTextWriter(jsonFile) )
-                {
-                    jsonWriter.Formatting = Formatting.Indented;
-                    jsonWriter.WriteStartArray();
 
-                    // get the requested resultSet from query
+            ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
-                    Batch selectedBatch = result.Batches[saveParams.BatchIndex];
+            if (!selectedResultSet.IsBeingDisposed)
-                    ResultSet selectedResultSet = selectedBatch.ResultSets.ToList()[saveParams.ResultSetIndex];
-                    int rowCount = 0;
-                    int rowStartIndex = 0;
-                    int columnStartIndex = 0;
-                    int columnEndIndex = 0;
-
-                    // set column, row counts depending on whether save request is for entire result set or a subset
-                    if (SaveResults.isSaveSelection(saveParams))
             {
-                       
+                // Create SaveResults object and add success and error handlers to respective events
-                        rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
+                SaveResults saveAsJson = new SaveResults();
-                        rowStartIndex = saveParams.RowStartIndex.Value;
+                SaveResults.AsyncSaveEventHandler successHandler = async message =>
-                        columnStartIndex = saveParams.ColumnStartIndex.Value;
-                        columnEndIndex = saveParams.ColumnEndIndex.Value + 1 ; // include the last column
-                    }
-                    else 
                 {
-                        rowCount = (int)selectedResultSet.RowCount;
+                    selectedResultSet.RemoveSaveTask(saveParams.FilePath);
-                        columnEndIndex = selectedResultSet.Columns.Length;
+                    await requestContext.SendResult(new SaveResultRequestResult { Messages = message });
+                };
+                saveAsJson.SaveCompleted += successHandler;
+                SaveResults.AsyncSaveEventHandler errorHandler = async message =>
+                {
+                    selectedResultSet.RemoveSaveTask(saveParams.FilePath);
+                    await requestContext.SendError(message);
+                };
+                saveAsJson.SaveFailed += errorHandler;
+
+                saveAsJson.SaveResultSetAsJson(saveParams, requestContext, result);
+
+                // Associate the ResultSet with the save task
+                selectedResultSet.AddSaveTask(saveParams.FilePath, saveAsJson.SaveTask);
             }
 
-                    // retrieve rows and write as json
-                    ResultSetSubset resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex, rowCount);
-                    foreach (var row in resultSubset.Rows)
-                    {
-                        jsonWriter.WriteStartObject();
-                        for (int i = columnStartIndex ; i < columnEndIndex; i++)
-                        {
-                            //get column name
-                            DbColumnWrapper col = selectedResultSet.Columns[i];
-                            string val = row[i];
-                            jsonWriter.WritePropertyName(col.ColumnName);
-                            if (val == null)
-                            {
-                                jsonWriter.WriteNull();
-                            }
-                            else
-                            {
-                                jsonWriter.WriteValue(val);
-                            }
-                        }
-                        jsonWriter.WriteEndObject();
-                    }
-                    jsonWriter.WriteEndArray();
-                }
-
-                await requestContext.SendResult(new SaveResultRequestResult { Messages = null });
-            }
-            catch(Exception ex)
-            {
-                // Delete file when exception occurs
-                if (File.Exists(saveParams.FilePath))
-                {
-                    File.Delete(saveParams.FilePath);
-                }
-                await requestContext.SendError(ex.Message);
-            }
         }
 
         #endregion

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

@@ -4,6 +4,7 @@
 //
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Data.Common;
 using System.Linq;
@@ -44,12 +45,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private readonly IFileStreamFactory fileStreamFactory;
 
-        /// <summary>
-        /// File stream reader that will be reused to make rapid-fire retrieval of result subsets
-        /// quick and low perf impact.
-        /// </summary>
-        private IFileStreamReader fileStreamReader;
-
         /// <summary>
         /// Whether or not the result set has been read in from the database
         /// </summary>
@@ -65,6 +60,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private readonly string outputFileName;
 
+        /// <summary>
+        /// Whether the resultSet is in the process of being disposed
+        /// </summary>
+        private bool isBeingDisposed;
+
+        /// <summary>
+        /// All save tasks currently saving this ResultSet
+        /// </summary>
+        private ConcurrentDictionary<string, Task> saveTasks;
+
         #endregion
 
         /// <summary>
@@ -86,10 +91,23 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             // Store the factory
             fileStreamFactory = factory;
             hasBeenRead = false;
+            saveTasks = new ConcurrentDictionary<string, Task>();
         }
 
         #region Properties
 
+        /// <summary>
+        /// Whether the resultSet is in the process of being disposed
+        /// </summary>
+        /// <returns></returns>
+        internal bool IsBeingDisposed
+        {
+            get
+            {
+                return isBeingDisposed;
+            }
+        }
+
         /// <summary>
         /// The columns for this result set
         /// </summary>
@@ -120,18 +138,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public long RowCount { get; private set; }
 
-        /// <summary>
-        /// The rows of this result set
-        /// </summary>
-        public IEnumerable<string[]> Rows
-        {
-            get
-            {
-                return FileOffsets.Select(
-                    offset => fileStreamReader.ReadRow(offset, Columns).Select(cell => cell.DisplayValue).ToArray());
-            }
-        }
-
         #endregion
 
         #region Public Methods
@@ -145,7 +151,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         public Task<ResultSetSubset> GetSubset(int startRow, int rowCount)
         {
             // Sanity check to make sure that the results have been read beforehand
-            if (!hasBeenRead || fileStreamReader == null)
+            if (!hasBeenRead)
             {
                 throw new InvalidOperationException(SR.QueryServiceResultSetNotRead);
             }
@@ -164,14 +170,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             {
 
                 string[][] rows;
+
+                using (IFileStreamReader fileStreamReader = fileStreamFactory.GetReader(outputFileName))
+                {
                     // If result set is 'for xml' or 'for json',
                     // Concatenate all the rows together into one row
                     if (isSingleColumnXmlJsonResultSet)
                     {
                         // Iterate over all the rows and process them into a list of string builders
-                    IEnumerable<StringBuilder> sbRows = FileOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)
+                        IEnumerable<string> rowValues = FileOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)[0].DisplayValue);
-                        .Select(cell => cell.DisplayValue).Aggregate(new StringBuilder(), (sb, value) => sb.Append(value)));
+                        rows = new[] { new[] { string.Join(string.Empty, rowValues) } };
-                    rows = new[] { new[] { string.Join(string.Empty, sbRows) } };
 
                     }
                     else
@@ -185,7 +193,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                             .ToArray();
 
                     }
-
+                }
                 // Retrieve the subset of the results as per the request
                 return new ResultSetSubset
                 {
@@ -203,7 +211,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         {
             // Mark that result has been read
             hasBeenRead = true;
-            fileStreamReader = fileStreamFactory.GetReader(outputFileName);
 
             // Open a writer for the file
             using (IFileStreamWriter fileWriter = fileStreamFactory.GetWriter(outputFileName, MaxCharsToStore, MaxXmlCharsToStore))
@@ -244,13 +251,31 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 return;
             }
 
+            isBeingDisposed = true;
+            // Check if saveTasks are running for this ResultSet
+            if (!saveTasks.IsEmpty)
+            {
+                // Wait for tasks to finish before disposing ResultSet
+                Task.WhenAll(saveTasks.Values.ToArray()).ContinueWith((antecedent) =>
+                {
                     if (disposing)
                     {
-                fileStreamReader?.Dispose();
                         fileStreamFactory.DisposeFile(outputFileName);
                     }
-
                     disposed = true;
+                    isBeingDisposed = false;
+                });
+            }
+            else
+            {
+                // If saveTasks is empty, continue with dispose
+                if (disposing)
+                {
+                    fileStreamFactory.DisposeFile(outputFileName);
+                }
+                disposed = true;
+                isBeingDisposed = false;
+            }
         }
 
         #endregion
@@ -283,5 +308,25 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         }
 
         #endregion
+
+        #region Internal Methods to Add and Remove save tasks
+        internal void AddSaveTask(string key, Task saveTask)
+        {
+            saveTasks.TryAdd(key, saveTask);
+        }
+
+        internal void RemoveSaveTask(string key)
+        {
+            Task completedTask;
+            saveTasks.TryRemove(key, out completedTask);
+        }
+
+        internal Task GetSaveTask(string key)
+        {
+            Task completedTask;
+            saveTasks.TryRemove(key, out completedTask);
+            return completedTask;
+        }
+        #endregion
     }
 }

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

@@ -3,15 +3,47 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 using System;
+using System.IO;
+using System.Linq;
 using System.Text;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 
+using Newtonsoft.Json;
+
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
-    internal class SaveResults{
+    internal class SaveResults
+    {
+        /// <summary>
+        /// Number of rows being read from the ResultSubset in one read
+        /// </summary>
+        private const int BatchSize = 1000;
+
+        /// <summary>
+        /// Save Task that asynchronously writes ResultSet to file
+        /// </summary>
+        internal Task SaveTask { get; set; }
+
+        /// <summary>
+        /// Event Handler for save events
+        /// </summary>
+        /// <param name="message"> Message to be returned to client</param>
+        /// <returns></returns>
+        internal delegate Task AsyncSaveEventHandler(string message);
+
+        /// <summary>
+        /// A successful save event
+        /// </summary>
+        internal event AsyncSaveEventHandler SaveCompleted;
+
+        /// <summary>
+        /// A failed save event
+        /// </summary>
+        internal event AsyncSaveEventHandler SaveFailed;
 
         /// Method ported from SSMS
-
         /// <summary>
         /// Encodes a single field for inserting into a CSV record. The following rules are applied:
         /// <list type="bullet">
@@ -81,11 +113,208 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return ret;
         }
 
-        internal static bool isSaveSelection(SaveResultsRequestParams saveParams)
+        /// <summary>
+        /// Check if request is a subset of result set or whole result set
+        /// </summary>
+        /// <param name="saveParams"> Parameters from the request </param>
+        /// <returns></returns>
+        internal static bool IsSaveSelection(SaveResultsRequestParams saveParams)
         {
             return (saveParams.ColumnStartIndex != null && saveParams.ColumnEndIndex != null
                 && saveParams.RowEndIndex != null && saveParams.RowEndIndex != null);
         }
+
+        /// <summary>
+        /// Save results as JSON format to the file specified in saveParams
+        /// </summary>
+        /// <param name="saveParams"> Parameters from the request </param>
+        /// <param name="requestContext"> Request context for save results </param>
+        /// <param name="result"> Result query object </param>
+        /// <returns></returns>
+        internal void SaveResultSetAsJson(SaveResultsAsJsonRequestParams saveParams, RequestContext<SaveResultRequestResult> requestContext, Query result)
+        {
+            // Run in a separate thread
+            SaveTask = Task.Run(async () =>
+            {
+                try
+                {
+                    using (StreamWriter jsonFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create, FileAccess.ReadWrite, FileShare.Read)))
+                    using (JsonWriter jsonWriter = new JsonTextWriter(jsonFile))
+                    {
+
+                        int rowCount = 0;
+                        int rowStartIndex = 0;
+                        int columnStartIndex = 0;
+                        int columnEndIndex = 0;
+
+                        jsonWriter.Formatting = Formatting.Indented;
+                        jsonWriter.WriteStartArray();
+
+                        // Get the requested resultSet from query
+                        Batch selectedBatch = result.Batches[saveParams.BatchIndex];
+                        ResultSet selectedResultSet = selectedBatch.ResultSets[saveParams.ResultSetIndex];
+
+                        // Set column, row counts depending on whether save request is for entire result set or a subset
+                        if (IsSaveSelection(saveParams))
+                        {
+
+                            rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
+                            rowStartIndex = saveParams.RowStartIndex.Value;
+                            columnStartIndex = saveParams.ColumnStartIndex.Value;
+                            columnEndIndex = saveParams.ColumnEndIndex.Value + 1; // include the last column
+                        }
+                        else
+                        {
+                            rowCount = (int)selectedResultSet.RowCount;
+                            columnEndIndex = selectedResultSet.Columns.Length;
+                        }
+
+                        // Split rows into batches
+                        for (int count = 0; count < (rowCount / BatchSize) + 1; count++)
+                        {
+                            int numberOfRows = (count < rowCount / BatchSize) ? BatchSize : (rowCount % BatchSize);
+                            if (numberOfRows == 0)
+                            {
+                                break;
+                            }
+
+                            // Retrieve rows and write as json
+                            ResultSetSubset resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex + count * BatchSize, numberOfRows);
+                            foreach (var row in resultSubset.Rows)
+                            {
+                                jsonWriter.WriteStartObject();
+                                for (int i = columnStartIndex; i < columnEndIndex; i++)
+                                {
+                                    // Write columnName, value pair
+                                    DbColumnWrapper col = selectedResultSet.Columns[i];
+                                    string val = row[i]?.ToString();
+                                    jsonWriter.WritePropertyName(col.ColumnName);
+                                    if (val == null)
+                                    {
+                                        jsonWriter.WriteNull();
+                                    }
+                                    else
+                                    {
+                                        jsonWriter.WriteValue(val);
+                                    }
+                                }
+                                jsonWriter.WriteEndObject();
+                            }
+
+                        }
+                        jsonWriter.WriteEndArray();
+                    }
+
+                    // Successfully wrote file, send success result
+                    if (SaveCompleted != null)
+                    {
+                        await SaveCompleted(null);
+                    }
+
+                    
+                }
+                catch (Exception ex)
+                {
+                    // Delete file when exception occurs
+                    if (FileUtils.SafeFileExists(saveParams.FilePath))
+                    {
+                        FileUtils.SafeFileDelete(saveParams.FilePath);
+                    }
+                    if (SaveFailed != null)
+                    {
+                        await SaveFailed(ex.ToString());
+                    }
+                }
+            });
+        }
+
+        /// <summary>
+        /// Save results as CSV format to the file specified in saveParams
+        /// </summary>
+        /// <param name="saveParams"> Parameters from the request </param>
+        /// <param name="requestContext">  Request context for save results </param>
+        /// <param name="result"> Result query object </param>
+        /// <returns></returns>
+        internal void SaveResultSetAsCsv(SaveResultsAsCsvRequestParams saveParams, RequestContext<SaveResultRequestResult> requestContext, Query result)
+        {
+            // Run in a separate thread
+            SaveTask = Task.Run(async () =>
+            {
+                try
+                {
+                    using (StreamWriter csvFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create, FileAccess.ReadWrite, FileShare.Read)))
+                    {
+                        ResultSetSubset resultSubset;
+                        int columnCount = 0;
+                        int rowCount = 0;
+                        int columnStartIndex = 0;
+                        int rowStartIndex = 0;
+
+                        // Get the requested resultSet from query
+                        Batch selectedBatch = result.Batches[saveParams.BatchIndex];
+                        ResultSet selectedResultSet = (selectedBatch.ResultSets)[saveParams.ResultSetIndex];
+                        // Set column, row counts depending on whether save request is for entire result set or a subset
+                        if (IsSaveSelection(saveParams))
+                        {
+                            columnCount = saveParams.ColumnEndIndex.Value - saveParams.ColumnStartIndex.Value + 1;
+                            rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
+                            columnStartIndex = saveParams.ColumnStartIndex.Value;
+                            rowStartIndex = saveParams.RowStartIndex.Value;
+                        }
+                        else
+                        {
+                            columnCount = selectedResultSet.Columns.Length;
+                            rowCount = (int)selectedResultSet.RowCount;
+                        }
+
+                        // Write column names if include headers option is chosen
+                        if (saveParams.IncludeHeaders)
+                        {
+                            csvFile.WriteLine(string.Join(",", selectedResultSet.Columns.Skip(columnStartIndex).Take(columnCount).Select(column =>
+                                            EncodeCsvField(column.ColumnName) ?? string.Empty)));
+                        }
+
+                        for (int i = 0; i < (rowCount / BatchSize) + 1; i++)
+                        {
+                            int numberOfRows = (i < rowCount / BatchSize) ? BatchSize : (rowCount % BatchSize);
+                            if (numberOfRows == 0)
+                            {
+                                break;
+                            }
+                            // Retrieve rows and write as csv
+                            resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex + i * BatchSize, numberOfRows);
+
+                            foreach (var row in resultSubset.Rows)
+                            {
+                                csvFile.WriteLine(string.Join(",", row.Skip(columnStartIndex).Take(columnCount).Select(field =>
+                                                EncodeCsvField((field != null) ? field.ToString() : "NULL"))));
+                            }
+                        }
+                    }
+
+                    // Successfully wrote file, send success result
+                    if (SaveCompleted != null)
+                    {
+                        await SaveCompleted(null);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    // Delete file when exception occurs
+                    if (FileUtils.SafeFileExists(saveParams.FilePath))
+                    {
+                        FileUtils.SafeFileDelete(saveParams.FilePath);
+                    }
+
+                    if (SaveFailed != null)
+                    {
+                        await SaveFailed(ex.Message);
+                    }
+                }
+            });
+        }
+
+
     }
 
 }

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs

@@ -3,10 +3,12 @@
 //
 
 using System;
+using System.Linq;
 using System.IO;
 using System.Threading.Tasks;
 using System.Runtime.InteropServices;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
@@ -46,7 +48,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             SaveResultRequestResult result = null;
             var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object).Wait();
+
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object);
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);          
+            await saveTask;
 
             // Expect to see a file successfully created in filepath and a success message
             Assert.Null(result.Messages);
@@ -89,7 +96,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             SaveResultRequestResult result = null;
             var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object).Wait();
+
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object);
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);         
+            await saveTask;
 
             // Expect to see a file successfully created in filepath and a success message
             Assert.Null(result.Messages);
@@ -128,7 +140,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             string errMessage = null;
             var saveRequest = GetSaveResultsContextMock( null, err => errMessage = (string) err);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object).Wait();
+
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object);           
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);          
+            await saveTask;
 
             // Expect to see error message
             Assert.NotNull(errMessage);
@@ -188,7 +205,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             SaveResultRequestResult result = null;
             var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object).Wait();
+            
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object);
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);         
+            await saveTask;
+
+
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
             // Expect to see a file successfully created in filepath and a success message
             Assert.Null(result.Messages);
@@ -223,14 +247,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 BatchIndex = 0,
                 FilePath = "testwrite_5.json",          
                 RowStartIndex = 0,
-                RowEndIndex = 0,
+                RowEndIndex = 1,
                 ColumnStartIndex = 0,
-                ColumnEndIndex = 0             
+                ColumnEndIndex = 1             
             };
             SaveResultRequestResult result = null;
             var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object).Wait();
+
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object);
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);         
+            await saveTask;
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                         
             // Expect to see a file successfully created in filepath and a success message
             Assert.Null(result.Messages);
@@ -265,11 +294,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                 BatchIndex = 0,
                 FilePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "G:\\test.json" : "/test.json"
             };
-            // SaveResultRequestResult result = null;
+
+
             string errMessage = null;
             var saveRequest = GetSaveResultsContextMock( null, err => errMessage = (string) err);
             queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
-            queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object).Wait();
+            
+            // Call save results and wait on the save task
+            await queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object);
+            ResultSet selectedResultSet = queryService.ActiveQueries[saveParams.OwnerUri].Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
+            Task saveTask = selectedResultSet.GetSaveTask(saveParams.FilePath);         
+            await saveTask;
 
             // Expect to see error message
             Assert.NotNull(errMessage);