Fix handling when there is no NextRowKey header (#19891)

* Fix handling when there is no NextRowKey header

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

@@ -8,9 +8,13 @@ Thank you to our developer community members who helped to make Azure Tables bet
 
 - Joel Verhagen _([GitHub](https://github.com/joelverhagen))_
 
+### Key Bug Fixes
+
+- 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 on the subsequent page query (A community contribution, courtesy of _[joelverhagen](https://github.com/joelverhagen)_)
+
 ### 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))_
+- 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)_)
 
 ## 12.0.0-beta.6 (2021-03-09)
 

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

@@ -20,6 +20,8 @@ namespace Azure.Data.Tables
     /// </summary>
     public class TableClient
     {
+        private static readonly char[] ContinuationTokenSplit = new[] { ' ' };
+
         private readonly string _table;
         private readonly ClientDiagnostics _diagnostics;
         private readonly TableRestClient _tableOperations;
@@ -1116,7 +1118,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 +1130,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(ContinuationTokenSplit, 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; }