Skip to content

Commit 4e61da5

Browse files
quartzmoquartzmo
committed
Update QueryPartition::List to add final QueryPartition only at end of all results
1 parent a8388a2 commit 4e61da5

File tree

8 files changed

+142
-100
lines changed

8 files changed

+142
-100
lines changed

‎google-cloud-firestore/acceptance/firestore/collection_group_test.rb‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -114,33 +114,37 @@
114114

115115
describe "#partitions" do
116116
it "queries a collection group using partitions" do
117+
rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
118+
117119
document_ids = ["a", "b", "c"].map do |prefix|
118120
# Minimum partition size is 128.
119121
128.times.map do |i|
120122
"#{prefix}#{(i+1).to_s.rjust(3, '0')}"
121123
end
122124
end.flatten # "a001", "a002", ... "c128"
123-
124-
rand_col = firestore.col "#{root_path}/query/#{SecureRandom.hex(4)}"
125125
firestore.batch do |b|
126126
document_ids.each do |id|
127127
doc_ref = rand_col.document id
128-
b.set doc_ref, {}
128+
b.set doc_ref, { foo: id }
129129
end
130130
end
131131

132132
collection_group = firestore.collection_group(rand_col.collection_id)
133133

134-
partitions = collection_group.partitions 3
134+
partitions = collection_group.partitions 6, max: 1
135135
_(partitions).must_be_kind_of Google::Cloud::Firestore::QueryPartition::List
136-
_(partitions.count).must_equal 3
136+
_(partitions.count).must_equal 1
137137

138138
_(partitions[0]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
139139
_(partitions[0].start_at).must_be :nil?
140140
_(partitions[0].end_before).must_be_kind_of Array
141141
_(partitions[0].end_before[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
142142
_(document_ids).must_include partitions[0].end_before[0].document_id
143143

144+
partitions += partitions.next
145+
_(partitions).must_be_kind_of Array
146+
_(partitions.count).must_equal 3
147+
144148
_(partitions[1]).must_be_kind_of Google::Cloud::Firestore::QueryPartition
145149
_(partitions[1].start_at).must_be_kind_of Array
146150
_(partitions[1].start_at[0]).must_be_kind_of Google::Cloud::Firestore::DocumentReference
@@ -159,7 +163,10 @@
159163
_(queries.count).must_equal 3
160164
results = queries.map do |query|
161165
_(query).must_be_kind_of Google::Cloud::Firestore::Query
162-
query.get.map(&:document_id)
166+
query.get.map do |snp|
167+
_(snp).must_be_kind_of Google::Cloud::Firestore::DocumentSnapshot
168+
snp.document_id
169+
end
163170
end
164171
results.each { |result| _(result).wont_be :empty? }
165172
# Verify all document IDs have been returned, in original order.

‎google-cloud-firestore/lib/google/cloud/firestore/collection_group.rb‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@ module Firestore
2424
# # CollectionGroup
2525
#
2626
# A collection group object is used for adding documents, getting
27-
# document references, and querying for documents (See {Query}).
27+
# document references, and querying for documents, including with partitions.
28+
#
29+
# See {Client#col_group} and {Query}.
2830
#
2931
# @example
3032
# require "google/cloud/firestore"

‎google-cloud-firestore/lib/google/cloud/firestore/document_reference/list.rb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ def next
109109
#
110110
# @return [Enumerator]
111111
#
112-
# @example Iterating each document reference by passing a block:
112+
# @example Iterating each document reference by passing a block or proc:
113113
# require "google/cloud/firestore"
114114
#
115115
# firestore = Google::Cloud::Firestore.new

‎google-cloud-firestore/lib/google/cloud/firestore/query.rb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1105,7 +1105,7 @@ def values_to_cursor values, query
11051105
order_field_paths = order_by_field_paths query
11061106
if values.count > order_field_paths.count
11071107
# raise if too many values provided for the cursor
1108-
raise ArgumentError, "There cannot be more values than Order By fields"
1108+
raise ArgumentError, "There cannot be more cursor values than order by fields"
11091109
end
11101110

11111111
values = values.zip(order_field_paths).map do |value, field_path|

‎google-cloud-firestore/lib/google/cloud/firestore/query_partition.rb‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,20 @@ module Firestore
2626
# The cursors returned by {#start_at} and {#end_before} can only be used in a query that matches the constraint of
2727
# the query that produced this partition.
2828
#
29-
# See {CollectionGroup#partitions}.
29+
# See {CollectionGroup#partitions} and {Query}.
3030
#
3131
# @!attribute [r] start_at
32-
# The cursor that defines the first result for this partition, or `nil` if this is the first partition.
33-
# @return [Array] a cursor value that can be used with {Query#start_at} or `nil` if this is the first partition.
32+
# The cursor values that define the first result for this partition, or `nil` if this is the first partition.
33+
# Returns an array of values that represent a position, in the order they appear in the order by clause of the
34+
# query. Can contain fewer values than specified in the order by clause. Will be used in the query returned by
35+
# {#create_query}.
36+
# @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
3437
# @!attribute [r] end_before
35-
# The cursor that defines the first result after this partition, or `nil` if this is the last partition.
36-
# @return [Array] a cursor value that can be used with {Query#end_before} or `nil` if this is the last
37-
# partition.
38+
# The cursor values that define the first result after this partition, or `nil` if this is the last partition.
39+
# Returns an array of values that represent a position, in the order they appear in the order by clause of the
40+
# query. Can contain fewer values than specified in the order by clause. Will be used in the query returned by
41+
# {#create_query}.
42+
# @return [Array<Object>, nil] Typically, the values are {DocumentSnapshot} objects.
3843
#
3944
# @example
4045
# require "google/cloud/firestore"
@@ -60,9 +65,10 @@ def initialize query, start_at, end_before
6065
end
6166

6267
##
63-
# Creates a query that only returns the documents for this partition.
68+
# Creates a new query that only returns the documents for this partition, using the cursor valuess from
69+
# {#start_at} and {#end_before}.
6470
#
65-
# @return [Query] query.
71+
# @return [Query] The query for the partition.
6672
#
6773
def create_query
6874
base_query = @query

‎google-cloud-firestore/lib/google/cloud/firestore/query_partition/list.rb‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def initialize arr = []
4747
end
4848

4949
##
50-
# Whether there is a next page of document references.
50+
# Whether there is a next page of query partitions.
5151
#
5252
# @return [Boolean]
5353
#
@@ -60,15 +60,15 @@ def initialize arr = []
6060
#
6161
# partitions = col_group.partitions 3
6262
# if partitions.next?
63-
# next_documents = partitions.next
63+
# next_partitions = partitions.next
6464
# end
6565
#
6666
def next?
6767
!token.nil?
6868
end
6969

7070
##
71-
# Retrieve the next page of document references.
71+
# Retrieve the next page of query partitions.
7272
#
7373
# @return [QueryPartition::List]
7474
#
@@ -81,14 +81,14 @@ def next?
8181
#
8282
# partitions = col_group.partitions 3
8383
# if partitions.next?
84-
# next_documents = partitions.next
84+
# next_partitions = partitions.next
8585
# end
8686
#
8787
def next
8888
return nil unless next?
8989
ensure_client!
9090
grpc = @client.service.partition_query @parent, @query.query, @partition_count, token: token, max: @max
91-
self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max
91+
self.class.from_grpc grpc, @client, @parent, @query, @partition_count, max: @max, start_at: @start_at
9292
end
9393

9494
##
@@ -103,14 +103,13 @@ def next
103103
# over the results returned by a single API call.) Use with caution.
104104
#
105105
# @param [Integer] request_limit The upper limit of API requests to
106-
# make to load all document references. Default is no limit.
107-
# @yield [document] The block for accessing each document.
108-
# @yieldparam [QueryPartition] document The document reference
109-
# object.
106+
# make to load all query partitions. Default is no limit.
107+
# @yield [partition] The block for accessing each partition.
108+
# @yieldparam [QueryPartition] partition The query partition object.
110109
#
111110
# @return [Enumerator]
112111
#
113-
# @example Iterating each document reference by passing a block:
112+
# @example Iterating each query partition by passing a block or proc:
114113
# require "google/cloud/firestore"
115114
#
116115
# firestore = Google::Cloud::Firestore.new
@@ -119,7 +118,7 @@ def next
119118
#
120119
# partitions = col_group.partitions 3
121120
#
122-
# queries = partitions.map(&:create_query)
121+
# queries = partitions.all(&:create_query)
123122
#
124123
# @example Using the enumerator by not passing a block:
125124
# require "google/cloud/firestore"
@@ -130,9 +129,9 @@ def next
130129
#
131130
# partitions = col_group.partitions 3
132131
#
133-
# queries = partitions.map(&:create_query)
132+
# queries = partitions.all.map(&:create_query)
134133
#
135-
# @example Limit the number of API calls made:
134+
# @example Limit the number of API calls made by `#all`:
136135
# require "google/cloud/firestore"
137136
#
138137
# firestore = Google::Cloud::Firestore.new
@@ -141,7 +140,7 @@ def next
141140
#
142141
# partitions = col_group.partitions 3
143142
#
144-
# queries = partitions.map(&:create_query)
143+
# queries = partitions.all(request_limit: 10, &:create_query)
145144
#
146145
def all request_limit: nil, &block
147146
request_limit = request_limit.to_i if request_limit
@@ -163,9 +162,10 @@ def all request_limit: nil, &block
163162
##
164163
# @private New QueryPartition::List from a
165164
# Google::Cloud::Firestore::V1::PartitionQueryResponse object.
166-
def self.from_grpc grpc, client, parent, query, partition_count, max: nil
167-
# TODO: Should the logic adding the last partition be applied to the entire result set, not the page?
168-
start_at = nil
165+
def self.from_grpc grpc, client, parent, query, partition_count, max: nil, start_at: nil
166+
token = grpc.next_page_token
167+
token = nil if token == ""
168+
169169
partitions = List.new(Array(grpc.partitions).map do |cursor|
170170
end_before = cursor.values.map do |value|
171171
Convert.value_to_raw value, client
@@ -174,16 +174,15 @@ def self.from_grpc grpc, client, parent, query, partition_count, max: nil
174174
start_at = end_before
175175
partition
176176
end)
177-
# We are always returning an extra partition (with en empty endBefore cursor).
178-
partitions << QueryPartition.new(query, start_at, nil)
177+
# Return an extra partition (with en empty endBefore cursor) at the end of the results.
178+
partitions << QueryPartition.new(query, start_at, nil) unless token
179179
partitions.instance_variable_set :@parent, parent
180180
partitions.instance_variable_set :@query, query
181181
partitions.instance_variable_set :@partition_count, partition_count
182-
token = grpc.next_page_token
183-
token = nil if token == ""
184182
partitions.instance_variable_set :@token, token
185183
partitions.instance_variable_set :@client, client
186184
partitions.instance_variable_set :@max, max
185+
partitions.instance_variable_set :@start_at, start_at
187186
partitions
188187
end
189188

0 commit comments

Comments
 (0)