邀请码系统实现与用户注册流程优化#

1
# 统一时区处理，所有时间使用UTC+8本地时间
2
def get_local_time():
3
    return datetime.now(timezone(timedelta(hours=8)))
4

5
class M_InvitationCodes(UserBase):
6
  __tablename__ = "invitation_codes"
7

8
  id = Column(Integer, primary_key=True, index=True)
9
  user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)  # 关联创建者
10
  created_at = Column(DateTime, nullable=False, default=get_local_time)
11
  code = Column(String, nullable=False, unique=True, index=True)  # 邀请码唯一索引
12
  expire_at = Column(DateTime, nullable=False, index=True)  # 过期时间
13
  remaining_uses = Column(Integer, nullable=False, default=1, index=True)  # 剩余可用次数
14
  is_used = Column(Boolean, nullable=False, default=False, index=True)  # 是否已被使用过
15

16
  # 关联创建者，用于审计追踪
17
  user = relationship("M_Users", back_populates="invitation_codes")
18

19
  # 针对高频查询场景创建复合索引，优化查询性能
20
  __table_args__ = (
21
      Index('idx_expire_uses', 'expire_at', 'remaining_uses'),
22
      Index('idx_used_expire', 'is_used', 'expire_at'),
23
  )

1
def UseInvitationCode(db: Session, code: str):
2
    """
3
    使用邀请码，仅当邀请码有效、未过期、还有剩余次数时才会成功
4
    整个操作在数据库事务中执行，保证原子性
5
    """
6
    invitation_code = db.query(M_InvitationCodes).filter(
7
        M_InvitationCodes.code == code,
8
        M_InvitationCodes.remaining_uses > 0,
9
        M_InvitationCodes.expire_at > get_local_time()
10
    ).first()
11

12
    if not invitation_code:
13
        return False
14

15
    # 消耗次数，标记为已使用
16
    invitation_code.remaining_uses -= 1
17
    invitation_code.is_used = True
18
    db.commit()
19
    return True

1
def CleanupExpiredInvitationCodes(db: Session):
2
    # 清理规则：仅删除未被使用过的过期邀请码
3
    expired_codes = db.query(M_InvitationCodes).filter(
4
        M_InvitationCodes.is_used == False,
5
        M_InvitationCodes.expire_at < get_local_time()
6
    ).all()
7

8
    # 批量删除过期数据
9
    deleted_count = db.query(M_InvitationCodes).filter(
10
        M_InvitationCodes.id.in_([c.id for c in expired_codes])
11
    ).delete()
12
    db.commit()
13
    return deleted_count

1
def generate_invitation_code(length: int = 10) -> str:
2
    """
3
    使用secrets模块生成加密安全的随机邀请码
4
    相比普通随机数，不可预测，避免邀请码被恶意猜测
5
    """
6
    alphabet = string.ascii_letters + string.digits
7
    return ''.join(secrets.choice(alphabet) for _ in range(length))
8

9
# 频率限制：同一个IP每分钟最多创建5个邀请码
10
# 开发环境使用内存版，生产环境建议替换为Redis实现分布式限流
11
def check_rate_limit(ip: str) -> bool:
12
    # ... 限流逻辑

1
async def register_user(body: UserItem, db: Session):
2
    # 1. 第一步：验证邀请码有效性
3
    invitation_code = Db.GetInvitationCode(db, body.invitation_code)
4
    if not invitation_code or invitation_code.is_used or invitation_code.expire_at < get_local_time():
5
        return CommonOut(msg="邀请码不存在或已失效")
6

7
    # 2. 第二步：检查用户名是否重复
8
    if Db.GetUsers(db, username=body.username):
9
        return CommonOut(msg="用户名已存在")
10

11
    # 3. 第三步：使用邀请码并注册用户，整个过程在同一个事务中
12
    try:
13
        # 消耗邀请码
14
        use_success = Db.UseInvitationCode(db, body.invitation_code)
15
        # 注册用户
16
        new_user = Db.RegisterUser(db, **body.model_dump())
17
        return CommonOut(data=new_user)
18
    except Exception as e:
19
        # 异常时回滚事务，避免数据不一致
20
        db.rollback()
21
        return CommonOut(msg="服务器内部错误")

1
<el-form-item v-if="mode === 'register'" label="邀请码" prop="invitationCode">
2
  <el-input v-model="form.invitationCode" placeholder="请输入邀请码"></el-input>
3
</el-form-item>

1
invitationCode: [
2
  { required: true, message: '请输入邀请码', trigger: 'blur' },
3
  { min: 9, max: 11, message: '邀请码长度必须为9-11个字符', trigger: 'blur' }
4
]

1
// 仅管理员可见生成按钮
2
const isRootUser = computed(() => localStorage.getItem('role') === 'root');
3

4
// 生成邀请码的核心逻辑
5
const handleGenerateInvitationCode = async () => {
6
  const response = await api.post('/api/invitation-codes', invitationForm.value);
7
  if (response.code === 201) {
8
    generatedInvitationCode.value = response.data.code;
9
  }
10
};
11

12
// 一键复制邀请码
13
const copyInvitationCode = async () => {
14
  await navigator.clipboard.writeText(generatedInvitationCode.value);
15
  ElMessage.success('邀请码已复制');
16
};