Add prefix_with clause to insert of a particular class

Question

How can I customize the prefix_with of each model class in SQLAlchemy such that each one can have a different insert statement?

I actually want the OR

Answer 1

The ORM doesn't have hooks into how it produces insert() constructs, so the best you can do here is intercept the insert() statements at the Table level, which is probably good enough here seeing as you want to do the "ignore" thing across the board for those tables, here's a recipe that uses a class decorator to make it generic. We are making use of the before_execute event here to rewrite certain insert() constructs:

from sqlalchemy import event
from sqlalchemy.engine import Engine
from sqlalchemy.sql import Insert

_ignore_tables = set()

@event.listens_for(Engine, "before_execute", retval=True)
def _ignore_insert(conn, element, multiparams, params):
    if isinstance(element, Insert) and \
        element.table.name in _ignore_tables:
        element = element.prefix_with("IGNORE")
    return element, multiparams, params

def ignore_inserts(cls):
    _ignore_tables.add(cls.__table__.name)
    return cls

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

    Base = declarative_base()

    class A(Base):
        __tablename__ = 'a'

        id = Column(Integer, primary_key=True)

    @ignore_inserts
    class B(Base):
        __tablename__ = 'b'

        id = Column(Integer, primary_key=True)

    @ignore_inserts
    class C(Base):
        __tablename__ = 'c'

        id = Column(Integer, primary_key=True)

    e = create_engine("mysql://scott:tiger@localhost/test", echo=True)
    Base.metadata.drop_all(e)
    Base.metadata.create_all(e)

    s = Session(e)

    s.add_all([A(), B(), C()])
    s.commit()

Having it across the board like that would make me nervous, here's a different version so that you can use a context manager to set up the rule for specific tables with a specific Session:

from sqlalchemy import event
from sqlalchemy.engine import Engine
from sqlalchemy.sql import Insert
from contextlib import contextmanager

@event.listens_for(Engine, "before_execute", retval=True)
def _ignore_insert(conn, element, multiparams, params):
    if isinstance(element, Insert) and \
        'ignore_tables' in conn.info and \
        element.table.name in conn.info['ignore_tables']:
        element = element.prefix_with("IGNORE")
    return element, multiparams, params


@contextmanager
def ignore_inserts(session, names):
    conn = session.connection()
    info = conn.info   # hold onto info so we can still
                       # get to it when the Connection is closed
    previous = info.get('ignore_tables', ())
    try:
        info['ignore_tables'] = set(names)
        yield
    finally:
        info['ignore_tables'] = previous

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

    Base = declarative_base()

    class A(Base):
        __tablename__ = 'a'

        id = Column(Integer, primary_key=True)

    class B(Base):
        __tablename__ = 'b'

        id = Column(Integer, primary_key=True)

    class C(Base):
        __tablename__ = 'c'

        id = Column(Integer, primary_key=True)

    e = create_engine("mysql://scott:tiger@localhost/test", echo=True)
    Base.metadata.drop_all(e)
    Base.metadata.create_all(e)

    s = Session(e)

    with ignore_inserts(s, ['b']):
        s.add_all([A(), B(), C()])
        s.commit()

    with ignore_inserts(s, ['a', 'c']):
        s.add_all([A(), B(), C()])
        s.commit()