Fix handling when there is no NextRowKey header

sdk/tables/Azure.Data.Tables/CHANGELOG.md

@@ -11,6 +11,7 @@ Thank you to our developer community members who helped to make Azure Tables bet
 ### Changed
 
 - Removed the `Timestamp` property from the serialized entity when sending it to the service as it is ignored by the service (A community contribution, courtesy of _[joelverhagen](https://github.com/joelverhagen))_
+- Fixed handling of paging headers when Table Storage returned a `x-ms-continuation-NextPartitionKey` but no `x-ms-continuation-NextRowKey`. This was causing an HTTP 400 Bad Request on the subsequent page query.
 
 ## 12.0.0-beta.6 (2021-03-09)
 

sdk/tables/Azure.Data.Tables/src/TableClient.cs

@@ -1116,7 +1116,7 @@ internal static string Bind(Expression expression)
             return parser.FilterString == "true" ? null : parser.FilterString;
         }
 
-        private static string CreateContinuationTokenFromHeaders(TableQueryEntitiesHeaders headers)
+        internal static string CreateContinuationTokenFromHeaders(TableQueryEntitiesHeaders headers)
         {
             if (headers.XMsContinuationNextPartitionKey == null && headers.XMsContinuationNextRowKey == null)
             {
@@ -1128,16 +1128,16 @@ private static string CreateContinuationTokenFromHeaders(TableQueryEntitiesHeade
             }
         }
 
-        private static (string NextPartitionKey, string NextRowKey) ParseContinuationToken(string continuationToken)
+        internal static (string NextPartitionKey, string NextRowKey) ParseContinuationToken(string continuationToken)
         {
             // There were no headers passed and the continuation token contains just the space delimiter
-            if (continuationToken?.Length <= 1)
+            if (continuationToken is null || continuationToken.Length <= 1)
             {
                 return (null, null);
             }
 
-            var tokens = continuationToken.Split(' ');
-            return (tokens[0], tokens.Length > 1 ? tokens[1] : null);
+            var tokens = continuationToken.Split(new[] { ' ' }, 2);
+            return (tokens[0], tokens.Length > 1 && tokens[1].Length > 0 ? tokens[1] : null);
         }
     }
 }

sdk/tables/Azure.Data.Tables/tests/TableClientTests.cs

@@ -3,8 +3,8 @@
 
 using System;
 using System.Net;
+using Azure.Core;
 using Azure.Core.TestFramework;
-using Azure.Data.Tables;
 using Azure.Data.Tables.Sas;
 using NUnit.Framework;
 using Parms = Azure.Data.Tables.TableConstants.Sas.Parameters;
@@ -183,6 +183,61 @@ public void EnumPropertiesAreDeSerializedProperly()
             Assert.That(dictEntity.TryGetValue(TableConstants.PropertyNames.Timestamp, out var _), Is.False, "Only PK, RK, and user properties should be sent");
         }
 
+        [Test]
+        public void RoundTripContinuationTokenWithPartitionKeyAndRowKey()
+        {
+            var response = new MockResponse(200);
+            (string NextPartitionKey, string NextRowKey) expected = ("next-pk", "next-rk");
+            response.AddHeader(new HttpHeader("x-ms-continuation-NextPartitionKey", expected.NextPartitionKey));
+            response.AddHeader(new HttpHeader("x-ms-continuation-NextRowKey", expected.NextRowKey));
+            var headers = new TableQueryEntitiesHeaders(response);
+
+            var continuationToken = TableClient.CreateContinuationTokenFromHeaders(headers);
+            var actual = TableClient.ParseContinuationToken(continuationToken);
+
+            Assert.That(continuationToken, Is.EqualTo("next-pk next-rk"));
+            Assert.That(actual, Is.EqualTo(expected));
+        }
+
+        [Test]
+        public void RoundTripContinuationTokenWithPartitionKeyAndNoRowKey()
+        {
+            var response = new MockResponse(200);
+            (string NextPartitionKey, string NextRowKey) expected = ("next-pk", null);
+            response.AddHeader(new HttpHeader("x-ms-continuation-NextPartitionKey", expected.NextPartitionKey));
+            var headers = new TableQueryEntitiesHeaders(response);
+
+            var continuationToken = TableClient.CreateContinuationTokenFromHeaders(headers);
+            var actual = TableClient.ParseContinuationToken(continuationToken);
+
+            Assert.That(continuationToken, Is.EqualTo("next-pk "));
+            Assert.That(actual, Is.EqualTo(expected));
+        }
+
+        [Test]
+        public void NullContinuationTokenReturnsWhenWithNoPartitionKeyAndNoRowKey()
+        {
+            var response = new MockResponse(200);
+            (string NextPartitionKey, string NextRowKey) expected = (null, null);
+            var headers = new TableQueryEntitiesHeaders(response);
+
+            var continuationToken = TableClient.CreateContinuationTokenFromHeaders(headers);
+            var actual = TableClient.ParseContinuationToken(continuationToken);
+
+            Assert.That(continuationToken, Is.Null);
+            Assert.That(actual, Is.EqualTo(expected));
+        }
+
+        [Test]
+        public void HandlesEmptyStringContinuationToken()
+        {
+            (string NextPartitionKey, string NextRowKey) expected = (null, null);
+
+            var actual = TableClient.ParseContinuationToken(" ");
+
+            Assert.That(actual, Is.EqualTo(expected));
+        }
+
         public class EnumEntity : ITableEntity
         {
             public string PartitionKey { get; set; }