ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-15 17:23:32 -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

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;