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) last_payment = models.DateTimeField(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): plan_list = [plan.name for plan in self.plans.all()] return plan in plan_list class Session(models.Model): user = models.ForeignKey(User, on_delete=models.CASCADE) request = models.CharField(max_length=255, null=True, blank=True) session = models.CharField(max_length=255, null=True, blank=True) subscription_id = models.CharField(max_length=255, null=True, blank=True) plan = models.ForeignKey(Plan, null=True, blank=True, on_delete=models.CASCADE) class ContentBlock(models.Model): user = models.ForeignKey(User, on_delete=models.PROTECT) position = models.IntegerField() title = models.CharField(max_length=255, null=True, blank=True) column1 = models.TextField(null=True, blank=True) column2 = models.TextField(null=True, blank=True) column3 = models.TextField(null=True, blank=True) image1 = models.CharField(max_length=255, null=True, blank=True) image2 = models.CharField(max_length=255, null=True, blank=True) image3 = models.CharField(max_length=255, null=True, blank=True) def save(self, *args, **kwargs): """ Override the save function to blank fields. """ if self.column1 == "NONE": self.column1 = None if self.column2 == "NONE": self.column2 = None if self.column3 == "NONE": self.column3 = None super().save(*args, **kwargs)