Cosmos DB: Add UserAgentSuffix (#763)

Adding user agent suffix configuration on host.json
Azure · Feb 8, 2022 · 003b07b · 003b07b
1 parent 51f31b3
commit 003b07b
Show file tree

Hide file tree

Showing 4 changed files with 137 additions and 3 deletions.
diff --git a/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs b/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs
@@ -105,6 +105,11 @@ internal Task<IValueBinder> BindForItemAsync(CosmosDBAttribute attribute, Type t
         internal CosmosClient GetService(string connection, string preferredLocations = "", string userAgent = "")
         {
             string cacheKey = BuildCacheKey(connection, preferredLocations);
+            if (!string.IsNullOrEmpty(_options.UserAgentSuffix))
+            {
+                userAgent += _options.UserAgentSuffix;
+            }
+
             CosmosClientOptions cosmosClientOptions = CosmosDBUtility.BuildClientOptions(_options.ConnectionMode, _cosmosSerializerFactory.CreateSerializer(), preferredLocations, userAgent);
             return ClientCache.GetOrAdd(cacheKey, (c) => _cosmosDBServiceFactory.CreateService(connection, cosmosClientOptions));
         }

diff --git a/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBOptions.cs b/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBOptions.cs
@@ -16,6 +16,11 @@ public class CosmosDBOptions : IOptionsFormatter
         /// <remarks>Default is Gateway mode.</remarks>
         public ConnectionMode? ConnectionMode { get; set; }
 
+        /// <summary>
+        /// Gets or sets a string to be included in the User Agent for all operations by Cosmos DB bindings and triggers.
+        /// </summary>
+        public string UserAgentSuffix { get; set; }
+
         public string Format()
         {
             StringWriter sw = new StringWriter();
@@ -26,6 +31,12 @@ public string Format()
                 writer.WritePropertyName(nameof(this.ConnectionMode));
                 writer.WriteValue(this.ConnectionMode);
 
+                if (!string.IsNullOrEmpty(UserAgentSuffix))
+                {
+                    writer.WritePropertyName(nameof(this.UserAgentSuffix));
+                    writer.WriteValue(this.UserAgentSuffix);
+                }
+
                 writer.WriteEndObject();
             }
 

diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBHostBuilderExtensionsTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBHostBuilderExtensionsTests.cs
@@ -21,6 +21,33 @@ namespace Microsoft.Azure.WebJobs.Extensions.CosmosDB.Tests
 {
     public class CosmosDBHostBuilderExtensionsTests
     {
+        [Fact]
+        public void ConfigurationBindsToDefaults()
+        {
+            IHost host = new HostBuilder()
+                 .ConfigureAppConfiguration(c =>
+                 {
+                     c.Sources.Clear();
+
+                     var source = new MemoryConfigurationSource
+                     {
+                         InitialData = new Dictionary<string, string>()
+                     };
+
+                     c.Add(source);
+                 })
+                 .ConfigureWebJobs(builder =>
+                 {
+                     builder.AddCosmosDB();
+                 })
+                .Build();
+
+            var options = host.Services.GetService<IOptions<CosmosDBOptions>>().Value;
+
+            Assert.Null(options.ConnectionMode);
+            Assert.Null(options.UserAgentSuffix);
+        }
+
         [Fact]
         public void ConfigurationBindsToOptions()
         {
@@ -33,7 +60,8 @@ public void ConfigurationBindsToOptions()
                      {
                          InitialData = new Dictionary<string, string>
                         {
-                            { "AzureWebJobs:extensions:cosmosDB:ConnectionMode", "Direct" }
+                            { "AzureWebJobs:extensions:cosmosDB:ConnectionMode", "Direct" },
+                            { "AzureWebJobs:extensions:cosmosDB:UserAgentSuffix", "randomtext" }
                         }
                      };
 
@@ -48,6 +76,7 @@ public void ConfigurationBindsToOptions()
             var options = host.Services.GetService<IOptions<CosmosDBOptions>>().Value;
 
             Assert.Equal(ConnectionMode.Direct, options.ConnectionMode);
+            Assert.Equal("randomtext", options.UserAgentSuffix);
         }
 
         [Fact]
@@ -61,13 +90,18 @@ public void ConfigurationBindsToOptions_WithConfigureServices()
                 .ConfigureServices(s =>
                 {
                     // Verifies that you can modify the bound options
-                    s.Configure<CosmosDBOptions>(o => o.ConnectionMode = ConnectionMode.Direct);
+                    s.Configure<CosmosDBOptions>(o =>
+                    {
+                        o.ConnectionMode = ConnectionMode.Direct;
+                        o.UserAgentSuffix = "randomtext";
+                    });
                 })
                 .Build();
 
             var options = host.Services.GetService<IOptions<CosmosDBOptions>>().Value;
 
             Assert.Equal(ConnectionMode.Direct, options.ConnectionMode);
+            Assert.Equal("randomtext", options.UserAgentSuffix);
         }
 
         [Fact]
@@ -113,7 +147,7 @@ public void ConfigurationBindsToOptions_WithSerializer_Before()
                     c.Sources.Clear();
                     c.AddInMemoryCollection(new Dictionary<string, string>
                     {
-                { Constants.DefaultConnectionStringName, "AccountEndpoint=https://defaultUri;AccountKey=c29tZV9rZXk=;" }
+                        { Constants.DefaultConnectionStringName, "AccountEndpoint=https://defaultUri;AccountKey=c29tZV9rZXk=;" }
                     });
                 })
                 .ConfigureServices(s =>
@@ -136,6 +170,61 @@ public void ConfigurationBindsToOptions_WithSerializer_Before()
             Assert.True(customFactory.CreateWasCalled);
         }
 
+        [Fact]
+        public void ConfigurationGetService_WithUserAgentOverride()
+        {
+            IHost host = new HostBuilder()
+                 .ConfigureAppConfiguration(c =>
+                 {
+                     c.Sources.Clear();
+                     c.AddInMemoryCollection(new Dictionary<string, string>
+                     {
+                         { Constants.DefaultConnectionStringName, "AccountEndpoint=https://defaultUri;AccountKey=c29tZV9rZXk=;" },
+                         { "AzureWebJobs:extensions:cosmosDB:UserAgentSuffix", "randomtext" }
+                     });
+                 })
+                 .ConfigureWebJobs(builder =>
+                 {
+                     builder.AddCosmosDB();
+                 })
+                .Build();
+
+            var extensionConfig = host.Services.GetServices<IExtensionConfigProvider>().Single();
+            Assert.NotNull(extensionConfig);
+            Assert.IsType<CosmosDBExtensionConfigProvider>(extensionConfig);
+
+            CosmosDBExtensionConfigProvider cosmosDBExtensionConfigProvider = (CosmosDBExtensionConfigProvider)extensionConfig;
+            CosmosClient dummyClient = cosmosDBExtensionConfigProvider.GetService(Constants.DefaultConnectionStringName, userAgent: "knownSuffix");
+            Assert.Equal(dummyClient.ClientOptions.ApplicationName, "knownSuffix" + "randomtext");
+        }
+
+        [Fact]
+        public void ConfigurationGetService_WithoutUserAgentOverride()
+        {
+            IHost host = new HostBuilder()
+                 .ConfigureAppConfiguration(c =>
+                 {
+                     c.Sources.Clear();
+                     c.AddInMemoryCollection(new Dictionary<string, string>
+                     {
+                         { Constants.DefaultConnectionStringName, "AccountEndpoint=https://defaultUri;AccountKey=c29tZV9rZXk=;" }
+                     });
+                 })
+                 .ConfigureWebJobs(builder =>
+                 {
+                     builder.AddCosmosDB();
+                 })
+                .Build();
+
+            var extensionConfig = host.Services.GetServices<IExtensionConfigProvider>().Single();
+            Assert.NotNull(extensionConfig);
+            Assert.IsType<CosmosDBExtensionConfigProvider>(extensionConfig);
+
+            CosmosDBExtensionConfigProvider cosmosDBExtensionConfigProvider = (CosmosDBExtensionConfigProvider)extensionConfig;
+            CosmosClient dummyClient = cosmosDBExtensionConfigProvider.GetService(Constants.DefaultConnectionStringName, userAgent: "knownSuffix");
+            Assert.Equal(dummyClient.ClientOptions.ApplicationName, "knownSuffix");
+        }
+
         private class CustomFactory : ICosmosDBSerializerFactory
         {
             public bool CreateWasCalled { get; private set; } = false;

diff --git a/...WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs b/...WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs
@@ -212,6 +212,35 @@ public async Task ValidParametersWithEnvironment_ConnectionMode_Succeed(Paramete
             Assert.Equal(new Uri("https://fromEnvironment"), binding.MonitoredContainer.Database.Client.Endpoint);
             Assert.Equal(new Uri("https://fromEnvironment"), binding.LeaseContainer.Database.Client.Endpoint);
             Assert.Equal(ConnectionMode.Direct, binding.LeaseContainer.Database.Client.ClientOptions.ConnectionMode);
+            Assert.Equal(ConnectionMode.Direct, binding.MonitoredContainer.Database.Client.ClientOptions.ConnectionMode);
+            Assert.Equal("CosmosDBTriggerFunctions", binding.LeaseContainer.Database.Client.ClientOptions.ApplicationName);
+            Assert.Equal("CosmosDBTriggerFunctions", binding.MonitoredContainer.Database.Client.ClientOptions.ApplicationName);
+        }
+
+        [Theory]
+        [MemberData(nameof(ValidCosmosDBTriggerBindingsWithEnvironmentParameters))]
+        public async Task ValidParametersWithEnvironment_UserAgent_Succeed(ParameterInfo parameter)
+        {
+            var nameResolver = new TestNameResolver();
+            var config = new ConfigurationBuilder()
+                .AddInMemoryCollection(new Dictionary<string, string>
+                {
+                    { "ConnectionStrings:CosmosDB", "AccountEndpoint=https://fromEnvironment;AccountKey=c29tZV9rZXk=;" },
+                    { "ConnectionStrings:CosmosDBConnectionString", "AccountEndpoint=https://fromSettings;AccountKey=c29tZV9rZXk=;" }
+                })
+                .Build();
+
+            _options.UserAgentSuffix = "randomtext";
+
+            CosmosDBTriggerAttributeBindingProvider<dynamic> provider = new CosmosDBTriggerAttributeBindingProvider<dynamic>(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory);
+
+            CosmosDBTriggerBinding<dynamic> binding = (CosmosDBTriggerBinding<dynamic>)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None));
+
+            Assert.Equal(typeof(IReadOnlyCollection<dynamic>), binding.TriggerValueType);
+            Assert.Equal(new Uri("https://fromEnvironment"), binding.MonitoredContainer.Database.Client.Endpoint);
+            Assert.Equal(new Uri("https://fromEnvironment"), binding.LeaseContainer.Database.Client.Endpoint);
+            Assert.Equal("CosmosDBTriggerFunctionsrandomtext", binding.LeaseContainer.Database.Client.ClientOptions.ApplicationName);
+            Assert.Equal("CosmosDBTriggerFunctionsrandomtext", binding.MonitoredContainer.Database.Client.ClientOptions.ApplicationName);
         }
 
         [Theory]