Support continuation tokens for tables

sdk/tables/Azure.Data.Tables/api/Azure.Data.Tables.netstandard2.0.cs

@@ -26,7 +26,7 @@ public TableServiceClient(System.Uri endpoint, Azure.Data.Tables.TablesSharedKey
         public virtual System.Threading.Tasks.Task<Azure.Data.Tables.Models.TableResponse> CreateTableAsync(string tableName, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual Azure.Response DeleteTable(string tableName, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual System.Threading.Tasks.Task<Azure.Response> DeleteTableAsync(string tableName, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
-        public Azure.Data.Tables.TableClient GetTableClient(string tableName) { throw null; }
+        public virtual Azure.Data.Tables.TableClient GetTableClient(string tableName) { throw null; }
         public virtual Azure.Pageable<Azure.Data.Tables.Models.TableResponseProperties> GetTables(string select = null, string filter = null, int? top = default(int?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public virtual Azure.AsyncPageable<Azure.Data.Tables.Models.TableResponseProperties> GetTablesAsync(string select = null, string filter = null, int? top = default(int?), System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
@@ -109,14 +109,6 @@ public LoggingSettings(string version, bool delete, bool read, bool write, Azure
         public static bool operator !=(Azure.Data.Tables.Models.OdataMetadataFormat left, Azure.Data.Tables.Models.OdataMetadataFormat right) { throw null; }
         public override string ToString() { throw null; }
     }
-    public partial class QueryOptions
-    {
-        public QueryOptions() { }
-        public string Filter { get { throw null; } set { } }
-        public Azure.Data.Tables.Models.OdataMetadataFormat? Format { get { throw null; } set { } }
-        public string Select { get { throw null; } set { } }
-        public int? Top { get { throw null; } set { } }
-    }
     public partial class RequestMetrics
     {
         public RequestMetrics(bool enabled) { }

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

@@ -9,6 +9,7 @@ namespace Azure.Data.Tables.Models
 {
     /// <summary> Azure Analytics Logging settings. </summary>
     [CodeGenModel("Logging")]
+    // Should be unnecessary when https://github.com/Azure/azure-rest-api-specs/pull/8151/files#r411681931 is implemented.
     public partial class LoggingSettings
     {
     }

sdk/tables/Azure.Data.Tables/src/swagger/OdataMetadataFormat.cs → sdk/tables/Azure.Data.Tables/src/OdataMetadataFormat.cs

@@ -7,8 +7,9 @@
 
 namespace Azure.Data.Tables.Models
 {
-    /// <summary> The OdataMetadataFormat. </summary>
+    /// <summary> The OdataMetadataFormat.</summary>
     [CodeGenModel("Enum0")]
+    // Should be unnecessary when https://github.com/Azure/azure-rest-api-specs/pull/8151/files#r411679368 is implemented.
     public readonly partial struct OdataMetadataFormat
     {
     }

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

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Azure.Data.Tables.Models
+{
+    ///<inheritdoc/>
+    internal partial class QueryOptions
+    {
+        /// <summary> The continuation token for a Tables request. </summary>
+        public string NextTableName { get; set; }
+        /// <summary> The partition key continuation token for an Entity request. </summary>
+        public string NextPartitionKey { get; set; }
+        /// <summary> The row key continuation token for an Entity request. </summary>
+        public string NextRowKey { get; set; }
+    }
+}

sdk/tables/Azure.Data.Tables/src/swagger/RequestMetrics.cs → sdk/tables/Azure.Data.Tables/src/RequestMetrics.cs

@@ -9,6 +9,7 @@ namespace Azure.Data.Tables.Models
 {
     /// <summary> The RequestMetrics. </summary>
     [CodeGenModel("Metrics")]
+    // Should be unnecessary when https://github.com/Azure/azure-rest-api-specs/pull/8151/files#r411684457 is implemented.
     public partial class RequestMetrics
     {
     }

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

@@ -42,7 +42,10 @@ protected TableClient()
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>The inserted Table entity.</returns>
         public virtual async Task<Response<IReadOnlyDictionary<string, object>>> InsertAsync(IDictionary<string, object> entity, CancellationToken cancellationToken = default) =>
-            await InsertInternalAsync(true, entity, cancellationToken).ConfigureAwait(false);
+            await _tableOperations.InsertEntityAsync(_table,
+                                                     tableEntityProperties: entity,
+                                                     queryOptions: new QueryOptions() { Format = _format },
+                                                     cancellationToken: cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Inserts a Table Entity into the Table.
@@ -51,7 +54,10 @@ public virtual async Task<Response<IReadOnlyDictionary<string, object>>> InsertA
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>The inserted Table entity.</returns>
         public virtual Response<IReadOnlyDictionary<string, object>> Insert(IDictionary<string, object> entity, CancellationToken cancellationToken = default) =>
-            InsertInternalAsync(false, entity, cancellationToken).EnsureCompleted();
+            _tableOperations.InsertEntity(_table,
+                                          tableEntityProperties: entity,
+                                          queryOptions: new QueryOptions() { Format = _format },
+                                          cancellationToken: cancellationToken);
 
         /// <summary>
         /// Updates the specified table entity.
@@ -62,7 +68,12 @@ public virtual Response<IReadOnlyDictionary<string, object>> Insert(IDictionary<
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>The <see cref="Response"/> indicating the result of the operation.</returns>
         public virtual async Task<Response> UpdateAsync(string partitionKey, string rowKey, IDictionary<string, object> entity, CancellationToken cancellationToken = default) =>
-            await UpdateInternalAsync(true, partitionKey, rowKey, entity, cancellationToken).ConfigureAwait(false);
+            await _tableOperations.UpdateEntityAsync(_table,
+                                                     partitionKey,
+                                                     rowKey,
+                                                     tableEntityProperties: entity,
+                                                     queryOptions: new QueryOptions() { Format = _format },
+                                                     cancellationToken: cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Updates the specified table entity.
@@ -73,7 +84,12 @@ public virtual async Task<Response> UpdateAsync(string partitionKey, string rowK
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>The <see cref="Response"/> indicating the result of the operation.</returns>
         public virtual Response Update(string partitionKey, string rowKey, IDictionary<string, object> entity, CancellationToken cancellationToken = default) =>
-            UpdateInternalAsync(false, partitionKey, rowKey, entity, cancellationToken).EnsureCompleted();
+            _tableOperations.UpdateEntity(_table,
+                                          partitionKey,
+                                          rowKey,
+                                          tableEntityProperties: entity,
+                                          queryOptions: new QueryOptions() { Format = _format },
+                                          cancellationToken: cancellationToken);
 
         /// <summary>
         /// Queries entities in the table.
@@ -85,16 +101,29 @@ public virtual Response Update(string partitionKey, string rowKey, IDictionary<s
         /// <returns></returns>
         public virtual AsyncPageable<IDictionary<string, object>> QueryAsync(string select = null, string filter = null, int? top = null, CancellationToken cancellationToken = default)
         {
-            //TODO: support continuation tokens
-
-            return PageableHelpers.CreateAsyncEnumerable(async tableName =>
+            return PageableHelpers.CreateAsyncEnumerable(async _ =>
             {
-                var response = await _tableOperations.RestClient.QueryEntitiesAsync(_table,
+                var response = await _tableOperations.RestClient.QueryEntitiesAsync(
+                    _table,
                     queryOptions: new QueryOptions() { Format = _format, Top = top, Filter = filter, Select = @select },
-                    cancellationToken: cancellationToken)
-                    .ConfigureAwait(false);
-                return Page.FromValues(response.Value.Value, response.Headers.XMsContinuationNextPartitionKey, response.GetRawResponse());
-            }, (_, __) => throw new NotImplementedException());
+                    cancellationToken: cancellationToken).ConfigureAwait(false);
+
+                return Page.FromValues(response.Value.Value,
+                                       CreateContinuationTokenFromHeaders(response.Headers),
+                                       response.GetRawResponse());
+            }, async (continuationToken, _) =>
+            {
+                var (NextPartitionKey, NextRowKey) = ParseContinuationToken(continuationToken);
+
+                var response = await _tableOperations.RestClient.QueryEntitiesAsync(
+                    _table,
+                    queryOptions: new QueryOptions() { Format = _format, Top = top, Filter = filter, Select = @select, NextPartitionKey = NextPartitionKey, NextRowKey = NextRowKey },
+                    cancellationToken: cancellationToken).ConfigureAwait(false);
+
+                return Page.FromValues(response.Value.Value,
+                                       CreateContinuationTokenFromHeaders(response.Headers),
+                                       response.GetRawResponse());
+            });
         }
 
         /// <summary>
@@ -106,60 +135,52 @@ public virtual AsyncPageable<IDictionary<string, object>> QueryAsync(string sele
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         public virtual Pageable<IDictionary<string, object>> Query(string select = null, string filter = null, int? top = null, CancellationToken cancellationToken = default)
         {
-            //TODO: support continuation tokens
-
-            return PageableHelpers.CreateEnumerable(tableName =>
+            return PageableHelpers.CreateEnumerable(_ =>
             {
                 var response = _tableOperations.RestClient.QueryEntities(_table,
                     queryOptions: new QueryOptions() { Format = _format, Top = top, Filter = filter, Select = @select },
                     cancellationToken: cancellationToken);
-                return Page.FromValues(response.Value.Value, response.Headers.XMsContinuationNextPartitionKey, response.GetRawResponse());
-            }, (_, __) => throw new NotImplementedException());
-        }
-
-        internal async Task<Response<IReadOnlyDictionary<string, object>>> InsertInternalAsync(bool async, IDictionary<string, object> entity, CancellationToken cancellationToken)
-        {
-            Response<IReadOnlyDictionary<string, object>> response;
-
-            if (async)
+                return Page.FromValues(
+                    response.Value.Value,
+                    CreateContinuationTokenFromHeaders(response.Headers),
+                    response.GetRawResponse());
+            }, (continuationToken, _) =>
             {
-                response = await _tableOperations.InsertEntityAsync(
-                    _table,
-                    tableEntityProperties: entity,
-                    queryOptions: new QueryOptions() { Format = _format },
-                    cancellationToken: cancellationToken).ConfigureAwait(false);
-            }
-            {
-                response = _tableOperations.InsertEntity(
+                var (NextPartitionKey, NextRowKey) = ParseContinuationToken(continuationToken);
+
+                var response = _tableOperations.RestClient.QueryEntities(
                     _table,
-                    tableEntityProperties: entity,
-                    queryOptions: new QueryOptions() { Format = _format },
+                    queryOptions: new QueryOptions() { Format = _format, Top = top, Filter = filter, Select = @select, NextPartitionKey = NextPartitionKey, NextRowKey = NextRowKey },
                     cancellationToken: cancellationToken);
-            }
 
-            return response;
+                return Page.FromValues(response.Value.Value,
+                                       CreateContinuationTokenFromHeaders(response.Headers),
+                                       response.GetRawResponse());
+            });
         }
 
-        internal async Task<Response> UpdateInternalAsync(bool async, string partitionKey, string rowKey, IDictionary<string, object> entity, CancellationToken cancellationToken = default)
+        private static string CreateContinuationTokenFromHeaders(TableQueryEntitiesHeaders headers)
         {
-            if (async)
+            if (headers.XMsContinuationNextPartitionKey == null && headers.XMsContinuationNextRowKey == null)
             {
-                return await _tableOperations.UpdateEntityAsync(_table,
-                                                                partitionKey,
-                                                                rowKey,
-                                                                tableEntityProperties: entity,
-                                                                queryOptions: new QueryOptions() { Format = _format },
-                                                                cancellationToken: cancellationToken).ConfigureAwait(false);
+                return null;
             }
             else
             {
-                return _tableOperations.UpdateEntity(_table,
-                                                     partitionKey,
-                                                     rowKey,
-                                                     tableEntityProperties: entity,
-                                                     queryOptions: new QueryOptions() { Format = _format },
-                                                     cancellationToken: cancellationToken);
+                return $"{headers.XMsContinuationNextPartitionKey} {headers.XMsContinuationNextRowKey}";
             }
         }
+
+        private 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)
+            {
+                return (null, null);
+            }
+
+            var tokens = continuationToken.Split(' ');
+            return (tokens[0], tokens.Length > 1 ? tokens[1] : null);
+        }
     }
 }

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

@@ -11,5 +11,12 @@ internal static class HeaderNames
             public const string SharedKey = "SharedKeyLite";
             public const string Authorization = "Authorization";
         }
+
+        internal static class QueryParameterNames
+        {
+            public const string NextTableName = "NextTableName";
+            public const string NextPartitionKey = "NextPartitionKey";
+            public const string NextRowKey = "NextRowKey";
+        }
     }
 }