ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-13 17:23:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

Support updating access token when found expired for OE (#1772)

Browse Source
This commit is contained in:
Cheena Malhotra
2022-11-28 17:42:05 -08:00
committed by GitHub
parent 4866515bde
commit 5cc3e6f657
15 changed files with 194 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs
View File

@@ -18,6 +18,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
/// </summary>
public class ConnectionInfo
{
private static readonly DateTime UnixEpochUtc = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
/// <summary>
/// Constructor
/// </summary>
@@ -27,7 +29,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
OwnerUri = ownerUri;
ConnectionDetails = details;
ConnectionId = Guid.NewGuid();
IntellisenseMetrics = new InteractionMetrics<double>(new int[] {50, 100, 200, 500, 1000, 2000});
IntellisenseMetrics = new InteractionMetrics<double>(new int[] { 50, 100, 200, 500, 1000, 2000 });
}
/// <summary>
@@ -72,11 +74,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
/// Returns true if the db connection is to a SQL db instance
/// </summary>
public bool IsSqlDb { get; set; }
/// <summary>
/// Returns true if the sql connection is to a DW instance
/// </summary>
public bool IsSqlDW { get; set; }
/// <summary>
/// Returns true if Authentication mode is AzureMFA, determines if Access token is in use.
/// </summary>
public bool IsAzureAuth { get; set; }
/// <summary>
/// Returns the connection Engine Edition
/// </summary>
@@ -158,8 +166,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
public void RemoveConnection(string connectionType)
{
Validate.IsNotNullOrEmptyString("Connection Type", connectionType);
DbConnection connection;
ConnectionTypeToConnectionMap.TryRemove(connectionType, out connection);
ConnectionTypeToConnectionMap.TryRemove(connectionType, out _);
}
/// <summary>
@@ -169,18 +176,44 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
{
foreach (var type in AllConnectionTypes)
{
DbConnection connection;
ConnectionTypeToConnectionMap.TryRemove(type, out connection);
ConnectionTypeToConnectionMap.TryRemove(type, out _);
}
}
}
/// <summary>
/// Updates the Auth Token and Expires On fields
/// </summary>
public void UpdateAuthToken(string token, int expiresOn)
public bool TryUpdateAccessToken(SecurityToken? securityToken)
{
ConnectionDetails.AzureAccountToken = token;
ConnectionDetails.ExpiresOn = expiresOn;
if (securityToken != null && !string.IsNullOrEmpty(securityToken.Token) && IsAzureAuth && IsAccessTokenExpired)
{
ConnectionDetails.AzureAccountToken = securityToken.Token;
ConnectionDetails.ExpiresOn = securityToken.ExpiresOn;
return true;
}
return false;
}
/// <summary>
/// Returns true if Access token saved in connection details is expired or about to expire in 2 minutes.
/// </summary>
private bool IsAccessTokenExpired
{
get
{
if (IsAzureAuth && ConnectionDetails.ExpiresOn != null && double.TryParse(ConnectionDetails.ExpiresOn.ToString(), out var expiresOn))
{
DateTime dateTime = UnixEpochUtc.AddSeconds(expiresOn);
// Check if access token is already expired or shall expire in 2 minutes.
if (dateTime <= DateTime.UtcNow.AddMinutes(2))
{
Logger.Verbose("Access token found expired.");
return true;
}
}
return false;
}
}
}
}

4
src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
View File

@@ -304,7 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
return;
}
this.TokenUpdateUris.Remove(tokenRefreshedParams.Uri, out var result);
connection.UpdateAuthToken(tokenRefreshedParams.Token, tokenRefreshedParams.ExpiresOn);
connection.TryUpdateAccessToken(new SecurityToken() { Token = tokenRefreshedParams.Token, ExpiresOn = tokenRefreshedParams.ExpiresOn });
}
/// <summary>
@@ -585,6 +585,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
connectionInfo.MajorVersion = serverInfo.ServerMajorVersion;
connectionInfo.IsSqlDb = serverInfo.EngineEditionId == (int)DatabaseEngineEdition.SqlDatabase;
connectionInfo.IsSqlDW = (serverInfo.EngineEditionId == (int)DatabaseEngineEdition.SqlDataWarehouse);
// Determines that access token is used for creating connection.
connectionInfo.IsAzureAuth = connectionInfo.ConnectionDetails.AuthenticationType == "AzureMFA";
connectionInfo.EngineEdition = (DatabaseEngineEdition)serverInfo.EngineEditionId;
// Azure Data Studio supports SQL Server 2014 and later releases.
response.IsSupportedVersion = serverInfo.IsCloud || serverInfo.ServerMajorVersion >= 12;

8
src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/RefreshTokenNotification.cs
View File

@@ -7,7 +7,7 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
{
class RefreshTokenParams
public class RefreshTokenParams
{
/// <summary>
/// ID of the tenant
@@ -39,14 +39,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
/// <summary>
/// Refresh token request mapping entry
/// </summary>
class RefreshTokenNotification
public class RefreshTokenNotification
{
public static readonly
EventType<RefreshTokenParams> Type =
EventType<RefreshTokenParams>.Create("account/refreshToken");
}
class TokenRefreshedParams
public class TokenRefreshedParams
{
/// <summary>
/// Gets or sets the refresh token.
@@ -64,7 +64,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
public string Uri { get; set; }
}
class TokenRefreshedNotification
public class TokenRefreshedNotification
{
public static readonly
EventType<TokenRefreshedParams> Type =

25
src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/SecurityToken.cs Normal file
View File

@@ -0,0 +1,25 @@
//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
//
namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
{
public class SecurityToken
{
/// <summary>
/// Gets or sets the refresh token.
/// </summary>
public string Token { get; set; }
/// <summmary>
/// Gets or sets the token expiration, a Unix epoch
/// </summary>
public int? ExpiresOn { get; set; }
/// <summmary>
/// Gets or sets the token type, e.g. 'Bearer'
/// </summary>
public string? TokenType { get; set; }
}
}

2
src/Microsoft.SqlTools.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs
View File

@@ -7,7 +7,7 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
namespace Microsoft.SqlTools.ServiceLayer.Connection.Contracts
{
class RequestSecurityTokenParams
class RequestSecurityTokenParams
{
/// <summary>
/// Gets or sets the address of the authority to issue token.

6
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/ExpandRequest.cs
View File

@@ -4,6 +4,7 @@
//
using Microsoft.SqlTools.Hosting.Protocol.Contracts;
using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
{
@@ -49,6 +50,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
/// Path identifying the node to expand. See <see cref="NodeInfo.NodePath"/> for details
/// </summary>
public string NodePath { get; set; }
/// <summary>
/// Security token for AzureMFA authentication for refresing access token on connection.
/// </summary>
public SecurityToken? SecurityToken { get; set; }
}
/// <summary>

2
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/RefreshRequest.cs
View File

@@ -10,7 +10,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
/// <summary>
/// Parameters to the <see cref="ExpandRequest"/>.
/// </summary>
public class RefreshParams: ExpandParams
public class RefreshParams : ExpandParams
{
}

17
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs
View File

@@ -240,14 +240,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
/// Expands this node and returns its children
/// </summary>
/// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
public IList<TreeNode> Expand(string name, CancellationToken cancellationToken)
public IList<TreeNode> Expand(string name, CancellationToken cancellationToken, string? accessToken = null)
{
// TODO consider why solution explorer has separate Children and Items options
if (children.IsInitialized)
{
return children;
}
PopulateChildren(false, name, cancellationToken);
PopulateChildren(false, name, cancellationToken, accessToken);
return children;
}
@@ -255,19 +255,19 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
/// Expands this node and returns its children
/// </summary>
/// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
public IList<TreeNode> Expand(CancellationToken cancellationToken)
public IList<TreeNode> Expand(CancellationToken cancellationToken, string? accessToken = null)
{
return Expand(null, cancellationToken);
return Expand(null, cancellationToken, accessToken);
}
/// <summary>
/// Refresh this node and returns its children
/// </summary>
/// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
public virtual IList<TreeNode> Refresh(CancellationToken cancellationToken)
public virtual IList<TreeNode> Refresh(CancellationToken cancellationToken, string? accessToken = null)
{
// TODO consider why solution explorer has separate Children and Items options
PopulateChildren(true, null, cancellationToken);
PopulateChildren(true, null, cancellationToken, accessToken);
return children;
}
@@ -319,7 +319,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
return Parent as T;
}
protected virtual void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken)
protected virtual void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken, string? accessToken = null)
{
Logger.Write(TraceEventType.Verbose, string.Format(CultureInfo.InvariantCulture, "Populating oe node :{0}", this.GetNodePath()));
Debug.Assert(IsAlwaysLeaf == false);
@@ -335,6 +335,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
children.Clear();
BeginChildrenInit();
// Update access token for future queries
context.UpdateAccessToken(accessToken);
try
{
ErrorMessage = null;

22
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs
View File

@@ -212,7 +212,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
}
else
{
RunExpandTask(session, expandParams);
bool refreshNeeded = session.ConnectionInfo.TryUpdateAccessToken(expandParams.SecurityToken);
RunExpandTask(session, expandParams, refreshNeeded);
return true;
}
};
@@ -239,6 +240,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
}
else
{
session.ConnectionInfo.TryUpdateAccessToken(refreshParams.SecurityToken);
RunExpandTask(session, refreshParams, true);
}
await context.SendResult(true);
@@ -373,12 +375,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
}
internal Task<ExpandResponse> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
internal Task<ExpandResponse> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false, SecurityToken? securityToken = null)
{
return Task.Run(() => QueueExpandNodeRequest(session, nodePath, forceRefresh));
return Task.Run(() => QueueExpandNodeRequest(session, nodePath, forceRefresh, securityToken));
}
internal ExpandResponse QueueExpandNodeRequest(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
internal ExpandResponse QueueExpandNodeRequest(ObjectExplorerSession session, string nodePath, bool forceRefresh = false, SecurityToken? securityToken = null)
{
NodeInfo[] nodes = null;
TreeNode node = session.Root.FindNodeByPath(nodePath);
@@ -432,15 +434,21 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
waitForLockTimeout: timeout,
bindOperation: (bindingContext, cancelToken) =>
{
if (!session.ConnectionInfo.IsAzureAuth)
{
// explicitly set null here to prevent setting access token for non-Azure auth modes.
securityToken = null;
}
if (forceRefresh)
{
Logger.Verbose($"Forcing refresh for {nodePath}");
nodes = node.Refresh(cancelToken).Select(x => x.ToNodeInfo()).ToArray();
nodes = node.Refresh(cancelToken, securityToken?.Token).Select(x => x.ToNodeInfo()).ToArray();
}
else
{
Logger.Verbose($"Expanding {nodePath}");
nodes = node.Expand(cancelToken).Select(x => x.ToNodeInfo()).ToArray();
nodes = node.Expand(cancelToken, securityToken?.Token).Select(x => x.ToNodeInfo()).ToArray();
}
response.Nodes = nodes;
response.ErrorMessage = node.ErrorMessage;
@@ -635,7 +643,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
private async Task ExpandNodeAsync(ObjectExplorerSession session, ExpandParams expandParams, CancellationToken cancellationToken, bool forceRefresh = false)
{
ExpandResponse response = null;
response = await ExpandNode(session, expandParams.NodePath, forceRefresh);
response = await ExpandNode(session, expandParams.NodePath, forceRefresh, expandParams.SecurityToken);
if (cancellationToken.IsCancellationRequested)
{
Logger.Write(TraceEventType.Verbose, "OE expand canceled");

6
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs
View File

@@ -16,7 +16,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
{
internal partial class DatabaseTreeNode
{
public DatabaseTreeNode(ServerNode serverNode, string databaseName): this()
public DatabaseTreeNode(ServerNode serverNode, string databaseName) : this()
{
Parent = serverNode;
NodeValue = databaseName;
@@ -52,12 +52,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
}
}
protected override void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken)
protected override void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken, string? accessToken = null)
{
var smoQueryContext = this.GetContextAs<SmoQueryContext>();
if (IsAccessible(smoQueryContext))
{
base.PopulateChildren(refresh, name, cancellationToken);
base.PopulateChildren(refresh, name, cancellationToken, accessToken);
}
else
{

27
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoExtensions.cs Normal file
View File

@@ -0,0 +1,27 @@
//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
//
using Microsoft.SqlServer.Management.Smo;
using Microsoft.SqlTools.ServiceLayer.Connection;
namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
{
internal static class SmoExtensions
{
/// <summary>
/// Updates access token on the connection context of <paramref name="sqlObj"/> instance.
/// </summary>
/// <param name="sqlObj">(this) SMO SQL Object containing connection context.</param>
/// <param name="accessToken">Access token</param>
public static void UpdateAccessToken(this SqlSmoObject sqlObj, string accessToken)
{
if (sqlObj != null && !string.IsNullOrEmpty(accessToken)
&& sqlObj.ExecutionManager != null
&& sqlObj.ExecutionManager.ConnectionContext != null)
{
sqlObj.ExecutionManager.ConnectionContext.AccessToken = new AzureAccessToken(accessToken);
}
}
}
}

42
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoQueryContext.cs
View File

@@ -46,17 +46,19 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
/// <summary>
/// The server SMO will query against
/// </summary>
public Server Server {
public Server Server
{
get
{
return GetObjectWithOpenedConnection(server);
}
}
}
/// <summary>
/// Optional Database context object to query against
/// </summary>
public Database Database {
public Database Database
{
get
{
return GetObjectWithOpenedConnection(database);
@@ -70,7 +72,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
/// <summary>
/// Parent of a give node to use for queries
/// </summary>
public SmoObjectBase Parent {
public SmoObjectBase Parent
{
get
{
return GetObjectWithOpenedConnection(parent);
@@ -86,7 +89,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
/// for specific SMO types
/// </summary>
public IMultiServiceProvider ServiceProvider { get; private set; }
/// <summary>
/// Helper method to cast a parent to a specific type
/// </summary>
@@ -116,7 +119,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
ObjectExplorerService service = ServiceProvider.GetService<ObjectExplorerService>();
if (service == null)
{
throw new InvalidOperationException(string.Format(CultureInfo.CurrentCulture,
throw new InvalidOperationException(string.Format(CultureInfo.CurrentCulture,
SqlTools.Hosting.SR.ServiceNotFound, nameof(ObjectExplorerService)));
}
@@ -147,7 +150,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
{
get
{
if(validFor == 0)
if (validFor == 0)
{
validFor = ServerVersionHelper.GetValidForFlag(SqlServerType, Database);
}
@@ -169,6 +172,31 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
return smoObj;
}
/// <summary>
/// Updates access token on parent connection context.
/// </summary>
/// <param name="accessToken">Acquired access token</param>
public void UpdateAccessToken(string? accessToken)
{
if (!string.IsNullOrEmpty(accessToken))
{
// Update all applicable nodes that could contain access token
// to prevent stale token from being re-used.
if (server != null)
{
(server as SqlSmoObject).UpdateAccessToken(accessToken);
}
if (database != null)
{
(database as SqlSmoObject).UpdateAccessToken(accessToken);
}
if (parent != null)
{
(parent as SqlSmoObject).UpdateAccessToken(accessToken);
}
}
}
/// <summary>
/// Ensures the server objects connection context is open. This is used by all child objects, and
/// the only way to easily access is via the server object. This should be called during access of

14
src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoWrapper.cs
View File

@@ -14,11 +14,21 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
/// </summary>
internal class SmoWrapper
{
/// <summary>
/// Creates instance of <see cref="Server"/> from provided <paramref name="serverConn"/> instance.
/// </summary>
/// <param name="serverConn">Server connection instance.</param>
/// <returns>Server instance.</returns>
public virtual Server CreateServer(ServerConnection serverConn)
{
return serverConn == null ? null : new Server(serverConn);
}
/// <summary>
/// Checks if connection is open on the <paramref name="smoObj"/> instance.
/// </summary>
/// <param name="smoObj">SMO Object containing connection context.</param>
/// <returns>True if connection is open, otherwise false.</returns>
public virtual bool IsConnectionOpen(SmoObjectBase smoObj)
{
SqlSmoObject sqlObj = smoObj as SqlSmoObject;
@@ -28,6 +38,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
&& sqlObj.ExecutionManager.ConnectionContext.IsOpen;
}
/// <summary>
/// Opens connection on the connection context of <paramref name="smoObj"/> instance.
/// </summary>
/// <param name="smoObj">SMO Object containing connection context.</param>
public virtual void OpenConnection(SmoObjectBase smoObj)
{
SqlSmoObject sqlObj = smoObj as SqlSmoObject;

11
test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Utility/LiveConnectionHelper.cs
View File

@@ -18,9 +18,10 @@ using System.Threading;
namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility
{
public class LiveConnectionException : Exception {
public LiveConnectionException(string message)
: base(message) { }
public class LiveConnectionException : Exception
{
public LiveConnectionException(string message)
: base(message) { }
}
public class LiveConnectionHelper
@@ -48,7 +49,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility
public static TestConnectionResult InitLiveConnectionInfo(string databaseName = null, string ownerUri = null)
=> InitLiveConnectionInfoAsync(databaseName, ownerUri, ServiceLayer.Connection.ConnectionType.Default).ConfigureAwait(false).GetAwaiter().GetResult();
public static async Task<TestConnectionResult> InitLiveConnectionInfoAsync(string databaseName = "master", string ownerUri = null,
public static async Task<TestConnectionResult> InitLiveConnectionInfoAsync(string databaseName = "master", string ownerUri = null,
string connectionType = ServiceLayer.Connection.ConnectionType.Default, TestServerType serverType = TestServerType.OnPrem)
{
ScriptFile scriptFile = null;
@@ -58,7 +59,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility
scriptFile = TestServiceProvider.Instance.WorkspaceService.Workspace.GetFile(ownerUri);
ownerUri = scriptFile.ClientUri;
}
if (string.IsNullOrEmpty(databaseName))
if (string.IsNullOrEmpty(databaseName))
{
databaseName = "master";
}

2
test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs
View File

@@ -295,7 +295,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
public void FindNodeCanExpandParentNodes()
{
var mockTreeNode = new Mock<TreeNode>();
object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>(), new CancellationToken() };
object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>(), new CancellationToken(), ItExpr.IsNull<string>() };
mockTreeNode.Protected().Setup("PopulateChildren", populateChildrenArguments);
mockTreeNode.Object.IsAlwaysLeaf = false;