generateur_v3/backend/api/services/database.py

30 lines
1.1 KiB
Python
Raw Normal View History

2022-09-16 21:50:55 +02:00
import random
import string
2022-10-10 01:34:38 +02:00
from sqlmodel import select, Session, Field
2022-09-16 21:50:55 +02:00
from sqlmodel import SQLModel
2022-10-10 01:34:38 +02:00
from typing import Optional
from pydantic import validator
from database.db import get_session
2022-09-18 22:43:04 +02:00
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")
2022-09-16 21:50:55 +02:00
while True:
code = ''.join(random.choices(string.ascii_uppercase, k=length))
is_unique = s.exec(select(model).where(
2022-09-18 22:43:04 +02:00
getattr(model, field_name) == code)).first() == None
2022-09-16 21:50:55 +02:00
if is_unique:
break
return code
2022-10-10 01:34:38 +02:00
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