ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-30 09:35:38 -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

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