From c3623b039f51968bcff3816d50925f62295e3063 Mon Sep 17 00:00:00 2001 From: Bianca Yang <21014901+xrdt@users.noreply.github.com> Date: Wed, 31 Jan 2024 13:37:40 -0800 Subject: [PATCH] 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 | 18 ++++ ee/billing/test/test_quota_limiting.py | 89 ++++++++++++++++++- .../test/__snapshots__/test_property.ambr | 2 +- ...ickhouse_experiment_secondary_results.ambr | 2 +- 4 files changed, 108 insertions(+), 3 deletions(-) diff --git a/ee/billing/quota_limiting.py b/ee/billing/quota_limiting.py index 0809266b1db64..0476ce5d8bb2a 100644 --- a/ee/billing/quota_limiting.py +++ b/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 diff --git a/ee/billing/test/test_quota_limiting.py b/ee/billing/test/test_quota_limiting.py index b8e68c235b2c5..bffed0f9f3578 100644 --- a/ee/billing/test/test_quota_limiting.py +++ b/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") ] diff --git a/ee/clickhouse/models/test/__snapshots__/test_property.ambr b/ee/clickhouse/models/test/__snapshots__/test_property.ambr index d804c6ca9363e..0f4e20f1888ab 100644 --- a/ee/clickhouse/models/test/__snapshots__/test_property.ambr +++ b/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([ diff --git a/ee/clickhouse/views/test/__snapshots__/test_clickhouse_experiment_secondary_results.ambr b/ee/clickhouse/views/test/__snapshots__/test_clickhouse_experiment_secondary_results.ambr index e5c434f501653..9012dcd46ecd9 100644 --- a/ee/clickhouse/views/test/__snapshots__/test_clickhouse_experiment_secondary_results.ambr +++ b/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