From ef17734ecc6995e8005738e1a8c4596deaa14b4f Mon Sep 17 00:00:00 2001
From: Ankur Charan <ankurcharan98@gmail.com>
Date: Mon, 30 Dec 2024 13:26:55 +0000
Subject: [PATCH] add cosmos mi

---
 .../ServiceCollectionExtensions.cs            | 19 +++++++++++++++++++
 .../Services/CosmosClientFactory.cs           |  5 +++++
 2 files changed, 24 insertions(+)

diff --git a/src/providers/WorkflowCore.Providers.Azure/ServiceCollectionExtensions.cs b/src/providers/WorkflowCore.Providers.Azure/ServiceCollectionExtensions.cs
index d6aec277d..aa3d5f106 100644
--- a/src/providers/WorkflowCore.Providers.Azure/ServiceCollectionExtensions.cs
+++ b/src/providers/WorkflowCore.Providers.Azure/ServiceCollectionExtensions.cs
@@ -46,5 +46,24 @@ public static WorkflowOptions UseCosmosDbPersistence(
             options.UsePersistence(sp => new CosmosDbPersistenceProvider(sp.GetService<ICosmosClientFactory>(), databaseId, sp.GetService<ICosmosDbProvisioner>(), cosmosDbStorageOptions));
             return options;
         }
+
+        public static WorkflowOptions UseCosmosDbPersistence(
+            this WorkflowOptions options,
+            CosmosClient client,
+            string databaseId,
+            CosmosDbStorageOptions cosmosDbStorageOptions = null,
+            CosmosClientOptions clientOptions = null)
+        {
+            if (cosmosDbStorageOptions == null)
+            {
+                cosmosDbStorageOptions = new CosmosDbStorageOptions();
+            }
+
+            options.Services.AddSingleton<ICosmosClientFactory>(sp => new CosmosClientFactory(client));
+            options.Services.AddTransient<ICosmosDbProvisioner>(sp => new CosmosDbProvisioner(sp.GetService<ICosmosClientFactory>(), cosmosDbStorageOptions));
+            options.Services.AddSingleton<IWorkflowPurger>(sp => new WorkflowPurger(sp.GetService<ICosmosClientFactory>(), databaseId, cosmosDbStorageOptions));
+            options.UsePersistence(sp => new CosmosDbPersistenceProvider(sp.GetService<ICosmosClientFactory>(), databaseId, sp.GetService<ICosmosDbProvisioner>(), cosmosDbStorageOptions));
+            return options;
+        }
     }
 }
diff --git a/src/providers/WorkflowCore.Providers.Azure/Services/CosmosClientFactory.cs b/src/providers/WorkflowCore.Providers.Azure/Services/CosmosClientFactory.cs
index ac5dc7f4a..2128d45e0 100644
--- a/src/providers/WorkflowCore.Providers.Azure/Services/CosmosClientFactory.cs
+++ b/src/providers/WorkflowCore.Providers.Azure/Services/CosmosClientFactory.cs
@@ -15,6 +15,11 @@ public CosmosClientFactory(string connectionString, CosmosClientOptions clientOp
             _client = new CosmosClient(connectionString, clientOptions);
         }
 
+        public CosmosClientFactory(CosmosClient client)
+        {
+            _client = client;
+        }
+
         public CosmosClient GetCosmosClient()
         {
             return this._client;