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();