From 367d55205267f86e65aa12f2dc5c4f55d94ef0ca Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Sun, 24 Dec 2023 15:56:33 +0100 Subject: [PATCH] billing: Improve make_end_of_cycle_... interactions with audit logs. - Make `self.write_to_audit_log` support a `background_update: bool=False` parameter that can be passed when code that might have an acting user happens to trigger a background update. - Make `make_end_of_cycle_updates_if_needed` pass that parameter for its direct audit log writes. - Audit code that `make_end_of_cycle_updates_if_needed` calls and make sure those write audit logs this way too. - Pass the user in the `billing_page` code that had to avoid it as a workaround: ``` # BUG: This should pass the acting_user; this is just working # around that make_end_of_cycle_updates_if_needed doesn't do audit # logging not using the session user properly. billing_session = RealmBillingSession(user=None, realm=user.realm) ``` --- corporate/lib/stripe.py | 92 ++++++++++++++++++++++++--------- corporate/tests/test_stripe.py | 3 ++ corporate/views/billing_page.py | 5 +- 3 files changed, 73 insertions(+), 27 deletions(-) diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py index 7f2899e4b7..fb486b84a1 100644 --- a/corporate/lib/stripe.py +++ b/corporate/lib/stripe.py @@ -692,6 +692,7 @@ class BillingSession(ABC): event_type: AuditLogEventType, event_time: datetime, *, + background_update: bool = False, extra_data: Optional[Dict[str, Any]] = None, ) -> None: pass @@ -745,11 +746,13 @@ class BillingSession(ABC): pass @abstractmethod - def do_change_plan_type(self, *, tier: Optional[int], is_sponsored: bool = False) -> None: + def do_change_plan_type( + self, *, tier: Optional[int], is_sponsored: bool = False, background_update: bool = False + ) -> None: pass @abstractmethod - def process_downgrade(self, plan: CustomerPlan) -> None: + def process_downgrade(self, plan: CustomerPlan, background_update: bool = False) -> None: pass @abstractmethod @@ -1576,7 +1579,7 @@ class BillingSession(ABC): ) new_plan.status = CustomerPlan.ACTIVE new_plan.save(update_fields=["status"]) - self.do_change_plan_type(tier=new_plan.tier) + self.do_change_plan_type(tier=new_plan.tier, background_update=True) return None, LicenseLedger.objects.create( plan=new_plan, is_renewal=True, @@ -1630,6 +1633,7 @@ class BillingSession(ABC): "monthly_plan_id": plan.id, "annual_plan_id": new_plan.id, }, + background_update=True, ) return new_plan, new_plan_ledger_entry @@ -1678,14 +1682,17 @@ class BillingSession(ABC): "annual_plan_id": plan.id, "monthly_plan_id": new_plan.id, }, + background_update=True, ) return new_plan, new_plan_ledger_entry if plan.status == CustomerPlan.DOWNGRADE_AT_END_OF_FREE_TRIAL: - self.downgrade_now_without_creating_additional_invoices(plan) + self.downgrade_now_without_creating_additional_invoices( + plan, background_update=True + ) if plan.status == CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE: - self.process_downgrade(plan) + self.process_downgrade(plan, background_update=True) return None, None return None, last_ledger_entry @@ -2047,6 +2054,7 @@ class BillingSession(ABC): def downgrade_now_without_creating_additional_invoices( self, plan: Optional[CustomerPlan] = None, + background_update: bool = False, ) -> None: if plan is None: customer = self.get_customer() @@ -2056,7 +2064,7 @@ class BillingSession(ABC): if plan is None: return # nocoverage - self.process_downgrade(plan) + self.process_downgrade(plan, background_update=background_update) plan.invoiced_through = LicenseLedger.objects.filter(plan=plan).order_by("id").last() plan.next_invoice_date = next_invoice_date(plan) plan.save(update_fields=["invoiced_through", "next_invoice_date"]) @@ -2912,6 +2920,7 @@ class RealmBillingSession(BillingSession): event_type: AuditLogEventType, event_time: datetime, *, + background_update: bool = False, extra_data: Optional[Dict[str, Any]] = None, ) -> None: audit_log_event = self.get_audit_log_event(event_type) @@ -2924,7 +2933,7 @@ class RealmBillingSession(BillingSession): if extra_data: audit_log_data["extra_data"] = extra_data - if self.user is not None: + if self.user is not None and not background_update: audit_log_data["acting_user"] = self.user RealmAuditLog.objects.create(**audit_log_data) @@ -2980,7 +2989,9 @@ class RealmBillingSession(BillingSession): return customer @override - def do_change_plan_type(self, *, tier: Optional[int], is_sponsored: bool = False) -> None: + def do_change_plan_type( + self, *, tier: Optional[int], is_sponsored: bool = False, background_update: bool = False + ) -> None: from zerver.actions.realm_settings import do_change_realm_plan_type # This function needs to translate between the different @@ -2996,14 +3007,25 @@ class RealmBillingSession(BillingSession): plan_type = Realm.PLAN_TYPE_PLUS else: raise AssertionError("Unexpected tier") - do_change_realm_plan_type(self.realm, plan_type, acting_user=self.user) + + acting_user = None + if not background_update: + acting_user = self.user + + do_change_realm_plan_type(self.realm, plan_type, acting_user=acting_user) @override - def process_downgrade(self, plan: CustomerPlan) -> None: + def process_downgrade(self, plan: CustomerPlan, background_update: bool = False) -> None: from zerver.actions.realm_settings import do_change_realm_plan_type + acting_user = None + if not background_update: + acting_user = self.user + assert plan.customer.realm is not None - do_change_realm_plan_type(plan.customer.realm, Realm.PLAN_TYPE_LIMITED, acting_user=None) + do_change_realm_plan_type( + plan.customer.realm, Realm.PLAN_TYPE_LIMITED, acting_user=acting_user + ) plan.status = CustomerPlan.ENDED plan.save(update_fields=["status"]) @@ -3260,6 +3282,7 @@ class RemoteRealmBillingSession(BillingSession): event_type: AuditLogEventType, event_time: datetime, *, + background_update: bool = False, extra_data: Optional[Dict[str, Any]] = None, ) -> None: # These audit logs don't use all the fields of `RemoteRealmAuditLog`: @@ -3273,12 +3296,18 @@ class RemoteRealmBillingSession(BillingSession): "remote_realm": self.remote_realm, "event_type": audit_log_event, "event_time": event_time, - # At most one of these should be set, but we may - # not want an assert for that yet: - "acting_support_user": self.support_staff, - "acting_remote_user": self.remote_billing_user, } + if not background_update: + log_data.update( + { + # At most one of these should be set, but we may + # not want an assert for that yet: + "acting_support_user": self.support_staff, + "acting_remote_user": self.remote_billing_user, + } + ) + if extra_data: log_data["extra_data"] = extra_data @@ -3336,7 +3365,7 @@ class RemoteRealmBillingSession(BillingSession): @override @transaction.atomic def do_change_plan_type( - self, *, tier: Optional[int], is_sponsored: bool = False + self, *, tier: Optional[int], is_sponsored: bool = False, background_update: bool = False ) -> None: # nocoverage if is_sponsored: plan_type = RemoteRealm.PLAN_TYPE_COMMUNITY @@ -3357,6 +3386,7 @@ class RemoteRealmBillingSession(BillingSession): event_type=AuditLogEventType.BILLING_ENTITY_PLAN_TYPE_CHANGED, event_time=timezone_now(), extra_data={"old_value": old_plan_type, "new_value": plan_type}, + background_update=background_update, ) @override @@ -3431,7 +3461,9 @@ class RemoteRealmBillingSession(BillingSession): ) @override - def process_downgrade(self, plan: CustomerPlan) -> None: # nocoverage + def process_downgrade( + self, plan: CustomerPlan, background_update: bool = False + ) -> None: # nocoverage with transaction.atomic(): old_plan_type = self.remote_realm.plan_type new_plan_type = RemoteRealm.PLAN_TYPE_SELF_MANAGED @@ -3441,6 +3473,7 @@ class RemoteRealmBillingSession(BillingSession): event_type=AuditLogEventType.BILLING_ENTITY_PLAN_TYPE_CHANGED, event_time=timezone_now(), extra_data={"old_value": old_plan_type, "new_value": new_plan_type}, + background_update=background_update, ) plan.status = CustomerPlan.ENDED @@ -3659,6 +3692,7 @@ class RemoteServerBillingSession(BillingSession): event_type: AuditLogEventType, event_time: datetime, *, + background_update: bool = False, extra_data: Optional[Dict[str, Any]] = None, ) -> None: audit_log_event = self.get_audit_log_event(event_type) @@ -3666,12 +3700,18 @@ class RemoteServerBillingSession(BillingSession): "server": self.remote_server, "event_type": audit_log_event, "event_time": event_time, - # At most one of these should be set, but we may - # not want an assert for that yet: - "acting_support_user": self.support_staff, - "acting_remote_user": self.remote_billing_user, } + if not background_update: + log_data.update( + { + # At most one of these should be set, but we may + # not want an assert for that yet: + "acting_support_user": self.support_staff, + "acting_remote_user": self.remote_billing_user, + } + ) + if extra_data: log_data["extra_data"] = extra_data @@ -3728,7 +3768,9 @@ class RemoteServerBillingSession(BillingSession): @override @transaction.atomic - def do_change_plan_type(self, *, tier: Optional[int], is_sponsored: bool = False) -> None: + def do_change_plan_type( + self, *, tier: Optional[int], is_sponsored: bool = False, background_update: bool = False + ) -> None: # This function needs to translate between the different # formats of CustomerPlan.tier and RealmZulipServer.plan_type. if is_sponsored: @@ -3750,6 +3792,7 @@ class RemoteServerBillingSession(BillingSession): event_type=AuditLogEventType.BILLING_ENTITY_PLAN_TYPE_CHANGED, event_time=timezone_now(), extra_data={"old_value": old_plan_type, "new_value": plan_type}, + background_update=background_update, ) @override @@ -3797,7 +3840,9 @@ class RemoteServerBillingSession(BillingSession): return f"Sponsorship approved for {self.billing_entity_display_name}" @override - def process_downgrade(self, plan: CustomerPlan) -> None: # nocoverage + def process_downgrade( + self, plan: CustomerPlan, background_update: bool = False + ) -> None: # nocoverage with transaction.atomic(): old_plan_type = self.remote_server.plan_type new_plan_type = RemoteZulipServer.PLAN_TYPE_SELF_MANAGED @@ -3807,6 +3852,7 @@ class RemoteServerBillingSession(BillingSession): event_type=AuditLogEventType.BILLING_ENTITY_PLAN_TYPE_CHANGED, event_time=timezone_now(), extra_data={"old_value": old_plan_type, "new_value": new_plan_type}, + background_update=background_update, ) plan.status = CustomerPlan.ENDED diff --git a/corporate/tests/test_stripe.py b/corporate/tests/test_stripe.py index a12469c8e6..e7543e2260 100644 --- a/corporate/tests/test_stripe.py +++ b/corporate/tests/test_stripe.py @@ -2354,6 +2354,9 @@ class StripeTest(StripeTestCase): .first(), (20, 20), ) + realm_audit_log = RealmAuditLog.objects.latest("id") + self.assertEqual(realm_audit_log.event_type, RealmAuditLog.REALM_PLAN_TYPE_CHANGED) + self.assertEqual(realm_audit_log.acting_user, None) # Verify that we don't write LicenseLedger rows once we've downgraded with patch("corporate.lib.stripe.get_latest_seat_count", return_value=40): diff --git a/corporate/views/billing_page.py b/corporate/views/billing_page.py index 0a45f49d0b..b58a1d397a 100644 --- a/corporate/views/billing_page.py +++ b/corporate/views/billing_page.py @@ -52,10 +52,7 @@ def billing_page( user = request.user assert user.is_authenticated - # BUG: This should pass the acting_user; this is just working - # around that make_end_of_cycle_updates_if_needed doesn't do audit - # logging not using the session user properly. - billing_session = RealmBillingSession(user=None, realm=user.realm) + billing_session = RealmBillingSession(user=user, realm=user.realm) context: Dict[str, Any] = { "admin_access": user.has_billing_access,