From 934df0556ae7068c40e9104249950867fc2b9a26 Mon Sep 17 00:00:00 2001
From: Alan Ren <alanren@microsoft.com>
Date: Fri, 9 Jun 2023 09:15:08 -0700
Subject: [PATCH] support copy result in STS (#2096)

* use new sql parser

* copy results to clipboard

* fix parameter name
---
 Packages.props                                |   1 +
 .../Microsoft.SqlTools.ServiceLayer.csproj    |   1 +
 .../Contracts/CopyResultsRequest.cs           |  56 ++++++++
 .../QueryExecution/Query.cs                   |  34 +++--
 .../QueryExecution/QueryExecutionService.cs   | 125 +++++++++++++++++-
 5 files changed, 204 insertions(+), 13 deletions(-)
 create mode 100644 src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/CopyResultsRequest.cs

diff --git a/Packages.props b/Packages.props
index b74c8ac7..f2e85ab4 100644
--- a/Packages.props
+++ b/Packages.props
@@ -47,6 +47,7 @@
 		<PackageReference Update="Azure.Storage.Blobs" Version="12.10.0" />
 		<PackageReference Update="coverlet.collector" Version="3.1.2" />
 		<PackageReference Update="coverlet.msbuild" Version="3.1.2" />
+		<PackageReference Update="TextCopy" Version="6.2.1" />
 	</ItemGroup>
 
 	<!-- When updating version of Dependencies in the below section, please also update the version in the following files:
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Microsoft.SqlTools.ServiceLayer.csproj b/src/Microsoft.SqlTools.ServiceLayer/Microsoft.SqlTools.ServiceLayer.csproj
index b0ced717..971600b3 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/Microsoft.SqlTools.ServiceLayer.csproj
+++ b/src/Microsoft.SqlTools.ServiceLayer/Microsoft.SqlTools.ServiceLayer.csproj
@@ -53,6 +53,7 @@
 		<PackageReference Include="Microsoft.SqlServer.TransactSql.ScriptDom.NRT">
 			<Aliases>ASAScriptDom</Aliases>
 		</PackageReference>
+		<PackageReference Include="TextCopy" />
 	</ItemGroup>
 	<ItemGroup>
 		<ProjectReference Include="../Microsoft.SqlTools.Hosting/Microsoft.SqlTools.Hosting.csproj" />
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/CopyResultsRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/CopyResultsRequest.cs
new file mode 100644
index 00000000..f89a5825
--- /dev/null
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/CopyResultsRequest.cs
@@ -0,0 +1,56 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+#nullable disable
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    public class TableSelectionRange
+    {
+        public int FromRow { get; set; }
+        public int ToRow { get; set; }
+        public int FromColumn { get; set; }
+        public int ToColumn { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters for the copy results request
+    /// </summary>
+    public class CopyResultsRequestParams : SubsetParams
+    {
+        /// <summary>
+        /// Whether to remove the line break from cell values.
+        /// </summary>
+        public bool RemoveNewLines { get; set; }
+
+        /// <summary>
+        /// Whether to include the column headers.
+        /// </summary>
+        public bool IncludeHeaders { get; set; }
+
+        /// <summary>
+        /// The selections.
+        /// </summary>
+        public TableSelectionRange[] Selections { get; set; }
+    }
+
+    /// <summary>
+    /// Result for the copy results request
+    /// </summary>
+    public class CopyResultsRequestResult
+    {
+    }
+
+    /// <summary>
+    /// Copy Results Request
+    /// </summary>
+    public class CopyResultsRequest
+    {
+        public static readonly RequestType<CopyResultsRequestParams, CopyResultsRequestResult> Type =
+            RequestType<CopyResultsRequestParams, CopyResultsRequestResult>.Create("query/copy");
+    }
+}
\ No newline at end of file
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
index 62cd0adc..a15e5177 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
@@ -143,10 +143,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 .Select((batchDefinition, index) =>
                     new Batch(batchDefinition.BatchText,
                         new SelectionData(
-                            batchDefinition.StartLine-1,
-                            batchDefinition.StartColumn-1,
-                            batchDefinition.EndLine-1,
-                            batchDefinition.EndColumn-1),
+                            batchDefinition.StartLine - 1,
+                            batchDefinition.StartColumn - 1,
+                            batchDefinition.EndLine - 1,
+                            batchDefinition.EndColumn - 1),
                         index, outputFactory,
                         batchDefinition.SqlCmdCommand,
                         batchDefinition.BatchExecutionCount,
@@ -347,6 +347,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return Batches[batchIndex].GetSubset(resultSetIndex, startRow, rowCount);
         }
 
+        /// <summary>
+        /// Retrieves the column names for a result set inside a batch.
+        /// </summary>
+        /// <param name="batchIndex">The index for selecting the batch item</param>
+        /// <param name="resultSetIndex">The index for selecting the result set</param>
+        /// <returns>The column names</returns>
+        public List<string> GetColumnNames(int batchIndex, int resultSetIndex)
+        {
+            if (batchIndex < 0 || batchIndex >= Batches.Length)
+            {
+                throw new ArgumentOutOfRangeException(nameof(batchIndex));
+            }
+            return Batches[batchIndex].ResultSets[resultSetIndex].Columns.Select(c => c.ColumnName).ToList();
+        }
+
         /// <summary>
         /// Retrieves a subset of the result sets
         /// </summary>
@@ -388,7 +403,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <summary>
         /// Changes the OwnerURI for the editor connection.
         /// </summary>
-        public String ConnectionOwnerURI { 
+        public String ConnectionOwnerURI
+        {
             get
             {
                 return this.editorConnection.OwnerUri;
@@ -433,7 +449,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
                 // Locate and setup the connection
                 queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
-                onErrorAction = OnErrorAction.Ignore; 
+                onErrorAction = OnErrorAction.Ignore;
                 sqlConn = queryConnection as ReliableSqlConnection;
                 if (sqlConn != null)
                 {
@@ -642,12 +658,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 if (settings.StatisticsIO)
                 {
                     builderBefore.AppendFormat("{0} ", helper.GetSetStatisticsIOString(true));
-                    builderAfter.AppendFormat("{0} ", helper.GetSetStatisticsIOString (false));
+                    builderAfter.AppendFormat("{0} ", helper.GetSetStatisticsIOString(false));
                 }
 
                 if (settings.StatisticsTime)
                 {
-                    builderBefore.AppendFormat("{0} ", helper.GetSetStatisticsTimeString (true));
+                    builderBefore.AppendFormat("{0} ", helper.GetSetStatisticsTimeString(true));
                     builderAfter.AppendFormat("{0} ", helper.GetSetStatisticsTimeString(false));
                 }
             }
@@ -660,7 +676,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             // append first part of exec options
             builderBefore.AppendFormat("{0} {1} {2}",
-                helper.SetRowCountString,  helper.SetTextSizeString,  helper.SetNoCountString);
+                helper.SetRowCountString, helper.SetTextSizeString, helper.SetNoCountString);
 
             if (!connection.IsSqlDW)
             {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
index 14b622f3..594cdac4 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
@@ -6,11 +6,14 @@
 #nullable disable
 
 using System;
-using System.IO;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
+using System.Linq;
+using System.Text;
 using System.Threading.Tasks;
+using TextCopy;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection;
@@ -75,9 +78,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             get
             {
                 BufferFileStreamFactory ??= new ServiceBufferFileStreamFactory
-                    {
-                        QueryExecutionSettings = Settings.QueryExecutionSettings
-                    };
+                {
+                    QueryExecutionSettings = Settings.QueryExecutionSettings
+                };
                 return BufferFileStreamFactory;
             }
         }
@@ -187,6 +190,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(QueryExecutionPlanRequest.Type, HandleExecutionPlanRequest, true);
             serviceHost.SetRequestHandler(SimpleExecuteRequest.Type, HandleSimpleExecuteRequest, true);
             serviceHost.SetRequestHandler(QueryExecutionOptionsRequest.Type, HandleQueryExecutionOptionsRequest, true);
+            serviceHost.SetRequestHandler(CopyResultsRequest.Type, HandleCopyResultsRequest, true);
 
             // Register the file open update handler
             WorkspaceService<SqlToolsSettings>.Instance.RegisterTextDocCloseCallback(HandleDidCloseTextDocumentNotification);
@@ -713,6 +717,90 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return Task.CompletedTask;
         }
 
+        /// <summary>
+        /// Handles the copy results.
+        /// </summary>
+        internal async Task HandleCopyResultsRequest(CopyResultsRequestParams requestParams, RequestContext<CopyResultsRequestResult> requestContext)
+        {
+            var valueSeparator = "\t";
+            var columnRanges = this.MergeRanges(requestParams.Selections.Select(selection => new Range() { Start = selection.FromColumn, End = selection.ToColumn }).ToList());
+            var rowRanges = this.MergeRanges(requestParams.Selections.Select(selection => new Range() { Start = selection.FromRow, End = selection.ToRow }).ToList());
+            var lastColumnIndex = columnRanges.Last().End;
+            var lastRowIndex = rowRanges.Last().End;
+            var builder = new StringBuilder();
+            var pageSize = 200;
+            if (requestParams.IncludeHeaders)
+            {
+                Validate.IsNotNullOrEmptyString(nameof(requestParams.OwnerUri), requestParams.OwnerUri);
+                Query query;
+                if (!ActiveQueries.TryGetValue(requestParams.OwnerUri, out query))
+                {
+                    throw new ArgumentOutOfRangeException(SR.QueryServiceRequestsNoQuery);
+                }
+                var columnNames = query.GetColumnNames(requestParams.BatchIndex, requestParams.ResultSetIndex);
+                var selectedColumns = new List<string>();
+                for (int i = 0; i < columnNames.Count; i++)
+                {
+                    if (columnRanges.Any(range => i >= range.Start && i <= range.End))
+                    {
+                        selectedColumns.Add(columnNames[i]);
+                    }
+                }
+                builder.Append(string.Join(valueSeparator, selectedColumns));
+                builder.Append(Environment.NewLine);
+            }
+
+            for (int rowRangeIndex = 0; rowRangeIndex < rowRanges.Count; rowRangeIndex++)
+            {
+                var rowRange = rowRanges[rowRangeIndex];
+                var pageStartRowIndex = rowRange.Start;
+                // Read the rows in batches to avoid holding all rows in memory
+                do
+                {
+                    var rowsToFetch = Math.Min(pageSize, rowRange.End - pageStartRowIndex + 1);
+                    ResultSetSubset subset = await InterServiceResultSubset(new SubsetParams()
+                    {
+                        OwnerUri = requestParams.OwnerUri,
+                        ResultSetIndex = requestParams.ResultSetIndex,
+                        BatchIndex = requestParams.BatchIndex,
+                        RowsStartIndex = pageStartRowIndex,
+                        RowsCount = rowsToFetch
+                    });
+                    for (int rowIndex = 0; rowIndex < subset.Rows.Length; rowIndex++)
+                    {
+                        var row = subset.Rows[rowIndex];
+                        for (int columnRangeIndex = 0; columnRangeIndex < columnRanges.Count; columnRangeIndex++)
+                        {
+                            var columnRange = columnRanges[columnRangeIndex];
+                            for (int columnIndex = columnRange.Start; columnIndex <= columnRange.End; columnIndex++)
+                            {
+                                if (requestParams.Selections.Any(selection =>
+                                selection.FromRow <= rowIndex + rowRange.Start &&
+                                selection.ToRow >= rowIndex + rowRange.Start &&
+                                selection.FromColumn <= columnIndex &&
+                                selection.ToColumn >= columnIndex))
+                                {
+                                    builder.Append(requestParams.RemoveNewLines ? row[columnIndex].DisplayValue.ReplaceLineEndings(" ") : row[columnIndex].DisplayValue);
+                                }
+                                if (columnIndex != lastColumnIndex)
+                                {
+                                    builder.Append(valueSeparator);
+                                }
+                            }
+                        }
+                        // Add line break if this is not the last row in all selections.
+                        if (rowIndex + pageStartRowIndex != lastRowIndex)
+                        {
+                            builder.Append(Environment.NewLine);
+                        }
+                    }
+                    pageStartRowIndex += rowsToFetch;
+                } while (pageStartRowIndex < rowRange.End);
+            }
+            await ClipboardService.SetTextAsync(builder.ToString());
+            await requestContext.SendResult(new CopyResultsRequestResult());
+        }
+
         #endregion
 
         #region Private Helpers
@@ -1056,6 +1144,35 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return Task.FromResult(0);
         }
 
+        internal class Range
+        {
+            public int Start { get; set; }
+            public int End { get; set; }
+        }
+
+        internal List<Range> MergeRanges(List<Range> ranges)
+        {
+            var mergedRanges = new List<Range>();
+            ranges.Sort((range1, range2) => (range1.Start - range2.Start));
+            foreach (var range in ranges)
+            {
+                bool merged = false;
+                foreach (var mergedRange in mergedRanges)
+                {
+                    if (range.Start <= mergedRange.End)
+                    {
+                        mergedRange.End = Math.Max(range.End, mergedRange.End);
+                        merged = true;
+                        break;
+                    }
+                }
+                if (!merged)
+                {
+                    mergedRanges.Add(range);
+                }
+            }
+            return mergedRanges;
+        }
         #endregion
 
         #region IDisposable Implementation