Add flag so that full column schema is only fetched with EditData queries. (#615)

src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs

@@ -313,6 +313,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             ExecuteStringParams executeParams = new ExecuteStringParams
             {
                 Query = query,
+                GetFullColumnSchema = true,
                 OwnerUri = ownerUri
             };
             await queryExecutionService.InterServiceExecuteQuery(executeParams, null, eventSender,

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

@@ -61,10 +61,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private readonly SpecialAction specialAction;
 
+        /// <summary>
+        /// Flag indicating whether a separate KeyInfo query should be run
+        /// to get the full ColumnSchema metadata.
+        /// </summary>
+        private readonly bool getFullColumnSchema;
+
         #endregion
 
         internal Batch(string batchText, SelectionData selection, int ordinalId,
-            IFileStreamFactory outputFileFactory, int executionCount = 1)
+            IFileStreamFactory outputFileFactory, int executionCount = 1, bool getFullColumnSchema = false)
         {
             // Sanity check for input
             Validate.IsNotNullOrEmptyString(nameof(batchText), batchText);
@@ -81,6 +87,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             this.outputFileFactory = outputFileFactory;
             specialAction = new SpecialAction();
             BatchExecutionCount = executionCount > 0 ? executionCount : 1;
+
+            this.getFullColumnSchema = getFullColumnSchema;
         }
 
         #region Events
@@ -347,18 +355,22 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 dbCommand.CommandTimeout = 0;
                 executionStartTime = DateTime.Now;
 
-                // Fetch schema info separately, since CommandBehavior.KeyInfo will include primary
+                List<DbColumn[]> columnSchemas = null;
-                // key columns in the result set, even if they weren't part of the select statement.
+                if (getFullColumnSchema)
-                // Extra key columns get added to the end, so just correlate via Column Ordinal.
-                List<DbColumn[]> columnSchemas = new List<DbColumn[]>();
-                using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken))
                 {
-                    if (reader != null && reader.CanGetColumnSchema())
+                    // Fetch schema info separately, since CommandBehavior.KeyInfo will include primary
+                    // key columns in the result set, even if they weren't part of the select statement.
+                    // Extra key columns get added to the end, so just correlate via Column Ordinal.
+                    columnSchemas = new List<DbColumn[]>();
+                    using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken))
                     {
-                        do
+                        if (reader != null && reader.CanGetColumnSchema())
                         {
-                            columnSchemas.Add(reader.GetColumnSchema().ToArray());
+                            do
-                        } while (await reader.NextResultAsync(cancellationToken));
+                            {
+                                columnSchemas.Add(reader.GetColumnSchema().ToArray());
+                            } while (await reader.NextResultAsync(cancellationToken));
+                        }
                     }
                 }
 

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/ExecuteRequestParamsBase.cs

@@ -19,5 +19,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteReques
         /// Execution plan options
         /// </summary>
         public ExecutionPlanOptions ExecutionPlanOptions { get; set; }
+
+        /// <summary>
+        /// Flag to get full column schema via additional queries.
+        /// </summary>
+        public bool GetFullColumnSchema { get; set; }
     }
 }

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

@@ -93,7 +93,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <param name="connection">The information of the connection to use to execute the query</param>
         /// <param name="settings">Settings for how to execute the query, from the user</param>
         /// <param name="outputFactory">Factory for creating output files</param>
-        public Query(string queryText, ConnectionInfo connection, QueryExecutionSettings settings, IFileStreamFactory outputFactory)
+        public Query(string queryText, ConnectionInfo connection, QueryExecutionSettings settings, IFileStreamFactory outputFactory, bool getFullColumnSchema = false)
         {
             // Sanity check for input
             Validate.IsNotNull(nameof(queryText), queryText);
@@ -119,7 +119,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                             batchDefinition.EndLine-1,
                             batchDefinition.EndColumn-1),                       
                         index, outputFactory,
-                        batchDefinition.BatchExecutionCount));
+                        batchDefinition.BatchExecutionCount,
+                        getFullColumnSchema));
 
             Batches = batchSelection.ToArray();
 

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

@@ -599,7 +599,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             settings.ExecutionPlanOptions = executeParams.ExecutionPlanOptions;
 
             // If we can't add the query now, it's assumed the query is in progress
-            Query newQuery = new Query(GetSqlText(executeParams), connectionInfo, settings, BufferFileFactory);
+            Query newQuery = new Query(GetSqlText(executeParams), connectionInfo, settings, BufferFileFactory, executeParams.GetFullColumnSchema);
             if (!ActiveQueries.TryAdd(executeParams.OwnerUri, newQuery))
             {
                 newQuery.Dispose();