feat: use feature flag to control whether to quota limit or not

ee/billing/quota_limiting.py

@@ -4,6 +4,7 @@
 from typing import Dict, List, Mapping, Optional, Sequence, TypedDict, cast
 
 import dateutil.parser
+import posthoganalytics
 from django.db.models import Q
 from django.utils import timezone
 from sentry_sdk import capture_exception
@@ -23,6 +24,8 @@
 
 QUOTA_LIMITER_CACHE_KEY = "@posthog/quota-limits/"
 
+QUOTA_LIMIT_DATA_RETENTION_FLAG = "retain-data-past-quota-limit"
+
 
 class QuotaResource(Enum):
     EVENTS = "events"
@@ -89,6 +92,12 @@ def org_quota_limited_until(organization: Organization, resource: QuotaResource)
     if is_quota_limited and organization.never_drop_data:
         return None
 
+    if is_quota_limited and posthoganalytics.feature_enabled(QUOTA_LIMIT_DATA_RETENTION_FLAG, organization.id):
+        # Don't drop data for this org but record that they __would have__ been
+        # limited.
+        report_organization_action(organization, "quota_limiting_suspended", properties={"current_usage": todays_usage})
+        return None
+
     if is_quota_limited and billing_period_end:
         return billing_period_end
 

ee/billing/test/test_quota_limiting.py

@@ -1,4 +1,5 @@
 import time
+from unittest.mock import patch
 from uuid import uuid4
 
 from dateutil.relativedelta import relativedelta
@@ -7,6 +8,7 @@
 from freezegun import freeze_time
 
 from ee.billing.quota_limiting import (
+    QUOTA_LIMIT_DATA_RETENTION_FLAG,
     QUOTA_LIMITER_CACHE_KEY,
     QuotaResource,
     list_limited_team_attributes,
@@ -26,6 +28,68 @@ def setUp(self) -> None:
         super().setUp()
         self.redis_client = get_client()
 
+    @patch("posthoganalytics.capture")
+    @patch("posthoganalytics.feature_enabled", return_value=True)
+    def test_dont_quota_limit_feature_flag_enabled(self, patch_feature_enabled, patch_capture) -> None:
+        with self.settings(USE_TZ=False):
+            self.organization.usage = {
+                "events": {"usage": 99, "limit": 100},
+                "recordings": {"usage": 1, "limit": 100},
+                "rows_synced": {"usage": 5, "limit": 100},
+                "period": ["2021-01-01T00:00:00Z", "2021-01-31T23:59:59Z"],
+            }
+            self.organization.save()
+
+            distinct_id = str(uuid4())
+
+            # add a bunch of events so that the organization is over the limit
+            # Because the feature flag is enabled
+            for _ in range(0, 10):
+                _create_event(
+                    distinct_id=distinct_id,
+                    event="$event1",
+                    properties={"$lib": "$web"},
+                    timestamp=now() - relativedelta(hours=1),
+                    team=self.team,
+                )
+        time.sleep(1)
+        result = update_all_org_billing_quotas()
+        patch_feature_enabled.assert_called_with(QUOTA_LIMIT_DATA_RETENTION_FLAG, self.organization.id)
+        patch_capture.called_once_with(self.organization.id, "quota_limiting_suspended", {"current_usage": 109})
+        assert result["events"] == {}
+        assert result["recordings"] == {}
+        assert result["rows_synced"] == {}
+
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}events", 0, -1) == []
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}recordings", 0, -1) == []
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}rows_synced", 0, -1) == []
+
+    @patch("posthoganalytics.capture")
+    @patch("posthoganalytics.feature_enabled", return_value=True)
+    def test_quota_limit_feature_flag_not_on(self, patch_feature_enabled, patch_capture) -> None:
+        # Confirm that we don't send an event if they weren't going to be limited.
+        with self.settings(USE_TZ=False):
+            self.organization.usage = {
+                "events": {"usage": 60, "limit": 100},
+                "recordings": {"usage": 1, "limit": 100},
+                "rows_synced": {"usage": 5, "limit": 100},
+                "period": ["2021-01-01T00:00:00Z", "2021-01-31T23:59:59Z"],
+            }
+            self.organization.save()
+
+        time.sleep(1)
+        result = update_all_org_billing_quotas()
+        # Shouldn't be called due to lazy evaluation of the conditional
+        patch_feature_enabled.assert_not_called()
+        patch_capture.assert_not_called()
+        assert result["events"] == {}
+        assert result["recordings"] == {}
+        assert result["rows_synced"] == {}
+
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}events", 0, -1) == []
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}recordings", 0, -1) == []
+        assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}rows_synced", 0, -1) == []
+
     def test_billing_rate_limit_not_set_if_missing_org_usage(self) -> None:
         with self.settings(USE_TZ=False):
             self.organization.usage = {}

ee/clickhouse/models/test/__snapshots__/test_property.ambr

@@ -147,7 +147,7 @@
       ))
     ''',
     dict({
-      'global_cohort_id_0': 47,
+      'global_cohort_id_0': 1,
       'global_version_0': None,
     }),
   )

ee/clickhouse/views/test/__snapshots__/test_clickhouse_experiment_secondary_results.ambr

@@ -1,7 +1,7 @@
 # serializer version: 1
 # name: ClickhouseTestExperimentSecondaryResults.test_basic_secondary_metric_results
   '''
-  /* user_id:132 celery:posthog.tasks.tasks.sync_insight_caching_state */
+  /* user_id:134 celery:posthog.tasks.tasks.sync_insight_caching_state */
   SELECT team_id,
          date_diff('second', max(timestamp), now()) AS age
   FROM events