feat: implement robust referral system and fix User relations

Models:
- Refactored InviteUser: changed  to OneToOneField and enforced unique=True on .
- Implemented auto-generation of unique 6-character invite codes in InviteUser.save().
- Added post_save signal to auto-create InviteUser profiles for new users.
- Fixed  in User model: changed on_delete to SET_NULL, removed invalid default values, and fixed lazy loading references.
- Removed incompatible  override in UserManager.

Serializers:
- Updated UserRegistrationSerializer and GoogleRegister to handle invite codes via SlugRelatedField.
- Wrapped user creation methods in @transaction.atomic for data integrity.
- Implemented race-condition-safe increment for  using F() expressions upon successful registration.

Author: hardusss

backend/NextVibeAPI/user/admin.py

@@ -1,5 +1,5 @@
 from django.contrib import admin
-from .models import User, HistorySearch, Notification, UserOnlineSession
+from .models import User, HistorySearch, Notification, UserOnlineSession, InviteUser
 from django.contrib import admin
 
 @admin.register(User)
@@ -16,3 +16,4 @@ def get_queryset(self, request):
 admin.site.register(HistorySearch)
 admin.site.register(Notification)
 admin.site.register(UserOnlineSession)
+admin.site.register(InviteUser)

backend/NextVibeAPI/user/models.py

@@ -2,6 +2,9 @@
 from django.db import models
 from django.utils import timezone
 from datetime import timedelta
+from django.utils.crypto import get_random_string
+from django.db.models.signals import post_save
+from django.dispatch import receiver
 
 class UserQuerySet(models.QuerySet):
     def visible(self):
@@ -26,9 +29,6 @@ def create_superuser(self, email, username, password=None, **extra_fields):
 
         return self.create_user(email, username, password, **extra_fields)
 
-    def get_by_natural_key(self, email):
-        return self.get(email=email)
-
 
 class User(AbstractBaseUser):
     user_id = models.AutoField(primary_key=True, unique=True)
@@ -58,6 +58,7 @@ class User(AbstractBaseUser):
     last_activity = models.DateTimeField(default=timezone.now)
     wallet_address = models.CharField(max_length=50, blank=True, null=True)
     expo_push_token = models.CharField(max_length=100, blank=True, null=True)
+    from_invite_code = models.ForeignKey("InviteUser", on_delete=models.SET_NULL, null=True, blank=True, related_name="invited_users")
 
     USERNAME_FIELD = 'email'
     REQUIRED_FIELDS = ['username']
@@ -73,6 +74,26 @@ def has_perm(self, perm, obj=None):
     def has_module_perms(self, app_label):
         return self.is_superuser
 
+class InviteUser(models.Model):
+    owner = models.OneToOneField(User, on_delete=models.CASCADE)
+    invite_code = models.CharField(max_length=6, unique=True)
+    invited_count = models.IntegerField(default=0, null=True, blank=True)
+    
+    def save(self, *args, **kwargs):
+        # If code not created
+        if not self.invite_code:
+            # Generate str with 6 length
+            new_code = get_random_string(length=6)
+
+            while InviteUser.objects.filter(invite_code=new_code).exists():
+                new_code = get_random_string(length=6)
+                
+            self.invite_code = new_code
+            
+        super().save(*args, **kwargs)
+
+    def __str__(self):
+        return f"Invite profile of {self.owner} ({self.invite_code}) have {self.invited_count} refs"
 
 class HistorySearch(models.Model):
     user = models.ForeignKey(User, on_delete=models.CASCADE, related_name="user_history")
@@ -139,8 +160,10 @@ def session_duration(self) -> timedelta:
         return end_time - self.connected_at
 
 
-
-
-
-
-
+@receiver(post_save, sender=User)
+def create_user_invite_profile(sender, instance, created, **kwargs):
+    """
+    Thats signal works every time, when save User model
+    """
+    if created:
+        InviteUser.objects.get_or_create(owner=instance)

backend/NextVibeAPI/user/serializers_pac/google_registration.py

@@ -4,17 +4,26 @@
 import requests
 from django.core.files.base import ContentFile
 import uuid
+from django.db import transaction
+from django.db.models import F
+from users.models import InviteUser 
 
 User = get_user_model()
 
-
 class GoogleRegister(serializers.ModelSerializer):
     token = serializers.SerializerMethodField()
     avatar_url = serializers.URLField(write_only=True, required=False)
+    
+    from_invite_code = serializers.SlugRelatedField(
+        slug_field='invite_code', 
+        queryset=InviteUser.objects.all(),
+        required=False,
+        allow_null=True
+    )
 
     class Meta:
         model = User
-        fields = ("user_id", "email", "username", "token", "avatar", "avatar_url")
+        fields = ("user_id", "email", "username", "token", "avatar", "avatar_url", "from_invite_code")
 
     def get_token(self, obj):
         refresh = RefreshToken.for_user(obj)
@@ -23,18 +32,32 @@ def get_token(self, obj):
             "access": str(refresh.access_token),
         }
 
+    @transaction.atomic
     def create(self, validated_data):
         avatar_url = validated_data.pop("avatar_url", None)  
-        user, created = User.objects.get_or_create(email=validated_data["email"], defaults={"username": validated_data["username"]})
-
-        if created and avatar_url:
-            try:
-                response = requests.get(avatar_url, timeout=5)
-                response.raise_for_status()
-                file_name = f"user_{user.user_id}_{uuid.uuid4().hex}.jpg"
-                user.avatar.save(file_name, ContentFile(response.content), save=True)
-            except requests.RequestException:
-                pass
-
-        return user
-    
+        invite_obj = validated_data.pop("from_invite_code", None)
+
+        defaults = {"username": validated_data["username"]}
+        if invite_obj:
+            defaults["from_invite_code"] = invite_obj
+
+        user, created = User.objects.get_or_create(
+            email=validated_data["email"], 
+            defaults=defaults
+        )
+
+        if created:
+            if avatar_url:
+                try:
+                    response = requests.get(avatar_url, timeout=5)
+                    response.raise_for_status()
+                    file_name = f"user_{user.user_id}_{uuid.uuid4().hex}.jpg"
+                    user.avatar.save(file_name, ContentFile(response.content), save=True)
+                except requests.RequestException:
+                    pass
+
+            if invite_obj:
+                invite_obj.invited_count = F('invited_count') + 1
+                invite_obj.save(update_fields=['invited_count'])
+
+        return user

backend/NextVibeAPI/user/serializers_pac/registration.py

@@ -1,29 +1,46 @@
 from rest_framework import serializers
 from django.contrib.auth import get_user_model
 from rest_framework_simplejwt.tokens import RefreshToken
-
-
+from user.models import InviteUser
+from django.db import transaction
 
 class UserRegistrationSerializer(serializers.ModelSerializer):
     avatar_url = serializers.SerializerMethodField()
     password = serializers.CharField(write_only=True)
     token = serializers.SerializerMethodField()
+    
+    from_invite_code = serializers.SlugRelatedField(
+        slug_field='invite_code', 
+        queryset=InviteUser.objects.all(),
+        required=True, 
+        allow_null=False
+    )
     class Meta:
-        fields = ("user_id", "email", "username", "password", "token", "avatar_url")
+        fields = ("user_id", "email", "username", "password", "token", "avatar_url", "from_invite_code")
         model = get_user_model()
         extra_kwargs = {
             'password': {'write_only': True},
             'email': {'required': True},
             'username': {'required': True},
+            'from_invite_code': {'required': True},
         }
-        
+    
+    @transaction.atomic
     def create(self, validated_data):
+        invite_obj = validated_data.get('from_invite_code')
+        
         user = get_user_model()(
             email=validated_data['email'],
             username=validated_data['username'],
+            from_invite_code=invite_obj
         )
         user.set_password(validated_data['password'])  # Hash password
         user.save()
+        
+        if invite_obj:
+            invite_obj.invited_count += 1
+            invite_obj.save(update_fields=['invited_count'])
+            
         return user
 
     def get_token(self, obj):