SQLAlchemy 多对多关系：唯一约束失败

Question

我在用SQLAlchemy定义一个多对多的关系，代码是这样的：

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, sessionmaker
from sqlalchemy import Column, Integer, String, ForeignKey, UniqueConstraint, Table, create_engine
from sqlalchemy.orm import relationship, registry


mapper_registry = registry()
Base = declarative_base()


bridge_category = Table(
    "bridge_category",
    Base.metadata,
    Column("video_id", ForeignKey("video.id"), primary_key=True),
    Column("category_id", ForeignKey("category.id"), primary_key=True),
    UniqueConstraint("video_id", "category_id"),
)
class BridgeCategory: pass
mapper_registry.map_imperatively(BridgeCategory, bridge_category)


class Video(Base):
    __tablename__ = 'video'

    id = Column(Integer, primary_key=True)
    title = Column(String)
    categories = relationship("Category", secondary=bridge_category, back_populates="videos")


class Category(Base):
    __tablename__ = 'category'

    id = Column(Integer, primary_key=True)
    text = Column(String, unique=True)
    videos = relationship("Video", secondary=bridge_category, back_populates="categories")


engine = create_engine('sqlite:///:memory:', echo=True)
Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)

with Session() as s:

    v1 = Video(title='A', categories=[Category(text='blue'), Category(text='red')])
    v2 = Video(title='B', categories=[Category(text='green'), Category(text='red')])
    v3 = Video(title='C', categories=[Category(text='grey'), Category(text='red')])
    videos = [v1, v2, v3]

    s.add_all(videos)
    s.commit()

因为在Category.text上有唯一性约束，所以我遇到了以下错误：

sqlalchemy.exc.IntegrityError: (sqlite3.IntegrityError) UNIQUE constraint failed: category.text
[SQL: INSERT INTO category (text) VALUES (?) RETURNING id]
[parameters: ('red',)]

我在想，处理这个问题的最好方法是什么。我的程序会处理很多视频对象，每个视频都有一个独特的分类对象列表。这些分类的文本重复问题出现在所有这些视频对象中。

我可以遍历所有视频和所有分类，形成一个分类集合，但这样做有点麻烦。而且我还得对我的视频对象的其他12个多对多关系做同样的事情，这样效率看起来真的不高。

有没有什么类似“插入忽略”的标志可以设置呢？我在网上找了很久，都没有找到关于这种情况的解决方案。

Answer 1

在SQLAlchemy的维护者的大力帮助下，我想出了一个通用的代码实现。这段代码几乎不需要任何配置，也不需要重复步骤，就可以处理一个包含多个多对多关系的单一SA模型对象。

from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import ForeignKey
from sqlalchemy import inspect
from sqlalchemy import Integer
from sqlalchemy import select
from sqlalchemy import String
from sqlalchemy import Table
from sqlalchemy import UniqueConstraint
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import registry
from sqlalchemy.orm import relationship
from sqlalchemy.orm import RelationshipDirection
from sqlalchemy.orm import Session
from sqlalchemy.orm import sessionmaker


mapper_registry = registry()
Base = declarative_base()


bridge_category = Table(
    "bridge_category",
    Base.metadata,
    Column("video_id", ForeignKey("video.id"), primary_key=True),
    Column("category_id", ForeignKey("category.id"), primary_key=True),
    UniqueConstraint("video_id", "category_id"),
)
class BridgeCategory: pass
mapper_registry.map_imperatively(BridgeCategory, bridge_category)


bridge_format = Table(
    "bridge_format",
    Base.metadata,
    Column("video_id", ForeignKey("video.id"), primary_key=True),
    Column("format_id", ForeignKey("format.id"), primary_key=True),
    UniqueConstraint("video_id", "format_id"),
)
class BridgeFormat: pass
mapper_registry.map_imperatively(BridgeFormat, bridge_format)


class Video(Base):
    __tablename__ = "video"

    id = Column(Integer, primary_key=True)
    title = Column(String)
    categories = relationship("Category", secondary=bridge_category, back_populates="videos")
    formats = relationship("Format", secondary=bridge_format, back_populates="videos")

class Category(Base):
    __tablename__ = "category"

    id = Column(Integer, primary_key=True)
    text = Column(String, unique=True)
    videos = relationship("Video", secondary=bridge_category, back_populates="categories")

class Format(Base):
    __tablename__ = "format"
    id = Column(Integer, primary_key=True, index=True)
    text = Column(String, unique=True)
    videos = relationship("Video", back_populates="formats", secondary=bridge_format)


def unique_robs(session_or_factory, main_obj, rob_unique_col):
    """Unique related objects"""

    def _unique_robs(session, robs, rob_name):
        if not robs:
            return robs
        
        rob_type = type(robs[0])

        with session.no_autoflush:
            local_existing_robs = session.info.get(rob_name, None)
            if local_existing_robs is None:
                session.info[rob_name] = local_existing_robs = {}

            unique_vals = []
            for r in robs:
                unique_val = getattr(r, rob_unique_col)
                if unique_val not in local_existing_robs:
                    unique_vals.append(unique_val)

            existing_categories = {}
            unique_col = getattr(rob_type, rob_unique_col)
            for r in session.scalars(select(rob_type).where(unique_col.in_(unique_vals))):
                existing_categories[getattr(r, rob_unique_col)] = r

            local_existing_robs.update(existing_categories)

            result = []
            for r in robs:
                if getattr(r, rob_unique_col) in local_existing_robs:
                    result.append(local_existing_robs[getattr(r, rob_unique_col)])
                    continue

                local_existing_robs[getattr(r, rob_unique_col)] = r
                result.append(r)

            return result

    @event.listens_for(session_or_factory, "before_attach", retval=True)
    def before_attach(session, obj):
        """Uniquifies all `main_obj` many-to-many relationships."""
        if isinstance(obj, main_obj):
            for r in inspect(obj).mapper.relationships:
                if r.direction.value == RelationshipDirection.MANYTOMANY.value:
                    rob_name = r.class_attribute.key
                    robs = getattr(obj, rob_name, None)
                    if isinstance(robs, list):
                        setattr(obj, rob_name, _unique_robs(session, robs, rob_name))


if __name__ == "__main__":
    engine = create_engine("sqlite:///test.db", echo=True)
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    unique_robs(Session, Video, 'text')

    v1 = Video(title="A", categories=[Category(text="blue"), Category(text="red")])
    v2 = Video(title="B", categories=[Category(text="green"), Category(text="red")], formats=[Format(text='h264')])
    v3 = Video(title="C", categories=[Category(text="grey"), Category(text="red")], formats=[Format(text='h264'), Format(text='vp9')])
    videos = [v1, v2, v3]

    with Session() as s:

        s.add_all(videos)
        s.commit()