Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Apply query execution options once per connection session #2379

Merged
merged 9 commits into from
Jul 16, 2024
Merged

Apply query execution options once per connection session #2379

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
734b249
Temp changes
kburtram Jul 9, 2024
42b1848
Merge branch 'main' into dev/karlb/isolev
kburtram Jul 9, 2024
73a9ffc
Test WIP
kburtram Jul 11, 2024
c4fd66d
Hook up session level query execution options
kburtram Jul 15, 2024
8e984df
Sort using statements
kburtram Jul 15, 2024
a3d3c73
Add nullref check for sqlconnection
kburtram Jul 16, 2024
54f174d
Add a test case that to confirm options persist across queries on sam…
kburtram Jul 16, 2024
a973427
Handle document rename scenario
kburtram Jul 16, 2024
19c8d59
Add a test for the SQL document rename scenario
kburtram Jul 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 13 additions & 38 deletions src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
View file
Open in desktop
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
Open in desktop
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