diff --git a/referral/models.py b/referral/models.py index a597eec..8f89ef6 100644 --- a/referral/models.py +++ b/referral/models.py @@ -15,7 +15,7 @@ from django.utils.translation import ugettext_lazy as _ # Create your models here. from registration.signals import user_activated -from core.models import AbstractStatusModel, STATUS_ACTIVE +from core.models import AbstractStatusModel, STATUS_ACTIVE, STATUS_CHOICES from referral.mixins import StatsFormatterMixin from referral.utils import get_cookie @@ -23,20 +23,22 @@ logger = logging.getLogger(__name__) # --------------------------------- Referral status list ----------------------------# -REFERRAL_DEFAULT_STATUS = STATUS_ACTIVE +REFERRAL_DEFAULT_CHOICE = STATUS_ACTIVE # @TODO: translate into english and use translation class Referral(AbstractStatusModel): CODE_LENGTH = 8 - user = models.OneToOneField(get_user_model(),verbose_name=_('пользователь'), on_delete=models.CASCADE, primary_key=True - ) + referral_owner = models.OneToOneField(get_user_model(), verbose_name=_('пользователь'), on_delete=models.CASCADE, + primary_key=True + ) code = models.CharField(_('код'), max_length=255) + status = models.SmallIntegerField(_('статус'), default=REFERRAL_DEFAULT_CHOICE, choices=STATUS_CHOICES) @classmethod def create(cls, user, code): - referral = cls(user=user, code=code) + referral = cls(referral_owner=user, code=code) return referral @staticmethod @@ -93,8 +95,8 @@ def create_user_referral(sender, instance, created, **kwargs): @receiver(post_save, sender=Referral) def init_referral_stuff(sender, instance, created, *args, **kwargs): if created: - if PartnerStats.objects.filter(user_id=instance.user_id).count() == 0: - PartnerStats.objects.create(user=instance.user).save() + if PartnerStats.objects.filter(user=instance.referral_owner).count() == 0: + PartnerStats.objects.create(user=instance.referral_owner).save() if ReferralStats.active.filter(referral=instance).count() == 0: ReferralStats.active.create(referral=instance).save() @@ -109,8 +111,8 @@ def update_ref_stats(sender, user, request, **kwargs): referral.referralstats.save() referral.user.partnerstats.total_regs += 1 referral.user.partnerstats.save() - user.referral = referral - user.referral_user = referral.user + user.referral_code = referral + user.referral_user = referral.referral_owner user.save() else: logger.warning("Missing referral code in database: " + referral_code)