How to handle Microsoft.Data.SqlClient new AccessTokenCallback

[kevinharing]

In version 5.2.0 of Microsoft.Data.SqlClient a new callback was added to the SqlConnection class in order to retrieve an access token. What is the best place to set this property. I don’t think it’s necessary to set it on each execution of the IDbConnection interceptor like I did before with the AccessToken property, right?

More info on the new functionality here.

Previous:

public override async ValueTask<InterceptionResult> ConnectionOpeningAsync(DbConnection connection, ConnectionEventData eventData, InterceptionResult result, CancellationToken cancellationToken = default)
{
    var sqlConnection = (SqlConnection)connection;
    var accessToken = await _credential.GetTokenAsync(AzureSqlScope, cancellationToken: cancellationToken);

    sqlConnection.AccessToken = accessToken.Token;

    return ValueTask.FromResult(result);
}

New:

private readonly Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> _accessTokenCallback;

public AzureAdAuthenticationDbConnectionInterceptor()
{
    _accessTokenCallback = AccessTokenCallback;
}

public override ValueTask<InterceptionResult> ConnectionOpeningAsync(DbConnection connection, ConnectionEventData eventData, InterceptionResult result, CancellationToken cancellationToken = default)
{
    if (connection is SqlConnection sqlConnection && sqlConnection.AccessTokenCallback is null)
    {
        sqlConnection.AccessTokenCallback = _accessTokenCallback;
    }

    return ValueTask.FromResult(result);
}

private async Task<SqlAuthenticationToken> AccessTokenCallback(SqlAuthenticationParameters sqlAuthenticationParameters, CancellationToken cancellationToken)
{
    var accessToken = await _credential.GetTokenAsync(AzureSqlScope, cancellationToken: cancellationToken);

    return new SqlAuthenticationToken(accessToken.Token, accessToken.ExpiresOn);
}

[ajcvickers]

Note for team triage: looks like we may want to add some sugar here, especially for cases where the EF is creating the SqlConnection object.

[ajcvickers]

@kevinharing What you have looks fine, but given that setting the callback doesn't need to be done asynchronously, and if it only needs to be done when EF creates the SqlConnection (i.e. you never create one and pass it to EF, or, if you do, it already has the callback set), then using ConnectionCreated might be better. For example:

public class Interceptor(
    Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> accessTokenCallback)
    : DbConnectionInterceptor
{
    public override DbConnection ConnectionCreated(ConnectionCreatedEventData eventData, DbConnection result)
    {
        ((SqlConnection)result).AccessTokenCallback = accessTokenCallback;
        
        return result;
    }
}

Can you let me know what you end up with and how it is working out for you to help drive design for the sugar?

[kevinharing]

@kevinharing What you have looks fine, but given that setting the callback doesn't need to be done asynchronously, and if it only needs to be done when EF creates the SqlConnection (i.e. you never create one and pass it to EF, or, if you do, it already has the callback set), then using ConnectionCreated might be better. For example:

public class Interceptor(
    Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> accessTokenCallback)
    : DbConnectionInterceptor
{
    public override DbConnection ConnectionCreated(ConnectionCreatedEventData eventData, DbConnection result)
    {
        ((SqlConnection)result).AccessTokenCallback = accessTokenCallback;
        
        return result;
    }
}

Can you let me know what you end up with and how it is working out for you to help drive design for the sugar?

Thank you for your reply. That sounds indeed as a better option. Will test this out soon.

[chidhambaram55]

@kevinharing , Iam also trying to implement the AccessToken Callback in our system since we noticed token expiry errors when using sqlConnection.AccessToken = accessToken.Token;

Currently we are also using the DbConnectionInterceptor interceptor for refreshing the token, and updating sqlConnection.AccessToken in both ConnectionOpeningAsync and ConnectionOpening methods.

For using AccessTokenCallback, i had one query:

1. since we are using both ConnectionOpeningAsync and ConnectionOpening methods for refreshing the token, is it fine to assign async callback function inside a non async ConnectionOpening method?

public class AADAuthEADBInterceptor : DbConnectionInterceptor
{
    private readonly Func<SqlAuthenticationParameters, CancellationToken, Task<SqlAuthenticationToken>> _accessTokenCallback;

    public AADAuthEADBInterceptor()
    {
        _accessTokenCallback = AccessTokenCallback;
    }

    public override async ValueTask<InterceptionResult> ConnectionOpeningAsync(
        DbConnection connection,
        ConnectionEventData eventData,
        InterceptionResult result,
        CancellationToken cancellationToken)
    {
        var sqlConnection = (SqlConnection)connection;
         
        var connectionStringBuilder = new SqlConnectionStringBuilder(sqlConnection.ConnectionString);
        if (connectionStringBuilder.DataSource.Contains("database.windows.net", StringComparison.OrdinalIgnoreCase) && string.IsNullOrEmpty(connectionStringBuilder.UserID))
        {
 
            if (_options.CurrentValue.EnableAccessTokenCallbackFlow)
            {
                sqlConnection.AccessTokenCallback = _accessTokenCallback;
            }
            else
            {
                sqlConnection.AccessToken = await GetAzureSqlAccessTokenAsync();
            }
        }

        return await base.ConnectionOpeningAsync(connection, eventData, result, cancellationToken);
    }

    public override InterceptionResult ConnectionOpening(
        DbConnection connection,
        ConnectionEventData eventData,
        InterceptionResult result)
    {
        var sqlConnection = (SqlConnection)connection;
         
        var connectionStringBuilder = new SqlConnectionStringBuilder(sqlConnection.ConnectionString);

        if (connectionStringBuilder.DataSource.Contains("database.windows.net", StringComparison.OrdinalIgnoreCase) && string.IsNullOrEmpty(connectionStringBuilder.UserID))
        { 
            if (_options.CurrentValue.EnableAccessTokenCallbackFlow)
            {
                sqlConnection.AccessTokenCallback = _accessTokenCallback;
            }
            else
            {
                sqlConnection.AccessToken = GetAzureSqlAccessToken();
            }
        }

        return base.ConnectionOpening(connection, eventData, result);
    }
}

private async Task<SqlAuthenticationToken> AccessTokenCallback(SqlAuthenticationParameters sqlAuthenticationParameters, CancellationToken cancellationToken)
{
     
    var applicationBuilder = _confidentialClientAppBuilder.GetConfidentialClientApp();

    try
    {
        var scopes = new[] { "https://database.windows.net/.default" };
        AuthenticationResult authenticationResult = await applicationBuilder
            .AcquireTokenForClient(scopes)
            .ExecuteAsync()
            .ConfigureAwait(false);
       
        DateTimeOffset validFrom = DateTimeOffset.UtcNow.AddMinutes(30); 
        DateTimeOffset expiresOn = authenticationResult.ExpiresOn.AddMinutes(-5); 

        if (DateTimeOffset.Compare(validFrom, expiresOn) < 0)
        {
            expiresOn = validFrom;
        }

        return new SqlAuthenticationToken(authenticationResult.AccessToken, expiresOn);
    }
    catch (Exception ex)
    {
        _logger.LogError($"Enable to fetch AAD Azure token failed with following exception: {ex}");
        throw;
    }
    finally
    {
        sw.Stop();
        _logger.LogInformation($"Finished executing GetAzureSqlAccessTokenAsync ElapsedTime: {sw.ElapsedMilliseconds}");
    }
}

[cleferman]

@chidhambaram55 thanks a lot for your interceptor sample. I was trying to get the authentication to work but I was creating the SqlConnection object and passing it to the UseSqlServer callback. It seems EF doesn't really like that and its SqlConnection implementation is a bit different than the custom one I was creating.

The interceptor works now, I've used the newest auth method which is a tad less verbose than your version. Leaving it below in case anyone else needs it:

public class IdentityAuthDbInterceptor : DbConnectionInterceptor
{
    private readonly TokenCredential tokenCredential;

    public IdentityAuthDbInterceptor(bool isLocal)
    {
        tokenCredential = isLocal 
                ? new ChainedTokenCredential(new AzureCliCredential(), new VisualStudioCredential(), new VisualStudioCodeCredential())
                : new ChainedTokenCredential(new WorkloadIdentityCredential());
    }

    public override async ValueTask<InterceptionResult> ConnectionOpeningAsync(
        DbConnection connection,
        ConnectionEventData eventData,
        InterceptionResult result,
        CancellationToken cancellationToken)
    {
        var sqlConnection = (SqlConnection)connection;
        sqlConnection.AccessTokenCallback = AccessTokenCallback;

        return await base.ConnectionOpeningAsync(sqlConnection, eventData, result, cancellationToken);
    }

    public override InterceptionResult ConnectionOpening(
        DbConnection connection,
        ConnectionEventData eventData,
        InterceptionResult result)
    {
        var sqlConnection = (SqlConnection)connection;
        sqlConnection.AccessTokenCallback = AccessTokenCallback;

        return base.ConnectionOpening(sqlConnection, eventData, result);
    }

    private async Task<SqlAuthenticationToken> AccessTokenCallback(SqlAuthenticationParameters sqlAuthenticationParameters, CancellationToken cancellationToken)
    {
        var result = await tokenCredential.GetTokenAsync(new TokenRequestContext(["https://database.windows.net/.default"]), cancellationToken);

        return new SqlAuthenticationToken(result.Token, result.ExpiresOn);
    }
}