47 lines
No EOL
1.9 KiB
Python
47 lines
No EOL
1.9 KiB
Python
from django.contrib.auth.models import AbstractUser
|
|
from django.db import models
|
|
from django.core.exceptions import ValidationError
|
|
|
|
class CustomUser(AbstractUser):
|
|
default_friend_code = models.ForeignKey("FriendCode", on_delete=models.SET_NULL, null=True, blank=True)
|
|
|
|
def __str__(self):
|
|
return self.email
|
|
|
|
def set_default_friend_code(self, friend_code):
|
|
"""Set a friend code as default if it belongs to the user."""
|
|
if friend_code.user != self:
|
|
raise ValidationError("Friend code does not belong to this user.")
|
|
self.default_friend_code = friend_code
|
|
self.save(update_fields=["default_friend_code"])
|
|
|
|
def remove_default_friend_code(self, friend_code):
|
|
"""
|
|
If the given friend code is the current default,
|
|
assign another of the user's friend codes (if any) as default.
|
|
"""
|
|
if self.default_friend_code == friend_code:
|
|
other_codes = self.friend_codes.exclude(pk=friend_code.pk)
|
|
self.default_friend_code = other_codes.first() if other_codes.exists() else None
|
|
self.save(update_fields=["default_friend_code"])
|
|
|
|
class FriendCode(models.Model):
|
|
friend_code = models.CharField(max_length=19)
|
|
user = models.ForeignKey(CustomUser, on_delete=models.PROTECT, related_name='friend_codes')
|
|
created_at = models.DateTimeField(auto_now_add=True)
|
|
updated_at = models.DateTimeField(auto_now=True)
|
|
|
|
def save(self, *args, **kwargs):
|
|
"""
|
|
When a new friend code is saved,
|
|
if the user has no default friend code yet,
|
|
automatically set this as the default.
|
|
"""
|
|
is_new = self.pk is None
|
|
super().save(*args, **kwargs)
|
|
if is_new and not self.user.default_friend_code:
|
|
self.user.default_friend_code = self
|
|
self.user.save(update_fields=["default_friend_code"])
|
|
|
|
def __str__(self):
|
|
return self.friend_code |