Cosmos DB - Add Support for AAD authentication

src/WebJobs.Extensions.CosmosDB/Bindings/CosmosDBClientBuilder.cs

@@ -22,9 +22,8 @@ public CosmosClient Convert(CosmosDBAttribute attribute)
                 throw new ArgumentNullException(nameof(attribute));
             }
 
-            string resolvedConnectionString = _configProvider.ResolveConnectionString(attribute.Connection);
             return _configProvider.GetService(
-                connectionString: resolvedConnectionString, 
+                connection: attribute.Connection ?? Constants.DefaultConnectionStringName, 
                 preferredLocations: attribute.PreferredLocations);
         }
     }

src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs

@@ -23,7 +23,6 @@ namespace Microsoft.Azure.WebJobs.Extensions.CosmosDB
     [Extension("CosmosDB")]
     internal class CosmosDBExtensionConfigProvider : IExtensionConfigProvider
     {
-        private readonly IConfiguration _configuration;
         private readonly ICosmosDBServiceFactory _cosmosDBServiceFactory;
         private readonly ICosmosDBSerializerFactory _cosmosSerializerFactory;
         private readonly INameResolver _nameResolver;
@@ -34,11 +33,9 @@ public CosmosDBExtensionConfigProvider(
             IOptions<CosmosDBOptions> options, 
             ICosmosDBServiceFactory cosmosDBServiceFactory, 
             ICosmosDBSerializerFactory cosmosSerializerFactory,
-            IConfiguration configuration, 
             INameResolver nameResolver, 
             ILoggerFactory loggerFactory)
         {
-            _configuration = configuration;
             _cosmosDBServiceFactory = cosmosDBServiceFactory;
             _cosmosSerializerFactory = cosmosSerializerFactory;
             _nameResolver = nameResolver;
@@ -77,29 +74,19 @@ public void Initialize(ExtensionConfigContext context)
 
             // Trigger
             var rule2 = context.AddBindingRule<CosmosDBTriggerAttribute>();
-            rule2.BindToTrigger(new CosmosDBTriggerAttributeBindingProviderGenerator(_configuration, _nameResolver, _options, this, _loggerFactory));
+            rule2.BindToTrigger(new CosmosDBTriggerAttributeBindingProviderGenerator(_nameResolver, _options, this, _loggerFactory));
         }
 
         internal void ValidateConnection(CosmosDBAttribute attribute, Type paramType)
         {
-            if (string.IsNullOrEmpty(_options.ConnectionString) &&
-                string.IsNullOrEmpty(attribute.Connection))
+            if (attribute.Connection == string.Empty)
             {
                 string attributeProperty = $"{nameof(CosmosDBAttribute)}.{nameof(CosmosDBAttribute.Connection)}";
-                string optionsProperty = $"{nameof(CosmosDBOptions)}.{nameof(CosmosDBOptions.ConnectionString)}";
                 throw new InvalidOperationException(
-                    $"The CosmosDB connection string must be set either via the '{Constants.DefaultConnectionStringName}' IConfiguration connection string, via the {attributeProperty} property or via {optionsProperty}.");
+                    $"The {attributeProperty} property cannot be an empty value.");
             }
         }
 
-        internal CosmosClient BindForClient(CosmosDBAttribute attribute)
-        {
-            string resolvedConnectionString = ResolveConnectionString(attribute.Connection);
-            return GetService(
-                connectionString: resolvedConnectionString, 
-                preferredLocations: attribute.PreferredLocations);
-        }
-
         internal Task<IValueBinder> BindForItemAsync(CosmosDBAttribute attribute, Type type)
         {
             if (string.IsNullOrEmpty(attribute.Id))
@@ -115,31 +102,17 @@ internal Task<IValueBinder> BindForItemAsync(CosmosDBAttribute attribute, Type t
             return Task.FromResult(binder);
         }
 
-        internal string ResolveConnectionString(string attributeConnectionString)
+        internal CosmosClient GetService(string connection, string preferredLocations = "", string userAgent = "")
         {
-            // First, try the Attribute's string.
-            if (!string.IsNullOrEmpty(attributeConnectionString))
-            {
-                return attributeConnectionString;
-            }
-
-            // Then use the options.
-            return _options.ConnectionString;
-        }
-
-        internal CosmosClient GetService(string connectionString, string preferredLocations = "", string userAgent = "")
-        {
-            string cacheKey = BuildCacheKey(connectionString, preferredLocations);
+            string cacheKey = BuildCacheKey(connection, preferredLocations);
             CosmosClientOptions cosmosClientOptions = CosmosDBUtility.BuildClientOptions(_options.ConnectionMode, _cosmosSerializerFactory.CreateSerializer(), preferredLocations, userAgent);
-            return ClientCache.GetOrAdd(cacheKey, (c) => _cosmosDBServiceFactory.CreateService(connectionString, cosmosClientOptions));
+            return ClientCache.GetOrAdd(cacheKey, (c) => _cosmosDBServiceFactory.CreateService(connection, cosmosClientOptions));
         }
 
         internal CosmosDBContext CreateContext(CosmosDBAttribute attribute)
         {
-            string resolvedConnectionString = ResolveConnectionString(attribute.Connection);
-
             CosmosClient service = GetService(
-                connectionString: resolvedConnectionString, 
+                connection: attribute.Connection ?? Constants.DefaultConnectionStringName, 
                 preferredLocations: attribute.PreferredLocations);
 
             return new CosmosDBContext

src/WebJobs.Extensions.CosmosDB/Config/CosmosDBOptions.cs

@@ -10,11 +10,6 @@ namespace Microsoft.Azure.WebJobs.Extensions.CosmosDB
 {
     public class CosmosDBOptions : IOptionsFormatter
     {
-        /// <summary>
-        /// Gets or sets the CosmosDB connection string.
-        /// </summary>
-        public string ConnectionString { get; set; }
-
         /// <summary>
         /// Gets or sets the ConnectionMode used in the CosmosClient instances.
         /// </summary>

src/WebJobs.Extensions.CosmosDB/Config/CosmosDBWebJobsBuilderExtensions.cs

@@ -28,8 +28,6 @@ public static IWebJobsBuilder AddCosmosDB(this IWebJobsBuilder builder)
             builder.AddExtension<CosmosDBExtensionConfigProvider>()               
                 .ConfigureOptions<CosmosDBOptions>((config, path, options) =>
                 {
-                    options.ConnectionString = config.GetConnectionStringOrSetting(Constants.DefaultConnectionStringName);
-
                     IConfigurationSection section = config.GetSection(path);
                     section.Bind(options);
                 });                

src/WebJobs.Extensions.CosmosDB/Config/DefaultCosmosDBServiceFactory.cs

@@ -1,15 +1,92 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System;
+using Azure.Core;
 using Microsoft.Azure.Cosmos;
+using Microsoft.Extensions.Azure;
+using Microsoft.Extensions.Configuration;
 
 namespace Microsoft.Azure.WebJobs.Extensions.CosmosDB
 {
     internal class DefaultCosmosDBServiceFactory : ICosmosDBServiceFactory
     {
-        public CosmosClient CreateService(string connectionString, CosmosClientOptions cosmosClientOptions)
+        private readonly IConfiguration _configuration;
+        private readonly AzureComponentFactory _componentFactory;
+
+        public DefaultCosmosDBServiceFactory(
+            IConfiguration configuration,
+            AzureComponentFactory componentFactory)
+        {
+            this._configuration = configuration;
+            this._componentFactory = componentFactory;
+        }
+
+        public CosmosClient CreateService(string connectionName, CosmosClientOptions cosmosClientOptions)
+        {
+            CosmosConnectionInformation cosmosConnectionInformation = this.ResolveConnectionInformation(connectionName);
+            if (cosmosConnectionInformation.UsesConnectionString)
+            {
+                // Connection string based auth
+                return new CosmosClient(cosmosConnectionInformation.ConnectionString, cosmosClientOptions);
+            }
+
+            // AAD auth
+            return new CosmosClient(cosmosConnectionInformation.AccountEndpoint, cosmosConnectionInformation.Credential, cosmosClientOptions);
+        }
+
+        private CosmosConnectionInformation ResolveConnectionInformation(string connection)
+        {
+            var connectionSetting = connection ?? Constants.DefaultConnectionStringName;
+            IConfigurationSection connectionSection = WebJobsConfigurationExtensions.GetWebJobsConnectionStringSection(this._configuration, connectionSetting);
+            if (!connectionSection.Exists())
+            {
+                // Not found
+                throw new InvalidOperationException($"Cosmos DB connection configuration '{connectionSetting}' does not exist. " +
+                                                    $"Make sure that it is a defined App Setting.");
+            }
+
+            if (!string.IsNullOrWhiteSpace(connectionSection.Value))
+            {
+                return new CosmosConnectionInformation(connectionSection.Value);
+            }
+            else
+            {
+                string accountEndpoint = connectionSection["accountEndpoint"];
+                if (string.IsNullOrWhiteSpace(accountEndpoint))
+                {
+                    // Not found
+                    throw new InvalidOperationException($"Connection should have an 'accountEndpoint' property or be a " +
+                        $"string representing a connection string.");
+                }
+
+                TokenCredential credential = _componentFactory.CreateTokenCredential(connectionSection);
+                return new CosmosConnectionInformation(accountEndpoint, credential);
+            }
+        }
+
+        private class CosmosConnectionInformation
         {
-            return new CosmosClient(connectionString, cosmosClientOptions);
+            public CosmosConnectionInformation(string connectionString)
+            {
+                this.ConnectionString = connectionString;
+                this.UsesConnectionString = true;
+            }
+
+            public CosmosConnectionInformation(string accountEndpoint, TokenCredential tokenCredential)
+            {
+                this.AccountEndpoint = accountEndpoint;
+                this.Credential = tokenCredential;
+                this.UsesConnectionString = false;
+            }
+
+            public bool UsesConnectionString { get; }
+
+            public string ConnectionString { get; }
+
+            public string AccountEndpoint { get; }
+
+            public TokenCredential Credential { get; }
         }
     }
 }

src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs

@@ -8,7 +8,6 @@
 using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Host.Triggers;
 using Microsoft.Azure.WebJobs.Logging;
-using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.Azure.WebJobs.Extensions.CosmosDB
@@ -19,16 +18,14 @@ internal class CosmosDBTriggerAttributeBindingProvider<T>
         private const string SharedThroughputRequirementException = "Shared throughput collection should have a partition key";
         private const string LeaseCollectionRequiredPartitionKey = "/id";
         private const string LeaseCollectionRequiredPartitionKeyFromGremlin = "/partitionKey";
-        private readonly IConfiguration _configuration;
         private readonly INameResolver _nameResolver;
         private readonly CosmosDBOptions _options;
         private readonly ILogger _logger;
         private readonly CosmosDBExtensionConfigProvider _configProvider;
 
-        public CosmosDBTriggerAttributeBindingProvider(IConfiguration configuration, INameResolver nameResolver, CosmosDBOptions options,
+        public CosmosDBTriggerAttributeBindingProvider(INameResolver nameResolver, CosmosDBOptions options,
             CosmosDBExtensionConfigProvider configProvider, ILoggerFactory loggerFactory)
         {
-            _configuration = configuration;
             _nameResolver = nameResolver;
             _options = options;
             _configProvider = configProvider;
@@ -61,16 +58,16 @@ public async Task<ITriggerBinding> TryCreateAsync(TriggerBindingProviderContext
 
             try
             {
-                string triggerConnectionString = ResolveAttributeConnectionString(attribute);
-                if (string.IsNullOrEmpty(triggerConnectionString))
+                string triggerConnection = ResolveAttributeConnection(attribute);
+                if (string.IsNullOrEmpty(triggerConnection))
                 {
-                    throw new InvalidOperationException("The connection string for the monitored container is in an invalid format, please use AccountEndpoint=XXXXXX;AccountKey=XXXXXX;.");
+                    throw new InvalidOperationException($"The attribute {nameof(attribute.Connection)} for the monitored container is in an invalid format, please use AccountEndpoint=XXXXXX;AccountKey=XXXXXX; or a node representing token authentication information.");
                 }
 
-                string leasesConnectionString = ResolveAttributeLeasesConnectionString(attribute);
-                if (string.IsNullOrEmpty(leasesConnectionString))
+                string leasesConnection = ResolveAttributeLeasesConnection(attribute);
+                if (string.IsNullOrEmpty(leasesConnection))
                 {
-                    throw new InvalidOperationException("The connection string for the leases container is in an invalid format, please use AccountEndpoint=XXXXXX;AccountKey=XXXXXX;.");
+                    throw new InvalidOperationException($"The attribute {nameof(attribute.LeaseConnection)} for the leases container is in an invalid format, please use AccountEndpoint=XXXXXX;AccountKey=XXXXXX;. or a node representing token authentication information.");
                 }
 
                 if (string.IsNullOrEmpty(monitoredDatabaseName)
@@ -81,19 +78,19 @@ public async Task<ITriggerBinding> TryCreateAsync(TriggerBindingProviderContext
                     throw new InvalidOperationException("Cannot establish database and container values. If you are using environment and configuration values, please ensure these are correctly set.");
                 }
 
-                if (triggerConnectionString.Equals(leasesConnectionString, StringComparison.InvariantCultureIgnoreCase)
+                if (triggerConnection.Equals(leasesConnection, StringComparison.InvariantCultureIgnoreCase)
                     && monitoredDatabaseName.Equals(leasesDatabaseName, StringComparison.InvariantCultureIgnoreCase)
                     && monitoredCollectionName.Equals(leasesCollectionName, StringComparison.InvariantCultureIgnoreCase))
                 {
                     throw new InvalidOperationException("The monitored container cannot be the same as the container storing the leases.");
                 }
 
                 CosmosClient monitoredCosmosDBService = _configProvider.GetService(
-                    connectionString: triggerConnectionString, 
+                    connection: triggerConnection, 
                     preferredLocations: preferredLocations, 
                     userAgent: CosmosDBTriggerUserAgentSuffix);
                 CosmosClient leaseCosmosDBService = _configProvider.GetService(
-                    connectionString: leasesConnectionString, 
+                    connection: leasesConnection, 
                     preferredLocations: preferredLocations, 
                     userAgent: CosmosDBTriggerUserAgentSuffix);
 
@@ -148,9 +145,9 @@ private static async Task CreateLeaseCollectionIfNotExistsAsync(CosmosClient cos
             }
         }
 
-        private string ResolveAttributeConnectionString(CosmosDBTriggerAttribute attribute)
+        private string ResolveAttributeConnection(CosmosDBTriggerAttribute attribute)
         {
-            string connectionString = ResolveConnectionString(attribute.Connection, nameof(CosmosDBTriggerAttribute.Connection));
+            string connectionString = attribute.Connection ?? Constants.DefaultConnectionStringName;
 
             if (string.IsNullOrEmpty(connectionString))
             {
@@ -160,7 +157,7 @@ private string ResolveAttributeConnectionString(CosmosDBTriggerAttribute attribu
             return connectionString;
         }
 
-        private string ResolveAttributeLeasesConnectionString(CosmosDBTriggerAttribute attribute)
+        private string ResolveAttributeLeasesConnection(CosmosDBTriggerAttribute attribute)
         {
             // If the lease connection string is not set, use the trigger's
             string keyToResolve = attribute.LeaseConnection;
@@ -169,7 +166,7 @@ private string ResolveAttributeLeasesConnectionString(CosmosDBTriggerAttribute a
                 keyToResolve = attribute.Connection;
             }
 
-            string connectionString = ResolveConnectionString(keyToResolve, nameof(CosmosDBTriggerAttribute.LeaseConnection));
+            string connectionString = keyToResolve ?? Constants.DefaultConnectionStringName;
 
             if (string.IsNullOrEmpty(connectionString))
             {
@@ -185,31 +182,10 @@ private void ThrowMissingConnectionStringException(bool isLeaseConnectionString
                 $"{nameof(CosmosDBTriggerAttribute)}.{nameof(CosmosDBTriggerAttribute.LeaseConnection)}" :
                 $"{nameof(CosmosDBTriggerAttribute)}.{nameof(CosmosDBTriggerAttribute.Connection)}";
 
-            string optionsProperty = $"{nameof(CosmosDBOptions)}.{nameof(CosmosDBOptions.ConnectionString)}";
-
             string leaseString = isLeaseConnectionString ? "lease " : string.Empty;
 
             throw new InvalidOperationException(
-                $"The CosmosDBTrigger {leaseString}connection string must be set either via a '{Constants.DefaultConnectionStringName}' configuration connection string, via the {attributeProperty} property or via {optionsProperty}.");
-        }
-
-        internal string ResolveConnectionString(string unresolvedConnectionString, string propertyName)
-        {
-            // First, resolve the string.
-            if (!string.IsNullOrEmpty(unresolvedConnectionString))
-            {
-                string resolvedString = _configuration.GetConnectionStringOrSetting(unresolvedConnectionString);
-
-                if (string.IsNullOrEmpty(resolvedString))
-                {
-                    throw new InvalidOperationException($"Unable to resolve app setting for property '{nameof(CosmosDBTriggerAttribute)}.{propertyName}'. Make sure the app setting exists and has a valid value.");
-                }
-
-                return resolvedString;
-            }
-
-            // If that didn't exist, fall back to options.
-            return _options.ConnectionString;
+                $"The CosmosDBTrigger {leaseString}connection must be set either via a '{Constants.DefaultConnectionStringName}' configuration or via the {attributeProperty} property.");
         }
 
         private string ResolveAttributeValue(string attributeValue)