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

populate_billing_realms: Add method to populate remote realms. 

We also avoid deleting all RemoteZulipServer at start so that
we don't delete the registered server for remote realms.
This commit is contained in:
Aman Agrawal 2023-12-05 11:49:22 +00:00 committed by Tim Abbott
parent cfb62e7ffa
commit 2c7afa19ba
1 changed files with 111 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

124
zilencer/management/commands/populate_billing_realms.py
View File

@ -5,12 +5,14 @@ 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,
@ -22,9 +24,11 @@ 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 RemoteZulipServer
from zilencer.models import RemoteRealm, RemoteZulipServer
from zilencer.views import update_remote_realm_data_for_server
from zproject.config import get_secret
current_time = timezone_now().strftime(TIMESTAMP_FORMAT)
@ -60,6 +64,12 @@ class Command(BaseCommand):
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
@ -129,7 +139,8 @@ class Command(BaseCommand):
unique_id="sponsored",
is_sponsored=True,
billing_schedule=CustomerPlan.BILLING_SCHEDULE_MONTHLY,
tier=CustomerPlan.TIER_CLOUD_STANDARD,
# Customer plan might not exist for sponsored realms.
tier=CustomerPlan.TIER_CLOUD_COMMUNITY,
),
CustomerProfile(
unique_id="free-trial",
@ -173,22 +184,54 @@ class Command(BaseCommand):
),
CustomerProfile(
unique_id="server-sponsored",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
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,
),
]
# Delete all existing remote servers
RemoteZulipServer.objects.all().delete()
flush_cache(None)
servers = []
remote_realms = []
for customer_profile in customer_profiles:
if customer_profile.is_remote_server:
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 not options.get("only_remote_server"):
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)
@ -197,6 +240,11 @@ class Command(BaseCommand):
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
@ -247,11 +295,14 @@ def create_plan_for_customer(customer: Customer, customer_profile: CustomerProfi
)
def populate_realm(customer_profile: CustomerProfile) -> Optional[Realm]:
unique_id = customer_profile.unique_id
if customer_profile.tier is None:
if customer_profile.is_remote_realm:
plan_type = Realm.PLAN_TYPE_SELF_HOSTED
elif customer_profile.tier is None:
plan_type = Realm.PLAN_TYPE_LIMITED
elif (
customer_profile.tier == CustomerPlan.TIER_CLOUD_STANDARD and customer_profile.is_sponsored
customer_profile.tier == CustomerPlan.TIER_CLOUD_COMMUNITY and customer_profile.is_sponsored
):
plan_type = Realm.PLAN_TYPE_STANDARD_FREE
elif customer_profile.tier == CustomerPlan.TIER_CLOUD_STANDARD:
@ -290,6 +341,7 @@ def create_plan_for_customer(customer: Customer, customer_profile: CustomerProfi
full_name,
role=UserProfile.ROLE_REALM_OWNER,
acting_user=None,
tos_version=settings.TERMS_OF_SERVICE_VERSION,
)
stream, _ = create_stream_if_needed(
@ -299,15 +351,19 @@ def create_plan_for_customer(customer: Customer, customer_profile: CustomerProfi
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
return realm
if customer_profile.tier is None:
return
return realm
billing_session = RealmBillingSession(user)
customer = billing_session.update_or_create_stripe_customer()
@ -335,6 +391,11 @@ def populate_remote_server(customer_profile: CustomerProfile) -> Dict[str, str]:
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,
@ -390,3 +451,40 @@ def populate_remote_server(customer_profile: CustomerProfile) -> Dict[str, str]:
"server_uuid": server_uuid,
"api_key": api_key,
}
def populate_remote_realms(customer_profile: CustomerProfile) -> Dict[str, str]:
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)
billing_session = RemoteRealmBillingSession(remote_realm)
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/",
}