ref: Plan query optimizations (#1131)

Author: ajay-sentry

api/internal/owner/serializers.py

@@ -339,11 +339,13 @@ def update(self, instance: Owner, validated_data: Dict[str, Any]) -> object:
                 instance, desired_plan
             )
 
-            sentry_plans = Plan.objects.filter(
-                tier__tier_name=TierName.SENTRY.value, is_active=True
-            ).values_list("name", flat=True)
+            plan = (
+                Plan.objects.select_related("tier")
+                .filter(name=desired_plan["value"])
+                .first()
+            )
 
-            if desired_plan["value"] in sentry_plans:
+            if plan and plan.tier.tier_name == TierName.SENTRY.value:
                 current_owner = self.context["view"].request.current_owner
                 send_sentry_webhook(current_owner, instance)
 

services/billing.py

@@ -316,17 +316,33 @@ def get_schedule(self, owner: Owner):
         return stripe.SubscriptionSchedule.retrieve(subscription_schedule_id)
 
     @_log_stripe_error
-    def modify_subscription(self, owner, desired_plan):
+    def modify_subscription(self, owner: Owner, desired_plan: dict):
+        desired_plan_info = Plan.objects.filter(name=desired_plan["value"]).first()
+        if not desired_plan_info:
+            log.error(
+                f"Plan {desired_plan['value']} not found",
+                extra=dict(owner_id=owner.ownerid),
+            )
+            return
+
         subscription = stripe.Subscription.retrieve(owner.stripe_subscription_id)
-        proration_behavior = self._get_proration_params(owner, desired_plan)
+        proration_behavior = self._get_proration_params(
+            owner,
+            desired_plan_info=desired_plan_info,
+            desired_quantity=desired_plan["quantity"],
+        )
         subscription_schedule_id = subscription.schedule
 
         # proration_behavior indicates whether we immediately invoice a user or not. We only immediately
         # invoice a user if the user increases the number of seats or if the plan changes from monthly to yearly.
         # An increase in seats and/or plan implies the user is upgrading, hence 'is_upgrading' is a consequence
         # of proration_behavior providing an invoice, in this case, != "none"
         # TODO: change this to "self._is_upgrading_seats(owner, desired_plan) or self._is_extending_term(owner, desired_plan)"
-        is_upgrading = True if proration_behavior != "none" else False
+        is_upgrading = (
+            True
+            if proration_behavior != "none" and desired_plan_info.stripe_id
+            else False
+        )
 
         # Divide logic bw immediate updates and scheduled updates
         # Immediate updates: when user upgrades seats or plan
@@ -335,15 +351,6 @@ def modify_subscription(self, owner, desired_plan):
         # Scheduled updates: when the user decreases seats or plan
         #   If the user is not in a schedule, create a schedule
         #   If the user is in a schedule, update the existing schedule
-
-        plan = Plan.objects.filter(name=desired_plan["value"]).first()
-        if not plan or not plan.stripe_id:
-            log.error(
-                f"Plan {desired_plan['value']} not found",
-                extra=dict(owner_id=owner.ownerid),
-            )
-            return
-
         if is_upgrading:
             if subscription_schedule_id:
                 log.info(
@@ -360,7 +367,7 @@ def modify_subscription(self, owner, desired_plan):
                 items=[
                     {
                         "id": subscription["items"]["data"][0]["id"],
-                        "plan": plan.stripe_id,
+                        "plan": desired_plan_info.stripe_id,
                         "quantity": desired_plan["quantity"],
                     }
                 ],
@@ -404,7 +411,11 @@ def modify_subscription(self, owner, desired_plan):
             )
 
     def _modify_subscription_schedule(
-        self, owner: Owner, subscription, subscription_schedule_id, desired_plan
+        self,
+        owner: Owner,
+        subscription: stripe.Subscription,
+        subscription_schedule_id: str,
+        desired_plan: dict,
     ):
         current_subscription_start_date = subscription["current_period_start"]
         current_subscription_end_date = subscription["current_period_end"]
@@ -453,20 +464,18 @@ def _modify_subscription_schedule(
             metadata=self._get_checkout_session_and_subscription_metadata(owner),
         )
 
-    def _is_upgrading_seats(self, owner: Owner, desired_plan: dict) -> bool:
+    def _is_upgrading_seats(self, owner: Owner, desired_quantity: int) -> bool:
         """
         Returns `True` if purchasing more seats.
         """
-        return bool(
-            owner.plan_user_count and owner.plan_user_count < desired_plan["quantity"]
-        )
+        return bool(owner.plan_user_count and owner.plan_user_count < desired_quantity)
 
-    def _is_extending_term(self, owner: Owner, desired_plan: dict) -> bool:
+    def _is_extending_term(
+        self, current_plan_info: Plan, desired_plan_info: Plan
+    ) -> bool:
         """
         Returns `True` if switching from monthly to yearly plan.
         """
-        current_plan_info = Plan.objects.get(name=owner.plan)
-        desired_plan_info = Plan.objects.get(name=desired_plan["value"])
 
         return bool(
             current_plan_info
@@ -475,23 +484,24 @@ def _is_extending_term(self, owner: Owner, desired_plan: dict) -> bool:
             and desired_plan_info.billing_rate == PlanBillingRate.YEARLY.value
         )
 
-    def _is_similar_plan(self, owner: Owner, desired_plan: dict) -> bool:
+    def _is_similar_plan(
+        self,
+        owner: Owner,
+        current_plan_info: Plan,
+        desired_plan_info: Plan,
+        desired_quantity: int,
+    ) -> bool:
         """
         Returns `True` if switching to a plan with similar term and seats.
         """
-        current_plan_info = Plan.objects.select_related("tier").get(name=owner.plan)
-        desired_plan_info = Plan.objects.select_related("tier").get(
-            name=desired_plan["value"]
-        )
-
         is_same_term = (
             current_plan_info
             and desired_plan_info
             and current_plan_info.billing_rate == desired_plan_info.billing_rate
         )
 
         is_same_seats = (
-            owner.plan_user_count and owner.plan_user_count == desired_plan["quantity"]
+            owner.plan_user_count and owner.plan_user_count == desired_quantity
         )
         # If from PRO to TEAM, then not a similar plan
         if (
@@ -508,25 +518,35 @@ def _is_similar_plan(self, owner: Owner, desired_plan: dict) -> bool:
 
         return bool(is_same_term and is_same_seats)
 
-    def _get_proration_params(self, owner: Owner, desired_plan: dict) -> str:
+    def _get_proration_params(
+        self, owner: Owner, desired_plan_info: Plan, desired_quantity: int
+    ) -> str:
+        current_plan_info = Plan.objects.select_related("tier").get(name=owner.plan)
         if (
-            self._is_upgrading_seats(owner, desired_plan)
-            or self._is_extending_term(owner, desired_plan)
-            or self._is_similar_plan(owner, desired_plan)
+            self._is_upgrading_seats(owner=owner, desired_quantity=desired_quantity)
+            or self._is_extending_term(
+                current_plan_info=current_plan_info, desired_plan_info=desired_plan_info
+            )
+            or self._is_similar_plan(
+                owner=owner,
+                current_plan_info=current_plan_info,
+                desired_plan_info=desired_plan_info,
+                desired_quantity=desired_quantity,
+            )
         ):
             return "always_invoice"
         else:
             return "none"
 
-    def _get_success_and_cancel_url(self, owner):
+    def _get_success_and_cancel_url(self, owner: Owner):
         short_services = {"github": "gh", "bitbucket": "bb", "gitlab": "gl"}
         base_path = f"/plan/{short_services[owner.service]}/{owner.username}"
         success_url = f"{settings.CODECOV_DASHBOARD_URL}{base_path}?success"
         cancel_url = f"{settings.CODECOV_DASHBOARD_URL}{base_path}?cancel"
         return success_url, cancel_url
 
     @_log_stripe_error
-    def create_checkout_session(self, owner: Owner, desired_plan):
+    def create_checkout_session(self, owner: Owner, desired_plan: dict):
         success_url, cancel_url = self._get_success_and_cancel_url(owner)
         log.info(
             "Creating Stripe Checkout Session for owner",