-
-
Notifications
You must be signed in to change notification settings - Fork 60
Add-on minute checkout session setup #2944
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
5f91e6a
b6f8979
22bbc95
aab6616
514a98f
c82e442
34fd646
0d51e2e
d95bd8f
9eadcb9
dd5add9
bf9ae05
1a8f005
4243fe6
db29e57
a7f9d39
f86c437
4f7b483
fe41b70
1f2f3c6
bb1ddae
7135f10
17228b8
edb74d9
a6ff231
bf9a3d5
6aff6e5
7fdc812
a162c1e
c9aed22
febe01f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,8 +12,19 @@ | |
| from uuid import UUID, uuid4 | ||
| from tempfile import NamedTemporaryFile | ||
|
|
||
| from typing import Optional, TYPE_CHECKING, Dict, Callable, List, AsyncGenerator, Any | ||
| from typing import ( | ||
| Awaitable, | ||
| Optional, | ||
| TYPE_CHECKING, | ||
| Dict, | ||
| Callable, | ||
| List, | ||
| Literal, | ||
| AsyncGenerator, | ||
| Any, | ||
| ) | ||
|
|
||
| from motor.motor_asyncio import AsyncIOMotorDatabase | ||
| from pydantic import ValidationError | ||
| from pymongo import ReturnDocument | ||
| from pymongo.errors import AutoReconnect, DuplicateKeyError | ||
|
|
@@ -546,9 +557,15 @@ async def update_subscription_data( | |
|
|
||
| org = Organization.from_dict(org_data) | ||
| if update.quotas: | ||
| # don't change gifted minutes here | ||
| # don't change gifted or extra minutes here | ||
| update.quotas.giftedExecMinutes = None | ||
| await self.update_quotas(org, update.quotas) | ||
| update.quotas.extraExecMinutes = None | ||
| await self.update_quotas( | ||
| org, | ||
| update.quotas, | ||
| mode="set", | ||
| context=f"subscription_change:{update.planId}", | ||
| ) | ||
|
|
||
| return org | ||
|
|
||
|
|
@@ -599,9 +616,17 @@ async def update_proxies(self, org: Organization, proxies: OrgProxies) -> None: | |
| }, | ||
| ) | ||
|
|
||
| async def update_quotas(self, org: Organization, quotas: OrgQuotasIn) -> None: | ||
| async def update_quotas( | ||
| self, | ||
| org: Organization, | ||
| quotas: OrgQuotasIn, | ||
| mode: Literal["set", "add"], | ||
| context: str | None = None, | ||
| ) -> None: | ||
| """update organization quotas""" | ||
|
|
||
| quotas.context = None | ||
|
|
||
| previous_extra_mins = ( | ||
| org.quotas.extraExecMinutes | ||
| if (org.quotas and org.quotas.extraExecMinutes) | ||
|
|
@@ -613,51 +638,65 @@ async def update_quotas(self, org: Organization, quotas: OrgQuotasIn) -> None: | |
| else 0 | ||
| ) | ||
|
|
||
| update = quotas.dict( | ||
| exclude_unset=True, exclude_defaults=True, exclude_none=True | ||
| ) | ||
| if mode == "add": | ||
| increment_update: dict[str, Any] = { | ||
| "$inc": {}, | ||
| } | ||
|
|
||
| quota_updates = [] | ||
| for prev_update in org.quotaUpdates or []: | ||
| quota_updates.append(prev_update.dict()) | ||
| quota_updates.append(OrgQuotaUpdate(update=update, modified=dt_now()).dict()) | ||
| 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 | ||
|
|
||
| await self.orgs.find_one_and_update( | ||
| {"_id": org.id}, | ||
| { | ||
| "$set": { | ||
| "quotas": update, | ||
| "quotaUpdates": quota_updates, | ||
| } | ||
| updated_org = await self.orgs.find_one_and_update( | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is needed because we want to store the absolute quota update in quotaUpdates, so it needs to add first, then store the absolute entry in quotaUpdates, right?
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yeah, we hadn't previously had this come up because we were previously just receiving new absolute values. It might be possible to do it all within one update with |
||
| {"_id": org.id}, | ||
| increment_update, | ||
| projection={"quotas": True}, | ||
| return_document=ReturnDocument.AFTER, | ||
| ) | ||
| quotas = OrgQuotasIn(**updated_org["quotas"]) | ||
|
|
||
| update: dict[str, dict[str, dict[str, Any] | int]] = { | ||
| "$push": { | ||
| "quotaUpdates": OrgQuotaUpdate( | ||
| modified=dt_now(), | ||
| 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 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}}, | ||
| ) | ||
| update["$set"]["extraExecSecondsAvailable"] = 0 | ||
| else: | ||
| await self.orgs.find_one_and_update( | ||
| {"_id": org.id}, | ||
| {"$inc": {"extraExecSecondsAvailable": extra_secs_diff}}, | ||
| ) | ||
| update["$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}}, | ||
| ) | ||
| update["$set"]["giftedExecSecondsAvailable"] = 0 | ||
| else: | ||
| await self.orgs.find_one_and_update( | ||
| {"_id": org.id}, | ||
| {"$inc": {"giftedExecSecondsAvailable": gifted_secs_diff}}, | ||
| ) | ||
| update["$inc"]["giftedExecSecondsAvailable"] = gifted_secs_diff | ||
|
|
||
| await self.orgs.find_one_and_update({"_id": org.id}, update) | ||
|
|
||
| async def update_event_webhook_urls( | ||
| self, org: Organization, urls: OrgWebhookUrls | ||
|
|
@@ -1127,7 +1166,7 @@ async def json_items_gen( | |
| yield b"\n" | ||
| doc_index += 1 | ||
|
|
||
| yield f']{"" if skip_closing_comma else ","}\n'.encode("utf-8") | ||
| yield f"]{'' if skip_closing_comma else ','}\n".encode("utf-8") | ||
|
|
||
| async def json_closing_gen() -> AsyncGenerator: | ||
| """Async generator to close JSON document""" | ||
|
|
@@ -1439,10 +1478,12 @@ async def delete_org_and_data( | |
| async def recalculate_storage(self, org: Organization) -> dict[str, bool]: | ||
| """Recalculate org storage use""" | ||
| try: | ||
| total_crawl_size, crawl_size, upload_size = ( | ||
| await self.base_crawl_ops.calculate_org_crawl_file_storage( | ||
| org.id, | ||
| ) | ||
| ( | ||
| total_crawl_size, | ||
| crawl_size, | ||
| upload_size, | ||
| ) = await self.base_crawl_ops.calculate_org_crawl_file_storage( | ||
| org.id, | ||
| ) | ||
| profile_size = await self.profile_ops.calculate_org_profile_file_storage( | ||
| org.id | ||
|
|
@@ -1499,12 +1540,13 @@ async def inc_org_bytes_stored_field(self, oid: UUID, field: str, size: int): | |
| # ============================================================================ | ||
| # pylint: disable=too-many-statements, too-many-arguments | ||
| def init_orgs_api( | ||
| app, | ||
| mdb, | ||
| app: APIRouter, | ||
| mdb: AsyncIOMotorDatabase[Any], | ||
| user_manager: UserManager, | ||
| crawl_manager: CrawlManager, | ||
| invites: InviteOps, | ||
| user_dep: Callable, | ||
| user_dep: Callable[[str], Awaitable[User]], | ||
| superuser_or_shared_secret_dep: Callable[[str], Awaitable[User]], | ||
| ): | ||
| """Init organizations api router for /orgs""" | ||
| # pylint: disable=too-many-locals,invalid-name | ||
|
|
@@ -1523,6 +1565,20 @@ async def org_dep(oid: UUID, user: User = Depends(user_dep)): | |
|
|
||
| return org | ||
|
|
||
| async def org_superuser_or_shared_secret_dep( | ||
| oid: UUID, user: User = Depends(superuser_or_shared_secret_dep) | ||
| ): | ||
| org = await ops.get_org_for_user_by_id(oid, user) | ||
| if not org: | ||
| raise HTTPException(status_code=404, detail="org_not_found") | ||
| if not org.is_viewer(user): | ||
| raise HTTPException( | ||
| status_code=403, | ||
| detail="User does not have permission to view this organization", | ||
| ) | ||
|
|
||
| return org | ||
|
|
||
| async def org_crawl_dep( | ||
| org: Organization = Depends(org_dep), user: User = Depends(user_dep) | ||
| ): | ||
|
|
@@ -1659,7 +1715,18 @@ async def update_quotas( | |
| if not user.is_superuser: | ||
| raise HTTPException(status_code=403, detail="Not Allowed") | ||
|
|
||
| await ops.update_quotas(org, quotas) | ||
| await ops.update_quotas(org, quotas, mode="set", context=quotas.context) | ||
|
|
||
| return {"updated": True} | ||
|
|
||
| @app.post( | ||
| "/orgs/{oid}/quotas/add", tags=["organizations"], response_model=UpdatedResponse | ||
| ) | ||
| async def update_quotas_add( | ||
| quotas: OrgQuotasIn, | ||
| org: Organization = Depends(org_superuser_or_shared_secret_dep), | ||
| ): | ||
| await ops.update_quotas(org, quotas, mode="add", context=quotas.context) | ||
|
|
||
| return {"updated": True} | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is to avoid setting the quota to be <0, right?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Exactly, it clips values below the the field's current value * -1 to that. Do you think we should error instead?