diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs index ad1139ad..ec29013f 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs +++ b/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, diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs index 96f85cac..9c43fb81 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs @@ -61,10 +61,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution /// private readonly SpecialAction specialAction; + /// + /// Flag indicating whether a separate KeyInfo query should be run + /// to get the full ColumnSchema metadata. + /// + 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 - // 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. - List columnSchemas = new List(); - using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken)) + List columnSchemas = null; + if (getFullColumnSchema) { - 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(); + using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken)) { - do + if (reader != null && reader.CanGetColumnSchema()) { - columnSchemas.Add(reader.GetColumnSchema().ToArray()); - } while (await reader.NextResultAsync(cancellationToken)); + do + { + columnSchemas.Add(reader.GetColumnSchema().ToArray()); + } while (await reader.NextResultAsync(cancellationToken)); + } } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/ExecuteRequestParamsBase.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/ExecuteRequestParamsBase.cs index 77e82320..a062eeb1 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/ExecuteRequestParamsBase.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecuteRequests/ExecuteRequestParamsBase.cs @@ -19,5 +19,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteReques /// Execution plan options /// public ExecutionPlanOptions ExecutionPlanOptions { get; set; } + + /// + /// Flag to get full column schema via additional queries. + /// + public bool GetFullColumnSchema { get; set; } } } diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs index ad9e1a5a..5c49f560 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs @@ -93,7 +93,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution /// The information of the connection to use to execute the query /// Settings for how to execute the query, from the user /// Factory for creating output files - 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(); diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs index 03807853..e8b76a6a 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs +++ b/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();