I wrote django referral system

Hello, for work I needed to write a referral system on django!
What the requirements:

user can sign up only with exists referral code
one referral code = one user
when user sign up we give 3 referral code for his friends What the technical requirements:
django
django rest framework

Also I used:

djoser

github repo:
https://github.com/mur4ik18/django-referral-system

First what I needed to write - create 2 app "referral_system" and "user".

My referral_system models:
models.py

from django.db import models
from django.conf import settings

class ReferralRelationship(models.Model):
    # who invite 
    employer = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        related_name='inviter',
        verbose_name="inviter",
        on_delete=models.CASCADE,
    )
    # who connected 
    employee = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        related_name='invited',
        verbose_name="invited",
        on_delete=models.CASCADE,
    )
    # referral code
    refer_token = models.ForeignKey(
        "ReferralCode",
        verbose_name="referral_code",
        on_delete=models.CASCADE,
    )
    def __str__(self) -> str:
        return f"{self.referrer}_{self.referred}"

class ReferralCode(models.Model):
    token = models.CharField(unique=True, max_length=150)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL, verbose_name="code_master", on_delete=models.CASCADE
    )
    def __str__(self) -> str:
        return f"{self.user}_{self.token}"

In the project I override standard user django model, cause we need to add new field and required field- referral_token
Also we need override standard UserManager cause we should add checking and writing tokens for new users and update referral relationships
user/models.py

from django.db import models
from django.contrib.auth.models import AbstractUser
from .manager import UserManager

class User(AbstractUser):
    referral_token = models.CharField(max_length=255)
    # redefining standard user manager
    objects = UserManager()
    # add new required field "referral_token"
    REQUIRED_FIELDS = ["referral_token", "email"]

In 'manager.py' I added new required argument - referral_token

user/manager.py

from referral_system.models import ReferralRelationship, ReferralCode
import secrets
from django.contrib.auth.models import BaseUserManager

class UserManager(BaseUserManager):
    def create_user(self, username, email, password=None, referral_token=None):
        # checking if user add all fields
        if not referral_token:
            raise ValueError("Please use your token!")
        if not email:
            raise ValueError("Users must have an email address!")
        # getting referral code and if code is not exist return error
        ref_code = ReferralCode.objects.filter(token=referral_token)
        if not ref_code:
            raise ValueError("Your token is not valid!")

        print(ref_code[0].user)

        # check code before using
        usages_token = ReferralRelationship.objects.filter(refer_token=ref_code[0])
        if not usages_token:
            # create user
            user = self.model(
                username=username, email=self.normalize_email(email), referral_token=referral_token
            )
            user.set_password(password)
            user.save(using=self._db)
            # create relationship for inviter and invited persone
            ReferralRelationship(
                employer=ref_code[0].user, employee=user, refer_token=ref_code[0]
            ).save()
            # create for new user 3 referall code (for frends)
            for i in range(3):
                self.create_reftoken(user)

        else:
            raise ValueError("This token is used!")

        return user

    def create_superuser(self, username, email, password=None, referral_token=None):
        # admin can use defunct referral code
        user = self.model(username=username, email=email, referral_token=referral_token)
        # set password and andmin things
        user.set_password(password)
        user.is_admin = True
        user.is_superuser = True
        user.is_staff = True
        user.save(using=self._db)
        # create 5 referral codes
        for i in range(5):
            self.create_reftoken(user)
        return user

    # method for creating tokens
    def create_reftoken(self, user):
        token = secrets.token_urlsafe(20)
        ReferralCode(token=token, user=user).save()

*If you know better how to write this piece of code, I would be grateful if you could tell me *

for i in range(5):
    self.create_reftoken(user)

def create_reftoken(self, user):
    token = secrets.token_urlsafe(20)
    ReferralCode(token=token, user=user).save()

Then
I created simple m2o (many to one) serializer

referral_system/serializers.py

import secrets
from rest_framework import serializers
from .models import ReferralCode, ReferralRelationship

class ReferralSerializer(serializers.ModelSerializer):
    class Meta:
        model = ReferralRelationship
        fields = [
            "employer",
            "employee", 
            "refer_token"]

class RefferCodeSerializer(serializers.ModelSerializer):
    refer_relations = ReferralSerializer(many=True, default="")
    class Meta:
        model = ReferralCode
        fields = [
            "token", 
            "user", 
            "refer_relations"]

and views.py
referral/views.py

from rest_framework import status
from rest_framework.response import Response
from rest_framework.generics import ListAPIView
from .models import ReferralCode
from .serializers import RefferCodeSerializer

class RefferCodeJsonListView(ListAPIView):
    def get(self, request):
        # getting querry set with referral codes for user who make request
        queryset = ReferralCode.objects.filter(user=request.user)
        data = RefferCodeSerializer(queryset, many=True)
        return Response(data.data, status=status.HTTP_200_OK)

How my system works
Admin signup - my system create 5 referral codes. Then admin send 5 codes for first users. Every new user get 3 referral codes for friends. The number of users will grow exponentially (boom grow)

I am open to criticism and if you wanna contrib this project you are welcome:
https://github.com/mur4ik18/django-referral-system

Blog

I wrote django referral system

Alex Kotov

Join Our Newsletter. No Spam, Only the good stuff.

Related