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

commands: Add script to create servers on legacy plan. 

Also adds `SWITCH_PLAN_TIER_AT_PLAN_END` for `CustomerPlan`
which will be used to mark status of remote server legacy
plans which are scheduled for an upgrade.
This commit is contained in:
Aman Agrawal 2023-12-04 13:03:24 +00:00 committed by Tim Abbott
parent c651c4f668
commit 7d83508235
4 changed files with 218 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
corporate/lib/stripe.py
View File

@ -1932,7 +1932,7 @@ class BillingSession(ABC):
# TODO: Implement downgrade that is a change from and to a paid plan
# tier. This should keep the same billing cycle schedule and change
# the plan when it's next invoiced vs immediately. Note this will need
# new CustomerPlan.status value, e.g. SWITCH_PLAN_TIER_AT_END_OF_CYCLE.
# new CustomerPlan.status value, e.g. SWITCH_PLAN_TIER_AT_PLAN_END.
assert type_of_tier_change == PlanTierChangeType.DOWNGRADE # nocoverage
return "" # nocoverage
@ -2938,6 +2938,57 @@ class RemoteServerBillingSession(BillingSession): # nocoverage
self.remote_server.org_type = org_type
self.remote_server.save(update_fields=["org_type"])
def add_server_to_legacy_plan(
self,
renewal_date: datetime,
end_date: datetime,
) -> None:
# Set stripe_customer_id to None to avoid customer being charged without a payment method.
customer = Customer.objects.create(
remote_server=self.remote_server, stripe_customer_id=None
)
# Servers on legacy plan which are scheduled to be upgraded have 2 plans.
# This plan will be used to track the current status of SWITCH_PLAN_TIER_AT_PLAN_END
# and will not charge the customer. The other plan will be used to track the new plan
# customer will move to the end of this plan.
legacy_plan_anchor = renewal_date
legacy_plan_params = {
"billing_cycle_anchor": legacy_plan_anchor,
"status": CustomerPlan.ACTIVE,
"tier": CustomerPlan.TIER_SELF_HOSTED_LEGACY,
# End when the new plan starts.
"end_date": end_date,
# The primary mechanism for preventing charges under this
# plan is setting a null `next_invoice_date`, but setting
# a 0 price is useful defense in depth here.
"next_invoice_date": None,
"price_per_license": 0,
"billing_schedule": CustomerPlan.BILLING_SCHEDULE_ANNUAL,
"automanage_licenses": True,
}
legacy_plan = CustomerPlan.objects.create(
customer=customer,
**legacy_plan_params,
)
# Create a ledger entry for the legacy plan for tracking purposes.
billed_licenses = self.current_count_for_billed_licenses()
ledger_entry = LicenseLedger.objects.create(
plan=legacy_plan,
is_renewal=True,
event_time=legacy_plan_anchor,
licenses=billed_licenses,
licenses_at_next_renewal=billed_licenses,
)
legacy_plan.invoiced_through = ledger_entry
legacy_plan.save(update_fields=["invoiced_through"])
self.write_to_audit_log(
event_type=AuditLogEventType.CUSTOMER_PLAN_CREATED,
event_time=legacy_plan_anchor,
extra_data=legacy_plan_params,
)
def stripe_customer_has_credit_card_as_default_payment_method(
stripe_customer: stripe.Customer,
@ -3138,6 +3189,7 @@ def invoice_plans_as_needed(event_time: Optional[datetime] = None) -> None:
for plan in CustomerPlan.objects.filter(next_invoice_date__lte=event_time):
if plan.customer.realm is not None:
RealmBillingSession(realm=plan.customer.realm).invoice_plan(plan, event_time)
# TODO: Assert that we never invoice legacy plans.
def is_realm_on_free_trial(realm: Realm) -> bool:

1
corporate/models.py
View File

@ -277,6 +277,7 @@ class CustomerPlan(models.Model):
SWITCH_PLAN_TIER_NOW = 5
SWITCH_TO_MONTHLY_AT_END_OF_CYCLE = 6
DOWNGRADE_AT_END_OF_FREE_TRIAL = 7
SWITCH_PLAN_TIER_AT_PLAN_END = 8
# "Live" plans should have a value < LIVE_STATUS_THRESHOLD.
# There should be at most one live plan per customer.
LIVE_STATUS_THRESHOLD = 10

63
zilencer/management/commands/add_server_to_legacy_plan.py Normal file
View File

@ -0,0 +1,63 @@
import datetime
from typing import Any
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 RemoteServerBillingSession
from scripts.lib.zulip_tools import TIMESTAMP_FORMAT
from zilencer.models import RemoteZulipServer
class Command(BaseCommand):
help = "Assigns an existing RemoteZulipServer to the legacy plan"
@override
def add_arguments(self, parser: CommandParser) -> None:
parser.add_argument(
"server_id",
type=int,
help="ID of the RemoteZulipServer to be assigned to the legacy plan",
)
parser.add_argument(
"renewal_date",
type=str,
help="Billing cycle renewal date in the format YYYY-MM-DD-HH-MM-SS",
)
parser.add_argument(
"end_date",
type=str,
help="Billing cycle end date in the format YYYY-MM-DD-HH-MM-SS",
)
@override
def handle(self, *args: Any, **options: Any) -> None:
server_id = options["server_id"]
renewal_date_str = options.get("renewal_date")
if renewal_date_str is None:
renewal_date = timezone_now()
else:
renewal_date = datetime.datetime.strptime(renewal_date_str, TIMESTAMP_FORMAT).replace(
tzinfo=datetime.timezone.utc
)
end_date_str = options.get("end_date")
if end_date_str is None:
raise ValueError("end_date must be provided")
end_date = datetime.datetime.strptime(end_date_str, TIMESTAMP_FORMAT).replace(
tzinfo=datetime.timezone.utc
)
server = RemoteZulipServer.objects.get(id=server_id)
self.add_server_to_legacy_plan(server, renewal_date, end_date)
def add_server_to_legacy_plan(
self,
server: RemoteZulipServer,
renewal_date: datetime.datetime,
end_date: datetime.datetime,
) -> None:
billing_schedule = RemoteServerBillingSession(server)
billing_schedule.add_server_to_legacy_plan(renewal_date, end_date)

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

@ -1,39 +1,56 @@
import contextlib
import datetime
import uuid
from dataclasses import dataclass
from typing import Any, Optional
from typing import Any, Dict, Optional
import stripe
from django.core.management.base import BaseCommand
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, add_months
from corporate.lib.stripe import RealmBillingSession, RemoteServerBillingSession, add_months
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.streams import create_stream_if_needed
from zerver.models import Realm, UserProfile, get_realm
from zilencer.models import RemoteZulipServer
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
renewal_date: str = current_time
end_date: str = "2030-10-10-01-10-10"
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",
)
@override
def handle(self, *args: Any, **options: Any) -> None:
# Create a realm for each plan type
@ -110,11 +127,44 @@ class Command(BaseCommand):
tier=CustomerPlan.TIER_CLOUD_STANDARD,
status=CustomerPlan.FREE_TRIAL,
),
# Use `server` keyword in the unique_id to indicate that this is a profile for remote server.
CustomerProfile(
unique_id="legacy-server",
tier=CustomerPlan.TIER_SELF_HOSTED_LEGACY,
),
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,
),
CustomerProfile(
unique_id="business-server",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
),
CustomerProfile(
unique_id="business-server-payment-starts-in-future",
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
),
]
# Create a realm for each customer profile
# Delete all existing remote servers
RemoteZulipServer.objects.all().delete()
flush_cache(None)
servers = []
for customer_profile in customer_profiles:
populate_realm(customer_profile)
if "server" in customer_profile.unique_id:
server_conf = populate_remote_server(customer_profile)
servers.append(server_conf)
elif not options.get("only_remote_server"):
populate_realm(customer_profile)
print("-" * 40)
for server in servers:
for key, value in server.items():
print(f"{key}: {value}")
print("-" * 40)
def populate_realm(customer_profile: CustomerProfile) -> None:
@ -228,3 +278,49 @@ def populate_realm(customer_profile: CustomerProfile) -> None:
is_renewal=True,
plan=customer_plan,
)
def populate_remote_server(customer_profile: CustomerProfile) -> Dict[str, str]:
unique_id = customer_profile.unique_id
if 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
else:
raise AssertionError("Unexpected tier!")
server_uuid = str(uuid.uuid4())
api_key = server_uuid
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,
)
if customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_LEGACY:
# Create customer plan for these servers for temporary period.
billing_session = RemoteServerBillingSession(remote_server)
renewal_date = renewal_date = datetime.datetime.strptime(
customer_profile.renewal_date, TIMESTAMP_FORMAT
).replace(tzinfo=datetime.timezone.utc)
end_date = datetime.datetime.strptime(customer_profile.end_date, TIMESTAMP_FORMAT).replace(
tzinfo=datetime.timezone.utc
)
billing_session.add_server_to_legacy_plan(renewal_date, end_date)
elif customer_profile.tier == CustomerPlan.TIER_SELF_HOSTED_BUSINESS:
# TBD
pass
if customer_profile.status == CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END:
# TBD
pass
return {
"unique_id": unique_id,
"server_uuid": server_uuid,
"api_key": api_key,
}