from django.db import models from django.utils import timezone from django.contrib.auth.base_user import AbstractBaseUser from django.contrib.auth.models import PermissionsMixin, UserManager from django.contrib.auth.hashers import make_password from django.utils.translation import gettext_lazy as _ INDUSTRY_UNKNOWN = "Unknown" class LabsUserManager(UserManager): def _create_user(self, email, password, **extra_fields): """ Create and save a user with the given email, and password. """ if not email: raise ValueError("The given email must be set") email = self.normalize_email(email) user = self.model(email=email, **extra_fields) user.password = make_password(password) user.save(using=self._db) return user def create_user(self, email=None, password=None, **extra_fields): extra_fields.setdefault("is_staff", False) extra_fields.setdefault("is_superuser", False) return self._create_user(email, password, **extra_fields) def create_superuser(self, email=None, password=None, **extra_fields): extra_fields.setdefault("is_staff", True) extra_fields.setdefault("is_superuser", True) if extra_fields.get("is_staff") is not True: raise ValueError("Superuser must have is_staff=True.") if extra_fields.get("is_superuser") is not True: raise ValueError("Superuser must have is_superuser=True.") return self._create_user(email, password, **extra_fields) class User(AbstractBaseUser, PermissionsMixin): email = models.EmailField(_("email address"), unique=True) is_staff = models.BooleanField( _("staff status"), default=False, help_text=_("Designates whether the user can log into this admin site."), ) is_active = models.BooleanField( _("active"), default=True, help_text=_( "Designates whether this user should be treated as active. " "Unselect this instead of deleting accounts." ), ) date_joined = models.DateTimeField(_("date joined"), default=timezone.now) updated_at = models.DateTimeField(auto_now=True) name = models.CharField(max_length=100, null=True, blank=True) first_name = models.CharField(max_length=100, null=True, blank=True) last_name = models.CharField(max_length=100, null=True, blank=True) language_code = models.CharField(max_length=2, default="en") last_lab_code = models.CharField(max_length=20, null=True, blank=True) last_region_name = models.CharField(max_length=40, null=True, blank=True) last_tidb_version = models.CharField(max_length=20, null=True, blank=True) organization = models.CharField(max_length=100, null=True, blank=True) job_role = models.CharField(max_length=60, null=True, blank=True) country = models.CharField(max_length=60, null=True, blank=True) industry = models.CharField( max_length=60, default=INDUSTRY_UNKNOWN, ) EMAIL_FIELD = "email" USERNAME_FIELD = "email" REQUIRED_FIELDS = [] objects = LabsUserManager() class Meta: db_table = "auth_user" def __str__(self): return f"User: {self.email}" def clean(self): super().clean() self.email = self.__class__.objects.normalize_email(self.email) def get_full_name(self): """ Return the first_name plus the last_name, with a space in between. """ return self.name def get_short_name(self): """Return the short name for the user.""" return self.name @property def is_pingcap_employee(self) -> bool: return self.email.strip().lower().endswith("@pingcap.com")