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

post_analytics: Migrate plan from server to realm after upgrade.

This commit is contained in:
Aman Agrawal 2023-12-11 17:00:42 +00:00 committed by Tim Abbott
parent 64517a7ad3
commit 23d712391e
6 changed files with 152 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
corporate/lib/stripe.py
View File

@ -1113,8 +1113,8 @@ class BillingSession(ABC):
def ensure_current_plan_is_upgradable(
self, customer: Customer, new_plan_tier: int
) -> None: # nocoverage
# Upgrade for customers with an existing plan is only supported for remote servers right now.
if not hasattr(self, "remote_server"):
# Upgrade for customers with an existing plan is only supported for remote realm / server right now.
if isinstance(self, RealmBillingSession):
ensure_customer_does_not_have_active_plan(customer)
return
@ -3032,6 +3032,7 @@ class RemoteRealmBillingSession(BillingSession): # nocoverage
raise AssertionError("Unexpected tier")
# TODO: Audit logging and set usage limits.
# TODO: Set the usage limit in handle_customer_migration_from_server_to_realms.
self.remote_realm.plan_type = plan_type
self.remote_realm.save(update_fields=["plan_type"])
@ -3117,6 +3118,7 @@ class RemoteRealmBillingSession(BillingSession): # nocoverage
self, current_plan_tier: int, new_plan_tier: int
) -> PlanTierChangeType:
valid_plan_tiers = [
CustomerPlan.TIER_SELF_HOSTED_LEGACY,
CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
CustomerPlan.TIER_SELF_HOSTED_PLUS,
]
@ -3131,6 +3133,11 @@ class RemoteRealmBillingSession(BillingSession): # nocoverage
and new_plan_tier == CustomerPlan.TIER_SELF_HOSTED_PLUS
):
return PlanTierChangeType.UPGRADE
elif current_plan_tier == CustomerPlan.TIER_SELF_HOSTED_LEGACY and new_plan_tier in (
CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
CustomerPlan.TIER_SELF_HOSTED_PLUS,
):
return PlanTierChangeType.UPGRADE
else:
assert current_plan_tier == CustomerPlan.TIER_SELF_HOSTED_PLUS
assert new_plan_tier == CustomerPlan.TIER_SELF_HOSTED_BUSINESS

12
corporate/views/billing_page.py
View File

@ -117,12 +117,16 @@ def remote_realm_billing_page(
# If the realm is on a paid plan, show the sponsorship pending message
context["sponsorship_pending"] = True
if billing_session.remote_realm.plan_type == RemoteRealm.PLAN_TYPE_SELF_HOSTED:
if (
customer is None
or get_current_plan_by_customer(customer) is None
or (
billing_session.get_legacy_remote_server_next_plan_name(customer) is None
and billing_session.remote_realm.plan_type == RemoteRealm.PLAN_TYPE_SELF_HOSTED
)
):
return HttpResponseRedirect(reverse("remote_realm_plans_page", args=(realm_uuid,)))
if customer is None or get_current_plan_by_customer(customer) is None:
return HttpResponseRedirect(reverse("remote_realm_upgrade_page", args=(realm_uuid,)))
main_context = billing_session.get_billing_page_context()
if main_context:
context.update(main_context)

14
corporate/views/portico.py
View File

@ -149,7 +149,21 @@ def remote_realm_plans_page(
if context.customer_plan is None:
context.on_free_tier = not context.is_sponsored
else:
context.on_free_tier = context.customer_plan.tier in (
CustomerPlan.TIER_SELF_HOSTED_LEGACY,
CustomerPlan.TIER_SELF_HOSTED_BASE,
)
context.on_free_trial = is_customer_on_free_trial(context.customer_plan)
context.is_legacy_server_with_scheduled_upgrade = (
context.customer_plan.status == CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END
)
if context.is_legacy_server_with_scheduled_upgrade:
assert context.customer_plan.end_date is not None
context.legacy_server_new_plan = CustomerPlan.objects.get(
customer=customer,
billing_cycle_anchor=context.customer_plan.end_date,
status=CustomerPlan.NEVER_STARTED,
)
context.is_new_customer = (
not context.on_free_tier and context.customer_plan is None and not context.is_sponsored

5
corporate/views/upgrade.py
View File

@ -95,6 +95,7 @@ def remote_realm_upgrade(
default=None, str_validator=check_string_in(VALID_LICENSE_MANAGEMENT_VALUES)
),
licenses: Optional[int] = REQ(json_validator=check_int, default=None),
remote_server_plan_start_date: Optional[str] = REQ(default=None),
) -> HttpResponse: # nocoverage
try:
upgrade_request = UpgradeRequest(
@ -106,7 +107,7 @@ def remote_realm_upgrade(
licenses=licenses,
# TODO: tier should be a passed parameter.
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
remote_server_plan_start_date=None,
remote_server_plan_start_date=remote_server_plan_start_date,
)
data = billing_session.do_upgrade(upgrade_request)
return json_success(request, data)
@ -212,10 +213,12 @@ def remote_realm_upgrade_page(
billing_session: RemoteRealmBillingSession,
*,
manual_license_management: Json[bool] = False,
success_message: str = "",
) -> HttpResponse: # nocoverage
initial_upgrade_request = InitialUpgradeRequest(
manual_license_management=manual_license_management,
tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS,
success_message=success_message,
)
redirect_url, context = billing_session.get_initial_upgrade_context(initial_upgrade_request)

1
zerver/models.py
View File

@ -4858,6 +4858,7 @@ class AbstractRealmAuditLog(models.Model):
# RemoteRealm model resulting from modified realm information sent to us
# via send_analytics_to_push_bouncer.
REMOTE_REALM_VALUE_UPDATED = 20001
REMOTE_PLAN_TRANSFERRED_SERVER_TO_REALM = 20002
event_type = models.PositiveSmallIntegerField()

117
zilencer/views.py
View File

@ -5,6 +5,7 @@ from typing import Any, Dict, List, Optional, Type, TypeVar, Union
from uuid import UUID
import orjson
from django.conf import settings
from django.core.exceptions import ValidationError
from django.core.validators import URLValidator, validate_email
from django.db import IntegrityError, transaction
@ -23,7 +24,11 @@ from analytics.lib.counts import (
REMOTE_INSTALLATION_COUNT_STATS,
do_increment_logging_stat,
)
from corporate.lib.stripe import RemoteRealmBillingSession, do_deactivate_remote_server
from corporate.lib.stripe import (
RemoteRealmBillingSession,
RemoteServerBillingSession,
do_deactivate_remote_server,
)
from corporate.models import CustomerPlan, get_current_plan_by_customer
from zerver.decorator import require_post
from zerver.lib.exceptions import JsonableError, RemoteRealmServerMismatchError
@ -673,6 +678,101 @@ def update_remote_realm_data_for_server(
RemoteRealmAuditLog.objects.bulk_create(remote_realm_audit_logs)
def get_human_user_realm_uuids(realms: List[RealmDataForAnalytics]) -> List[UUID]: # nocoverage
billable_realm_uuids = []
for realm in realms:
# TODO: Remove the `zulipinternal` string_id check once no server is on 8.0-beta.
if (
realm.is_system_bot_realm
or realm.host.startswith("zulipinternal.")
or (settings.DEVELOPMENT and realm.host.startswith("analytics."))
):
continue
billable_realm_uuids.append(realm.uuid)
return billable_realm_uuids
@transaction.atomic
def handle_customer_migration_from_server_to_realms(
server: RemoteZulipServer, realms: List[RealmDataForAnalytics]
) -> None: # nocoverage
server_billing_session = RemoteServerBillingSession(server)
server_customer = server_billing_session.get_customer()
if server_customer is None:
return
server_plan = get_current_plan_by_customer(server_customer)
realm_uuids = get_human_user_realm_uuids(realms)
if not realm_uuids:
return
event_time = timezone_now()
remote_realm_audit_logs = []
if (
server_plan is not None
and server_plan.tier == CustomerPlan.TIER_SELF_HOSTED_LEGACY
and server_plan.status == CustomerPlan.ACTIVE
):
assert server.plan_type == RemoteZulipServer.PLAN_TYPE_SELF_HOSTED
assert server_plan.end_date is not None
remote_realms = RemoteRealm.objects.filter(
uuid__in=realm_uuids, server=server, plan_type=RemoteRealm.PLAN_TYPE_SELF_HOSTED
)
# Verify that all the realms are on self hosted plan.
assert remote_realms.count() == len(realm_uuids)
# End existing plan for server.
server_plan.status = CustomerPlan.ENDED
server_plan.save(update_fields=["status"])
# Create new legacy plan for each remote realm.
for remote_realm in remote_realms:
RemoteRealmBillingSession(remote_realm).migrate_customer_to_legacy_plan(
server_plan.billing_cycle_anchor, server_plan.end_date
)
remote_realm_audit_logs.append(
RemoteRealmAuditLog(
server=server,
remote_realm=remote_realm,
event_type=RemoteRealmAuditLog.REMOTE_PLAN_TRANSFERRED_SERVER_TO_REALM,
event_time=event_time,
# No extra_data since there was no real change in any RemoteRealm attribute.
)
)
# We only do this migration if there is only one realm besides the system bot realm.
elif len(realm_uuids) == 1:
remote_realm = RemoteRealm.objects.get(
uuid=realm_uuids[0], plan_type=RemoteRealm.PLAN_TYPE_SELF_HOSTED
)
# Migrate customer from server to remote realm if there is only one realm.
server_customer.remote_realm = remote_realm
server_customer.remote_server = None
server_customer.save(update_fields=["remote_realm", "remote_server"])
# TODO: Set usage limits for remote realm and server.
# Might be better to call do_change_plan_type here.
remote_realm.plan_type = server.plan_type
remote_realm.save(update_fields=["plan_type"])
server.plan_type = RemoteZulipServer.PLAN_TYPE_SELF_HOSTED
server.save(update_fields=["plan_type"])
remote_realm_audit_logs.append(
RemoteRealmAuditLog(
server=server,
remote_realm=remote_realm,
event_type=RemoteRealmAuditLog.REMOTE_PLAN_TRANSFERRED_SERVER_TO_REALM,
event_time=event_time,
extra_data={
"attr_name": "plan_type",
"old_value": RemoteRealm.PLAN_TYPE_SELF_HOSTED,
"new_value": remote_realm.plan_type,
},
)
)
RemoteRealmAuditLog.objects.bulk_create(remote_realm_audit_logs)
@typed_endpoint
@transaction.atomic
def remote_server_post_analytics(
@ -725,6 +825,21 @@ def remote_server_post_analytics(
if remote_server_version_updated:
fix_remote_realm_foreign_keys(server, realms)
try:
handle_customer_migration_from_server_to_realms(server, realms)
except Exception: # nocoverage
logger.exception(
"%s: Failed to migrate customer from server (id: %s) to realms",
request.path,
server.id,
stack_info=True,
)
raise JsonableError(
_(
"Failed to migrate customer from server to realms. Please contact support for assistance."
)
)
realm_id_to_remote_realm = build_realm_id_to_remote_realm_dict(server, realms)
remote_realm_counts = [