Merge pull request #3055 from webrecorder/additional-minutes-addon-purchase--add-sub-add-minutes-events

Author: emma-sg

.github/workflows/microk8s-ci.yaml

@@ -17,6 +17,12 @@ jobs:
   btrix-microk8s-test:
     runs-on: ubuntu-latest
     steps:
+      - name: Initial Disk Cleanup
+        uses: mathio/gha-cleanup@v1
+        with:
+          remove-browsers: true
+          verbose: true
+
       - uses: balchua/[email protected]
         with:
           channel: "1.25/stable"

backend/btrixcloud/main.py

@@ -181,7 +181,6 @@ def main() -> None:
         crawl_manager,
         invites,
         current_active_user,
-        shared_secret_or_superuser,
     )
 
     init_subs_api(app, mdb, org_ops, user_manager, shared_secret_or_superuser)

backend/btrixcloud/models.py

@@ -1829,6 +1829,8 @@ class S3Storage(BaseModel):
 REASON_PAUSED = "subscriptionPaused"
 REASON_CANCELED = "subscriptionCanceled"
 
+SubscriptionEventType = Literal["create", "import", "update", "cancel", "add-minutes"]
+
 
 # ============================================================================
 class OrgQuotas(BaseModel):
@@ -1857,8 +1859,6 @@ class OrgQuotasIn(BaseModel):
     extraExecMinutes: Optional[int] = None
     giftedExecMinutes: Optional[int] = None
 
-    context: str | None = None
-
 
 # ============================================================================
 class Plan(BaseModel):
@@ -1883,6 +1883,7 @@ class SubscriptionEventOut(BaseModel):
 
     oid: UUID
     timestamp: datetime
+    type: SubscriptionEventType
 
 
 # ============================================================================
@@ -1948,18 +1949,55 @@ class SubscriptionCancel(BaseModel):
 
 
 # ============================================================================
-class SubscriptionTrialEndReminder(BaseModel):
-    """Email reminder that subscription will end soon"""
+class SubscriptionCancelOut(SubscriptionCancel, SubscriptionEventOut):
+    """Output model for subscription cancellation event"""
 
-    subId: str
-    behavior_on_trial_end: Literal["cancel", "continue", "read-only"]
+    type: Literal["cancel"] = "cancel"
 
 
 # ============================================================================
-class SubscriptionCancelOut(SubscriptionCancel, SubscriptionEventOut):
-    """Output model for subscription cancellation event"""
+class SubscriptionAddMinutes(BaseModel):
+    """Represents a purchase of additional minutes"""
+
+    oid: UUID
+    minutes: int
+    total_price: float
+    currency: str
+
+    context: str
 
-    type: Literal["cancel"] = "cancel"
+
+# ============================================================================
+class SubscriptionAddMinutesOut(SubscriptionAddMinutes, SubscriptionEventOut):
+    """SubscriptionAddMinutes output model"""
+
+    type: Literal["add-minutes"] = "add-minutes"
+
+
+# ============================================================================
+SubscriptionEventAny = Union[
+    SubscriptionCreate,
+    SubscriptionUpdate,
+    SubscriptionCancel,
+    SubscriptionImport,
+    SubscriptionAddMinutes,
+]
+
+SubscriptionEventAnyOut = Union[
+    SubscriptionCreateOut,
+    SubscriptionUpdateOut,
+    SubscriptionCancelOut,
+    SubscriptionImportOut,
+    SubscriptionAddMinutesOut,
+]
+
+
+# ============================================================================
+class SubscriptionTrialEndReminder(BaseModel):
+    """Email reminder that subscription will end soon"""
+
+    subId: str
+    behavior_on_trial_end: Literal["cancel", "continue", "read-only"]
 
 
 # ============================================================================
@@ -3126,14 +3164,7 @@ class PaginatedProfileResponse(PaginatedResponse):
 class PaginatedSubscriptionEventResponse(PaginatedResponse):
     """Response model for paginated subscription events"""
 
-    items: List[
-        Union[
-            SubscriptionCreateOut,
-            SubscriptionUpdateOut,
-            SubscriptionCancelOut,
-            SubscriptionImportOut,
-        ]
-    ]
+    items: List[SubscriptionEventAnyOut]
 
 
 # ============================================================================

backend/btrixcloud/orgs.py

@@ -625,8 +625,6 @@ async def update_quotas(
     ) -> None:
         """update organization quotas"""
 
-        quotas.context = None
-
         previous_extra_mins = (
             org.quotas.extraExecMinutes
             if (org.quotas and org.quotas.extraExecMinutes)
@@ -1546,7 +1544,6 @@ def init_orgs_api(
     crawl_manager: CrawlManager,
     invites: InviteOps,
     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
@@ -1701,20 +1698,7 @@ async def update_quotas(
         if not user.is_superuser:
             raise HTTPException(status_code=403, detail="Not Allowed")
 
-        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(
-        oid: UUID,
-        quotas: OrgQuotasIn,
-        _user: User = Depends(superuser_or_shared_secret_dep),
-    ):
-        org = await ops.get_org_by_id(oid)
-        await ops.update_quotas(org, quotas, mode="add", context=quotas.context)
+        await ops.update_quotas(org, quotas, mode="set")
 
         return {"updated": True}
 

backend/btrixcloud/subs.py

@@ -2,13 +2,13 @@
 Subscription API handling
 """
 
-from typing import Awaitable, Callable, Union, Any, Optional, Tuple, List
+from typing import Awaitable, Callable, Any, Optional, Tuple, List, Annotated
 import os
 import asyncio
 from uuid import UUID
 from datetime import datetime
 
-from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request, Query
 import aiohttp
 from motor.motor_asyncio import AsyncIOMotorDatabase
 
@@ -23,16 +23,22 @@
     SubscriptionImport,
     SubscriptionUpdate,
     SubscriptionCancel,
+    SubscriptionAddMinutes,
+    SubscriptionEventAny,
     SubscriptionCreateOut,
     SubscriptionImportOut,
     SubscriptionUpdateOut,
     SubscriptionCancelOut,
+    SubscriptionAddMinutesOut,
+    SubscriptionEventAnyOut,
+    SubscriptionEventType,
     Subscription,
     SubscriptionPortalUrlRequest,
     SubscriptionPortalUrlResponse,
     SubscriptionCanceledResponse,
     SubscriptionTrialEndReminder,
     Organization,
+    OrgQuotasIn,
     InviteToOrgRequest,
     InviteAddedResponse,
     User,
@@ -228,15 +234,22 @@ async def send_trial_end_reminder(
 
         return SuccessResponse(success=True)
 
+    async def add_sub_minutes(self, add_min: SubscriptionAddMinutes):
+        """add extra minutes for subscription"""
+        org = await self.org_ops.get_org_by_id(add_min.oid)
+        quotas = OrgQuotasIn(extraExecMinutes=add_min.minutes)
+        await self.org_ops.update_quotas(
+            org, quotas, mode="add", context=add_min.context
+        )
+
+        await self.add_sub_event("add-minutes", add_min, add_min.oid)
+
+        return {"updated": True}
+
     async def add_sub_event(
         self,
-        type_: str,
-        event: Union[
-            SubscriptionCreate,
-            SubscriptionImport,
-            SubscriptionUpdate,
-            SubscriptionCancel,
-        ],
+        type_: SubscriptionEventType,
+        event: SubscriptionEventAny,
         oid: UUID,
     ) -> None:
         """add a subscription event to the db"""
@@ -246,20 +259,22 @@ async def add_sub_event(
         data["oid"] = oid
         await self.subs.insert_one(data)
 
-    def _get_sub_by_type_from_data(self, data: dict[str, object]) -> Union[
-        SubscriptionCreateOut,
-        SubscriptionImportOut,
-        SubscriptionUpdateOut,
-        SubscriptionCancelOut,
-    ]:
+    def _get_sub_by_type_from_data(
+        self, data: dict[str, object]
+    ) -> SubscriptionEventAnyOut:
         """convert dict to propert background job type"""
         if data["type"] == "create":
             return SubscriptionCreateOut(**data)
         if data["type"] == "import":
             return SubscriptionImportOut(**data)
         if data["type"] == "update":
             return SubscriptionUpdateOut(**data)
-        return SubscriptionCancelOut(**data)
+        if data["type"] == "cancel":
+            return SubscriptionCancelOut(**data)
+        if data["type"] == "add-minutes":
+            return SubscriptionAddMinutesOut(**data)
+
+        raise HTTPException(status_code=500, detail="unknown sub event")
 
     # pylint: disable=too-many-arguments
     async def list_sub_events(
@@ -268,19 +283,13 @@ async def list_sub_events(
         sub_id: Optional[str] = None,
         oid: Optional[UUID] = None,
         plan_id: Optional[str] = None,
+        type_: Optional[SubscriptionEventType] = None,
         page_size: int = DEFAULT_PAGE_SIZE,
         page: int = 1,
         sort_by: Optional[str] = None,
         sort_direction: Optional[int] = -1,
     ) -> Tuple[
-        List[
-            Union[
-                SubscriptionCreateOut,
-                SubscriptionImportOut,
-                SubscriptionUpdateOut,
-                SubscriptionCancelOut,
-            ]
-        ],
+        List[SubscriptionEventAnyOut],
         int,
     ]:
         """list subscription events"""
@@ -298,6 +307,8 @@ async def list_sub_events(
             query["planId"] = plan_id
         if oid:
             query["oid"] = oid
+        if type_:
+            query["type"] = type_
 
         aggregate = [{"$match": query}]
 
@@ -515,6 +526,15 @@ async def send_trial_end_reminder(
     ):
         return await ops.send_trial_end_reminder(reminder)
 
+    @app.post(
+        "/subscriptions/add-minutes",
+        tags=["subscriptions"],
+        dependencies=[Depends(superuser_or_shared_secret_dep)],
+        response_model=UpdatedResponse,
+    )
+    async def add_sub_minutes(add_min: SubscriptionAddMinutes):
+        return await ops.add_sub_minutes(add_min)
+
     assert org_ops.router
 
     @app.get(
@@ -540,6 +560,7 @@ async def get_sub_events(
         subId: Optional[str] = None,
         oid: Optional[UUID] = None,
         planId: Optional[str] = None,
+        type_: Annotated[Optional[SubscriptionEventType], Query(alias="type")] = None,
         pageSize: int = DEFAULT_PAGE_SIZE,
         page: int = 1,
         sortBy: Optional[str] = "timestamp",
@@ -551,6 +572,7 @@ async def get_sub_events(
             oid=oid,
             plan_id=planId,
             page_size=pageSize,
+            type_=type_,
             page=page,
             sort_by=sortBy,
             sort_direction=sortDirection,