Save as XML feature added (#684)

Similar approach used like Save as JSON or Save as CSV.

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

@@ -118,6 +118,22 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         //TODO: define config for save as JSON
     }
 
+    /// <summary>
+    /// Parameters to save results as XML
+    /// </summary>
+    public class SaveResultsAsXmlRequestParams: SaveResultsRequestParams
+    {
+        /// <summary>
+        /// Formatting of the XML file
+        /// </summary>
+        public bool Formatted { get; set; }
+        
+        /// <summary>
+        /// Encoding of the XML file
+        /// </summary>
+        public string Encoding { get; set; }
+    }
+    
     /// <summary>
     /// Parameters for the save results result
     /// </summary>
@@ -158,5 +174,15 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
             RequestType<SaveResultsAsJsonRequestParams, SaveResultRequestResult> Type =
             RequestType<SaveResultsAsJsonRequestParams, SaveResultRequestResult>.Create("query/saveJson");
     }
+    
+    /// <summary>
+    /// Request type to save results as XML
+    /// </summary>
+    public class SaveResultsAsXmlRequest
+    {
+        public static readonly
+            RequestType<SaveResultsAsXmlRequestParams, SaveResultRequestResult> Type =
+                RequestType<SaveResultsAsXmlRequestParams, SaveResultRequestResult>.Create("query/saveXml");
+    }
 
 }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsXmlFileStreamFactory.cs

@@ -0,0 +1,71 @@
+// 
+// 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;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage
+{
+    public class SaveAsXmlFileStreamFactory : IFileStreamFactory
+    {
+
+        #region Properties
+
+        /// <summary>
+        /// Settings for query execution
+        /// </summary>
+        public QueryExecutionSettings QueryExecutionSettings { get; set; }
+
+        /// <summary>
+        /// Parameters for the save as XML request
+        /// </summary>
+        public SaveResultsAsXmlRequestParams SaveRequestParams { get; set; }
+
+        #endregion
+
+        /// <summary>
+        /// File names are not meant to be created with this factory.
+        /// </summary>
+        /// <exception cref="NotImplementedException">Thrown all times</exception>
+        [Obsolete]
+        public string CreateFile()
+        {
+            throw new InvalidOperationException();
+        }
+
+        /// <summary>
+        /// Returns a new service buffer reader for reading results back in from the temporary buffer files
+        /// </summary>
+        /// <param name="fileName">Path to the temp buffer file</param>
+        /// <returns>Stream reader</returns>
+        public IFileStreamReader GetReader(string fileName)
+        {
+            return new ServiceBufferFileStreamReader(new FileStream(fileName, FileMode.Open, FileAccess.Read), QueryExecutionSettings);
+        }
+
+        /// <summary>
+        /// Returns a new XML writer for writing results to a XML file
+        /// </summary>
+        /// <param name="fileName">Path to the XML output file</param>
+        /// <returns>Stream writer</returns>
+        public IFileStreamWriter GetWriter(string fileName)
+        {
+            return new SaveAsXmlFileStreamWriter(new FileStream(fileName, FileMode.Create, FileAccess.ReadWrite), SaveRequestParams);
+        }
+
+        /// <summary>
+        /// Safely deletes the file
+        /// </summary>
+        /// <param name="fileName">Path to the file to delete</param>
+        public void DisposeFile(string fileName)
+        {
+            FileUtilities.SafeFileDelete(fileName);
+        }
+
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/SaveAsXmlFileStreamWriter.cs

@@ -0,0 +1,142 @@
+// 
+// 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.Xml;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage
+{
+    /// <summary>
+    /// Writer for writing rows of results to a XML file.
+    /// </summary>
+    /// <remarks>
+    /// This implements its own IDisposable because the cleanup logic closes the element 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.
+    /// </remarks>
+    public class SaveAsXmlFileStreamWriter : SaveAsStreamWriter, IDisposable
+    {
+        // Root element name for the output XML
+        private const string RootElementTag = "data";
+        
+        // Item element name which will be used for every row
+        private const string ItemElementTag = "row";
+        
+        #region Member Variables
+
+        private readonly XmlTextWriter xmlTextWriter;
+
+        #endregion
+
+        /// <summary>
+        /// Constructor, writes the header to the file, chains into the base constructor
+        /// </summary>
+        /// <param name="stream">FileStream to access the JSON file output</param>
+        /// <param name="requestParams">XML save as request parameters</param>
+        public SaveAsXmlFileStreamWriter(Stream stream, SaveResultsAsXmlRequestParams requestParams)
+            : base(stream, requestParams)
+        {
+            // Setup the internal state
+            var encoding = GetEncoding(requestParams);
+            xmlTextWriter = new XmlTextWriter(stream, encoding);
+            xmlTextWriter.Formatting = requestParams.Formatted ? Formatting.Indented : Formatting.None;
+
+            //Start the document and the root element
+            xmlTextWriter.WriteStartDocument();
+            xmlTextWriter.WriteStartElement(RootElementTag);
+        }
+
+        /// <summary>
+        /// Writes a row of data as a XML object
+        /// </summary>
+        /// <param name="row">The data of the row to output to the file</param>
+        /// <param name="columns">
+        /// The entire list of columns for the result set. They will be filtered down as per the
+        /// request params.
+        /// </param>
+        public override void WriteRow(IList<DbCellValue> row, IList<DbColumnWrapper> columns)
+        {
+            // Write the header for the object
+            xmlTextWriter.WriteStartElement(ItemElementTag);
+
+            // Write the items out as properties
+            int columnStart = ColumnStartIndex ?? 0;
+            int columnEnd = ColumnEndIndex + 1 ?? columns.Count;
+            for (int i = columnStart; i < columnEnd; i++)
+            {
+                // Write the column name as item tag
+                xmlTextWriter.WriteStartElement(columns[i].ColumnName);
+                
+                if (row[i].RawObject != null)
+                {
+                    xmlTextWriter.WriteString(row[i].DisplayValue);
+                }
+
+                // End the item tag
+                xmlTextWriter.WriteEndElement();
+            }
+
+            // Write the footer for the object
+            xmlTextWriter.WriteEndElement();
+        }
+
+        /// <summary>
+        /// Get the encoding for the XML file according to <param name="requestParams"></param>
+        /// </summary>
+        /// <param name="requestParams">XML save as request parameters</param>
+        /// <returns></returns>
+        private Encoding GetEncoding(SaveResultsAsXmlRequestParams requestParams)
+        {
+            Encoding.RegisterProvider(CodePagesEncodingProvider.Instance);
+            Encoding encoding;
+            try
+            {
+                if (int.TryParse(requestParams.Encoding, out var codepage))
+                {
+                    encoding = Encoding.GetEncoding(codepage);
+                }
+                else
+                {
+                    encoding = Encoding.GetEncoding(requestParams.Encoding);
+                }
+            }
+            catch
+            {
+                // Fallback encoding when specified codepage is invalid
+                encoding = Encoding.GetEncoding("utf-8");
+            }
+
+            return encoding;
+        }
+
+        private bool disposed = false;
+
+        /// <summary>
+        /// Disposes the writer by closing up the element that contains the row objects
+        /// </summary>
+        protected override void Dispose(bool disposing)
+        {
+            if (disposed)
+                return;
+
+            if (disposing)
+            {
+                // Write the footer of the file
+                xmlTextWriter.WriteEndElement();
+                xmlTextWriter.WriteEndDocument();
+
+                xmlTextWriter.Close();
+                xmlTextWriter.Dispose();
+            }
+
+            disposed = true;
+            base.Dispose(disposing);
+        }
+    }
+}

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

@@ -2,9 +2,9 @@
 // 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.Collections.Concurrent;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Hosting.Protocol;
@@ -98,6 +98,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         internal IFileStreamFactory JsonFileFactory { get; set; }
 
+        /// <summary>
+        /// File factory to be used to create XML files from result sets. Set to internal in order
+        /// to allow overriding in unit testing
+        /// </summary>
+        internal IFileStreamFactory XmlFileFactory { get; set; }
+
         /// <summary>
         /// The collection of active queries
         /// </summary>
@@ -151,6 +157,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(SaveResultsAsCsvRequest.Type, HandleSaveResultsAsCsvRequest);
             serviceHost.SetRequestHandler(SaveResultsAsExcelRequest.Type, HandleSaveResultsAsExcelRequest);
             serviceHost.SetRequestHandler(SaveResultsAsJsonRequest.Type, HandleSaveResultsAsJsonRequest);
+            serviceHost.SetRequestHandler(SaveResultsAsXmlRequest.Type, HandleSaveResultsAsXmlRequest);
             serviceHost.SetRequestHandler(QueryExecutionPlanRequest.Type, HandleExecutionPlanRequest);
             serviceHost.SetRequestHandler(SimpleExecuteRequest.Type, HandleSimpleExecuteRequest);
 
@@ -448,6 +455,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             await SaveResultsHelper(saveParams, requestContext, jsonFactory);
         }
 
+        /// <summary>
+        /// Process request to save a resultSet to a file in XML format
+        /// </summary>
+        internal async Task HandleSaveResultsAsXmlRequest(SaveResultsAsXmlRequestParams saveParams,
+            RequestContext<SaveResultRequestResult> requestContext)
+        {
+            // Use the default XML file factory if we haven't overridden it
+            IFileStreamFactory xmlFactory = XmlFileFactory ?? new SaveAsXmlFileStreamFactory
+            {
+                SaveRequestParams = saveParams,
+                QueryExecutionSettings = Settings.QueryExecutionSettings
+            };
+            await SaveResultsHelper(saveParams, requestContext, xmlFactory);
+        }
+
         #endregion
 
         #region Inter-Service API Handlers

test/Microsoft.SqlTools.ServiceLayer.Test.Common/TestServiceDriverProvider.cs

@@ -523,6 +523,23 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
             return result;
         }
 
+        /// <summary>
+        /// Request to save query results as XML
+        /// </summary>
+        public async Task<SaveResultRequestResult> SaveAsXml(string ownerUri, string filename, int batchIndex, int resultSetIndex)
+        {
+            var saveParams = new SaveResultsAsXmlRequestParams
+            {
+                OwnerUri = ownerUri,
+                BatchIndex = batchIndex,
+                ResultSetIndex = resultSetIndex,
+                FilePath = filename
+            };
+
+            var result = await Driver.SendRequest(SaveResultsAsXmlRequest.Type, saveParams);
+            return result;
+        }
+        
         /// <summary>
         /// Request a subset of results from a query
         /// </summary>

test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/SaveResults/ServiceIntegrationTests.cs

@@ -244,6 +244,117 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
 
         #endregion
         
+        #region XML tests
+
+        [Fact]
+        public async Task SaveResultsXmlNonExistentQuery()
+        {
+            // Given: A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(null);
+            QueryExecutionService qes = Common.GetPrimedExecutionService(null, false, false, false, ws);
+
+            // If: I attempt to save a result set from a query that doesn't exist
+            SaveResultsAsXmlRequestParams saveParams = new SaveResultsAsXmlRequestParams
+            {
+                OwnerUri = Constants.OwnerUri // Won't exist because nothing has executed
+            };
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardErrorValidation()
+                .Complete();
+            await qes.HandleSaveResultsAsXmlRequest(saveParams, efv.Object);
+
+            // Then:
+            // ... An error event should have been fired
+            // ... No success event should have been fired
+            efv.Validate();
+        }
+
+        [Fact]
+        public async Task SaveResultAsXmlFailure()
+        {
+            // Given: 
+            // ... A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(Constants.StandardQuery);
+            ConcurrentDictionary<string, byte[]> storage;
+            QueryExecutionService qes = Common.GetPrimedExecutionService(Common.StandardTestDataSet, true, false, false, ws, out storage);
+
+            // ... The query execution service has executed a query with results
+            var executeParams = new ExecuteDocumentSelectionParams {QuerySelection = null, OwnerUri = Constants.OwnerUri};
+            var executeRequest = RequestContextMocks.Create<ExecuteRequestResult>(null);
+            await qes.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await qes.ActiveQueries[Constants.OwnerUri].ExecutionTask;
+
+            // If: I attempt to save a result set and get it to throw because of invalid column selection
+            SaveResultsAsXmlRequestParams saveParams = new SaveResultsAsXmlRequestParams
+            {
+                BatchIndex = 0,
+                FilePath = "qqq",
+                OwnerUri = Constants.OwnerUri,
+                ResultSetIndex = 0,
+                ColumnStartIndex = -1,
+                ColumnEndIndex = 100,
+                RowStartIndex = 0,
+                RowEndIndex = 5
+            };
+            qes.XmlFileFactory = GetXmlStreamFactory(storage, saveParams);
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardErrorValidation()
+                .Complete();
+            await qes.HandleSaveResultsAsXmlRequest(saveParams, efv.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
+            efv.Validate();
+        }
+
+        [Fact]
+        public async Task SaveResultsAsXmlSuccess()
+        {
+            // Given: 
+            // ... A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(Constants.StandardQuery);
+            ConcurrentDictionary<string, byte[]> storage;
+            QueryExecutionService qes = Common.GetPrimedExecutionService(Common.StandardTestDataSet, true, false, false, ws, out storage);
+
+            // ... The query execution service has executed a query with results
+            var executeParams = new ExecuteDocumentSelectionParams {QuerySelection = null, OwnerUri = Constants.OwnerUri};
+            var executeRequest = RequestContextMocks.Create<ExecuteRequestResult>(null);
+            await qes.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await qes.ActiveQueries[Constants.OwnerUri].ExecutionTask;
+
+            // If: I attempt to save a result set from a query
+            SaveResultsAsXmlRequestParams saveParams = new SaveResultsAsXmlRequestParams
+            {
+                OwnerUri = Constants.OwnerUri,
+                FilePath = "qqq",
+                BatchIndex = 0,
+                ResultSetIndex = 0,
+                Formatted = true
+            };
+            qes.XmlFileFactory = GetXmlStreamFactory(storage, saveParams);
+            
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardResultValidator()
+                .Complete();
+            await qes.HandleSaveResultsAsXmlRequest(saveParams, efv.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
+            efv.Validate();
+        }
+
+        #endregion
+        
         #region Excel Tests 
         
         [Fact]
@@ -385,6 +496,22 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
             return mock.Object;
         }
 
+        private static IFileStreamFactory GetXmlStreamFactory(IDictionary<string, byte[]> storage,
+            SaveResultsAsXmlRequestParams saveParams)
+        {
+            Mock<IFileStreamFactory> mock = new Mock<IFileStreamFactory>();
+            mock.Setup(fsf => fsf.GetReader(It.IsAny<string>()))
+                .Returns<string>(output => new ServiceBufferFileStreamReader(new MemoryStream(storage[output]), new QueryExecutionSettings()));
+            mock.Setup(fsf => fsf.GetWriter(It.IsAny<string>()))
+                .Returns<string>(output =>
+                {
+                    storage.Add(output, new byte[8192]);
+                    return new SaveAsXmlFileStreamWriter(new MemoryStream(storage[output]), saveParams);
+                });
+
+            return mock.Object;
+        }
+        
         private static IFileStreamFactory GetExcelStreamFactory(IDictionary<string, byte[]> storage,
             SaveResultsAsExcelRequestParams saveParams)
         {