From 037541af8e06c8e58376f15d8383b2ea2ce9d912 Mon Sep 17 00:00:00 2001 From: Raquel Smith Date: Tue, 6 Feb 2024 15:59:20 -0800 Subject: [PATCH] fix: quota limiting memory usage (#20161) * only get the team data we need * make sure we got all the fields (with assert num queries) --- ee/billing/quota_limiting.py | 11 +++++++++-- ee/billing/test/test_quota_limiting.py | 4 +++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/ee/billing/quota_limiting.py b/ee/billing/quota_limiting.py index e492d3338cf7d..f8021237a4792 100644 --- a/ee/billing/quota_limiting.py +++ b/ee/billing/quota_limiting.py @@ -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", ) ) diff --git a/ee/billing/test/test_quota_limiting.py b/ee/billing/test/test_quota_limiting.py index ce71caa5a5697..ec9e5d360234f 100644 --- a/ee/billing/test/test_quota_limiting.py +++ b/ee/billing/test/test_quota_limiting.py @@ -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, @@ -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()