diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs index c7c379c9..0ee7d550 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs @@ -368,6 +368,34 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution return targetResultSet.GetSubset(startRow, rowCount); } + /// + /// Saves a result to a file format selected by the user + /// + /// Parameters for the save as request + /// + /// Factory for creating the reader/writer pair for outputing to the selected format + /// + /// Delegate to call when request successfully completes + /// Delegate to call if the request fails + public void SaveAs(SaveResultsRequestParams saveParams, IFileStreamFactory fileFactory, + ResultSet.SaveAsAsyncEventHandler successHandler, ResultSet.SaveAsFailureAsyncEventHandler failureHandler) + { + // Get the result set to save + ResultSet resultSet; + lock (resultSets) + { + // Sanity check to make sure we have a valid result set + if (saveParams.ResultSetIndex < 0 || saveParams.ResultSetIndex >= resultSets.Count) + { + throw new ArgumentOutOfRangeException(nameof(saveParams.BatchIndex), SR.QueryServiceSubsetResultSetOutOfRange); + } + + + resultSet = resultSets[saveParams.ResultSetIndex]; + } + resultSet.SaveAs(saveParams, fileFactory, successHandler, failureHandler); + } + #endregion #region Private Helpers diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs index 85f87b9d..a30a8e59 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs @@ -54,6 +54,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts /// /// public int? ColumnEndIndex { get; set; } + + /// + /// Check if request is a subset of result set or whole result set + /// + /// + internal bool IsSaveSelection + { + get + { + return ColumnStartIndex.HasValue && ColumnEndIndex.HasValue + && RowStartIndex.HasValue && RowEndIndex.HasValue; + } + } } /// diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamFactory.cs index 6cb50095..2bd4f96b 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamFactory.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamFactory.cs @@ -14,7 +14,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage IFileStreamReader GetReader(string fileName); - IFileStreamWriter GetWriter(string fileName, int maxCharsToStore, int maxXmlCharsToStore); + IFileStreamWriter GetWriter(string fileName); void DisposeFile(string fileName); diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamWriter.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamWriter.cs index 7cfffee8..9e919dfb 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamWriter.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/IFileStreamWriter.cs @@ -4,7 +4,8 @@ // using System; -using System.Data.SqlTypes; +using System.Collections.Generic; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage { @@ -14,24 +15,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage public interface IFileStreamWriter : IDisposable { int WriteRow(StorageDataReader dataReader); - int WriteNull(); - int WriteInt16(short val); - int WriteInt32(int val); - int WriteInt64(long val); - int WriteByte(byte val); - int WriteChar(char val); - int WriteBoolean(bool val); - int WriteSingle(float val); - int WriteDouble(double val); - int WriteDecimal(decimal val); - int WriteSqlDecimal(SqlDecimal val); - int WriteDateTime(DateTime val); - int WriteDateTimeOffset(DateTimeOffset dtoVal); - int WriteTimeSpan(TimeSpan val); - int WriteString(string val); - int WriteBytes(byte[] bytes); - int WriteGuid(Guid val); - int WriteMoney(SqlMoney val); + void WriteRow(IList row, IList columns); void FlushBuffer(); } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamFactory.cs new file mode 100644 index 00000000..becd5588 --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamFactory.cs @@ -0,0 +1,66 @@ +// +// 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; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage +{ + /// + /// Factory for creating a reader/writer pair that will read from the temporary buffer file + /// and output to a CSV file. + /// + public class SaveAsCsvFileStreamFactory : IFileStreamFactory + { + #region Properties + + /// + /// Parameters for the save as CSV request + /// + public SaveResultsAsCsvRequestParams SaveRequestParams { get; set; } + + #endregion + + /// + /// File names are not meant to be created with this factory. + /// + /// Thrown all times + [Obsolete] + public string CreateFile() + { + throw new NotImplementedException(); + } + + /// + /// Returns a new service buffer reader for reading results back in from the temporary buffer files + /// + /// Path to the temp buffer file + /// Stream reader + public IFileStreamReader GetReader(string fileName) + { + return new ServiceBufferFileStreamReader(new FileStream(fileName, FileMode.Open, FileAccess.Read)); + } + + /// + /// Returns a new CSV writer for writing results to a CSV file + /// + /// Path to the CSV output file + /// Stream writer + public IFileStreamWriter GetWriter(string fileName) + { + return new SaveAsCsvFileStreamWriter(new FileStream(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite), SaveRequestParams); + } + + /// + /// Safely deletes the file + /// + /// Path to the file to delete + public void DisposeFile(string fileName) + { + FileUtils.SafeFileDelete(fileName); + } + } +} diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamWriter.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamWriter.cs new file mode 100644 index 00000000..a23b2cbe --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsCsvFileStreamWriter.cs @@ -0,0 +1,118 @@ +// +// 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.Collections.Generic; +using System.IO; +using System.Linq; +using System.Text; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage +{ + /// + /// Writer for writing rows of results to a CSV file + /// + public class SaveAsCsvFileStreamWriter : SaveAsStreamWriter + { + + #region Member Variables + + private readonly SaveResultsAsCsvRequestParams saveParams; + private bool headerWritten; + + #endregion + + /// + /// Constructor, stores the CSV specific request params locally, chains into the base + /// constructor + /// + /// FileStream to access the CSV file output + /// CSV save as request parameters + public SaveAsCsvFileStreamWriter(Stream stream, SaveResultsAsCsvRequestParams requestParams) + : base(stream, requestParams) + { + saveParams = requestParams; + } + + /// + /// Writes a row of data as a CSV row. If this is the first row and the user has requested + /// it, the headers for the column will be emitted as well. + /// + /// The data of the row to output to the file + /// + /// The entire list of columns for the result set. They will be filtered down as per the + /// request params. + /// + public override void WriteRow(IList row, IList columns) + { + // Write out the header if we haven't already and the user chose to have it + if (saveParams.IncludeHeaders && !headerWritten) + { + // Build the string + var selectedColumns = columns.Skip(ColumnStartIndex ?? 0).Take(ColumnCount ?? columns.Count) + .Select(c => EncodeCsvField(c.ColumnName) ?? string.Empty); + string headerLine = string.Join(",", selectedColumns); + + // Encode it and write it out + byte[] headerBytes = Encoding.UTF8.GetBytes(headerLine + Environment.NewLine); + FileStream.Write(headerBytes, 0, headerBytes.Length); + + headerWritten = true; + } + + // Build the string for the row + var selectedCells = row.Skip(ColumnStartIndex ?? 0) + .Take(ColumnCount ?? columns.Count) + .Select(c => EncodeCsvField(c.DisplayValue)); + string rowLine = string.Join(",", selectedCells); + + // Encode it and write it out + byte[] rowBytes = Encoding.UTF8.GetBytes(rowLine + Environment.NewLine); + FileStream.Write(rowBytes, 0, rowBytes.Length); + } + + /// + /// Encodes a single field for inserting into a CSV record. The following rules are applied: + /// + /// All double quotes (") are replaced with a pair of consecutive double quotes + /// + /// The entire field is also surrounded by a pair of double quotes if any of the following conditions are met: + /// + /// The field begins or ends with a space + /// The field begins or ends with a tab + /// The field contains the ListSeparator string + /// The field contains the '\n' character + /// The field contains the '\r' character + /// The field contains the '"' character + /// + /// + /// The field to encode + /// The CSV encoded version of the original field + internal static string EncodeCsvField(string field) + { + // Special case for nulls + if (field == null) + { + return "NULL"; + } + + // Whether this field has special characters which require it to be embedded in quotes + bool embedInQuotes = field.IndexOfAny(new[] {',', '\r', '\n', '"'}) >= 0 // Contains special characters + || field.StartsWith(" ") || field.EndsWith(" ") // Start/Ends with space + || field.StartsWith("\t") || field.EndsWith("\t"); // Starts/Ends with tab + + //Replace all quotes in the original field with double quotes + string ret = field.Replace("\"", "\"\""); + + if (embedInQuotes) + { + ret = $"\"{ret}\""; + } + + return ret; + } + } +} \ No newline at end of file diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamFactory.cs new file mode 100644 index 00000000..87c7ec94 --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamFactory.cs @@ -0,0 +1,64 @@ +// +// 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; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage +{ + public class SaveAsJsonFileStreamFactory : IFileStreamFactory + { + + #region Properties + + /// + /// Parameters for the save as JSON request + /// + public SaveResultsAsJsonRequestParams SaveRequestParams { get; set; } + + #endregion + + /// + /// File names are not meant to be created with this factory. + /// + /// Thrown all times + [Obsolete] + public string CreateFile() + { + throw new NotImplementedException(); + } + + /// + /// Returns a new service buffer reader for reading results back in from the temporary buffer files + /// + /// Path to the temp buffer file + /// Stream reader + public IFileStreamReader GetReader(string fileName) + { + return new ServiceBufferFileStreamReader(new FileStream(fileName, FileMode.Open, FileAccess.Read)); + } + + /// + /// Returns a new JSON writer for writing results to a JSON file + /// + /// Path to the JSON output file + /// Stream writer + public IFileStreamWriter GetWriter(string fileName) + { + return new SaveAsJsonFileStreamWriter(new FileStream(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite), SaveRequestParams); + } + + /// + /// Safely deletes the file + /// + /// Path to the file to delete + public void DisposeFile(string fileName) + { + FileUtils.SafeFileDelete(fileName); + } + + } +} diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamWriter.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamWriter.cs new file mode 100644 index 00000000..d25277ec --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsJsonFileStreamWriter.cs @@ -0,0 +1,93 @@ +// +// 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.Collections.Generic; +using System.IO; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Newtonsoft.Json; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage +{ + /// + /// Writer for writing rows of results to a JSON file. + /// + /// + /// This implements its own IDisposable because the cleanup logic closes the array that was + /// created when the writer was created. Since this behavior is different than the standard + /// file stream cleanup, the extra Dispose method was added. + /// + public class SaveAsJsonFileStreamWriter : SaveAsStreamWriter, IDisposable + { + #region Member Variables + + private readonly StreamWriter streamWriter; + private readonly JsonWriter jsonWriter; + + #endregion + + /// + /// Constructor, writes the header to the file, chains into the base constructor + /// + /// FileStream to access the JSON file output + /// JSON save as request parameters + public SaveAsJsonFileStreamWriter(Stream stream, SaveResultsRequestParams requestParams) + : base(stream, requestParams) + { + // Setup the internal state + streamWriter = new StreamWriter(stream); + jsonWriter = new JsonTextWriter(streamWriter); + + // Write the header of the file + jsonWriter.WriteStartArray(); + } + + /// + /// Writes a row of data as a JSON object + /// + /// The data of the row to output to the file + /// + /// The entire list of columns for the result set. They will be filtered down as per the + /// request params. + /// + public override void WriteRow(IList row, IList columns) + { + // Write the header for the object + jsonWriter.WriteStartObject(); + + // Write the items out as properties + int columnStart = ColumnStartIndex ?? 0; + int columnEnd = ColumnEndIndex ?? columns.Count; + for (int i = columnStart; i < columnEnd; i++) + { + jsonWriter.WritePropertyName(columns[i].ColumnName); + if (row[i].RawObject == null) + { + jsonWriter.WriteNull(); + } + else + { + jsonWriter.WriteValue(row[i].DisplayValue); + } + } + + // Write the footer for the object + jsonWriter.WriteEndObject(); + } + + /// + /// Disposes the writer by closing up the array that contains the row objects + /// + public new void Dispose() + { + // Write the footer of the file + jsonWriter.WriteEndArray(); + + jsonWriter.Close(); + streamWriter.Dispose(); + base.Dispose(); + } + } +} diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsWriterBase.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsWriterBase.cs new file mode 100644 index 00000000..38ad7eda --- /dev/null +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsWriterBase.cs @@ -0,0 +1,113 @@ +// +// 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.Collections.Generic; +using System.IO; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; + +namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage +{ + /// + /// Abstract class for implementing writers that save results to file. Stores some basic info + /// that all save as writer would need. + /// + public abstract class SaveAsStreamWriter : IFileStreamWriter + { + /// + /// Stores the internal state for the writer that will be necessary for any writer. + /// + /// The stream that will be written to + /// The SaveAs request parameters + protected SaveAsStreamWriter(Stream stream, SaveResultsRequestParams requestParams) + { + FileStream = stream; + var saveParams = requestParams; + if (requestParams.IsSaveSelection) + { + // ReSharper disable PossibleInvalidOperationException IsSaveSelection verifies these values exist + ColumnStartIndex = saveParams.ColumnStartIndex.Value; + ColumnEndIndex = saveParams.ColumnEndIndex.Value; + ColumnCount = saveParams.ColumnEndIndex.Value - saveParams.ColumnStartIndex.Value + 1; + // ReSharper restore PossibleInvalidOperationException + } + } + + #region Properties + + /// + /// Index of the first column to write to the output file + /// + protected int? ColumnStartIndex { get; private set; } + + /// + /// Number of columns to write to the output file + /// + protected int? ColumnCount { get; private set; } + + /// + /// Index of the last column to write to the output file + /// + protected int? ColumnEndIndex { get; private set; } + + /// + /// The file stream to use to write the output file + /// + protected Stream FileStream { get; private set; } + + #endregion + + /// + /// Not implemented, do not use. + /// + [Obsolete] + public int WriteRow(StorageDataReader dataReader) + { + throw new InvalidOperationException("This type of writer is meant to write values from a list of cell values only."); + } + + /// + /// Writes a row of data to the output file using the format provided by the implementing class. + /// + /// The row of data to output + /// The list of columns to output + public abstract void WriteRow(IList row, IList columns); + + /// + /// Flushes the file stream buffer + /// + public void FlushBuffer() + { + FileStream.Flush(); + } + + #region IDisposable Implementation + + private bool disposed; + + /// + /// Disposes the instance by flushing and closing the file stream + /// + /// + private void Dispose(bool disposing) + { + if (disposed || !disposing) + { + disposed = true; + return; + } + + FileStream.Flush(); + FileStream.Dispose(); + } + public virtual void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + #endregion + } +} diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamFactory.cs index 573a62d4..ab436851 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamFactory.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamFactory.cs @@ -12,6 +12,20 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// public class ServiceBufferFileStreamFactory : IFileStreamFactory { + #region Properties + + /// + /// The maximum number of characters to store from long text fields + /// + public int MaxCharsToStore { get; set; } + + /// + /// The maximum number of characters to store from xml fields + /// + public int MaxXmlCharsToStore { get; set; } + + #endregion + /// /// Creates a new temporary file /// @@ -37,12 +51,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// SSMS formatted buffer file /// /// The file to write values to - /// The maximum number of characters to store from long text fields - /// The maximum number of characters to store from xml fields /// A - public IFileStreamWriter GetWriter(string fileName, int maxCharsToStore, int maxXmlCharsToStore) + public IFileStreamWriter GetWriter(string fileName) { - return new ServiceBufferFileStreamWriter(new FileStream(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite), maxCharsToStore, maxXmlCharsToStore); + return new ServiceBufferFileStreamWriter(new FileStream(fileName, FileMode.OpenOrCreate, FileAccess.ReadWrite), MaxCharsToStore, MaxXmlCharsToStore); } /// diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamReader.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamReader.cs index e1e27635..c73bd73c 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamReader.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamReader.cs @@ -150,6 +150,58 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage return results; } + #endregion + + #region Private Helpers + + /// + /// Creates a new buffer that is of the specified length if the buffer is not already + /// at least as long as specified. + /// + /// The minimum buffer size + private void AssureBufferLength(int newBufferLength) + { + if (buffer.Length < newBufferLength) + { + buffer = new byte[newBufferLength]; + } + } + + /// + /// Reads the value of a cell from the file wrapper, checks to see if it null using + /// , and converts it to the proper output type using + /// . + /// + /// Offset into the file to read from + /// Function to use to convert the buffer to the target type + /// + /// If provided, this function will be used to determine if the value is null + /// + /// Optional function to use to convert the object to a string. + /// The expected type of the cell. Used to keep the code honest + /// The object, a display value, and the length of the value + its length + private FileStreamReadResult ReadCellHelper(long offset, Func convertFunc, Func isNullFunc = null, Func toStringFunc = null) + { + LengthResult length = ReadLength(offset); + DbCellValue result = new DbCellValue(); + + if (isNullFunc == null ? length.ValueLength == 0 : isNullFunc(length.TotalLength)) + { + result.RawObject = null; + result.DisplayValue = null; + } + else + { + AssureBufferLength(length.ValueLength); + fileStream.Read(buffer, 0, length.ValueLength); + T resultObject = convertFunc(length.ValueLength); + result.RawObject = resultObject; + result.DisplayValue = toStringFunc == null ? result.RawObject.ToString() : toStringFunc(resultObject); + } + + return new FileStreamReadResult(result, length.TotalLength); + } + /// /// Reads a short from the file at the offset provided /// @@ -317,7 +369,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage return ReadCellHelper(offset, length => { long dtTicks = BitConverter.ToInt64(buffer, 0); long dtOffset = BitConverter.ToInt64(buffer, 8); - return new DateTimeOffset(new DateTime(dtTicks), new TimeSpan(dtOffset)); + return new DateTimeOffset(new DateTime(dtTicks), new TimeSpan(dtOffset)); }); } @@ -408,7 +460,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// /// Offset into the file to read the field length from /// A LengthResult - internal LengthResult ReadLength(long offset) + private LengthResult ReadLength(long offset) { // read in length information int lengthValue; @@ -428,59 +480,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage lengthValue = BitConverter.ToInt32(buffer, 0); } - return new LengthResult {LengthLength = lengthLength, ValueLength = lengthValue}; - } - - #endregion - - #region Private Helpers - - /// - /// Creates a new buffer that is of the specified length if the buffer is not already - /// at least as long as specified. - /// - /// The minimum buffer size - private void AssureBufferLength(int newBufferLength) - { - if (buffer.Length < newBufferLength) - { - buffer = new byte[newBufferLength]; - } - } - - /// - /// Reads the value of a cell from the file wrapper, checks to see if it null using - /// , and converts it to the proper output type using - /// . - /// - /// Offset into the file to read from - /// Function to use to convert the buffer to the target type - /// - /// If provided, this function will be used to determine if the value is null - /// - /// Optional function to use to convert the object to a string. - /// The expected type of the cell. Used to keep the code honest - /// The object, a display value, and the length of the value + its length - private FileStreamReadResult ReadCellHelper(long offset, Func convertFunc, Func isNullFunc = null, Func toStringFunc = null) - { - LengthResult length = ReadLength(offset); - DbCellValue result = new DbCellValue(); - - if (isNullFunc == null ? length.ValueLength == 0 : isNullFunc(length.TotalLength)) - { - result.RawObject = null; - result.DisplayValue = null; - } - else - { - AssureBufferLength(length.ValueLength); - fileStream.Read(buffer, 0, length.ValueLength); - T resultObject = convertFunc(length.ValueLength); - result.RawObject = resultObject; - result.DisplayValue = toStringFunc == null ? result.RawObject.ToString() : toStringFunc(resultObject); - } - - return new FileStreamReadResult(result, length.TotalLength); + return new LengthResult { LengthLength = lengthLength, ValueLength = lengthValue }; } #endregion diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamWriter.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamWriter.cs index e9f280a4..c8eb536e 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamWriter.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/ServiceBufferFileStreamWriter.cs @@ -207,115 +207,133 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage return rowBytes; } + [Obsolete] + public void WriteRow(IList row, IList columns) + { + throw new InvalidOperationException("This type of writer is meant to write values from a DbDataReader only."); + } + + /// + /// Flushes the internal buffer to the file stream + /// + public void FlushBuffer() + { + fileStream.Flush(); + } + + #endregion + + #region Private Helpers + /// /// Writes null to the file as one 0x00 byte /// /// Number of bytes used to store the null - public int WriteNull() + internal int WriteNull() { byteBuffer[0] = 0x00; - return WriteHelper(byteBuffer, 1); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 1); } /// /// Writes a short to the file /// /// Number of bytes used to store the short - public int WriteInt16(short val) + internal int WriteInt16(short val) { byteBuffer[0] = 0x02; // length shortBuffer[0] = val; Buffer.BlockCopy(shortBuffer, 0, byteBuffer, 1, 2); - return WriteHelper(byteBuffer, 3); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 3); } /// /// Writes a int to the file /// /// Number of bytes used to store the int - public int WriteInt32(int val) + internal int WriteInt32(int val) { byteBuffer[0] = 0x04; // length intBuffer[0] = val; Buffer.BlockCopy(intBuffer, 0, byteBuffer, 1, 4); - return WriteHelper(byteBuffer, 5); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 5); } /// /// Writes a long to the file /// /// Number of bytes used to store the long - public int WriteInt64(long val) + internal int WriteInt64(long val) { byteBuffer[0] = 0x08; // length longBuffer[0] = val; Buffer.BlockCopy(longBuffer, 0, byteBuffer, 1, 8); - return WriteHelper(byteBuffer, 9); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 9); } /// /// Writes a char to the file /// /// Number of bytes used to store the char - public int WriteChar(char val) + internal int WriteChar(char val) { byteBuffer[0] = 0x02; // length charBuffer[0] = val; Buffer.BlockCopy(charBuffer, 0, byteBuffer, 1, 2); - return WriteHelper(byteBuffer, 3); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 3); } /// /// Writes a bool to the file /// /// Number of bytes used to store the bool - public int WriteBoolean(bool val) + internal int WriteBoolean(bool val) { byteBuffer[0] = 0x01; // length byteBuffer[1] = (byte) (val ? 0x01 : 0x00); - return WriteHelper(byteBuffer, 2); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 2); } /// /// Writes a byte to the file /// /// Number of bytes used to store the byte - public int WriteByte(byte val) + internal int WriteByte(byte val) { byteBuffer[0] = 0x01; // length byteBuffer[1] = val; - return WriteHelper(byteBuffer, 2); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 2); } /// /// Writes a float to the file /// /// Number of bytes used to store the float - public int WriteSingle(float val) + internal int WriteSingle(float val) { byteBuffer[0] = 0x04; // length floatBuffer[0] = val; Buffer.BlockCopy(floatBuffer, 0, byteBuffer, 1, 4); - return WriteHelper(byteBuffer, 5); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 5); } /// /// Writes a double to the file /// /// Number of bytes used to store the double - public int WriteDouble(double val) + internal int WriteDouble(double val) { byteBuffer[0] = 0x08; // length doubleBuffer[0] = val; Buffer.BlockCopy(doubleBuffer, 0, byteBuffer, 1, 8); - return WriteHelper(byteBuffer, 9); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 9); } /// /// Writes a SqlDecimal to the file /// /// Number of bytes used to store the SqlDecimal - public int WriteSqlDecimal(SqlDecimal val) + internal int WriteSqlDecimal(SqlDecimal val) { int[] arrInt32 = val.Data; int iLen = 3 + (arrInt32.Length * 4); @@ -332,7 +350,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage // data value Buffer.BlockCopy(arrInt32, 0, byteBuffer, 3, iLen - 3); - iTotalLen += WriteHelper(byteBuffer, iLen); + iTotalLen += FileUtils.WriteWithLength(fileStream, byteBuffer, iLen); return iTotalLen; // len+data } @@ -340,7 +358,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// Writes a decimal to the file /// /// Number of bytes used to store the decimal - public int WriteDecimal(decimal val) + internal int WriteDecimal(decimal val) { int[] arrInt32 = decimal.GetBits(val); @@ -348,7 +366,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage int iTotalLen = WriteLength(iLen); // length Buffer.BlockCopy(arrInt32, 0, byteBuffer, 0, iLen); - iTotalLen += WriteHelper(byteBuffer, iLen); + iTotalLen += FileUtils.WriteWithLength(fileStream, byteBuffer, iLen); return iTotalLen; // len+data } @@ -366,7 +384,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// Writes a DateTimeOffset to the file /// /// Number of bytes used to store the DateTimeOffset - public int WriteDateTimeOffset(DateTimeOffset dtoVal) + internal int WriteDateTimeOffset(DateTimeOffset dtoVal) { // Write the length, which is the 2*sizeof(long) byteBuffer[0] = 0x10; // length (16) @@ -376,14 +394,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage longBufferOffset[0] = dtoVal.Ticks; longBufferOffset[1] = dtoVal.Offset.Ticks; Buffer.BlockCopy(longBufferOffset, 0, byteBuffer, 1, 16); - return WriteHelper(byteBuffer, 17); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 17); } /// /// Writes a TimeSpan to the file /// /// Number of bytes used to store the TimeSpan - public int WriteTimeSpan(TimeSpan timeSpan) + internal int WriteTimeSpan(TimeSpan timeSpan) { return WriteInt64(timeSpan.Ticks); } @@ -392,7 +410,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// Writes a string to the file /// /// Number of bytes used to store the string - public int WriteString(string sVal) + internal int WriteString(string sVal) { Validate.IsNotNull(nameof(sVal), sVal); @@ -408,7 +426,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage byteBuffer[3] = 0x00; byteBuffer[4] = 0x00; - iTotalLen = WriteHelper(byteBuffer, 5); + iTotalLen = FileUtils.WriteWithLength(fileStream, byteBuffer, 5); } else { @@ -417,7 +435,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage // convert char array into byte array and write it out iTotalLen = WriteLength(bytes.Length); - iTotalLen += WriteHelper(bytes, bytes.Length); + iTotalLen += FileUtils.WriteWithLength(fileStream, bytes, bytes.Length); } return iTotalLen; // len+data } @@ -426,7 +444,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// Writes a byte[] to the file /// /// Number of bytes used to store the byte[] - public int WriteBytes(byte[] bytesVal) + internal int WriteBytes(byte[] bytesVal) { Validate.IsNotNull(nameof(bytesVal), bytesVal); @@ -440,12 +458,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage byteBuffer[3] = 0x00; byteBuffer[4] = 0x00; - iTotalLen = WriteHelper(byteBuffer, 5); + iTotalLen = FileUtils.WriteWithLength(fileStream, byteBuffer, 5); } else { iTotalLen = WriteLength(bytesVal.Length); - iTotalLen += WriteHelper(bytesVal, bytesVal.Length); + iTotalLen += FileUtils.WriteWithLength(fileStream, bytesVal, bytesVal.Length); } return iTotalLen; // len+data } @@ -455,7 +473,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// /// The GUID to write to the file /// Number of bytes written to the file - public int WriteGuid(Guid val) + internal int WriteGuid(Guid val) { byte[] guidBytes = val.ToByteArray(); return WriteBytes(guidBytes); @@ -466,23 +484,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage /// /// The SqlMoney value to write to the file /// Number of bytes written to the file - public int WriteMoney(SqlMoney val) + internal int WriteMoney(SqlMoney val) { return WriteDecimal(val.Value); } - /// - /// Flushes the internal buffer to the file stream - /// - public void FlushBuffer() - { - fileStream.Flush(); - } - - #endregion - - #region Private Helpers - /// /// Creates a new buffer that is of the specified length if the buffer is not already /// at least as long as specified. @@ -509,7 +515,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage int iTmp = iLen & 0x000000FF; byteBuffer[0] = Convert.ToByte(iTmp); - return WriteHelper(byteBuffer, 1); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 1); } // The length won't fit in 1 byte, so we need to use 1 byte to signify that the length // is a full 4 bytes. @@ -518,7 +524,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage // convert int32 into array of bytes intBuffer[0] = iLen; Buffer.BlockCopy(intBuffer, 0, byteBuffer, 1, 4); - return WriteHelper(byteBuffer, 5); + return FileUtils.WriteWithLength(fileStream, byteBuffer, 5); } /// @@ -534,12 +540,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage return val.IsNull ? WriteNull() : valueWriteFunc(val); } - private int WriteHelper(byte[] buffer, int length) - { - fileStream.Write(buffer, 0, length); - return length; - } - #endregion #region IDisposable Implementation diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs index 3c5bae8c..d5948a8a 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs @@ -209,6 +209,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution cancellationSource.Cancel(); } + /// + /// Launches the asynchronous process for executing the query + /// public void Execute() { ExecutionTask = Task.Run(ExecuteInternal); @@ -233,6 +236,27 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution return Batches[batchIndex].GetSubset(resultSetIndex, startRow, rowCount); } + /// + /// Saves the requested results to a file format of the user's choice + /// + /// Parameters for the save as request + /// + /// Factory for creating the reader/writer pair for the requested output format + /// + /// Delegate to call when the request completes successfully + /// Delegate to call if the request fails + public void SaveAs(SaveResultsRequestParams saveParams, IFileStreamFactory fileFactory, + ResultSet.SaveAsAsyncEventHandler successHandler, ResultSet.SaveAsFailureAsyncEventHandler failureHandler) + { + // Sanity check to make sure that the batch is within bounds + if (saveParams.BatchIndex < 0 || saveParams.BatchIndex >= Batches.Length) + { + throw new ArgumentOutOfRangeException(nameof(saveParams.BatchIndex), SR.QueryServiceSubsetBatchOutOfRange); + } + + Batches[saveParams.BatchIndex].SaveAs(saveParams, fileFactory, successHandler, failureHandler); + } + #endregion #region Private Helpers diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs index 4bd9f1fa..b65808cc 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs @@ -4,6 +4,7 @@ // using System; using System.Collections.Concurrent; +using System.IO; using System.Linq; using System.Threading.Tasks; using Microsoft.SqlTools.ServiceLayer.Connection; @@ -64,9 +65,28 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution /// private IFileStreamFactory BufferFileFactory { - get { return BufferFileStreamFactory ?? (BufferFileStreamFactory = new ServiceBufferFileStreamFactory()); } + get + { + return BufferFileStreamFactory ?? (BufferFileStreamFactory = new ServiceBufferFileStreamFactory + { + MaxCharsToStore = Settings.SqlTools.QueryExecutionSettings.MaxCharsToStore, + MaxXmlCharsToStore = Settings.SqlTools.QueryExecutionSettings.MaxXmlCharsToStore + }); + } } + /// + /// File factory to be used to create CSV files from result sets. Set to internal in order + /// to allow overriding in unit testing + /// + internal IFileStreamFactory CsvFileFactory { get; set; } + + /// + /// File factory to be used to create JSON files from result sets. Set to internal in order + /// to allow overriding in unit testing + /// + internal IFileStreamFactory JsonFileFactory { get; set; } + /// /// The collection of active queries /// @@ -124,6 +144,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution #region Request Handlers + /// + /// Handles request to execute the query + /// public async Task HandleExecuteRequest(QueryExecuteParams executeParams, RequestContext requestContext) { @@ -134,6 +157,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution await ExecuteAndCompleteQuery(executeParams, requestContext, newQuery); } + /// + /// Handles a request to get a subset of the results of this query + /// public async Task HandleResultSubsetRequest(QueryExecuteSubsetParams subsetParams, RequestContext requestContext) { @@ -182,6 +208,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } } + /// + /// Handles a request to dispose of this query + /// public async Task HandleDisposeRequest(QueryDisposeParams disposeParams, RequestContext requestContext) { @@ -213,6 +242,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } } + /// + /// Handles a request to cancel this query if it is in progress + /// public async Task HandleCancelRequest(QueryCancelParams cancelParams, RequestContext requestContext) { @@ -253,43 +285,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution internal async Task HandleSaveResultsAsCsvRequest(SaveResultsAsCsvRequestParams saveParams, RequestContext requestContext) { - // retrieve query for OwnerUri - Query result; - if (!ActiveQueries.TryGetValue(saveParams.OwnerUri, out result)) + // Use the default CSV file factory if we haven't overridden it + IFileStreamFactory csvFactory = CsvFileFactory ?? new SaveAsCsvFileStreamFactory { - await requestContext.SendResult(new SaveResultRequestResult - { - Messages = SR.QueryServiceRequestsNoQuery - }); - return; - } - - - ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex]; - if (!selectedResultSet.IsBeingDisposed) - { - // Create SaveResults object and add success and error handlers to respective events - SaveResults saveAsCsv = new SaveResults(); - - SaveResults.AsyncSaveEventHandler successHandler = async message => - { - selectedResultSet.RemoveSaveTask(saveParams.FilePath); - await requestContext.SendResult(new SaveResultRequestResult { Messages = message }); - }; - saveAsCsv.SaveCompleted += successHandler; - SaveResults.AsyncSaveEventHandler errorHandler = async message => - { - selectedResultSet.RemoveSaveTask(saveParams.FilePath); - await requestContext.SendError(new SaveResultRequestError { message = message }); - }; - saveAsCsv.SaveFailed += errorHandler; - - saveAsCsv.SaveResultSetAsCsv(saveParams, requestContext, result); - - // Associate the ResultSet with the save task - selectedResultSet.AddSaveTask(saveParams.FilePath, saveAsCsv.SaveTask); - - } + SaveRequestParams = saveParams + }; + await SaveResultsHelper(saveParams, requestContext, csvFactory); } /// @@ -298,41 +299,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution internal async Task HandleSaveResultsAsJsonRequest(SaveResultsAsJsonRequestParams saveParams, RequestContext requestContext) { - // retrieve query for OwnerUri - Query result; - if (!ActiveQueries.TryGetValue(saveParams.OwnerUri, out result)) + // Use the default JSON file factory if we haven't overridden it + IFileStreamFactory jsonFactory = JsonFileFactory ?? new SaveAsJsonFileStreamFactory { - await requestContext.SendResult(new SaveResultRequestResult - { - Messages = "Failed to save results, ID not found." - }); - return; - } - - ResultSet selectedResultSet = result.Batches[saveParams.BatchIndex].ResultSets[saveParams.ResultSetIndex]; - if (!selectedResultSet.IsBeingDisposed) - { - // Create SaveResults object and add success and error handlers to respective events - SaveResults saveAsJson = new SaveResults(); - SaveResults.AsyncSaveEventHandler successHandler = async message => - { - selectedResultSet.RemoveSaveTask(saveParams.FilePath); - await requestContext.SendResult(new SaveResultRequestResult { Messages = message }); - }; - saveAsJson.SaveCompleted += successHandler; - SaveResults.AsyncSaveEventHandler errorHandler = async message => - { - selectedResultSet.RemoveSaveTask(saveParams.FilePath); - await requestContext.SendError(new SaveResultRequestError { message = message }); - }; - saveAsJson.SaveFailed += errorHandler; - - saveAsJson.SaveResultSetAsJson(saveParams, requestContext, result); - - // Associate the ResultSet with the save task - selectedResultSet.AddSaveTask(saveParams.FilePath, saveAsJson.SaveTask); - } - + SaveRequestParams = saveParams + }; + await SaveResultsHelper(saveParams, requestContext, jsonFactory); } #endregion @@ -404,7 +376,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution await requestContext.SendError(e.Message); return null; } - // Any other exceptions will fall through here and be collected at the end } private static async Task ExecuteAndCompleteQuery(QueryExecuteParams executeParams, RequestContext requestContext, Query query) @@ -494,6 +465,42 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution }); } + private async Task SaveResultsHelper(SaveResultsRequestParams saveParams, + RequestContext requestContext, IFileStreamFactory fileFactory) + { + // retrieve query for OwnerUri + Query query; + if (!ActiveQueries.TryGetValue(saveParams.OwnerUri, out query)) + { + await requestContext.SendError(new SaveResultRequestError + { + message = SR.QueryServiceQueryInvalidOwnerUri + }); + return; + } + + //Setup the callback for completion of the save task + ResultSet.SaveAsAsyncEventHandler successHandler = async parameters => + { + await requestContext.SendResult(new SaveResultRequestResult()); + }; + ResultSet.SaveAsFailureAsyncEventHandler errorHandler = async (parameters, reason) => + { + string message = SR.QueryServiceSaveAsFail(Path.GetFileName(parameters.FilePath), reason); + await requestContext.SendError(new SaveResultRequestError { message = message }); + }; + + try + { + // Launch the task + query.SaveAs(saveParams, fileFactory, successHandler, errorHandler); + } + catch (Exception e) + { + await errorHandler(saveParams, e.Message); + } + } + #endregion #region IDisposable Implementation diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs index e3ecb890..8037178a 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs @@ -24,16 +24,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution { #region Constants - private const int DefaultMaxCharsToStore = 65535; // 64 KB - QE default - - // xml is a special case so number of chars to store is usually greater than for other long types - private const int DefaultMaxXmlCharsToStore = 2097152; // 2 MB - QE default - // Column names of 'for xml' and 'for json' queries private const string NameOfForXMLColumn = "XML_F52E2B61-18A1-11d1-B105-00805F49916B"; private const string NameOfForJSONColumn = "JSON_F52E2B61-18A1-11d1-B105-00805F49916B"; - #endregion #region Member Variables @@ -74,11 +68,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution /// private readonly string outputFileName; - /// - /// All save tasks currently saving this ResultSet - /// - private readonly ConcurrentDictionary saveTasks; - #endregion /// @@ -104,11 +93,24 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution // Store the factory fileStreamFactory = factory; hasBeenRead = false; - saveTasks = new ConcurrentDictionary(); + SaveTasks = new ConcurrentDictionary(); } #region Properties + /// + /// Asynchronous handler for when saving query results succeeds + /// + /// Request parameters for identifying the request + public delegate Task SaveAsAsyncEventHandler(SaveResultsRequestParams parameters); + + /// + /// Asynchronous handler for when saving query results fails + /// + /// Request parameters for identifying the request + /// Message to send back describing why the request failed + public delegate Task SaveAsFailureAsyncEventHandler(SaveResultsRequestParams parameters, string message); + /// /// Asynchronous handler for when a resultset has completed /// @@ -141,21 +143,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution /// public int BatchId { get; private set; } - /// - /// Maximum number of characters to store for a field - /// - public int MaxCharsToStore { get { return DefaultMaxCharsToStore; } } - - /// - /// Maximum number of characters to store for an XML field - /// - public int MaxXmlCharsToStore { get { return DefaultMaxXmlCharsToStore; } } - /// /// The number of rows for this result set /// public long RowCount { get; private set; } + /// + /// All save tasks currently saving this ResultSet + /// + internal ConcurrentDictionary SaveTasks { get; set; } + /// /// Generates a summary of this result set /// @@ -251,7 +248,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution hasBeenRead = true; // Open a writer for the file - var fileWriter = fileStreamFactory.GetWriter(outputFileName, MaxCharsToStore, MaxCharsToStore); + var fileWriter = fileStreamFactory.GetWriter(outputFileName); using (fileWriter) { // If we can initialize the columns using the column schema, use that @@ -282,6 +279,89 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } } + public void SaveAs(SaveResultsRequestParams saveParams, IFileStreamFactory fileFactory, + SaveAsAsyncEventHandler successHandler, SaveAsFailureAsyncEventHandler failureHandler) + { + // Sanity check the save params and file factory + Validate.IsNotNull(nameof(saveParams), saveParams); + Validate.IsNotNull(nameof(fileFactory), fileFactory); + + // Make sure the resultset has finished being read + if (!hasBeenRead) + { + throw new InvalidOperationException(SR.QueryServiceSaveAsResultSetNotComplete); + } + + // Make sure there isn't a task for this file already + Task existingTask; + if (SaveTasks.TryGetValue(saveParams.FilePath, out existingTask)) + { + if (existingTask.IsCompleted) + { + // The task has completed, so let's attempt to remove it + if (!SaveTasks.TryRemove(saveParams.FilePath, out existingTask)) + { + throw new InvalidOperationException(SR.QueryServiceSaveAsMiscStartingError); + } + } + else + { + // The task hasn't completed, so we shouldn't continue + throw new InvalidOperationException(SR.QueryServiceSaveAsInProgress); + } + } + + // Create the new task + Task saveAsTask = new Task(async () => + { + try + { + // Set row counts depending on whether save request is for entire set or a subset + long rowEndIndex = RowCount; + int rowStartIndex = 0; + if (saveParams.IsSaveSelection) + { + // ReSharper disable PossibleInvalidOperationException IsSaveSelection verifies these values exist + rowEndIndex = saveParams.RowEndIndex.Value + 1; + rowStartIndex = saveParams.RowStartIndex.Value; + // ReSharper restore PossibleInvalidOperationException + } + + using (var fileReader = fileFactory.GetReader(outputFileName)) + using (var fileWriter = fileFactory.GetWriter(saveParams.FilePath)) + { + // Iterate over the rows that are in the selected row set + for (long i = rowStartIndex; i < rowEndIndex; ++i) + { + var row = fileReader.ReadRow(fileOffsets[i], Columns); + fileWriter.WriteRow(row, Columns); + } + if (successHandler != null) + { + await successHandler(saveParams); + } + } + } + catch (Exception e) + { + fileFactory.DisposeFile(saveParams.FilePath); + if (failureHandler != null) + { + await failureHandler(saveParams, e.Message); + } + } + }); + + // If saving the task fails, return a failure + if (!SaveTasks.TryAdd(saveParams.FilePath, saveAsTask)) + { + throw new InvalidOperationException(SR.QueryServiceSaveAsMiscStartingError); + } + + // Task was saved, so start up the task + saveAsTask.Start(); + } + #endregion #region IDisposable Implementation @@ -301,10 +381,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution IsBeingDisposed = true; // Check if saveTasks are running for this ResultSet - if (!saveTasks.IsEmpty) + if (!SaveTasks.IsEmpty) { // Wait for tasks to finish before disposing ResultSet - Task.WhenAll(saveTasks.Values.ToArray()).ContinueWith((antecedent) => + Task.WhenAll(SaveTasks.Values.ToArray()).ContinueWith((antecedent) => { if (disposing) { @@ -357,25 +437,5 @@ 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 } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs deleted file mode 100644 index 0490f6d4..00000000 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs +++ /dev/null @@ -1,319 +0,0 @@ -// -// 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; -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 - { - /// - /// Number of rows being read from the ResultSubset in one read - /// - private const int BatchSize = 1000; - - /// - /// Save Task that asynchronously writes ResultSet to file - /// - internal Task SaveTask { get; set; } - - /// - /// Event Handler for save events - /// - /// Message to be returned to client - /// - internal delegate Task AsyncSaveEventHandler(string message); - - /// - /// A successful save event - /// - internal event AsyncSaveEventHandler SaveCompleted; - - /// - /// A failed save event - /// - internal event AsyncSaveEventHandler SaveFailed; - - /// Method ported from SSMS - /// - /// Encodes a single field for inserting into a CSV record. The following rules are applied: - /// - /// All double quotes (") are replaced with a pair of consecutive double quotes - /// - /// The entire field is also surrounded by a pair of double quotes if any of the following conditions are met: - /// - /// The field begins or ends with a space - /// The field begins or ends with a tab - /// The field contains the ListSeparator string - /// The field contains the '\n' character - /// The field contains the '\r' character - /// The field contains the '"' character - /// - /// - /// The field to encode - /// The CSV encoded version of the original field - 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; - - //Check for leading/trailing spaces - if (sbField.Length > 0 && - (sbField[0] == ' ' || - sbField[0] == '\t' || - sbField[sbField.Length - 1] == ' ' || - sbField[sbField.Length - 1] == '\t')) - { - embedInQuotes = true; - } - else - { //List separator being in the field will require quotes - if (field.Contains(",")) - { - embedInQuotes = true; - } - else - { - for (int i = 0; i < sbField.Length; ++i) - { - //Check whether this character is a special character - if (sbField[i] == '\r' || - sbField[i] == '\n' || - sbField[i] == '"') - { //If even one character requires embedding the whole field will - //be embedded in quotes so we can just break out now - embedInQuotes = true; - break; - } - } - } - } - - //Replace all quotes in the original field with double quotes - sbField.Replace("\"", "\"\""); - string ret = sbField.ToString(); - - if (embedInQuotes) - { - ret = "\"" + ret + "\""; - } - - return ret; - } - - /// - /// Check if request is a subset of result set or whole result set - /// - /// Parameters from the request - /// - internal static bool IsSaveSelection(SaveResultsRequestParams saveParams) - { - return (saveParams.ColumnStartIndex != null && saveParams.ColumnEndIndex != null - && saveParams.RowStartIndex != null && saveParams.RowEndIndex != null); - } - - /// - /// Save results as JSON format to the file specified in saveParams - /// - /// Parameters from the request - /// Request context for save results - /// Result query object - /// - internal void SaveResultSetAsJson(SaveResultsAsJsonRequestParams saveParams, RequestContext 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.Message); - } - } - }); - } - - /// - /// Save results as CSV format to the file specified in saveParams - /// - /// Parameters from the request - /// Request context for save results - /// Result query object - /// - internal void SaveResultSetAsCsv(SaveResultsAsCsvRequestParams saveParams, RequestContext 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); - } - } - }); - } - - - } - -} \ No newline at end of file diff --git a/src/Microsoft.SqlTools.ServiceLayer/SqlContext/QueryExecutionSettings.cs b/src/Microsoft.SqlTools.ServiceLayer/SqlContext/QueryExecutionSettings.cs index 4934a4da..a573240c 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/SqlContext/QueryExecutionSettings.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/SqlContext/QueryExecutionSettings.cs @@ -1,8 +1,6 @@ -// +// // Copyright (c) Microsoft. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. -// - namespace Microsoft.SqlTools.ServiceLayer.SqlContext { /// @@ -10,13 +8,38 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext /// public class QueryExecutionSettings { + #region Constants + /// /// Default value for batch separator (de facto standard as per SSMS) /// private const string DefaultBatchSeparator = "GO"; + /// + /// Default number of chars to store for long text fields (de facto standard as per SSMS) + /// + private const int DefaultMaxCharsToStore = 65535; // 64 KB - QE default + + /// + /// Default number of chars to store of XML values (de facto standard as per SSMS) + /// xml is a special case so number of chars to store is usually greater than for other long types + /// + private const int DefaultMaxXmlCharsToStore = 2097152; // 2 MB - QE default + + #endregion + + #region Member Variables + private string batchSeparator; + private int? maxCharsToStore; + + private int? maxXmlCharsToStore; + + #endregion + + #region Properties + /// /// The configured batch separator, will use a default if a value was not configured /// @@ -26,6 +49,22 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext set { batchSeparator = value; } } + public int MaxCharsToStore + { + get { return maxCharsToStore ?? DefaultMaxCharsToStore; } + set { maxCharsToStore = value; } + } + + public int MaxXmlCharsToStore + { + get { return maxXmlCharsToStore ?? DefaultMaxXmlCharsToStore; } + set { maxXmlCharsToStore = value; } + } + + #endregion + + #region Public Methods + /// /// Update the current settings with the new settings /// @@ -33,6 +72,10 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext public void Update(QueryExecutionSettings newSettings) { BatchSeparator = newSettings.BatchSeparator; + MaxCharsToStore = newSettings.MaxCharsToStore; + MaxXmlCharsToStore = newSettings.MaxXmlCharsToStore; } + + #endregion } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/Utility/FileUtils.cs b/src/Microsoft.SqlTools.ServiceLayer/Utility/FileUtils.cs index 6f8a3fb8..75290fd7 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Utility/FileUtils.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Utility/FileUtils.cs @@ -82,6 +82,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution } } + internal static int WriteWithLength(Stream stream, byte[] buffer, int length) + { + stream.Write(buffer, 0, length); + return length; + } + /// /// Checks if file exists and swallows exceptions, if any /// diff --git a/src/Microsoft.SqlTools.ServiceLayer/sr.cs b/src/Microsoft.SqlTools.ServiceLayer/sr.cs index 9c162533..1fba1d66 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/sr.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/sr.cs @@ -13,7 +13,7 @@ namespace Microsoft.SqlTools.ServiceLayer [System.Runtime.CompilerServices.CompilerGeneratedAttribute()] internal class SR { - internal SR() + protected SR() { } public static CultureInfo Culture @@ -309,6 +309,30 @@ namespace Microsoft.SqlTools.ServiceLayer } } + public static string QueryServiceSaveAsResultSetNotComplete + { + get + { + return Keys.GetString(Keys.QueryServiceSaveAsResultSetNotComplete); + } + } + + public static string QueryServiceSaveAsMiscStartingError + { + get + { + return Keys.GetString(Keys.QueryServiceSaveAsMiscStartingError); + } + } + + public static string QueryServiceSaveAsInProgress + { + get + { + return Keys.GetString(Keys.QueryServiceSaveAsInProgress); + } + } + public static string QueryServiceResultSetNotRead { get @@ -416,6 +440,11 @@ namespace Microsoft.SqlTools.ServiceLayer return Keys.GetString(Keys.QueryServiceQueryFailed, message); } + public static string QueryServiceSaveAsFail(string fileName, string message) + { + return Keys.GetString(Keys.QueryServiceSaveAsFail, fileName, message); + } + public static string PeekDefinitionAzureError(string errorMessage) { return Keys.GetString(Keys.PeekDefinitionAzureError, errorMessage); @@ -570,6 +599,18 @@ namespace Microsoft.SqlTools.ServiceLayer public const string QueryServiceResultSetReaderNull = "QueryServiceResultSetReaderNull"; + public const string QueryServiceSaveAsResultSetNotComplete = "QueryServiceSaveAsResultSetNotComplete"; + + + public const string QueryServiceSaveAsMiscStartingError = "QueryServiceSaveAsMiscStartingError"; + + + public const string QueryServiceSaveAsInProgress = "QueryServiceSaveAsInProgress"; + + + public const string QueryServiceSaveAsFail = "QueryServiceSaveAsFail"; + + public const string QueryServiceResultSetNotRead = "QueryServiceResultSetNotRead"; @@ -609,7 +650,7 @@ namespace Microsoft.SqlTools.ServiceLayer public const string WorkspaceServiceBufferPositionOutOfOrder = "WorkspaceServiceBufferPositionOutOfOrder"; - internal Keys() + private Keys() { } public static CultureInfo Culture @@ -636,6 +677,12 @@ namespace Microsoft.SqlTools.ServiceLayer } + public static string GetString(string key, object arg0, object arg1) + { + return string.Format(global::System.Globalization.CultureInfo.CurrentCulture, resourceManager.GetString(key, _culture), arg0, arg1); + } + + public static string GetString(string key, object arg0, object arg1, object arg2, object arg3) { return string.Format(global::System.Globalization.CultureInfo.CurrentCulture, resourceManager.GetString(key, _culture), arg0, arg1, arg2, arg3); diff --git a/src/Microsoft.SqlTools.ServiceLayer/sr.resx b/src/Microsoft.SqlTools.ServiceLayer/sr.resx index bf5ed353..955b2e6e 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/sr.resx +++ b/src/Microsoft.SqlTools.ServiceLayer/sr.resx @@ -292,6 +292,23 @@ Reader cannot be null + + Result cannot be saved until query execution has completed + + + + Internal error occurred while starting save task + + + + A save request to the same path is in progress + + + + Failed to save {0}: {1} + . + Parameters: 0 - fileName (string), 1 - message (string) + Cannot read subset unless the results have been read from the server diff --git a/src/Microsoft.SqlTools.ServiceLayer/sr.strings b/src/Microsoft.SqlTools.ServiceLayer/sr.strings index fee7f907..ac86e2b6 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/sr.strings +++ b/src/Microsoft.SqlTools.ServiceLayer/sr.strings @@ -129,7 +129,17 @@ QueryServiceMessageSenderNotSql = Sender for OnInfoMessage event must be a SqlCo QueryServiceResultSetReaderNull = Reader cannot be null -### MSC +### Save As Requests + +QueryServiceSaveAsResultSetNotComplete = Result cannot be saved until query execution has completed + +QueryServiceSaveAsMiscStartingError = Internal error occurred while starting save task + +QueryServiceSaveAsInProgress = A save request to the same path is in progress + +QueryServiceSaveAsFail(string fileName, string message) = Failed to save {0}: {1} + +### MISC QueryServiceResultSetNotRead = Cannot read subset unless the results have been read from the server diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs index 7923a554..4a5ead4b 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs @@ -123,9 +123,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution }); mock.Setup(fsf => fsf.GetReader(It.IsAny())) .Returns(output => new ServiceBufferFileStreamReader(new MemoryStream(storage[output]))); - mock.Setup(fsf => fsf.GetWriter(It.IsAny(), It.IsAny(), It.IsAny())) - .Returns((output, chars, xml) => new ServiceBufferFileStreamWriter( - new MemoryStream(storage[output]), chars, xml)); + mock.Setup(fsf => fsf.GetWriter(It.IsAny())) + .Returns(output => new ServiceBufferFileStreamWriter(new MemoryStream(storage[output]), 1024, 1024)); return mock.Object; } @@ -188,10 +187,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution #region Service Mocking - public static QueryExecutionService GetPrimedExecutionService(Dictionary[][] data, bool isConnected, bool throwOnRead, WorkspaceService workspaceService) + public static QueryExecutionService GetPrimedExecutionService(Dictionary[][] data, + bool isConnected, bool throwOnRead, WorkspaceService workspaceService, + out Dictionary storage) { // Create a place for the temp "files" to be written - Dictionary storage = new Dictionary(); + storage = new Dictionary(); // Create the connection factory with the dataset var factory = CreateTestConnectionInfo(data, throwOnRead).Factory; @@ -205,7 +206,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution .OutCallback((string owner, out ConnectionInfo connInfo) => connInfo = isConnected ? ci : null) .Returns(isConnected); - return new QueryExecutionService(connectionService.Object, workspaceService) {BufferFileStreamFactory = GetFileStreamFactory(storage)}; + return new QueryExecutionService(connectionService.Object, workspaceService) { BufferFileStreamFactory = GetFileStreamFactory(storage) }; + } + + public static QueryExecutionService GetPrimedExecutionService(Dictionary[][] data, bool isConnected, bool throwOnRead, WorkspaceService workspaceService) + { + Dictionary storage; + return GetPrimedExecutionService(data, isConnected, throwOnRead, workspaceService, out storage); } public static WorkspaceService GetPrimedWorkspaceService(string query) diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsCsvFileStreamWriterTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsCsvFileStreamWriterTests.cs new file mode 100644 index 00000000..55c1a75b --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsCsvFileStreamWriterTests.cs @@ -0,0 +1,216 @@ +// +// 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.Collections.Generic; +using System.IO; +using System.Text; +using System.Text.RegularExpressions; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage; +using Microsoft.SqlTools.ServiceLayer.Test.Utility; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.DataStorage +{ + public class SaveAsCsvFileStreamWriterTests + { + [Theory] + [InlineData("Something\rElse")] + [InlineData("Something\nElse")] + [InlineData("Something\"Else")] + [InlineData("Something,Else")] + [InlineData("\tSomething")] + [InlineData("Something\t")] + [InlineData(" Something")] + [InlineData("Something ")] + [InlineData(" \t\r\n\",\r\n\"\r ")] + public void EncodeCsvFieldShouldWrap(string field) + { + // If: I CSV encode a field that has forbidden characters in it + string output = SaveAsCsvFileStreamWriter.EncodeCsvField(field); + + // Then: It should wrap it in quotes + Assert.True(Regex.IsMatch(output, "^\".*") + && Regex.IsMatch(output, ".*\"$")); + } + + [Theory] + [InlineData("Something")] + [InlineData("Something valid.")] + [InlineData("Something\tvalid")] + public void EncodeCsvFieldShouldNotWrap(string field) + { + // If: I CSV encode a field that does not have forbidden characters in it + string output = SaveAsCsvFileStreamWriter.EncodeCsvField(field); + + // Then: It should not wrap it in quotes + Assert.False(Regex.IsMatch(output, "^\".*\"$")); + } + + [Fact] + public void EncodeCsvFieldReplace() + { + // If: I CSV encode a field that has a double quote in it, + string output = SaveAsCsvFileStreamWriter.EncodeCsvField("Some\"thing"); + + // Then: It should be replaced with double double quotes + Assert.Equal("\"Some\"\"thing\"", output); + } + + [Fact] + public void EncodeCsvFieldNull() + { + // If: I CSV encode a null + string output = SaveAsCsvFileStreamWriter.EncodeCsvField(null); + + // Then: there should be a string version of null returned + Assert.Equal("NULL", output); + } + + [Fact] + public void WriteRowWithoutColumnSelectionOrHeader() + { + // Setup: + // ... Create a request params that has no selection made + // ... Create a set of data to write + // ... Create a memory location to store the data + var requestParams = new SaveResultsAsCsvRequestParams(); + List data = new List + { + new DbCellValue { DisplayValue = "item1" }, + new DbCellValue { DisplayValue = "item2" } + }; + List columns = new List + { + new DbColumnWrapper(new TestDbColumn("column1")), + new DbColumnWrapper(new TestDbColumn("column2")) + }; + byte[] output = new byte[8192]; + + // If: I write a row + SaveAsCsvFileStreamWriter writer = new SaveAsCsvFileStreamWriter(new MemoryStream(output), requestParams); + using (writer) + { + writer.WriteRow(data, columns); + } + + // Then: It should write one line with 2 items, comma delimited + string outputString = Encoding.UTF8.GetString(output).TrimEnd('\0', '\r', '\n'); + string[] lines = outputString.Split(new[] {Environment.NewLine}, StringSplitOptions.None); + Assert.Equal(1, lines.Length); + string[] values = lines[0].Split(','); + Assert.Equal(2, values.Length); + } + + [Fact] + public void WriteRowWithHeader() + { + // Setup: + // ... Create a request params that has no selection made, headers should be printed + // ... Create a set of data to write + // ... Create a memory location to store the data + var requestParams = new SaveResultsAsCsvRequestParams + { + IncludeHeaders = true + }; + List data = new List + { + new DbCellValue { DisplayValue = "item1" }, + new DbCellValue { DisplayValue = "item2" } + }; + List columns = new List + { + new DbColumnWrapper(new TestDbColumn("column1")), + new DbColumnWrapper(new TestDbColumn("column2")) + }; + byte[] output = new byte[8192]; + + // If: I write a row + SaveAsCsvFileStreamWriter writer = new SaveAsCsvFileStreamWriter(new MemoryStream(output), requestParams); + using (writer) + { + writer.WriteRow(data, columns); + } + + // Then: + // ... It should have written two lines + string outputString = Encoding.UTF8.GetString(output).TrimEnd('\0', '\r', '\n'); + string[] lines = outputString.Split(new[] { Environment.NewLine }, StringSplitOptions.None); + Assert.Equal(2, lines.Length); + + // ... It should have written a header line with two, comma separated names + string[] headerValues = lines[0].Split(','); + Assert.Equal(2, headerValues.Length); + for (int i = 0; i < columns.Count; i++) + { + Assert.Equal(columns[i].ColumnName, headerValues[i]); + } + + // Note: No need to check values, it is done as part of the previous test + } + + [Fact] + public void WriteRowWithColumnSelection() + { + // Setup: + // ... Create a request params that selects n-1 columns from the front and back + // ... Create a set of data to write + // ... Create a memory location to store the data + var requestParams = new SaveResultsAsCsvRequestParams + { + ColumnStartIndex = 1, + ColumnEndIndex = 2, + RowStartIndex = 0, // Including b/c it is required to be a "save selection" + RowEndIndex = 10, + IncludeHeaders = true // Including headers to test both column selection logic + }; + List data = new List + { + new DbCellValue { DisplayValue = "item1" }, + new DbCellValue { DisplayValue = "item2" }, + new DbCellValue { DisplayValue = "item3" }, + new DbCellValue { DisplayValue = "item4" } + }; + List columns = new List + { + new DbColumnWrapper(new TestDbColumn("column1")), + new DbColumnWrapper(new TestDbColumn("column2")), + new DbColumnWrapper(new TestDbColumn("column3")), + new DbColumnWrapper(new TestDbColumn("column4")) + }; + byte[] output = new byte[8192]; + + // If: I write a row + SaveAsCsvFileStreamWriter writer = new SaveAsCsvFileStreamWriter(new MemoryStream(output), requestParams); + using (writer) + { + writer.WriteRow(data, columns); + } + + // Then: + // ... It should have written two lines + string outputString = Encoding.UTF8.GetString(output).TrimEnd('\0', '\r', '\n'); + string[] lines = outputString.Split(new[] { Environment.NewLine }, StringSplitOptions.None); + Assert.Equal(2, lines.Length); + + // ... It should have written a header line with two, comma separated names + string[] headerValues = lines[0].Split(','); + Assert.Equal(2, headerValues.Length); + for (int i = 1; i <= 2; i++) + { + Assert.Equal(columns[i].ColumnName, headerValues[i-1]); + } + + // ... The second line should have two, comma separated values + string[] dataValues = lines[1].Split(','); + Assert.Equal(2, dataValues.Length); + for (int i = 1; i <= 2; i++) + { + Assert.Equal(data[i].DisplayValue, dataValues[i-1]); + } + } + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsJsonFileStreamWriterTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsJsonFileStreamWriterTests.cs new file mode 100644 index 00000000..c6fc26a9 --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DataStorage/SaveAsJsonFileStreamWriterTests.cs @@ -0,0 +1,146 @@ +// +// Copyright (c) Microsoft. All rights reserved. +// Licensed under the MIT license. See LICENSE file in the project root for full license information. +// + +using System.Collections.Generic; +using System.IO; +using System.Text; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage; +using Microsoft.SqlTools.ServiceLayer.Test.Utility; +using Newtonsoft.Json; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.DataStorage +{ + public class SaveAsJsonFileStreamWriterTests + { + [Fact] + public void ArrayWrapperTest() + { + // Setup: + // ... Create storage for the output + byte[] output = new byte[8192]; + SaveResultsAsJsonRequestParams saveParams = new SaveResultsAsJsonRequestParams(); + + // If: + // ... I create and then destruct a json writer + var jsonWriter = new SaveAsJsonFileStreamWriter(new MemoryStream(output), saveParams); + jsonWriter.Dispose(); + + // Then: + // ... The output should be an empty array + string outputString = Encoding.UTF8.GetString(output).TrimEnd('\0'); + object[] outputArray = JsonConvert.DeserializeObject(outputString); + Assert.Equal(0, outputArray.Length); + } + + [Fact] + public void WriteRowWithoutColumnSelection() + { + // Setup: + // ... Create a request params that has no selection made + // ... Create a set of data to write + // ... Create storage for the output + SaveResultsAsJsonRequestParams saveParams = new SaveResultsAsJsonRequestParams(); + List data = new List + { + new DbCellValue {DisplayValue = "item1", RawObject = "item1"}, + new DbCellValue {DisplayValue = "null", RawObject = null} + }; + List columns = new List + { + new DbColumnWrapper(new TestDbColumn("column1")), + new DbColumnWrapper(new TestDbColumn("column2")) + }; + byte[] output = new byte[8192]; + + // If: + // ... I write two rows + var jsonWriter = new SaveAsJsonFileStreamWriter(new MemoryStream(output), saveParams); + using (jsonWriter) + { + jsonWriter.WriteRow(data, columns); + jsonWriter.WriteRow(data, columns); + } + + // Then: + // ... Upon deserialization to an array of dictionaries + string outputString = Encoding.UTF8.GetString(output).TrimEnd('\0'); + Dictionary[] outputObject = + JsonConvert.DeserializeObject[]>(outputString); + + // ... There should be 2 items in the array, + // ... The item should have two fields, and two values, assigned appropriately + Assert.Equal(2, outputObject.Length); + foreach (var item in outputObject) + { + Assert.Equal(2, item.Count); + for (int i = 0; i < columns.Count; i++) + { + Assert.True(item.ContainsKey(columns[i].ColumnName)); + Assert.Equal(data[i].RawObject == null ? null : data[i].DisplayValue, item[columns[i].ColumnName]); + } + } + } + + [Fact] + public void WriteRowWithColumnSelection() + { + // Setup: + // ... Create a request params that selects n-1 columns from the front and back + // ... Create a set of data to write + // ... Create a memory location to store the data + var saveParams = new SaveResultsAsJsonRequestParams + { + ColumnStartIndex = 1, + ColumnEndIndex = 3, + RowStartIndex = 0, // Including b/c it is required to be a "save selection" + RowEndIndex = 10 + }; + List data = new List + { + new DbCellValue { DisplayValue = "item1", RawObject = "item1"}, + new DbCellValue { DisplayValue = "item2", RawObject = "item2"}, + new DbCellValue { DisplayValue = "null", RawObject = null}, + new DbCellValue { DisplayValue = "null", RawObject = null} + }; + List columns = new List + { + new DbColumnWrapper(new TestDbColumn("column1")), + new DbColumnWrapper(new TestDbColumn("column2")), + new DbColumnWrapper(new TestDbColumn("column3")), + new DbColumnWrapper(new TestDbColumn("column4")) + }; + byte[] output = new byte[8192]; + + // If: I write two rows + var jsonWriter = new SaveAsJsonFileStreamWriter(new MemoryStream(output), saveParams); + using (jsonWriter) + { + jsonWriter.WriteRow(data, columns); + jsonWriter.WriteRow(data, columns); + } + + // Then: + // ... Upon deserialization to an array of dictionaries + string outputString = Encoding.UTF8.GetString(output).Trim('\0'); + Dictionary[] outputObject = + JsonConvert.DeserializeObject[]>(outputString); + + // ... There should be 2 items in the array + // ... The items should have 2 fields and values + Assert.Equal(2, outputObject.Length); + foreach (var item in outputObject) + { + Assert.Equal(2, item.Count); + for (int i = 1; i <= 2; i++) + { + Assert.True(item.ContainsKey(columns[i].ColumnName)); + Assert.Equal(data[i].RawObject == null ? null : data[i].DisplayValue, item[columns[i].ColumnName]); + } + } + } + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/BatchTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/BatchTests.cs new file mode 100644 index 00000000..1d5da988 --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/BatchTests.cs @@ -0,0 +1,28 @@ +// +// 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 Microsoft.SqlTools.ServiceLayer.QueryExecution; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.SaveResults +{ + public class BatchTests + { + [Theory] + [InlineData(-1)] + [InlineData(100)] + public void SaveAsFailsOutOfRangeResultSet(int resultSetIndex) + { + // If: I attempt to save results for an invalid result set index + // Then: I should get an ArgumentOutOfRange exception + Batch batch = Common.GetBasicExecutedBatch(); + SaveResultsRequestParams saveParams = new SaveResultsRequestParams {ResultSetIndex = resultSetIndex}; + Assert.Throws(() => + batch.SaveAs(saveParams, null, null, null)); + } + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/QueryTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/QueryTests.cs new file mode 100644 index 00000000..856109ec --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/QueryTests.cs @@ -0,0 +1,28 @@ +// +// 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 Microsoft.SqlTools.ServiceLayer.QueryExecution; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.SaveResults +{ + public class QueryTests + { + [Theory] + [InlineData(-1)] + [InlineData(100)] + public void SaveAsFailsOutOfRangeBatch(int batchIndex) + { + // If: I save a basic query's results with out of range batch index + // Then: I should get an out of range exception + Query query = Common.GetBasicExecutedQuery(); + SaveResultsRequestParams saveParams = new SaveResultsRequestParams {BatchIndex = batchIndex}; + Assert.Throws(() => + query.SaveAs(saveParams, null, null, null)); + } + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ResultSetTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ResultSetTests.cs new file mode 100644 index 00000000..a97e69f4 --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ResultSetTests.cs @@ -0,0 +1,190 @@ +// +// 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.Collections.Generic; +using System.Data.Common; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.SqlTools.ServiceLayer.Connection; +using Microsoft.SqlTools.ServiceLayer.QueryExecution; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage; +using Moq; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.SaveResults +{ + public class ResultSetTests + { + [Fact] + public void SaveAsNullParams() + { + // If: I attempt to save with a null set of params + // Then: I should get a null argument exception + ResultSet rs = new ResultSet( + GetReader(null, false, Common.NoOpQuery), Common.Ordinal, Common.Ordinal, + Common.GetFileStreamFactory(new Dictionary())); + Assert.Throws(() => rs.SaveAs( + null, + Common.GetFileStreamFactory(new Dictionary()), + null, null)); + } + + [Fact] + public void SaveAsNullFactory() + { + // If: I attempt to save with a null set of params + // Then: I should get a null argument exception + ResultSet rs = new ResultSet( + GetReader(null, false, Common.NoOpQuery), Common.Ordinal, Common.Ordinal, + Common.GetFileStreamFactory(new Dictionary())); + Assert.Throws(() => rs.SaveAs( + new SaveResultsRequestParams(), + null, null, null)); + } + + [Fact] + public void SaveAsFailedIncomplete() + { + // If: I attempt to save a result set that hasn't completed execution + // Then: I should get an invalid operation exception + ResultSet rs = new ResultSet( + GetReader(null, false, Common.NoOpQuery), Common.Ordinal, Common.Ordinal, + Common.GetFileStreamFactory(new Dictionary())); + Assert.Throws(() => rs.SaveAs( + new SaveResultsRequestParams(), + Common.GetFileStreamFactory(new Dictionary()), + null, null)); + } + + [Fact] + public void SaveAsFailedExistingTaskInProgress() + { + // Setup: + // ... Create a result set that has been executed + ResultSet rs = new ResultSet( + GetReader(new[] { Common.StandardTestData }, false, Common.StandardQuery), + Common.Ordinal, Common.Ordinal, + Common.GetFileStreamFactory(new Dictionary())); + + // ... Insert a non-started task into the save as tasks + rs.SaveTasks.AddOrUpdate(Common.OwnerUri, new Task(() => { }), (s, t) => null); + + // If: I attempt to save results with the same name as the non-completed task + // Then: I should get an invalid operation exception + var requestParams = new SaveResultsRequestParams {FilePath = Common.OwnerUri}; + Assert.Throws(() => rs.SaveAs( + requestParams, GetMockFactory(GetMockWriter().Object, null), + null, null)); + } + + [Fact] + public async Task SaveAsWithoutRowSelection() + { + // Setup: + // ... Create a fake place to store data + Dictionary mockFs = new Dictionary(); + + // ... Create a mock reader/writer for reading the result + IFileStreamFactory resultFactory = Common.GetFileStreamFactory(mockFs); + + // ... Create a result set with dummy data and read to the end + ResultSet rs = new ResultSet( + GetReader(new[] {Common.StandardTestData}, false, Common.StandardQuery), + Common.Ordinal, Common.Ordinal, + resultFactory); + await rs.ReadResultToEnd(CancellationToken.None); + + // ... Create a mock writer for writing the save as file + Mock saveWriter = GetMockWriter(); + IFileStreamFactory saveFactory = GetMockFactory(saveWriter.Object, resultFactory.GetReader); + + // If: I attempt to save results and await completion + rs.SaveAs(new SaveResultsRequestParams {FilePath = Common.OwnerUri}, saveFactory, null, null); + Assert.True(rs.SaveTasks.ContainsKey(Common.OwnerUri)); + await rs.SaveTasks[Common.OwnerUri]; + + // Then: + // ... The task should have completed successfully + Assert.Equal(TaskStatus.RanToCompletion, rs.SaveTasks[Common.OwnerUri].Status); + + // ... All the rows should have been written successfully + saveWriter.Verify( + w => w.WriteRow(It.IsAny>(), It.IsAny>()), + Times.Exactly(Common.StandardRows)); + } + + [Fact] + public async Task SaveAsWithRowSelection() + { + // Setup: + // ... Create a fake place to store data + Dictionary mockFs = new Dictionary(); + + // ... Create a mock reader/writer for reading the result + IFileStreamFactory resultFactory = Common.GetFileStreamFactory(mockFs); + + // ... Create a result set with dummy data and read to the end + ResultSet rs = new ResultSet( + GetReader(new[] { Common.StandardTestData }, false, Common.StandardQuery), + Common.Ordinal, Common.Ordinal, + resultFactory); + await rs.ReadResultToEnd(CancellationToken.None); + + // ... Create a mock writer for writing the save as file + Mock saveWriter = GetMockWriter(); + IFileStreamFactory saveFactory = GetMockFactory(saveWriter.Object, resultFactory.GetReader); + + // If: I attempt to save results that has a selection made + var saveParams = new SaveResultsRequestParams + { + FilePath = Common.OwnerUri, + RowStartIndex = 1, + RowEndIndex = Common.StandardRows - 2, + ColumnStartIndex = 0, // Column start/end doesn't matter, but are required to be + ColumnEndIndex = 10 // considered a "save selection" + }; + rs.SaveAs(saveParams, saveFactory, null, null); + Assert.True(rs.SaveTasks.ContainsKey(Common.OwnerUri)); + await rs.SaveTasks[Common.OwnerUri]; + + // Then: + // ... The task should have completed successfully + Assert.Equal(TaskStatus.RanToCompletion, rs.SaveTasks[Common.OwnerUri].Status); + + // ... All the rows should have been written successfully + saveWriter.Verify( + w => w.WriteRow(It.IsAny>(), It.IsAny>()), + Times.Exactly(Common.StandardRows - 2)); + } + + private static Mock GetMockWriter() + { + var mockWriter = new Mock(); + mockWriter.Setup(w => w.WriteRow(It.IsAny>(), It.IsAny>())); + return mockWriter; + } + + private static IFileStreamFactory GetMockFactory(IFileStreamWriter writer, Func readerGenerator) + { + var mockFactory = new Mock(); + mockFactory.Setup(f => f.GetWriter(It.IsAny())) + .Returns(writer); + mockFactory.Setup(f => f.GetReader(It.IsAny())) + .Returns(readerGenerator); + return mockFactory.Object; + } + + private static DbDataReader GetReader(Dictionary[][] dataSet, bool throwOnRead, string query) + { + var info = Common.CreateTestConnectionInfo(dataSet, throwOnRead); + var connection = info.Factory.CreateSqlConnection(ConnectionService.BuildConnectionString(info.ConnectionDetails)); + var command = connection.CreateCommand(); + command.CommandText = query; + return command.ExecuteReader(); + } + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ServiceIntegrationTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ServiceIntegrationTests.cs new file mode 100644 index 00000000..4421bfe7 --- /dev/null +++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResults/ServiceIntegrationTests.cs @@ -0,0 +1,299 @@ +// +// Copyright (c) Microsoft. All rights reserved. +// Licensed under the MIT license. See LICENSE file in the project root for full license information. +// + +using System.Collections.Generic; +using System.IO; +using System.Threading.Tasks; +using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol; +using Microsoft.SqlTools.ServiceLayer.QueryExecution; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; +using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage; +using Microsoft.SqlTools.ServiceLayer.SqlContext; +using Microsoft.SqlTools.ServiceLayer.Test.Utility; +using Microsoft.SqlTools.ServiceLayer.Workspace; +using Moq; +using Xunit; + +namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.SaveResults +{ + public class ServiceIntegrationTests + { + #region CSV Tests + + [Fact] + public async Task SaveResultsCsvNonExistentQuery() + + { + // Given: A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(null); + QueryExecutionService qes = Common.GetPrimedExecutionService(null, false, false, ws); + + // If: I attempt to save a result set from a query that doesn't exist + SaveResultsAsCsvRequestParams saveParams = new SaveResultsAsCsvRequestParams + { + OwnerUri = Common.OwnerUri // Won't exist because nothing has executed + }; + object error = null; + var requestContext = RequestContextMocks.Create(null) + .AddErrorHandling(o => error = o); + await qes.HandleSaveResultsAsCsvRequest(saveParams, requestContext.Object); + + // Then: + // ... An error event should have been fired + // ... No success event should have been fired + VerifyResponseCalls(requestContext, false, true); + Assert.IsType(error); + Assert.NotNull(error); + Assert.NotNull(((SaveResultRequestError)error).message); + } + + [Fact] + public async void SaveResultAsCsvFailure() + { + // Given: + // ... A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(Common.StandardQuery); + Dictionary storage; + QueryExecutionService qes = Common.GetPrimedExecutionService(new[] {Common.StandardTestData}, true, false, ws, out storage); + + // ... The query execution service has executed a query with results + var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri }; + var executeRequest = RequestContextMocks.Create(null); + await qes.HandleExecuteRequest(executeParams, executeRequest.Object); + await qes.ActiveQueries[Common.OwnerUri].ExecutionTask; + + // If: I attempt to save a result set and get it to throw because of invalid column selection + SaveResultsAsCsvRequestParams saveParams = new SaveResultsAsCsvRequestParams + { + BatchIndex = 0, + FilePath = "qqq", + OwnerUri = Common.OwnerUri, + ResultSetIndex = 0, + ColumnStartIndex = -1, + ColumnEndIndex = 100, + RowStartIndex = 0, + RowEndIndex = 5 + }; + qes.CsvFileFactory = GetCsvStreamFactory(storage, saveParams); + object error = null; + var requestContext = RequestContextMocks.Create(null) + .AddErrorHandling(e => error = e); + await qes.HandleSaveResultsAsCsvRequest(saveParams, requestContext.Object); + await qes.ActiveQueries[saveParams.OwnerUri] + .Batches[saveParams.BatchIndex] + .ResultSets[saveParams.ResultSetIndex] + .SaveTasks[saveParams.FilePath]; + + // Then: + // ... An error event should have been fired + // ... No success event should have been fired + VerifyResponseCalls(requestContext, false, true); + Assert.IsType(error); + Assert.NotNull(error); + Assert.NotNull(((SaveResultRequestError)error).message); + } + + [Fact] + public async void SaveResultsAsCsvSuccess() + { + // Given: + // ... A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(Common.StandardQuery); + Dictionary storage; + QueryExecutionService qes = Common.GetPrimedExecutionService(new[] {Common.StandardTestData}, true, false, ws, out storage); + + // ... The query execution service has executed a query with results + var executeParams = new QueryExecuteParams {QuerySelection = null, OwnerUri = Common.OwnerUri}; + var executeRequest = RequestContextMocks.Create(null); + await qes.HandleExecuteRequest(executeParams, executeRequest.Object); + await qes.ActiveQueries[Common.OwnerUri].ExecutionTask; + + // If: I attempt to save a result set from a query + SaveResultsAsCsvRequestParams saveParams = new SaveResultsAsCsvRequestParams + { + OwnerUri = Common.OwnerUri, + FilePath = "qqq", + BatchIndex = 0, + ResultSetIndex = 0 + }; + qes.CsvFileFactory = GetCsvStreamFactory(storage, saveParams); + SaveResultRequestResult result = null; + var requestContext = RequestContextMocks.Create(r => result = r); + await qes.HandleSaveResultsAsCsvRequest(saveParams, requestContext.Object); + await qes.ActiveQueries[saveParams.OwnerUri] + .Batches[saveParams.BatchIndex] + .ResultSets[saveParams.ResultSetIndex] + .SaveTasks[saveParams.FilePath]; + + // Then: + // ... I should have a successful result + // ... There should not have been an error + VerifyResponseCalls(requestContext, true, false); + Assert.NotNull(result); + Assert.Null(result.Messages); + } + + #endregion + + #region JSON tests + + [Fact] + public async Task SaveResultsJsonNonExistentQuery() + + { + // Given: A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(null); + QueryExecutionService qes = Common.GetPrimedExecutionService(null, false, false, ws); + + // If: I attempt to save a result set from a query that doesn't exist + SaveResultsAsJsonRequestParams saveParams = new SaveResultsAsJsonRequestParams + { + OwnerUri = Common.OwnerUri // Won't exist because nothing has executed + }; + object error = null; + var requestContext = RequestContextMocks.Create(null) + .AddErrorHandling(o => error = o); + await qes.HandleSaveResultsAsJsonRequest(saveParams, requestContext.Object); + + // Then: + // ... An error event should have been fired + // ... No success event should have been fired + VerifyResponseCalls(requestContext, false, true); + Assert.IsType(error); + Assert.NotNull(error); + Assert.NotNull(((SaveResultRequestError)error).message); + } + + [Fact] + public async void SaveResultAsJsonFailure() + { + // Given: + // ... A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(Common.StandardQuery); + Dictionary storage; + QueryExecutionService qes = Common.GetPrimedExecutionService(new[] { Common.StandardTestData }, true, false, ws, out storage); + + // ... The query execution service has executed a query with results + var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri }; + var executeRequest = RequestContextMocks.Create(null); + await qes.HandleExecuteRequest(executeParams, executeRequest.Object); + await qes.ActiveQueries[Common.OwnerUri].ExecutionTask; + + // If: I attempt to save a result set and get it to throw because of invalid column selection + SaveResultsAsJsonRequestParams saveParams = new SaveResultsAsJsonRequestParams + { + BatchIndex = 0, + FilePath = "qqq", + OwnerUri = Common.OwnerUri, + ResultSetIndex = 0, + ColumnStartIndex = -1, + ColumnEndIndex = 100, + RowStartIndex = 0, + RowEndIndex = 5 + }; + qes.JsonFileFactory = GetJsonStreamFactory(storage, saveParams); + object error = null; + var requestContext = RequestContextMocks.Create(null) + .AddErrorHandling(e => error = e); + await qes.HandleSaveResultsAsJsonRequest(saveParams, requestContext.Object); + await qes.ActiveQueries[saveParams.OwnerUri] + .Batches[saveParams.BatchIndex] + .ResultSets[saveParams.ResultSetIndex] + .SaveTasks[saveParams.FilePath]; + + // Then: + // ... An error event should have been fired + // ... No success event should have been fired + VerifyResponseCalls(requestContext, false, true); + Assert.IsType(error); + Assert.NotNull(error); + Assert.NotNull(((SaveResultRequestError)error).message); + } + + [Fact] + public async void SaveResultsAsJsonSuccess() + { + // Given: + // ... A working query and workspace service + WorkspaceService ws = Common.GetPrimedWorkspaceService(Common.StandardQuery); + Dictionary storage; + QueryExecutionService qes = Common.GetPrimedExecutionService(new[] { Common.StandardTestData }, true, false, ws, out storage); + + // ... The query execution service has executed a query with results + var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri }; + var executeRequest = RequestContextMocks.Create(null); + await qes.HandleExecuteRequest(executeParams, executeRequest.Object); + await qes.ActiveQueries[Common.OwnerUri].ExecutionTask; + + // If: I attempt to save a result set from a query + SaveResultsAsJsonRequestParams saveParams = new SaveResultsAsJsonRequestParams + { + OwnerUri = Common.OwnerUri, + FilePath = "qqq", + BatchIndex = 0, + ResultSetIndex = 0 + }; + qes.JsonFileFactory = GetJsonStreamFactory(storage, saveParams); + SaveResultRequestResult result = null; + var requestContext = RequestContextMocks.Create(r => result = r); + await qes.HandleSaveResultsAsJsonRequest(saveParams, requestContext.Object); + await qes.ActiveQueries[saveParams.OwnerUri] + .Batches[saveParams.BatchIndex] + .ResultSets[saveParams.ResultSetIndex] + .SaveTasks[saveParams.FilePath]; + + // Then: + // ... I should have a successful result + // ... There should not have been an error + VerifyResponseCalls(requestContext, true, false); + Assert.NotNull(result); + Assert.Null(result.Messages); + } + + #endregion + + #region Private Helpers + + private static void VerifyResponseCalls(Mock> requestContext, bool successCalled, bool errorCalled) + { + requestContext.Verify(rc => rc.SendResult(It.IsAny()), + successCalled ? Times.Once() : Times.Never()); + requestContext.Verify(rc => rc.SendError(It.IsAny()), + errorCalled ? Times.Once() : Times.Never()); + } + + private static IFileStreamFactory GetCsvStreamFactory(IDictionary storage, SaveResultsAsCsvRequestParams saveParams) + { + Mock mock = new Mock(); + mock.Setup(fsf => fsf.GetReader(It.IsAny())) + .Returns(output => new ServiceBufferFileStreamReader(new MemoryStream(storage[output]))); + mock.Setup(fsf => fsf.GetWriter(It.IsAny())) + .Returns(output => + { + storage.Add(output, new byte[8192]); + return new SaveAsCsvFileStreamWriter(new MemoryStream(storage[output]), saveParams); + }); + + return mock.Object; + } + + private static IFileStreamFactory GetJsonStreamFactory(IDictionary storage, SaveResultsAsJsonRequestParams saveParams) + { + Mock mock = new Mock(); + mock.Setup(fsf => fsf.GetReader(It.IsAny())) + .Returns(output => new ServiceBufferFileStreamReader(new MemoryStream(storage[output]))); + mock.Setup(fsf => fsf.GetWriter(It.IsAny())) + .Returns(output => + { + storage.Add(output, new byte[8192]); + return new SaveAsJsonFileStreamWriter(new MemoryStream(storage[output]), saveParams); + }); + + return mock.Object; + } + + #endregion + } +} diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs deleted file mode 100644 index c97e84b8..00000000 --- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs +++ /dev/null @@ -1,338 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. -// Licensed under the MIT license. See LICENSE file in the project root for full license information. -// - -using System.IO; -using System.Threading.Tasks; -using System.Runtime.InteropServices; -using Microsoft.SqlTools.ServiceLayer.QueryExecution; -using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; -using Microsoft.SqlTools.ServiceLayer.Test.Utility; -using Xunit; - -namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution -{ - /// - /// Tests for saving a result set to a file - /// - public class SaveResultsTests - { - /// - /// Test save results to a file as CSV with correct parameters - /// - [Fact] - public async void SaveResultsAsCsvSuccessTest() - { - // Execute a query - var workplaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new [] {Common.StandardTestData}, true, false, workplaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as csv with correct parameters - var saveParams = new SaveResultsAsCsvRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_1.csv", - IncludeHeaders = true - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.Null(r.Messages); - }).Complete(); - - // 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]; - await selectedResultSet.GetSaveTask(saveParams.FilePath); - - // Expect to see a file successfully created in filepath and a success message - saveRequest.Validate(); - Assert.True(File.Exists(saveParams.FilePath)); - - // Delete temp file after test - if (File.Exists(saveParams.FilePath)) - { - File.Delete(saveParams.FilePath); - } - } - - /// - /// Test save results to a file as CSV with a selection of cells and correct parameters - /// - [Fact] - public async void SaveResultsAsCsvWithSelectionSuccessTest() - { - // Execute a query - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new []{Common.StandardTestData}, true, false, workspaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument , OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as csv with correct parameters - var saveParams = new SaveResultsAsCsvRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_2.csv", - IncludeHeaders = true, - RowStartIndex = 0, - RowEndIndex = 0, - ColumnStartIndex = 0, - ColumnEndIndex = 0 - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.Null(r.Messages); - }).Complete(); - - // 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 - saveRequest.Validate(); - Assert.True(File.Exists(saveParams.FilePath)); - - // Delete temp file after test - if (File.Exists(saveParams.FilePath)) - { - File.Delete(saveParams.FilePath); - } - } - - /// - /// Test handling exception in saving results to CSV file - /// - [Fact] - public async void SaveResultsAsCsvExceptionTest() - { - // Execute a query - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new[] {Common.StandardTestData}, true, false, workspaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as csv with incorrect filepath - var saveParams = new SaveResultsAsCsvRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "G:\\test.csv" : "/test.csv" - }; - var saveRequest = new EventFlowValidator() - .AddErrorValidation(e => - { - Assert.False(string.IsNullOrWhiteSpace(e.message)); - }).Complete(); - - // 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]; - await selectedResultSet.GetSaveTask(saveParams.FilePath); - - // Expect to see error message - saveRequest.Validate(); - Assert.False(File.Exists(saveParams.FilePath)); - } - - /// - /// Test saving results to CSV file when the requested result set is no longer active - /// - [Fact] - public async Task SaveResultsAsCsvQueryNotFoundTest() - { - // Create a query execution service - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService); - - // Request to save the results as csv with query that is no longer active - var saveParams = new SaveResultsAsCsvRequestParams - { - OwnerUri = "falseuri", - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_3.csv" - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.NotNull(r.Messages); - }).Complete(); - await queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object); - - // Expect message that save failed - saveRequest.Validate(); - Assert.False(File.Exists(saveParams.FilePath)); - } - - /// - /// Test save results to a file as JSON with correct parameters - /// - [Fact] - public async void SaveResultsAsJsonSuccessTest() - { - // Execute a query - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new[] {Common.StandardTestData}, true, false, workspaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as json with correct parameters - var saveParams = new SaveResultsAsJsonRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_4.json" - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.Null(r.Messages); - }).Complete(); - - // 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]; - await selectedResultSet.GetSaveTask(saveParams.FilePath); - - // Expect to see a file successfully created in filepath and a success message - saveRequest.Validate(); - Assert.True(File.Exists(saveParams.FilePath)); - - // Delete temp file after test - if (File.Exists(saveParams.FilePath)) - { - File.Delete(saveParams.FilePath); - } - } - - /// - /// Test save results to a file as JSON with a selection of cells and correct parameters - /// - [Fact] - public async void SaveResultsAsJsonWithSelectionSuccessTest() - { - // Execute a query - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new[] { Common.StandardTestData }, true, false, workspaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument , OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as json with correct parameters - var saveParams = new SaveResultsAsJsonRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_5.json", - RowStartIndex = 0, - RowEndIndex = 1, - ColumnStartIndex = 0, - ColumnEndIndex = 1 - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.Null(r.Messages); - }).Complete(); - - // 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]; - await selectedResultSet.GetSaveTask(saveParams.FilePath); - - // Expect to see a file successfully created in filepath and a success message - saveRequest.Validate(); - Assert.True(File.Exists(saveParams.FilePath)); - - // Delete temp file after test - if (File.Exists(saveParams.FilePath)) - { - File.Delete(saveParams.FilePath); - } - } - - /// - /// Test handling exception in saving results to JSON file - /// - [Fact] - public async void SaveResultsAsJsonExceptionTest() - { - // Execute a query - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(new [] {Common.StandardTestData}, true, false, workspaceService); - var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri }; - var executeRequest = RequestContextMocks.Create(null); - await Common.AwaitExecution(queryService, executeParams, executeRequest.Object); - - // Request to save the results as json with incorrect filepath - var saveParams = new SaveResultsAsJsonRequestParams - { - OwnerUri = Common.OwnerUri, - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "G:\\test.json" : "/test.json" - }; - var saveRequest = new EventFlowValidator() - .AddErrorValidation(e => - { - Assert.False(string.IsNullOrWhiteSpace(e.message)); - }).Complete(); - queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch(); - - // 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]; - await selectedResultSet.GetSaveTask(saveParams.FilePath); - - // Expect to see error message - saveRequest.Validate(); - Assert.False(File.Exists(saveParams.FilePath)); - } - - /// - /// Test saving results to JSON file when the requested result set is no longer active - /// - [Fact] - public async Task SaveResultsAsJsonQueryNotFoundTest() - { - // Create a query service - var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery); - var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService); - - // Request to save the results as json with query that is no longer active - var saveParams = new SaveResultsAsJsonRequestParams - { - OwnerUri = "falseuri", - ResultSetIndex = 0, - BatchIndex = 0, - FilePath = "testwrite_6.json" - }; - var saveRequest = new EventFlowValidator() - .AddResultValidation(r => - { - Assert.Equal("Failed to save results, ID not found.", r.Messages); - }).Complete(); - await queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object); - - // Expect message that save failed - saveRequest.Validate(); - Assert.False(File.Exists(saveParams.FilePath)); - } - } -}