From a3095331eb77410ab62ee3913aeb12b7f9b8c154 Mon Sep 17 00:00:00 2001 From: Eeshan Garg Date: Sat, 20 Nov 2021 19:20:14 +0530 Subject: [PATCH] corporate/models: Modify Customer to accommodate self-hosted customers. This is a part of our efforts to introduce billing for our on-premise customers. --- corporate/lib/stripe.py | 10 +++++- corporate/lib/stripe_event_handler.py | 2 ++ .../0016_customer_add_remote_server_field.py | 32 +++++++++++++++++++ corporate/models.py | 20 +++++++++++- 4 files changed, 62 insertions(+), 2 deletions(-) create mode 100644 corporate/migrations/0016_customer_add_remote_server_field.py diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py index 10732b8dff..72fd1f059c 100644 --- a/corporate/lib/stripe.py +++ b/corporate/lib/stripe.py @@ -454,8 +454,11 @@ def make_end_of_cycle_updates_if_needed( licenses_at_next_renewal=licenses_at_next_renewal, ) + realm = new_plan.customer.realm + assert realm is not None + RealmAuditLog.objects.create( - realm=new_plan.customer.realm, + realm=realm, event_time=event_time, event_type=RealmAuditLog.CUSTOMER_SWITCHED_FROM_MONTHLY_TO_ANNUAL_PLAN, extra_data=orjson.dumps( @@ -629,6 +632,7 @@ def process_initial_upgrade( realm = user.realm customer = update_or_create_stripe_customer(user) assert customer.stripe_customer_id is not None # for mypy + assert customer.realm is not None ensure_realm_does_not_have_active_plan(customer.realm) ( billing_cycle_anchor, @@ -722,12 +726,14 @@ def update_license_ledger_for_manual_plan( licenses_at_next_renewal: Optional[int] = None, ) -> None: if licenses is not None: + assert plan.customer.realm is not None assert get_latest_seat_count(plan.customer.realm) <= licenses assert licenses > plan.licenses() LicenseLedger.objects.create( plan=plan, event_time=event_time, licenses=licenses, licenses_at_next_renewal=licenses ) elif licenses_at_next_renewal is not None: + assert plan.customer.realm is not None assert get_latest_seat_count(plan.customer.realm) <= licenses_at_next_renewal LicenseLedger.objects.create( plan=plan, @@ -779,6 +785,7 @@ def invoice_plan(plan: CustomerPlan, event_time: datetime) -> None: if plan.invoicing_status == CustomerPlan.STARTED: raise NotImplementedError("Plan with invoicing_status==STARTED needs manual resolution.") if not plan.customer.stripe_customer_id: + assert plan.customer.realm is not None raise BillingError( f"Realm {plan.customer.realm.string_id} has a paid plan without a Stripe customer." ) @@ -981,6 +988,7 @@ def do_change_plan_status(plan: CustomerPlan, status: int) -> None: def process_downgrade(plan: CustomerPlan) -> None: from zerver.lib.actions import do_change_plan_type + assert plan.customer.realm is not None do_change_plan_type(plan.customer.realm, Realm.PLAN_TYPE_LIMITED, acting_user=None) plan.status = CustomerPlan.ENDED plan.save(update_fields=["status"]) diff --git a/corporate/lib/stripe_event_handler.py b/corporate/lib/stripe_event_handler.py index d9e5668afc..81bb240d18 100644 --- a/corporate/lib/stripe_event_handler.py +++ b/corporate/lib/stripe_event_handler.py @@ -71,6 +71,7 @@ def handle_checkout_session_completed_event( stripe_setup_intent = stripe.SetupIntent.retrieve(stripe_session.setup_intent) stripe_customer = stripe.Customer.retrieve(stripe_setup_intent.customer) + assert session.customer.realm is not None user = get_user_by_delivery_email(stripe_customer.email, session.customer.realm) payment_method = stripe_setup_intent.payment_method @@ -116,6 +117,7 @@ def handle_payment_intent_succeeded_event( payment_intent.status = PaymentIntent.SUCCEEDED payment_intent.save() metadata = stripe_payment_intent.metadata + assert payment_intent.customer.realm is not None user = get_user_by_delivery_email(metadata["user_email"], payment_intent.customer.realm) description = "" diff --git a/corporate/migrations/0016_customer_add_remote_server_field.py b/corporate/migrations/0016_customer_add_remote_server_field.py new file mode 100644 index 0000000000..ea04cc88c0 --- /dev/null +++ b/corporate/migrations/0016_customer_add_remote_server_field.py @@ -0,0 +1,32 @@ +# Generated by Django 3.2.9 on 2021-11-27 00:08 + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("zilencer", "0018_remoterealmauditlog"), + ("zerver", "0370_realm_enable_spectator_access"), + ("corporate", "0015_event_paymentintent_session"), + ] + + operations = [ + migrations.AddField( + model_name="customer", + name="remote_server", + field=models.OneToOneField( + null=True, + on_delete=django.db.models.deletion.CASCADE, + to="zilencer.remotezulipserver", + ), + ), + migrations.AlterField( + model_name="customer", + name="realm", + field=models.OneToOneField( + null=True, on_delete=django.db.models.deletion.CASCADE, to="zerver.realm" + ), + ), + ] diff --git a/corporate/models.py b/corporate/models.py index fa0c20ece4..f5aacc5a6b 100644 --- a/corporate/models.py +++ b/corporate/models.py @@ -8,6 +8,7 @@ from django.db import models from django.db.models import CASCADE from zerver.models import Realm, UserProfile +from zilencer.models import RemoteZulipServer class Customer(models.Model): @@ -17,7 +18,10 @@ class Customer(models.Model): and the active plan, if any. """ - realm: Realm = models.OneToOneField(Realm, on_delete=CASCADE) + realm: Optional[Realm] = models.OneToOneField(Realm, on_delete=CASCADE, null=True) + remote_server: Optional[RemoteZulipServer] = models.OneToOneField( + RemoteZulipServer, on_delete=CASCADE, null=True + ) stripe_customer_id: Optional[str] = models.CharField(max_length=255, null=True, unique=True) sponsorship_pending: bool = models.BooleanField(default=False) # A percentage, like 85. @@ -30,6 +34,20 @@ class Customer(models.Model): # they purchased. exempt_from_from_license_number_check: bool = models.BooleanField(default=False) + @property + def is_self_hosted(self) -> bool: + is_self_hosted = self.remote_server is not None + if is_self_hosted: + assert self.realm is None + return is_self_hosted + + @property + def is_cloud(self) -> bool: + is_cloud = self.realm is not None + if is_cloud: + assert self.remote_server is None + return is_cloud + def __str__(self) -> str: return f""