fix(query): Guard against ordering by non unique columns in RangeQuerySetWrapper (#69367)

Follow-up from inc-709. This incident was caused by
`RangeQuerySetWrapper` looping infinitely and slamming queries into the
database. It did this because we ordered on a non-unique column, and
currently the logic here will loop forever when the last two rows of the
result set have the same value for the sort col.

Author: wedamija

src/sentry/tasks/auto_ongoing_issues.py

@@ -140,6 +140,7 @@ def get_total_count(results):
                 limit=ITERATOR_CHUNK * CHILD_TASK_COUNT,
                 callbacks=[get_total_count],
                 order_by="first_seen",
+                override_unique_safety_check=True,
             ),
             ITERATOR_CHUNK,
         ):

src/sentry/utils/query.py

@@ -90,6 +90,7 @@ def __init__(
         order_by="pk",
         callbacks=(),
         result_value_getter=None,
+        override_unique_safety_check=True,
     ):
         # Support for slicing
         if queryset.query.low_mark == 0 and not (
@@ -114,6 +115,16 @@ def __init__(
         self.callbacks = callbacks
         self.result_value_getter = result_value_getter
 
+        order_by_col = queryset.model._meta.get_field(order_by if order_by != "pk" else "id")
+        if not override_unique_safety_check and not order_by_col.unique:
+            # TODO: Ideally we could fix this bug and support ordering by a non unique col
+            raise InvalidQuerySetError(
+                "Order by column must be unique, otherwise this wrapper can get "
+                "stuck in an infinite loop. If you're sure your data is unique, "
+                "you can disable this by passing "
+                "`override_unique_safety_check=True`"
+            )
+
     def __iter__(self):
         if self.min_value is not None:
             cur_value = self.min_value

tests/sentry/utils/test_query.py

@@ -1,10 +1,13 @@
+import pytest
+
 from sentry.db.models.query import in_iexact
 from sentry.models.organization import Organization
 from sentry.models.user import User
 from sentry.models.userreport import UserReport
 from sentry.testutils.cases import TestCase
 from sentry.testutils.silo import no_silo_test
 from sentry.utils.query import (
+    InvalidQuerySetError,
     RangeQuerySetWrapper,
     RangeQuerySetWrapperWithProgressBar,
     RangeQuerySetWrapperWithProgressBarApprox,
@@ -54,6 +57,20 @@ def test_empty(self):
         qs = User.objects.all()
         assert len(list(self.range_wrapper(qs, step=2))) == 0
 
+    def test_order_by_non_unique_fails(self):
+        qs = User.objects.all()
+        with pytest.raises(InvalidQuerySetError):
+            self.range_wrapper(qs, order_by="name", override_unique_safety_check=False)
+
+        # Shouldn't error if the safety check is disabled
+        self.range_wrapper(qs, order_by="name", override_unique_safety_check=True)
+
+    def test_order_by_unique(self):
+        self.create_user()
+        qs = User.objects.all()
+        self.range_wrapper(qs, order_by="username", override_unique_safety_check=False)
+        assert len(list(self.range_wrapper(qs, order_by="username", step=2))) == 1
+
 
 @no_silo_test
 class RangeQuerySetWrapperWithProgressBarTest(RangeQuerySetWrapperTest):