Manage Stripe customers from the User model

This commit is contained in:
Mark Veidemanis 2022-07-21 13:48:14 +01:00
parent 4234669ad9
commit e0390f383c
Signed by: m
GPG Key ID: 5ACFCEED46C0904F
2 changed files with 83 additions and 1 deletions

58
core/lib/customers.py Normal file
View File

@ -0,0 +1,58 @@
import stripe
def expand_name(first_name, last_name):
"""
Convert two name variables into one.
Last name without a first name is ignored.
"""
name = None
if first_name:
name = first_name
# We only want to put the last name if we have a first name
if last_name:
name += f" {last_name}"
return name
def get_or_create(email, first_name, last_name):
"""
Get a customer ID from Stripe if one with the given email exists.
Create a customer if one does not.
Raise an exception if two or more customers matching the given email exist.
"""
# Let's see if we're just missing the ID
matching_customers = stripe.Customer.list(email=email, limit=2)
if len(matching_customers) == 2:
# Something is horribly wrong
print("Two customers match!")
raise Exception(f"Two customers found for email {email}")
elif len(matching_customers) == 1:
# We found a customer. Let's copy the ID
customer = matching_customers["data"][0]
customer_id = customer["id"]
return customer_id
else:
# We didn't find anything. Create the customer
# Create a name, since we have 2 variables which could be null
name = expand_name(first_name, last_name)
cast = {"email": email}
if name:
cast["name"] = name
customer = stripe.Customer.create(**cast)
return customer.id
def update_customer_fields(stripe_id, email=None, first_name=None, last_name=None):
"""
Update the customer fields in Stripe.
"""
print("Update customer fields called")
if email:
print("Email modified")
stripe.Customer.modify(stripe_id, email=email)
if first_name or last_name:
print("Name modified")
name = expand_name(first_name, last_name)
stripe.Customer.modify(stripe_id, name=name)

View File

@ -1,7 +1,7 @@
from django.contrib.auth.models import AbstractUser
from django.db import models
# Create your models here.
from core.lib.customers import get_or_create, update_customer_fields
class Plan(models.Model):
@ -24,6 +24,30 @@ class User(AbstractUser):
paid = models.BooleanField(null=True, blank=True)
plans = models.ManyToManyField(Plan, blank=True)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._original = self
def save(self, *args, **kwargs):
"""
Override the save function to create a Stripe customer.
"""
if not self.stripe_id: # stripe ID not stored
self.stripe_id = get_or_create(self.email, self.first_name, self.last_name)
to_update = {}
if self.email != self._original.email:
to_update["email"] = self.email
if self.first_name != self._original.first_name:
to_update["first_name"] = self.first_name
if self.last_name != self._original.last_name:
to_update["last_name"] = self.last_name
update_customer_fields(self.stripe_id, **to_update)
super().save(*args, **kwargs)
def has_plan(self, plan):
if not self.paid: # We can't have any plans if we haven't paid
return False