Mirror
/
zulip
mirror of https://github.com/zulip/zulip.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
zulip/zilencer/management/commands/populate_billing_realms.py

524 lines
20 KiB
Python
Raw Blame History

import contextlib
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, Optional
import stripe
from django.conf import settings
from django.core.management.base import BaseCommand, CommandParser
from django.utils.timezone import now as timezone_now
from typing_extensions import override
from corporate.lib.stripe import (
RealmBillingSession,
RemoteRealmBillingSession,
RemoteServerBillingSession,
UpgradeRequest,
add_months,
sign_string,
)
from corporate.models import Customer, CustomerPlan, LicenseLedger
from scripts.lib.zulip_tools import TIMESTAMP_FORMAT
from zerver.actions.create_realm import do_create_realm
from zerver.actions.create_user import do_create_user
from zerver.actions.streams import bulk_add_subscriptions
from zerver.apps import flush_cache
from zerver.lib.remote_server import get_realms_info_for_push_bouncer
from zerver.lib.streams import create_stream_if_needed
from zerver.models import Realm, UserProfile, get_realm
from zilencer.models import (
RemoteRealm,
RemoteRealmBillingUser,
RemoteServerBillingUser,
RemoteZulipServer,
)
from zilencer.views import update_remote_realm_data_for_server
from zproject.config import get_secret
current_time = timezone_now().strftime(TIMESTAMP_FORMAT)
@dataclass
class CustomerProfile:
unique_id: str
billing_schedule: int = CustomerPlan.BILLING_SCHEDULE_ANNUAL
tier: Optional[int] = None
new_plan_tier: Optional[int] = None
automanage_licenses: bool = False
status: int = CustomerPlan.ACTIVE
sponsorship_pending: bool = False
is_sponsored: bool = False
card: str = ""
charge_automatically: bool = True
is_remote_realm: bool = False
is_remote_server: bool = False
renewal_date: str = current_time
# Use (timezone_now() + timedelta(minutes=1)).strftime(TIMESTAMP_FORMAT) as `end_date` for testing.
# `invoice_plan` is not implemented yet for remote servers and realms so no payment is generated in stripe.
end_date: str = "2030-10-10-01-10-10"
remote_server_plan_start_date: str = "billing_cycle_end_date"
class Command(BaseCommand):
help = "Populate database with different types of realms that can exist."
@override
def add_arguments(self, parser: CommandParser) -> None:
parser.add_argument(
"--only-remote-server",
action="store_true",
help="Whether to only run for remote servers",
)
parser.add_argument(
"--only-remote-realm",
action="store_true",
help="Whether to only run for remote realms",
)
@override
def handle(self, *args: Any, **options: Any) -> None:
# Create a realm for each plan type
customer_profiles = [
# NOTE: The unique_id has to be less than 40 characters.
CustomerProfile(unique_id="sponsorship-pending", sponsorship_pending=True),
CustomerProfile(
unique_id="annual-free",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL,
),
CustomerProfile(
unique_id="annual-standard",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
),
CustomerProfile(
unique_id="annual-plus",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL,
tier=CustomerPlan.TIER_CLOUD_PLUS,
),
CustomerProfile(
unique_id="monthly-free",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
),
CustomerProfile(
unique_id="monthly-standard",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
),
CustomerProfile(
unique_id="monthly-plus",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_PLUS,
),
CustomerProfile(
unique_id="downgrade-end-of-cycle",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
status=CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE,
),
CustomerProfile(
unique_id="standard-automanage-licenses",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
automanage_licenses=True,
),
CustomerProfile(
unique_id="standard-automatic-card",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
card="pm_card_visa",
),
CustomerProfile(
unique_id="standard-invoice-payment",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
charge_automatically=False,
),
CustomerProfile(
unique_id="standard-switch-to-annual-eoc",
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
status=CustomerPlan.SWITCH_TO_ANNUAL_AT_END_OF_CYCLE,
),
CustomerProfile(
unique_id="sponsored",
is_sponsored=True,
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
# Customer plan might not exist for sponsored realms.
tier=CustomerPlan.TIER_CLOUD_STANDARD,
),
CustomerProfile(
unique_id="free-trial",
tier=CustomerPlan.TIER_CLOUD_STANDARD,
status=CustomerPlan.FREE_TRIAL,
),
CustomerProfile(
unique_id="legacy-server",
tier=CustomerPlan.TIER_SELF_HOSTED_LEGACY,
is_remote_server=True,
),
CustomerProfile(
unique_id="legacy-server-upgrade-scheduled",
tier=CustomerPlan.TIER_SELF_HOSTED_LEGACY,
status=CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END,
new_plan_tier=CustomerPlan.TIER_SELF_HOSTED_PLUS,
is_remote_server=True,
),
CustomerProfile(
unique_id="business-server",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
is_remote_server=True,
),
CustomerProfile(
unique_id="business-server-free-trial",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
status=CustomerPlan.FREE_TRIAL,
is_remote_server=True,
),
CustomerProfile(
unique_id="business-server-sponsorship-pending",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
sponsorship_pending=True,
is_remote_server=True,
),
CustomerProfile(
unique_id="self-hosted-server-sponsorship-pending",
tier=CustomerPlan.TIER_SELF_HOSTED_BASE,
sponsorship_pending=True,
is_remote_server=True,
),
CustomerProfile(
unique_id="server-sponsored",
tier=CustomerPlan.TIER_SELF_HOSTED_COMMUNITY,
is_sponsored=True,
is_remote_server=True,
),
CustomerProfile(
unique_id="free-tier-remote-realm",
is_remote_realm=True,
),
CustomerProfile(
unique_id="business-remote-realm",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
is_remote_realm=True,
),
CustomerProfile(
unique_id="business-remote-free-trial",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
status=CustomerPlan.FREE_TRIAL,
is_remote_realm=True,
),
CustomerProfile(
unique_id="business-remote-sponsorship-pending",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
sponsorship_pending=True,
is_remote_realm=True,
),
CustomerProfile(
unique_id="remote-sponsorship-pending",
sponsorship_pending=True,
is_remote_realm=True,
),
CustomerProfile(
unique_id="remote-realm-sponsored",
tier=CustomerPlan.TIER_SELF_HOSTED_COMMUNITY,
is_sponsored=True,
is_remote_realm=True,
),
]
servers = []
remote_realms = []
for customer_profile in customer_profiles:
if customer_profile.is_remote_server and not options.get("only_remote_realm"):
server_conf = populate_remote_server(customer_profile)
servers.append(server_conf)
elif customer_profile.is_remote_realm and not options.get("only_remote_server"):
remote_realm_conf = populate_remote_realms(customer_profile)
remote_realms.append(remote_realm_conf)
elif not options.get("only_remote_server") and not options.get("only_remote_realm"):
populate_realm(customer_profile)
print("-" * 40)
for server in servers:
for key, value in server.items():
print(f"{key}: {value}")
print("-" * 40)
for remote_realm_conf in remote_realms:
for key, value in remote_realm_conf.items():
print(f"{key}: {value}")
print("-" * 40)
def add_card_to_customer(customer: Customer) -> None:
assert customer.stripe_customer_id is not None
# Set the Stripe API key
stripe.api_key = get_secret("stripe_secret_key")
# Create a card payment method and attach it to the customer
payment_method = stripe.PaymentMethod.create(
type="card",
card={"token": "tok_visa"},
)
# Attach the payment method to the customer
stripe.PaymentMethod.attach(payment_method.id, customer=customer.stripe_customer_id)
# Set the default payment method for the customer
stripe.Customer.modify(
customer.stripe_customer_id,
invoice_settings={"default_payment_method": payment_method.id},
)
def create_plan_for_customer(customer: Customer, customer_profile: CustomerProfile) -> None:
assert customer_profile.tier is not None
if customer_profile.status == CustomerPlan.FREE_TRIAL:
# 2 months free trial.
next_invoice_date = add_months(timezone_now(), 2)
else:
months = 12
if customer_profile.billing_schedule == CustomerPlan.BILLING_SCHEDULE_MONTHLY:
months = 1
next_invoice_date = add_months(timezone_now(), months)
customer_plan = CustomerPlan.objects.create(
customer=customer,
billing_cycle_anchor=timezone_now(),
billing_schedule=customer_profile.billing_schedule,
tier=customer_profile.tier,
price_per_license=1200,
automanage_licenses=customer_profile.automanage_licenses,
status=customer_profile.status,
charge_automatically=customer_profile.charge_automatically,
next_invoice_date=next_invoice_date,
)
LicenseLedger.objects.create(
licenses=10,
licenses_at_next_renewal=10,
event_time=timezone_now(),
is_renewal=True,
plan=customer_plan,
)
def populate_realm(customer_profile: CustomerProfile) -> Optional[Realm]:
unique_id = customer_profile.unique_id
if customer_profile.is_remote_realm:
plan_type = Realm.PLAN_TYPE_SELF_HOSTED
elif customer_profile.is_sponsored:
plan_type = Realm.PLAN_TYPE_STANDARD_FREE
elif customer_profile.tier is None:
plan_type = Realm.PLAN_TYPE_LIMITED
elif customer_profile.tier == CustomerPlan.TIER_CLOUD_STANDARD:
plan_type = Realm.PLAN_TYPE_STANDARD
elif customer_profile.tier == CustomerPlan.TIER_CLOUD_PLUS:
plan_type = Realm.PLAN_TYPE_PLUS
else:
raise AssertionError("Unexpected tier!")
plan_name = Realm.ALL_PLAN_TYPES[plan_type]
# Delete existing realm with this name
with contextlib.suppress(Realm.DoesNotExist):
get_realm(unique_id).delete()
# Because we just deleted a bunch of objects in the database
# directly (rather than deleting individual objects in Django,
# in which case our post_save hooks would have flushed the
# individual objects from memcached for us), we need to flush
# memcached in order to ensure deleted objects aren't still
# present in the memcached cache.
flush_cache(None)
realm = do_create_realm(
string_id=unique_id,
name=unique_id,
description=unique_id,
plan_type=plan_type,
)
# Create a user with billing access
full_name = f"{plan_name}-admin"
email = f"{full_name}@zulip.com"
user = do_create_user(
email,
full_name,
realm,
full_name,
role=UserProfile.ROLE_REALM_OWNER,
acting_user=None,
tos_version=settings.TERMS_OF_SERVICE_VERSION,
)
stream, _ = create_stream_if_needed(
realm,
"all",
)
bulk_add_subscriptions(realm, [stream], [user], acting_user=None)
if customer_profile.is_remote_realm:
# Remote realm billing data on their local server is irrelevant.
return realm
if customer_profile.sponsorship_pending:
customer = Customer.objects.create(
realm=realm,
sponsorship_pending=customer_profile.sponsorship_pending,
)
return realm
if customer_profile.tier is None:
return realm
billing_session = RealmBillingSession(user)
customer = billing_session.update_or_create_stripe_customer()
assert customer.stripe_customer_id is not None
if customer_profile.card:
add_card_to_customer(customer)
create_plan_for_customer(customer, customer_profile)
return realm
def populate_remote_server(customer_profile: CustomerProfile) -> Dict[str, str]:
unique_id = customer_profile.unique_id
if (
customer_profile.is_sponsored
and customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_COMMUNITY
):
plan_type = RemoteZulipServer.PLAN_TYPE_COMMUNITY
elif customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_LEGACY:
plan_type = RemoteZulipServer.PLAN_TYPE_SELF_HOSTED
elif customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_BUSINESS:
plan_type = RemoteZulipServer.PLAN_TYPE_BUSINESS
elif customer_profile.tier is CustomerPlan.TIER_SELF_HOSTED_BASE:
plan_type = RemoteZulipServer.PLAN_TYPE_SELF_HOSTED
else:
raise AssertionError("Unexpected tier!")
server_uuid = str(uuid.uuid4())
api_key = server_uuid
hostname = f"{unique_id}.example.com"
# Delete existing remote server.
RemoteZulipServer.objects.filter(hostname=hostname).delete()
flush_cache(None)
remote_server = RemoteZulipServer.objects.create(
uuid=server_uuid,
api_key=api_key,
hostname=f"{unique_id}.example.com",
contact_email=f"{unique_id}@example.com",
plan_type=plan_type,
# TODO: Save property audit log data for server.
last_audit_log_update=timezone_now(),
)
billing_user = RemoteServerBillingUser.objects.create(
full_name="Server user",
remote_server=remote_server,
email=f"{unique_id}@example.com",
)
billing_session = RemoteServerBillingSession(remote_server, billing_user)
if customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_LEGACY:
# Create customer plan for these servers for temporary period.
renewal_date = datetime.strptime(customer_profile.renewal_date, TIMESTAMP_FORMAT).replace(
tzinfo=timezone.utc
)
end_date = datetime.strptime(customer_profile.end_date, TIMESTAMP_FORMAT).replace(
tzinfo=timezone.utc
)
billing_session.migrate_customer_to_legacy_plan(renewal_date, end_date)
# Scheduled server to upgrade to business plan.
if customer_profile.status == CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END:
# This attaches stripe_customer_id to customer.
customer = billing_session.update_or_create_stripe_customer()
add_card_to_customer(customer)
seat_count = 10
signed_seat_count, salt = sign_string(str(seat_count))
upgrade_request = UpgradeRequest(
billing_modality="charge_automatically",
schedule="annual",
signed_seat_count=signed_seat_count,
salt=salt,
license_management="automatic",
licenses=seat_count,
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
remote_server_plan_start_date=customer_profile.remote_server_plan_start_date,
)
billing_session.do_upgrade(upgrade_request)
elif customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_BUSINESS:
customer = billing_session.update_or_create_stripe_customer()
assert customer.stripe_customer_id is not None
add_card_to_customer(customer)
create_plan_for_customer(customer, customer_profile)
if customer_profile.sponsorship_pending:
billing_session.update_customer_sponsorship_status(True)
elif customer_profile.is_sponsored:
billing_session.do_change_plan_type(tier=None, is_sponsored=True)
return {
"unique_id": unique_id,
"server_uuid": server_uuid,
"api_key": api_key,
}
def populate_remote_realms(customer_profile: CustomerProfile) -> Dict[str, str]:
# Delete existing remote realm.
RemoteRealm.objects.filter(name=customer_profile.unique_id).delete()
flush_cache(None)
local_realm = populate_realm(customer_profile)
assert local_realm is not None
remote_server_uuid = settings.ZULIP_ORG_ID
assert remote_server_uuid is not None
remote_server = RemoteZulipServer.objects.filter(
uuid=remote_server_uuid,
).first()
if remote_server is None:
raise AssertionError("Remote server not found! Please run manage.py register_server")
update_remote_realm_data_for_server(
remote_server, get_realms_info_for_push_bouncer(local_realm.id)
)
remote_realm = RemoteRealm.objects.get(uuid=local_realm.uuid)
user = UserProfile.objects.filter(realm=local_realm).first()
assert user is not None
billing_user = RemoteRealmBillingUser.objects.create(
full_name=user.full_name,
remote_realm=remote_realm,
user_uuid=user.uuid,
email=user.email,
)
billing_session = RemoteRealmBillingSession(remote_realm, billing_user)
# TODO: Save property audit log data for server.
remote_realm.server.last_audit_log_update = timezone_now()
remote_realm.server.save(update_fields=["last_audit_log_update"])
customer = billing_session.update_or_create_stripe_customer()
assert customer.stripe_customer_id is not None
add_card_to_customer(customer)
if customer_profile.tier is not None:
billing_session.do_change_plan_type(
tier=customer_profile.tier, is_sponsored=customer_profile.is_sponsored
)
create_plan_for_customer(customer, customer_profile)
if customer_profile.sponsorship_pending:
billing_session.update_customer_sponsorship_status(True)
return {
"unique_id": customer_profile.unique_id,
"login_url": local_realm.uri + "/self-hosted-billing/",
}
Reference in New Issue View Git Blame Copy Permalink