Firestore: impose required semantics for snapshots as cursors.

firestore/google/cloud/firestore_v1beta1/query.py

@@ -261,6 +261,14 @@ def where(self, field_path, op_string, value):
             end_at=self._end_at,
         )
 
+    @staticmethod
+    def _make_order(field_path, direction):
+        """Helper for :meth:`order_by`."""
+        return query_pb2.StructuredQuery.Order(
+            field=query_pb2.StructuredQuery.FieldReference(field_path=field_path),
+            direction=_enum_from_direction(direction),
+        )
+
     def order_by(self, field_path, direction=ASCENDING):
         """Modify the query to add an order clause on a specific field.
 
@@ -290,10 +298,7 @@ def order_by(self, field_path, direction=ASCENDING):
         """
         _helpers.split_field_path(field_path)  # raises
 
-        order_pb = query_pb2.StructuredQuery.Order(
-            field=query_pb2.StructuredQuery.FieldReference(field_path=field_path),
-            direction=_enum_from_direction(direction),
-        )
+        order_pb = self._make_order(field_path, direction)
 
         new_orders = self._orders + (order_pb,)
         return self.__class__(
@@ -387,7 +392,10 @@ def _cursor_helper(self, document_fields, before, start):
         if isinstance(document_fields, tuple):
             document_fields = list(document_fields)
         elif isinstance(document_fields, document.DocumentSnapshot):
-            document_fields = document_fields.to_dict()
+            if document_fields.reference._path[:-1] != self._parent._path:
+                raise ValueError(
+                    "Cannot use snapshot from another collection as a cursor."
+                )
         else:
             # NOTE: We copy so that the caller can't modify after calling.
             document_fields = copy.deepcopy(document_fields)
@@ -563,6 +571,40 @@ def _normalize_projection(projection):
 
         return projection
 
+    def _normalize_orders(self):
+        """Helper:  adjust orders based on cursors, where clauses."""
+        orders = list(self._orders)
+        _has_snapshot_cursor = False
+
+        if self._start_at:
+            if isinstance(self._start_at[0], document.DocumentSnapshot):
+                _has_snapshot_cursor = True
+
+        if self._end_at:
+            if isinstance(self._end_at[0], document.DocumentSnapshot):
+                _has_snapshot_cursor = True
+
+        if _has_snapshot_cursor:
+            should_order = [
+                _enum_from_op_string(key)
+                for key in _COMPARISON_OPERATORS
+                if key not in (_EQ_OP, "array_contains")
+            ]
+            order_keys = [order.field.field_path for order in orders]
+            for filter_ in self._field_filters:
+                field = filter_.field.field_path
+                if filter_.op in should_order and field not in order_keys:
+                    orders.append(self._make_order(field, "ASCENDING"))
+            if not orders:
+                orders.append(self._make_order("__name__", "ASCENDING"))
+            else:
+                order_keys = [order.field.field_path for order in orders]
+                if "__name__" not in order_keys:
+                    direction = orders[-1].direction  # enum?
+                    orders.append(self._make_order("__name__", direction))
+
+        return orders
+
     def _normalize_cursor(self, cursor, orders):
         """Helper: convert cursor to a list of values based on orders."""
         if cursor is None:
@@ -575,6 +617,11 @@ def _normalize_cursor(self, cursor, orders):
 
         order_keys = [order.field.field_path for order in orders]
 
+        if isinstance(document_fields, document.DocumentSnapshot):
+            snapshot = document_fields
+            document_fields = snapshot.to_dict()
+            document_fields["__name__"] = snapshot.reference
+
         if isinstance(document_fields, dict):
             # Transform to list using orders
             values = []
@@ -615,8 +662,9 @@ def _to_protobuf(self):
             query protobuf.
         """
         projection = self._normalize_projection(self._projection)
-        start_at = self._normalize_cursor(self._start_at, self._orders)
-        end_at = self._normalize_cursor(self._end_at, self._orders)
+        orders = self._normalize_orders()
+        start_at = self._normalize_cursor(self._start_at, orders)
+        end_at = self._normalize_cursor(self._end_at, orders)
 
         query_kwargs = {
             "select": projection,
@@ -626,7 +674,7 @@ def _to_protobuf(self):
                 )
             ],
             "where": self._filters_pb(),
-            "order_by": self._orders,
+            "order_by": orders,
             "start_at": _cursor_pb(start_at),
             "end_at": _cursor_pb(end_at),
         }
@@ -824,6 +872,9 @@ def _enum_from_direction(direction):
     Raises:
         ValueError: If ``direction`` is not a valid direction.
     """
+    if isinstance(direction, int):
+        return direction
+
     if direction == Query.ASCENDING:
         return enums.StructuredQuery.Direction.ASCENDING
     elif direction == Query.DESCENDING:

firestore/tests/unit/test_query.py

@@ -46,7 +46,7 @@ def test_constructor_defaults(self):
         self.assertIsNone(query._start_at)
         self.assertIsNone(query._end_at)
 
-    def _make_one_all_fields(self, limit=9876, offset=12, skip_fields=()):
+    def _make_one_all_fields(self, limit=9876, offset=12, skip_fields=(), parent=None):
         kwargs = {
             "projection": mock.sentinel.projection,
             "field_filters": mock.sentinel.filters,
@@ -58,7 +58,9 @@ def _make_one_all_fields(self, limit=9876, offset=12, skip_fields=()):
         }
         for field in skip_fields:
             kwargs.pop(field)
-        return self._make_one(mock.sentinel.parent, **kwargs)
+        if parent is None:
+            parent = mock.sentinel.parent
+        return self._make_one(parent, **kwargs)
 
     def test_constructor_explicit(self):
         limit = 234
@@ -289,10 +291,22 @@ def test_offset(self):
         self._compare_queries(query2, query3, "_offset")
 
     @staticmethod
-    def _make_snapshot(values):
-        from google.cloud.firestore_v1beta1.document import DocumentSnapshot
+    def _make_collection(*path, **kw):
+        from google.cloud.firestore_v1beta1 import collection
+
+        return collection.CollectionReference(*path, **kw)
 
-        return DocumentSnapshot(None, values, True, None, None, None)
+    @staticmethod
+    def _make_docref(*path, **kw):
+        from google.cloud.firestore_v1beta1 import document
+
+        return document.DocumentReference(*path, **kw)
+
+    @staticmethod
+    def _make_snapshot(docref, values):
+        from google.cloud.firestore_v1beta1 import document
+
+        return document.DocumentSnapshot(docref, values, True, None, None, None)
 
     def test__cursor_helper_w_dict(self):
         values = {"a": 7, "b": "foo"}
@@ -349,15 +363,26 @@ def test__cursor_helper_w_list(self):
         self.assertIsNot(cursor, values)
         self.assertTrue(before)
 
-    def test__cursor_helper_w_snapshot(self):
+    def test__cursor_helper_w_snapshot_wrong_collection(self):
+        values = {"a": 7, "b": "foo"}
+        docref = self._make_docref("there", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query = self._make_one(collection)
 
+        with self.assertRaises(ValueError):
+            query._cursor_helper(snapshot, False, False)
+
+    def test__cursor_helper_w_snapshot(self):
         values = {"a": 7, "b": "foo"}
-        snapshot = self._make_snapshot(values)
-        query1 = self._make_one(mock.sentinel.parent)
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query1 = self._make_one(collection)
 
         query2 = query1._cursor_helper(snapshot, False, False)
 
-        self.assertIs(query2._parent, mock.sentinel.parent)
+        self.assertIs(query2._parent, collection)
         self.assertIsNone(query2._projection)
         self.assertEqual(query2._field_filters, ())
         self.assertEqual(query2._orders, ())
@@ -367,11 +392,12 @@ def test__cursor_helper_w_snapshot(self):
 
         cursor, before = query2._end_at
 
-        self.assertEqual(cursor, values)
+        self.assertIs(cursor, snapshot)
         self.assertFalse(before)
 
     def test_start_at(self):
-        query1 = self._make_one_all_fields(skip_fields=("orders",))
+        collection = self._make_collection("here")
+        query1 = self._make_one_all_fields(parent=collection, skip_fields=("orders",))
         query2 = query1.order_by("hi")
 
         document_fields3 = {"hi": "mom"}
@@ -384,15 +410,17 @@ def test_start_at(self):
         # Make sure it overrides.
         query4 = query3.order_by("bye")
         values5 = {"hi": "zap", "bye": 88}
-        document_fields5 = self._make_snapshot(values5)
+        docref = self._make_docref("here", "doc_id")
+        document_fields5 = self._make_snapshot(docref, values5)
         query5 = query4.start_at(document_fields5)
         self.assertIsNot(query5, query4)
         self.assertIsInstance(query5, self._get_target_class())
-        self.assertEqual(query5._start_at, (values5, True))
+        self.assertEqual(query5._start_at, (document_fields5, True))
         self._compare_queries(query4, query5, "_start_at")
 
     def test_start_after(self):
-        query1 = self._make_one_all_fields(skip_fields=("orders",))
+        collection = self._make_collection("here")
+        query1 = self._make_one_all_fields(parent=collection, skip_fields=("orders",))
         query2 = query1.order_by("down")
 
         document_fields3 = {"down": 99.75}
@@ -405,15 +433,17 @@ def test_start_after(self):
         # Make sure it overrides.
         query4 = query3.order_by("out")
         values5 = {"down": 100.25, "out": b"\x00\x01"}
-        document_fields5 = self._make_snapshot(values5)
+        docref = self._make_docref("here", "doc_id")
+        document_fields5 = self._make_snapshot(docref, values5)
         query5 = query4.start_after(document_fields5)
         self.assertIsNot(query5, query4)
         self.assertIsInstance(query5, self._get_target_class())
-        self.assertEqual(query5._start_at, (values5, False))
+        self.assertEqual(query5._start_at, (document_fields5, False))
         self._compare_queries(query4, query5, "_start_at")
 
     def test_end_before(self):
-        query1 = self._make_one_all_fields(skip_fields=("orders",))
+        collection = self._make_collection("here")
+        query1 = self._make_one_all_fields(parent=collection, skip_fields=("orders",))
         query2 = query1.order_by("down")
 
         document_fields3 = {"down": 99.75}
@@ -426,15 +456,18 @@ def test_end_before(self):
         # Make sure it overrides.
         query4 = query3.order_by("out")
         values5 = {"down": 100.25, "out": b"\x00\x01"}
-        document_fields5 = self._make_snapshot(values5)
+        docref = self._make_docref("here", "doc_id")
+        document_fields5 = self._make_snapshot(docref, values5)
         query5 = query4.end_before(document_fields5)
         self.assertIsNot(query5, query4)
         self.assertIsInstance(query5, self._get_target_class())
-        self.assertEqual(query5._end_at, (values5, True))
+        self.assertEqual(query5._end_at, (document_fields5, True))
+        self._compare_queries(query4, query5, "_end_at")
         self._compare_queries(query4, query5, "_end_at")
 
     def test_end_at(self):
-        query1 = self._make_one_all_fields(skip_fields=("orders",))
+        collection = self._make_collection("here")
+        query1 = self._make_one_all_fields(parent=collection, skip_fields=("orders",))
         query2 = query1.order_by("hi")
 
         document_fields3 = {"hi": "mom"}
@@ -447,11 +480,12 @@ def test_end_at(self):
         # Make sure it overrides.
         query4 = query3.order_by("bye")
         values5 = {"hi": "zap", "bye": 88}
-        document_fields5 = self._make_snapshot(values5)
+        docref = self._make_docref("here", "doc_id")
+        document_fields5 = self._make_snapshot(docref, values5)
         query5 = query4.end_at(document_fields5)
         self.assertIsNot(query5, query4)
         self.assertIsInstance(query5, self._get_target_class())
-        self.assertEqual(query5._end_at, (values5, False))
+        self.assertEqual(query5._end_at, (document_fields5, False))
         self._compare_queries(query4, query5, "_end_at")
 
     def test__filters_pb_empty(self):
@@ -530,6 +564,67 @@ def test__normalize_projection_non_empty(self):
         query = self._make_one(mock.sentinel.parent)
         self.assertIs(query._normalize_projection(projection), projection)
 
+    def test__normalize_orders_wo_orders_wo_cursors(self):
+        query = self._make_one(mock.sentinel.parent)
+        expected = []
+        self.assertEqual(query._normalize_orders(), expected)
+
+    def test__normalize_orders_w_orders_wo_cursors(self):
+        query = self._make_one(mock.sentinel.parent).order_by("a")
+        expected = [query._make_order("a", "ASCENDING")]
+        self.assertEqual(query._normalize_orders(), expected)
+
+    def test__normalize_orders_wo_orders_w_snapshot_cursor(self):
+        values = {"a": 7, "b": "foo"}
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query = self._make_one(collection).start_at(snapshot)
+        expected = [query._make_order("__name__", "ASCENDING")]
+        self.assertEqual(query._normalize_orders(), expected)
+
+    def test__normalize_orders_w_name_orders_w_snapshot_cursor(self):
+        values = {"a": 7, "b": "foo"}
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query = (
+            self._make_one(collection)
+            .order_by("__name__", "DESCENDING")
+            .start_at(snapshot)
+        )
+        expected = [query._make_order("__name__", "DESCENDING")]
+        self.assertEqual(query._normalize_orders(), expected)
+
+    def test__normalize_orders_wo_orders_w_snapshot_cursor_w_neq_exists(self):
+        values = {"a": 7, "b": "foo"}
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query = (
+            self._make_one(collection)
+            .where("c", "<=", 20)
+            .order_by("c", "DESCENDING")
+            .start_at(snapshot)
+        )
+        expected = [
+            query._make_order("c", "DESCENDING"),
+            query._make_order("__name__", "DESCENDING"),
+        ]
+        self.assertEqual(query._normalize_orders(), expected)
+
+    def test__normalize_orders_wo_orders_w_snapshot_cursor_w_neq_where(self):
+        values = {"a": 7, "b": "foo"}
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        collection = self._make_collection("here")
+        query = self._make_one(collection).where("c", "<=", 20).end_at(snapshot)
+        expected = [
+            query._make_order("c", "ASCENDING"),
+            query._make_order("__name__", "ASCENDING"),
+        ]
+        self.assertEqual(query._normalize_orders(), expected)
+
     def test__normalize_cursor_none(self):
         query = self._make_one(mock.sentinel.parent)
         self.assertIsNone(query._normalize_cursor(None, query._orders))
@@ -603,6 +698,16 @@ def test__normalize_cursor_as_dict_hit(self):
 
         self.assertEqual(query._normalize_cursor(cursor, query._orders), ([1], True))
 
+    def test__normalize_cursor_as_snapshot_hit(self):
+        values = {"b": 1}
+        docref = self._make_docref("here", "doc_id")
+        snapshot = self._make_snapshot(docref, values)
+        cursor = (snapshot, True)
+        collection = self._make_collection("here")
+        query = self._make_one(collection).order_by("b", "ASCENDING")
+
+        self.assertEqual(query._normalize_cursor(cursor, query._orders), ([1], True))
+
     def test__normalize_cursor_w___name___w_slash(self):
         db_string = "projects/my-project/database/(default)"
         client = mock.Mock(spec=["_database_string"])
@@ -1206,6 +1311,10 @@ def test_success(self):
         self.assertEqual(self._call_fut(Query.ASCENDING), dir_class.ASCENDING)
         self.assertEqual(self._call_fut(Query.DESCENDING), dir_class.DESCENDING)
 
+        # Ints pass through
+        self.assertEqual(self._call_fut(dir_class.ASCENDING), dir_class.ASCENDING)
+        self.assertEqual(self._call_fut(dir_class.DESCENDING), dir_class.DESCENDING)
+
     def test_failure(self):
         with self.assertRaises(ValueError):
             self._call_fut("neither-ASCENDING-nor-DESCENDING")