Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 22 commits into from
Jan 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions ee/billing/quota_limiting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand Down Expand Up @@ -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

Expand Down
89 changes: 88 additions & 1 deletion ee/billing/test/test_quota_limiting.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import time
from unittest.mock import patch
from uuid import uuid4

from dateutil.relativedelta import relativedelta
Expand All @@ -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,
Expand All @@ -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:
xrdt marked this conversation as resolved.
Show resolved Hide resolved
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)
xrdt marked this conversation as resolved.
Show resolved Hide resolved
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:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unsure about this test, flag is on it seems? Is there another reason nothing is captured?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The conditional on line 95 of the quota_limiting.py file should lazy eval and fail since the org hasn't exceed their quota limit.

# 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 = {}
Expand Down Expand Up @@ -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},
Expand Down Expand Up @@ -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")
]
Expand Down
2 changes: 1 addition & 1 deletion ee/clickhouse/models/test/__snapshots__/test_property.ambr
Original file line number Diff line number Diff line change
Expand Up @@ -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([
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading