Use djstripe models for organization subscriptions (#9486)

Author: stsewd

readthedocs/organizations/migrations/0010_add_stripe_customer.py

@@ -0,0 +1,40 @@
+# Generated by Django 3.2.13 on 2022-08-10 02:21
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("djstripe", "0010_alter_customer_balance"),
+        ("organizations", "0009_update_meta_options"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="historicalorganization",
+            name="stripe_customer",
+            field=models.ForeignKey(
+                blank=True,
+                db_constraint=False,
+                null=True,
+                on_delete=django.db.models.deletion.DO_NOTHING,
+                related_name="+",
+                to="djstripe.customer",
+                verbose_name="Stripe customer",
+            ),
+        ),
+        migrations.AddField(
+            model_name="organization",
+            name="stripe_customer",
+            field=models.OneToOneField(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.SET_NULL,
+                related_name="rtd_organization",
+                to="djstripe.customer",
+                verbose_name="Stripe customer",
+            ),
+        ),
+    ]

readthedocs/organizations/models.py

@@ -93,6 +93,14 @@ class Organization(models.Model):
         blank=True,
         null=True,
     )
+    stripe_customer = models.OneToOneField(
+        "djstripe.Customer",
+        verbose_name=_("Stripe customer"),
+        on_delete=models.SET_NULL,
+        related_name="rtd_organization",
+        null=True,
+        blank=True,
+    )
 
     # Managers
     objects = OrganizationQuerySet.as_manager()
@@ -122,8 +130,18 @@ def save(self, *args, **kwargs):  # pylint: disable=signature-differs
         if not self.slug:
             self.slug = slugify(self.name)
 
+        if self.stripe_customer:
+            self.stripe_id = self.stripe_customer.id
+
         super().save(*args, **kwargs)
 
+    def get_stripe_metadata(self):
+        """Get metadata for the stripe account."""
+        return {
+            "org:id": self.id,
+            "org:slug": self.slug,
+        }
+
     # pylint: disable=no-self-use
     def add_member(self, user, team):
         """

readthedocs/settings/base.py

@@ -797,6 +797,12 @@ def DOCKER_LIMITS(self):
     DJSTRIPE_FOREIGN_KEY_TO_FIELD = "id"
     DJSTRIPE_USE_NATIVE_JSONFIELD = True  # We recommend setting to True for new installations
 
+    # Disable adding djstripe metadata to the Customer objects.
+    # We are managing the subscriber relationship by ourselves,
+    # since we have subscriptions attached to an organization or gold user
+    # we can't make use of the DJSTRIPE_SUBSCRIBER_MODEL setting.
+    DJSTRIPE_SUBSCRIBER_CUSTOMER_KEY = None
+
     # Do Not Track support
     DO_NOT_TRACK_ENABLED = False
 

readthedocs/subscriptions/managers.py

@@ -2,11 +2,11 @@
 
 from datetime import datetime
 
-import stripe
 import structlog
 from django.conf import settings
 from django.db import models
 from django.utils import timezone
+from djstripe.enums import SubscriptionStatus
 
 from readthedocs.core.history import set_change_reason
 from readthedocs.subscriptions.utils import get_or_create_stripe_customer
@@ -44,33 +44,33 @@ def get_or_create_default_subscription(self, organization):
             return None
 
         stripe_customer = get_or_create_stripe_customer(organization)
-        stripe_subscription = stripe.Subscription.create(
-            customer=stripe_customer.id,
-            items=[{"price": plan.stripe_id}],
-            trial_period_days=plan.trial,
-        )
-        # Stripe renamed ``start`` to ``start_date``,
-        # our API calls will return the new object,
-        # but webhooks will still return the old object
-        # till we change the default version.
-        # TODO: use stripe_subscription.start_date after the webhook version has been updated.
-        start_date = getattr(
-            stripe_subscription, "start", getattr(stripe_subscription, "start_date")
+        stripe_subscriptions = stripe_customer.subscriptions.exclude(
+            status=SubscriptionStatus.canceled
         )
+        if stripe_subscriptions.count() > 1:
+            log.warning(
+                "Customer with more than one active subscription.",
+                stripe_customer=stripe_customer.id,
+            )
+
+        stripe_subscription = stripe_subscriptions.last()
+        if not stripe_subscription:
+            # TODO: djstripe 2.6.x doesn't return the subscription object
+            # on subscribe(), but 2.7.x (unreleased) does!
+            stripe_customer.subscribe(
+                items=[{"price": plan.stripe_id}],
+                trial_period_days=plan.trial,
+            )
+            stripe_subscription = stripe_customer.subscriptions.latest()
+
         return self.create(
             plan=plan,
             organization=organization,
             stripe_id=stripe_subscription.id,
             status=stripe_subscription.status,
-            start_date=timezone.make_aware(
-                datetime.fromtimestamp(int(start_date)),
-            ),
-            end_date=timezone.make_aware(
-                datetime.fromtimestamp(int(stripe_subscription.current_period_end)),
-            ),
-            trial_end_date=timezone.make_aware(
-                datetime.fromtimestamp(int(stripe_subscription.trial_end)),
-            ),
+            start_date=stripe_subscription.start_date,
+            end_date=stripe_subscription.current_period_end,
+            trial_end_date=stripe_subscription.trial_end,
         )
 
     def update_from_stripe(self, *, rtd_subscription, stripe_subscription):

readthedocs/subscriptions/tests/test_views.py

@@ -1,12 +1,14 @@
 from unittest import mock
 
 import requests_mock
-import stripe
 from django.conf import settings
 from django.contrib.auth.models import User
 from django.test import TestCase, override_settings
 from django.urls import reverse
+from django.utils import timezone
 from django_dynamic_fixture import get
+from djstripe import models as djstripe
+from djstripe.enums import SubscriptionStatus
 
 from readthedocs.organizations.models import Organization
 from readthedocs.subscriptions.models import Plan, Subscription
@@ -56,65 +58,81 @@ def test_manage_subscription(self, mock_request):
             fetch_redirect_response=False,
         )
 
-    @mock.patch("readthedocs.subscriptions.managers.stripe.Subscription.create")
-    @mock.patch("readthedocs.subscriptions.utils.stripe.Customer.retrieve")
+    @mock.patch(
+        "readthedocs.subscriptions.utils.stripe.Customer.modify", new=mock.MagicMock
+    )
+    @mock.patch("readthedocs.subscriptions.utils.djstripe.Customer._get_or_retrieve")
     @mock.patch("readthedocs.subscriptions.utils.stripe.Customer.create")
     def test_user_without_subscription(
-        self, customer_create_mock, customer_retrieve_mock, subscription_create_mock
+        self, customer_create_mock, customer_retrieve_mock
     ):
-        subscription_create_mock.return_value = stripe.Subscription.construct_from(
-            values={
-                "id": "sub_a1b2c3",
-                "start_date": 1610532715.085267,
-                "current_period_end": 1610532715.085267,
-                "trial_end": 1610532715.085267,
-                "status": "active",
-            },
-            key=None,
+        stripe_customer = get(
+            djstripe.Customer,
+            id="cus_a1b2c3",
         )
-        customer_retrieve_mock.return_value = stripe.Customer.construct_from(
-            values={"id": "cus_a1b2c3"},
-            key=None,
+        stripe_subscription = get(
+            djstripe.Subscription,
+            id="sub_a1b2c3",
+            start_date=timezone.now(),
+            current_period_end=timezone.now() + timezone.timedelta(days=30),
+            trial_end=timezone.now() + timezone.timedelta(days=30),
+            status=SubscriptionStatus.active,
+            customer=stripe_customer,
         )
+        stripe_customer.subscribe = mock.MagicMock()
+        stripe_customer.subscribe.return_value = stripe_subscription
+        customer_retrieve_mock.return_value = stripe_customer
+
         self.subscription.delete()
         self.organization.refresh_from_db()
         self.assertFalse(hasattr(self.organization, 'subscription'))
+        self.assertIsNone(self.organization.stripe_customer)
+
         resp = self.client.get(reverse('subscription_detail', args=[self.organization.slug]))
         self.assertEqual(resp.status_code, 200)
         self.organization.refresh_from_db()
         subscription = self.organization.subscription
         self.assertEqual(subscription.status, 'active')
         self.assertEqual(subscription.stripe_id, 'sub_a1b2c3')
+        self.assertEqual(self.organization.stripe_customer, stripe_customer)
         customer_retrieve_mock.assert_called_once()
         customer_create_mock.assert_not_called()
 
-    @mock.patch("readthedocs.subscriptions.managers.stripe.Subscription.create")
-    @mock.patch("readthedocs.subscriptions.utils.stripe.Customer.retrieve")
+    @mock.patch(
+        "readthedocs.subscriptions.utils.djstripe.Customer.sync_from_stripe_data"
+    )
+    @mock.patch("readthedocs.subscriptions.utils.djstripe.Customer._get_or_retrieve")
     @mock.patch("readthedocs.subscriptions.utils.stripe.Customer.create")
     def test_user_without_subscription_and_customer(
-        self, customer_create_mock, customer_retrieve_mock, subscription_create_mock
+        self, customer_create_mock, customer_retrieve_mock, sync_from_stripe_data_mock
     ):
-        subscription_create_mock.return_value = stripe.Subscription.construct_from(
-            values={
-                "id": "sub_a1b2c3",
-                "start_date": 1610532715.085267,
-                "current_period_end": 1610532715.085267,
-                "trial_end": 1610532715.085267,
-                "status": "active",
-            },
-            key=None,
+        stripe_customer = get(
+            djstripe.Customer,
+            id="cus_a1b2c3",
         )
-        customer_create_mock.return_value = stripe.Customer.construct_from(
-            values={"id": "cus_a1b2c3"},
-            key=None,
+        stripe_subscription = get(
+            djstripe.Subscription,
+            id="sub_a1b2c3",
+            start_date=timezone.now(),
+            current_period_end=timezone.now() + timezone.timedelta(days=30),
+            trial_end=timezone.now() + timezone.timedelta(days=30),
+            status=SubscriptionStatus.active,
+            customer=stripe_customer,
         )
+        stripe_customer.subscribe = mock.MagicMock()
+        stripe_customer.subscribe.return_value = stripe_subscription
+        customer_retrieve_mock.return_value = None
+        sync_from_stripe_data_mock.return_value = stripe_customer
+
         # When stripe_id is None, a new customer is created.
         self.organization.stripe_id = None
         self.organization.save()
         self.subscription.delete()
         self.organization.refresh_from_db()
         self.assertFalse(hasattr(self.organization, 'subscription'))
         self.assertIsNone(self.organization.stripe_id)
+        self.assertIsNone(self.organization.stripe_customer)
+
         customer_retrieve_mock.reset_mock()
         resp = self.client.get(reverse('subscription_detail', args=[self.organization.slug]))
         self.assertEqual(resp.status_code, 200)
@@ -123,9 +141,9 @@ def test_user_without_subscription_and_customer(
         self.assertEqual(subscription.status, 'active')
         self.assertEqual(subscription.stripe_id, 'sub_a1b2c3')
         self.assertEqual(self.organization.stripe_id, 'cus_a1b2c3')
+        self.assertEqual(self.organization.stripe_customer, stripe_customer)
         customer_create_mock.assert_called_once()
-        # Called from a signal of .save()
-        customer_retrieve_mock.assert_called_once()
+        customer_retrieve_mock.assert_not_called()
 
     def test_user_with_canceled_subscription(self):
         self.subscription.status = 'canceled'

readthedocs/subscriptions/utils.py

@@ -1,53 +1,70 @@
 """Utilities to interact with subscriptions and stripe."""
 
-import structlog
-
 import stripe
+import structlog
+from djstripe import models as djstripe
 from stripe.error import InvalidRequestError
 
 log = structlog.get_logger(__name__)
 
 
 def create_stripe_customer(organization):
     """Create a stripe customer for organization."""
-    stripe_customer = stripe.Customer.create(
+    stripe_data = stripe.Customer.create(
         email=organization.email,
         description=organization.name,
+        metadata=organization.get_stripe_metadata(),
     )
+    stripe_customer = djstripe.Customer.sync_from_stripe_data(stripe_data)
     if organization.stripe_id:
         log.warning(
-            'Overriding existing stripe customer. ',
+            "Overriding existing stripe customer.",
             organization_slug=organization.slug,
             previous_stripe_customer=organization.stripe_id,
             stripe_customer=stripe_customer.id,
         )
-    organization.stripe_id = stripe_customer.id
+    organization.stripe_customer = stripe_customer
     organization.save()
     return stripe_customer
 
 
 def get_or_create_stripe_customer(organization):
     """
-    Retrieve the stripe customer from `organization.stripe_id` or create a new one.
+    Retrieve the stripe customer or create a new one.
 
-    If `organization.stripe_id` is `None` or if the existing customer
-    doesn't exist in stripe, a new customer is created.
+    If `organization.stripe_customer` is `None`, a new customer is created.
+    Not all models are migrated to use djstripe yet,
+    so we retrieve the customer from the stripe_id attribute if the model has one.
     """
     log.bind(
         organization_slug=organization.slug,
         stripe_customer=organization.stripe_id,
     )
-    if not organization.stripe_id:
-        log.info('No stripe customer found, creating one.')
-        return create_stripe_customer(organization)
-
-    try:
-        log.debug('Retrieving existing stripe customer.')
-        stripe_customer = stripe.Customer.retrieve(organization.stripe_id)
-        return stripe_customer
-    except InvalidRequestError as exc:
-        if exc.code == 'resource_missing':
-            log.info('Invalid stripe customer, creating new one.')
+    stripe_customer = organization.stripe_customer
+    if not stripe_customer:
+        stripe_customer = None
+        if organization.stripe_id:
+            try:
+                # TODO: Don't fully trust on djstripe yet,
+                # the customer may not be in our DB yet.
+                # pylint: disable=protected-access
+                stripe_customer = djstripe.Customer._get_or_retrieve(
+                    organization.stripe_id
+                )
+            except InvalidRequestError as exc:
+                if exc.code == "resource_missing":
+                    log.info("Invalid stripe customer, creating new one.")
+
+        if stripe_customer:
+            metadata = stripe_customer.metadata or {}
+            metadata.update(organization.get_stripe_metadata())
+            stripe.Customer.modify(
+                stripe_customer.id,
+                metadata=metadata,
+            )
+            organization.stripe_customer = stripe_customer
+            organization.save()
+        else:
+            log.info("No stripe customer found, creating one.")
             return create_stripe_customer(organization)
-        log.exception('Error while retrieving stripe customer.')
-        raise
+    return stripe_customer