Cosmos: Support AAD RBAC via TokenCredential

.github/workflows/TestCosmos.yaml

@@ -33,3 +33,10 @@ jobs:
       - name: Test on Cosmos
         run: test.cmd /p:Projects=${{ github.workspace }}\test\EFCore.Cosmos.FunctionalTests\EFCore.Cosmos.FunctionalTests.csproj
         shell: cmd
+
+      - name: Publish Test Results
+        uses: actions/upload-artifact@v3
+        if: always()
+        with:
+          name: test-results
+          path: artifacts/TestResults/Debug/*

src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs

@@ -26,8 +26,9 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade)
         => GetService<ISingletonCosmosClientWrapper>(databaseFacade).Client;
 
     private static TService GetService<TService>(IInfrastructure<IServiceProvider> databaseFacade)
+        where TService : class
     {
-        var service = databaseFacade.Instance.GetService<TService>();
+        var service = databaseFacade.GetService<TService>();
         if (service == null)
         {
             throw new InvalidOperationException(CosmosStrings.CosmosNotInUse);
@@ -36,6 +37,25 @@ private static TService GetService<TService>(IInfrastructure<IServiceProvider> d
         return service;
     }
 
+    /// <summary>
+    ///     Gets the configured database name for this <see cref="DbContext" />.
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <returns>The database name.</returns>
+    public static string GetCosmosDatabaseId(this DatabaseFacade databaseFacade)
+    {
+        var cosmosOptions = databaseFacade.GetService<IDbContextOptions>().FindExtension<CosmosOptionsExtension>();
+        if (cosmosOptions == null)
+        {
+            throw new InvalidOperationException(CosmosStrings.CosmosNotInUse);
+        }
+
+        return cosmosOptions.DatabaseName;
+    }
+
     /// <summary>
     ///     Returns <see langword="true" /> if the database provider currently in use is the Cosmos provider.
     /// </summary>

src/EFCore.Cosmos/Extensions/CosmosDbContextOptionsExtensions.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Azure.Core;
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 
 // ReSharper disable once CheckNamespace
@@ -83,6 +84,74 @@ public static DbContextOptionsBuilder UseCosmos(
         return optionsBuilder;
     }
 
+    /// <summary>
+    ///     Configures the context to connect to an Azure Cosmos database.
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
+    ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <typeparam name="TContext">The type of context to be configured.</typeparam>
+    /// <param name="optionsBuilder">The builder being used to configure the context.</param>
+    /// <param name="accountEndpoint">The account end-point to connect to.</param>
+    /// <param name="tokenCredential">The Azure authentication token.</param>
+    /// <param name="databaseName">The database name.</param>
+    /// <param name="cosmosOptionsAction">An optional action to allow additional Cosmos-specific configuration.</param>
+    /// <returns>The options builder so that further configuration can be chained.</returns>
+    public static DbContextOptionsBuilder<TContext> UseCosmos<TContext>(
+        this DbContextOptionsBuilder<TContext> optionsBuilder,
+        string accountEndpoint,
+        TokenCredential tokenCredential,
+        string databaseName,
+        Action<CosmosDbContextOptionsBuilder>? cosmosOptionsAction = null)
+        where TContext : DbContext
+        => (DbContextOptionsBuilder<TContext>)UseCosmos(
+            (DbContextOptionsBuilder)optionsBuilder,
+            accountEndpoint,
+            tokenCredential,
+            databaseName,
+            cosmosOptionsAction);
+
+    /// <summary>
+    ///     Configures the context to connect to an Azure Cosmos database.
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
+    ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <param name="optionsBuilder">The builder being used to configure the context.</param>
+    /// <param name="accountEndpoint">The account end-point to connect to.</param>
+    /// <param name="tokenCredential">The Azure authentication token.</param>
+    /// <param name="databaseName">The database name.</param>
+    /// <param name="cosmosOptionsAction">An optional action to allow additional Cosmos-specific configuration.</param>
+    /// <returns>The options builder so that further configuration can be chained.</returns>
+    public static DbContextOptionsBuilder UseCosmos(
+        this DbContextOptionsBuilder optionsBuilder,
+        string accountEndpoint,
+        TokenCredential tokenCredential,
+        string databaseName,
+        Action<CosmosDbContextOptionsBuilder>? cosmosOptionsAction = null)
+    {
+        Check.NotNull(optionsBuilder, nameof(optionsBuilder));
+        Check.NotNull(accountEndpoint, nameof(accountEndpoint));
+        Check.NotNull(tokenCredential, nameof(tokenCredential));
+        Check.NotEmpty(databaseName, nameof(databaseName));
+
+        var extension = optionsBuilder.Options.FindExtension<CosmosOptionsExtension>()
+            ?? new CosmosOptionsExtension();
+
+        extension = extension
+            .WithAccountEndpoint(accountEndpoint)
+            .WithTokenCredential(tokenCredential)
+            .WithDatabaseName(databaseName);
+
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
+
+        cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));
+
+        return optionsBuilder;
+    }
+
     /// <summary>
     ///     Configures the context to connect to an Azure Cosmos database.
     /// </summary>

src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs

@@ -4,6 +4,7 @@
 using System.Globalization;
 using System.Net;
 using System.Text;
+using Azure.Core;
 using Microsoft.EntityFrameworkCore.Cosmos.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
@@ -18,6 +19,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension
 {
     private string? _accountEndpoint;
     private string? _accountKey;
+    private TokenCredential? _tokenCredential;
     private string? _connectionString;
     private string? _databaseName;
     private string? _region;
@@ -55,6 +57,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom)
     {
         _accountEndpoint = copyFrom._accountEndpoint;
         _accountKey = copyFrom._accountKey;
+        _tokenCredential = copyFrom._tokenCredential;
         _databaseName = copyFrom._databaseName;
         _connectionString = copyFrom._connectionString;
         _region = copyFrom._region;
@@ -138,6 +141,35 @@ public virtual CosmosOptionsExtension WithAccountKey(string? accountKey)
         return clone;
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual TokenCredential? TokenCredential
+        => _tokenCredential;
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual CosmosOptionsExtension WithTokenCredential(TokenCredential? tokenCredential)
+    {
+        if (tokenCredential is not null && _connectionString is not null)
+        {
+            throw new InvalidOperationException(CosmosStrings.ConnectionStringConflictingConfiguration);
+        }
+
+        var clone = Clone();
+
+        clone._tokenCredential = tokenCredential;
+
+        return clone;
+    }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -155,7 +187,7 @@ public virtual string? ConnectionString
     /// </summary>
     public virtual CosmosOptionsExtension WithConnectionString(string? connectionString)
     {
-        if (connectionString is not null && (_accountEndpoint != null || _accountKey != null))
+        if (connectionString is not null && (_accountEndpoint != null || _accountKey != null || _tokenCredential != null))
         {
             throw new InvalidOperationException(CosmosStrings.ConnectionStringConflictingConfiguration);
         }
@@ -565,6 +597,7 @@ public override int GetServiceProviderHashCode()
                 {
                     hashCode.Add(Extension._accountEndpoint);
                     hashCode.Add(Extension._accountKey);
+                    hashCode.Add(Extension._tokenCredential);
                 }
 
                 hashCode.Add(Extension._region);
@@ -591,6 +624,7 @@ public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo
                 && Extension._connectionString == otherInfo.Extension._connectionString
                 && Extension._accountEndpoint == otherInfo.Extension._accountEndpoint
                 && Extension._accountKey == otherInfo.Extension._accountKey
+                && Extension._tokenCredential == otherInfo.Extension._tokenCredential
                 && Extension._region == otherInfo.Extension._region
                 && Extension._connectionMode == otherInfo.Extension._connectionMode
                 && Extension._limitToEndpoint == otherInfo.Extension._limitToEndpoint
@@ -615,8 +649,17 @@ public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
             {
                 debugInfo["Cosmos:" + nameof(AccountEndpoint)] =
                     (Extension._accountEndpoint?.GetHashCode() ?? 0L).ToString(CultureInfo.InvariantCulture);
-                debugInfo["Cosmos:" + nameof(AccountKey)] =
-                    (Extension._accountKey?.GetHashCode() ?? 0L).ToString(CultureInfo.InvariantCulture);
+
+                if (Extension._accountKey == null)
+                {
+                    debugInfo["Cosmos:" + nameof(TokenCredential)] =
+                        (Extension._tokenCredential?.GetHashCode() ?? 0L).ToString(CultureInfo.InvariantCulture);
+                }
+                else
+                {
+                    debugInfo["Cosmos:" + nameof(AccountKey)] =
+                        (Extension._accountKey?.GetHashCode() ?? 0L).ToString(CultureInfo.InvariantCulture);
+                }
             }
 
             debugInfo["Cosmos:" + nameof(CosmosDbContextOptionsBuilder.Region)] =

src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net;
+using Azure.Core;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 
@@ -29,6 +30,14 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions
     /// </summary>
     public virtual string? AccountKey { get; private set; }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual TokenCredential? TokenCredential { get; private set; }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -146,6 +155,7 @@ public virtual void Initialize(IDbContextOptions options)
         {
             AccountEndpoint = cosmosOptions.AccountEndpoint;
             AccountKey = cosmosOptions.AccountKey;
+            TokenCredential = cosmosOptions.TokenCredential;
             ConnectionString = cosmosOptions.ConnectionString;
             Region = cosmosOptions.Region;
             LimitToEndpoint = cosmosOptions.LimitToEndpoint;
@@ -175,6 +185,7 @@ public virtual void Validate(IDbContextOptions options)
         if (cosmosOptions != null
             && (AccountEndpoint != cosmosOptions.AccountEndpoint
                 || AccountKey != cosmosOptions.AccountKey
+                || TokenCredential != cosmosOptions.TokenCredential
                 || ConnectionString != cosmosOptions.ConnectionString
                 || Region != cosmosOptions.Region
                 || LimitToEndpoint != cosmosOptions.LimitToEndpoint

src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Net;
+using Azure.Core;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 
@@ -35,6 +36,14 @@ public interface ICosmosSingletonOptions : ISingletonOptions
     /// </summary>
     string? AccountKey { get; }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    TokenCredential? TokenCredential { get; }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore.Cosmos/Metadata/Conventions/Internal/CosmosConventionSetBuilder.cs

@@ -36,7 +36,7 @@ public override ConventionSet CreateConventionSet()
         conventionSet.Add(new ContextContainerConvention(Dependencies));
         conventionSet.Add(new ETagPropertyConvention());
         conventionSet.Add(new StoreKeyConvention(Dependencies));
-        
+
         conventionSet.Replace<ValueGenerationConvention>(new CosmosValueGenerationConvention(Dependencies));
         conventionSet.Replace<KeyDiscoveryConvention>(new CosmosKeyDiscoveryConvention(Dependencies));
         conventionSet.Replace<InversePropertyAttributeConvention>(new CosmosInversePropertyAttributeConvention(Dependencies));

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -124,7 +124,7 @@
     <value>The Cosmos database does not support 'CanConnect' or 'CanConnectAsync'.</value>
   </data>
   <data name="ConnectionStringConflictingConfiguration" xml:space="preserve">
-    <value>Both the connection string and account key or account endpoint were specified. Specify only one set of connection details.</value>
+    <value>Both the connection string and CredentialToken, account key or account endpoint were specified. Specify only one set of connection details.</value>
   </data>
   <data name="CosmosNotInUse" xml:space="preserve">
     <value>Cosmos-specific methods can only be used when the context is using the Cosmos provider.</value>

src/EFCore.Cosmos/Storage/Internal/SingletonCosmosClientWrapper.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Azure.Core;
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
@@ -18,6 +19,7 @@ public class SingletonCosmosClientWrapper : ISingletonCosmosClientWrapper
     private readonly string? _endpoint;
     private readonly string? _key;
     private readonly string? _connectionString;
+    private readonly TokenCredential? _tokenCredential;
     private CosmosClient? _client;
 
     /// <summary>
@@ -31,6 +33,7 @@ public SingletonCosmosClientWrapper(ICosmosSingletonOptions options)
         _endpoint = options.AccountEndpoint;
         _key = options.AccountKey;
         _connectionString = options.ConnectionString;
+        _tokenCredential = options.TokenCredential;
         var configuration = new CosmosClientOptions { ApplicationName = UserAgent, Serializer = new JsonCosmosSerializer() };
 
         if (options.Region != null)
@@ -99,7 +102,9 @@ public SingletonCosmosClientWrapper(ICosmosSingletonOptions options)
     /// </summary>
     public virtual CosmosClient Client
         => _client ??= string.IsNullOrEmpty(_connectionString)
-            ? new CosmosClient(_endpoint, _key, _options)
+            ? _tokenCredential == null
+                ? new CosmosClient(_endpoint, _key, _options)
+                : new CosmosClient(_endpoint, _tokenCredential, _options)
             : new CosmosClient(_connectionString, _options);
 
     /// <summary>

src/EFCore.Relational/Metadata/Conventions/Infrastructure/RelationalConventionSetBuilder.cs

@@ -80,9 +80,9 @@ public override ConventionSet CreateConventionSet()
         conventionSet.Add(new RelationalMapToJsonConvention(Dependencies, RelationalDependencies));
 
         conventionSet.Replace<ValueGenerationConvention>(
-            new RelationalValueGenerationConvention(Dependencies, RelationalDependencies));        
+            new RelationalValueGenerationConvention(Dependencies, RelationalDependencies));
         conventionSet.Replace<QueryFilterRewritingConvention>(
-            new RelationalQueryFilterRewritingConvention(Dependencies, RelationalDependencies));        
+            new RelationalQueryFilterRewritingConvention(Dependencies, RelationalDependencies));
         conventionSet.Replace<RuntimeModelConvention>(new RelationalRuntimeModelConvention(Dependencies, RelationalDependencies));
 
         return conventionSet;

src/EFCore.SqlServer/Metadata/Conventions/SqlServerConventionSetBuilder.cs

@@ -61,15 +61,15 @@ public override ConventionSet CreateConventionSet()
         conventionSet.Replace<ValueGenerationConvention>(
             new SqlServerValueGenerationConvention(Dependencies, RelationalDependencies));
         conventionSet.Replace<RuntimeModelConvention>(new SqlServerRuntimeModelConvention(Dependencies, RelationalDependencies));
-        
+
         var sqlServerTemporalConvention = new SqlServerTemporalConvention(Dependencies, RelationalDependencies);
         ConventionSet.AddBefore(
             conventionSet.EntityTypeAnnotationChangedConventions,
             sqlServerTemporalConvention,
             typeof(SqlServerValueGenerationConvention));
         conventionSet.SkipNavigationForeignKeyChangedConventions.Add(sqlServerTemporalConvention);
         conventionSet.ModelFinalizingConventions.Add(sqlServerTemporalConvention);
-        
+
         return conventionSet;
     }