import logging import stripe from django.contrib.auth.models import AbstractUser from django.db import models from core.lib.customers import get_or_create, update_customer_fields logger = logging.getLogger(__name__) class Plan(models.Model): name = models.CharField(max_length=255, unique=True) description = models.CharField(max_length=1024, null=True, blank=True) cost = models.IntegerField() product_id = models.CharField(max_length=255, unique=True, null=True, blank=True) image = models.CharField(max_length=1024, null=True, blank=True) def __str__(self): return f"{self.name} (£{self.cost})" class User(AbstractUser): # Stripe customer ID stripe_id = models.CharField(max_length=255, null=True, blank=True) subscription_id = models.CharField(max_length=255, null=True, blank=True) subscription_active = models.BooleanField(null=True, blank=True) last_payment = models.DateTimeField(null=True, blank=True) paid = models.BooleanField(null=True, blank=True) plans = models.ManyToManyField(Plan, blank=True) email = models.EmailField(unique=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 delete(self, *args, **kwargs): if self.stripe_id: stripe.Customer.delete(self.stripe_id) logger.info(f"Deleted Stripe customer {self.stripe_id}") super().delete(*args, **kwargs) def has_plan(self, plan): if not self.paid: # We can't have any plans if we haven't paid return False plan_list = [plan.name for plan in self.plans.all()] return plan in plan_list class Session(models.Model): email = models.EmailField() session = models.CharField(max_length=255)