Support for Default Azure Identity

Signed-off-by: Piotr Karpala <[email protected]>
kedacore · Jun 11, 2024 · 24eae81 · 24eae81
1 parent a00c272
commit 24eae81
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 9 deletions.
diff --git a/src/Scaler.Demo/OrderGenerator/Keda.CosmosDb.Scaler.Demo.OrderGenerator.csproj b/src/Scaler.Demo/OrderGenerator/Keda.CosmosDb.Scaler.Demo.OrderGenerator.csproj
@@ -6,6 +6,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="Azure.Identity" Version="1.11.4" />
     <PackageReference Include="Bogus" Version="34.0.2" />
     <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.40.0" />
     <PackageReference Include="Microsoft.Extensions.Hosting" Version="6.0.1" />

diff --git a/src/Scaler.Demo/OrderGenerator/Program.cs b/src/Scaler.Demo/OrderGenerator/Program.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Linq;
 using System.Threading.Tasks;
+using Azure.Identity;
 using Bogus;
 using Bogus.DataSets;
 using Keda.CosmosDb.Scaler.Demo.Shared;
@@ -85,7 +86,11 @@ private static bool ReadIsSingleArticle()
 
         private static async Task CreateOrdersAsync(int count, bool isSingleArticle)
         {
-            Container container = new CosmosClient(_cosmosDbConfig.Connection)
+            using var cosmosClient = _cosmosDbConfig.Connection.Contains("AccountKey")
+                ? new CosmosClient(_cosmosDbConfig.Connection, new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway })
+                : new CosmosClient(_cosmosDbConfig.Connection, new DefaultAzureCredential(), new CosmosClientOptions { ConnectionMode = ConnectionMode.Direct });
+
+            Container container = cosmosClient
                 .GetContainer(_cosmosDbConfig.DatabaseId, _cosmosDbConfig.ContainerId);
 
             int remainingCount = count;

diff --git a/src/Scaler/Keda.CosmosDb.Scaler.csproj b/src/Scaler/Keda.CosmosDb.Scaler.csproj
@@ -9,6 +9,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <PackageReference Include="Azure.Identity" Version="1.11.4" />
     <PackageReference Include="Grpc.AspNetCore" Version="2.63.0" />
     <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.40.0" />
     <PackageReference Include="Microsoft.Extensions.Logging" Version="6.0.0" />

diff --git a/src/Scaler/Services/CosmosDbFactory.cs b/src/Scaler/Services/CosmosDbFactory.cs
@@ -1,4 +1,5 @@
 using System.Collections.Concurrent;
+using Azure.Identity;
 using Microsoft.Azure.Cosmos;
 
 namespace Keda.CosmosDb.Scaler
@@ -9,14 +10,13 @@ internal sealed class CosmosDbFactory
         // maintain a single instance of CosmosClient per lifetime of the application.
         private readonly ConcurrentDictionary<string, CosmosClient> _cosmosClientCache = new();
 
-        public CosmosClient GetCosmosClient(string connection)
-        {
-            return _cosmosClientCache.GetOrAdd(connection, CreateCosmosClient);
-        }
+        public CosmosClient GetCosmosClient(string connection) =>
+            _cosmosClientCache.GetOrAdd(connection, CreateCosmosClient);
+
+        private CosmosClient CreateCosmosClient(string connection) =>
+            connection.Contains("Accountkey") ?
+                new CosmosClient(connection, new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway }) :
+                new CosmosClient(connection, new DefaultAzureCredential(), new CosmosClientOptions { ConnectionMode = ConnectionMode.Direct });
 
-        private CosmosClient CreateCosmosClient(string connection)
-        {
-            return new CosmosClient(connection, new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway });
-        }
     }
 }
diff --git a/src/Scaler/Services/CosmosDbMetricProvider.cs b/src/Scaler/Services/CosmosDbMetricProvider.cs
@@ -3,6 +3,7 @@
 using System.Net;
 using System.Net.Http;
 using System.Threading.Tasks;
+using Azure.Identity;
 using Microsoft.Azure.Cosmos;
 using Microsoft.Extensions.Logging;