Support manually handling continuationTokens (#18179)

* Support manually handling continuationTokens

* Use PageSettings and update version

* Update samples and page type

* Address PR Comments

* Undo changes to generated samples

* Fix format

sdk/tables/data-tables/CHANGELOG.md

@@ -1,8 +1,9 @@
 # Release History
 
-## 12.1.3 (Unreleased)
+## 12.2.0 (Unreleased)
 
 ### Features Added
+- Take `continuationToken` as a `PageSetting` and expose it in the page when listing entities `byPage`. [#](https://github.com/Azure/azure-sdk-for-js/pull/)
 
 ### Breaking Changes
 

sdk/tables/data-tables/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@azure/data-tables",
-  "version": "12.1.3",
+  "version": "12.2.0",
   "description": "An isomorphic client library for the Azure Tables service.",
   "sdk-type": "client",
   "main": "dist/index.js",

sdk/tables/data-tables/review/data-tables.api.md

@@ -238,7 +238,7 @@ export class TableClient {
     static fromConnectionString(connectionString: string, tableName: string, options?: TableServiceClientOptions): TableClient;
     getAccessPolicy(options?: OperationOptions): Promise<GetAccessPolicyResponse>;
     getEntity<T extends object = Record<string, unknown>>(partitionKey: string, rowKey: string, options?: GetTableEntityOptions): Promise<GetTableEntityResponse<TableEntityResult<T>>>;
-    listEntities<T extends object = Record<string, unknown>>(options?: ListTableEntitiesOptions): PagedAsyncIterableIterator<TableEntityResult<T>, TableEntityResult<T>[]>;
+    listEntities<T extends object = Record<string, unknown>>(options?: ListTableEntitiesOptions): PagedAsyncIterableIterator<TableEntityResult<T>, TableEntityResultPage<T>>;
     pipeline: Pipeline;
     setAccessPolicy(tableAcl: SignedIdentifier[], options?: OperationOptions): Promise<SetAccessPolicyResponse>;
     submitTransaction(actions: TransactionAction[]): Promise<TableTransactionResponse>;
@@ -293,6 +293,11 @@ export type TableEntityResult<T> = T & {
     timestamp?: string;
 };
 
+// @public
+export type TableEntityResultPage<T> = Array<TableEntityResult<T>> & {
+    continuationToken?: string;
+};
+
 // @public
 export interface TableGetAccessPolicyHeaders {
     clientRequestId?: string;
@@ -448,7 +453,6 @@ export type UpdateTableEntityOptions = OperationOptions & {
 // @public
 export type UpsertEntityResponse = TableMergeEntityHeaders;
 
-
 // (No @packageDocumentation comment for this package)
 
 ```

sdk/tables/data-tables/samples-dev/manualEntityPageQuery.ts

@@ -0,0 +1,66 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+/**
+ * This sample demonstrates how to query entities in a table by page by manually handling the continuation token
+ *
+ * @summary queries entities in a table by page manually handling continuation tokens
+ * @azsdk-weight 50
+ */
+
+import { TableClient, AzureSASCredential, TransactionAction } from "@azure/data-tables";
+
+// Load the .env file if it exists
+import * as dotenv from "dotenv";
+dotenv.config();
+
+const tablesUrl = process.env["TABLES_URL"] || "";
+const sasToken = process.env["SAS_TOKEN"] || "";
+
+async function listEntitiesPage() {
+  const tableName = `manualListByPage`;
+
+  // See authenticationMethods sample for other options of creating a new client
+  const client = new TableClient(tablesUrl, tableName, new AzureSASCredential(sasToken));
+  // Create the table
+  await client.createTable();
+
+  let actions: TransactionAction[] = [];
+
+  // Create 100 entities
+  for (let i = 0; i < 100; i++) {
+    actions.push(["create", { partitionKey: `one`, rowKey: `${i}`, foo: i }]);
+  }
+  await client.submitTransaction(actions);
+
+  // Limit the size to 2 entities by page
+  let iterator = client.listEntities().byPage({ maxPageSize: 2 });
+
+  // Iterating the pages to find the page that contains row key 50
+  let interestingPage: string | undefined;
+  for await (const page of iterator) {
+    if (page.some((p) => p.rowKey === "50")) {
+      interestingPage = page.continuationToken;
+    }
+  }
+
+  if (!interestingPage) {
+    console.error("Didn't find entity with rowKey = 50");
+    return;
+  }
+
+  // Fetch only the page that contains rowKey 50;
+  const page = await client
+    .listEntities()
+    .byPage({ maxPageSize: 2, continuationToken: interestingPage })
+    .next();
+
+  if (!page.done) {
+    for (const entity of page.value) {
+      console.log(entity.rowKey);
+    }
+  }
+}
+
+listEntitiesPage().catch((err) => {
+  console.error("The sample encountered an error:", err);
+});

sdk/tables/data-tables/src/TableClient.ts

@@ -16,7 +16,8 @@ import {
   TransactionAction,
   TableTransactionResponse,
   SignedIdentifier,
-  GetAccessPolicyResponse
+  GetAccessPolicyResponse,
+  TableEntityResultPage
 } from "./models";
 import {
   UpdateEntityResponse,
@@ -24,7 +25,10 @@ import {
   DeleteTableEntityResponse,
   SetAccessPolicyResponse
 } from "./generatedModels";
-import { QueryOptions as GeneratedQueryOptions } from "./generated/models";
+import {
+  QueryOptions as GeneratedQueryOptions,
+  TableQueryEntitiesOptionalParams
+} from "./generated/models";
 import { getClientParamsFromConnectionString } from "./utils/connectionString";
 import {
   isNamedKeyCredential,
@@ -64,6 +68,7 @@ import { isCredential } from "./utils/isCredential";
 import { tablesSASTokenPolicy } from "./tablesSASTokenPolicy";
 import { isCosmosEndpoint } from "./utils/isCosmosEndpoint";
 import { cosmosPatchPolicy } from "./cosmosPathPolicy";
+import { decodeContinuationToken, encodeContinuationToken } from "./utils/continuationToken";
 
 /**
  * A TableClient represents a Client to the Azure Tables service allowing you
@@ -438,7 +443,7 @@ export class TableClient {
   public listEntities<T extends object = Record<string, unknown>>(
     // eslint-disable-next-line @azure/azure-sdk/ts-naming-options
     options: ListTableEntitiesOptions = {}
-  ): PagedAsyncIterableIterator<TableEntityResult<T>, TableEntityResult<T>[]> {
+  ): PagedAsyncIterableIterator<TableEntityResult<T>, TableEntityResultPage<T>> {
     const tableName = this.tableName;
     const iter = this.listEntitiesAll<T>(tableName, options);
 
@@ -454,6 +459,11 @@ export class TableClient {
           ...options,
           queryOptions: { ...options.queryOptions, top: settings?.maxPageSize }
         };
+
+        if (settings?.continuationToken) {
+          pageOptions.continuationToken = settings.continuationToken;
+        }
+
         return this.listEntitiesPage(tableName, pageOptions);
       }
     };
@@ -464,13 +474,11 @@ export class TableClient {
     options?: InternalListTableEntitiesOptions
   ): AsyncIterableIterator<TableEntityResult<T>> {
     const firstPage = await this._listEntities<T>(tableName, options);
-    const { nextPartitionKey, nextRowKey } = firstPage;
     yield* firstPage;
-    if (nextRowKey && nextPartitionKey) {
+    if (firstPage.continuationToken) {
       const optionsWithContinuation: InternalListTableEntitiesOptions = {
         ...options,
-        nextPartitionKey,
-        nextRowKey
+        continuationToken: firstPage.continuationToken
       };
       for await (const page of this.listEntitiesPage<T>(tableName, optionsWithContinuation)) {
         yield* page;
@@ -489,12 +497,12 @@ export class TableClient {
 
       yield result;
 
-      while (result.nextPartitionKey && result.nextRowKey) {
+      while (result.continuationToken) {
         const optionsWithContinuation: InternalListTableEntitiesOptions = {
           ...updatedOptions,
-          nextPartitionKey: result.nextPartitionKey,
-          nextRowKey: result.nextRowKey
+          continuationToken: result.continuationToken
         };
+
         result = await this._listEntities(tableName, optionsWithContinuation);
 
         yield result;
@@ -513,27 +521,40 @@ export class TableClient {
   private async _listEntities<T extends object>(
     tableName: string,
     options: InternalListTableEntitiesOptions = {}
-  ): Promise<ListEntitiesResponse<TableEntityResult<T>>> {
+  ): Promise<TableEntityResultPage<T>> {
     const { disableTypeConversion = false } = options;
     const queryOptions = this.convertQueryOptions(options.queryOptions || {});
+    const listEntitiesOptions: TableQueryEntitiesOptionalParams = {
+      ...options,
+      queryOptions
+    };
+
+    // If a continuation token is used, decode it and set the next row and partition key
+    if (options.continuationToken) {
+      const continuationToken = decodeContinuationToken(options.continuationToken);
+      listEntitiesOptions.nextRowKey = continuationToken.nextRowKey;
+      listEntitiesOptions.nextPartitionKey = continuationToken.nextPartitionKey;
+    }
+
     const {
       xMsContinuationNextPartitionKey: nextPartitionKey,
       xMsContinuationNextRowKey: nextRowKey,
       value
-    } = await this.table.queryEntities(tableName, {
-      ...options,
-      queryOptions
-    });
+    } = await this.table.queryEntities(tableName, listEntitiesOptions);
 
     const tableEntities = deserializeObjectsArray<TableEntityResult<T>>(
       value ?? [],
       disableTypeConversion
     );
 
-    return Object.assign([...tableEntities], {
-      nextPartitionKey,
-      nextRowKey
+    // Encode nextPartitionKey and nextRowKey as a single continuation token and add it as a
+    // property to the page.
+    const continuationToken = encodeContinuationToken(nextPartitionKey, nextRowKey);
+    const page: TableEntityResultPage<T> = Object.assign([...tableEntities], {
+      continuationToken
     });
+
+    return page;
   }
 
   /**
@@ -945,11 +966,7 @@ interface InternalListTableEntitiesOptions extends ListTableEntitiesOptions {
   /**
    * An entity query continuation token from a previous call.
    */
-  nextPartitionKey?: string;
-  /**
-   * An entity query continuation token from a previous call.
-   */
-  nextRowKey?: string;
+  continuationToken?: string;
   /**
    * If true, automatic type conversion will be disabled and entity properties will
    * be represented by full metadata types. For example, an Int32 value will be \{value: "123", type: "Int32"\} instead of 123.

sdk/tables/data-tables/src/generated/generatedClientContext.ts

@@ -32,7 +32,7 @@ export class GeneratedClientContext extends coreClient.ServiceClient {
       requestContentType: "application/json; charset=utf-8"
     };
 
-    const packageDetails = `azsdk-js-data-tables/12.1.3`;
+    const packageDetails = `azsdk-js-data-tables/12.2.0`;
     const userAgentPrefix =
       options.userAgentOptions && options.userAgentOptions.userAgentPrefix
         ? `${options.userAgentOptions.userAgentPrefix} ${packageDetails}`

sdk/tables/data-tables/src/models.ts

@@ -112,6 +112,16 @@ export type TableEntityResult<T> = T & {
   timestamp?: string;
 };
 
+/**
+ * Output page type for query operations
+ */
+export type TableEntityResultPage<T> = Array<TableEntityResult<T>> & {
+  /**
+   * Continuation token to get the next page
+   */
+  continuationToken?: string;
+};
+
 /**
  * List entities optional parameters.
  */

sdk/tables/data-tables/src/utils/bufferSerializer.browser.ts

@@ -3,12 +3,16 @@
 
 /**
  * Encodes a byte array in base64 format.
- * @param value - The Uint8Aray to encode
+ * @param value - The Uint8Aray or string to encode
  */
-export function base64Encode(value: Uint8Array): string {
+export function base64Encode(value: Uint8Array | string): string {
   let str = "";
-  for (let i = 0; i < value.length; i++) {
-    str += String.fromCharCode(value[i]);
+  if (typeof value === "string") {
+    str = value;
+  } else {
+    for (let i = 0; i < value.length; i++) {
+      str += String.fromCharCode(value[i]);
+    }
   }
   return btoa(str);
 }

sdk/tables/data-tables/src/utils/bufferSerializer.ts

@@ -3,11 +3,15 @@
 
 /**
  * Encodes a byte array in base64 format.
- * @param value - The Uint8Aray to encode
+ * @param value - The Uint8Aray or string to encode
  */
-export function base64Encode(value: Uint8Array): string {
-  const bufferValue = value instanceof Buffer ? value : Buffer.from(value.buffer);
-  return bufferValue.toString("base64");
+export function base64Encode(value: Uint8Array | string): string {
+  if (value instanceof Uint8Array) {
+    const bufferValue = value instanceof Buffer ? value : Buffer.from(value.buffer);
+    return bufferValue.toString("base64");
+  } else {
+    return Buffer.from(value).toString("base64");
+  }
 }
 
 /**

sdk/tables/data-tables/src/utils/continuationToken.ts

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { base64Encode, base64Decode } from "./bufferSerializer";
+
+export interface ContinuationToken {
+  nextPartitionKey: string;
+  nextRowKey: string;
+}
+
+/**
+ * Encodes the nextPartitionKey and nextRowKey into a single continuation token
+ */
+export function encodeContinuationToken(
+  nextPartitionKey: string = "",
+  nextRowKey: string = ""
+): string | undefined {
+  if (!nextPartitionKey && !nextRowKey) {
+    return undefined;
+  }
+
+  const continuationToken = JSON.stringify({
+    nextPartitionKey,
+    nextRowKey
+  });
+
+  return base64Encode(continuationToken);
+}
+
+/**
+ * Decodes a continuationToken into an object containing a nextPartitionKey and nextRowKey
+ */
+export function decodeContinuationToken(encodedToken: string): ContinuationToken {
+  const decodedToken = base64Decode(encodedToken);
+  let tokenStr = "";
+
+  for (const byte of decodedToken) {
+    tokenStr += String.fromCharCode(byte);
+  }
+  const continuationToken: ContinuationToken = JSON.parse(tokenStr);
+
+  return continuationToken;
+}

sdk/tables/data-tables/src/utils/internalModels.ts

@@ -45,13 +45,9 @@ export type ListTableItemsResponse = Array<TableItem> & {
  */
 export type ListEntitiesResponse<T extends object> = Array<TableEntityResult<T>> & {
   /**
-   * Contains the continuation token value for partition key.
+   * Contains the continuation token value for the next page.
    */
-  nextPartitionKey?: string;
-  /**
-   * Contains the continuation token value for row key.
-   */
-  nextRowKey?: string;
+  continuationToken?: string;
 };
 
 export interface ClientParamsFromConnectionString {

sdk/tables/data-tables/swagger/README.md

@@ -6,7 +6,7 @@
 
 ```yaml
 v3: true
-package-version: 12.1.3
+package-version: 12.2.0
 package-name: "@azure/data-tables"
 title: TablesClient
 description: Tables Client