merge quota set & add methods into one quota update method

Author: emma-sg

backend/btrixcloud/orgs.py

@@ -19,6 +19,7 @@
     Dict,
     Callable,
     List,
+    Literal,
     AsyncGenerator,
     Any,
 )
@@ -611,71 +612,16 @@ async def update_proxies(self, org: Organization, proxies: OrgProxies) -> None:
         )
 
     async def update_quotas(
-        self, org: Organization, quotas: OrgQuotasIn, context: str | None = None
+        self,
+        org: Organization,
+        quotas: OrgQuotasIn,
+        mode: Literal["set", "add"],
+        context: str | None = None,
     ) -> None:
         """update organization quotas"""
 
-        previous_extra_mins = (
-            org.quotas.extraExecMinutes
-            if (org.quotas and org.quotas.extraExecMinutes)
-            else 0
-        )
-        previous_gifted_mins = (
-            org.quotas.giftedExecMinutes
-            if (org.quotas and org.quotas.giftedExecMinutes)
-            else 0
-        )
-
         quotas.context = None
 
-        update = quotas.dict(
-            exclude_unset=True, exclude_defaults=True, exclude_none=True
-        )
-
-        await self.orgs.find_one_and_update(
-            {"_id": org.id},
-            {
-                "$set": {
-                    "quotas": update,
-                },
-                "$push": {
-                    "quotaUpdates": OrgQuotaUpdate(
-                        update=OrgQuotas(**update), modified=dt_now(), context=context
-                    ).model_dump()
-                },
-            },
-        )
-
-        # Inc org available fields for extra/gifted execution time as needed
-        if quotas.extraExecMinutes is not None:
-            extra_secs_diff = (quotas.extraExecMinutes - previous_extra_mins) * 60
-            if org.extraExecSecondsAvailable + extra_secs_diff <= 0:
-                await self.orgs.find_one_and_update(
-                    {"_id": org.id},
-                    {"$set": {"extraExecSecondsAvailable": 0}},
-                )
-            else:
-                await self.orgs.find_one_and_update(
-                    {"_id": org.id},
-                    {"$inc": {"extraExecSecondsAvailable": extra_secs_diff}},
-                )
-
-        if quotas.giftedExecMinutes is not None:
-            gifted_secs_diff = (quotas.giftedExecMinutes - previous_gifted_mins) * 60
-            if org.giftedExecSecondsAvailable + gifted_secs_diff <= 0:
-                await self.orgs.find_one_and_update(
-                    {"_id": org.id},
-                    {"$set": {"giftedExecSecondsAvailable": 0}},
-                )
-            else:
-                await self.orgs.find_one_and_update(
-                    {"_id": org.id},
-                    {"$inc": {"giftedExecSecondsAvailable": gifted_secs_diff}},
-                )
-
-    async def add_to_org_quotas(
-        self, org: Organization, quotas: OrgQuotasIn, context: str | None
-    ):
         previous_extra_mins = (
             org.quotas.extraExecMinutes
             if (org.quotas and org.quotas.extraExecMinutes)
@@ -687,56 +633,65 @@ async def add_to_org_quotas(
             else 0
         )
 
-        update: dict[str, Any] = {
-            "$inc": {},
-        }
+        if mode == "add":
+            increment_update: dict[str, Any] = {
+                "$inc": {},
+            }
 
-        for field, value in quotas.model_dump(exclude_unset=True).items():
-            if field == "context" or value is None:
-                continue
-            inc = max(value, -org.quotas.model_dump().get(field, 0))
-            update["$inc"][f"quotas.{field}"] = inc
+            for field, value in quotas.model_dump(
+                exclude_unset=True, exclude_defaults=True, exclude_none=True
+            ).items():
+                if field == "context" or value is None:
+                    continue
+                inc = max(value, -org.quotas.model_dump().get(field, 0))
+                increment_update["$inc"][f"quotas.{field}"] = inc
 
-        updated_org = await self.orgs.find_one_and_update(
-            {"_id": org.id},
-            update,
-            projection={"quotas": True},
-            return_document=ReturnDocument.AFTER,
-        )
-        updated_quotas = OrgQuotas(**updated_org["quotas"])
+            updated_org = await self.orgs.find_one_and_update(
+                {"_id": org.id},
+                increment_update,
+                projection={"quotas": True},
+                return_document=ReturnDocument.AFTER,
+            )
+            quotas = OrgQuotasIn(**updated_org["quotas"])
 
-        quotaUpdate: dict[str, dict[str, dict[str, Any] | int]] = {
+        update: dict[str, dict[str, dict[str, Any] | int]] = {
             "$push": {
                 "quotaUpdates": OrgQuotaUpdate(
                     modified=dt_now(),
-                    update=updated_quotas,
+                    update=OrgQuotas(
+                        **quotas.model_dump(
+                            exclude_unset=True, exclude_defaults=True, exclude_none=True
+                        )
+                    ),
                     context=context,
                 ).model_dump()
             },
             "$inc": {},
             "$set": {},
         }
 
+        if mode == "set":
+            increment_update = quotas.model_dump(
+                exclude_unset=True, exclude_defaults=True, exclude_none=True
+            )
+            update["$set"]["quotas"] = increment_update
+
         # Inc org available fields for extra/gifted execution time as needed
-        if updated_quotas.extraExecMinutes is not None:
-            extra_secs_diff = (
-                updated_quotas.extraExecMinutes - previous_extra_mins
-            ) * 60
+        if quotas.extraExecMinutes is not None:
+            extra_secs_diff = (quotas.extraExecMinutes - previous_extra_mins) * 60
             if org.extraExecSecondsAvailable + extra_secs_diff <= 0:
-                quotaUpdate["$set"]["extraExecSecondsAvailable"] = 0
+                update["$set"]["extraExecSecondsAvailable"] = 0
             else:
-                quotaUpdate["$inc"]["extraExecSecondsAvailable"] = extra_secs_diff
+                update["$inc"]["extraExecSecondsAvailable"] = extra_secs_diff
 
-        if updated_quotas.giftedExecMinutes is not None:
-            gifted_secs_diff = (
-                updated_quotas.giftedExecMinutes - previous_gifted_mins
-            ) * 60
+        if quotas.giftedExecMinutes is not None:
+            gifted_secs_diff = (quotas.giftedExecMinutes - previous_gifted_mins) * 60
             if org.giftedExecSecondsAvailable + gifted_secs_diff <= 0:
-                quotaUpdate["$set"]["giftedExecSecondsAvailable"] = 0
+                update["$set"]["giftedExecSecondsAvailable"] = 0
             else:
-                quotaUpdate["$inc"]["giftedExecSecondsAvailable"] = gifted_secs_diff
+                update["$inc"]["giftedExecSecondsAvailable"] = gifted_secs_diff
 
-        await self.orgs.find_one_and_update({"_id": org.id}, quotaUpdate)
+        await self.orgs.find_one_and_update({"_id": org.id}, update)
 
     async def update_event_webhook_urls(
         self, org: Organization, urls: OrgWebhookUrls
@@ -1751,7 +1706,7 @@ async def update_quotas(
         if not user.is_superuser:
             raise HTTPException(status_code=403, detail="Not Allowed")
 
-        await ops.update_quotas(org, quotas, quotas.context)
+        await ops.update_quotas(org, quotas, mode="set", context=quotas.context)
 
         return {"updated": True}
 
@@ -1762,7 +1717,7 @@ async def update_quotas_add(
         quotas: OrgQuotasIn,
         org: Organization = Depends(org_superuser_or_shared_secret_dep),
     ):
-        await ops.add_to_org_quotas(org, quotas, quotas.context)
+        await ops.update_quotas(org, quotas, mode="add", context=quotas.context)
 
         return {"updated": True}