Extend LeaseCollection allowed partition key paths to "\partitionKey" to support gremlin accounts

src/DocumentDB.ChangeFeedProcessor.IntegrationTests/IntegrationTest.cs

@@ -90,12 +90,16 @@ protected DocumentCollectionInfo MonitoredCollectionInfo
 
         protected readonly bool IsPartitionedLeaseCollection;
 
+        protected readonly bool isPartitionedByLeasePk;
+
         public IntegrationTest(
             bool isPartitionedMonitoredCollection = true,
-            bool isPartitionedLeaseCollection = false)
+            bool isPartitionedLeaseCollection = false,
+            bool isPartitionedByLeasePk = false)
         {
             this.IsPartitionedMonitoredCollection = isPartitionedMonitoredCollection;
             this.IsPartitionedLeaseCollection = isPartitionedLeaseCollection;
+            this.isPartitionedByLeasePk = isPartitionedByLeasePk;
         }
 
         public async Task InitializeAsync()
@@ -119,7 +123,7 @@ public async Task InitializeAsync()
 
             if (this.IsPartitionedLeaseCollection)
             {
-                leaseCollection.PartitionKey = new PartitionKeyDefinition { Paths = { "/id" } };
+                leaseCollection.PartitionKey = this.isPartitionedByLeasePk ? new PartitionKeyDefinition { Paths = { "/leasepk" } } : new PartitionKeyDefinition { Paths = { "/id" } };
             }
 
             using (var client = new DocumentClient(this.LeaseCollectionInfo.Uri, this.LeaseCollectionInfo.MasterKey, this.LeaseCollectionInfo.ConnectionPolicy))
@@ -160,7 +164,7 @@ private async Task CreateMonitoredCollectionAsync(string monitoredCollectionName
 
             if (this.IsPartitionedMonitoredCollection)
             {
-                monitoredCollection.PartitionKey = new PartitionKeyDefinition { Paths = { "/id" } };
+                monitoredCollection.PartitionKey = new PartitionKeyDefinition { Paths = { "/partitionId" } };
             }
             else
             {

src/DocumentDB.ChangeFeedProcessor.IntegrationTests/LeasePkLeaseCollectionTests.cs

@@ -0,0 +1,21 @@
+//----------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  Licensed under the MIT license.
+//----------------------------------------------------------------
+
+namespace Microsoft.Azure.Documents.ChangeFeedProcessor.IntegrationTests
+{
+    using Xunit;
+
+    /// <summary>
+    /// To truly test this class, run emulator with /EnableGremlinEndpoint
+    /// </summary>
+    [Trait("Category", "Integration")]
+    [Collection("Integration tests")]
+    public class LeasePkLeaseCollectionTests : StaticCollectionTests
+    {
+        public LeasePkLeaseCollectionTests() :
+            base(true,true)
+        {
+        }
+    }
+}

src/DocumentDB.ChangeFeedProcessor.IntegrationTests/StaticCollectionTests.cs

@@ -33,6 +33,12 @@ public StaticCollectionTests(
         {
         }
 
+        public StaticCollectionTests(
+            bool isPartitionedLeaseCollection, bool isPartitionedByLeasePk) :
+            base(isPartitionedLeaseCollection: isPartitionedLeaseCollection, isPartitionedByLeasePk: isPartitionedByLeasePk)
+        {
+        }
+
         [Fact]
         public async Task CountDocumentsInCollection_NormalCase()
         {

src/DocumentDB.ChangeFeedProcessor.UnitTests/ChangeFeedProcessorBuilderTests.cs

@@ -304,6 +304,49 @@ public async Task BuildPassesPartitionKey_WhenLeaseCollectionIsPartitionedById()
             Assert.Equal(typeof(LeaseLostException), exception.GetType());
         }
 
+        [Fact]
+        public async Task BuildPassesPartitionKey_WhenLeaseCollectionIsPartitionedByLeasePk()
+        {
+            var leaseCollection = MockHelpers.CreateCollection(
+                "collectionId",
+                "collectionRid",
+                new PartitionKeyDefinition { Paths = { "/leasepk" } },
+                collectionLink);
+
+            var lease = new DocumentServiceLease()
+            {
+                LeasePartitionKey = "leasePk",
+                Id = "Id"
+            };
+
+            var leaseClient = this.CreateMockDocumentClient(collection);
+            Mock.Get(leaseClient)
+                .Setup(c => c.ReadDocumentCollectionAsync(
+                    It.IsAny<Uri>(),
+                    It.IsAny<RequestOptions>()))
+                .ReturnsAsync(new ResourceResponse<DocumentCollection>(leaseCollection));
+            Mock.Get(leaseClient)
+                .Setup(c => c.ReadDocumentAsync(
+                    It.IsAny<Uri>(),
+                    It.IsAny<RequestOptions>(),
+                    It.IsAny<CancellationToken>()))
+                .Callback((Uri uri, RequestOptions options, CancellationToken token) =>
+                {
+                    if (new PartitionKey("leasePk").Equals(options.PartitionKey))
+                        throw DocumentExceptionHelpers.CreateNotFoundException();   // Success code path: cause lease lost.
+                    throw new Exception("Failure");
+                });
+
+            this.builder
+                .WithFeedDocumentClient(this.CreateMockDocumentClient())
+                .WithLeaseDocumentClient(leaseClient)
+                .WithObserverFactory(Mock.Of<IChangeFeedObserverFactory>());
+            await this.builder.BuildAsync();
+
+            Exception exception = await Record.ExceptionAsync(() => this.builder.LeaseStoreManager.ReleaseAsync(lease));
+            Assert.Equal(typeof(LeaseLostException), exception.GetType());
+        }
+
         [Fact]
         public async Task BuildThrowsWhenBothPartitionProcessorFactoriesSpecified()
         {

src/DocumentDB.ChangeFeedProcessor.UnitTests/LeaseManagement/DocumentServiceLeaseTests.cs

@@ -8,6 +8,8 @@ namespace Microsoft.Azure.Documents.ChangeFeedProcessor.UnitTests.LeaseManagemen
     using System.IO;
     using System.Runtime.Serialization.Formatters.Binary;
     using Microsoft.Azure.Documents.ChangeFeedProcessor.LeaseManagement;
+    using Microsoft.Azure.Documents.ChangeFeedProcessor.PartitionManagement;
+    using Newtonsoft.Json;
     using Xunit;
 
     [Trait("Category", "Gated")]
@@ -55,6 +57,7 @@ public void ValidateProperties()
             Assert.Equal(timestamp, lease.Timestamp);
             Assert.Equal(value, lease.Properties[key]);
             Assert.Equal(etag, lease.ConcurrencyToken);
+            Assert.Null(lease.LeasePartitionKey);
         }
 
         [Fact]
@@ -63,6 +66,7 @@ public void ValidateSerialization_AllFields()
             DocumentServiceLease originalLease = new DocumentServiceLease
             {
                 Id = "id",
+                LeasePartitionKey="leasePk",
                 ETag = "etag",
                 PartitionId = "0",
                 Owner = "owner",
@@ -80,6 +84,7 @@ public void ValidateSerialization_AllFields()
             var lease = (DocumentServiceLease)formatter.Deserialize(stream2);
 
             Assert.Equal(originalLease.Id, lease.Id);
+            Assert.Equal(originalLease.LeasePartitionKey, lease.LeasePartitionKey);
             Assert.Equal(originalLease.ETag, lease.ETag);
             Assert.Equal(originalLease.PartitionId, lease.PartitionId);
             Assert.Equal(originalLease.Owner, lease.Owner);
@@ -102,12 +107,121 @@ public void ValidateSerialization_NullFields()
             var lease = (DocumentServiceLease)formatter.Deserialize(stream2);
 
             Assert.Null(lease.Id);
+            Assert.Null(lease.LeasePartitionKey);
             Assert.Null(lease.ETag);
             Assert.Null(lease.PartitionId);
             Assert.Null(lease.Owner);
             Assert.Null(lease.ContinuationToken);
             Assert.Equal(new DocumentServiceLease().Timestamp, lease.Timestamp);
             Assert.Empty(lease.Properties);
         }
+
+
+        #region Compat_Tests
+        // this class doesnt contain LeaseId property
+
+        [Serializable]
+        class DocumentServiceLeaseV1
+        {
+            private static readonly DateTime UnixStartTime = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
+
+            [JsonProperty("id")]
+            public string Id { get; set; }
+
+            [JsonProperty("_etag")]
+            public string ETag { get; set; }
+
+            [JsonProperty("PartitionId")]
+            public string PartitionId { get; set; }
+
+            [JsonProperty("Owner")]
+            public string Owner { get; set; }
+
+            /// <summary>
+            /// Gets or sets the current value for the offset in the stream.
+            /// </summary>
+            [JsonProperty("ContinuationToken")]
+            public string ContinuationToken { get; set; }
+
+            [JsonIgnore]
+            public DateTime Timestamp
+            {
+                get { return this.ExplicitTimestamp ?? UnixStartTime.AddSeconds(this.TS); }
+                set { this.ExplicitTimestamp = value; }
+            }
+
+            [JsonIgnore]
+            public string ConcurrencyToken => this.ETag;
+
+            [JsonProperty("properties")]
+            public Dictionary<string, string> Properties { get; set; } = new Dictionary<string, string>();
+
+            [JsonIgnore]
+            public LeaseAcquireReason AcquireReason { get; set; }
+
+            [JsonProperty("timestamp")]
+            private DateTime? ExplicitTimestamp { get; set; }
+
+            [JsonProperty("_ts")]
+            private long TS { get; set; }
+
+        }
+
+        [Fact]
+        public void ValidateBackwardCompat_OldLeaseFormat()
+        {
+
+            DocumentServiceLeaseV1 originalLease = new DocumentServiceLeaseV1
+            {
+                Id = "id",
+                ETag = "etag",
+                PartitionId = "0",
+                Owner = "owner",
+                ContinuationToken = "continuation",
+                Timestamp = DateTime.Now - TimeSpan.FromSeconds(5),
+                Properties = new Dictionary<string, string> { { "key", "value" } }
+            };
+
+            var serializedV1Lease = JsonConvert.SerializeObject(originalLease);
+            var lease = JsonConvert.DeserializeObject<DocumentServiceLease>(serializedV1Lease);
+
+            Assert.Equal(originalLease.Id, lease.Id);
+            Assert.Equal(null, lease.LeasePartitionKey);
+            Assert.Equal(originalLease.ETag, lease.ETag);
+            Assert.Equal(originalLease.PartitionId, lease.PartitionId);
+            Assert.Equal(originalLease.Owner, lease.Owner);
+            Assert.Equal(originalLease.ContinuationToken, lease.ContinuationToken);
+            Assert.Equal(originalLease.Timestamp, lease.Timestamp);
+            Assert.Equal(originalLease.Properties["key"], lease.Properties["key"]);
+        }
+
+        [Fact]
+        public void ValidateForwardCompat_OldLeaseFormat()
+        {
+
+            DocumentServiceLease originalLease = new DocumentServiceLease
+            {
+                Id = "id",
+                ETag = "etag",
+                PartitionId = "0",
+                Owner = "owner",
+                ContinuationToken = "continuation",
+                Timestamp = DateTime.Now - TimeSpan.FromSeconds(5),
+                Properties = new Dictionary<string, string> { { "key", "value" } }
+            };
+
+            var serializedLease = JsonConvert.SerializeObject(originalLease);
+            var lease = JsonConvert.DeserializeObject<DocumentServiceLeaseV1>(serializedLease);
+
+            Assert.Equal(originalLease.Id, lease.Id);
+            Assert.Equal(originalLease.ETag, lease.ETag);
+            Assert.Equal(originalLease.PartitionId, lease.PartitionId);
+            Assert.Equal(originalLease.Owner, lease.Owner);
+            Assert.Equal(originalLease.ContinuationToken, lease.ContinuationToken);
+            Assert.Equal(originalLease.Timestamp, lease.Timestamp);
+            Assert.Equal(originalLease.Properties["key"], lease.Properties["key"]);
+        }
+        #endregion
+
     }
 }

src/DocumentDB.ChangeFeedProcessor/ChangeFeedProcessorBuilder.cs

@@ -374,7 +374,7 @@ public async Task<IChangeFeedProcessor> BuildAsync()
 
             await this.InitializeCollectionPropertiesForBuildAsync().ConfigureAwait(false);
 
-            ILeaseStoreManager leaseStoreManager = await this.GetLeaseStoreManagerAsync(this.leaseCollectionLocation, true).ConfigureAwait(false);
+            ILeaseStoreManager leaseStoreManager = await this.GetLeaseStoreManagerAsync(this.leaseCollectionLocation).ConfigureAwait(false);
             IPartitionManager partitionManager = this.BuildPartitionManager(leaseStoreManager);
             return new ChangeFeedProcessor(partitionManager);
         }
@@ -397,7 +397,7 @@ public async Task<IRemainingWorkEstimator> BuildEstimatorAsync()
 
             await this.InitializeCollectionPropertiesForBuildAsync().ConfigureAwait(false);
 
-            var leaseStoreManager = await this.GetLeaseStoreManagerAsync(this.leaseCollectionLocation, true).ConfigureAwait(false);
+            var leaseStoreManager = await this.GetLeaseStoreManagerAsync(this.leaseCollectionLocation).ConfigureAwait(false);
 
             IRemainingWorkEstimator remainingWorkEstimator = new RemainingWorkEstimator(
                 leaseStoreManager,
@@ -461,8 +461,7 @@ private IPartitionManager BuildPartitionManager(ILeaseStoreManager leaseStoreMan
         }
 
         private async Task<ILeaseStoreManager> GetLeaseStoreManagerAsync(
-            DocumentCollectionInfo collectionInfo,
-            bool isPartitionKeyByIdRequiredIfPartitioned)
+            DocumentCollectionInfo collectionInfo)
         {
             if (this.LeaseStoreManager == null)
             {
@@ -473,15 +472,26 @@ private async Task<ILeaseStoreManager> GetLeaseStoreManagerAsync(
                     collection.PartitionKey != null &&
                     collection.PartitionKey.Paths != null &&
                     collection.PartitionKey.Paths.Count > 0;
-                if (isPartitioned && isPartitionKeyByIdRequiredIfPartitioned &&
-                    (collection.PartitionKey.Paths.Count != 1 || collection.PartitionKey.Paths[0] != "/id"))
+
+                if (isPartitioned &&
+                    (collection.PartitionKey.Paths.Count != 1 || !(collection.PartitionKey.Paths[0].Equals($"/{DocumentServiceLease.IdPropertyName}", StringComparison.OrdinalIgnoreCase) ||
+                                                                   collection.PartitionKey.Paths[0].Equals($"/{DocumentServiceLease.LeasePartitionKeyPropertyName}", StringComparison.OrdinalIgnoreCase))))
                 {
-                    throw new ArgumentException("The lease collection, if partitioned, must have partition key equal to id.");
+                    throw new ArgumentException($"The lease collection, if partitioned, must have partition key equal to {DocumentServiceLease.IdPropertyName} or {DocumentServiceLease.LeasePartitionKeyPropertyName}.");
                 }
 
-                var requestOptionsFactory = isPartitioned ?
-                    (IRequestOptionsFactory)new PartitionedByIdCollectionRequestOptionsFactory() :
-                    (IRequestOptionsFactory)new SinglePartitionRequestOptionsFactory();
+                IRequestOptionsFactory requestOptionsFactory = null;
+                if (isPartitioned)
+                {
+                    // we allow only id or leasePk partitioning so check only flag
+                    requestOptionsFactory = collection.PartitionKey.Paths[0].Equals($"/{DocumentServiceLease.IdPropertyName}", StringComparison.OrdinalIgnoreCase) ?
+                                            (IRequestOptionsFactory)new PartitionedByIdCollectionRequestOptionsFactory() :
+                                            (IRequestOptionsFactory)new PartitionedByLeasePkCollectionRequestOptionsFactory();
+                }
+                else
+                {
+                    requestOptionsFactory = (IRequestOptionsFactory)new SinglePartitionRequestOptionsFactory();
+                }
 
                 string leasePrefix = this.GetLeasePrefix();
                 var leaseStoreManagerBuilder = new DocumentServiceLeaseStoreManagerBuilder()

src/DocumentDB.ChangeFeedProcessor/LeaseManagement/DocumentServiceLease.cs

@@ -14,6 +14,9 @@ namespace Microsoft.Azure.Documents.ChangeFeedProcessor.LeaseManagement
     [Serializable]
     internal class DocumentServiceLease : ILease, ILeaseAcquireReasonProvider
     {
+        internal const string IdPropertyName = "id";
+        internal const string LeasePartitionKeyPropertyName = "leasepk";
+
         private static readonly DateTime UnixStartTime = new DateTime(1970, 1, 1, 0, 0, 0, 0, DateTimeKind.Utc);
 
         public DocumentServiceLease()
@@ -32,9 +35,17 @@ public DocumentServiceLease(DocumentServiceLease other)
             this.Properties = other.Properties;
         }
 
-        [JsonProperty("id")]
+        [JsonProperty(IdPropertyName)]
         public string Id { get; set; }
 
+        /// <summary>
+        /// Gets or sets property to be used as partition key path for lease collections.
+        /// This is clone of existing Id property to maintain backward compat.
+        /// This property name is compatible to both GremlinAccounts and SqlAccounts
+        /// </summary>
+        [JsonProperty(LeasePartitionKeyPropertyName)]
+        public string LeasePartitionKey { get; set; }
+
         [JsonProperty("_etag")]
         public string ETag { get; set; }