Skip to content

Commit

Permalink
Apply query execution options once per connection session (#2379)
Browse files Browse the repository at this point in the history 
* Temp changes

* Test WIP

* Hook up session level query execution options

* Sort using statements

* Add nullref check for sqlconnection

* Add a test case that to confirm options persist across queries on same connection

* Handle document rename scenario

* Add a test for the SQL document rename scenario
  • Loading branch information
@kburtram
kburtram authored Jul 16, 2024
1 parent a995b07 commit 2259db5
Show file tree
Hide file tree
Showing 3 changed files with 303 additions and 48 deletions.
51 changes: 13 additions & 38 deletions src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,23 @@
#nullable disable

using System;
using System.Collections.Generic;
using System.Data;
using System.Data.Common;
using Microsoft.Data.SqlClient;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Data.SqlClient;
using Microsoft.SqlTools.ServiceLayer.BatchParser;
using Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode;
using Microsoft.SqlTools.ServiceLayer.Connection;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
using Microsoft.SqlTools.ServiceLayer.SqlContext;
using Microsoft.SqlTools.Utility;
using Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode;
using System.Collections.Generic;
using Microsoft.SqlTools.ServiceLayer.Utility;
using System.Text;
using Microsoft.SqlTools.Utility;

namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
Expand Down Expand Up @@ -163,6 +164,8 @@ public Query(
{
ApplyExecutionSettings(connection, settings, outputFactory);
}

this.Settings = settings;
}

#region Events
Expand Down Expand Up @@ -298,6 +301,11 @@ internal set
/// </summary>
public string QueryText { get; }

/// <summary>
/// Query execution settings
/// </summary>
public QueryExecutionSettings Settings { get; private set; }

#endregion

#region Public Methods
Expand Down Expand Up @@ -700,41 +708,8 @@ private void ApplyExecutionSettings(
builderAfter.AppendFormat("{0} ", helper.GetSetParseOnlyString(false));
}

// append first part of exec options
builderBefore.AppendFormat("{0} {1} {2}",
helper.SetRowCountString, helper.SetTextSizeString, helper.SetNoCountString);

if (!connection.IsSqlDW)
{
// append second part of exec options
builderBefore.AppendFormat(" {0} {1} {2} {3} {4} {5} {6}",
helper.SetConcatenationNullString,
helper.SetArithAbortString,
helper.SetLockTimeoutString,
helper.SetQueryGovernorCostString,
helper.SetDeadlockPriorityString,
helper.SetTransactionIsolationLevelString,
// We treat XACT_ABORT special in that we don't add anything if the option
// isn't checked. This is because we don't want to be overwriting the server
// if it has a default of ON since that's something people would specifically
// set and having a client change it could be dangerous (the reverse is much
// less risky)

// The full fix would probably be to make the options tri-state instead of
// just on/off, where the default is to use the servers default. Until that
// happens though this is the best solution we came up with. See TFS#7937925

// Note that users can always specifically add SET XACT_ABORT OFF to their
// queries if they do truly want to set it off. We just don't want to
// do it silently (since the default is going to be off)
settings.XactAbortOn ? helper.SetXactAbortString : string.Empty);

// append Ansi options
builderBefore.AppendFormat(" {0} {1} {2} {3} {4} {5} {6}",
helper.SetAnsiNullsString, helper.SetAnsiNullDefaultString, helper.SetAnsiPaddingString,
helper.SetAnsiWarningsString, helper.SetCursorCloseOnCommitString,
helper.SetImplicitTransactionString, helper.SetQuotedIdentifierString);

// "set noexec on" should be the very last command, cause everything after it is not
// being executed unitl "set noexec off" is encounered
// NOEXEC is not currently supported by SqlOnDemand servers
Expand Down
131 changes: 121 additions & 10 deletions src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,23 +8,27 @@
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Data;
using System.Data.Common;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Microsoft.Data.SqlClient;
using Microsoft.SqlTools.Hosting.Protocol;
using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
using Microsoft.SqlTools.ServiceLayer.Connection;
using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
using Microsoft.SqlTools.ServiceLayer.ExecutionPlan;
using Microsoft.SqlTools.ServiceLayer.ExecutionPlan.Contracts;
using Microsoft.SqlTools.ServiceLayer.Hosting;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteRequests;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
using Microsoft.SqlTools.ServiceLayer.ExecutionPlan;
using Microsoft.SqlTools.ServiceLayer.SqlContext;
using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
using Microsoft.SqlTools.ServiceLayer.Workspace;
using Microsoft.SqlTools.Utility;
using Microsoft.SqlTools.ServiceLayer.ExecutionPlan.Contracts;

namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
Expand Down Expand Up @@ -123,6 +127,11 @@ private IFileStreamFactory BufferFileFactory
/// </summary>
internal ConcurrentDictionary<string, QueryExecutionSettings> ActiveQueryExecutionSettings => queryExecutionSettings.Value;

/// <summary>
/// The collection of query session execution options tracking flags
/// </summary>
internal ConcurrentDictionary<string, bool> QuerySessionSettingsApplied => querySessionSettingsApplied.Value;

/// <summary>
/// Internal task for testability
/// </summary>
Expand All @@ -147,6 +156,12 @@ private IFileStreamFactory BufferFileFactory
private readonly Lazy<ConcurrentDictionary<string, QueryExecutionSettings>> queryExecutionSettings =
new Lazy<ConcurrentDictionary<string, QueryExecutionSettings>>(() => new ConcurrentDictionary<string, QueryExecutionSettings>());

/// <summary>
/// Internal storage of tracking flags for whether query sessions settings have been applied
/// </summary>
private readonly Lazy<ConcurrentDictionary<string, bool>> querySessionSettingsApplied =
new Lazy<ConcurrentDictionary<string, bool>>(() => new ConcurrentDictionary<string, bool>());

/// <summary>
/// Settings that will be used to execute queries. Internal for unit testing
/// </summary>
Expand Down Expand Up @@ -204,6 +219,9 @@ public void InitializeService(ServiceHost serviceHost)

// Register a handler for when the configuration changes
WorkspaceService.RegisterConfigChangeCallback(UpdateSettings);

// Register a callback for when a connection is created
ConnectionService.RegisterOnConnectionTask(OnNewConnection);
}

#region Request Handlers
Expand Down Expand Up @@ -363,17 +381,26 @@ internal Task HandleConnectionUriChangedNotification(ConnectionUriChangedParams
{
try
{
string OriginalOwnerUri = changeUriParams.OriginalOwnerUri;
string NewOwnerUri = changeUriParams.NewOwnerUri;
ConnectionService.ReplaceUri(OriginalOwnerUri, NewOwnerUri);
string originalOwnerUri = changeUriParams.OriginalOwnerUri;
string newOwnerUri = changeUriParams.NewOwnerUri;
ConnectionService.ReplaceUri(originalOwnerUri, newOwnerUri);
// Attempt to load the query
Query query;
if (!ActiveQueries.TryRemove(OriginalOwnerUri, out query))
if (!ActiveQueries.TryRemove(originalOwnerUri, out query))
{
throw new Exception("Uri: " + OriginalOwnerUri + " is not associated with an active query.");
throw new Exception("Uri: " + originalOwnerUri + " is not associated with an active query.");
}
query.ConnectionOwnerURI = NewOwnerUri;
ActiveQueries.TryAdd(NewOwnerUri, query);
query.ConnectionOwnerURI = newOwnerUri;
ActiveQueries.TryAdd(newOwnerUri, query);

// Update the session query execution options applied map
bool settingsApplied;
if (!QuerySessionSettingsApplied.TryRemove(originalOwnerUri, out settingsApplied))
{
throw new Exception("Uri: " + originalOwnerUri + " is not associated with an active query settings.");
}
QuerySessionSettingsApplied.TryAdd(newOwnerUri, settingsApplied);

return Task.FromResult(true);
}
catch (Exception ex)
Expand Down Expand Up @@ -627,6 +654,23 @@ public async Task InterServiceExecuteQuery(ExecuteRequestParamsBase executeParam
return;
}

if (connInfo == null)
{
ConnectionService.TryFindConnection(executeParams.OwnerUri, out connInfo);
}

bool sessionSettingsApplied;
if (!this.QuerySessionSettingsApplied.TryGetValue(connInfo.OwnerUri, out sessionSettingsApplied))
{
sessionSettingsApplied = false;
}

if (!sessionSettingsApplied)
{
ApplySessionQueryExecutionOptions(connInfo, newQuery.Settings);
this.QuerySessionSettingsApplied.AddOrUpdate(connInfo.OwnerUri, true, (key, oldValue) => true);
}

// Execute the query asynchronously
ExecuteAndCompleteQuery(executeParams.OwnerUri, newQuery, queryEventSender, querySuccessFunc, queryFailureFunc);
}
Expand Down Expand Up @@ -717,6 +761,18 @@ public Task HandleDidCloseTextDocumentNotification(
return Task.CompletedTask;
}

/// <summary>
/// Runs UpdateLanguageServiceOnConnection as a background task
/// </summary>
/// <param name="info">Connection Info</param>
/// <returns></returns>
private Task OnNewConnection(ConnectionInfo info)
{
this.QuerySessionSettingsApplied.AddOrUpdate(info.OwnerUri, false, (key, oldValue) => false);

return Task.CompletedTask;
}

/// <summary>
/// Handles the copy results.
/// </summary>
Expand Down Expand Up @@ -1197,6 +1253,61 @@ internal List<Range> MergeRanges(List<Range> ranges)
}
return mergedRanges;
}

private async void ApplySessionQueryExecutionOptions(ConnectionInfo connection, QueryExecutionSettings settings)
{
QuerySettingsHelper helper = new QuerySettingsHelper(settings);

StringBuilder sqlBuilder = new StringBuilder(512);

// append first part of exec options
sqlBuilder.AppendFormat("{0} {1} {2}",
helper.SetRowCountString, helper.SetTextSizeString, helper.SetNoCountString);

if (!connection.IsSqlDW)
{
// append second part of exec options
sqlBuilder.AppendFormat(" {0} {1} {2} {3} {4} {5} {6}",
helper.SetConcatenationNullString,
helper.SetArithAbortString,
helper.SetLockTimeoutString,
helper.SetQueryGovernorCostString,
helper.SetDeadlockPriorityString,
helper.SetTransactionIsolationLevelString,
// We treat XACT_ABORT special in that we don't add anything if the option
// isn't checked. This is because we don't want to be overwriting the server
// if it has a default of ON since that's something people would specifically
// set and having a client change it could be dangerous (the reverse is much
// less risky)

// The full fix would probably be to make the options tri-state instead of
// just on/off, where the default is to use the servers default. Until that
// happens though this is the best solution we came up with. See TFS#7937925

// Note that users can always specifically add SET XACT_ABORT OFF to their
// queries if they do truly want to set it off. We just don't want to
// do it silently (since the default is going to be off)
settings.XactAbortOn ? helper.SetXactAbortString : string.Empty);

// append Ansi options
sqlBuilder.AppendFormat(" {0} {1} {2} {3} {4} {5} {6}",
helper.SetAnsiNullsString, helper.SetAnsiNullDefaultString, helper.SetAnsiPaddingString,
helper.SetAnsiWarningsString, helper.SetCursorCloseOnCommitString,
helper.SetImplicitTransactionString, helper.SetQuotedIdentifierString);
}

DbConnection dbConnection = await ConnectionService.GetOrOpenConnection(connection.OwnerUri, ConnectionType.Default);
ReliableSqlConnection reliableSqlConnection = dbConnection as ReliableSqlConnection;
if (reliableSqlConnection != null)
{
using (SqlCommand cmd = new SqlCommand(sqlBuilder.ToString(), reliableSqlConnection.GetUnderlyingConnection()))
{
cmd.CommandType = CommandType.Text;
cmd.ExecuteNonQuery();
}
}
}

#endregion

#region IDisposable Implementation
Expand Down
Loading

0 comments on commit 2259db5

Please sign in to comment.