Skip to content

Commit

Permalink
feat: support limit to last feature (#57)
Browse files Browse the repository at this point in the history
* feat: implement limit to last feature

* simplify test

* add more tests and de-deprecate Collection.get() method

* fix unit tests

* reverse should be used in get() only when limit_to_last set

* cosmetic changes

* use bool instead of int for _limit_to_last attr

* fix comments

* fix comments

* Apply suggestions from code review

Co-authored-by: BenWhitehead <BenWhitehead@users.noreply.github.com>

* add notes about mutually exclusivity

Co-authored-by: BenWhitehead <BenWhitehead@users.noreply.github.com>
Co-authored-by: Christopher Wilcox <crwilcox@google.com>
  • Loading branch information
3 people committed Jun 26, 2020
1 parent b9e2ab5 commit 8c75e21
Show file tree
Hide file tree
Showing 4 changed files with 229 additions and 46 deletions.
55 changes: 47 additions & 8 deletions google/cloud/firestore_v1/collection.py
Expand Up @@ -14,7 +14,6 @@

"""Classes for representing collections for the Google Cloud Firestore API."""
import random
import warnings

import six

Expand Down Expand Up @@ -257,6 +256,11 @@ def order_by(self, field_path, **kwargs):
def limit(self, count):
"""Create a limited query with this collection as parent.
.. note::
`limit` and `limit_to_last` are mutually exclusive.
Setting `limit` will drop previously set `limit_to_last`.
See
:meth:`~google.cloud.firestore_v1.query.Query.limit` for
more information on this method.
Expand All @@ -272,6 +276,29 @@ def limit(self, count):
query = query_mod.Query(self)
return query.limit(count)

def limit_to_last(self, count):
"""Create a limited to last query with this collection as parent.
.. note::
`limit` and `limit_to_last` are mutually exclusive.
Setting `limit_to_last` will drop previously set `limit`.
See
:meth:`~google.cloud.firestore_v1.query.Query.limit_to_last`
for more information on this method.
Args:
count (int): Maximum number of documents to return that
match the query.
Returns:
:class:`~google.cloud.firestore_v1.query.Query`:
A limited to last query.
"""
query = query_mod.Query(self)
return query.limit_to_last(count)

def offset(self, num_to_skip):
"""Skip to an offset in a query with this collection as parent.
Expand Down Expand Up @@ -375,13 +402,25 @@ def end_at(self, document_fields):
return query.end_at(document_fields)

def get(self, transaction=None):
"""Deprecated alias for :meth:`stream`."""
warnings.warn(
"'Collection.get' is deprecated: please use 'Collection.stream' instead.",
DeprecationWarning,
stacklevel=2,
)
return self.stream(transaction=transaction)
"""Read the documents in this collection.
This sends a ``RunQuery`` RPC and returns a list of documents
returned in the stream of ``RunQueryResponse`` messages.
Args:
transaction
(Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
An existing transaction that this query will run in.
If a ``transaction`` is used and it already has write operations
added, this method cannot be used (i.e. read-after-write is not
allowed).
Returns:
list: The documents in this collection that match the query.
"""
query = query_mod.Query(self)
return query.get(transaction=transaction)

def stream(self, transaction=None):
"""Read the documents in this collection.
Expand Down
101 changes: 91 additions & 10 deletions google/cloud/firestore_v1/query.py
Expand Up @@ -20,7 +20,6 @@
"""
import copy
import math
import warnings

from google.protobuf import wrappers_pb2
import six
Expand Down Expand Up @@ -86,6 +85,8 @@ class Query(object):
The "order by" entries to use in the query.
limit (Optional[int]):
The maximum number of documents the query is allowed to return.
limit_to_last (Optional[bool]):
Denotes whether a provided limit is applied to the end of the result set.
offset (Optional[int]):
The number of results to skip.
start_at (Optional[Tuple[dict, bool]]):
Expand Down Expand Up @@ -134,6 +135,7 @@ def __init__(
field_filters=(),
orders=(),
limit=None,
limit_to_last=False,
offset=None,
start_at=None,
end_at=None,
Expand All @@ -144,6 +146,7 @@ def __init__(
self._field_filters = field_filters
self._orders = orders
self._limit = limit
self._limit_to_last = limit_to_last
self._offset = offset
self._start_at = start_at
self._end_at = end_at
Expand All @@ -158,6 +161,7 @@ def __eq__(self, other):
and self._field_filters == other._field_filters
and self._orders == other._orders
and self._limit == other._limit
and self._limit_to_last == other._limit_to_last
and self._offset == other._offset
and self._start_at == other._start_at
and self._end_at == other._end_at
Expand Down Expand Up @@ -212,6 +216,7 @@ def select(self, field_paths):
field_filters=self._field_filters,
orders=self._orders,
limit=self._limit,
limit_to_last=self._limit_to_last,
offset=self._offset,
start_at=self._start_at,
end_at=self._end_at,
Expand Down Expand Up @@ -281,6 +286,7 @@ def where(self, field_path, op_string, value):
field_filters=new_filters,
orders=self._orders,
limit=self._limit,
limit_to_last=self._limit_to_last,
offset=self._offset,
start_at=self._start_at,
end_at=self._end_at,
Expand Down Expand Up @@ -333,16 +339,55 @@ def order_by(self, field_path, direction=ASCENDING):
field_filters=self._field_filters,
orders=new_orders,
limit=self._limit,
limit_to_last=self._limit_to_last,
offset=self._offset,
start_at=self._start_at,
end_at=self._end_at,
all_descendants=self._all_descendants,
)

def limit(self, count):
"""Limit a query to return a fixed number of results.
"""Limit a query to return at most `count` matching results.
If the current query already has a limit set, this will overwrite it.
If the current query already has a `limit` set, this will override it.
.. note::
`limit` and `limit_to_last` are mutually exclusive.
Setting `limit` will drop previously set `limit_to_last`.
Args:
count (int): Maximum number of documents to return that match
the query.
Returns:
:class:`~google.cloud.firestore_v1.query.Query`:
A limited query. Acts as a copy of the current query, modified
with the newly added "limit" filter.
"""
return self.__class__(
self._parent,
projection=self._projection,
field_filters=self._field_filters,
orders=self._orders,
limit=count,
limit_to_last=False,
offset=self._offset,
start_at=self._start_at,
end_at=self._end_at,
all_descendants=self._all_descendants,
)

def limit_to_last(self, count):
"""Limit a query to return the last `count` matching results.
If the current query already has a `limit_to_last`
set, this will override it.
.. note::
`limit` and `limit_to_last` are mutually exclusive.
Setting `limit_to_last` will drop previously set `limit`.
Args:
count (int): Maximum number of documents to return that match
Expand All @@ -359,6 +404,7 @@ def limit(self, count):
field_filters=self._field_filters,
orders=self._orders,
limit=count,
limit_to_last=True,
offset=self._offset,
start_at=self._start_at,
end_at=self._end_at,
Expand Down Expand Up @@ -386,6 +432,7 @@ def offset(self, num_to_skip):
field_filters=self._field_filters,
orders=self._orders,
limit=self._limit,
limit_to_last=self._limit_to_last,
offset=num_to_skip,
start_at=self._start_at,
end_at=self._end_at,
Expand Down Expand Up @@ -729,13 +776,42 @@ def _to_protobuf(self):
return query_pb2.StructuredQuery(**query_kwargs)

def get(self, transaction=None):
"""Deprecated alias for :meth:`stream`."""
warnings.warn(
"'Query.get' is deprecated: please use 'Query.stream' instead.",
DeprecationWarning,
stacklevel=2,
)
return self.stream(transaction=transaction)
"""Read the documents in the collection that match this query.
This sends a ``RunQuery`` RPC and returns a list of documents
returned in the stream of ``RunQueryResponse`` messages.
Args:
transaction
(Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
An existing transaction that this query will run in.
If a ``transaction`` is used and it already has write operations
added, this method cannot be used (i.e. read-after-write is not
allowed).
Returns:
list: The documents in the collection that match this query.
"""
is_limited_to_last = self._limit_to_last

if self._limit_to_last:
# In order to fetch up to `self._limit` results from the end of the
# query flip the defined ordering on the query to start from the
# end, retrieving up to `self._limit` results from the backend.
for order in self._orders:
order.direction = _enum_from_direction(
self.DESCENDING
if order.direction == self.ASCENDING
else self.ASCENDING
)
self._limit_to_last = False

result = self.stream(transaction=transaction)
if is_limited_to_last:
result = reversed(list(result))

return list(result)

def stream(self, transaction=None):
"""Read the documents in the collection that match this query.
Expand Down Expand Up @@ -764,6 +840,11 @@ def stream(self, transaction=None):
:class:`~google.cloud.firestore_v1.document.DocumentSnapshot`:
The next document that fulfills the query.
"""
if self._limit_to_last:
raise ValueError(
"Query results for queries that include limit_to_last() "
"constraints cannot be streamed. Use Query.get() instead."
)
parent_path, expected_prefix = self._parent._parent_info()
response_iterator = self._client._firestore_api.run_query(
parent_path,
Expand Down
36 changes: 18 additions & 18 deletions tests/unit/v1/test_collection.py
Expand Up @@ -371,6 +371,18 @@ def test_limit(self):
self.assertIs(query._parent, collection)
self.assertEqual(query._limit, limit)

def test_limit_to_last(self):
from google.cloud.firestore_v1.query import Query

LIMIT = 15
collection = self._make_one("collection")
query = collection.limit_to_last(LIMIT)

self.assertIsInstance(query, Query)
self.assertIs(query._parent, collection)
self.assertEqual(query._limit, LIMIT)
self.assertTrue(query._limit_to_last)

def test_offset(self):
from google.cloud.firestore_v1.query import Query

Expand Down Expand Up @@ -484,38 +496,26 @@ def test_list_documents_w_page_size(self):

@mock.patch("google.cloud.firestore_v1.query.Query", autospec=True)
def test_get(self, query_class):
import warnings

collection = self._make_one("collection")
with warnings.catch_warnings(record=True) as warned:
get_response = collection.get()
get_response = collection.get()

query_class.assert_called_once_with(collection)
query_instance = query_class.return_value
self.assertIs(get_response, query_instance.stream.return_value)
query_instance.stream.assert_called_once_with(transaction=None)

# Verify the deprecation
self.assertEqual(len(warned), 1)
self.assertIs(warned[0].category, DeprecationWarning)
self.assertIs(get_response, query_instance.get.return_value)
query_instance.get.assert_called_once_with(transaction=None)

@mock.patch("google.cloud.firestore_v1.query.Query", autospec=True)
def test_get_with_transaction(self, query_class):
import warnings

collection = self._make_one("collection")
transaction = mock.sentinel.txn
with warnings.catch_warnings(record=True) as warned:
get_response = collection.get(transaction=transaction)
get_response = collection.get(transaction=transaction)

query_class.assert_called_once_with(collection)
query_instance = query_class.return_value
self.assertIs(get_response, query_instance.stream.return_value)
query_instance.stream.assert_called_once_with(transaction=transaction)

# Verify the deprecation
self.assertEqual(len(warned), 1)
self.assertIs(warned[0].category, DeprecationWarning)
self.assertIs(get_response, query_instance.get.return_value)
query_instance.get.assert_called_once_with(transaction=transaction)

@mock.patch("google.cloud.firestore_v1.query.Query", autospec=True)
def test_stream(self, query_class):
Expand Down

0 comments on commit 8c75e21

Please sign in to comment.