30 lines
1.1 KiB
Python
30 lines
1.1 KiB
Python
import random
|
|
import string
|
|
from sqlmodel import select, Session, Field
|
|
from sqlmodel import SQLModel
|
|
from typing import Optional
|
|
from pydantic import validator
|
|
from database.db import get_session
|
|
def generate_unique_code(model: SQLModel, s: Session, field_name='id_code', length: int = 6):
|
|
if getattr(model, field_name, None) is None:
|
|
raise KeyError("Invalid field name")
|
|
while True:
|
|
code = ''.join(random.choices(string.ascii_uppercase, k=length))
|
|
is_unique = s.exec(select(model).where(
|
|
getattr(model, field_name) == code)).first() == None
|
|
if is_unique:
|
|
break
|
|
return code
|
|
|
|
class SQLModelWithIdCode(SQLModel):
|
|
id: Optional[int] = Field(default=None, primary_key=True)
|
|
id_code: Optional[str] = Field(default=None, unique=True, index=True)
|
|
|
|
@validator('id_code', always=True)
|
|
def generate_id_code(cls, value, values):
|
|
if value is None:
|
|
sessions = get_session()
|
|
session = next(sessions)
|
|
code = generate_unique_code(cls, s=session)
|
|
return code
|
|
return value |