Firestore: make cursor-related 'Query' methods accept lists

firestore/google/cloud/firestore_v1beta1/collection.py

@@ -279,10 +279,10 @@ def start_at(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor.
@@ -299,10 +299,10 @@ def start_after(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor.
@@ -319,10 +319,10 @@ def end_before(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor.
@@ -339,10 +339,10 @@ def end_at(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor.

firestore/google/cloud/firestore_v1beta1/query.py

@@ -56,6 +56,8 @@
     'When defining a cursor with one of ``start_at()`` / ``start_after()`` / '
     '``end_before()`` / ``end_at()``, all fields in the cursor must '
     'come from fields set in ``order_by()``.')
+_MISMATCH_CURSOR_W_ORDER_BY = (
+    'The cursor {!r} does not match the order fields {!r}.')
 _EMPTY_DOC_TEMPLATE = (
     'Unexpected server response. All responses other than the first must '
     'contain a document. The response at index {} was\n{}.')
@@ -349,10 +351,10 @@ def _cursor_helper(self, document_fields, before, start):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
             before (bool): Flag indicating if the document in
                 ``document_fields`` should (:data:`False`) or
                 shouldn't (:data:`True`) be included in the result set.
@@ -364,12 +366,13 @@ def _cursor_helper(self, document_fields, before, start):
             a copy of the current query, modified with the newly added
             "start at" cursor.
         """
-        if isinstance(document_fields, dict):
+        if isinstance(document_fields, tuple):
+            document_fields = list(document_fields)
+        elif isinstance(document_fields, document.DocumentSnapshot):
+            document_fields = document_fields.to_dict()
+        else:
             # NOTE: We copy so that the caller can't modify after calling.
             document_fields = copy.deepcopy(document_fields)
-        else:
-            # NOTE: This **assumes** a DocumentSnapshot.
-            document_fields = document_fields.to_dict()
 
         cursor_pair = document_fields, before
         query_kwargs = {
@@ -405,10 +408,10 @@ def start_at(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor. Acts as
@@ -434,10 +437,10 @@ def start_after(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor. Acts as
@@ -463,10 +466,10 @@ def end_before(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor. Acts as
@@ -492,10 +495,10 @@ def end_at(self, document_fields):
 
         Args:
             document_fields (Union[~.firestore_v1beta1.\
-                document.DocumentSnapshot, dict]): Either a document snapshot
-                or a dictionary of fields representing a query results
-                cursor. A cursor is a collection of values that represent a
-                position in a query result set.
+                document.DocumentSnapshot, dict, list, tuple]): a document
+                snapshot or a dictionary/list/tuple of fields representing a
+                query results cursor. A cursor is a collection of values that
+                represent a position in a query result set.
 
         Returns:
             ~.firestore_v1beta1.query.Query: A query with cursor. Acts as
@@ -530,13 +533,48 @@ def _filters_pb(self):
             return query_pb2.StructuredQuery.Filter(
                 composite_filter=composite_filter)
 
+    @staticmethod
+    def _normalize_cursor(cursor, orders):
+        """Helper: convert cursor to a list of values based on orders."""
+        if cursor is None:
+            return
+
+        if not orders:
+            raise ValueError(_NO_ORDERS_FOR_CURSOR)
+
+        document_fields, before = cursor
+
+        order_keys = [order.field.field_path for order in orders]
+
+        if isinstance(document_fields, dict):
+            # Transform to list using orders
+            values = []
+            data = document_fields
+            for order_key in order_keys:
+                try:
+                    values.append(_helpers.get_nested_value(order_key, data))
+                except KeyError:
+                    msg = _MISSING_ORDER_BY.format(order_key, data)
+                    raise ValueError(msg)
+            document_fields = values
+
+        if len(document_fields) != len(orders):
+            msg = _MISMATCH_CURSOR_W_ORDER_BY.format(
+                document_fields, order_keys)
+            raise ValueError(msg)
+
+        return document_fields, before
+
     def _to_protobuf(self):
         """Convert the current query into the equivalent protobuf.
 
         Returns:
             google.cloud.firestore_v1beta1.types.StructuredQuery: The
             query protobuf.
         """
+        start_at = self._normalize_cursor(self._start_at, self._orders)
+        end_at = self._normalize_cursor(self._end_at, self._orders)
+
         query_kwargs = {
             'select': self._projection,
             'from': [
@@ -546,8 +584,8 @@ def _to_protobuf(self):
             ],
             'where': self._filters_pb(),
             'order_by': self._orders,
-            'start_at': _cursor_pb(self._start_at, self._orders),
-            'end_at': _cursor_pb(self._end_at, self._orders),
+            'start_at': _cursor_pb(start_at),
+            'end_at': _cursor_pb(end_at),
         }
         if self._offset is not None:
             query_kwargs['offset'] = self._offset
@@ -775,54 +813,25 @@ def _filter_pb(field_or_unary):
             'Unexpected filter type', type(field_or_unary), field_or_unary)
 
 
-def _cursor_pb(cursor_pair, orders):
+def _cursor_pb(cursor_pair):
     """Convert a cursor pair to a protobuf.
 
     If ``cursor_pair`` is :data:`None`, just returns :data:`None`.
 
     Args:
-        cursor_pair (Optional[Tuple[dict, bool]]): Two-tuple of
+        cursor_pair (Optional[Tuple[list, bool]]): Two-tuple of
 
-            * a mapping of fields. Any field that is present in this mapping
-              must also be present in ``orders``
+            * a list of field values.
             * a ``before`` flag
 
-        orders (Tuple[google.cloud.proto.firestore.v1beta1.\
-            query_pb2.StructuredQuery.Order, ...]]): The "order by" entries
-            to use for a query. (We use this rather than a list of field path
-            strings just because it is how a query stores calls
-            to ``order_by``.)
-
     Returns:
         Optional[google.cloud.firestore_v1beta1.types.Cursor]: A
         protobuf cursor corresponding to the values.
-
-    Raises:
-        ValueError: If ``cursor_pair`` is not :data:`None`, but there are
-            no ``orders``.
-        ValueError: If one of the field paths in ``orders`` is not contained
-            in the ``data`` (i.e. the first component of ``cursor_pair``).
     """
-    if cursor_pair is None:
-        return None
-
-    if len(orders) == 0:
-        raise ValueError(_NO_ORDERS_FOR_CURSOR)
-
-    data, before = cursor_pair
-    value_pbs = []
-    for order in orders:
-        field_path = order.field.field_path
-        try:
-            value = _helpers.get_nested_value(field_path, data)
-        except KeyError:
-            msg = _MISSING_ORDER_BY.format(field_path, data)
-            raise ValueError(msg)
-
-        value_pb = _helpers.encode_value(value)
-        value_pbs.append(value_pb)
-
-    return query_pb2.Cursor(values=value_pbs, before=before)
+    if cursor_pair is not None:
+        data, before = cursor_pair
+        value_pbs = [_helpers.encode_value(value) for value in data]
+        return query_pb2.Cursor(values=value_pbs, before=before)
 
 
 def _query_response_to_snapshot(response_pb, collection, expected_prefix):

firestore/tests/system.py

@@ -565,9 +565,8 @@ def test_query_get(client, cleanup):
     assert expected_ab_pairs == ab_pairs2
 
     # 3. Use a start and end cursor.
-    query3 = collection.start_at({'a': num_vals - 2})
-    query3 = query3.order_by('a')
-    query3 = query3.end_before({'a': num_vals - 1})
+    query3 = collection.order_by(
+        'a').start_at({'a': num_vals - 2}).end_before({'a': num_vals - 1})
     values3 = [
         (snapshot.id, snapshot.to_dict())
         for snapshot in query3.get()