From 3530ec780a9eeede689114ee9734743bf5b3dcfb Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 4 May 2021 09:44:49 -0600
Subject: [PATCH 01/22] feat(firestore): Add support for Query Partitions

---
 .../lib/google/cloud/firestore/client.rb      |  19 +--
 .../cloud/firestore/collection_group.rb       | 113 ++++++++++++++++++
 .../lib/google/cloud/firestore/service.rb     |   9 ++
 3 files changed, 128 insertions(+), 13 deletions(-)
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index 737cfa4cd157..45702b5a5d5a 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -139,7 +139,7 @@ def col collection_path
         alias collection col
 
         ##
-        # Creates and returns a new Query that includes all documents in the
+        # Creates and returns a new collection group that includes all documents in the
         # database that are contained in a collection or subcollection with the
         # given collection_id.
         #
@@ -147,7 +147,7 @@ def col collection_path
         #   over. Every collection or subcollection with this ID as the last
         #   segment of its path will be included. Cannot contain a slash (`/`).
         #
-        # @return [Query] The created Query.
+        # @return [CollectionGroup] The created collection group.
         #
         # @example
         #   require "google/cloud/firestore"
@@ -155,9 +155,9 @@ def col collection_path
         #   firestore = Google::Cloud::Firestore.new
         #
         #   # Get the cities collection group query
-        #   query = firestore.col_group "cities"
+        #   col_group = firestore.col_group "cities"
         #
-        #   query.get do |city|
+        #   col_group.get do |city|
         #     puts "#{city.document_id} has #{city[:population]} residents."
         #   end
         #
@@ -166,15 +166,8 @@ def col_group collection_id
             raise ArgumentError, "Invalid collection_id: '#{collection_id}', " \
               "must not contain '/'."
           end
-          query = Google::Cloud::Firestore::V1::StructuredQuery.new(
-            from: [
-              Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
-                collection_id: collection_id, all_descendants: true
-              )
-            ]
-          )
-
-          Query.start query, service.documents_path, self
+
+          CollectionGroup.from_collection_id service.documents_path, collection_id, self
         end
         alias collection_group col_group
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
new file mode 100644
index 000000000000..c45e1b52b00c
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -0,0 +1,113 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "google/cloud/firestore/v1"
+require "google/cloud/firestore/document_reference"
+require "google/cloud/firestore/document_snapshot"
+require "google/cloud/firestore/query"
+require "google/cloud/firestore/generate"
+require "google/cloud/firestore/collection_reference_list"
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # # CollectionGroup
+      #
+      # A collection group object is used for adding documents, getting
+      # document references, and querying for documents (See {Query}).
+      #
+      # @example
+      #   require "google/cloud/firestore"
+      #
+      #   firestore = Google::Cloud::Firestore.new
+      #
+      #   # Get a collection group
+      #   cities_col = firestore.col "cities"
+      #
+      #   # Get and print all city documents
+      #   cities_col.get do |city|
+      #     puts "#{city.document_id} has #{city[:population]} residents."
+      #   end
+      #
+      class CollectionGroup < Query
+        ##
+        # @private The gRPC service object.
+        attr_accessor :service
+
+        ##
+        # Retrieves a list of document references for the documents in this
+        # collection.
+        #
+        # The document references returned may include references to "missing
+        # documents", i.e. document locations that have no document present but
+        # which contain subcollections with documents. Attempting to read such a
+        # document reference (e.g. via {DocumentReference#get}) will return
+        # a {DocumentSnapshot} whose `exists?` method returns false.
+        #
+        # @param [String] token A previously-returned page token representing
+        #   part of the larger set of results to view.
+        # @param [Integer] max Maximum number of results to return.
+        #
+        # @return [Array<DocumentReference>] An array of document references.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   # Get the cities collection group query
+        #   col_group = firestore.col_group "cities"
+        #
+        #   col_group.partitions(3).each do |query_partition|
+        #     puts query_partition.create_query
+        #   end
+        #
+        def partitions partition_count, token: nil, max: nil
+          ensure_service!
+
+          unless block_given?
+            return enum_for :partitions, partition_count
+          end
+
+          results = service.partition_query parent_path, query, partition_count
+          results.each do |result|
+            next if result.result.nil?
+            yield result # TODO: QueryPartition.from_result result, self
+          end
+        end
+
+        ##
+        # @private New Collection group object from a path.
+        def self.from_collection_id parent_path, collection_id, client
+          query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+            from: [
+              Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
+                collection_id: collection_id,
+                all_descendants: true
+              )
+            ]
+          )
+
+          new.tap do |q|
+            q.instance_variable_set :@query, query
+            q.instance_variable_set :@parent_path, parent_path
+            q.instance_variable_set :@client, client
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index afc2fec7f628..49ea3edd8d46 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -96,6 +96,15 @@ def list_collections parent, token: nil, max: nil
           )
         end
 
+        def partition_query parent, structured_query, partition_count, token: nil, max: nil
+          paged_enum = firestore.partition_query parent: parent,
+                                                 structured_query: structured_query,
+                                                 partition_count: partition_count,
+                                                 page_token: token,
+                                                 page_size: max
+          paged_enum.response
+        end
+
         def run_query path, query_grpc, transaction: nil
           run_query_req = {
             parent:           path,

From d9747de9b09c7a572c65833d3b6fe9b07a5daa55 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 18 May 2021 17:50:43 -0600
Subject: [PATCH 02/22] Add constructors to Query and CollectionReference

---
 .../lib/google/cloud/firestore/client.rb      |  1 +
 .../cloud/firestore/collection_group.rb       | 24 ++-------
 .../cloud/firestore/collection_reference.rb   | 13 +++--
 .../lib/google/cloud/firestore/query.rb       | 20 +++++---
 .../support/doctest_helper.rb                 | 29 +++++++++++
 .../cloud/firestore/client/col_group_test.rb  | 12 ++---
 .../cloud/firestore/collection_group_test.rb  | 51 +++++++++++++++++++
 google-cloud-firestore/test/helper.rb         | 28 ++++++++++
 8 files changed, 139 insertions(+), 39 deletions(-)
 create mode 100644 google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index 45702b5a5d5a..d1845719874c 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -17,6 +17,7 @@
 require "google/cloud/firestore/service"
 require "google/cloud/firestore/field_path"
 require "google/cloud/firestore/field_value"
+require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/collection_reference"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index c45e1b52b00c..a2a740c11efc 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -43,10 +43,6 @@ module Firestore
       #   end
       #
       class CollectionGroup < Query
-        ##
-        # @private The gRPC service object.
-        attr_accessor :service
-
         ##
         # Retrieves a list of document references for the documents in this
         # collection.
@@ -72,21 +68,14 @@ class CollectionGroup < Query
         #   col_group = firestore.col_group "cities"
         #
         #   col_group.partitions(3).each do |query_partition|
-        #     puts query_partition.create_query
+        #     # puts query_partition.create_query
         #   end
         #
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          unless block_given?
-            return enum_for :partitions, partition_count
-          end
-
-          results = service.partition_query parent_path, query, partition_count
-          results.each do |result|
-            next if result.result.nil?
-            yield result # TODO: QueryPartition.from_result result, self
-          end
+          resp_gapi = service.partition_query parent_path, query, partition_count, token: token, max: max
+          resp_gapi.partitions
         end
 
         ##
@@ -100,12 +89,7 @@ def self.from_collection_id parent_path, collection_id, client
               )
             ]
           )
-
-          new.tap do |q|
-            q.instance_variable_set :@query, query
-            q.instance_variable_set :@parent_path, parent_path
-            q.instance_variable_set :@client, client
-          end
+          CollectionGroup.new query, parent_path, client
         end
       end
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
index 3a8b7d3c97cf..9d724b8b7d15 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
@@ -47,6 +47,13 @@ class CollectionReference < Query
         # @private The firestore client object.
         attr_accessor :client
 
+        ##
+        # @private Creates a new CollectionReference.
+        def initialize path, client, query
+          super query, nil, client
+          @path = path
+        end
+
         ##
         # The collection identifier for the collection resource.
         #
@@ -257,11 +264,7 @@ def self.from_path path, client
             ]
           )
 
-          new.tap do |c|
-            c.client = client
-            c.instance_variable_set :@path, path
-            c.instance_variable_set :@query, query
-          end
+          CollectionReference.new path, client, query
         end
 
         protected
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 757f83363dd9..0bb0e910f8db 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -74,6 +74,16 @@ class Query
         # @private The firestore client object.
         attr_accessor :client
 
+        ##
+        # @private Creates a new Query.
+        def initialize query, parent_path, client, limit_type: nil
+          query ||= StructuredQuery.new
+          @query = query
+          @parent_path = parent_path
+          @limit_type = limit_type
+          @client = client
+        end
+
         ##
         # Restricts documents matching the query to return only data for the
         # provided fields.
@@ -965,13 +975,7 @@ def listen &callback
         ##
         # @private Start a new Query.
         def self.start query, parent_path, client, limit_type: nil
-          query ||= StructuredQuery.new
-          Query.new.tap do |q|
-            q.instance_variable_set :@query, query
-            q.instance_variable_set :@parent_path, parent_path
-            q.instance_variable_set :@limit_type, limit_type
-            q.instance_variable_set :@client, client
-          end
+          new query, parent_path, client, limit_type: limit_type
         end
 
         protected
@@ -1216,7 +1220,7 @@ def doc_id_path
         ##
         # @private Raise an error unless an database available.
         def ensure_client!
-          raise "Must have active connection to service" unless client
+          raise "Must have active connection to client" unless client
         end
 
         ##
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index baef8699aa44..2b84fe33d482 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -178,6 +178,18 @@ def mock_firestore
     end
   end
 
+  doctest.before "Google::Cloud::Firestore::CollectionGroup" do
+    mock_firestore do |mock|
+      mock.expect :run_query, run_query_resp, run_query_args
+    end
+  end
+
+  doctest.before "Google::Cloud::Firestore::CollectionGroup#partitions" do
+    mock_firestore do |mock|
+      mock.expect :partition_query, partition_query_resp, [Hash]
+    end
+  end
+
   doctest.skip "Google::Cloud::Firestore::FieldPath" do
     mock_firestore do |mock|
       mock.expect :batch_get_documents, batch_get_resp_users, batch_get_args
@@ -515,3 +527,20 @@ def documents_resp token: nil
   response.next_page_token = token if token
   paged_enum_struct response
 end
+
+def cursor_gapi
+  Google::Cloud::Firestore::V1::Cursor.new(
+    values: [
+      Google::Cloud::Firestore::V1::Value.new(string_value: "a")
+    ],
+    before: false
+  )
+end
+
+def partition_query_resp count: 3, token: nil
+  response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
+    partitions: count.times.map { cursor_gapi },
+    next_page_token: token
+  )
+  paged_enum_struct response
+end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
index 273252f45157..800961e43d47 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
@@ -19,10 +19,10 @@
   let(:collection_id_bad) { "a/b/my-collection-id" }
 
   it "creates a collection group query" do
-    query = firestore.col_group(collection_id).where "foo", "==", "bar"
+    collection_group = firestore.col_group(collection_id)
 
-    _(query).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = query.query
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
     _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
     _(query_gapi.from.size).must_equal 1
     _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
@@ -31,10 +31,10 @@
   end
 
   it "creates a collection group query using collection_group alias" do
-    query = firestore.collection_group(collection_id).where "foo", "==", "bar"
+    collection_group = firestore.collection_group(collection_id)
 
-    _(query).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = query.query
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
     _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
     _(query_gapi.from.size).must_equal 1
     _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
new file mode 100644
index 000000000000..5f78596df987
--- /dev/null
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -0,0 +1,51 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "helper"
+
+describe Google::Cloud::Firestore::CollectionGroup, :mock_firestore do
+  let(:collection_id) { "my-collection-id" }
+  let(:collection_group) do
+    Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
+  end
+
+  it "creates a query" do
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
+
+    _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
+    _(query_gapi.from.size).must_equal 1
+    _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
+    _(query_gapi.from.first.all_descendants).must_equal true
+    _(query_gapi.from.first.collection_id).must_equal collection_id
+  end
+focus
+  it "creates partitions" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [
+        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "my-collection-id", all_descendants: true)
+      ]
+    )
+    num_partitions = 3
+
+    list_res = paged_enum_struct partition_query_resp(num_partitions)
+    firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
+
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    partitions = collection_group.partitions 3
+
+    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+    _(partitions.count).must_equal 3
+  end
+end
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 56651a21a255..0d78920f379e 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -182,6 +182,34 @@ def run_query_args query,
     [req, default_options]
   end
 
+  def partition_query_args query,
+                           parent: "projects/#{project}/databases/(default)/documents",
+                           partition_count: 3
+    {
+      parent: parent,
+      structured_query: query,
+      partition_count: partition_count,
+      page_token: nil,
+      page_size: nil
+    }
+  end
+
+  def partition_query_resp count = 3, token = nil
+    Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
+      partitions: count.times.map { cursor_gapi },
+      next_page_token: token
+    )
+  end
+
+  def cursor_gapi
+    Google::Cloud::Firestore::V1::Cursor.new(
+      values: [
+        Google::Cloud::Firestore::V1::Value.new(string_value: "a")
+      ],
+      before: false
+    )
+  end
+
   def paged_enum_struct response
     OpenStruct.new response: response
   end

From 2580944db68706499f69f1b3f1ba7f0adc90d316 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Wed, 19 May 2021 17:49:31 -0600
Subject: [PATCH 03/22] Add QueryPartition and QueryPartition::List

---
 .../lib/google/cloud/firestore/client.rb      |   2 +-
 .../cloud/firestore/collection_group.rb       |  16 +-
 .../cloud/firestore/collection_reference.rb   |   2 +-
 .../google/cloud/firestore/query_partition.rb |  66 ++++++
 .../cloud/firestore/query_partition/list.rb   | 202 ++++++++++++++++++
 .../lib/google/cloud/firestore/service.rb     |   2 +
 .../support/doctest_helper.rb                 |  29 ++-
 .../cloud/firestore/collection_group_test.rb  |   4 +-
 google-cloud-firestore/test/helper.rb         |  15 +-
 9 files changed, 307 insertions(+), 31 deletions(-)
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index d1845719874c..b6e8459074d3 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -17,10 +17,10 @@
 require "google/cloud/firestore/service"
 require "google/cloud/firestore/field_path"
 require "google/cloud/firestore/field_value"
-require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/collection_reference"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
+require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/batch"
 require "google/cloud/firestore/transaction"
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index a2a740c11efc..b09e8b4f4f8e 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -14,11 +14,8 @@
 
 
 require "google/cloud/firestore/v1"
-require "google/cloud/firestore/document_reference"
-require "google/cloud/firestore/document_snapshot"
 require "google/cloud/firestore/query"
-require "google/cloud/firestore/generate"
-require "google/cloud/firestore/collection_reference_list"
+require "google/cloud/firestore/query_partition"
 
 module Google
   module Cloud
@@ -64,18 +61,19 @@ class CollectionGroup < Query
         #
         #   firestore = Google::Cloud::Firestore.new
         #
-        #   # Get the cities collection group query
         #   col_group = firestore.col_group "cities"
         #
-        #   col_group.partitions(3).each do |query_partition|
-        #     # puts query_partition.create_query
+        #   partitions = col_group.partitions 3
+        #   partitions.each do |partition|
+        #     puts partition.create_query
         #   end
         #
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          resp_gapi = service.partition_query parent_path, query, partition_count, token: token, max: max
-          resp_gapi.partitions
+          grpc = service.partition_query parent_path, query, partition_count, token: token, max: max
+
+          QueryPartition::List.from_grpc grpc, client, parent_path, query, partition_count, max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
index 9d724b8b7d15..ab8a1a377a7b 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
@@ -14,9 +14,9 @@
 
 
 require "google/cloud/firestore/v1"
+require "google/cloud/firestore/query"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
-require "google/cloud/firestore/query"
 require "google/cloud/firestore/generate"
 require "google/cloud/firestore/collection_reference_list"
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
new file mode 100644
index 000000000000..51e4fa191c1e
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -0,0 +1,66 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "google/cloud/firestore/query_partition/list"
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # # QueryPartition
+      #
+      # Represents a split point that can be used in a query as a starting and/or end point for the query results.
+      #
+      # The cursors returned by {#start_at} and {#end_before} can only be used in a query that matches the constraint of
+      # the query that produced this partition.
+      #
+      # See {CollectionGroup#partitions}.
+      #
+      # @example
+      #   require "google/cloud/firestore"
+      #
+      #   firestore = Google::Cloud::Firestore.new
+      #
+      #   col_group = firestore.col_group "cities"
+      #
+      #   partitions = col_group.partitions 3
+      #   partitions.each do |partition|
+      #     puts partition.create_query
+      #   end
+      #
+      class QueryPartition
+        attr_reader :start_at
+        attr_reader :end_before
+
+        ##
+        # @private New QueryPartition from query and Cursor
+        def initialize query, start_at, end_before
+          @query = query
+          @start_at = start_at
+          @end_before = end_before
+        end
+
+        ##
+        # Creates a query that only returns the documents for this partition.
+        #
+        # @return [Query] query.
+        #
+        def create_query
+          @query
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
new file mode 100644
index 000000000000..127861ec5a3b
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -0,0 +1,202 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "delegate"
+
+module Google
+  module Cloud
+    module Firestore
+      class QueryPartition
+        ##
+        # QueryPartition::List is a special case Array with additional values.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   col_group = firestore.col_group "cities"
+        #
+        #   partitions = col_group.partitions 3
+        #   partitions.each do |partition|
+        #     puts partition.create_query
+        #   end
+        #
+        class List < DelegateClass(::Array)
+          ##
+          # If not empty, indicates that there are more records that match
+          # the request and this value should be passed to continue.
+          attr_accessor :token
+
+          ##
+          # @private Create a new QueryPartition::List with an array of
+          # QueryPartition instances.
+          def initialize arr = []
+            super arr
+          end
+
+          ##
+          # Whether there is a next page of document references.
+          #
+          # @return [Boolean]
+          #
+          # @example
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   if partitions.next?
+          #     next_documents = partitions.next
+          #   end
+          #
+          def next?
+            !token.nil?
+          end
+
+          ##
+          # Retrieve the next page of document references.
+          #
+          # @return [QueryPartition::List]
+          #
+          # @example
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   if partitions.next?
+          #     next_documents = partitions.next
+          #   end
+          #
+          def next
+            return nil unless next?
+            ensure_client!
+            grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
+            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, @max
+          end
+
+          ##
+          # Retrieves remaining results by repeatedly invoking {#next} until
+          # {#next?} returns `false`. Calls the given block once for each
+          # result, which is passed as the argument to the block.
+          #
+          # An Enumerator is returned if no block is given.
+          #
+          # This method will make repeated API calls until all remaining results
+          # are retrieved. (Unlike `#each`, for example, which merely iterates
+          # over the results returned by a single API call.) Use with caution.
+          #
+          # @param [Integer] request_limit The upper limit of API requests to
+          #   make to load all document references. Default is no limit.
+          # @yield [document] The block for accessing each document.
+          # @yieldparam [QueryPartition] document The document reference
+          #   object.
+          #
+          # @return [Enumerator]
+          #
+          # @example Iterating each document reference by passing a block:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   partitions.all do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          # @example Using the enumerator by not passing a block:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   all_queries = partitions.all.map do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          # @example Limit the number of API calls made:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   partitions.all(request_limit: 10) do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          def all request_limit: nil, &block
+            request_limit = request_limit.to_i if request_limit
+            unless block_given?
+              return enum_for :all, request_limit: request_limit
+            end
+            results = self
+            loop do
+              results.each(&block)
+              if request_limit
+                request_limit -= 1
+                break if request_limit.negative?
+              end
+              break unless results.next?
+              results = results.next
+            end
+          end
+
+          ##
+          # @private New QueryPartition::List from a
+          # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
+          def self.from_grpc grpc, client, parent, structured_query, partition_count, max: nil
+            start_at = nil
+            partitions = List.new(Array(grpc.partitions).map do |cursor|
+              end_before = cursor.values.map do |value|
+                Convert.value_to_raw value, client
+              end
+              partition = QueryPartition.new structured_query, start_at, end_before
+              start_at = end_before
+              partition
+            end)
+            partitions.instance_variable_set :@parent, parent
+            partitions.instance_variable_set :@structured_query, structured_query
+            partitions.instance_variable_set :@partition_count, partition_count
+            token = grpc.next_page_token
+            token = nil if token == ""
+            partitions.instance_variable_set :@token, token
+            partitions.instance_variable_set :@client, client
+            partitions.instance_variable_set :@max, max
+            partitions
+          end
+
+          protected
+
+          ##
+          # Raise an error unless an active client is available.
+          def ensure_client!
+            raise "Must have active connection" unless @client
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index 49ea3edd8d46..6417ac4d0e25 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -96,6 +96,8 @@ def list_collections parent, token: nil, max: nil
           )
         end
 
+        ##
+        # Returns Google::Cloud::Firestore::V1::PartitionQueryResponse
         def partition_query parent, structured_query, partition_count, token: nil, max: nil
           paged_enum = firestore.partition_query parent: parent,
                                                  structured_query: structured_query,
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index 2b84fe33d482..cd380e3c3123 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -298,6 +298,12 @@ def mock_firestore
     end
   end
 
+  doctest.before "Google::Cloud::Firestore::QueryPartition" do
+    mock_firestore do |mock|
+      mock.expect :partition_query, partition_query_resp, [Hash]
+    end
+  end
+
   doctest.before "Google::Cloud::Firestore::CollectionReference" do
     mock_firestore do |mock|
       mock.expect :run_query, run_query_resp, run_query_args
@@ -528,19 +534,20 @@ def documents_resp token: nil
   paged_enum_struct response
 end
 
-def cursor_gapi
-  Google::Cloud::Firestore::V1::Cursor.new(
-    values: [
-      Google::Cloud::Firestore::V1::Value.new(string_value: "a")
-    ],
-    before: false
-  )
-end
-
 def partition_query_resp count: 3, token: nil
   response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-    partitions: count.times.map { cursor_gapi },
-    next_page_token: token
+    partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] }
   )
+  response.next_page_token = token if token
   paged_enum_struct response
 end
+
+def cursor_grpc values: [1,2], before: false
+  converted_values = values.map do |val|
+    Google::Cloud::Firestore::Convert.raw_to_value val
+  end
+  Google::Cloud::Firestore::V1::Cursor.new(
+    values: converted_values,
+    before: before
+  )
+end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index 5f78596df987..39022b622bd0 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -30,7 +30,7 @@
     _(query_gapi.from.first.all_descendants).must_equal true
     _(query_gapi.from.first.collection_id).must_equal collection_id
   end
-focus
+
   it "creates partitions" do
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [
@@ -39,7 +39,7 @@
     )
     num_partitions = 3
 
-    list_res = paged_enum_struct partition_query_resp(num_partitions)
+    list_res = paged_enum_struct partition_query_resp(count: num_partitions)
     firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
 
     _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 0d78920f379e..3ba1b94e7c32 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -194,19 +194,20 @@ def partition_query_args query,
     }
   end
 
-  def partition_query_resp count = 3, token = nil
+  def partition_query_resp count: 3, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      partitions: count.times.map { cursor_gapi },
+      partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
       next_page_token: token
     )
   end
 
-  def cursor_gapi
+  def cursor_grpc values: [1,2], before: false
+    converted_values = values.map do |val|
+      Google::Cloud::Firestore::Convert.raw_to_value val
+    end
     Google::Cloud::Firestore::V1::Cursor.new(
-      values: [
-        Google::Cloud::Firestore::V1::Value.new(string_value: "a")
-      ],
-      before: false
+      values: converted_values,
+      before: before
     )
   end
 

From ce6aed97bf03cfbe866489def88cee47e1214a4d Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Fri, 21 May 2021 15:56:50 -0600
Subject: [PATCH 04/22] Add order by to CollectionGroup#partitions

---
 .../acceptance/firestore/query_test.rb        |  97 ------------
 .../acceptance/firestore/watch_test.rb        |   2 +-
 .../cloud/firestore/collection_group.rb       |  10 +-
 .../cloud/firestore/query_partition/list.rb   |   2 +-
 .../cloud/firestore/collection_group_test.rb  | 138 +++++++++++++++---
 .../list_documents_test.rb                    |  15 --
 google-cloud-firestore/test/helper.rb         |  22 +--
 7 files changed, 142 insertions(+), 144 deletions(-)

diff --git a/google-cloud-firestore/acceptance/firestore/query_test.rb b/google-cloud-firestore/acceptance/firestore/query_test.rb
index 9e318a7e1dd8..1fdf6e4d6e03 100644
--- a/google-cloud-firestore/acceptance/firestore/query_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/query_test.rb
@@ -284,101 +284,4 @@
     _(results.map(&:document_id)).must_equal ["doc1", "doc2"]
     _(results.map { |doc| doc[:foo] }).must_equal ["a", "b"]
   end
-
-  describe "Collection Group" do
-    it "queries a collection group" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-          "abc/123/#{collection_group}/cg-doc1",
-          "abc/123/#{collection_group}/cg-doc2",
-          "#{collection_group}/cg-doc3",
-          "#{collection_group}/cg-doc4",
-          "def/456/#{collection_group}/cg-doc5",
-          "#{collection_group}/virtual-doc/nested-coll/not-cg-doc",
-          "x#{collection_group}/not-cg-doc",
-          "#{collection_group}x/not-cg-doc",
-          "abc/123/#{collection_group}x/not-cg-doc",
-          "abc/123/x#{collection_group}/not-cg-doc",
-          "abc/#{collection_group}"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group collection_group
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc1", "cg-doc2", "cg-doc3", "cg-doc4", "cg-doc5"]
-    end
-
-    it "queries a collection group with start_at and end_at" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-        "a/a/#{collection_group}/cg-doc1",
-        "a/b/a/b/#{collection_group}/cg-doc2",
-        "a/b/#{collection_group}/cg-doc3",
-        "a/b/c/d/#{collection_group}/cg-doc4",
-        "a/c/#{collection_group}/cg-doc5",
-        "#{collection_group}/cg-doc6",
-        "a/b/nope/nope"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group(collection_group)
-        .order_by("__name__")
-        .start_at(firestore.document("a/b"))
-        .end_at(firestore.document("a/b0"))
-
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
-
-      query = firestore.collection_group(collection_group)
-        .order_by("__name__")
-        .start_after(firestore.document("a/b"))
-        .end_before(firestore.document("a/b/#{collection_group}/cg-doc3"))
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
-    end
-
-    it "queries a collection group with filters" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-        "a/a/#{collection_group}/cg-doc1",
-        "a/b/a/b/#{collection_group}/cg-doc2",
-        "a/b/#{collection_group}/cg-doc3",
-        "a/b/c/d/#{collection_group}/cg-doc4",
-        "a/c/#{collection_group}/cg-doc5",
-        "#{collection_group}/cg-doc6",
-        "a/b/nope/nope"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group(collection_group)
-        .where("__name__", ">=", firestore.document("a/b"))
-        .where("__name__", "<=", firestore.document("a/b0"))
-
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
-
-      query = firestore.collection_group(collection_group)
-        .where("__name__", ">", firestore.document("a/b"))
-        .where(
-          "__name__", "<", firestore.document("a/b/#{collection_group}/cg-doc3")
-        )
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
-    end
-  end
 end
diff --git a/google-cloud-firestore/acceptance/firestore/watch_test.rb b/google-cloud-firestore/acceptance/firestore/watch_test.rb
index c8b77cc6b49e..b385f641847c 100644
--- a/google-cloud-firestore/acceptance/firestore/watch_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/watch_test.rb
@@ -188,7 +188,7 @@
   def wait_until &block
     wait_count = 0
     until block.call
-      fail "wait_until criterial was not met" if wait_count > 200
+      fail "wait_until criteria was not met" if wait_count > 200
       wait_count += 1
       sleep 0.01
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index b09e8b4f4f8e..af724c196612 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -71,9 +71,15 @@ class CollectionGroup < Query
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          grpc = service.partition_query parent_path, query, partition_count, token: token, max: max
+          # Partition queries require explicit ordering by __name__.
+          query_with_default_order = order("__name__").query
+          # Since we are always returning an extra partition (with en empty endBefore cursor), we reduce the desired
+          # partition count by one.
+          partition_count -= 1
 
-          QueryPartition::List.from_grpc grpc, client, parent_path, query, partition_count, max: max
+          grpc = service.partition_query parent_path, query_with_default_order, partition_count, token: token, max: max
+
+          QueryPartition::List.from_grpc grpc, client, parent_path, query_with_default_order, partition_count, max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index 127861ec5a3b..3bc8dd5259f2 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -89,7 +89,7 @@ def next
             return nil unless next?
             ensure_client!
             grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, @max
+            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, max: @max
           end
 
           ##
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index 39022b622bd0..cf0507da5fd5 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -19,33 +19,133 @@
   let(:collection_group) do
     Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
   end
-
-  it "creates a query" do
-    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = collection_group.query
-
-    _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
-    _(query_gapi.from.size).must_equal 1
-    _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
-    _(query_gapi.from.first.all_descendants).must_equal true
-    _(query_gapi.from.first.collection_id).must_equal collection_id
-  end
-
-  it "creates partitions" do
-    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+  let(:expected_query) do
+    Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [
-        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "my-collection-id", all_descendants: true)
+        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
+          collection_id: "my-collection-id",
+          all_descendants: true
+        )
+      ],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "__name__"),
+          direction: :ASCENDING
+        )
       ]
     )
-    num_partitions = 3
+  end
 
-    list_res = paged_enum_struct partition_query_resp(count: num_partitions)
-    firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
+  it "lists partitions" do
+    list_res = paged_enum_struct partition_query_resp(count: 3)
+    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
 
-    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
     partitions = collection_group.partitions 3
 
+    firestore_mock.verify
+
+    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+    _(partitions.count).must_equal 3
+    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+  end
+
+  it "paginates partitions" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    first_partitions = collection_group.partitions 6
+    second_partitions = collection_group.partitions 6, token: first_partitions.token
+
+    firestore_mock.verify
+
+    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(first_partitions.count).must_equal 3
+    _(first_partitions.token).must_equal "next_page_token"
+
+    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(second_partitions.count).must_equal 2
+    _(second_partitions.token).must_be :nil?
+  end
+
+  it "paginates partitions using and next" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    first_partitions = collection_group.partitions 6
+    second_partitions = first_partitions.next
+
+    firestore_mock.verify
+
+    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(first_partitions.count).must_equal 3
+    _(first_partitions.token).must_equal "next_page_token"
+
+    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(second_partitions.count).must_equal 2
+    _(second_partitions.token).must_be :nil?
+  end
+
+  it "paginates partitions using all" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all.to_a
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 5
+  end
+
+  it "paginates partitions using all and Enumerator" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all.take 5
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 5
+  end
+
+  it "paginates partitions using all with request_limit set" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all(request_limit: 1).to_a
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 6
+  end
+
+  it "paginates partitions with max set" do
+    list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+
+    partitions = collection_group.partitions 6, max: 3
+
+    firestore_mock.verify
+
     _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
     _(partitions.count).must_equal 3
+    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
   end
 end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
index a73a097f7e78..9de697a5bf41 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
@@ -137,21 +137,6 @@
     _(documents.token).must_equal "next_page_token"
   end
 
-  it "paginates documents without max set" do
-    list_res = paged_enum_struct list_documents_gapi(3, "next_page_token")
-
-    firestore_mock.expect :list_documents, list_res, list_documents_args
-
-    documents = collection.list_documents
-
-    firestore_mock.verify
-
-    documents.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::DocumentReference }
-    _(documents.count).must_equal 3
-    _(documents.token).wont_be :nil?
-    _(documents.token).must_equal "next_page_token"
-  end
-
   def document_gapi
     Google::Cloud::Firestore::V1::Document.new(
       name: "projects/#{project}/databases/(default)/documents/my-document",
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 3ba1b94e7c32..637ab76556d0 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -127,7 +127,7 @@ class MockFirestore < Minitest::Spec
   def wait_until &block
     wait_count = 0
     until block.call
-      fail "wait_until criterial was not met" if wait_count > 100
+      fail "wait_until criteria was not met" if wait_count > 100
       wait_count += 1
       sleep 0.01
     end
@@ -184,14 +184,18 @@ def run_query_args query,
 
   def partition_query_args query,
                            parent: "projects/#{project}/databases/(default)/documents",
-                           partition_count: 3
-    {
-      parent: parent,
-      structured_query: query,
-      partition_count: partition_count,
-      page_token: nil,
-      page_size: nil
-    }
+                           partition_count: 2,
+                           page_token: nil,
+                           page_size: nil
+    [
+      {
+        parent: parent,
+        structured_query: query,
+        partition_count: partition_count,
+        page_token: page_token,
+        page_size: page_size
+      }
+    ]
   end
 
   def partition_query_resp count: 3, token: nil

From 3895e3f8ded3c4ec43d646afedc336454b05b340 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Fri, 21 May 2021 18:03:03 -0600
Subject: [PATCH 05/22] Fix CollectionGroup#partitions and
 QueryPartition::List.from_grpc

---
 .../firestore/collection_group_test.rb        | 133 ++++++++++++++++++
 .../acceptance/firestore_helper.rb            |  20 +--
 .../cloud/firestore/collection_group.rb       |  15 +-
 .../google/cloud/firestore/query_partition.rb |   5 +-
 .../cloud/firestore/query_partition/list.rb   |  13 +-
 .../lib/google/cloud/firestore/service.rb     |  15 +-
 .../support/doctest_helper.rb                 |   4 +-
 .../cloud/firestore/collection_group_test.rb  |  36 ++---
 google-cloud-firestore/test/helper.rb         |  14 +-
 9 files changed, 207 insertions(+), 48 deletions(-)
 create mode 100644 google-cloud-firestore/acceptance/firestore/collection_group_test.rb

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
new file mode 100644
index 000000000000..dda44fc346f3
--- /dev/null
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -0,0 +1,133 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "firestore_helper"
+
+describe Google::Cloud::Firestore::CollectionGroup, :firestore_acceptance do
+  describe "#get" do
+    it "queries a collection group" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+          "abc/123/#{collection_id}/cg-doc1",
+          "abc/123/#{collection_id}/cg-doc2",
+          "#{collection_id}/cg-doc3",
+          "#{collection_id}/cg-doc4",
+          "def/456/#{collection_id}/cg-doc5",
+          "#{collection_id}/virtual-doc/nested-coll/not-cg-doc",
+          "x#{collection_id}/not-cg-doc",
+          "#{collection_id}x/not-cg-doc",
+          "abc/123/#{collection_id}x/not-cg-doc",
+          "abc/123/x#{collection_id}/not-cg-doc",
+          "abc/#{collection_id}"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc1", "cg-doc2", "cg-doc3", "cg-doc4", "cg-doc5"]
+    end
+
+    it "queries a collection group with start_at and end_at" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+        "a/a/#{collection_id}/cg-doc1",
+        "a/b/a/b/#{collection_id}/cg-doc2",
+        "a/b/#{collection_id}/cg-doc3",
+        "a/b/c/d/#{collection_id}/cg-doc4",
+        "a/c/#{collection_id}/cg-doc5",
+        "#{collection_id}/cg-doc6",
+        "a/b/nope/nope"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+        .order_by("__name__")
+        .start_at(firestore.document("a/b"))
+        .end_at(firestore.document("a/b0"))
+
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
+
+      collection_group = firestore.collection_group(collection_id)
+        .order_by("__name__")
+        .start_after(firestore.document("a/b"))
+        .end_before(firestore.document("a/b/#{collection_id}/cg-doc3"))
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
+    end
+
+    it "queries a collection group with filters" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+        "a/a/#{collection_id}/cg-doc1",
+        "a/b/a/b/#{collection_id}/cg-doc2",
+        "a/b/#{collection_id}/cg-doc3",
+        "a/b/c/d/#{collection_id}/cg-doc4",
+        "a/c/#{collection_id}/cg-doc5",
+        "#{collection_id}/cg-doc6",
+        "a/b/nope/nope"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+        .where("__name__", ">=", firestore.document("a/b"))
+        .where("__name__", "<=", firestore.document("a/b0"))
+
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
+
+      collection_group = firestore.collection_group(collection_id)
+        .where("__name__", ">", firestore.document("a/b"))
+        .where(
+          "__name__", "<", firestore.document("a/b/#{collection_id}/cg-doc3")
+        )
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
+    end
+  end
+
+  describe "#partitions" do
+    it "queries a collection group with partitions" do
+      document_count = 2 * 128 + 127 # Minimum partition size is 128.
+      rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
+      firestore.batch do |b|
+        document_count.times do |i|
+          doc_ref = rand_col.document i
+          b.set doc_ref, {foo: i}
+        end
+      end
+
+      collection_group = firestore.collection_group(rand_col.collection_id)
+
+      partitions = collection_group.partitions 3
+      _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+      _(partitions.count).must_equal 3
+    end
+  end
+end
diff --git a/google-cloud-firestore/acceptance/firestore_helper.rb b/google-cloud-firestore/acceptance/firestore_helper.rb
index af21a6405b82..4b9d678dd745 100644
--- a/google-cloud-firestore/acceptance/firestore_helper.rb
+++ b/google-cloud-firestore/acceptance/firestore_helper.rb
@@ -65,16 +65,16 @@ def root_col
       addl.include? :firestore_acceptance
     end
 
-    def self.run_one_method klass, method_name, reporter
-      result = nil
-      reporter.prerecord klass, method_name
-      (1..3).each do |try|
-        result = Minitest.run_one_method(klass, method_name)
-        break if (result.passed? || result.skipped?)
-        puts "Retrying #{klass}##{method_name} (#{try})"
-      end
-      reporter.record result
-    end
+    # def self.run_one_method klass, method_name, reporter
+    #   result = nil
+    #   reporter.prerecord klass, method_name
+    #   (1..3).each do |try|
+    #     result = Minitest.run_one_method(klass, method_name)
+    #     break if (result.passed? || result.skipped?)
+    #     puts "Retrying #{klass}##{method_name} (#{try})"
+    #   end
+    #   reporter.record result
+    # end
   end
 end
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index af724c196612..efa21947d922 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -72,14 +72,23 @@ def partitions partition_count, token: nil, max: nil
           ensure_service!
 
           # Partition queries require explicit ordering by __name__.
-          query_with_default_order = order("__name__").query
+          query_with_default_order = order "__name__"
           # Since we are always returning an extra partition (with en empty endBefore cursor), we reduce the desired
           # partition count by one.
           partition_count -= 1
 
-          grpc = service.partition_query parent_path, query_with_default_order, partition_count, token: token, max: max
+          grpc = service.partition_query parent_path,
+                                         query_with_default_order.query,
+                                         partition_count,
+                                         token: token,
+                                         max: max
 
-          QueryPartition::List.from_grpc grpc, client, parent_path, query_with_default_order, partition_count, max: max
+          QueryPartition::List.from_grpc grpc,
+                                         client,
+                                         parent_path,
+                                         query_with_default_order,
+                                         partition_count,
+                                         max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 51e4fa191c1e..071a8cb71eb4 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -58,7 +58,10 @@ def initialize query, start_at, end_before
         # @return [Query] query.
         #
         def create_query
-          @query
+          base_query = @query
+          base_query = base_query.start_at start_at if start_at
+          base_query = base_query.end_before end_before if end_before
+          base_query
         end
       end
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index 3bc8dd5259f2..a068726f6be7 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -88,8 +88,8 @@ def next?
           def next
             return nil unless next?
             ensure_client!
-            grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, max: @max
+            grpc = @client.service.partition_query @parent, @query.query, @partition_count, token: token, max: @max
+            self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max
           end
 
           ##
@@ -167,18 +167,21 @@ def all request_limit: nil, &block
           ##
           # @private New QueryPartition::List from a
           # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
-          def self.from_grpc grpc, client, parent, structured_query, partition_count, max: nil
+          def self.from_grpc grpc, client, parent, query, partition_count, max: nil
+            # TODO: Should the logic adding the last partition be applied to the entire result set, not the page?
             start_at = nil
             partitions = List.new(Array(grpc.partitions).map do |cursor|
               end_before = cursor.values.map do |value|
                 Convert.value_to_raw value, client
               end
-              partition = QueryPartition.new structured_query, start_at, end_before
+              partition = QueryPartition.new query, start_at, end_before
               start_at = end_before
               partition
             end)
+            # We are always returning an extra partition (with en empty endBefore cursor).
+            partitions << QueryPartition.new(query, start_at, nil)
             partitions.instance_variable_set :@parent, parent
-            partitions.instance_variable_set :@structured_query, structured_query
+            partitions.instance_variable_set :@query, query
             partitions.instance_variable_set :@partition_count, partition_count
             token = grpc.next_page_token
             token = nil if token == ""
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index 6417ac4d0e25..54c4acf7d77b 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -98,12 +98,15 @@ def list_collections parent, token: nil, max: nil
 
         ##
         # Returns Google::Cloud::Firestore::V1::PartitionQueryResponse
-        def partition_query parent, structured_query, partition_count, token: nil, max: nil
-          paged_enum = firestore.partition_query parent: parent,
-                                                 structured_query: structured_query,
-                                                 partition_count: partition_count,
-                                                 page_token: token,
-                                                 page_size: max
+        def partition_query parent, query_grpc, partition_count, token: nil, max: nil
+          request = Google::Cloud::Firestore::V1::PartitionQueryRequest.new(
+            parent: parent,
+            structured_query: query_grpc,
+            partition_count: partition_count,
+            page_token: token,
+            page_size: max
+          )
+          paged_enum = firestore.partition_query request
           paged_enum.response
         end
 
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index cd380e3c3123..4b2e216a7fca 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -186,7 +186,7 @@ def mock_firestore
 
   doctest.before "Google::Cloud::Firestore::CollectionGroup#partitions" do
     mock_firestore do |mock|
-      mock.expect :partition_query, partition_query_resp, [Hash]
+      mock.expect :partition_query, partition_query_resp, [Google::Cloud::Firestore::V1::PartitionQueryRequest]
     end
   end
 
@@ -300,7 +300,7 @@ def mock_firestore
 
   doctest.before "Google::Cloud::Firestore::QueryPartition" do
     mock_firestore do |mock|
-      mock.expect :partition_query, partition_query_resp, [Hash]
+      mock.expect :partition_query, partition_query_resp, [Google::Cloud::Firestore::V1::PartitionQueryRequest]
     end
   end
 
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index cf0507da5fd5..f66859960d06 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -37,7 +37,10 @@
   end
 
   it "lists partitions" do
-    list_res = paged_enum_struct partition_query_resp(count: 3)
+    # grpc.partitions: 
+    # [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/32">], before: false>,
+    #  <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/89">], before: false>]
+    list_res = paged_enum_struct partition_query_resp
     firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
 
     partitions = collection_group.partitions 3
@@ -46,12 +49,17 @@
 
     _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
     _(partitions.count).must_equal 3
-    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+
+    _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[0].start_at).must_be :nil?
+    _(partitions[0].end_before).must_equal [0,1]
+    query_1 = partitions[0].create_query
+    # TODO: query assertions
   end
 
   it "paginates partitions" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+    first_list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -70,9 +78,9 @@
     _(second_partitions.token).must_be :nil?
   end
 
-  it "paginates partitions using and next" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+  it "paginates partitions using next" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -80,8 +88,6 @@
     first_partitions = collection_group.partitions 6
     second_partitions = first_partitions.next
 
-    firestore_mock.verify
-
     first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
     _(first_partitions.count).must_equal 3
     _(first_partitions.token).must_equal "next_page_token"
@@ -92,8 +98,8 @@
   end
 
   it "paginates partitions using all" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -122,22 +128,20 @@
   end
 
   it "paginates partitions using all with request_limit set" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2, token: "second_page_token")
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
 
     all_partitions = collection_group.partitions(6).all(request_limit: 1).to_a
 
-    firestore_mock.verify
-
     all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
     _(all_partitions.count).must_equal 6
   end
 
   it "paginates partitions with max set" do
-    list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
     firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
 
     partitions = collection_group.partitions 6, max: 3
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 637ab76556d0..cad0a5cbdf1b 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -182,23 +182,27 @@ def run_query_args query,
     [req, default_options]
   end
 
-  def partition_query_args query,
+  # expected partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>"next_page_token", :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "second_page_token">>
+  # got     [partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>nil, :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "next_page_token">>]
+
+
+  def partition_query_args query_grpc,
                            parent: "projects/#{project}/databases/(default)/documents",
                            partition_count: 2,
                            page_token: nil,
                            page_size: nil
     [
-      {
+      Google::Cloud::Firestore::V1::PartitionQueryRequest.new(
         parent: parent,
-        structured_query: query,
+        structured_query: query_grpc,
         partition_count: partition_count,
         page_token: page_token,
         page_size: page_size
-      }
+      )
     ]
   end
 
-  def partition_query_resp count: 3, token: nil
+  def partition_query_resp count: 2, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
       partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
       next_page_token: token

From 3f02954be8448a90f8b9b177dea0bb467dfceb99 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Mon, 24 May 2021 16:30:36 -0600
Subject: [PATCH 06/22] Fix array support in Query#start_at, #start_after,
 #end_before and #end_at

---
 .../lib/google/cloud/firestore/query.rb       |   6 +-
 .../support/doctest_helper.rb                 |   6 +-
 .../cloud/firestore/collection_group_test.rb  |   2 +-
 .../cloud/firestore/query/cursors_test.rb     | 101 +++++++++++++++++-
 google-cloud-firestore/test/helper.rb         |   6 +-
 5 files changed, 113 insertions(+), 8 deletions(-)

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 0bb0e910f8db..5bd38475efa4 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -1098,11 +1098,14 @@ def values_to_cursor values, query
             return snapshot_to_cursor values.first, query
           end
 
+          # The *values param in start_at, start_after, etc. will wrap an array argument in an array, so unwrap it here.
+          values = values.first if values.count == 1 && values.first.is_a?(Array)
+
           # pair values with their field_paths to ensure correct formatting
           order_field_paths = order_by_field_paths query
           if values.count > order_field_paths.count
             # raise if too many values provided for the cursor
-            raise ArgumentError, "too many values"
+            raise ArgumentError, "There cannot be more values than Order By fields"
           end
 
           values = values.zip(order_field_paths).map do |value, field_path|
@@ -1134,7 +1137,6 @@ def snapshot_to_cursor snapshot, query
               snapshot[field_path]
             end
           end
-
           values_to_cursor values, query
         end
 
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index 4b2e216a7fca..5d954e206195 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -536,13 +536,15 @@ def documents_resp token: nil
 
 def partition_query_resp count: 3, token: nil
   response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-    partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] }
+    # Minimum partition size is 128.
+    partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] }
   )
   response.next_page_token = token if token
   paged_enum_struct response
 end
 
-def cursor_grpc values: [1,2], before: false
+# Minimum partition size is 128.
+def cursor_grpc values: ["128"], before: false
   converted_values = values.map do |val|
     Google::Cloud::Firestore::Convert.raw_to_value val
   end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index f66859960d06..fc9e549a9e53 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -52,7 +52,7 @@
 
     _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
     _(partitions[0].start_at).must_be :nil?
-    _(partitions[0].end_before).must_equal [0,1]
+    _(partitions[0].end_before).must_equal ["128"]
     query_1 = partitions[0].create_query
     # TODO: query assertions
   end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb b/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
index a9346b8c3e2c..61648762d57a 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
@@ -305,7 +305,40 @@
     _(generated_query).must_equal expected_query
   end
 
-  it "Start/End with two values" do
+  it "StartAt/EndBefore with two values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: true
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: true
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_at(7, 8).end_before(9, 10).query
+    _(generated_query).must_equal expected_query
+  end
+
+  it "StartAfter/EndAt with two values" do
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
       order_by: [
@@ -338,6 +371,72 @@
     _(generated_query).must_equal expected_query
   end
 
+  it "StartAt/EndBefore with array values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: true
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: true
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_at([7, 8]).end_before([9, 10]).query
+    _(generated_query).must_equal expected_query
+  end
+
+  it "StartAfter/EndAt with array values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: false
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: false
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_after([7, 8]).end_at([9, 10]).query
+    _(generated_query).must_equal expected_query
+  end
+
   it "with __name__" do
     # TODO: Looks like we need to create the full path when paired with __name__
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index cad0a5cbdf1b..df251ad9fb41 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -204,12 +204,14 @@ def partition_query_args query_grpc,
 
   def partition_query_resp count: 2, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
+      # Minimum partition size is 128.
+      partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] },
       next_page_token: token
     )
   end
 
-  def cursor_grpc values: [1,2], before: false
+  # Minimum partition size is 128.
+  def cursor_grpc values: ["128"], before: false
     converted_values = values.map do |val|
       Google::Cloud::Firestore::Convert.raw_to_value val
     end

From 9d566d0c84460e5776525325a71f92c03829dcf1 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 4 May 2021 09:44:49 -0600
Subject: [PATCH 07/22] feat(firestore): Add support for Query Partitions

---
 .../lib/google/cloud/firestore/client.rb      |  19 +--
 .../cloud/firestore/collection_group.rb       | 113 ++++++++++++++++++
 .../lib/google/cloud/firestore/service.rb     |   9 ++
 3 files changed, 128 insertions(+), 13 deletions(-)
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index 737cfa4cd157..45702b5a5d5a 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -139,7 +139,7 @@ def col collection_path
         alias collection col
 
         ##
-        # Creates and returns a new Query that includes all documents in the
+        # Creates and returns a new collection group that includes all documents in the
         # database that are contained in a collection or subcollection with the
         # given collection_id.
         #
@@ -147,7 +147,7 @@ def col collection_path
         #   over. Every collection or subcollection with this ID as the last
         #   segment of its path will be included. Cannot contain a slash (`/`).
         #
-        # @return [Query] The created Query.
+        # @return [CollectionGroup] The created collection group.
         #
         # @example
         #   require "google/cloud/firestore"
@@ -155,9 +155,9 @@ def col collection_path
         #   firestore = Google::Cloud::Firestore.new
         #
         #   # Get the cities collection group query
-        #   query = firestore.col_group "cities"
+        #   col_group = firestore.col_group "cities"
         #
-        #   query.get do |city|
+        #   col_group.get do |city|
         #     puts "#{city.document_id} has #{city[:population]} residents."
         #   end
         #
@@ -166,15 +166,8 @@ def col_group collection_id
             raise ArgumentError, "Invalid collection_id: '#{collection_id}', " \
               "must not contain '/'."
           end
-          query = Google::Cloud::Firestore::V1::StructuredQuery.new(
-            from: [
-              Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
-                collection_id: collection_id, all_descendants: true
-              )
-            ]
-          )
-
-          Query.start query, service.documents_path, self
+
+          CollectionGroup.from_collection_id service.documents_path, collection_id, self
         end
         alias collection_group col_group
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
new file mode 100644
index 000000000000..c45e1b52b00c
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -0,0 +1,113 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "google/cloud/firestore/v1"
+require "google/cloud/firestore/document_reference"
+require "google/cloud/firestore/document_snapshot"
+require "google/cloud/firestore/query"
+require "google/cloud/firestore/generate"
+require "google/cloud/firestore/collection_reference_list"
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # # CollectionGroup
+      #
+      # A collection group object is used for adding documents, getting
+      # document references, and querying for documents (See {Query}).
+      #
+      # @example
+      #   require "google/cloud/firestore"
+      #
+      #   firestore = Google::Cloud::Firestore.new
+      #
+      #   # Get a collection group
+      #   cities_col = firestore.col "cities"
+      #
+      #   # Get and print all city documents
+      #   cities_col.get do |city|
+      #     puts "#{city.document_id} has #{city[:population]} residents."
+      #   end
+      #
+      class CollectionGroup < Query
+        ##
+        # @private The gRPC service object.
+        attr_accessor :service
+
+        ##
+        # Retrieves a list of document references for the documents in this
+        # collection.
+        #
+        # The document references returned may include references to "missing
+        # documents", i.e. document locations that have no document present but
+        # which contain subcollections with documents. Attempting to read such a
+        # document reference (e.g. via {DocumentReference#get}) will return
+        # a {DocumentSnapshot} whose `exists?` method returns false.
+        #
+        # @param [String] token A previously-returned page token representing
+        #   part of the larger set of results to view.
+        # @param [Integer] max Maximum number of results to return.
+        #
+        # @return [Array<DocumentReference>] An array of document references.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   # Get the cities collection group query
+        #   col_group = firestore.col_group "cities"
+        #
+        #   col_group.partitions(3).each do |query_partition|
+        #     puts query_partition.create_query
+        #   end
+        #
+        def partitions partition_count, token: nil, max: nil
+          ensure_service!
+
+          unless block_given?
+            return enum_for :partitions, partition_count
+          end
+
+          results = service.partition_query parent_path, query, partition_count
+          results.each do |result|
+            next if result.result.nil?
+            yield result # TODO: QueryPartition.from_result result, self
+          end
+        end
+
+        ##
+        # @private New Collection group object from a path.
+        def self.from_collection_id parent_path, collection_id, client
+          query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+            from: [
+              Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
+                collection_id: collection_id,
+                all_descendants: true
+              )
+            ]
+          )
+
+          new.tap do |q|
+            q.instance_variable_set :@query, query
+            q.instance_variable_set :@parent_path, parent_path
+            q.instance_variable_set :@client, client
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index afc2fec7f628..49ea3edd8d46 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -96,6 +96,15 @@ def list_collections parent, token: nil, max: nil
           )
         end
 
+        def partition_query parent, structured_query, partition_count, token: nil, max: nil
+          paged_enum = firestore.partition_query parent: parent,
+                                                 structured_query: structured_query,
+                                                 partition_count: partition_count,
+                                                 page_token: token,
+                                                 page_size: max
+          paged_enum.response
+        end
+
         def run_query path, query_grpc, transaction: nil
           run_query_req = {
             parent:           path,

From 47425a7d3593f07b7bc6cf447953a001c84121e2 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 18 May 2021 17:50:43 -0600
Subject: [PATCH 08/22] Add constructors to Query and CollectionReference

---
 .../lib/google/cloud/firestore/client.rb      |  1 +
 .../cloud/firestore/collection_group.rb       | 24 ++-------
 .../cloud/firestore/collection_reference.rb   | 13 +++--
 .../lib/google/cloud/firestore/query.rb       | 20 +++++---
 .../support/doctest_helper.rb                 | 29 +++++++++++
 .../cloud/firestore/client/col_group_test.rb  | 12 ++---
 .../cloud/firestore/collection_group_test.rb  | 51 +++++++++++++++++++
 google-cloud-firestore/test/helper.rb         | 28 ++++++++++
 8 files changed, 139 insertions(+), 39 deletions(-)
 create mode 100644 google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index 45702b5a5d5a..d1845719874c 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -17,6 +17,7 @@
 require "google/cloud/firestore/service"
 require "google/cloud/firestore/field_path"
 require "google/cloud/firestore/field_value"
+require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/collection_reference"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index c45e1b52b00c..a2a740c11efc 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -43,10 +43,6 @@ module Firestore
       #   end
       #
       class CollectionGroup < Query
-        ##
-        # @private The gRPC service object.
-        attr_accessor :service
-
         ##
         # Retrieves a list of document references for the documents in this
         # collection.
@@ -72,21 +68,14 @@ class CollectionGroup < Query
         #   col_group = firestore.col_group "cities"
         #
         #   col_group.partitions(3).each do |query_partition|
-        #     puts query_partition.create_query
+        #     # puts query_partition.create_query
         #   end
         #
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          unless block_given?
-            return enum_for :partitions, partition_count
-          end
-
-          results = service.partition_query parent_path, query, partition_count
-          results.each do |result|
-            next if result.result.nil?
-            yield result # TODO: QueryPartition.from_result result, self
-          end
+          resp_gapi = service.partition_query parent_path, query, partition_count, token: token, max: max
+          resp_gapi.partitions
         end
 
         ##
@@ -100,12 +89,7 @@ def self.from_collection_id parent_path, collection_id, client
               )
             ]
           )
-
-          new.tap do |q|
-            q.instance_variable_set :@query, query
-            q.instance_variable_set :@parent_path, parent_path
-            q.instance_variable_set :@client, client
-          end
+          CollectionGroup.new query, parent_path, client
         end
       end
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
index 3a8b7d3c97cf..9d724b8b7d15 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
@@ -47,6 +47,13 @@ class CollectionReference < Query
         # @private The firestore client object.
         attr_accessor :client
 
+        ##
+        # @private Creates a new CollectionReference.
+        def initialize path, client, query
+          super query, nil, client
+          @path = path
+        end
+
         ##
         # The collection identifier for the collection resource.
         #
@@ -257,11 +264,7 @@ def self.from_path path, client
             ]
           )
 
-          new.tap do |c|
-            c.client = client
-            c.instance_variable_set :@path, path
-            c.instance_variable_set :@query, query
-          end
+          CollectionReference.new path, client, query
         end
 
         protected
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 757f83363dd9..0bb0e910f8db 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -74,6 +74,16 @@ class Query
         # @private The firestore client object.
         attr_accessor :client
 
+        ##
+        # @private Creates a new Query.
+        def initialize query, parent_path, client, limit_type: nil
+          query ||= StructuredQuery.new
+          @query = query
+          @parent_path = parent_path
+          @limit_type = limit_type
+          @client = client
+        end
+
         ##
         # Restricts documents matching the query to return only data for the
         # provided fields.
@@ -965,13 +975,7 @@ def listen &callback
         ##
         # @private Start a new Query.
         def self.start query, parent_path, client, limit_type: nil
-          query ||= StructuredQuery.new
-          Query.new.tap do |q|
-            q.instance_variable_set :@query, query
-            q.instance_variable_set :@parent_path, parent_path
-            q.instance_variable_set :@limit_type, limit_type
-            q.instance_variable_set :@client, client
-          end
+          new query, parent_path, client, limit_type: limit_type
         end
 
         protected
@@ -1216,7 +1220,7 @@ def doc_id_path
         ##
         # @private Raise an error unless an database available.
         def ensure_client!
-          raise "Must have active connection to service" unless client
+          raise "Must have active connection to client" unless client
         end
 
         ##
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index baef8699aa44..2b84fe33d482 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -178,6 +178,18 @@ def mock_firestore
     end
   end
 
+  doctest.before "Google::Cloud::Firestore::CollectionGroup" do
+    mock_firestore do |mock|
+      mock.expect :run_query, run_query_resp, run_query_args
+    end
+  end
+
+  doctest.before "Google::Cloud::Firestore::CollectionGroup#partitions" do
+    mock_firestore do |mock|
+      mock.expect :partition_query, partition_query_resp, [Hash]
+    end
+  end
+
   doctest.skip "Google::Cloud::Firestore::FieldPath" do
     mock_firestore do |mock|
       mock.expect :batch_get_documents, batch_get_resp_users, batch_get_args
@@ -515,3 +527,20 @@ def documents_resp token: nil
   response.next_page_token = token if token
   paged_enum_struct response
 end
+
+def cursor_gapi
+  Google::Cloud::Firestore::V1::Cursor.new(
+    values: [
+      Google::Cloud::Firestore::V1::Value.new(string_value: "a")
+    ],
+    before: false
+  )
+end
+
+def partition_query_resp count: 3, token: nil
+  response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
+    partitions: count.times.map { cursor_gapi },
+    next_page_token: token
+  )
+  paged_enum_struct response
+end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
index 273252f45157..800961e43d47 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/client/col_group_test.rb
@@ -19,10 +19,10 @@
   let(:collection_id_bad) { "a/b/my-collection-id" }
 
   it "creates a collection group query" do
-    query = firestore.col_group(collection_id).where "foo", "==", "bar"
+    collection_group = firestore.col_group(collection_id)
 
-    _(query).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = query.query
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
     _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
     _(query_gapi.from.size).must_equal 1
     _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
@@ -31,10 +31,10 @@
   end
 
   it "creates a collection group query using collection_group alias" do
-    query = firestore.collection_group(collection_id).where "foo", "==", "bar"
+    collection_group = firestore.collection_group(collection_id)
 
-    _(query).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = query.query
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
     _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
     _(query_gapi.from.size).must_equal 1
     _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
new file mode 100644
index 000000000000..5f78596df987
--- /dev/null
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -0,0 +1,51 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "helper"
+
+describe Google::Cloud::Firestore::CollectionGroup, :mock_firestore do
+  let(:collection_id) { "my-collection-id" }
+  let(:collection_group) do
+    Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
+  end
+
+  it "creates a query" do
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    query_gapi = collection_group.query
+
+    _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
+    _(query_gapi.from.size).must_equal 1
+    _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
+    _(query_gapi.from.first.all_descendants).must_equal true
+    _(query_gapi.from.first.collection_id).must_equal collection_id
+  end
+focus
+  it "creates partitions" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [
+        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "my-collection-id", all_descendants: true)
+      ]
+    )
+    num_partitions = 3
+
+    list_res = paged_enum_struct partition_query_resp(num_partitions)
+    firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
+
+    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
+    partitions = collection_group.partitions 3
+
+    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+    _(partitions.count).must_equal 3
+  end
+end
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 56651a21a255..0d78920f379e 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -182,6 +182,34 @@ def run_query_args query,
     [req, default_options]
   end
 
+  def partition_query_args query,
+                           parent: "projects/#{project}/databases/(default)/documents",
+                           partition_count: 3
+    {
+      parent: parent,
+      structured_query: query,
+      partition_count: partition_count,
+      page_token: nil,
+      page_size: nil
+    }
+  end
+
+  def partition_query_resp count = 3, token = nil
+    Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
+      partitions: count.times.map { cursor_gapi },
+      next_page_token: token
+    )
+  end
+
+  def cursor_gapi
+    Google::Cloud::Firestore::V1::Cursor.new(
+      values: [
+        Google::Cloud::Firestore::V1::Value.new(string_value: "a")
+      ],
+      before: false
+    )
+  end
+
   def paged_enum_struct response
     OpenStruct.new response: response
   end

From fb1c55eb552da090e41966b59fc1f740a2998928 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Wed, 19 May 2021 17:49:31 -0600
Subject: [PATCH 09/22] Add QueryPartition and QueryPartition::List

---
 .../lib/google/cloud/firestore/client.rb      |   2 +-
 .../cloud/firestore/collection_group.rb       |  16 +-
 .../cloud/firestore/collection_reference.rb   |   2 +-
 .../google/cloud/firestore/query_partition.rb |  66 ++++++
 .../cloud/firestore/query_partition/list.rb   | 202 ++++++++++++++++++
 .../lib/google/cloud/firestore/service.rb     |   2 +
 .../support/doctest_helper.rb                 |  29 ++-
 .../cloud/firestore/collection_group_test.rb  |   4 +-
 google-cloud-firestore/test/helper.rb         |  15 +-
 9 files changed, 307 insertions(+), 31 deletions(-)
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/client.rb b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
index d1845719874c..b6e8459074d3 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/client.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/client.rb
@@ -17,10 +17,10 @@
 require "google/cloud/firestore/service"
 require "google/cloud/firestore/field_path"
 require "google/cloud/firestore/field_value"
-require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/collection_reference"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
+require "google/cloud/firestore/collection_group"
 require "google/cloud/firestore/batch"
 require "google/cloud/firestore/transaction"
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index a2a740c11efc..b09e8b4f4f8e 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -14,11 +14,8 @@
 
 
 require "google/cloud/firestore/v1"
-require "google/cloud/firestore/document_reference"
-require "google/cloud/firestore/document_snapshot"
 require "google/cloud/firestore/query"
-require "google/cloud/firestore/generate"
-require "google/cloud/firestore/collection_reference_list"
+require "google/cloud/firestore/query_partition"
 
 module Google
   module Cloud
@@ -64,18 +61,19 @@ class CollectionGroup < Query
         #
         #   firestore = Google::Cloud::Firestore.new
         #
-        #   # Get the cities collection group query
         #   col_group = firestore.col_group "cities"
         #
-        #   col_group.partitions(3).each do |query_partition|
-        #     # puts query_partition.create_query
+        #   partitions = col_group.partitions 3
+        #   partitions.each do |partition|
+        #     puts partition.create_query
         #   end
         #
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          resp_gapi = service.partition_query parent_path, query, partition_count, token: token, max: max
-          resp_gapi.partitions
+          grpc = service.partition_query parent_path, query, partition_count, token: token, max: max
+
+          QueryPartition::List.from_grpc grpc, client, parent_path, query, partition_count, max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
index 9d724b8b7d15..ab8a1a377a7b 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
@@ -14,9 +14,9 @@
 
 
 require "google/cloud/firestore/v1"
+require "google/cloud/firestore/query"
 require "google/cloud/firestore/document_reference"
 require "google/cloud/firestore/document_snapshot"
-require "google/cloud/firestore/query"
 require "google/cloud/firestore/generate"
 require "google/cloud/firestore/collection_reference_list"
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
new file mode 100644
index 000000000000..51e4fa191c1e
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -0,0 +1,66 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "google/cloud/firestore/query_partition/list"
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # # QueryPartition
+      #
+      # Represents a split point that can be used in a query as a starting and/or end point for the query results.
+      #
+      # The cursors returned by {#start_at} and {#end_before} can only be used in a query that matches the constraint of
+      # the query that produced this partition.
+      #
+      # See {CollectionGroup#partitions}.
+      #
+      # @example
+      #   require "google/cloud/firestore"
+      #
+      #   firestore = Google::Cloud::Firestore.new
+      #
+      #   col_group = firestore.col_group "cities"
+      #
+      #   partitions = col_group.partitions 3
+      #   partitions.each do |partition|
+      #     puts partition.create_query
+      #   end
+      #
+      class QueryPartition
+        attr_reader :start_at
+        attr_reader :end_before
+
+        ##
+        # @private New QueryPartition from query and Cursor
+        def initialize query, start_at, end_before
+          @query = query
+          @start_at = start_at
+          @end_before = end_before
+        end
+
+        ##
+        # Creates a query that only returns the documents for this partition.
+        #
+        # @return [Query] query.
+        #
+        def create_query
+          @query
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
new file mode 100644
index 000000000000..127861ec5a3b
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -0,0 +1,202 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+require "delegate"
+
+module Google
+  module Cloud
+    module Firestore
+      class QueryPartition
+        ##
+        # QueryPartition::List is a special case Array with additional values.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #
+        #   col_group = firestore.col_group "cities"
+        #
+        #   partitions = col_group.partitions 3
+        #   partitions.each do |partition|
+        #     puts partition.create_query
+        #   end
+        #
+        class List < DelegateClass(::Array)
+          ##
+          # If not empty, indicates that there are more records that match
+          # the request and this value should be passed to continue.
+          attr_accessor :token
+
+          ##
+          # @private Create a new QueryPartition::List with an array of
+          # QueryPartition instances.
+          def initialize arr = []
+            super arr
+          end
+
+          ##
+          # Whether there is a next page of document references.
+          #
+          # @return [Boolean]
+          #
+          # @example
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   if partitions.next?
+          #     next_documents = partitions.next
+          #   end
+          #
+          def next?
+            !token.nil?
+          end
+
+          ##
+          # Retrieve the next page of document references.
+          #
+          # @return [QueryPartition::List]
+          #
+          # @example
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   if partitions.next?
+          #     next_documents = partitions.next
+          #   end
+          #
+          def next
+            return nil unless next?
+            ensure_client!
+            grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
+            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, @max
+          end
+
+          ##
+          # Retrieves remaining results by repeatedly invoking {#next} until
+          # {#next?} returns `false`. Calls the given block once for each
+          # result, which is passed as the argument to the block.
+          #
+          # An Enumerator is returned if no block is given.
+          #
+          # This method will make repeated API calls until all remaining results
+          # are retrieved. (Unlike `#each`, for example, which merely iterates
+          # over the results returned by a single API call.) Use with caution.
+          #
+          # @param [Integer] request_limit The upper limit of API requests to
+          #   make to load all document references. Default is no limit.
+          # @yield [document] The block for accessing each document.
+          # @yieldparam [QueryPartition] document The document reference
+          #   object.
+          #
+          # @return [Enumerator]
+          #
+          # @example Iterating each document reference by passing a block:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   partitions.all do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          # @example Using the enumerator by not passing a block:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   all_queries = partitions.all.map do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          # @example Limit the number of API calls made:
+          #   require "google/cloud/firestore"
+          #
+          #   firestore = Google::Cloud::Firestore.new
+          #
+          #   col_group = firestore.col_group "cities"
+          #
+          #   partitions = col_group.partitions 3
+          #   partitions.all(request_limit: 10) do |partition|
+          #     puts partition.create_query
+          #   end
+          #
+          def all request_limit: nil, &block
+            request_limit = request_limit.to_i if request_limit
+            unless block_given?
+              return enum_for :all, request_limit: request_limit
+            end
+            results = self
+            loop do
+              results.each(&block)
+              if request_limit
+                request_limit -= 1
+                break if request_limit.negative?
+              end
+              break unless results.next?
+              results = results.next
+            end
+          end
+
+          ##
+          # @private New QueryPartition::List from a
+          # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
+          def self.from_grpc grpc, client, parent, structured_query, partition_count, max: nil
+            start_at = nil
+            partitions = List.new(Array(grpc.partitions).map do |cursor|
+              end_before = cursor.values.map do |value|
+                Convert.value_to_raw value, client
+              end
+              partition = QueryPartition.new structured_query, start_at, end_before
+              start_at = end_before
+              partition
+            end)
+            partitions.instance_variable_set :@parent, parent
+            partitions.instance_variable_set :@structured_query, structured_query
+            partitions.instance_variable_set :@partition_count, partition_count
+            token = grpc.next_page_token
+            token = nil if token == ""
+            partitions.instance_variable_set :@token, token
+            partitions.instance_variable_set :@client, client
+            partitions.instance_variable_set :@max, max
+            partitions
+          end
+
+          protected
+
+          ##
+          # Raise an error unless an active client is available.
+          def ensure_client!
+            raise "Must have active connection" unless @client
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index 49ea3edd8d46..6417ac4d0e25 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -96,6 +96,8 @@ def list_collections parent, token: nil, max: nil
           )
         end
 
+        ##
+        # Returns Google::Cloud::Firestore::V1::PartitionQueryResponse
         def partition_query parent, structured_query, partition_count, token: nil, max: nil
           paged_enum = firestore.partition_query parent: parent,
                                                  structured_query: structured_query,
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index 2b84fe33d482..cd380e3c3123 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -298,6 +298,12 @@ def mock_firestore
     end
   end
 
+  doctest.before "Google::Cloud::Firestore::QueryPartition" do
+    mock_firestore do |mock|
+      mock.expect :partition_query, partition_query_resp, [Hash]
+    end
+  end
+
   doctest.before "Google::Cloud::Firestore::CollectionReference" do
     mock_firestore do |mock|
       mock.expect :run_query, run_query_resp, run_query_args
@@ -528,19 +534,20 @@ def documents_resp token: nil
   paged_enum_struct response
 end
 
-def cursor_gapi
-  Google::Cloud::Firestore::V1::Cursor.new(
-    values: [
-      Google::Cloud::Firestore::V1::Value.new(string_value: "a")
-    ],
-    before: false
-  )
-end
-
 def partition_query_resp count: 3, token: nil
   response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-    partitions: count.times.map { cursor_gapi },
-    next_page_token: token
+    partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] }
   )
+  response.next_page_token = token if token
   paged_enum_struct response
 end
+
+def cursor_grpc values: [1,2], before: false
+  converted_values = values.map do |val|
+    Google::Cloud::Firestore::Convert.raw_to_value val
+  end
+  Google::Cloud::Firestore::V1::Cursor.new(
+    values: converted_values,
+    before: before
+  )
+end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index 5f78596df987..39022b622bd0 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -30,7 +30,7 @@
     _(query_gapi.from.first.all_descendants).must_equal true
     _(query_gapi.from.first.collection_id).must_equal collection_id
   end
-focus
+
   it "creates partitions" do
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [
@@ -39,7 +39,7 @@
     )
     num_partitions = 3
 
-    list_res = paged_enum_struct partition_query_resp(num_partitions)
+    list_res = paged_enum_struct partition_query_resp(count: num_partitions)
     firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
 
     _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 0d78920f379e..3ba1b94e7c32 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -194,19 +194,20 @@ def partition_query_args query,
     }
   end
 
-  def partition_query_resp count = 3, token = nil
+  def partition_query_resp count: 3, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      partitions: count.times.map { cursor_gapi },
+      partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
       next_page_token: token
     )
   end
 
-  def cursor_gapi
+  def cursor_grpc values: [1,2], before: false
+    converted_values = values.map do |val|
+      Google::Cloud::Firestore::Convert.raw_to_value val
+    end
     Google::Cloud::Firestore::V1::Cursor.new(
-      values: [
-        Google::Cloud::Firestore::V1::Value.new(string_value: "a")
-      ],
-      before: false
+      values: converted_values,
+      before: before
     )
   end
 

From e3a9d9f48a96094a201c6adff3c3196b970c8835 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Fri, 21 May 2021 15:56:50 -0600
Subject: [PATCH 10/22] Add order by to CollectionGroup#partitions

---
 .../acceptance/firestore/query_test.rb        |  97 ------------
 .../acceptance/firestore/watch_test.rb        |   2 +-
 .../cloud/firestore/collection_group.rb       |  10 +-
 .../cloud/firestore/query_partition/list.rb   |   2 +-
 .../cloud/firestore/collection_group_test.rb  | 138 +++++++++++++++---
 .../list_documents_test.rb                    |  15 --
 google-cloud-firestore/test/helper.rb         |  22 +--
 7 files changed, 142 insertions(+), 144 deletions(-)

diff --git a/google-cloud-firestore/acceptance/firestore/query_test.rb b/google-cloud-firestore/acceptance/firestore/query_test.rb
index 9e318a7e1dd8..1fdf6e4d6e03 100644
--- a/google-cloud-firestore/acceptance/firestore/query_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/query_test.rb
@@ -284,101 +284,4 @@
     _(results.map(&:document_id)).must_equal ["doc1", "doc2"]
     _(results.map { |doc| doc[:foo] }).must_equal ["a", "b"]
   end
-
-  describe "Collection Group" do
-    it "queries a collection group" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-          "abc/123/#{collection_group}/cg-doc1",
-          "abc/123/#{collection_group}/cg-doc2",
-          "#{collection_group}/cg-doc3",
-          "#{collection_group}/cg-doc4",
-          "def/456/#{collection_group}/cg-doc5",
-          "#{collection_group}/virtual-doc/nested-coll/not-cg-doc",
-          "x#{collection_group}/not-cg-doc",
-          "#{collection_group}x/not-cg-doc",
-          "abc/123/#{collection_group}x/not-cg-doc",
-          "abc/123/x#{collection_group}/not-cg-doc",
-          "abc/#{collection_group}"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group collection_group
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc1", "cg-doc2", "cg-doc3", "cg-doc4", "cg-doc5"]
-    end
-
-    it "queries a collection group with start_at and end_at" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-        "a/a/#{collection_group}/cg-doc1",
-        "a/b/a/b/#{collection_group}/cg-doc2",
-        "a/b/#{collection_group}/cg-doc3",
-        "a/b/c/d/#{collection_group}/cg-doc4",
-        "a/c/#{collection_group}/cg-doc5",
-        "#{collection_group}/cg-doc6",
-        "a/b/nope/nope"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group(collection_group)
-        .order_by("__name__")
-        .start_at(firestore.document("a/b"))
-        .end_at(firestore.document("a/b0"))
-
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
-
-      query = firestore.collection_group(collection_group)
-        .order_by("__name__")
-        .start_after(firestore.document("a/b"))
-        .end_before(firestore.document("a/b/#{collection_group}/cg-doc3"))
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
-    end
-
-    it "queries a collection group with filters" do
-      collection_group = "b-#{SecureRandom.hex(4)}"
-      doc_paths = [
-        "a/a/#{collection_group}/cg-doc1",
-        "a/b/a/b/#{collection_group}/cg-doc2",
-        "a/b/#{collection_group}/cg-doc3",
-        "a/b/c/d/#{collection_group}/cg-doc4",
-        "a/c/#{collection_group}/cg-doc5",
-        "#{collection_group}/cg-doc6",
-        "a/b/nope/nope"
-      ]
-      firestore.batch do |b|
-        doc_paths.each do |doc_path|
-          doc_ref = firestore.document doc_path
-          b.set doc_ref, {x: 1}
-        end
-      end
-
-      query = firestore.collection_group(collection_group)
-        .where("__name__", ">=", firestore.document("a/b"))
-        .where("__name__", "<=", firestore.document("a/b0"))
-
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
-
-      query = firestore.collection_group(collection_group)
-        .where("__name__", ">", firestore.document("a/b"))
-        .where(
-          "__name__", "<", firestore.document("a/b/#{collection_group}/cg-doc3")
-        )
-      snapshots = query.get
-      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
-    end
-  end
 end
diff --git a/google-cloud-firestore/acceptance/firestore/watch_test.rb b/google-cloud-firestore/acceptance/firestore/watch_test.rb
index c8b77cc6b49e..b385f641847c 100644
--- a/google-cloud-firestore/acceptance/firestore/watch_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/watch_test.rb
@@ -188,7 +188,7 @@
   def wait_until &block
     wait_count = 0
     until block.call
-      fail "wait_until criterial was not met" if wait_count > 200
+      fail "wait_until criteria was not met" if wait_count > 200
       wait_count += 1
       sleep 0.01
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index b09e8b4f4f8e..af724c196612 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -71,9 +71,15 @@ class CollectionGroup < Query
         def partitions partition_count, token: nil, max: nil
           ensure_service!
 
-          grpc = service.partition_query parent_path, query, partition_count, token: token, max: max
+          # Partition queries require explicit ordering by __name__.
+          query_with_default_order = order("__name__").query
+          # Since we are always returning an extra partition (with en empty endBefore cursor), we reduce the desired
+          # partition count by one.
+          partition_count -= 1
 
-          QueryPartition::List.from_grpc grpc, client, parent_path, query, partition_count, max: max
+          grpc = service.partition_query parent_path, query_with_default_order, partition_count, token: token, max: max
+
+          QueryPartition::List.from_grpc grpc, client, parent_path, query_with_default_order, partition_count, max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index 127861ec5a3b..3bc8dd5259f2 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -89,7 +89,7 @@ def next
             return nil unless next?
             ensure_client!
             grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, @max
+            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, max: @max
           end
 
           ##
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index 39022b622bd0..cf0507da5fd5 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -19,33 +19,133 @@
   let(:collection_group) do
     Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
   end
-
-  it "creates a query" do
-    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
-    query_gapi = collection_group.query
-
-    _(query_gapi).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery
-    _(query_gapi.from.size).must_equal 1
-    _(query_gapi.from.first).must_be_kind_of Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector
-    _(query_gapi.from.first.all_descendants).must_equal true
-    _(query_gapi.from.first.collection_id).must_equal collection_id
-  end
-
-  it "creates partitions" do
-    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+  let(:expected_query) do
+    Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [
-        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "my-collection-id", all_descendants: true)
+        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
+          collection_id: "my-collection-id",
+          all_descendants: true
+        )
+      ],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "__name__"),
+          direction: :ASCENDING
+        )
       ]
     )
-    num_partitions = 3
+  end
 
-    list_res = paged_enum_struct partition_query_resp(count: num_partitions)
-    firestore_mock.expect :partition_query, list_res, [partition_query_args(expected_query, partition_count: num_partitions)]
+  it "lists partitions" do
+    list_res = paged_enum_struct partition_query_resp(count: 3)
+    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
 
-    _(collection_group).must_be_kind_of Google::Cloud::Firestore::Query
     partitions = collection_group.partitions 3
 
+    firestore_mock.verify
+
+    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+    _(partitions.count).must_equal 3
+    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+  end
+
+  it "paginates partitions" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    first_partitions = collection_group.partitions 6
+    second_partitions = collection_group.partitions 6, token: first_partitions.token
+
+    firestore_mock.verify
+
+    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(first_partitions.count).must_equal 3
+    _(first_partitions.token).must_equal "next_page_token"
+
+    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(second_partitions.count).must_equal 2
+    _(second_partitions.token).must_be :nil?
+  end
+
+  it "paginates partitions using and next" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    first_partitions = collection_group.partitions 6
+    second_partitions = first_partitions.next
+
+    firestore_mock.verify
+
+    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(first_partitions.count).must_equal 3
+    _(first_partitions.token).must_equal "next_page_token"
+
+    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(second_partitions.count).must_equal 2
+    _(second_partitions.token).must_be :nil?
+  end
+
+  it "paginates partitions using all" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all.to_a
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 5
+  end
+
+  it "paginates partitions using all and Enumerator" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all.take 5
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 5
+  end
+
+  it "paginates partitions using all with request_limit set" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+
+    all_partitions = collection_group.partitions(6).all(request_limit: 1).to_a
+
+    firestore_mock.verify
+
+    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+    _(all_partitions.count).must_equal 6
+  end
+
+  it "paginates partitions with max set" do
+    list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+
+    partitions = collection_group.partitions 6, max: 3
+
+    firestore_mock.verify
+
     _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
     _(partitions.count).must_equal 3
+    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
   end
 end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
index a73a097f7e78..9de697a5bf41 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_reference/list_documents_test.rb
@@ -137,21 +137,6 @@
     _(documents.token).must_equal "next_page_token"
   end
 
-  it "paginates documents without max set" do
-    list_res = paged_enum_struct list_documents_gapi(3, "next_page_token")
-
-    firestore_mock.expect :list_documents, list_res, list_documents_args
-
-    documents = collection.list_documents
-
-    firestore_mock.verify
-
-    documents.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::DocumentReference }
-    _(documents.count).must_equal 3
-    _(documents.token).wont_be :nil?
-    _(documents.token).must_equal "next_page_token"
-  end
-
   def document_gapi
     Google::Cloud::Firestore::V1::Document.new(
       name: "projects/#{project}/databases/(default)/documents/my-document",
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 3ba1b94e7c32..637ab76556d0 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -127,7 +127,7 @@ class MockFirestore < Minitest::Spec
   def wait_until &block
     wait_count = 0
     until block.call
-      fail "wait_until criterial was not met" if wait_count > 100
+      fail "wait_until criteria was not met" if wait_count > 100
       wait_count += 1
       sleep 0.01
     end
@@ -184,14 +184,18 @@ def run_query_args query,
 
   def partition_query_args query,
                            parent: "projects/#{project}/databases/(default)/documents",
-                           partition_count: 3
-    {
-      parent: parent,
-      structured_query: query,
-      partition_count: partition_count,
-      page_token: nil,
-      page_size: nil
-    }
+                           partition_count: 2,
+                           page_token: nil,
+                           page_size: nil
+    [
+      {
+        parent: parent,
+        structured_query: query,
+        partition_count: partition_count,
+        page_token: page_token,
+        page_size: page_size
+      }
+    ]
   end
 
   def partition_query_resp count: 3, token: nil

From 7317f9b0f45bcf9be941417b89c0fc952f2c3b35 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Fri, 21 May 2021 18:03:03 -0600
Subject: [PATCH 11/22] Fix CollectionGroup#partitions and
 QueryPartition::List.from_grpc

---
 .../firestore/collection_group_test.rb        | 133 ++++++++++++++++++
 .../acceptance/firestore_helper.rb            |  20 +--
 .../cloud/firestore/collection_group.rb       |  15 +-
 .../google/cloud/firestore/query_partition.rb |   5 +-
 .../cloud/firestore/query_partition/list.rb   |  13 +-
 .../lib/google/cloud/firestore/service.rb     |  15 +-
 .../support/doctest_helper.rb                 |   4 +-
 .../cloud/firestore/collection_group_test.rb  |  36 ++---
 google-cloud-firestore/test/helper.rb         |  14 +-
 9 files changed, 207 insertions(+), 48 deletions(-)
 create mode 100644 google-cloud-firestore/acceptance/firestore/collection_group_test.rb

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
new file mode 100644
index 000000000000..dda44fc346f3
--- /dev/null
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -0,0 +1,133 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "firestore_helper"
+
+describe Google::Cloud::Firestore::CollectionGroup, :firestore_acceptance do
+  describe "#get" do
+    it "queries a collection group" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+          "abc/123/#{collection_id}/cg-doc1",
+          "abc/123/#{collection_id}/cg-doc2",
+          "#{collection_id}/cg-doc3",
+          "#{collection_id}/cg-doc4",
+          "def/456/#{collection_id}/cg-doc5",
+          "#{collection_id}/virtual-doc/nested-coll/not-cg-doc",
+          "x#{collection_id}/not-cg-doc",
+          "#{collection_id}x/not-cg-doc",
+          "abc/123/#{collection_id}x/not-cg-doc",
+          "abc/123/x#{collection_id}/not-cg-doc",
+          "abc/#{collection_id}"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc1", "cg-doc2", "cg-doc3", "cg-doc4", "cg-doc5"]
+    end
+
+    it "queries a collection group with start_at and end_at" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+        "a/a/#{collection_id}/cg-doc1",
+        "a/b/a/b/#{collection_id}/cg-doc2",
+        "a/b/#{collection_id}/cg-doc3",
+        "a/b/c/d/#{collection_id}/cg-doc4",
+        "a/c/#{collection_id}/cg-doc5",
+        "#{collection_id}/cg-doc6",
+        "a/b/nope/nope"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+        .order_by("__name__")
+        .start_at(firestore.document("a/b"))
+        .end_at(firestore.document("a/b0"))
+
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
+
+      collection_group = firestore.collection_group(collection_id)
+        .order_by("__name__")
+        .start_after(firestore.document("a/b"))
+        .end_before(firestore.document("a/b/#{collection_id}/cg-doc3"))
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
+    end
+
+    it "queries a collection group with filters" do
+      collection_id = "b-#{SecureRandom.hex(4)}"
+      doc_paths = [
+        "a/a/#{collection_id}/cg-doc1",
+        "a/b/a/b/#{collection_id}/cg-doc2",
+        "a/b/#{collection_id}/cg-doc3",
+        "a/b/c/d/#{collection_id}/cg-doc4",
+        "a/c/#{collection_id}/cg-doc5",
+        "#{collection_id}/cg-doc6",
+        "a/b/nope/nope"
+      ]
+      firestore.batch do |b|
+        doc_paths.each do |doc_path|
+          doc_ref = firestore.document doc_path
+          b.set doc_ref, {x: 1}
+        end
+      end
+
+      collection_group = firestore.collection_group(collection_id)
+        .where("__name__", ">=", firestore.document("a/b"))
+        .where("__name__", "<=", firestore.document("a/b0"))
+
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2", "cg-doc3", "cg-doc4"]
+
+      collection_group = firestore.collection_group(collection_id)
+        .where("__name__", ">", firestore.document("a/b"))
+        .where(
+          "__name__", "<", firestore.document("a/b/#{collection_id}/cg-doc3")
+        )
+      snapshots = collection_group.get
+      _(snapshots.map(&:document_id)).must_equal ["cg-doc2"]
+    end
+  end
+
+  describe "#partitions" do
+    it "queries a collection group with partitions" do
+      document_count = 2 * 128 + 127 # Minimum partition size is 128.
+      rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
+      firestore.batch do |b|
+        document_count.times do |i|
+          doc_ref = rand_col.document i
+          b.set doc_ref, {foo: i}
+        end
+      end
+
+      collection_group = firestore.collection_group(rand_col.collection_id)
+
+      partitions = collection_group.partitions 3
+      _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+      _(partitions.count).must_equal 3
+    end
+  end
+end
diff --git a/google-cloud-firestore/acceptance/firestore_helper.rb b/google-cloud-firestore/acceptance/firestore_helper.rb
index af21a6405b82..4b9d678dd745 100644
--- a/google-cloud-firestore/acceptance/firestore_helper.rb
+++ b/google-cloud-firestore/acceptance/firestore_helper.rb
@@ -65,16 +65,16 @@ def root_col
       addl.include? :firestore_acceptance
     end
 
-    def self.run_one_method klass, method_name, reporter
-      result = nil
-      reporter.prerecord klass, method_name
-      (1..3).each do |try|
-        result = Minitest.run_one_method(klass, method_name)
-        break if (result.passed? || result.skipped?)
-        puts "Retrying #{klass}##{method_name} (#{try})"
-      end
-      reporter.record result
-    end
+    # def self.run_one_method klass, method_name, reporter
+    #   result = nil
+    #   reporter.prerecord klass, method_name
+    #   (1..3).each do |try|
+    #     result = Minitest.run_one_method(klass, method_name)
+    #     break if (result.passed? || result.skipped?)
+    #     puts "Retrying #{klass}##{method_name} (#{try})"
+    #   end
+    #   reporter.record result
+    # end
   end
 end
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index af724c196612..efa21947d922 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -72,14 +72,23 @@ def partitions partition_count, token: nil, max: nil
           ensure_service!
 
           # Partition queries require explicit ordering by __name__.
-          query_with_default_order = order("__name__").query
+          query_with_default_order = order "__name__"
           # Since we are always returning an extra partition (with en empty endBefore cursor), we reduce the desired
           # partition count by one.
           partition_count -= 1
 
-          grpc = service.partition_query parent_path, query_with_default_order, partition_count, token: token, max: max
+          grpc = service.partition_query parent_path,
+                                         query_with_default_order.query,
+                                         partition_count,
+                                         token: token,
+                                         max: max
 
-          QueryPartition::List.from_grpc grpc, client, parent_path, query_with_default_order, partition_count, max: max
+          QueryPartition::List.from_grpc grpc,
+                                         client,
+                                         parent_path,
+                                         query_with_default_order,
+                                         partition_count,
+                                         max: max
         end
 
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 51e4fa191c1e..071a8cb71eb4 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -58,7 +58,10 @@ def initialize query, start_at, end_before
         # @return [Query] query.
         #
         def create_query
-          @query
+          base_query = @query
+          base_query = base_query.start_at start_at if start_at
+          base_query = base_query.end_before end_before if end_before
+          base_query
         end
       end
     end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index 3bc8dd5259f2..a068726f6be7 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -88,8 +88,8 @@ def next?
           def next
             return nil unless next?
             ensure_client!
-            grpc = @client.service.partition_query @parent, @structured_query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @structured_query, @partition_count, max: @max
+            grpc = @client.service.partition_query @parent, @query.query, @partition_count, token: token, max: @max
+            self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max
           end
 
           ##
@@ -167,18 +167,21 @@ def all request_limit: nil, &block
           ##
           # @private New QueryPartition::List from a
           # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
-          def self.from_grpc grpc, client, parent, structured_query, partition_count, max: nil
+          def self.from_grpc grpc, client, parent, query, partition_count, max: nil
+            # TODO: Should the logic adding the last partition be applied to the entire result set, not the page?
             start_at = nil
             partitions = List.new(Array(grpc.partitions).map do |cursor|
               end_before = cursor.values.map do |value|
                 Convert.value_to_raw value, client
               end
-              partition = QueryPartition.new structured_query, start_at, end_before
+              partition = QueryPartition.new query, start_at, end_before
               start_at = end_before
               partition
             end)
+            # We are always returning an extra partition (with en empty endBefore cursor).
+            partitions << QueryPartition.new(query, start_at, nil)
             partitions.instance_variable_set :@parent, parent
-            partitions.instance_variable_set :@structured_query, structured_query
+            partitions.instance_variable_set :@query, query
             partitions.instance_variable_set :@partition_count, partition_count
             token = grpc.next_page_token
             token = nil if token == ""
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/service.rb b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
index 6417ac4d0e25..54c4acf7d77b 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/service.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/service.rb
@@ -98,12 +98,15 @@ def list_collections parent, token: nil, max: nil
 
         ##
         # Returns Google::Cloud::Firestore::V1::PartitionQueryResponse
-        def partition_query parent, structured_query, partition_count, token: nil, max: nil
-          paged_enum = firestore.partition_query parent: parent,
-                                                 structured_query: structured_query,
-                                                 partition_count: partition_count,
-                                                 page_token: token,
-                                                 page_size: max
+        def partition_query parent, query_grpc, partition_count, token: nil, max: nil
+          request = Google::Cloud::Firestore::V1::PartitionQueryRequest.new(
+            parent: parent,
+            structured_query: query_grpc,
+            partition_count: partition_count,
+            page_token: token,
+            page_size: max
+          )
+          paged_enum = firestore.partition_query request
           paged_enum.response
         end
 
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index cd380e3c3123..4b2e216a7fca 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -186,7 +186,7 @@ def mock_firestore
 
   doctest.before "Google::Cloud::Firestore::CollectionGroup#partitions" do
     mock_firestore do |mock|
-      mock.expect :partition_query, partition_query_resp, [Hash]
+      mock.expect :partition_query, partition_query_resp, [Google::Cloud::Firestore::V1::PartitionQueryRequest]
     end
   end
 
@@ -300,7 +300,7 @@ def mock_firestore
 
   doctest.before "Google::Cloud::Firestore::QueryPartition" do
     mock_firestore do |mock|
-      mock.expect :partition_query, partition_query_resp, [Hash]
+      mock.expect :partition_query, partition_query_resp, [Google::Cloud::Firestore::V1::PartitionQueryRequest]
     end
   end
 
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index cf0507da5fd5..f66859960d06 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -37,7 +37,10 @@
   end
 
   it "lists partitions" do
-    list_res = paged_enum_struct partition_query_resp(count: 3)
+    # grpc.partitions: 
+    # [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/32">], before: false>,
+    #  <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/89">], before: false>]
+    list_res = paged_enum_struct partition_query_resp
     firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
 
     partitions = collection_group.partitions 3
@@ -46,12 +49,17 @@
 
     _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
     _(partitions.count).must_equal 3
-    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+
+    _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[0].start_at).must_be :nil?
+    _(partitions[0].end_before).must_equal [0,1]
+    query_1 = partitions[0].create_query
+    # TODO: query assertions
   end
 
   it "paginates partitions" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+    first_list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -70,9 +78,9 @@
     _(second_partitions.token).must_be :nil?
   end
 
-  it "paginates partitions using and next" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+  it "paginates partitions using next" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -80,8 +88,6 @@
     first_partitions = collection_group.partitions 6
     second_partitions = first_partitions.next
 
-    firestore_mock.verify
-
     first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
     _(first_partitions.count).must_equal 3
     _(first_partitions.token).must_equal "next_page_token"
@@ -92,8 +98,8 @@
   end
 
   it "paginates partitions using all" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 1)
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
@@ -122,22 +128,20 @@
   end
 
   it "paginates partitions using all with request_limit set" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
+    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2, token: "second_page_token")
 
     firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
     firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
 
     all_partitions = collection_group.partitions(6).all(request_limit: 1).to_a
 
-    firestore_mock.verify
-
     all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
     _(all_partitions.count).must_equal 6
   end
 
   it "paginates partitions with max set" do
-    list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
     firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
 
     partitions = collection_group.partitions 6, max: 3
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 637ab76556d0..cad0a5cbdf1b 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -182,23 +182,27 @@ def run_query_args query,
     [req, default_options]
   end
 
-  def partition_query_args query,
+  # expected partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>"next_page_token", :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "second_page_token">>
+  # got     [partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>nil, :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "next_page_token">>]
+
+
+  def partition_query_args query_grpc,
                            parent: "projects/#{project}/databases/(default)/documents",
                            partition_count: 2,
                            page_token: nil,
                            page_size: nil
     [
-      {
+      Google::Cloud::Firestore::V1::PartitionQueryRequest.new(
         parent: parent,
-        structured_query: query,
+        structured_query: query_grpc,
         partition_count: partition_count,
         page_token: page_token,
         page_size: page_size
-      }
+      )
     ]
   end
 
-  def partition_query_resp count: 3, token: nil
+  def partition_query_resp count: 2, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
       partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
       next_page_token: token

From 058e306088e706fa950d3658c36bf038c1a15731 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Mon, 24 May 2021 16:30:36 -0600
Subject: [PATCH 12/22] Fix array support in Query#start_at, #start_after,
 #end_before and #end_at

---
 .../lib/google/cloud/firestore/query.rb       |   6 +-
 .../support/doctest_helper.rb                 |   6 +-
 .../cloud/firestore/collection_group_test.rb  |   2 +-
 .../cloud/firestore/query/cursors_test.rb     | 101 +++++++++++++++++-
 google-cloud-firestore/test/helper.rb         |   6 +-
 5 files changed, 113 insertions(+), 8 deletions(-)

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 0bb0e910f8db..5bd38475efa4 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -1098,11 +1098,14 @@ def values_to_cursor values, query
             return snapshot_to_cursor values.first, query
           end
 
+          # The *values param in start_at, start_after, etc. will wrap an array argument in an array, so unwrap it here.
+          values = values.first if values.count == 1 && values.first.is_a?(Array)
+
           # pair values with their field_paths to ensure correct formatting
           order_field_paths = order_by_field_paths query
           if values.count > order_field_paths.count
             # raise if too many values provided for the cursor
-            raise ArgumentError, "too many values"
+            raise ArgumentError, "There cannot be more values than Order By fields"
           end
 
           values = values.zip(order_field_paths).map do |value, field_path|
@@ -1134,7 +1137,6 @@ def snapshot_to_cursor snapshot, query
               snapshot[field_path]
             end
           end
-
           values_to_cursor values, query
         end
 
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index 4b2e216a7fca..5d954e206195 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -536,13 +536,15 @@ def documents_resp token: nil
 
 def partition_query_resp count: 3, token: nil
   response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-    partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] }
+    # Minimum partition size is 128.
+    partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] }
   )
   response.next_page_token = token if token
   paged_enum_struct response
 end
 
-def cursor_grpc values: [1,2], before: false
+# Minimum partition size is 128.
+def cursor_grpc values: ["128"], before: false
   converted_values = values.map do |val|
     Google::Cloud::Firestore::Convert.raw_to_value val
   end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index f66859960d06..fc9e549a9e53 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -52,7 +52,7 @@
 
     _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
     _(partitions[0].start_at).must_be :nil?
-    _(partitions[0].end_before).must_equal [0,1]
+    _(partitions[0].end_before).must_equal ["128"]
     query_1 = partitions[0].create_query
     # TODO: query assertions
   end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb b/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
index a9346b8c3e2c..61648762d57a 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/query/cursors_test.rb
@@ -305,7 +305,40 @@
     _(generated_query).must_equal expected_query
   end
 
-  it "Start/End with two values" do
+  it "StartAt/EndBefore with two values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: true
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: true
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_at(7, 8).end_before(9, 10).query
+    _(generated_query).must_equal expected_query
+  end
+
+  it "StartAfter/EndAt with two values" do
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
       order_by: [
@@ -338,6 +371,72 @@
     _(generated_query).must_equal expected_query
   end
 
+  it "StartAt/EndBefore with array values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: true
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: true
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_at([7, 8]).end_before([9, 10]).query
+    _(generated_query).must_equal expected_query
+  end
+
+  it "StartAfter/EndAt with array values" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "C")],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "a"),
+          direction: :ASCENDING
+        ),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "b"),
+          direction: :DESCENDING
+        )
+      ],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(7),
+          Google::Cloud::Firestore::Convert.raw_to_value(8)
+        ],
+        before: false
+      ),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(
+        values: [
+          Google::Cloud::Firestore::Convert.raw_to_value(9),
+          Google::Cloud::Firestore::Convert.raw_to_value(10)
+        ],
+        before: false
+      )
+    )
+
+    generated_query = collection.order(:a).order(:b, :desc).start_after([7, 8]).end_at([9, 10]).query
+    _(generated_query).must_equal expected_query
+  end
+
   it "with __name__" do
     # TODO: Looks like we need to create the full path when paired with __name__
     expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index cad0a5cbdf1b..df251ad9fb41 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -204,12 +204,14 @@ def partition_query_args query_grpc,
 
   def partition_query_resp count: 2, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      partitions: count.times.map { |i| cursor_grpc values: [i*10, i*10+1] },
+      # Minimum partition size is 128.
+      partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] },
       next_page_token: token
     )
   end
 
-  def cursor_grpc values: [1,2], before: false
+  # Minimum partition size is 128.
+  def cursor_grpc values: ["128"], before: false
     converted_values = values.map do |val|
       Google::Cloud::Firestore::Convert.raw_to_value val
     end

From f7f80c50bb0026d9b4e2a653f43ab705199b1e49 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 25 May 2021 15:49:17 -0600
Subject: [PATCH 13/22] Update partitions samples

---
 .../cloud/firestore/collection_group.rb       |  5 ++---
 .../google/cloud/firestore/query_partition.rb |  5 ++---
 .../cloud/firestore/query_partition/list.rb   | 20 ++++++++-----------
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index efa21947d922..68c5df58450e 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -64,9 +64,8 @@ class CollectionGroup < Query
         #   col_group = firestore.col_group "cities"
         #
         #   partitions = col_group.partitions 3
-        #   partitions.each do |partition|
-        #     puts partition.create_query
-        #   end
+        #
+        #   queries = partitions.map(&:create_query)
         #
         def partitions partition_count, token: nil, max: nil
           ensure_service!
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 071a8cb71eb4..c179c1f02a30 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -36,9 +36,8 @@ module Firestore
       #   col_group = firestore.col_group "cities"
       #
       #   partitions = col_group.partitions 3
-      #   partitions.each do |partition|
-      #     puts partition.create_query
-      #   end
+      #
+      #   queries = partitions.map(&:create_query)
       #
       class QueryPartition
         attr_reader :start_at
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index a068726f6be7..4b66a941a655 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -30,9 +30,8 @@ class QueryPartition
         #   col_group = firestore.col_group "cities"
         #
         #   partitions = col_group.partitions 3
-        #   partitions.each do |partition|
-        #     puts partition.create_query
-        #   end
+        #
+        #   queries = partitions.map(&:create_query)
         #
         class List < DelegateClass(::Array)
           ##
@@ -119,9 +118,8 @@ def next
           #   col_group = firestore.col_group "cities"
           #
           #   partitions = col_group.partitions 3
-          #   partitions.all do |partition|
-          #     puts partition.create_query
-          #   end
+          #
+          #   queries = partitions.map(&:create_query)
           #
           # @example Using the enumerator by not passing a block:
           #   require "google/cloud/firestore"
@@ -131,9 +129,8 @@ def next
           #   col_group = firestore.col_group "cities"
           #
           #   partitions = col_group.partitions 3
-          #   all_queries = partitions.all.map do |partition|
-          #     puts partition.create_query
-          #   end
+          #
+          #   queries = partitions.map(&:create_query)
           #
           # @example Limit the number of API calls made:
           #   require "google/cloud/firestore"
@@ -143,9 +140,8 @@ def next
           #   col_group = firestore.col_group "cities"
           #
           #   partitions = col_group.partitions 3
-          #   partitions.all(request_limit: 10) do |partition|
-          #     puts partition.create_query
-          #   end
+          #
+          #   queries = partitions.map(&:create_query)
           #
           def all request_limit: nil, &block
             request_limit = request_limit.to_i if request_limit

From a8388a2630f4f0cc54a135339054bd1622d67c72 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 25 May 2021 15:52:04 -0600
Subject: [PATCH 14/22] Update documentation

---
 .../google/cloud/firestore/collection_group.rb  | 17 ++++++-----------
 .../lib/google/cloud/firestore/field_value.rb   |  9 +++++++--
 .../google/cloud/firestore/query_partition.rb   |  8 ++++++++
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index 68c5df58450e..fe77f7501493 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -41,20 +41,15 @@ module Firestore
       #
       class CollectionGroup < Query
         ##
-        # Retrieves a list of document references for the documents in this
-        # collection.
+        # Partitions a query by returning partition cursors that can be used to run the query in parallel. The returned
+        # partition cursors are split points that can be used as starting/end points for the query results.
         #
-        # The document references returned may include references to "missing
-        # documents", i.e. document locations that have no document present but
-        # which contain subcollections with documents. Attempting to read such a
-        # document reference (e.g. via {DocumentReference#get}) will return
-        # a {DocumentSnapshot} whose `exists?` method returns false.
-        #
-        # @param [String] token A previously-returned page token representing
-        #   part of the larger set of results to view.
+        # @param [Integer] partition_count The desired maximum number of partition points. The number must be strictly
+        #   positive. The actual number of partitions returned may be fewer.
+        # @param [String] token A previously-returned page token representing part of the larger set of results to view.
         # @param [Integer] max Maximum number of results to return.
         #
-        # @return [Array<DocumentReference>] An array of document references.
+        # @return [Array<QueryPartition>] An array of query partitions.
         #
         # @example
         #   require "google/cloud/firestore"
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/field_value.rb b/google-cloud-firestore/lib/google/cloud/firestore/field_value.rb
index 9a3855f51d93..97cd7ce661b2 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/field_value.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/field_value.rb
@@ -27,9 +27,14 @@ module Firestore
       #
       #   firestore = Google::Cloud::Firestore.new
       #
-      #   user_snap = firestore.doc("users/frank").get
+      #   # Get a document reference
+      #   nyc_ref = firestore.doc "cities/NYC"
       #
-      #   # TODO
+      #   # Set the population to increment by 1.
+      #   increment_value = Google::Cloud::Firestore::FieldValue.increment 1
+      #
+      #   nyc_ref.update({ name: "New York City",
+      #                    population: increment_value })
       #
       class FieldValue
         ##
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index c179c1f02a30..071013b161f9 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -28,6 +28,14 @@ module Firestore
       #
       # See {CollectionGroup#partitions}.
       #
+      # @!attribute [r] start_at
+      #   The cursor that defines the first result for this partition, or `nil` if this is the first partition.
+      #   @return [Array] a cursor value that can be used with {Query#start_at} or `nil` if this is the first partition.
+      # @!attribute [r] end_before
+      #   The cursor that defines the first result after this partition, or `nil` if this is the last partition.
+      #   @return [Array] a cursor value that can be used with {Query#end_before} or `nil` if this is the last
+      #     partition.
+      #
       # @example
       #   require "google/cloud/firestore"
       #

From 4e61da5890334a305c3c833a94575fee8ba11fd7 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Wed, 26 May 2021 13:08:10 -0600
Subject: [PATCH 15/22] Update QueryPartition::List to add final QueryPartition
 only at end of all results

---
 .../firestore/collection_group_test.rb        |  19 ++-
 .../cloud/firestore/collection_group.rb       |   4 +-
 .../firestore/document_reference/list.rb      |   2 +-
 .../lib/google/cloud/firestore/query.rb       |   2 +-
 .../google/cloud/firestore/query_partition.rb |  22 +--
 .../cloud/firestore/query_partition/list.rb   |  41 +++---
 .../cloud/firestore/collection_group_test.rb  | 136 ++++++++++--------
 google-cloud-firestore/test/helper.rb         |  16 ++-
 8 files changed, 142 insertions(+), 100 deletions(-)

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
index 07fe5249efa2..b9435c348ade 100644
--- a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -114,26 +114,26 @@
 
   describe "#partitions" do
     it "queries a collection group using partitions" do
+      rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
+
       document_ids = ["a", "b", "c"].map do |prefix|
         # Minimum partition size is 128.
         128.times.map do |i|
           "#{prefix}#{(i+1).to_s.rjust(3, '0')}"
         end
       end.flatten # "a001", "a002", ... "c128"
-
-      rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
       firestore.batch do |b|
         document_ids.each do |id|
           doc_ref = rand_col.document id
-          b.set doc_ref, {}
+          b.set doc_ref, { foo: id }
         end
       end
 
       collection_group = firestore.collection_group(rand_col.collection_id)
 
-      partitions = collection_group.partitions 3
+      partitions = collection_group.partitions 6, max: 1
       _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
-      _(partitions.count).must_equal 3
+      _(partitions.count).must_equal 1
 
       _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
       _(partitions[0].start_at).must_be :nil?
@@ -141,6 +141,10 @@
       _(partitions[0].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
       _(document_ids).must_include partitions[0].end_before[0].document_id
 
+      partitions += partitions.next
+      _(partitions).must_be_kind_of Array
+      _(partitions.count).must_equal 3
+
       _(partitions[1]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
       _(partitions[1].start_at).must_be_kind_of Array
       _(partitions[1].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
@@ -159,7 +163,10 @@
       _(queries.count).must_equal 3
       results = queries.map do |query|
         _(query).must_be_kind_of Google::Cloud::Firestore::Query
-        query.get.map(&:document_id)
+        query.get.map do |snp|
+          _(snp).must_be_kind_of Google::Cloud::Firestore::DocumentSnapshot
+          snp.document_id
+        end
       end
       results.each { |result| _(result).wont_be :empty? }
       # Verify all document IDs have been returned, in original order.
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index fe77f7501493..fe333a81735a 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -24,7 +24,9 @@ module Firestore
       # # CollectionGroup
       #
       # A collection group object is used for adding documents, getting
-      # document references, and querying for documents (See {Query}).
+      # document references, and querying for documents, including with partitions.
+      #
+      # See {Client#col_group} and {Query}.
       #
       # @example
       #   require "google/cloud/firestore"
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/document_reference/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/document_reference/list.rb
index 87ff400ca365..776cf76c8cf7 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/document_reference/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/document_reference/list.rb
@@ -109,7 +109,7 @@ def next
           #
           # @return [Enumerator]
           #
-          # @example Iterating each document reference by passing a block:
+          # @example Iterating each document reference by passing a block or proc:
           #   require "google/cloud/firestore"
           #
           #   firestore = Google::Cloud::Firestore.new
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 5bd38475efa4..7e226a191c5c 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -1105,7 +1105,7 @@ def values_to_cursor values, query
           order_field_paths = order_by_field_paths query
           if values.count > order_field_paths.count
             # raise if too many values provided for the cursor
-            raise ArgumentError, "There cannot be more values than Order By fields"
+            raise ArgumentError, "There cannot be more cursor values than order by fields"
           end
 
           values = values.zip(order_field_paths).map do |value, field_path|
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 071013b161f9..84d71dcf2a74 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -26,15 +26,20 @@ module Firestore
       # The cursors returned by {#start_at} and {#end_before} can only be used in a query that matches the constraint of
       # the query that produced this partition.
       #
-      # See {CollectionGroup#partitions}.
+      # See {CollectionGroup#partitions} and {Query}.
       #
       # @!attribute [r] start_at
-      #   The cursor that defines the first result for this partition, or `nil` if this is the first partition.
-      #   @return [Array] a cursor value that can be used with {Query#start_at} or `nil` if this is the first partition.
+      #   The cursor values that define the first result for this partition, or `nil` if this is the first partition.
+      #   Returns an array of values that represent a position, in the order they appear in the order by clause of the
+      #   query. Can contain fewer values than specified in the order by clause. Will be used in the query returned by
+      #   {#create_query}.
+      #   @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
       # @!attribute [r] end_before
-      #   The cursor that defines the first result after this partition, or `nil` if this is the last partition.
-      #   @return [Array] a cursor value that can be used with {Query#end_before} or `nil` if this is the last
-      #     partition.
+      #   The cursor values that define the first result after this partition, or `nil` if this is the last partition.
+      #   Returns an array of values that represent a position, in the order they appear in the order by clause of the
+      #   query.  Can contain fewer values than specified in the order by clause. Will be used in the query returned by
+      #   {#create_query}.
+      #   @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
       #
       # @example
       #   require "google/cloud/firestore"
@@ -60,9 +65,10 @@ def initialize query, start_at, end_before
         end
 
         ##
-        # Creates a query that only returns the documents for this partition.
+        # Creates a new query that only returns the documents for this partition, using the cursor valuess from
+        # {#start_at} and {#end_before}.
         #
-        # @return [Query] query.
+        # @return [Query] The query for the partition.
         #
         def create_query
           base_query = @query
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
index 4b66a941a655..135fc0fefb2e 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
@@ -47,7 +47,7 @@ def initialize arr = []
           end
 
           ##
-          # Whether there is a next page of document references.
+          # Whether there is a next page of query partitions.
           #
           # @return [Boolean]
           #
@@ -60,7 +60,7 @@ def initialize arr = []
           #
           #   partitions = col_group.partitions 3
           #   if partitions.next?
-          #     next_documents = partitions.next
+          #     next_partitions = partitions.next
           #   end
           #
           def next?
@@ -68,7 +68,7 @@ def next?
           end
 
           ##
-          # Retrieve the next page of document references.
+          # Retrieve the next page of query partitions.
           #
           # @return [QueryPartition::List]
           #
@@ -81,14 +81,14 @@ def next?
           #
           #   partitions = col_group.partitions 3
           #   if partitions.next?
-          #     next_documents = partitions.next
+          #     next_partitions = partitions.next
           #   end
           #
           def next
             return nil unless next?
             ensure_client!
             grpc = @client.service.partition_query @parent, @query.query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max
+            self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max, start_at: @start_at
           end
 
           ##
@@ -103,14 +103,13 @@ def next
           # over the results returned by a single API call.) Use with caution.
           #
           # @param [Integer] request_limit The upper limit of API requests to
-          #   make to load all document references. Default is no limit.
-          # @yield [document] The block for accessing each document.
-          # @yieldparam [QueryPartition] document The document reference
-          #   object.
+          #   make to load all query partitions. Default is no limit.
+          # @yield [partition] The block for accessing each partition.
+          # @yieldparam [QueryPartition] partition The query partition object.
           #
           # @return [Enumerator]
           #
-          # @example Iterating each document reference by passing a block:
+          # @example Iterating each query partition by passing a block or proc:
           #   require "google/cloud/firestore"
           #
           #   firestore = Google::Cloud::Firestore.new
@@ -119,7 +118,7 @@ def next
           #
           #   partitions = col_group.partitions 3
           #
-          #   queries = partitions.map(&:create_query)
+          #   queries = partitions.all(&:create_query)
           #
           # @example Using the enumerator by not passing a block:
           #   require "google/cloud/firestore"
@@ -130,9 +129,9 @@ def next
           #
           #   partitions = col_group.partitions 3
           #
-          #   queries = partitions.map(&:create_query)
+          #   queries = partitions.all.map(&:create_query)
           #
-          # @example Limit the number of API calls made:
+          # @example Limit the number of API calls made by `#all`:
           #   require "google/cloud/firestore"
           #
           #   firestore = Google::Cloud::Firestore.new
@@ -141,7 +140,7 @@ def next
           #
           #   partitions = col_group.partitions 3
           #
-          #   queries = partitions.map(&:create_query)
+          #   queries = partitions.all(request_limit: 10, &:create_query)
           #
           def all request_limit: nil, &block
             request_limit = request_limit.to_i if request_limit
@@ -163,9 +162,10 @@ def all request_limit: nil, &block
           ##
           # @private New QueryPartition::List from a
           # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
-          def self.from_grpc grpc, client, parent, query, partition_count, max: nil
-            # TODO: Should the logic adding the last partition be applied to the entire result set, not the page?
-            start_at = nil
+          def self.from_grpc grpc, client, parent, query, partition_count, max: nil, start_at: nil
+            token = grpc.next_page_token
+            token = nil if token == ""
+
             partitions = List.new(Array(grpc.partitions).map do |cursor|
               end_before = cursor.values.map do |value|
                 Convert.value_to_raw value, client
@@ -174,16 +174,15 @@ def self.from_grpc grpc, client, parent, query, partition_count, max: nil
               start_at = end_before
               partition
             end)
-            # We are always returning an extra partition (with en empty endBefore cursor).
-            partitions << QueryPartition.new(query, start_at, nil)
+            # Return an extra partition (with en empty endBefore cursor) at the end of the results.
+            partitions << QueryPartition.new(query, start_at, nil) unless token
             partitions.instance_variable_set :@parent, parent
             partitions.instance_variable_set :@query, query
             partitions.instance_variable_set :@partition_count, partition_count
-            token = grpc.next_page_token
-            token = nil if token == ""
             partitions.instance_variable_set :@token, token
             partitions.instance_variable_set :@client, client
             partitions.instance_variable_set :@max, max
+            partitions.instance_variable_set :@start_at, start_at
             partitions
           end
 
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index fc9e549a9e53..58f750528f46 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -20,20 +20,7 @@
     Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
   end
   let(:expected_query) do
-    Google::Cloud::Firestore::V1::StructuredQuery.new(
-      from: [
-        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
-          collection_id: "my-collection-id",
-          all_descendants: true
-        )
-      ],
-      order_by: [
-        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
-          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "__name__"),
-          direction: :ASCENDING
-        )
-      ]
-    )
+    collection_group_query
   end
 
   it "lists partitions" do
@@ -52,20 +39,50 @@
 
     _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
     _(partitions[0].start_at).must_be :nil?
-    _(partitions[0].end_before).must_equal ["128"]
+    _(partitions[0].end_before).must_be_kind_of Array
+    _(partitions[0].end_before.count).must_equal 1 # The array should have an element for each field in Order By.
+    _(partitions[0].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
+    _(partitions[0].end_before[0].path).must_equal document_path("10")
+
+    _(partitions[1]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[1].start_at).must_be_kind_of Array
+    _(partitions[1].start_at.count).must_equal 1 # The array should have an element for each field in Order By.
+    _(partitions[1].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
+    _(partitions[1].start_at[0].path).must_equal document_path("10")
+    _(partitions[1].end_before).must_be_kind_of Array
+    _(partitions[1].end_before.count).must_equal 1 # The array should have an element for each field in Order By.
+    _(partitions[1].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
+    _(partitions[1].end_before[0].path).must_equal document_path("20")
+
+    _(partitions[2]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[2].start_at).must_be_kind_of Array
+    _(partitions[2].start_at.count).must_equal 1 # The array should have an element for each field in Order By.
+    _(partitions[2].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
+    _(partitions[2].start_at[0].path).must_equal document_path("20")
+    _(partitions[2].end_before).must_be :nil?
+
     query_1 = partitions[0].create_query
-    # TODO: query assertions
+    _(query_1).must_be_kind_of Google::Cloud::Firestore::Query
+    _(query_1.query).must_equal collection_group_query(end_before: ["10"])
+
+    query_2 = partitions[1].create_query
+    _(query_2).must_be_kind_of Google::Cloud::Firestore::Query
+    _(query_2.query).must_equal collection_group_query(start_at: ["10"], end_before: ["20"])
+
+    query_3 = partitions[2].create_query
+    _(query_3).must_be_kind_of Google::Cloud::Firestore::Query
+    _(query_3.query).must_equal collection_group_query(start_at: ["20"])
   end
 
-  it "paginates partitions" do
-    first_list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 1)
+  it "paginates partitions with max" do
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
 
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
 
-    first_partitions = collection_group.partitions 6
-    second_partitions = collection_group.partitions 6, token: first_partitions.token
+    first_partitions = collection_group.partitions 6, max: 3
+    second_partitions = collection_group.partitions 6, max: 3, token: first_partitions.token
 
     firestore_mock.verify
 
@@ -74,18 +91,18 @@
     _(first_partitions.token).must_equal "next_page_token"
 
     second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(second_partitions.count).must_equal 2
+    _(second_partitions.count).must_equal 3
     _(second_partitions.token).must_be :nil?
   end
 
   it "paginates partitions using next" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 1)
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
 
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
 
-    first_partitions = collection_group.partitions 6
+    first_partitions = collection_group.partitions 6, max: 3
     second_partitions = first_partitions.next
 
     first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
@@ -93,63 +110,70 @@
     _(first_partitions.token).must_equal "next_page_token"
 
     second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(second_partitions.count).must_equal 2
+    _(second_partitions.count).must_equal 3
     _(second_partitions.token).must_be :nil?
   end
 
   it "paginates partitions using all" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 1)
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 2)
 
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
 
-    all_partitions = collection_group.partitions(6).all.to_a
+    all_partitions = collection_group.partitions(6, max: 3).all.to_a
 
     firestore_mock.verify
 
     all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(all_partitions.count).must_equal 5
+    _(all_partitions.count).must_equal 6
   end
 
   it "paginates partitions using all and Enumerator" do
     first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
     second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
 
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
 
-    all_partitions = collection_group.partitions(6).all.take 5
+    all_partitions = collection_group.partitions(6, max: 3).all.take 4
 
     firestore_mock.verify
 
     all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(all_partitions.count).must_equal 5
+    _(all_partitions.count).must_equal 4
   end
 
   it "paginates partitions using all with request_limit set" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 2, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2, token: "second_page_token")
+    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
+    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
 
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_token: "next_page_token")
+    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 8, page_size: 3)
+    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 8, page_size: 3, page_token: "next_page_token")
 
-    all_partitions = collection_group.partitions(6).all(request_limit: 1).to_a
+    all_partitions = collection_group.partitions(9, max: 3).all(request_limit: 1).to_a
 
     all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
     _(all_partitions.count).must_equal 6
   end
 
-  it "paginates partitions with max set" do
-    list_res = paged_enum_struct partition_query_resp(token: "next_page_token")
-    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
-
-    partitions = collection_group.partitions 6, max: 3
-
-    firestore_mock.verify
-
-    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
-    _(partitions.count).must_equal 3
-    partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
+  def collection_group_query start_at: nil, end_before: nil
+    query_grpc = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      from: [
+        Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(
+          collection_id: "my-collection-id",
+          all_descendants: true
+        )
+      ],
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "__name__"),
+          direction: :ASCENDING
+        )
+      ]
+    )
+    query_grpc.start_at = cursor_grpc(doc_ids: start_at) if start_at
+    query_grpc.end_at = cursor_grpc(doc_ids: end_before) if end_before
+    query_grpc
   end
 end
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index df251ad9fb41..6722b8048652 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -204,16 +204,16 @@ def partition_query_args query_grpc,
 
   def partition_query_resp count: 2, token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      # Minimum partition size is 128.
-      partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] },
+      partitions: count.times.map { |i| cursor_grpc doc_ids: [((i+1) * 10).to_s] },
       next_page_token: token
     )
   end
 
-  # Minimum partition size is 128.
-  def cursor_grpc values: ["128"], before: false
-    converted_values = values.map do |val|
-      Google::Cloud::Firestore::Convert.raw_to_value val
+  def cursor_grpc doc_ids: ["10"], before: true
+    converted_values = doc_ids.map do |doc_id|
+      Google::Cloud::Firestore::V1::Value.new(
+        reference_value: document_path(doc_id)
+      )
     end
     Google::Cloud::Firestore::V1::Cursor.new(
       values: converted_values,
@@ -224,6 +224,10 @@ def cursor_grpc values: ["128"], before: false
   def paged_enum_struct response
     OpenStruct.new response: response
   end
+
+  def document_path doc_id
+    "projects/#{project}/databases/(default)/documents/my-collection-id/#{doc_id}"
+  end
 end
 
 class WatchFirestore < MockFirestore

From ab513c3a7a73ee6f51893ea172d910c049ac7faf Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Fri, 28 May 2021 11:06:07 -0600
Subject: [PATCH 16/22] Update CollectionGroup#partitions to retrieve and sort
 all pages

---
 .../firestore/collection_group_test.rb        |  13 +-
 .../acceptance/firestore_helper.rb            |  20 +-
 .../cloud/firestore/collection_group.rb       |  53 +++--
 .../google/cloud/firestore/query_partition.rb |   2 -
 .../cloud/firestore/query_partition/list.rb   | 200 ------------------
 .../support/doctest_helper.rb                 |  14 +-
 .../cloud/firestore/collection_group_test.rb  |  85 +-------
 7 files changed, 65 insertions(+), 322 deletions(-)
 delete mode 100644 google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
index b9435c348ade..0787a6e45460 100644
--- a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -131,9 +131,9 @@
 
       collection_group = firestore.collection_group(rand_col.collection_id)
 
-      partitions = collection_group.partitions 6, max: 1
-      _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
-      _(partitions.count).must_equal 1
+      partitions = collection_group.partitions 6
+      _(partitions).must_be_kind_of Array
+      _(partitions.count).must_equal 3
 
       _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
       _(partitions[0].start_at).must_be :nil?
@@ -141,10 +141,6 @@
       _(partitions[0].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
       _(document_ids).must_include partitions[0].end_before[0].document_id
 
-      partitions += partitions.next
-      _(partitions).must_be_kind_of Array
-      _(partitions.count).must_equal 3
-
       _(partitions[1]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
       _(partitions[1].start_at).must_be_kind_of Array
       _(partitions[1].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
@@ -153,6 +149,9 @@
       _(partitions[1].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
       _(document_ids).must_include partitions[1].end_before[0].document_id
 
+      # Verify that partitions are sorted ascending order
+      _(partitions[0].end_before[0].document_id).must_be :<, partitions[1].end_before[0].document_id
+
       _(partitions[2]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
       _(partitions[2].start_at).must_be_kind_of Array
       _(partitions[2].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
diff --git a/google-cloud-firestore/acceptance/firestore_helper.rb b/google-cloud-firestore/acceptance/firestore_helper.rb
index af21a6405b82..4b9d678dd745 100644
--- a/google-cloud-firestore/acceptance/firestore_helper.rb
+++ b/google-cloud-firestore/acceptance/firestore_helper.rb
@@ -65,16 +65,16 @@ def root_col
       addl.include? :firestore_acceptance
     end
 
-    def self.run_one_method klass, method_name, reporter
-      result = nil
-      reporter.prerecord klass, method_name
-      (1..3).each do |try|
-        result = Minitest.run_one_method(klass, method_name)
-        break if (result.passed? || result.skipped?)
-        puts "Retrying #{klass}##{method_name} (#{try})"
-      end
-      reporter.record result
-    end
+    # def self.run_one_method klass, method_name, reporter
+    #   result = nil
+    #   reporter.prerecord klass, method_name
+    #   (1..3).each do |try|
+    #     result = Minitest.run_one_method(klass, method_name)
+    #     break if (result.passed? || result.skipped?)
+    #     puts "Retrying #{klass}##{method_name} (#{try})"
+    #   end
+    #   reporter.record result
+    # end
   end
 end
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index fe333a81735a..fee809786504 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -48,10 +48,8 @@ class CollectionGroup < Query
         #
         # @param [Integer] partition_count The desired maximum number of partition points. The number must be strictly
         #   positive. The actual number of partitions returned may be fewer.
-        # @param [String] token A previously-returned page token representing part of the larger set of results to view.
-        # @param [Integer] max Maximum number of results to return.
         #
-        # @return [Array<QueryPartition>] An array of query partitions.
+        # @return [Array<QueryPartition>] An ordered array of query partitions.
         #
         # @example
         #   require "google/cloud/firestore"
@@ -64,7 +62,7 @@ class CollectionGroup < Query
         #
         #   queries = partitions.map(&:create_query)
         #
-        def partitions partition_count, token: nil, max: nil
+        def partitions partition_count
           ensure_service!
 
           # Partition queries require explicit ordering by __name__.
@@ -73,18 +71,28 @@ def partitions partition_count, token: nil, max: nil
           # partition count by one.
           partition_count -= 1
 
-          grpc = service.partition_query parent_path,
-                                         query_with_default_order.query,
-                                         partition_count,
-                                         token: token,
-                                         max: max
+          # Retrieve all pages of the results, because order is not guaranteed and they must be sorted.
+          grpc_partitions = list_all partition_count, query_with_default_order
+          cursor_values = grpc_partitions.map do |cursor|
+            # Convert each cursor to a (single-element) array of Google::Cloud::Firestore::DocumentReference.
+            cursor.values.map do |value|
+              Convert.value_to_raw value, client
+            end
+          end
+          # Sort the values of the returned cursor, which right now should only contain a single reference value (which
+          # needs to be sorted one component at a time).
+          cursor_values.sort! do |a, b|
+            a.first.path <=> b.first.path
+          end
 
-          QueryPartition::List.from_grpc grpc,
-                                         client,
-                                         parent_path,
-                                         query_with_default_order,
-                                         partition_count,
-                                         max: max
+          start_at = nil
+          results = cursor_values.map do |end_before|
+            partition = QueryPartition.new query_with_default_order, start_at, end_before
+            start_at = end_before
+            partition
+          end
+          results << QueryPartition.new(query_with_default_order, start_at, nil)
+          results
         end
 
         ##
@@ -100,6 +108,21 @@ def self.from_collection_id parent_path, collection_id, client
           )
           CollectionGroup.new query, parent_path, client
         end
+
+        protected
+
+        def list_all partition_count, query_with_default_order
+          grpc_partitions = []
+          token = nil
+          loop do
+            grpc = service.partition_query parent_path, query_with_default_order.query, partition_count, token: token
+            grpc_partitions += Array(grpc.partitions)
+            token = grpc.next_page_token
+            token = nil if token == ""
+            break unless token
+          end
+          grpc_partitions
+        end
       end
     end
   end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 84d71dcf2a74..fd8e9c5bf0fd 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -13,8 +13,6 @@
 # limitations under the License.
 
 
-require "google/cloud/firestore/query_partition/list"
-
 module Google
   module Cloud
     module Firestore
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
deleted file mode 100644
index 135fc0fefb2e..000000000000
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb
+++ /dev/null
@@ -1,200 +0,0 @@
-# Copyright 2021 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-require "delegate"
-
-module Google
-  module Cloud
-    module Firestore
-      class QueryPartition
-        ##
-        # QueryPartition::List is a special case Array with additional values.
-        #
-        # @example
-        #   require "google/cloud/firestore"
-        #
-        #   firestore = Google::Cloud::Firestore.new
-        #
-        #   col_group = firestore.col_group "cities"
-        #
-        #   partitions = col_group.partitions 3
-        #
-        #   queries = partitions.map(&:create_query)
-        #
-        class List < DelegateClass(::Array)
-          ##
-          # If not empty, indicates that there are more records that match
-          # the request and this value should be passed to continue.
-          attr_accessor :token
-
-          ##
-          # @private Create a new QueryPartition::List with an array of
-          # QueryPartition instances.
-          def initialize arr = []
-            super arr
-          end
-
-          ##
-          # Whether there is a next page of query partitions.
-          #
-          # @return [Boolean]
-          #
-          # @example
-          #   require "google/cloud/firestore"
-          #
-          #   firestore = Google::Cloud::Firestore.new
-          #
-          #   col_group = firestore.col_group "cities"
-          #
-          #   partitions = col_group.partitions 3
-          #   if partitions.next?
-          #     next_partitions = partitions.next
-          #   end
-          #
-          def next?
-            !token.nil?
-          end
-
-          ##
-          # Retrieve the next page of query partitions.
-          #
-          # @return [QueryPartition::List]
-          #
-          # @example
-          #   require "google/cloud/firestore"
-          #
-          #   firestore = Google::Cloud::Firestore.new
-          #
-          #   col_group = firestore.col_group "cities"
-          #
-          #   partitions = col_group.partitions 3
-          #   if partitions.next?
-          #     next_partitions = partitions.next
-          #   end
-          #
-          def next
-            return nil unless next?
-            ensure_client!
-            grpc = @client.service.partition_query @parent, @query.query, @partition_count, token: token, max: @max
-            self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max, start_at: @start_at
-          end
-
-          ##
-          # Retrieves remaining results by repeatedly invoking {#next} until
-          # {#next?} returns `false`. Calls the given block once for each
-          # result, which is passed as the argument to the block.
-          #
-          # An Enumerator is returned if no block is given.
-          #
-          # This method will make repeated API calls until all remaining results
-          # are retrieved. (Unlike `#each`, for example, which merely iterates
-          # over the results returned by a single API call.) Use with caution.
-          #
-          # @param [Integer] request_limit The upper limit of API requests to
-          #   make to load all query partitions. Default is no limit.
-          # @yield [partition] The block for accessing each partition.
-          # @yieldparam [QueryPartition] partition The query partition object.
-          #
-          # @return [Enumerator]
-          #
-          # @example Iterating each query partition by passing a block or proc:
-          #   require "google/cloud/firestore"
-          #
-          #   firestore = Google::Cloud::Firestore.new
-          #
-          #   col_group = firestore.col_group "cities"
-          #
-          #   partitions = col_group.partitions 3
-          #
-          #   queries = partitions.all(&:create_query)
-          #
-          # @example Using the enumerator by not passing a block:
-          #   require "google/cloud/firestore"
-          #
-          #   firestore = Google::Cloud::Firestore.new
-          #
-          #   col_group = firestore.col_group "cities"
-          #
-          #   partitions = col_group.partitions 3
-          #
-          #   queries = partitions.all.map(&:create_query)
-          #
-          # @example Limit the number of API calls made by `#all`:
-          #   require "google/cloud/firestore"
-          #
-          #   firestore = Google::Cloud::Firestore.new
-          #
-          #   col_group = firestore.col_group "cities"
-          #
-          #   partitions = col_group.partitions 3
-          #
-          #   queries = partitions.all(request_limit: 10, &:create_query)
-          #
-          def all request_limit: nil, &block
-            request_limit = request_limit.to_i if request_limit
-            unless block_given?
-              return enum_for :all, request_limit: request_limit
-            end
-            results = self
-            loop do
-              results.each(&block)
-              if request_limit
-                request_limit -= 1
-                break if request_limit.negative?
-              end
-              break unless results.next?
-              results = results.next
-            end
-          end
-
-          ##
-          # @private New QueryPartition::List from a
-          # Google::Cloud::Firestore::V1::PartitionQueryResponse object.
-          def self.from_grpc grpc, client, parent, query, partition_count, max: nil, start_at: nil
-            token = grpc.next_page_token
-            token = nil if token == ""
-
-            partitions = List.new(Array(grpc.partitions).map do |cursor|
-              end_before = cursor.values.map do |value|
-                Convert.value_to_raw value, client
-              end
-              partition = QueryPartition.new query, start_at, end_before
-              start_at = end_before
-              partition
-            end)
-            # Return an extra partition (with en empty endBefore cursor) at the end of the results.
-            partitions << QueryPartition.new(query, start_at, nil) unless token
-            partitions.instance_variable_set :@parent, parent
-            partitions.instance_variable_set :@query, query
-            partitions.instance_variable_set :@partition_count, partition_count
-            partitions.instance_variable_set :@token, token
-            partitions.instance_variable_set :@client, client
-            partitions.instance_variable_set :@max, max
-            partitions.instance_variable_set :@start_at, start_at
-            partitions
-          end
-
-          protected
-
-          ##
-          # Raise an error unless an active client is available.
-          def ensure_client!
-            raise "Must have active connection" unless @client
-          end
-        end
-      end
-    end
-  end
-end
diff --git a/google-cloud-firestore/support/doctest_helper.rb b/google-cloud-firestore/support/doctest_helper.rb
index 5d954e206195..803fe40606b9 100644
--- a/google-cloud-firestore/support/doctest_helper.rb
+++ b/google-cloud-firestore/support/doctest_helper.rb
@@ -517,6 +517,10 @@ def list_documents_args
   ]
 end
 
+def document_path doc_id
+  "projects/my-project-id/databases/(default)/documents/#{doc_id}"
+end
+
 def document_gapi doc: "my-document", fields: {}
   Google::Cloud::Firestore::V1::Document.new(
     name: "projects/my-project-id/databases/(default)/documents/#{doc}",
@@ -537,16 +541,18 @@ def documents_resp token: nil
 def partition_query_resp count: 3, token: nil
   response = Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
     # Minimum partition size is 128.
-    partitions: count.times.map { |i| cursor_grpc values: [(i+1*128).to_s] }
+    partitions: count.times.map { |i| cursor_grpc doc_ids: [((i+1) * 10).to_s] }
   )
   response.next_page_token = token if token
   paged_enum_struct response
 end
 
 # Minimum partition size is 128.
-def cursor_grpc values: ["128"], before: false
-  converted_values = values.map do |val|
-    Google::Cloud::Firestore::Convert.raw_to_value val
+def cursor_grpc doc_ids: ["10"], before: true
+  converted_values = doc_ids.map do |doc_id|
+    Google::Cloud::Firestore::V1::Value.new(
+      reference_value: document_path(doc_id)
+    )
   end
   Google::Cloud::Firestore::V1::Cursor.new(
     values: converted_values,
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index 58f750528f46..a85abef5764c 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -34,7 +34,7 @@
 
     firestore_mock.verify
 
-    _(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
+    _(partitions).must_be_kind_of Array
     _(partitions.count).must_equal 3
 
     _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
@@ -74,89 +74,6 @@
     _(query_3.query).must_equal collection_group_query(start_at: ["20"])
   end
 
-  it "paginates partitions with max" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
-
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
-
-    first_partitions = collection_group.partitions 6, max: 3
-    second_partitions = collection_group.partitions 6, max: 3, token: first_partitions.token
-
-    firestore_mock.verify
-
-    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(first_partitions.count).must_equal 3
-    _(first_partitions.token).must_equal "next_page_token"
-
-    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(second_partitions.count).must_equal 3
-    _(second_partitions.token).must_be :nil?
-  end
-
-  it "paginates partitions using next" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
-
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
-
-    first_partitions = collection_group.partitions 6, max: 3
-    second_partitions = first_partitions.next
-
-    first_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(first_partitions.count).must_equal 3
-    _(first_partitions.token).must_equal "next_page_token"
-
-    second_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(second_partitions.count).must_equal 3
-    _(second_partitions.token).must_be :nil?
-  end
-
-  it "paginates partitions using all" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 2)
-
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
-
-    all_partitions = collection_group.partitions(6, max: 3).all.to_a
-
-    firestore_mock.verify
-
-    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(all_partitions.count).must_equal 6
-  end
-
-  it "paginates partitions using all and Enumerator" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
-
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 5, page_size: 3, page_token: "next_page_token")
-
-    all_partitions = collection_group.partitions(6, max: 3).all.take 4
-
-    firestore_mock.verify
-
-    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(all_partitions.count).must_equal 4
-  end
-
-  it "paginates partitions using all with request_limit set" do
-    first_list_res = paged_enum_struct partition_query_resp(count: 3, token: "next_page_token")
-    second_list_res = paged_enum_struct partition_query_resp(count: 3, token: "second_page_token")
-
-    firestore_mock.expect :partition_query, first_list_res, partition_query_args(expected_query, partition_count: 8, page_size: 3)
-    firestore_mock.expect :partition_query, second_list_res, partition_query_args(expected_query, partition_count: 8, page_size: 3, page_token: "next_page_token")
-
-    all_partitions = collection_group.partitions(9, max: 3).all(request_limit: 1).to_a
-
-    all_partitions.each { |m| _(m).must_be_kind_of Google::Cloud::Firestore::QueryPartition }
-    _(all_partitions.count).must_equal 6
-  end
-
   def collection_group_query start_at: nil, end_before: nil
     query_grpc = Google::Cloud::Firestore::V1::StructuredQuery.new(
       from: [

From f9a8d5ca6cbdd61fb7184e4a7dcaaca7cb982c0e Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Mon, 31 May 2021 16:31:34 -0600
Subject: [PATCH 17/22] Add ResourcePath and update sort in
 CollectionGroup#partitions

---
 .../cloud/firestore/collection_group.rb       |  2 +-
 .../cloud/firestore/document_reference.rb     |  7 +++
 .../google/cloud/firestore/resource_path.rb   | 58 +++++++++++++++++++
 .../cloud/firestore/collection_group_test.rb  | 24 ++++----
 .../cloud/firestore/resource_path_test.rb     | 58 +++++++++++++++++++
 google-cloud-firestore/test/helper.rb         |  8 +--
 6 files changed, 138 insertions(+), 19 deletions(-)
 create mode 100644 google-cloud-firestore/lib/google/cloud/firestore/resource_path.rb
 create mode 100644 google-cloud-firestore/test/google/cloud/firestore/resource_path_test.rb

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index fee809786504..e861cf4ceb74 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -82,7 +82,7 @@ def partitions partition_count
           # Sort the values of the returned cursor, which right now should only contain a single reference value (which
           # needs to be sorted one component at a time).
           cursor_values.sort! do |a, b|
-            a.first.path <=> b.first.path
+            a.first <=> b.first
           end
 
           start_at = nil
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/document_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/document_reference.rb
index c5da03ef4283..940230a88390 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/document_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/document_reference.rb
@@ -18,6 +18,7 @@
 require "google/cloud/firestore/collection_reference"
 require "google/cloud/firestore/document_listener"
 require "google/cloud/firestore/document_reference/list"
+require "google/cloud/firestore/resource_path"
 
 module Google
   module Cloud
@@ -459,6 +460,12 @@ def delete exists: nil, update_time: nil
 
         # @!endgroup
 
+        # @private
+        def <=> other
+          return nil unless other.is_a? self.class
+          ResourcePath.from_path(path) <=> ResourcePath.from_path(other.path)
+        end
+
         ##
         # @private New DocumentReference object from a path.
         def self.from_path path, client
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/resource_path.rb b/google-cloud-firestore/lib/google/cloud/firestore/resource_path.rb
new file mode 100644
index 000000000000..c5d4d1fb0e14
--- /dev/null
+++ b/google-cloud-firestore/lib/google/cloud/firestore/resource_path.rb
@@ -0,0 +1,58 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+module Google
+  module Cloud
+    module Firestore
+      ##
+      # @private
+      #
+      # Represents a resource path to the Firestore API.
+      #
+      class ResourcePath
+        include Comparable
+
+        RESOURCE_PATH_RE = %r{^projects/([^/]*)/databases/([^/]*)(?:/documents/)?([\s\S]*)$}.freeze
+
+        attr_reader :project_id
+        attr_reader :database_id
+        attr_reader :segments
+
+        ##
+        # Creates a resource path object.
+        #
+        # @param [Array<String>] segments One or more strings representing the resource path.
+        #
+        # @return [ResourcePath] The resource path object.
+        #
+        def initialize project_id, database_id, segments
+          @project_id = project_id
+          @database_id = database_id
+          @segments = segments.split "/"
+        end
+
+        def <=> other
+          return nil unless other.is_a? ResourcePath
+          [project_id, database_id, segments] <=> [other.project_id, other.database_id, other.segments]
+        end
+
+        def self.from_path path
+          data = RESOURCE_PATH_RE.match path
+          new data[1], data[2], data[3]
+        end
+      end
+    end
+  end
+end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index a85abef5764c..a7e4a145594a 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -23,11 +23,11 @@
     collection_group_query
   end
 
-  it "lists partitions" do
-    # grpc.partitions: 
-    # [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/32">], before: false>,
-    #  <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: reference_value: "projects/buoyant-ability-182823/databases/(default)/documents/gcloud-2021-05-21t23-38-03z-a581db45/query/dabf1479/89">], before: false>]
-    list_res = paged_enum_struct partition_query_resp
+  it "sorts and lists partitions" do
+    # Results should be sorted so that "alice" comes before "alice-"
+    # Use an ID ending in "-" to ensure correct sorting, since full path strings are sorted dash before slash
+    # See Google::Cloud::Firestore::ResourcePath
+    list_res = paged_enum_struct partition_query_resp(doc_ids: ["alice-", "alice"])
     firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
 
     partitions = collection_group.partitions 3
@@ -42,36 +42,36 @@
     _(partitions[0].end_before).must_be_kind_of Array
     _(partitions[0].end_before.count).must_equal 1 # The array should have an element for each field in Order By.
     _(partitions[0].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
-    _(partitions[0].end_before[0].path).must_equal document_path("10")
+    _(partitions[0].end_before[0].path).must_equal document_path("alice")
 
     _(partitions[1]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
     _(partitions[1].start_at).must_be_kind_of Array
     _(partitions[1].start_at.count).must_equal 1 # The array should have an element for each field in Order By.
     _(partitions[1].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
-    _(partitions[1].start_at[0].path).must_equal document_path("10")
+    _(partitions[1].start_at[0].path).must_equal document_path("alice")
     _(partitions[1].end_before).must_be_kind_of Array
     _(partitions[1].end_before.count).must_equal 1 # The array should have an element for each field in Order By.
     _(partitions[1].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
-    _(partitions[1].end_before[0].path).must_equal document_path("20")
+    _(partitions[1].end_before[0].path).must_equal document_path("alice-")
 
     _(partitions[2]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
     _(partitions[2].start_at).must_be_kind_of Array
     _(partitions[2].start_at.count).must_equal 1 # The array should have an element for each field in Order By.
     _(partitions[2].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
-    _(partitions[2].start_at[0].path).must_equal document_path("20")
+    _(partitions[2].start_at[0].path).must_equal document_path("alice-")
     _(partitions[2].end_before).must_be :nil?
 
     query_1 = partitions[0].create_query
     _(query_1).must_be_kind_of Google::Cloud::Firestore::Query
-    _(query_1.query).must_equal collection_group_query(end_before: ["10"])
+    _(query_1.query).must_equal collection_group_query(end_before: ["alice"])
 
     query_2 = partitions[1].create_query
     _(query_2).must_be_kind_of Google::Cloud::Firestore::Query
-    _(query_2.query).must_equal collection_group_query(start_at: ["10"], end_before: ["20"])
+    _(query_2.query).must_equal collection_group_query(start_at: ["alice"], end_before: ["alice-"])
 
     query_3 = partitions[2].create_query
     _(query_3).must_be_kind_of Google::Cloud::Firestore::Query
-    _(query_3.query).must_equal collection_group_query(start_at: ["20"])
+    _(query_3.query).must_equal collection_group_query(start_at: ["alice-"])
   end
 
   def collection_group_query start_at: nil, end_before: nil
diff --git a/google-cloud-firestore/test/google/cloud/firestore/resource_path_test.rb b/google-cloud-firestore/test/google/cloud/firestore/resource_path_test.rb
new file mode 100644
index 000000000000..34bdb70ee8fe
--- /dev/null
+++ b/google-cloud-firestore/test/google/cloud/firestore/resource_path_test.rb
@@ -0,0 +1,58 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require "helper"
+
+describe Google::Cloud::Firestore::ResourcePath do
+  let(:project_id) { "my-project" }
+  let(:database_id) { "my-database" }
+  let(:doc_id_a) { "alice" }
+  let(:doc_id_b) { "alice-" }
+  let(:doc_id_c) { "bob" }
+  let(:path_a) { "projects/#{project_id}/databases/(default)/documents/users/#{doc_id_a}/pets" }
+  let(:path_b) { "projects/#{project_id}/databases/(default)/documents/users/#{doc_id_b}/pets" }
+  let(:path_c) { "projects/#{project_id}/databases/(default)/documents/users/#{doc_id_c}/pets" }
+  let(:path_d) { "projects/#{project_id}/databases/(default)/documents/users/#{doc_id_c}" }
+  let(:resource_path_a) { Google::Cloud::Firestore::ResourcePath.from_path path_a }
+  let(:resource_path_b) { Google::Cloud::Firestore::ResourcePath.from_path path_b }
+  let(:resource_path_c) { Google::Cloud::Firestore::ResourcePath.from_path path_c }
+  let(:resource_path_d) { Google::Cloud::Firestore::ResourcePath.from_path path_d }
+
+  describe String, "<=>" do
+    it "correctly compares isolated resource IDs" do
+      # "alice" should come before "alice-"
+      _(doc_id_a.<=>(doc_id_b)).must_equal -1
+    end
+    it "incorrectly compares the resource ID, because full paths are sorted dash before slash" do
+      # "alice-" should not come before "alice"
+      _(path_a.<=>(path_b)).must_equal 1
+    end
+  end
+
+  it "parses project_id" do
+    _(resource_path_a.project_id).must_equal project_id
+  end
+
+  it "parses default database_id" do
+    _(resource_path_a.database_id).must_equal "(default)"
+  end
+
+  it "compares a dashed document ID reference" do
+    _(resource_path_a.<=>(resource_path_b)).must_equal -1
+  end
+
+  it "compares a simple document ID reference" do
+    _(resource_path_a.<=>(resource_path_c)).must_equal -1
+  end
+end
diff --git a/google-cloud-firestore/test/helper.rb b/google-cloud-firestore/test/helper.rb
index 6722b8048652..14b169b0f723 100644
--- a/google-cloud-firestore/test/helper.rb
+++ b/google-cloud-firestore/test/helper.rb
@@ -182,10 +182,6 @@ def run_query_args query,
     [req, default_options]
   end
 
-  # expected partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>"next_page_token", :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "second_page_token">>
-  # got     [partition_query({:parent=>"projects/projectID/databases/(default)/documents", :structured_query=><Google::Cloud::Firestore::V1::StructuredQuery: from: [<Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector: collection_id: "my-collection-id", all_descendants: true>], order_by: [<Google::Cloud::Firestore::V1::StructuredQuery::Order: field: <Google::Cloud::Firestore::V1::StructuredQuery::FieldReference: field_path: "__name__">, direction: :ASCENDING>], offset: 0>, :partition_count=>5, :page_token=>nil, :page_size=>nil}) => #<OpenStruct response=<Google::Cloud::Firestore::V1::PartitionQueryResponse: partitions: [<Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 0>, <Google::Cloud::Firestore::V1::Value: integer_value: 1>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 10>, <Google::Cloud::Firestore::V1::Value: integer_value: 11>], before: false>, <Google::Cloud::Firestore::V1::Cursor: values: [<Google::Cloud::Firestore::V1::Value: integer_value: 20>, <Google::Cloud::Firestore::V1::Value: integer_value: 21>], before: false>], next_page_token: "next_page_token">>]
-
-
   def partition_query_args query_grpc,
                            parent: "projects/#{project}/databases/(default)/documents",
                            partition_count: 2,
@@ -202,9 +198,9 @@ def partition_query_args query_grpc,
     ]
   end
 
-  def partition_query_resp count: 2, token: nil
+  def partition_query_resp doc_ids: ["10", "20"], token: nil
     Google::Cloud::Firestore::V1::PartitionQueryResponse.new(
-      partitions: count.times.map { |i| cursor_grpc doc_ids: [((i+1) * 10).to_s] },
+      partitions: doc_ids.map { |id| cursor_grpc doc_ids: [id] },
       next_page_token: token
     )
   end

From 85c5ffa66073009096b184bb6f2e1e7965607be0 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 1 Jun 2021 12:16:02 -0600
Subject: [PATCH 18/22] Rename QueryPartition#create_query to #to_query

* Fix return type docs for QueryPartition.
---
 .../acceptance/firestore/collection_group_test.rb    |  2 +-
 .../lib/google/cloud/firestore/collection_group.rb   |  2 +-
 .../lib/google/cloud/firestore/query_partition.rb    | 12 ++++++------
 .../google/cloud/firestore/collection_group_test.rb  |  6 +++---
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
index 0787a6e45460..dc0a943c8856 100644
--- a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -158,7 +158,7 @@
       _(document_ids).must_include partitions[2].start_at[0].document_id
       _(partitions[2].end_before).must_be :nil?
 
-      queries = partitions.map(&:create_query)
+      queries = partitions.map(&:to_query)
       _(queries.count).must_equal 3
       results = queries.map do |query|
         _(query).must_be_kind_of Google::Cloud::Firestore::Query
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index e861cf4ceb74..1733e5c51f55 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -60,7 +60,7 @@ class CollectionGroup < Query
         #
         #   partitions = col_group.partitions 3
         #
-        #   queries = partitions.map(&:create_query)
+        #   queries = partitions.map(&:to_query)
         #
         def partitions partition_count
           ensure_service!
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index fd8e9c5bf0fd..0ae236f0120b 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -30,14 +30,14 @@ module Firestore
       #   The cursor values that define the first result for this partition, or `nil` if this is the first partition.
       #   Returns an array of values that represent a position, in the order they appear in the order by clause of the
       #   query. Can contain fewer values than specified in the order by clause. Will be used in the query returned by
-      #   {#create_query}.
-      #   @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
+      #   {#to_query}.
+      #   @return [Array<Object>, nil] Typically, the values are {DocumentReference} objects.
       # @!attribute [r] end_before
       #   The cursor values that define the first result after this partition, or `nil` if this is the last partition.
       #   Returns an array of values that represent a position, in the order they appear in the order by clause of the
       #   query.  Can contain fewer values than specified in the order by clause. Will be used in the query returned by
-      #   {#create_query}.
-      #   @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
+      #   {#to_query}.
+      #   @return [Array<Object>, nil] Typically, the values are {DocumentReference} objects.
       #
       # @example
       #   require "google/cloud/firestore"
@@ -48,7 +48,7 @@ module Firestore
       #
       #   partitions = col_group.partitions 3
       #
-      #   queries = partitions.map(&:create_query)
+      #   queries = partitions.map(&:to_query)
       #
       class QueryPartition
         attr_reader :start_at
@@ -68,7 +68,7 @@ def initialize query, start_at, end_before
         #
         # @return [Query] The query for the partition.
         #
-        def create_query
+        def to_query
           base_query = @query
           base_query = base_query.start_at start_at if start_at
           base_query = base_query.end_before end_before if end_before
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index a7e4a145594a..a54d9c4df23d 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -61,15 +61,15 @@
     _(partitions[2].start_at[0].path).must_equal document_path("alice-")
     _(partitions[2].end_before).must_be :nil?
 
-    query_1 = partitions[0].create_query
+    query_1 = partitions[0].to_query
     _(query_1).must_be_kind_of Google::Cloud::Firestore::Query
     _(query_1.query).must_equal collection_group_query(end_before: ["alice"])
 
-    query_2 = partitions[1].create_query
+    query_2 = partitions[1].to_query
     _(query_2).must_be_kind_of Google::Cloud::Firestore::Query
     _(query_2.query).must_equal collection_group_query(start_at: ["alice"], end_before: ["alice-"])
 
-    query_3 = partitions[2].create_query
+    query_3 = partitions[2].to_query
     _(query_3).must_be_kind_of Google::Cloud::Firestore::Query
     _(query_3.query).must_equal collection_group_query(start_at: ["alice-"])
   end

From 345bfa945f308f9c8fb0e07849034820396ea722 Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 1 Jun 2021 13:05:21 -0600
Subject: [PATCH 19/22] Verify QueryPartition#start_at and #end_before can be
 used with a new Query

---
 .../acceptance/firestore/collection_group_test.rb         | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
index dc0a943c8856..a4c30b353d25 100644
--- a/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/collection_group_test.rb
@@ -170,6 +170,14 @@
       results.each { |result| _(result).wont_be :empty? }
       # Verify all document IDs have been returned, in original order.
       _(results.flatten).must_equal document_ids
+
+      # Verify QueryPartition#start_at and #end_before can be used with a new Query.
+      query = collection_group.order("__name__").start_at(partitions[1].start_at).end_before(partitions[1].end_before)
+      result = query.get.map do |snp|
+        _(snp).must_be_kind_of Google::Cloud::Firestore::DocumentSnapshot
+        snp.document_id
+      end
+      _(result).must_equal results[1]
     end
   end
 end

From dc263495d94b02d65615878490b9b461e27111ce Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 1 Jun 2021 13:06:29 -0600
Subject: [PATCH 20/22] Ensure an empty partition is returned when 1 is
 requested or when RPC returns none

---
 .../cloud/firestore/collection_group.rb       | 11 ++++--
 .../cloud/firestore/collection_group_test.rb  | 35 +++++++++++++++++--
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index 1733e5c51f55..ba0ca1009aaf 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -65,14 +65,20 @@ class CollectionGroup < Query
         def partitions partition_count
           ensure_service!
 
+          raise ArgumentError, "partition_count must be > 0" unless partition_count.positive?
+
           # Partition queries require explicit ordering by __name__.
           query_with_default_order = order "__name__"
           # Since we are always returning an extra partition (with en empty endBefore cursor), we reduce the desired
           # partition count by one.
           partition_count -= 1
 
-          # Retrieve all pages of the results, because order is not guaranteed and they must be sorted.
-          grpc_partitions = list_all partition_count, query_with_default_order
+          grpc_partitions = if partition_count.positive?
+                              # Retrieve all pages, since cursor order is not guaranteed and they must be sorted.
+                              list_all partition_count, query_with_default_order
+                            else
+                              [] # Ensure that a single, empty QueryPartition is returned.
+                            end
           cursor_values = grpc_partitions.map do |cursor|
             # Convert each cursor to a (single-element) array of Google::Cloud::Firestore::DocumentReference.
             cursor.values.map do |value|
@@ -91,6 +97,7 @@ def partitions partition_count
             start_at = end_before
             partition
           end
+          # Always add a final QueryPartition with an empty end_before value.
           results << QueryPartition.new(query_with_default_order, start_at, nil)
           results
         end
diff --git a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
index a54d9c4df23d..b39bda61278e 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/collection_group_test.rb
@@ -19,8 +19,39 @@
   let(:collection_group) do
     Google::Cloud::Firestore::CollectionGroup.from_collection_id documents_path, collection_id, firestore
   end
-  let(:expected_query) do
-    collection_group_query
+  let(:expected_query) { collection_group_query }
+
+  it "raises if partition_count is < 1" do
+    expect do
+      partitions = collection_group.partitions 0
+    end.must_raise ArgumentError
+  end
+
+  it "returns 1 empty partition if partition_count is 1, without RPC" do
+    partitions = collection_group.partitions 1
+
+    _(partitions).must_be_kind_of Array
+    _(partitions.count).must_equal 1
+
+    _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[0].start_at).must_be :nil?
+    _(partitions[0].end_before).must_be :nil?
+  end
+
+  it "returns 1 empty partition if RPC returns no partitions" do
+    list_res = paged_enum_struct partition_query_resp(doc_ids: [])
+    firestore_mock.expect :partition_query, list_res, partition_query_args(expected_query)
+
+    partitions = collection_group.partitions 3
+
+    firestore_mock.verify
+
+    _(partitions).must_be_kind_of Array
+    _(partitions.count).must_equal 1
+
+    _(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
+    _(partitions[0].start_at).must_be :nil?
+    _(partitions[0].end_before).must_be :nil?
   end
 
   it "sorts and lists partitions" do

From 7463a033aa568d34a8a2208f281eebc75de7b6be Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Tue, 1 Jun 2021 15:52:47 -0600
Subject: [PATCH 21/22] Updates for dazuma comments

---
 .../acceptance/firestore_helper.rb            | 20 +++++++++----------
 .../cloud/firestore/collection_group.rb       |  4 ++--
 .../cloud/firestore/collection_reference.rb   |  6 +++---
 .../lib/google/cloud/firestore/query.rb       |  2 +-
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/google-cloud-firestore/acceptance/firestore_helper.rb b/google-cloud-firestore/acceptance/firestore_helper.rb
index 4b9d678dd745..af21a6405b82 100644
--- a/google-cloud-firestore/acceptance/firestore_helper.rb
+++ b/google-cloud-firestore/acceptance/firestore_helper.rb
@@ -65,16 +65,16 @@ def root_col
       addl.include? :firestore_acceptance
     end
 
-    # def self.run_one_method klass, method_name, reporter
-    #   result = nil
-    #   reporter.prerecord klass, method_name
-    #   (1..3).each do |try|
-    #     result = Minitest.run_one_method(klass, method_name)
-    #     break if (result.passed? || result.skipped?)
-    #     puts "Retrying #{klass}##{method_name} (#{try})"
-    #   end
-    #   reporter.record result
-    # end
+    def self.run_one_method klass, method_name, reporter
+      result = nil
+      reporter.prerecord klass, method_name
+      (1..3).each do |try|
+        result = Minitest.run_one_method(klass, method_name)
+        break if (result.passed? || result.skipped?)
+        puts "Retrying #{klass}##{method_name} (#{try})"
+      end
+      reporter.record result
+    end
   end
 end
 
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
index ba0ca1009aaf..ecc779188b89 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb
@@ -34,10 +34,10 @@ module Firestore
       #   firestore = Google::Cloud::Firestore.new
       #
       #   # Get a collection group
-      #   cities_col = firestore.col "cities"
+      #   col_group = firestore.col_group "cities"
       #
       #   # Get and print all city documents
-      #   cities_col.get do |city|
+      #   col_group.get do |city|
       #     puts "#{city.document_id} has #{city[:population]} residents."
       #   end
       #
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
index ab8a1a377a7b..684c2206548e 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/collection_reference.rb
@@ -49,8 +49,8 @@ class CollectionReference < Query
 
         ##
         # @private Creates a new CollectionReference.
-        def initialize path, client, query
-          super query, nil, client
+        def initialize query, path, client
+          super query, nil, client # Pass nil parent_path arg since this class implements #parent_path
           @path = path
         end
 
@@ -264,7 +264,7 @@ def self.from_path path, client
             ]
           )
 
-          CollectionReference.new path, client, query
+          CollectionReference.new query, path, client
         end
 
         protected
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 7e226a191c5c..8358df4288da 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -1222,7 +1222,7 @@ def doc_id_path
         ##
         # @private Raise an error unless an database available.
         def ensure_client!
-          raise "Must have active connection to client" unless client
+          raise "Must have active connection to service" unless client
         end
 
         ##

From cbbea3679ce2fa4c5ec6c057f4749b8d5291b2cc Mon Sep 17 00:00:00 2001
From: Chris Smith <quartzmo@gmail.com>
Date: Wed, 2 Jun 2021 15:56:21 -0600
Subject: [PATCH 22/22] Add Query#to_json and Query.from_json

---
 .../acceptance/firestore/query_test.rb        | 33 ++++++++++
 .../lib/google/cloud/firestore/query.rb       | 62 +++++++++++++++++++
 .../google/cloud/firestore/query_partition.rb |  2 +-
 .../google/cloud/firestore/query/get_test.rb  | 37 +++++++++++
 4 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/google-cloud-firestore/acceptance/firestore/query_test.rb b/google-cloud-firestore/acceptance/firestore/query_test.rb
index 1fdf6e4d6e03..db837b36a478 100644
--- a/google-cloud-firestore/acceptance/firestore/query_test.rb
+++ b/google-cloud-firestore/acceptance/firestore/query_test.rb
@@ -284,4 +284,37 @@
     _(results.map(&:document_id)).must_equal ["doc1", "doc2"]
     _(results.map { |doc| doc[:foo] }).must_equal ["a", "b"]
   end
+
+  it "has to_json method and from_json class method" do
+    rand_query_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
+    rand_query_col.doc("doc1").create({foo: "a"})
+    rand_query_col.doc("doc2").create({foo: "b"})
+    rand_query_col.doc("doc3").create({foo: "c"})
+
+    original_query = rand_query_col.order(:foo).limit_to_last 2
+
+    json = original_query.to_json
+    _(json).must_be_instance_of String
+
+    query = Google::Cloud::Firestore::Query.from_json json, firestore
+    _(query).must_be_instance_of Google::Cloud::Firestore::Query
+
+    results_1 = []
+    query.get { |result| results_1 << result } # block directly to get, rpc
+    _(results_1.map(&:document_id)).must_equal ["doc2","doc3"]
+    _(results_1.map { |doc| doc[:foo] }).must_equal ["b","c"]
+
+    results_2 = []
+    query.get { |result| results_2 << result } # block directly to get, rpc
+    _(results_2.map(&:document_id)).must_equal ["doc2","doc3"]
+    _(results_2.map { |doc| doc[:foo] }).must_equal ["b","c"]
+
+    results_3 = query.get # enum_for :get
+    _(results_3.map(&:document_id)).must_equal ["doc2","doc3"] # rpc
+    _(results_3.map { |doc| doc[:foo] }).must_equal ["b","c"] # rpc
+
+    results_4 = query.get # enum_for :get
+    _(results_4.map(&:document_id)).must_equal ["doc2","doc3"] # rpc
+    _(results_4.map { |doc| doc[:foo] }).must_equal ["b","c"] # rpc
+  end
 end
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query.rb b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
index 8358df4288da..13af07448109 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query.rb
@@ -17,6 +17,7 @@
 require "google/cloud/firestore/document_snapshot"
 require "google/cloud/firestore/query_listener"
 require "google/cloud/firestore/convert"
+require "json"
 
 module Google
   module Cloud
@@ -972,6 +973,67 @@ def listen &callback
         end
         alias on_snapshot listen
 
+        ##
+        # Serializes the instance to a JSON text string. See also {Query.from_json}.
+        #
+        # @return [String] A JSON text string.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #   query = firestore.col(:cities).select(:population)
+        #
+        #   json = query.to_json
+        #
+        #   new_query = Google::Cloud::Firestore::Query.from_json json, firestore
+        #
+        #   new_query.get do |city|
+        #     puts "#{city.document_id} has #{city[:population]} residents."
+        #   end
+        #
+        def to_json options = nil
+          query_json = Google::Cloud::Firestore::V1::StructuredQuery.encode_json query
+          {
+            "query" => JSON.parse(query_json),
+            "parent_path" => parent_path,
+            "limit_type" => limit_type
+          }.to_json options
+        end
+
+        ##
+        # Deserializes a JSON text string serialized from this class and returns it as a new instance. See also
+        # {#to_json}.
+        #
+        # @param [String] json A JSON text string serialized using {#to_json}.
+        # @param [Google::Cloud::Firestore::Client] client A connected client instance.
+        #
+        # @return [Query] A new query equal to the original query used to create the JSON text string.
+        #
+        # @example
+        #   require "google/cloud/firestore"
+        #
+        #   firestore = Google::Cloud::Firestore.new
+        #   query = firestore.col(:cities).select(:population)
+        #
+        #   json = query.to_json
+        #
+        #   new_query = Google::Cloud::Firestore::Query.from_json json, firestore
+        #
+        #   new_query.get do |city|
+        #     puts "#{city.document_id} has #{city[:population]} residents."
+        #   end
+        #
+        def self.from_json json, client
+          raise ArgumentError, "client is required" unless client
+
+          json = JSON.parse json
+          query_json = json["query"]
+          raise ArgumentError, "Field 'query' is required" unless query_json
+          query = Google::Cloud::Firestore::V1::StructuredQuery.decode_json query_json.to_json
+          start query, json["parent_path"], client, limit_type: json["limit_type"]&.to_sym
+        end
+
         ##
         # @private Start a new Query.
         def self.start query, parent_path, client, limit_type: nil
diff --git a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
index 0ae236f0120b..362e54ec8590 100644
--- a/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
+++ b/google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb
@@ -63,7 +63,7 @@ def initialize query, start_at, end_before
         end
 
         ##
-        # Creates a new query that only returns the documents for this partition, using the cursor valuess from
+        # Creates a new query that only returns the documents for this partition, using the cursor values from
         # {#start_at} and {#end_before}.
         #
         # @return [Query] The query for the partition.
diff --git a/google-cloud-firestore/test/google/cloud/firestore/query/get_test.rb b/google-cloud-firestore/test/google/cloud/firestore/query/get_test.rb
index 29d2a6d6fb0a..4def253dbf5e 100644
--- a/google-cloud-firestore/test/google/cloud/firestore/query/get_test.rb
+++ b/google-cloud-firestore/test/google/cloud/firestore/query/get_test.rb
@@ -253,6 +253,43 @@
     assert_results_enum results_enum
   end
 
+  it "gets a complex query after serialization and deserialization" do
+    expected_query = Google::Cloud::Firestore::V1::StructuredQuery.new(
+      select: Google::Cloud::Firestore::V1::StructuredQuery::Projection.new(
+        fields: [Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "name")]),
+      from: [Google::Cloud::Firestore::V1::StructuredQuery::CollectionSelector.new(collection_id: "users", all_descendants: false)],
+      offset: 3,
+      limit: Google::Protobuf::Int32Value.new(value: 42),
+      order_by: [
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "name"),
+          direction: :ASCENDING),
+        Google::Cloud::Firestore::V1::StructuredQuery::Order.new(
+          field: Google::Cloud::Firestore::V1::StructuredQuery::FieldReference.new(field_path: "__name__"),
+          direction: :DESCENDING)],
+      start_at: Google::Cloud::Firestore::V1::Cursor.new(values: [Google::Cloud::Firestore::Convert.raw_to_value("foo")], before: false),
+      end_at: Google::Cloud::Firestore::V1::Cursor.new(values: [Google::Cloud::Firestore::Convert.raw_to_value("bar")], before: true)
+    )
+    firestore_mock.expect :run_query, query_results_enum, run_query_args(expected_query)
+
+    original_query = firestore.col(:users).select(:name).offset(3).limit(42).order(:name).order(firestore.document_id, :desc).start_after(:foo).end_before(:bar)
+
+    json = original_query.to_json
+    _(json).must_be_instance_of String
+
+    deserialized_query = Google::Cloud::Firestore::Query.from_json json, firestore
+    _(deserialized_query).must_be_instance_of Google::Cloud::Firestore::Query
+
+    _(deserialized_query.query).must_equal expected_query # Private field
+    _(deserialized_query.parent_path).must_equal original_query.parent_path # Private field
+    _(deserialized_query.limit_type).must_equal original_query.limit_type # Private field
+    _(deserialized_query.client).must_equal original_query.client # Private field
+
+    results_enum = deserialized_query.get
+
+    assert_results_enum results_enum
+  end
+
   def assert_results_enum enum
     _(enum).must_be_kind_of Enumerator