Ensure Subscription Update doesn't update the gifted quotas (#2012)

- add a separate OrgQuotasIn where all quota updates are optional
- ensure gifted quotas are never updated as part of org update
- update tests
This commit is contained in:
Ilya Kreymer 2024-08-20 13:15:03 -07:00 committed by GitHub
parent 351e92ae2f
commit 8c9a14b6a2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 48 additions and 29 deletions

View File

@ -1115,12 +1115,28 @@ REASON_CANCELED = "subscriptionCanceled"
class OrgQuotas(BaseModel):
"""Organization quotas (settable by superadmin)"""
maxConcurrentCrawls: Optional[int] = 0
maxPagesPerCrawl: Optional[int] = 0
storageQuota: Optional[int] = 0
maxExecMinutesPerMonth: Optional[int] = 0
extraExecMinutes: Optional[int] = 0
giftedExecMinutes: Optional[int] = 0
storageQuota: int = 0
maxExecMinutesPerMonth: int = 0
maxConcurrentCrawls: int = 0
maxPagesPerCrawl: int = 0
extraExecMinutes: int = 0
giftedExecMinutes: int = 0
# ============================================================================
class OrgQuotasIn(BaseModel):
"""Update for existing OrgQuotas"""
storageQuota: Optional[int] = None
maxExecMinutesPerMonth: Optional[int] = None
maxConcurrentCrawls: Optional[int] = None
maxPagesPerCrawl: Optional[int] = None
extraExecMinutes: Optional[int] = None
giftedExecMinutes: Optional[int] = None
# ============================================================================
@ -1176,7 +1192,7 @@ class SubscriptionUpdate(BaseModel):
planId: str
futureCancelDate: Optional[datetime] = None
quotas: Optional[OrgQuotas] = None
quotas: Optional[OrgQuotasIn] = None
# ============================================================================

View File

@ -32,6 +32,7 @@ from .models import (
Organization,
StorageRef,
OrgQuotas,
OrgQuotasIn,
OrgQuotaUpdate,
OrgReadOnlyUpdate,
OrgReadOnlyOnCancel,
@ -489,6 +490,8 @@ class OrgOps:
org = Organization.from_dict(org_data)
if update.quotas:
# don't change gifted minutes here
update.quotas.giftedExecMinutes = None
await self.update_quotas(org, update.quotas)
return org
@ -512,7 +515,7 @@ class OrgOps:
res = await self.orgs.find_one_and_update({"_id": org.id}, {"$set": set_dict})
return res is not None
async def update_quotas(self, org: Organization, quotas: OrgQuotas) -> None:
async def update_quotas(self, org: Organization, quotas: OrgQuotasIn) -> None:
"""update organization quotas"""
previous_extra_mins = (
@ -1458,7 +1461,7 @@ def init_orgs_api(
@router.post("/quotas", tags=["organizations"], response_model=UpdatedResponse)
async def update_quotas(
quotas: OrgQuotas,
quotas: OrgQuotasIn,
org: Organization = Depends(org_owner_dep),
user: User = Depends(user_dep),
):

View File

@ -479,10 +479,10 @@ def test_subscription_events_log(admin_auth_headers, non_default_org_id):
"quotas": {
"maxPagesPerCrawl": 50,
"storageQuota": 500000,
"extraExecMinutes": 0,
"giftedExecMinutes": 0,
"maxConcurrentCrawls": 0,
"maxExecMinutesPerMonth": 0,
"extraExecMinutes": None,
"giftedExecMinutes": None,
"maxConcurrentCrawls": None,
"maxExecMinutesPerMonth": None,
},
},
{"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@ -547,10 +547,10 @@ def test_subscription_events_log_filter_sub_id(admin_auth_headers):
"quotas": {
"maxPagesPerCrawl": 50,
"storageQuota": 500000,
"extraExecMinutes": 0,
"giftedExecMinutes": 0,
"maxConcurrentCrawls": 0,
"maxExecMinutesPerMonth": 0,
"extraExecMinutes": None,
"giftedExecMinutes": None,
"maxConcurrentCrawls": None,
"maxExecMinutesPerMonth": None,
},
},
{"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@ -608,10 +608,10 @@ def test_subscription_events_log_filter_oid(admin_auth_headers):
"quotas": {
"maxPagesPerCrawl": 50,
"storageQuota": 500000,
"extraExecMinutes": 0,
"giftedExecMinutes": 0,
"maxConcurrentCrawls": 0,
"maxExecMinutesPerMonth": 0,
"extraExecMinutes": None,
"giftedExecMinutes": None,
"maxConcurrentCrawls": None,
"maxExecMinutesPerMonth": None,
},
},
{"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@ -643,10 +643,10 @@ def test_subscription_events_log_filter_plan_id(admin_auth_headers):
"quotas": {
"maxPagesPerCrawl": 50,
"storageQuota": 500000,
"extraExecMinutes": 0,
"giftedExecMinutes": 0,
"maxConcurrentCrawls": 0,
"maxExecMinutesPerMonth": 0,
"extraExecMinutes": None,
"giftedExecMinutes": None,
"maxConcurrentCrawls": None,
"maxExecMinutesPerMonth": None,
},
}
]
@ -694,10 +694,10 @@ def test_subscription_events_log_filter_status(admin_auth_headers):
"quotas": {
"maxPagesPerCrawl": 50,
"storageQuota": 500000,
"extraExecMinutes": 0,
"giftedExecMinutes": 0,
"maxConcurrentCrawls": 0,
"maxExecMinutesPerMonth": 0,
"extraExecMinutes": None,
"giftedExecMinutes": None,
"maxConcurrentCrawls": None,
"maxExecMinutesPerMonth": None,
},
},
]