From 3ce3a05c7b1ee793a0924632de260e5dc3b12e27 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Sat, 2 May 2020 18:08:05 +0200
Subject: [PATCH] Add referral name

---
 app/models.py |  5 ++++-
 server.py     | 13 +++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/app/models.py b/app/models.py
index 2cedaf81..0f68b97e 100644
--- a/app/models.py
+++ b/app/models.py
@@ -158,7 +158,9 @@ class User(db.Model, ModelMixin, UserMixin):
         db.Boolean, default=True, nullable=False, server_default="1"
     )
 
-    referral_id = db.Column(db.ForeignKey("referral.id"), nullable=True, default=None)
+    referral_id = db.Column(
+        db.ForeignKey("referral.id", ondelete="SET NULL"), nullable=True, default=None
+    )
 
     referral = db.relationship("Referral", foreign_keys=[referral_id])
 
@@ -1153,6 +1155,7 @@ class Referral(db.Model, ModelMixin):
     """Referral code so user can invite others"""
 
     user_id = db.Column(db.ForeignKey(User.id, ondelete="cascade"), nullable=False)
+    name = db.Column(db.String(512), nullable=True, default=None)
 
     code = db.Column(db.String(128), unique=True, nullable=False)
 
diff --git a/server.py b/server.py
index 99ae24ba..3f1ad13f 100644
--- a/server.py
+++ b/server.py
@@ -51,6 +51,7 @@ from app.models import (
     DeletedAlias,
     Contact,
     EmailLog,
+    Referral,
 )
 from app.monitor.base import monitor_bp
 from app.oauth.base import oauth_bp
@@ -216,6 +217,18 @@ def fake_data():
     DeletedAlias.create(user_id=user.id, email="d2@ab.cd")
     db.session.commit()
 
+    referral = Referral.create(user_id=user.id, code="REFCODE", name="First referral")
+    db.session.commit()
+
+    User.create(
+        email="winston@continental.com",
+        name="Winston",
+        password="password",
+        activated=True,
+        referral_id=referral.id,
+    )
+    db.session.commit()
+
 
 @login_manager.user_loader
 def load_user(user_id):