ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-14 01:25:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

Make save result async (#107)

Browse Source 
* 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
This commit is contained in:
Sharon Ravindran
2016-10-21 20:07:21 -07:00
committed by GitHub
parent b389d275a2
commit 2a688cb87f
7 changed files with 468 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
View File

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

11
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs
View File

@@ -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>

46
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/FileUtils.cs Normal file
View File

@@ -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
}
}
}
}

178
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
View File

@@ -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
ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
if (!selectedResultSet.IsBeingDisposed)
{
using (StreamWriter csvFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create)))
// Create SaveResults object and add success and error handlers to respective events
SaveResults saveAsCsv = new SaveResults();
SaveResults.AsyncSaveEventHandler successHandler = async message =>
{
// 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
if (saveParams.IncludeHeaders)
{
await csvFile.WriteLineAsync( string.Join( ",", selectedResultSet.Columns.Skip(columnStartIndex).Take(columnCount).Select( column =>
SaveResults.EncodeCsvField(column.ColumnName) ?? string.Empty)));
}
// retrieve rows and write as csv
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 =>
SaveResults.EncodeCsvField((field != null) ? field.ToString(): "NULL"))));
}
}
// 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))
selectedResultSet.RemoveSaveTask(saveParams.FilePath);
await requestContext.SendResult(new SaveResultRequestResult { Messages = message });
};
saveAsCsv.SaveCompleted += successHandler;
SaveResults.AsyncSaveEventHandler errorHandler = async message =>
{
File.Delete(saveParams.FilePath);
}
await requestContext.SendError(ex.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);
}
}
/// <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
ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex];
if (!selectedResultSet.IsBeingDisposed)
{
using (StreamWriter jsonFile = new StreamWriter(File.Open(saveParams.FilePath, FileMode.Create)))
using (JsonWriter jsonWriter = new JsonTextWriter(jsonFile) )
// Create SaveResults object and add success and error handlers to respective events
SaveResults saveAsJson = new SaveResults();
SaveResults.AsyncSaveEventHandler successHandler = async message =>
{
jsonWriter.Formatting = Formatting.Indented;
jsonWriter.WriteStartArray();
// get the requested resultSet from query
Batch selectedBatch = result.Batches[saveParams.BatchIndex];
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))
{
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;
}
// 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))
selectedResultSet.RemoveSaveTask(saveParams.FilePath);
await requestContext.SendResult(new SaveResultRequestResult { Messages = message });
};
saveAsJson.SaveCompleted += successHandler;
SaveResults.AsyncSaveEventHandler errorHandler = async message =>
{
File.Delete(saveParams.FilePath);
}
await requestContext.SendError(ex.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);
}
}
#endregion
@@ -439,27 +367,27 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
string queryText;
if (executeParams.QuerySelection != null)
if (executeParams.QuerySelection != null)
{
string[] queryTextArray = queryFile.GetLinesInRange(
new BufferRange(
new BufferPosition(
executeParams.QuerySelection.StartLine + 1,
executeParams.QuerySelection.StartLine + 1,
executeParams.QuerySelection.StartColumn + 1
),
),
new BufferPosition(
executeParams.QuerySelection.EndLine + 1,
executeParams.QuerySelection.EndLine + 1,
executeParams.QuerySelection.EndColumn + 1
)
)
);
queryText = queryTextArray.Aggregate((a, b) => a + '\r' + '\n' + b);
}
else
}
else
{
queryText = queryFile.Contents;
}
// If we can't add the query now, it's assumed the query is in progress
Query newQuery = new Query(queryText, connectionInfo, settings, BufferFileFactory);
if (!ActiveQueries.TryAdd(executeParams.OwnerUri, newQuery))

133
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs
View File

@@ -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,28 +170,30 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
string[][] rows;
// If result set is 'for xml' or 'for json',
// Concatenate all the rows together into one row
if (isSingleColumnXmlJsonResultSet)
using (IFileStreamReader fileStreamReader = fileStreamFactory.GetReader(outputFileName))
{
// Iterate over all the rows and process them into a list of string builders
IEnumerable<StringBuilder> sbRows = FileOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)
.Select(cell => cell.DisplayValue).Aggregate(new StringBuilder(), (sb, value) => sb.Append(value)));
rows = new[] { new[] { string.Join(string.Empty, sbRows) } };
// 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<string> rowValues = FileOffsets.Select(rowOffset => fileStreamReader.ReadRow(rowOffset, Columns)[0].DisplayValue);
rows = new[] { new[] { string.Join(string.Empty, rowValues) } };
}
else
{
// Figure out which rows we need to read back
IEnumerable<long> rowOffsets = FileOffsets.Skip(startRow).Take(rowCount);
// Iterate over the rows we need and process them into output
rows = rowOffsets.Select(rowOffset =>
fileStreamReader.ReadRow(rowOffset, Columns).Select(cell => cell.DisplayValue).ToArray())
.ToArray();
}
}
else
{
// Figure out which rows we need to read back
IEnumerable<long> rowOffsets = FileOffsets.Skip(startRow).Take(rowCount);
// Iterate over the rows we need and process them into output
rows = rowOffsets.Select(rowOffset =>
fileStreamReader.ReadRow(rowOffset, Columns).Select(cell => cell.DisplayValue).ToArray())
.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;
}
if (disposing)
isBeingDisposed = true;
// Check if saveTasks are running for this ResultSet
if (!saveTasks.IsEmpty)
{
fileStreamReader?.Dispose();
fileStreamFactory.DisposeFile(outputFileName);
// Wait for tasks to finish before disposing ResultSet
Task.WhenAll(saveTasks.Values.ToArray()).ContinueWith((antecedent) =>
{
if (disposing)
{
fileStreamFactory.DisposeFile(outputFileName);
}
disposed = true;
isBeingDisposed = false;
});
}
else
{
// If saveTasks is empty, continue with dispose
if (disposing)
{
fileStreamFactory.DisposeFile(outputFileName);
}
disposed = true;
isBeingDisposed = false;
}
disposed = true;
}
#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
}
}

241
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs
View File

@@ -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">
@@ -32,7 +64,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
internal static String EncodeCsvField(String field)
{
StringBuilder sbField = new StringBuilder(field);
//Whether this field has special characters which require it to be embedded in quotes
bool embedInQuotes = false;
@@ -67,12 +99,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
}
}
//Replace all quotes in the original field with double quotes
sbField.Replace("\"", "\"\"");
String ret = sbField.ToString();
if (embedInQuotes)
{
ret = "\"" + ret + "\"";
@@ -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);
}
}
});
}
}
}

53
test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
View File

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