Integrate Lago with Stripe

core/models.py

@@ -1,8 +1,8 @@
 import uuid
 from datetime import timedelta
 
-# import stripe
-# from django.conf import settings
+import stripe
+from django.conf import settings
 from django.contrib.auth.models import AbstractUser
 from django.db import models
 
@@ -11,6 +11,7 @@ from core.exchanges.fake import FakeExchange
 from core.exchanges.oanda import OANDAExchange
 
 # from core.lib.customers import get_or_create, update_customer_fields
+from core.lib import billing
 from core.util import logs
 
 log = logs.get_logger(__name__)
@@ -93,21 +94,36 @@ ADJUST_CLOSE_NOTIFY_CHOICES = (
 
 class User(AbstractUser):
     # Stripe customer ID
-    # unique_id = models.UUIDField(
-    #     default=uuid.uuid4,
-    # )
-    # stripe_id = models.CharField(max_length=255, null=True, blank=True)
+    stripe_id = models.CharField(max_length=255, null=True, blank=True)
+    customer_id = models.UUIDField(default=uuid.uuid4, null=True, blank=True)
     billing_provider_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 delete(self, *args, **kwargs):
+        if settings.BILLING_ENABLED:
+            if self.stripe_id:
+                stripe.Customer.delete(self.stripe_id)
+                log.info(f"Deleted Stripe customer {self.stripe_id}")
+        super().delete(*args, **kwargs)
 
     # Override save to update attributes in Lago
     def save(self, *args, **kwargs):
+        if self.customer_id is None:
+            self.customer_id = uuid.uuid4()
+
+        if settings.BILLING_ENABLED:
+            if not self.stripe_id:  # stripe ID not stored
+                self.stripe_id = billing.get_or_create(
+                    self.email, self.first_name, self.last_name
+                )
+
+            billing_id = billing.create_or_update_customer(self)
+            self.billing_provider_id = billing_id
+
+            billing.update_customer_fields(self)
+
         super().save(*args, **kwargs)
 
     def get_notification_settings(self):