remote_billing: Make handle_customer_migration_... more robust.

The logic in the case where there's only one realm and the function
tries to migrate the server's plan to it, had two main unhandled edge
cases that would throw exceptions:
1.
```
        remote_realm = RemoteRealm.objects.get(
            uuid=realm_uuids[0], plan_type=RemoteRealm.PLAN_TYPE_SELF_MANAGED
        )
```

This could throw an exception if the RemoteRealm exists, but has an
active e.g. Legacy plan. Then there'd be no object matching the
plan_type in the query, raising RemoteRealm.DoesNotExist.

2. If the RemoteRealm had e.g. a Legacy plan in the past, that's now
   expired, then it'd have a Customer object. Meaning that the attempt
   to move the server's customer to the realm:
   `server_plan.customer = remote_realm_customer`
   would trigger an IntegrityError since a RemoteRealm can't have two
   Customer objects.

In simple cases the situation in (2) can still be easily migrated, by
moving the plan from the server's customer to the realm's customer.

corporate/tests/test_remote_billing.py

@@ -27,10 +27,12 @@ from zerver.actions.realm_settings import do_deactivate_realm
 from zerver.lib.exceptions import RemoteRealmServerMismatchError
 from zerver.lib.rate_limiter import RateLimitedIPAddr
 from zerver.lib.remote_server import send_server_data_to_push_bouncer
+from zerver.lib.send_email import FromAddress
 from zerver.lib.test_classes import BouncerTestCase
 from zerver.lib.test_helpers import ratelimit_rule
 from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.models import Realm, UserProfile
+from zerver.models.realms import get_realm
 from zilencer.models import (
     PreregistrationRemoteRealmBillingUser,
     PreregistrationRemoteServerBillingUser,
@@ -876,6 +878,117 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase):
             RemoteRealmBillingSession(remote_realm_with_plan).get_next_plan(plan), server_next_plan
         )
 
+    @responses.activate
+    def test_transfer_plan_from_server_to_realm_when_realm_has_customer(
+        self,
+    ) -> None:
+        self.login("desdemona")
+        desdemona = self.example_user("desdemona")
+        zulip_realm = get_realm("zulip")
+
+        server_billing_session = RemoteServerBillingSession(self.server)
+        server_customer = server_billing_session.update_or_create_customer(stripe_customer_id=None)
+        server_plan = CustomerPlan.objects.create(
+            customer=server_customer,
+            billing_cycle_anchor=timezone_now(),
+            billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL,
+            tier=CustomerPlan.TIER_SELF_HOSTED_COMMUNITY,
+            status=CustomerPlan.ACTIVE,
+        )
+        self.server.plan_type = RemoteZulipServer.PLAN_TYPE_COMMUNITY
+        self.server.save(update_fields=["plan_type"])
+
+        # Delete any existing remote realms.
+        RemoteRealm.objects.all().delete()
+
+        # We want there to be only a single (non-system bot) realm on the server for our setup.
+        Realm.objects.exclude(string_id__in=["zulip", "zulipinternal"]).update(deactivated=True)
+
+        # Send server data to push bouncer.
+        self.add_mock_response()
+        send_server_data_to_push_bouncer(consider_usage_statistics=False)
+
+        # Let's create a plan for the realm. This will conflict with the server plan.
+        remote_realm = RemoteRealm.objects.get(uuid=zulip_realm.uuid)
+        realm_billing_session = RemoteRealmBillingSession(remote_realm)
+        realm_customer = realm_billing_session.update_or_create_customer(stripe_customer_id=None)
+        realm_plan = CustomerPlan.objects.create(
+            customer=realm_customer,
+            billing_cycle_anchor=timezone_now(),
+            billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL,
+            tier=CustomerPlan.TIER_SELF_HOSTED_LEGACY,
+            status=CustomerPlan.ACTIVE,
+        )
+        remote_realm.plan_type = RemoteRealm.PLAN_TYPE_SELF_MANAGED_LEGACY
+        remote_realm.save(update_fields=["plan_type"])
+
+        with self.assertLogs("zilencer.views", "WARN") as mock_warn:
+            result = self.execute_remote_billing_authentication_flow(
+                desdemona, return_from_auth_url=True
+            )
+        self.assertEqual(
+            mock_warn.output,
+            [
+                f"WARNING:zilencer.views:Failed to migrate customer from server (id: {remote_realm.server.id}) to realm (id: {remote_realm.id}): "
+                "RemoteRealm customer already exists and plans can't be migrated automatically."
+            ],
+        )
+        self.assert_json_error(
+            result,
+            f"Couldn't reconcile billing data between server and realm. Please contact {FromAddress.SUPPORT}",
+        )
+
+        # If the realm's plan is ENDED, it's safe to move the server plan over.
+        realm_plan.status = CustomerPlan.ENDED
+        realm_plan.save(update_fields=["status"])
+        # However, not if the server's status indicates that there's some kind
+        # of plan change queued up after the plan, since that state would be
+        # harder and more risky to try to migrate.
+        server_plan.status = CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END
+        server_plan.save(update_fields=["status"])
+
+        with self.assertLogs("zilencer.views", "WARN") as mock_warn:
+            result = self.execute_remote_billing_authentication_flow(
+                desdemona, return_from_auth_url=True
+            )
+        self.assertEqual(
+            mock_warn.output,
+            [
+                f"WARNING:zilencer.views:Failed to migrate customer from server (id: {remote_realm.server.id}) to realm (id: {remote_realm.id}): "
+                "RemoteRealm customer already exists and plans can't be migrated automatically."
+            ],
+        )
+        self.assert_json_error(
+            result,
+            f"Couldn't reconcile billing data between server and realm. Please contact {FromAddress.SUPPORT}",
+        )
+
+        # Finally, we simulate a regular, ACTIVE plan for the server again. Combined with
+        # the ENDED plan for the realm, we now have a simple case, where the migration
+        # should proceed.
+        server_plan.status = CustomerPlan.ACTIVE
+        server_plan.save(update_fields=["status"])
+
+        result = self.execute_remote_billing_authentication_flow(
+            desdemona, return_from_auth_url=False
+        )
+        self.assertEqual(result.status_code, 302)
+
+        # Server plan status was reset
+        self.server.refresh_from_db()
+        self.assertEqual(self.server.plan_type, RemoteZulipServer.PLAN_TYPE_SELF_MANAGED)
+
+        # The Customer objects remain as they were.
+        self.assertEqual(get_customer_by_remote_realm(remote_realm), realm_customer)
+        self.assertEqual(get_customer_by_remote_server(self.server), server_customer)
+
+        # The plan that used to be for the server, has been migrated to the realm customer:
+        self.assertEqual(get_current_plan_by_customer(server_customer), None)
+        self.assertEqual(get_current_plan_by_customer(realm_customer), server_plan)
+
+        remote_realm.refresh_from_db()
+        self.assertEqual(remote_realm.plan_type, RemoteRealm.PLAN_TYPE_COMMUNITY)
+
     @responses.activate
     def test_transfer_business_plan_from_server_to_realm(
         self,

corporate/views/remote_billing_page.py

@@ -201,6 +201,10 @@ def remote_realm_billing_finalize_login(
 
     try:
         handle_customer_migration_from_server_to_realms(server=remote_server)
+    except JsonableError:
+        # JsonableError should be propagated up, as they are meant to convey
+        # a json error response to be returned.
+        raise
     except Exception:  # nocoverage
         billing_logger.exception(
             "%s: Failed to migrate customer from server (id: %s) to realms",

zilencer/views.py

@@ -32,7 +32,11 @@ from corporate.lib.stripe import (
     do_deactivate_remote_server,
     get_push_status_for_remote_request,
 )
-from corporate.models import CustomerPlan, get_current_plan_by_customer
+from corporate.models import (
+    CustomerPlan,
+    get_current_plan_by_customer,
+    get_customer_by_remote_realm,
+)
 from zerver.decorator import require_post
 from zerver.lib.email_validation import validate_disposable
 from zerver.lib.exceptions import (
@@ -56,6 +60,7 @@ from zerver.lib.remote_server import (
 )
 from zerver.lib.request import REQ, RequestNotes, has_request_variables
 from zerver.lib.response import json_success
+from zerver.lib.send_email import FromAddress
 from zerver.lib.timestamp import timestamp_to_datetime
 from zerver.lib.typed_endpoint import JsonBodyPayload, typed_endpoint
 from zerver.lib.types import RemoteRealmDictValue
@@ -1000,13 +1005,47 @@ def handle_customer_migration_from_server_to_realms(
     elif len(realm_uuids) == 1:
         # Here, we have exactly one non-system-bot realm, and some
         # sort of plan on the server; move it to the realm.
-        remote_realm = RemoteRealm.objects.get(
-            uuid=realm_uuids[0], plan_type=RemoteRealm.PLAN_TYPE_SELF_MANAGED
-        )
+        remote_realm = RemoteRealm.objects.get(uuid=realm_uuids[0], server=server)
+        remote_realm_customer = get_customer_by_remote_realm(remote_realm)
+
         # 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"])
+        if remote_realm_customer is None:
+            # In this case the migration is easy, since we can just move the customer
+            # object directly.
+            server_customer.remote_realm = remote_realm
+            server_customer.remote_server = None
+            server_customer.save(update_fields=["remote_realm", "remote_server"])
+        else:
+            # If there's a Customer object for the realm already, things are harder,
+            # because it's an unusual state and there may be a plan already active
+            # for the realm, or there may have been.
+            # In the simplest case, where the realm doesn't have an active plan and the
+            # server's plan state can easily be moved, we proceed with the migrations.
+            remote_realm_plan = get_current_plan_by_customer(remote_realm_customer)
+            if (
+                remote_realm_plan is None
+                and server_plan.status != CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END
+            ):
+                # This is a simple case where we don't have to worry about the realm already
+                # having an active plan, or the server having a next plan scheduled that we'd need
+                # to figure out how to migrate correctly as well.
+                # Any other case is too complex to handle here, and should be handled manually,
+                # especially since that should be extremely rare.
+                server_plan.customer = remote_realm_customer
+                server_plan.save(update_fields=["customer"])
+            else:
+                logger.warning(
+                    "Failed to migrate customer from server (id: %s) to realm (id: %s): RemoteRealm customer already exists "
+                    "and plans can't be migrated automatically.",
+                    server.id,
+                    remote_realm.id,
+                )
+                raise JsonableError(
+                    _(
+                        "Couldn't reconcile billing data between server and realm. Please contact {support_email}"
+                    ).format(support_email=FromAddress.SUPPORT)
+                )
+
         # TODO: Might be better to call do_change_plan_type here.
         remote_realm.plan_type = server.plan_type
         remote_realm.save(update_fields=["plan_type"])