fix: order normalization with descending query (#788)

Author: daniel-sanche

google/cloud/firestore_v1/base_query.py

@@ -85,6 +85,16 @@
  "not-in": _operator_enum.NOT_IN,
  "array_contains_any": _operator_enum.ARRAY_CONTAINS_ANY,
 }
+# set of operators that don't involve equlity comparisons
+# will be used in query normalization
+_INEQUALITY_OPERATORS = (
+ _operator_enum.LESS_THAN,
+ _operator_enum.LESS_THAN_OR_EQUAL,
+ _operator_enum.GREATER_THAN_OR_EQUAL,
+ _operator_enum.GREATER_THAN,
+ _operator_enum.NOT_EQUAL,
+ _operator_enum.NOT_IN,
+)
 _BAD_OP_STRING = "Operator string {!r} is invalid. Valid choices are: {}."
 _BAD_OP_NAN_NULL = 'Only an equality filter ("==") can be used with None or NaN values'
 _INVALID_WHERE_TRANSFORM = "Transforms cannot be used as where values."
@@ -858,28 +868,21 @@ def _normalize_orders(self) -> list:
  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")
- ]
+ # added orders should use direction of last order
+ last_direction = orders[-1].direction if orders else BaseQuery.ASCENDING
  order_keys = [order.field.field_path for order in orders]
  for filter_ in self._field_filters:
  # FieldFilter.Operator should not compare equal to
  # UnaryFilter.Operator, but it does
  if isinstance(filter_.op, StructuredQuery.FieldFilter.Operator):
  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))
+ # skip equality filters and filters on fields already ordered
+ if filter_.op in _INEQUALITY_OPERATORS and field not in order_keys:
+ orders.append(self._make_order(field, last_direction))
+ # add __name__ if not already in orders
+ if "__name__" not in [order.field.field_path for order in orders]:
+ orders.append(self._make_order("__name__", last_direction))
 
  return orders
 

tests/unit/v1/test_base_query.py

@@ -1207,6 +1207,61 @@ def test_basequery__normalize_orders_w_name_orders_w_none_cursor():
  assert query._normalize_orders() == expected
 
 
+def test_basequery__normalize_orders_w_cursor_descending():
+ """
+ Test case for b/306472103
+ """
+ from google.cloud.firestore_v1.base_query import FieldFilter
+
+ collection = _make_collection("here")
+ snapshot = _make_snapshot(_make_docref("here", "doc_id"), {"a": 1, "b": 2})
+ query = (
+ _make_base_query(collection)
+ .where(filter=FieldFilter("a", "==", 1))
+ .where(filter=FieldFilter("b", "in", [1, 2, 3]))
+ .order_by("c", "DESCENDING")
+ )
+ query_w_snapshot = query.start_after(snapshot)
+
+ normalized = query._normalize_orders()
+ expected = [query._make_order("c", "DESCENDING")]
+ assert normalized == expected
+
+ normalized_w_snapshot = query_w_snapshot._normalize_orders()
+ expected_w_snapshot = expected + [query._make_order("__name__", "DESCENDING")]
+ assert normalized_w_snapshot == expected_w_snapshot
+
+
+def test_basequery__normalize_orders_w_cursor_descending_w_inequality():
+ """
+ Test case for b/306472103, with extra ineuality filter in "where" clause
+ """
+ from google.cloud.firestore_v1.base_query import FieldFilter
+
+ collection = _make_collection("here")
+ snapshot = _make_snapshot(_make_docref("here", "doc_id"), {"a": 1, "b": 2})
+ query = (
+ _make_base_query(collection)
+ .where(filter=FieldFilter("a", "==", 1))
+ .where(filter=FieldFilter("b", "in", [1, 2, 3]))
+ .where(filter=FieldFilter("c", "not-in", [4, 5, 6]))
+ .order_by("d", "DESCENDING")
+ )
+ query_w_snapshot = query.start_after(snapshot)
+
+ normalized = query._normalize_orders()
+ expected = [query._make_order("d", "DESCENDING")]
+ assert normalized == expected
+
+ normalized_w_snapshot = query_w_snapshot._normalize_orders()
+ expected_w_snapshot = [
+ query._make_order("d", "DESCENDING"),
+ query._make_order("c", "DESCENDING"),
+ query._make_order("__name__", "DESCENDING"),
+ ]
+ assert normalized_w_snapshot == expected_w_snapshot
+
+
 def test_basequery__normalize_cursor_none():
  query = _make_base_query(mock.sentinel.parent)
  assert query._normalize_cursor(None, query._orders) is None