fix move to next page in paged response (#5625)

* fix move to next page in paged response

* clangs

* query tables continuation

* remove constructors

* fix with constructor

* jgjhghj

sdk/tables/azure-data-tables/inc/azure/data/tables/models.hpp

@@ -98,9 +98,6 @@ namespace Azure { namespace Data { namespace Tables {
     class QueryTablesPagedResponse final
         : public Azure::Core::PagedResponse<QueryTablesPagedResponse> {
 
-      friend class Azure::Data::Tables::TableServiceClient;
-      friend class Azure::Core::PagedResponse<QueryTablesPagedResponse>;
-
     public:
       /**
        * Service endpoint.
@@ -117,15 +114,16 @@ namespace Azure { namespace Data { namespace Tables {
        */
       std::vector<Models::Table> Tables;
 
-    public:
-      /**
-       * Table Service Client.
-       */
-      std::shared_ptr<TableServiceClient> m_tableServiceClient;
       /** Operation options */
       QueryTablesOptions m_operationOptions;
 
     private:
+      QueryTablesPagedResponse(std::shared_ptr<TableServiceClient> tableServiceClient)
+          : m_tableServiceClient(tableServiceClient){};
+
+      friend class Azure::Data::Tables::TableServiceClient;
+      friend class Azure::Core::PagedResponse<QueryTablesPagedResponse>;
+      std::shared_ptr<TableServiceClient> m_tableServiceClient;
       void OnNextPage(const Azure::Core::Context& context);
     };
 
@@ -743,6 +741,16 @@ namespace Azure { namespace Data { namespace Tables {
        *
        */
       std::string RowKey;
+      /**
+       * @brief The next Partition key.
+       *
+       */
+      std::string NextPartitionKey;
+      /**
+       * @brief The next row key.
+       *
+       */
+      std::string NextRowKey;
       /**
        * @brief The select query.
        *
@@ -774,11 +782,17 @@ namespace Azure { namespace Data { namespace Tables {
        * Table entities.
        */
       std::vector<Models::TableEntity> TableEntities;
+      /**
+       * Operation options
+       */
+      QueryEntitiesOptions m_operationOptions;
 
     private:
+      QueryEntitiesPagedResponse(std::shared_ptr<TableClient> tableClient)
+          : m_tableClient(tableClient){};
+
       std::shared_ptr<TableClient> m_tableClient;
-      QueryEntitiesOptions m_operationOptions;
-      friend class Azure::Data::Tables::TableServiceClient;
+      friend class Azure::Data::Tables::TableClient;
       friend class Azure::Core::PagedResponse<QueryEntitiesPagedResponse>;
 
       void OnNextPage(const Azure::Core::Context& context);

sdk/tables/azure-data-tables/src/tables_clients.cpp

@@ -512,6 +512,10 @@ Models::QueryTablesPagedResponse TableServiceClient::QueryTables(
   {
     request.GetUrl().AppendQueryParameter(IfMatch, options.Prefix.Value());
   }
+  if (options.ContinuationToken.HasValue())
+  {
+    request.GetUrl().AppendQueryParameter("NextTableName", options.ContinuationToken.Value());
+  }
   auto rawResponse = m_pipeline->Send(request, context);
   auto const httpStatusCode = rawResponse->GetStatusCode();
 
@@ -520,7 +524,7 @@ Models::QueryTablesPagedResponse TableServiceClient::QueryTables(
     throw Core::RequestFailedException(rawResponse);
   }
 
-  Models::QueryTablesPagedResponse response;
+  Models::QueryTablesPagedResponse response{std::make_shared<TableServiceClient>(*this)};
   {
     auto const& responseBody = rawResponse->GetBody();
     std::string responseString = std::string(responseBody.begin(), responseBody.end());
@@ -543,7 +547,6 @@ Models::QueryTablesPagedResponse TableServiceClient::QueryTables(
 
     response.ServiceEndpoint = url.GetAbsoluteUrl();
     response.Prefix = options.Prefix;
-    response.m_tableServiceClient = std::make_shared<TableServiceClient>(*this);
     response.m_operationOptions = options;
     response.CurrentPageToken = options.ContinuationToken.ValueOr(std::string());
     response.RawResponse = std::move(response.RawResponse);
@@ -817,8 +820,8 @@ Azure::Response<Models::UpsertEntityResult> TableClient::UpsertEntity(
 
 void Models::QueryEntitiesPagedResponse::OnNextPage(const Azure::Core::Context& context)
 {
-  m_operationOptions.PartitionKey = NextPartitionKey;
-  m_operationOptions.RowKey = NextRowKey;
+  m_operationOptions.NextPartitionKey = NextPartitionKey;
+  m_operationOptions.NextRowKey = NextRowKey;
   *this = m_tableClient->QueryEntities(m_operationOptions, context);
 }
 
@@ -860,18 +863,27 @@ Models::QueryEntitiesPagedResponse TableClient::QueryEntities(
     Core::Context const& context)
 {
   auto url = m_url;
-  std::string appendPath = m_tableName + "(";
-  if (!options.PartitionKey.empty())
+  if (!options.NextPartitionKey.empty() && !options.NextRowKey.empty())
   {
-    appendPath += "PartitionKey='" + Azure::Core::Url::Encode(options.PartitionKey) + "'";
+    url.AppendPath(m_tableName);
+    url.AppendQueryParameter("NextPartitionKey", options.NextPartitionKey);
+    url.AppendQueryParameter("NextRowKey", options.NextRowKey);
   }
-  if (!options.RowKey.empty())
+  else
   {
-    appendPath += ",RowKey='" + Azure::Core::Url::Encode(options.RowKey) + "'";
-  }
-  appendPath += ")";
+    std::string appendPath = m_tableName + "(";
+    if (!options.PartitionKey.empty())
+    {
+      appendPath += "PartitionKey='" + Azure::Core::Url::Encode(options.PartitionKey) + "'";
+    }
+    if (!options.RowKey.empty())
+    {
+      appendPath += ",RowKey='" + Azure::Core::Url::Encode(options.RowKey) + "'";
+    }
+    appendPath += ")";
 
-  url.AppendPath(appendPath);
+    url.AppendPath(appendPath);
+  }
 
   if (options.Filter.HasValue())
   {
@@ -892,7 +904,7 @@ Models::QueryEntitiesPagedResponse TableClient::QueryEntities(
     throw Core::RequestFailedException(rawResponse);
   }
 
-  Models::QueryEntitiesPagedResponse response{};
+  Models::QueryEntitiesPagedResponse response(std::make_shared<TableClient>(*this));
   {
     const auto& responseBody = rawResponse->GetBody();
     std::string responseString = std::string(responseBody.begin(), responseBody.end());
@@ -908,6 +920,12 @@ Models::QueryEntitiesPagedResponse TableClient::QueryEntities(
       response.NextRowKey = headers.at("x-ms-continuation-NextRowKey");
     }
 
+    if (!response.NextPartitionKey.empty() || !response.NextRowKey.empty())
+    {
+      response.NextPageToken = "true";
+    }
+
+    response.TableEntities.clear();
     auto const jsonRoot
         = Core::Json::_internal::json::parse(responseBody.begin(), responseBody.end());
 

sdk/tables/azure-data-tables/test/ut/table_client_test.cpp

@@ -506,6 +506,26 @@ namespace Azure { namespace Data { namespace Test {
     EXPECT_EQ(responseQuery.TableEntities.size(), 1);
   }
 
+  TEST_P(TablesClientTest, QueryEntityPagedResponse_LIVEONLY_)
+  {
+    auto createResponse = m_tableServiceClient->CreateTable(m_tableName);
+    for (int i = 0; i < 1010; i++)
+    {
+      auto entity = Azure::Data::Tables::Models::TableEntity();
+      entity.SetPartitionKey("partition");
+      entity.SetRowKey("rowKey" + std::to_string(i));
+      m_tableClient->AddEntity(entity);
+    }
+
+    Azure::Data::Tables::Models::QueryEntitiesOptions options;
+    auto response = m_tableClient->QueryEntities(options);
+    EXPECT_EQ(response.TableEntities.size(), 1000);
+    EXPECT_EQ(response.TableEntities[0].GetRowKey().Value, "rowKey0");
+
+    response.MoveToNextPage();
+    EXPECT_EQ(response.TableEntities.size(), 10);
+  }
+
   TEST_P(TablesClientTest, EntityGet)
   {
     Azure::Data::Tables::Models::TableEntity entity;