Skip to content

Commit

Permalink
fix: quota limiting memory usage (#20161)
Browse files Browse the repository at this point in the history
* only get the team data we need

* make sure we got all the fields (with assert num queries)
  • Loading branch information
raquelmsmith authored Feb 6, 2024
1 parent 0d55f92 commit 037541a
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 3 deletions.
11 changes: 9 additions & 2 deletions ee/billing/quota_limiting.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,15 @@ def update_all_org_billing_quotas(dry_run: bool = False) -> Dict[str, Dict[str,
)

teams: Sequence[Team] = list(
Team.objects.select_related("organization").exclude(
Q(organization__for_internal_metrics=True) | Q(is_demo=True)
Team.objects.select_related("organization")
.exclude(Q(organization__for_internal_metrics=True) | Q(is_demo=True))
.only(
"id",
"api_token",
"organization__id",
"organization__usage",
"organization__created_at",
"organization__never_drop_data",
)
)

Expand Down
4 changes: 3 additions & 1 deletion ee/billing/test/test_quota_limiting.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def test_dont_quota_limit_feature_flag_enabled(self, patch_feature_enabled, patc
team=self.team,
)
time.sleep(1)

result = update_all_org_billing_quotas()
patch_feature_enabled.assert_called_with(
QUOTA_LIMIT_DATA_RETENTION_FLAG,
Expand Down Expand Up @@ -89,7 +90,8 @@ def test_quota_limit_feature_flag_not_on(self, patch_feature_enabled, patch_capt
self.organization.save()

time.sleep(1)
result = update_all_org_billing_quotas()
with self.assertNumQueries(2):
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()
Expand Down

0 comments on commit 037541a

Please sign in to comment.