import random
import string
from sqlmodel import select, Session
from sqlmodel import SQLModel


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