diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ReliableConnection/ReliableSqlConnection.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ReliableConnection/ReliableSqlConnection.cs index 8b66ade8..37ae5cd6 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ReliableConnection/ReliableSqlConnection.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ReliableConnection/ReliableSqlConnection.cs @@ -31,6 +31,8 @@ using System.Data.SqlClient; using System.Diagnostics; using System.Globalization; using System.Text; +using System.Threading; +using System.Threading.Tasks; using Microsoft.SqlTools.ServiceLayer.Utility; namespace Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection @@ -260,7 +262,47 @@ SET NUMERIC_ROUNDABORT OFF;"; /// public override void Open() { - OpenConnection(); + // Check if retry policy was specified, if not, disable retries by executing the Open method using RetryPolicy.NoRetry. + _connectionRetryPolicy.ExecuteAction(() => + { + if (_underlyingConnection.State != ConnectionState.Open) + { + _underlyingConnection.Open(); + } + SetLockAndCommandTimeout(_underlyingConnection); + SetDefaultAnsiSettings(_underlyingConnection, IsSqlDwConnection(_underlyingConnection)); + }); + } + + /// + /// Opens a database connection with the settings specified by the ConnectionString + /// property of the provider-specific Connection object. + /// + public override Task OpenAsync(CancellationToken token) + { + // Make sure that the token isn't cancelled before we try + if (token.IsCancellationRequested) + { + return Task.FromCanceled(token); + } + + // Check if retry policy was specified, if not, disable retries by executing the Open method using RetryPolicy.NoRetry. + try + { + return _connectionRetryPolicy.ExecuteAction(async () => + { + if (_underlyingConnection.State != ConnectionState.Open) + { + await _underlyingConnection.OpenAsync(token); + } + SetLockAndCommandTimeout(_underlyingConnection); + SetDefaultAnsiSettings(_underlyingConnection, IsSqlDwConnection(_underlyingConnection)); + }); + } + catch (Exception e) + { + return Task.FromException(e); + } } /// @@ -356,26 +398,6 @@ SET NUMERIC_ROUNDABORT OFF;"; RetryPolicyUtils.RaiseSchemaAmbientRetryMessage(retryState, SqlSchemaModelErrorCodes.ServiceActions.ConnectionRetry, _azureSessionId); } - /// - /// Opens a database connection with the settings specified by the ConnectionString and ConnectionRetryPolicy properties. - /// - /// An object representing the open connection. - private SqlConnection OpenConnection() - { - // Check if retry policy was specified, if not, disable retries by executing the Open method using RetryPolicy.NoRetry. - _connectionRetryPolicy.ExecuteAction(() => - { - if (_underlyingConnection.State != ConnectionState.Open) - { - _underlyingConnection.Open(); - } - SetLockAndCommandTimeout(_underlyingConnection); - SetDefaultAnsiSettings(_underlyingConnection, IsSqlDwConnection(_underlyingConnection)); - }); - - return _underlyingConnection; - } - public void OnConnectionStateChange(object sender, StateChangeEventArgs e) { SqlConnection conn = (SqlConnection)sender;