diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs
index 721d13c9..1cf2390e 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/SaveResultsRequest.cs
@@ -32,6 +32,28 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
/// URI for the editor that called save results
///
public string OwnerUri { get; set; }
+
+ ///
+ /// Start index of the selected rows (inclusive)
+ ///
+ public int? RowStartIndex { get; set; }
+
+ ///
+ /// End index of the selected rows (inclusive)
+ ///
+ public int? RowEndIndex { get; set; }
+
+ ///
+ /// Start index of the selected columns (inclusive)
+ ///
+ ///
+ public int? ColumnStartIndex { get; set; }
+
+ ///
+ /// End index of the selected columns (inclusive)
+ ///
+ ///
+ public int? ColumnEndIndex { get; set; }
}
///
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
index 20bc4347..d8222972 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
@@ -287,20 +287,41 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
// get the requested resultSet from query
Batch selectedBatch = result.Batches[saveParams.BatchIndex];
- ResultSet selectedResultSet = selectedBatch.ResultSets.ToList()[saveParams.ResultSetIndex];
- if (saveParams.IncludeHeaders)
- {
- // write column names to csv
- await csvFile.WriteLineAsync(string.Join(",",
- selectedResultSet.Columns.Select(column => SaveResults.EncodeCsvField(column.ColumnName) ?? string.Empty)));
+ ResultSet selectedResultSet = (selectedBatch.ResultSets.ToList())[saveParams.ResultSetIndex];
+ int columnCount = 0;
+ int rowCount = 0;
+ int columnStartIndex = 0;
+ int rowStartIndex = 0;
+
+ // set column, row counts depending on whether save request is for entire result set or a subset
+ if (SaveResults.isSaveSelection(saveParams))
+ {
+ columnCount = saveParams.ColumnEndIndex.Value - saveParams.ColumnStartIndex.Value + 1;
+ rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
+ columnStartIndex = saveParams.ColumnStartIndex.Value;
+ rowStartIndex =saveParams.RowStartIndex.Value;
+ }
+ else
+ {
+ columnCount = selectedResultSet.Columns.Length;
+ rowCount = (int)selectedResultSet.RowCount;
}
- // write rows to csv
- foreach (var row in selectedResultSet.Rows)
+ // write column names if include headers option is chosen
+ if (saveParams.IncludeHeaders)
{
- await csvFile.WriteLineAsync(string.Join(",",
- row.Select(field => SaveResults.EncodeCsvField(field ?? string.Empty))));
+ await csvFile.WriteLineAsync( string.Join( ",", selectedResultSet.Columns.Skip(columnStartIndex).Take(columnCount).Select( column =>
+ SaveResults.EncodeCsvField(column.ColumnName) ?? string.Empty)));
}
+
+ // retrieve rows and write as csv
+ ResultSetSubset resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex, rowCount);
+ foreach (var row in resultSubset.Rows)
+ {
+ await csvFile.WriteLineAsync( string.Join( ",", row.Skip(columnStartIndex).Take(columnCount).Select( field =>
+ SaveResults.EncodeCsvField((field != null) ? field.ToString(): "NULL"))));
+ }
+
}
// Successfully wrote file, send success result
@@ -340,20 +361,40 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
jsonWriter.Formatting = Formatting.Indented;
jsonWriter.WriteStartArray();
-
+
// get the requested resultSet from query
Batch selectedBatch = result.Batches[saveParams.BatchIndex];
ResultSet selectedResultSet = selectedBatch.ResultSets.ToList()[saveParams.ResultSetIndex];
+ int rowCount = 0;
+ int rowStartIndex = 0;
+ int columnStartIndex = 0;
+ int columnEndIndex = 0;
- // write each row to JSON
- foreach (var row in selectedResultSet.Rows)
+ // set column, row counts depending on whether save request is for entire result set or a subset
+ if (SaveResults.isSaveSelection(saveParams))
+ {
+
+ rowCount = saveParams.RowEndIndex.Value - saveParams.RowStartIndex.Value + 1;
+ rowStartIndex = saveParams.RowStartIndex.Value;
+ columnStartIndex = saveParams.ColumnStartIndex.Value;
+ columnEndIndex = saveParams.ColumnEndIndex.Value + 1 ; // include the last column
+ }
+ else
+ {
+ rowCount = (int)selectedResultSet.RowCount;
+ columnEndIndex = selectedResultSet.Columns.Length;
+ }
+
+ // retrieve rows and write as json
+ ResultSetSubset resultSubset = await result.GetSubset(saveParams.BatchIndex, saveParams.ResultSetIndex, rowStartIndex, rowCount);
+ foreach (var row in resultSubset.Rows)
{
jsonWriter.WriteStartObject();
- for (int i = 0; i < row.Length; i++)
+ for (int i = columnStartIndex ; i < columnEndIndex; i++)
{
+ //get column name
DbColumnWrapper col = selectedResultSet.Columns[i];
- string val = row[i];
-
+ string val = row[i]?.ToString();
jsonWriter.WritePropertyName(col.ColumnName);
if (val == null)
{
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs
index f63f253f..9188d5b0 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SaveResults.cs
@@ -4,6 +4,7 @@
//
using System;
using System.Text;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
@@ -79,6 +80,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
return ret;
}
+
+ internal static bool isSaveSelection(SaveResultsRequestParams saveParams)
+ {
+ return (saveParams.ColumnStartIndex != null && saveParams.ColumnEndIndex != null
+ && saveParams.RowEndIndex != null && saveParams.RowEndIndex != null);
+ }
}
}
\ No newline at end of file
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
index 96e5742b..153e2977 100644
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
@@ -68,6 +68,57 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
}
}
+ ///
+ /// Test save results to a file as CSV with a selection of cells and correct parameters
+ ///
+ [Fact]
+ public void SaveResultsAsCsvWithSelectionSuccessTest()
+ {
+
+ // Set up file for returning the query
+ var fileMock = new Mock();
+ fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+ // Set up workspace mock
+ var workspaceService = new Mock>();
+ workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny()))
+ .Returns(fileMock.Object);
+
+ // Execute a query
+ var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+ var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument , OwnerUri = Common.OwnerUri };
+ var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
+ queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
+
+ // 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
+ };
+ SaveResultRequestResult result = null;
+ var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
+ queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
+ queryService.HandleSaveResultsAsCsvRequest(saveParams, saveRequest.Object).Wait();
+
+ // Expect to see a file successfully created in filepath and a success message
+ Assert.Null(result.Messages);
+ Assert.True(File.Exists(saveParams.FilePath));
+ VerifySaveResultsCallCount(saveRequest, Times.Once(), Times.Never());
+
+ // Delete temp file after test
+ if (File.Exists(saveParams.FilePath))
+ {
+ File.Delete(saveParams.FilePath);
+ }
+ }
+
///
/// Test handling exception in saving results to CSV file
///
@@ -167,13 +218,62 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
OwnerUri = Common.OwnerUri,
ResultSetIndex = 0,
BatchIndex = 0,
- FilePath = "testwrite_4.json"
+ FilePath = "testwrite_4.json"
};
SaveResultRequestResult result = null;
var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object).Wait();
+
+ // Expect to see a file successfully created in filepath and a success message
+ Assert.Null(result.Messages);
+ Assert.True(File.Exists(saveParams.FilePath));
+ VerifySaveResultsCallCount(saveRequest, Times.Once(), Times.Never());
+ // 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 void SaveResultsAsJsonWithSelectionSuccessTest()
+ {
+ // Set up file for returning the query
+ var fileMock = new Mock();
+ fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+ // Set up workspace mock
+ var workspaceService = new Mock>();
+ workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny()))
+ .Returns(fileMock.Object);
+
+ // Execute a query
+ var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+ var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument , OwnerUri = Common.OwnerUri };
+ var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
+ queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
+
+ // 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 = 0,
+ ColumnStartIndex = 0,
+ ColumnEndIndex = 0
+ };
+ SaveResultRequestResult result = null;
+ var saveRequest = GetSaveResultsContextMock(qcr => result = qcr, null);
+ queryService.ActiveQueries[Common.OwnerUri].Batches[0] = Common.GetBasicExecutedBatch();
+ queryService.HandleSaveResultsAsJsonRequest(saveParams, saveRequest.Object).Wait();
+
// Expect to see a file successfully created in filepath and a success message
Assert.Null(result.Messages);
Assert.True(File.Exists(saveParams.FilePath));