feat: use feature flag to control whether to quota limit or not (#19996)

* use feature flag to control whether to quota limit or not

* add simple test

* Update query snapshots

* fix up the feature flag activation logic

* confirm that the fixed feature flag logic works

* remove unnecessary change

* Update query snapshots

* Update query snapshots

* fix up tests

* another fix for tests

* Update query snapshots

* still fixing tests

* Update query snapshots

* pr feedback, more tests

* tune up ff enabled call, add capture check to test for quota limit changed

* i really need to run tests locally

---------

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

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,21 @@ 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,
+        groups={"organization": str(organization.id)},
+        group_properties={
+            "organization": str(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": 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,
@@ -22,10 +24,83 @@
 
 
 class TestQuotaLimiting(BaseTest):
+    CLASS_DATA_LEVEL_SETUP = False
+
     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,
+            groups={"organization": str(self.organization.id)},
+            group_properties={"organization": str(self.organization.id)},
+        )
+        patch_capture.assert_called_once_with(
+            str(self.organization.id),
+            "quota limiting suspended",
+            properties={"current_usage": 109},
+            groups={"instance": "http://localhost:8000", "organization": str(self.organization.id)},
+        )
+        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.
+        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()
+
+        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 = {}
@@ -53,7 +128,8 @@ def test_billing_rate_limit_not_set_if_missing_org_usage(self) -> None:
         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(self) -> None:
+    @patch("posthoganalytics.capture")
+    def test_billing_rate_limit(self, patch_capture) -> None:
         with self.settings(USE_TZ=False):
             self.organization.usage = {
                 "events": {"usage": 99, "limit": 100},
@@ -82,6 +158,17 @@ def test_billing_rate_limit(self) -> None:
         assert result["recordings"] == {}
         assert result["rows_synced"] == {}
 
+        patch_capture.assert_called_once_with(
+            org_id,
+            "organization quota limits changed",
+            properties={
+                "quota_limited_events": 1612137599,
+                "quota_limited_recordings": 1612137599,
+                "quota_limited_rows_synced": None,
+            },
+            groups={"instance": "http://localhost:8000", "organization": org_id},
+        )
+
         assert self.redis_client.zrange(f"{QUOTA_LIMITER_CACHE_KEY}events", 0, -1) == [
             self.team.api_token.encode("UTF-8")
         ]

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

@@ -58,7 +58,7 @@
 # ---
 # name: test_parse_groups_persons_edge_case_with_single_filter
   tuple(
-    'AND (  has(%(vglobalperson_0)s, "pmat_email"))',
+    'AND (  has(%(vglobalperson_0)s, replaceRegexpAll(JSONExtractRaw(person_props, %(kglobalperson_0)s), \'^"|"$\', \'\')))',
     dict({
       'kglobalperson_0': 'email',
       'vglobalperson_0': list([

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:129 celery:posthog.tasks.tasks.sync_insight_caching_state */
+  /* user_id:131 celery:posthog.tasks.tasks.sync_insight_caching_state */
   SELECT team_id,
          date_diff('second', max(timestamp), now()) AS age
   FROM events