Как гарантировать, что вставки принимают только значения, определенные в Enum?

Я пытаюсь создать простое приложение для записи платежей, в котором будут храниться сумма и тип платежа. Для типа оплаты у меня есть «наличные» и «кредит», которые определены в режиме:

from enum import Enum

class Mode(Enum):
    cash: str
    upi: str

Затем я определяю создаваемую таблицу с ее полем как:

from sqlalchemy import Column, Integer, Float, Enum
from utils.database import Base
from utils.paytype import Mode


class Records(Base):
    __tablename__ = "records"

    id = Column(Integer, primary_key=True, index=True)
    amount = Column(Float, nullable=False)
    payment_type = Column(Enum(Mode), nullable=False)

Когда я пытаюсь добавить запись в таблицу, я могу предоставить любую строку в качестве входных данных для payment_type, что не было идеей. Как гарантировать, что записи принимают только один из двух определенных?

🤔 А знаете ли вы, что...
В Python можно легко работать с базами данных, такими как SQLite и MySQL.


1
56
2

Ответы:

Вы на самом деле не можете. В sqlite нет типа Enum: https://www.sqlite.org/datatype3.html

Что вы можете сделать, так это использовать pydantic, чтобы проверить входные данные и в конечном итоге поднять ValidationError

from sqlalchemy import Column, Integer, Float, String
from pydantic import BaseModel
from utils.database import Base
from typing import Literal

class Records(Base):
    __tablename__ = "records"

    id = Column(Integer, primary_key=True, index=True)
    amount = Column(Float, nullable=False)
    payment_type = Column(String, nullable=False)

class RecordCreate(BaseModel):
    amount: float
    payment_type: Literal["cash", "credit"]

Сюда

record = RecordCreate(amount=3.14, payment_type = "credit") # this is OK
record = RecordCreate(amount=3.14, payment_type = "foo")    # this raises a ValidationError

Затем вам следует отправить в базу данных сериализованный объект RecordCreate, т.е.

**record.model_dump()  # with Pydantic v2
**record.dict()        # with Pydantic v1

и используйте Records для чтения данных из БД.


Решено

Ваше определение перечисления неверно — оно использует подсказки типов, которые фактически не создают никаких членов. Что вам нужно:

class Mode(Enum):
    cash = "cash"
    upi = "upi"

Чтобы сделать его доступным для использования в sqlite, добавьте метод адаптера в ваше перечисление:

class Mode(Enum):
    #
    cash = "cash"
    upi = "upi"
    #
    def __conform__(self, protocol):
        if protocol is sqlite3.PrepareProtocol:
            return self.name

Если у вас есть несколько перечислений, которые будут использоваться sqlite, создайте свой собственный базовый класс и используйте его:

import enum
import sqlite

class SqliteEnum(enum.Enum):
    def __conform__(self, protocol):
        if protocol is sqlite3.PrepareProtocol:
            return self.name

class Mode(SqliteEnum):
    cash = "cash"
    upi = "upi"