Allow specifying a custom Cosmos DB partition key path (#6)

* Updated to allow specifying Cosmos Db partition key path

* Updated to make backwards compatible with existing version. Kept previous version of the extension methods and also the default value of `/pk` for the partition key path.

* Revert the now obsolete documentation changes

---------

Co-authored-by: Andrew Moreno <[email protected]>
Co-authored-by: f1x3d <[email protected]>

DistributedLeaseManager.AzureCosmosDb/DistributedLeaseCosmosDb.cs

@@ -1,22 +1,24 @@
-using System.Net;
 using DistributedLeaseManager.Core;
 using Microsoft.Azure.Cosmos;
 using Microsoft.Extensions.Options;
 using Newtonsoft.Json.Linq;
+using System.Net;
 
 namespace DistributedLeaseManager.AzureCosmosDb;
 
 public class DistributedLeaseCosmosDb : IDistributedLeaseRepository
 {
     private readonly CosmosClient _cosmosClient;
     private readonly DistributedLeaseCosmosDbOptions _options;
+    private readonly string _partitionKeyPropertyName;
 
     public DistributedLeaseCosmosDb(
         CosmosClient cosmosClient,
         IOptions<DistributedLeaseCosmosDbOptions> options)
     {
         _cosmosClient = cosmosClient;
         _options = options.Value;
+        _partitionKeyPropertyName = _options.PartitionKeyPath[1..];
     }
 
     public async Task EnsureCreated()
@@ -25,25 +27,19 @@ public async Task EnsureCreated()
 
         await _cosmosClient
             .GetDatabase(_options.DatabaseName)
-            .CreateContainerIfNotExistsAsync(_options.ContainerName, "/pk");
+            .CreateContainerIfNotExistsAsync(_options.ContainerName, _options.PartitionKeyPath);
     }
 
     public async Task<bool> Add(DistributedLease lease)
     {
-        var cosmosLease = new
-        {
-            id = lease.ResourceId.ToString(),
-            category = lease.ResourceCategory,
-            expirationTime = lease.ExpirationTime,
-            pk = GetPartitionKey(lease),
-        };
+        var cosmosLease = CreateLease(lease);
 
         try
         {
             await _cosmosClient
                 .GetDatabase(_options.DatabaseName)
                 .GetContainer(_options.ContainerName)
-                .CreateItemAsync(cosmosLease, new(cosmosLease.pk));
+                .CreateItemAsync(cosmosLease, new (GetPartitionKey(lease)));
 
             return true;
         }
@@ -80,28 +76,22 @@ await _cosmosClient
 
     public async Task<bool> Update(DistributedLease lease)
     {
-        var cosmosLease = new
-        {
-            id = lease.ResourceId.ToString(),
-            category = lease.ResourceCategory,
-            expirationTime = lease.ExpirationTime,
-            pk = GetPartitionKey(lease),
-        };
+        var cosmosLease = CreateLease(lease);
 
         try
         {
             await _cosmosClient
                 .GetDatabase(_options.DatabaseName)
                 .GetContainer(_options.ContainerName)
-                .ReplaceItemAsync(cosmosLease, cosmosLease.id, new(cosmosLease.pk), new()
+                .ReplaceItemAsync(cosmosLease, cosmosLease["id"].ToString(), new(GetPartitionKey(lease)), new()
                 {
                     IfMatchEtag = lease.ETag
                 });
 
             return true;
         }
         catch (CosmosException ex)
-        when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
+            when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
         {
             return false;
         }
@@ -122,4 +112,18 @@ private static string GetPartitionKey(DistributedLease lease)
 
     private static string GetPartitionKey(string resourceCategory, Guid resourceId)
         => $"{resourceCategory}/{resourceId}";
+
+    private JObject CreateLease(DistributedLease lease) =>
+        CreateLease(lease.ResourceId, lease.ResourceCategory, lease.ExpirationTime, GetPartitionKey(lease));
+
+    private JObject CreateLease(Guid resourceId, string category, DateTimeOffset expirationTime, string partitionKey)
+    {
+        return new()
+        {
+            ["id"] = resourceId.ToString(),
+            ["category"] = category,
+            ["expirationTime"] = expirationTime,
+            [_partitionKeyPropertyName] = partitionKey,
+        };
+    }
 }

DistributedLeaseManager.AzureCosmosDb/DistributedLeaseCosmosDbOptions.cs

@@ -4,4 +4,31 @@ public class DistributedLeaseCosmosDbOptions
 {
     public string DatabaseName { get; set; } = string.Empty;
     public string ContainerName { get; set; } = string.Empty;
+
+    private string _partitionKeyPath = "/pk";
+
+    public string PartitionKeyPath
+    {
+        get => _partitionKeyPath;
+        set
+        {
+            if (string.IsNullOrWhiteSpace(value))
+            {
+                throw new ArgumentNullException(nameof(value));
+            }
+
+            if (!value.StartsWith("/"))
+            {
+                throw new ArgumentException("Must start with /.", nameof(value));
+            }
+
+            if (value.LastIndexOf('/') > 0)
+            {
+                throw new ArgumentException("Nested partition key paths are not supported at this time.",
+                    nameof(value));
+            }
+
+            _partitionKeyPath = value;
+        }
+    }
 }

DistributedLeaseManager.AzureCosmosDb/README.md

@@ -23,6 +23,19 @@ This library contains a lease storage implemented using the Azure Cosmos DB.
     ```csharp
     builder.Services.AddCosmosDbDistributedLeaseManager("ConnectionString", "DatabaseName", "DistributedLeases");
     ```
+
+    or utilize one of the overloads that accepts an action in order to also customize the partition key path:
+    ```csharp
+    builder.Services.AddCosmosDbDistributedLeaseManager(
+        "AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
+        options =>
+        {
+            options.DatabaseName = "DatabaseName";
+            options.ContainerName = "DistributedLeases";
+            options.PartitionKeyPath = "/partitionKey";
+        }); 
+    ```
+   Note: Partition Key Path must start with `/`. Nested partition key paths such as `/my/partitionKey` are not supported.
 
 1. Inside your controller/service inject the `IDistributedLeaseManager` and call the `TryAcquireLease` method. Verify if the result was successful - if it was then you can proceed with the operation; otherwise, someone else has acquired the lease:
     ```csharp

DistributedLeaseManager.AzureCosmosDb/ServiceCollectionExtensions.cs

@@ -23,11 +23,28 @@ public static IServiceCollection AddCosmosDbDistributedLeaseManager(
         string databaseName,
         string containerName)
     {
-        services.Configure<DistributedLeaseCosmosDbOptions>(options =>
+        return services.AddCosmosDbDistributedLeaseManager(options =>
         {
             options.DatabaseName = databaseName;
             options.ContainerName = containerName;
         });
+    }
+
+    public static IServiceCollection AddCosmosDbDistributedLeaseManager(
+        this IServiceCollection services,
+        string cosmosDbConnectionString,
+        Action<DistributedLeaseCosmosDbOptions> optionsConfiguration)
+    {
+        services.TryAddSingleton(new CosmosClient(cosmosDbConnectionString));
+
+        return services.AddCosmosDbDistributedLeaseManager(optionsConfiguration);
+    }
+
+    public static IServiceCollection AddCosmosDbDistributedLeaseManager(
+        this IServiceCollection services,
+        Action<DistributedLeaseCosmosDbOptions> optionsConfiguration)
+    {
+        services.Configure(optionsConfiguration);
 
         services.TryAddScoped<IDistributedLeaseRepository, DistributedLeaseCosmosDb>();
         services.TryAddScoped<IDistributedLeaseManager, DistributedLeaseManager.Core.DistributedLeaseManager>();

DistributedLeaseManager.Example/Program.cs

@@ -13,8 +13,12 @@
 // See https://learn.microsoft.com/en-us/azure/cosmos-db/emulator
 //builder.Services.AddCosmosDbDistributedLeaseManager(
 //    "AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
-//    "DatabaseName",
-//    "DistributedLeases");
+//    options =>
+//    {
+//        options.DatabaseName = "DatabaseName";
+//        options.ContainerName = "DistributedLeases";
+//        options.PartitionKeyPath = "/partitionKey";
+//    });
 
 // The following example uses a SQL Server Express LocalDB
 // See https://learn.microsoft.com/en-us/sql/database-engine/configure-windows/sql-server-express-localdb