Commit 43af5383 authored by ibuler's avatar ibuler

Add model user method: get_token set_token

parent ffed4617
...@@ -3,11 +3,16 @@ ...@@ -3,11 +3,16 @@
from __future__ import unicode_literals from __future__ import unicode_literals
import datetime import datetime
from django.conf import settings
from django.contrib.auth.hashers import make_password from django.contrib.auth.hashers import make_password
from django.utils import timezone from django.utils import timezone
from django.db import models from django.db import models
from django.contrib.auth.models import AbstractUser, Permission from django.contrib.auth.models import AbstractUser, Permission
from django.db import OperationalError from django.db.models.signals import post_save
from django.dispatch import receiver
from rest_framework.authtoken.models import Token
from django.db import IntegrityError
# class Role(models.Model): # class Role(models.Model):
...@@ -170,10 +175,28 @@ class User(AbstractUser): ...@@ -170,10 +175,28 @@ class User(AbstractUser):
self.groups.add(group) self.groups.add(group)
# super(User, self).save(*args, **kwargs) # super(User, self).save(*args, **kwargs)
@property
def token(self):
return self.get_token()
def get_token(self):
try:
token = Token.objects.get(user=self)
return token.key
except Token.DoesNotExist:
return ''
def set_token(self):
try:
return Token.objects.create(user=self)
except IntegrityError:
Token.objects.filter(user=self).delete()
return Token.objects.create(user=self)
class Meta: class Meta:
db_table = 'user' db_table = 'user'
#: Use this method #: Use this method initial user
@classmethod @classmethod
def initial(cls): def initial(cls):
user = cls(username='admin', user = cls(username='admin',
...@@ -222,3 +245,10 @@ def generate_fake(): ...@@ -222,3 +245,10 @@ def generate_fake():
for model in (UserGroup, User): for model in (UserGroup, User):
if hasattr(model, 'generate_fake'): if hasattr(model, 'generate_fake'):
model.generate_fake() model.generate_fake()
@receiver(post_save, sender=settings.AUTH_USER_MODEL)
def create_auth_token(sender, instance=None, created=False, **kwargs):
if created:
Token.objects.create(user=instance)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment