chore(data-warehouse): Updated billing limits to work per sync (#24752)

PostHog · Sep 4, 2024 · 9842e97 · 9842e97
1 parent e26da47
commit 9842e97
Show file tree

Hide file tree

Showing 22 changed files with 190 additions and 415 deletions.
diff --git a/ee/billing/quota_limiting.py b/ee/billing/quota_limiting.py
@@ -122,7 +122,7 @@ def org_quota_limited_until(
     if organization.never_drop_data or trust_score == 15:
         return None
 
-    team_tokens = get_team_attribute_by_quota_resource(organization, resource)
+    team_tokens = get_team_attribute_by_quota_resource(organization)
     team_being_limited = any(x in previously_quota_limited_team_tokens for x in team_tokens)
 
     if team_being_limited:
@@ -237,7 +237,7 @@ def sync_org_quota_limits(organization: Organization):
         previously_quota_limited_team_tokens = list_limited_team_attributes(
             resource, QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY
         )
-        team_attributes = get_team_attribute_by_quota_resource(organization, resource)
+        team_attributes = get_team_attribute_by_quota_resource(organization)
         result = org_quota_limited_until(organization, resource, previously_quota_limited_team_tokens)
 
         if result:
@@ -264,24 +264,14 @@ def sync_org_quota_limits(organization: Organization):
             remove_limited_team_tokens(resource, team_attributes, QuotaLimitingCaches.QUOTA_LIMITING_SUSPENDED_KEY)
 
 
-def get_team_attribute_by_quota_resource(organization: Organization, resource: QuotaResource):
-    if resource in [QuotaResource.EVENTS, QuotaResource.RECORDINGS]:
-        team_tokens: list[str] = [x for x in list(organization.teams.values_list("api_token", flat=True)) if x]
+def get_team_attribute_by_quota_resource(organization: Organization):
+    team_tokens: list[str] = [x for x in list(organization.teams.values_list("api_token", flat=True)) if x]
 
-        if not team_tokens:
-            capture_exception(Exception(f"quota_limiting: No team tokens found for organization: {organization.id}"))
-            return
+    if not team_tokens:
+        capture_exception(Exception(f"quota_limiting: No team tokens found for organization: {organization.id}"))
+        return
 
-        return team_tokens
-
-    if resource == QuotaResource.ROWS_SYNCED:
-        team_ids: list[str] = [x for x in list(organization.teams.values_list("id", flat=True)) if x]
-
-        if not team_ids:
-            capture_exception(Exception(f"quota_limiting: No team ids found for organization: {organization.id}"))
-            return
-
-        return team_ids
+    return team_tokens
 
 
 def set_org_usage_summary(

diff --git a/ee/billing/test/test_quota_limiting.py b/ee/billing/test/test_quota_limiting.py
@@ -92,7 +92,7 @@ def test_quota_limiting_feature_flag_enabled(self, patch_feature_enabled, patch_
 
         patch_capture.reset_mock()
         # Add this org to the redis cache.
-        team_tokens = get_team_attribute_by_quota_resource(self.organization, QuotaResource.EVENTS)
+        team_tokens = get_team_attribute_by_quota_resource(self.organization)
         add_limited_team_tokens(
             QuotaResource.EVENTS,
             {x: 1612137599 for x in team_tokens},
@@ -715,7 +715,7 @@ def test_sync_org_quota_limits(self):
             # rows_synced uses teams, not tokens
             assert sorted(
                 list_limited_team_attributes(QuotaResource.ROWS_SYNCED, QuotaLimitingCaches.QUOTA_LIMITER_CACHE_KEY)
-            ) == sorted(["1337", str(self.team.pk), str(other_team.pk)])
+            ) == sorted(["1337", str(self.team.api_token), str(other_team.api_token)])
 
             self.organization.usage["events"]["usage"] = 80
             self.organization.usage["rows_synced"]["usage"] = 36
@@ -748,7 +748,7 @@ def test_sync_org_quota_limits(self):
                 list_limited_team_attributes(
                     QuotaResource.ROWS_SYNCED, QuotaLimitingCaches.QUOTA_LIMITING_SUSPENDED_KEY
                 )
-            ) == sorted([str(self.team.pk), str(other_team.pk)])
+            ) == sorted([str(self.team.api_token), str(other_team.api_token)])
 
             self.organization.usage["events"]["usage"] = 80
             self.organization.usage["rows_synced"]["usage"] = 36

diff --git a/frontend/src/scenes/data-warehouse/settings/DataWarehouseManagedSourcesTable.tsx b/frontend/src/scenes/data-warehouse/settings/DataWarehouseManagedSourcesTable.tsx
@@ -27,6 +27,7 @@ const StatusTagSetting = {
     Completed: 'success',
     Error: 'danger',
     Failed: 'danger',
+    'Billing limits': 'danger',
 }
 
 export function DataWarehouseManagedSourcesTable(): JSX.Element {

diff --git a/frontend/src/scenes/data-warehouse/settings/source/Schemas.tsx b/frontend/src/scenes/data-warehouse/settings/source/Schemas.tsx
@@ -47,6 +47,7 @@ const StatusTagSetting = {
     Completed: 'success',
     Error: 'danger',
     Failed: 'danger',
+    'Billing limits': 'danger',
 }
 
 export const SchemaTable = ({ schemas, isLoading }: SchemaTableProps): JSX.Element => {

diff --git a/frontend/src/scenes/data-warehouse/settings/source/Syncs.tsx b/frontend/src/scenes/data-warehouse/settings/source/Syncs.tsx
@@ -11,7 +11,7 @@ const StatusTagSetting: Record<ExternalDataJob['status'], LemonTagType> = {
     Running: 'primary',
     Completed: 'success',
     Failed: 'danger',
-    Cancelled: 'default',
+    'Billing limits': 'danger',
 }
 
 interface SyncsProps {

diff --git a/frontend/src/types.ts b/frontend/src/types.ts
@@ -3924,7 +3924,7 @@ export interface ExternalDataSourceSchema extends SimpleExternalDataSourceSchema
 export interface ExternalDataJob {
     id: string
     created_at: string
-    status: 'Running' | 'Failed' | 'Completed' | 'Cancelled'
+    status: 'Running' | 'Failed' | 'Completed' | 'Billing limits'
     schema: SimpleExternalDataSourceSchema
     rows_synced: number
     latest_error: string