使SQLAlChemy错误更加用户友好和详细

2022-03-10 00:00:00 python sqlalchemy flask flask-sqlalchemy

问题描述

我有一个这样的模型:

class Company(db.Model):
    __tablename__ = "my_table"
    id = db.Column(db.Integer(), primary_key=True)
    name = db.Column(db.String(128), unique=True, nullable=False)
    slug = db.Column(db.String(128), unique=True, nullable=False)
如您所见,我使用的是Flask-SQLAlchemy的类和方法,但这不是重点, 假设我有一个执行以下行的视图:

c = Company("Test", "test")
try:
    db.session.add(c)
    db.session.commit()
    return "Added!"
except Exception as e:
    db.session.rollback()
    return f"{e}"

上面的代码创建Company的对象,尝试将其添加到数据库,在异常时回滚事务

问题出在这里,因为数据是硬编码的,所以它应该始终返回异常,SQLAlChemy将引发IntegrityError

IntegrityError太难看了,对用户毫无用处,例如:

(sqlite3.IntegrityError) UNIQUE constraint failed: my_table.name [SQL: 'INSERT INTO my_table (name, slug) VALUES (?, ?)'] [parameters: ('Test', 'tests')] (Background on this error at: http://sqlalche.me/e/gkpj)

我正在寻找一种使其美观和用户友好的方法,在此之前我使用了db.validates修饰器并在验证时检查重复数据,但我觉得这不太对劲

我最不需要的就是找出是哪个字段导致了问题,而不需要硬编码

DBAPI

SQLAlChemy包括一种机制,允许使用handle_error事件挂钩自定义推荐答案错误。我在OpenStackoslo.db中使用了此API,可以在以下文件中看到:https://github.com/openstack/oslo.db/blob/master/oslo_db/sqlalchemy/exc_filters.py。

由于Stackoverflow讨厌指向代码的链接,因此以下是基于上述链接方法的POC:

import collections
from sqlalchemy import event
from sqlalchemy import exc as sqla_exc
import re


class DuplicateKeyError(Exception):
    """Duplicate entry at unique column error."""

    def __init__(self, columns=None, inner_exception=None, value=None):
        self.columns = columns or []
        self.value = value
        self.inner_exception = inner_exception

    def __str__(self):
        return "Duplicate key for columns %s" % (
            self.columns,
        )


_registry = collections.defaultdict(lambda: collections.defaultdict(list))


def filters(ame, exception_type, regex):
    """Mark a function as receiving a filtered exception."""

    def _receive(fn):
        _registry[ame][exception_type].extend(
            (fn, re.compile(reg))
            for reg in ((regex,) if not isinstance(regex, tuple) else regex)
        )
        return fn

    return _receive


# each @filters() lists a database name, a SQLAlchemy exception to catch,
# and a list of regular expressions that will be matched.  If all the
# conditions match, the handler is called which then raises a nicer
# error message.

@filters(
    "sqlite",
    sqla_exc.IntegrityError,
    (
        r"^.*columns?(?P<columns>[^)]+)(is|are)s+nots+unique$",
        r"^.*UNIQUEs+constraints+failed:s+(?P<columns>.+)$",
        r"^.*PRIMARYs+KEYs+musts+bes+unique.*$",
    ),
)
def _sqlite_dupe_key_error(integrity_error, match, engine_name, is_disconnect):
    columns = []
    try:
        columns = match.group("columns")
        columns = [c.split(".")[-1] for c in columns.strip().split(", ")]
    except IndexError:
        pass

    raise DuplicateKeyError(columns, integrity_error)


@filters(
    "mysql",
    sqla_exc.IntegrityError,
    r"^.*1062.*Duplicate entry '(?P<value>.*)'"
    r" for key '(?P<columns>[^']+)'.*$",
)
@filters(
    "postgresql",
    sqla_exc.IntegrityError,
    (
        r'^.*duplicates+key.*"(?P<columns>[^"]+)"s*
.*'
        r"Keys+((?P<key>.*))=((?P<value>.*))s+alreadys+exists.*$",
        r"^.*duplicates+key.*"(?P<columns>[^"]+)"s*
.*$",
    ),
)
def _default_dupe_key_error(
    integrity_error, match, engine_name, is_disconnect
):
    columns = match.group("columns")
    uniqbase = "uniq_"
    if not columns.startswith(uniqbase):
        if engine_name == "postgresql":
            columns = [columns[columns.index("_") + 1 : columns.rindex("_")]]
        else:
            columns = [columns]
    else:
        columns = columns[len(uniqbase) :].split("0")[1:]

    value = match.groupdict().get("value")

    raise DuplicateKeyError(columns, integrity_error, value)


def handler(context):
    """Iterate through available filters and invoke those which match.
    The first one which raises wins.
    """

    def _dialect_registries(engine):
        if engine.dialect.name in _registry:
            yield _registry[engine.dialect.name]
        if "*" in _registry:
            yield _registry["*"]

    for per_dialect in _dialect_registries(context.engine):
        for exc in (context.sqlalchemy_exception, context.original_exception):
            for super_ in exc.__class__.__mro__:
                if super_ in per_dialect:
                    regexp_reg = per_dialect[super_]
                    for fn, regexp in regexp_reg:
                        match = regexp.match(exc.args[0])
                        if match:
                            fn(
                                exc,
                                match,
                                context.engine.dialect.name,
                                context.is_disconnect,
                            )


if __name__ == '__main__':
    from sqlalchemy import Column, Integer, String, create_engine
    from sqlalchemy.orm import Session
    from sqlalchemy.ext.declarative import declarative_base

    Base = declarative_base()


    class Company(Base):
        __tablename__ = "my_table"
        id = Column(Integer(), primary_key=True)
        name = Column(String(128), unique=True, nullable=False)
        slug = Column(String(128), unique=True, nullable=False)

        def __init__(self, name, slug):
            self.name = name
            self.slug = slug

    e = create_engine("sqlite://", echo=True)
    Base.metadata.create_all(e)
    event.listen(e, "handle_error", handler)

    s = Session(e)
    s.add(Company("Test", "test"))
    s.commit()


    s.add(Company("Test", "test"))
    s.commit()

运行它,我们看到:

2019-03-13 09:44:51,701 INFO sqlalchemy.engine.base.Engine INSERT INTO my_table (name, slug) VALUES (?, ?)
2019-03-13 09:44:51,701 INFO sqlalchemy.engine.base.Engine ('Test', 'test')

2019-03-13 09:44:53,387 INFO sqlalchemy.engine.base.Engine ROLLBACK
Traceback (most recent call last):
# ...
sqlite3.IntegrityError: UNIQUE constraint failed: my_table.slug

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
 # ...
__main__.DuplicateKeyError: Duplicate key for columns ['slug']

相关文章