From 24eae81419f1f02331839278bb2ccf73d2bf5168 Mon Sep 17 00:00:00 2001
From: Piotr Karpala <piotrkarpala@microsoft.com>
Date: Mon, 10 Jun 2024 23:09:35 -0400
Subject: [PATCH] Support for Default Azure Identity

Signed-off-by: Piotr Karpala <piotrkarpala@microsoft.com>
---
 ...da.CosmosDb.Scaler.Demo.OrderGenerator.csproj |  1 +
 src/Scaler.Demo/OrderGenerator/Program.cs        |  7 ++++++-
 src/Scaler/Keda.CosmosDb.Scaler.csproj           |  1 +
 src/Scaler/Services/CosmosDbFactory.cs           | 16 ++++++++--------
 src/Scaler/Services/CosmosDbMetricProvider.cs    |  1 +
 5 files changed, 17 insertions(+), 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
index 0cbfce7..4bb7efb 100644
--- 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
index 875094b..77bfac5 100644
--- 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
index 9c9c6e8..5ce777f 100644
--- 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
index 25cf364..249c361 100644
--- 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
index 61242dc..db9cd85 100644
--- 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;