Sqlalchemy: Truncate all tables

sqlalchemy-truncate_db.py

def truncate_db(engine):
    # delete all table data (but keep tables)
    # we do cleanup before test 'cause if previous test errored,
    # DB can contain dust
    meta = MetaData(bind=engine, reflect=True)
    con = engine.connect()
    trans = con.begin()
    con.execute('SET FOREIGN_KEY_CHECKS = 0;')
    for table in meta.sorted_tables:
        con.execute(table.delete())
    con.execute('SET FOREIGN_KEY_CHECKS = 1;')
    trans.commit()

Might be worth noting down that FOREIGN_KEY_CHECKS options is supported MySQL, though other databases commonly used with SQLAlchemy might not support it (e.g., PostgreSQL).

For Postgres, you can temporarily disable triggers, something like

def truncate_db(engine):
    # delete all table data (but keep tables)
    # we do cleanup before test 'cause if previous test errored,
    # DB can contain dust
    meta = MetaData(bind=engine, reflect=True)
    con = engine.connect()
    trans = con.begin()
    for table in meta.sorted_tables:
        con.execute(f'ALTER TABLE "{table.name}" DISABLE TRIGGER ALL;')
        con.execute(table.delete())
        con.execute(f'ALTER TABLE "{table.name}" ENABLE TRIGGER ALL;')
    trans.commit()

Surely there are ways to replicate the same functionality in other databases. For PostgreSQL specifically one might want to opt for using SET session_replication_role (https://www.postgresql.org/docs/12/runtime-config-client.html#GUC-SESSION-REPLICATION-ROLE, or https://stackoverflow.com/questions/3942258/how-do-i-temporarily-disable-triggers-in-postgresql).

@sloria replying to your snippet specifically as your code is effectively different from the original gist. Note you disable the triggers (and foreign key checks) only on a particular table, then deleting its content. In case a value from the table which is being deleted from is referenced as a foreign key constraint from any other table, the operation might fail. (Obviously this would be depending on ON DELETE option set on the other table but still it's error prone especially in a scenarios where the database model is managed by more people.)

It may be enough to disable a foreign key checks just for the current session:
con.execute('SET SESSION FOREIGN_KEY_CHECKS = ON')

from logging.config import fileConfig
from sqlalchemy import engine_from_config
from sqlalchemy import pool

from alembic import config

from sqlalchemy import MetaData, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base


config = config.Config("alembic.ini")

if config.config_file_name is not None:
    fileConfig(config.config_file_name)


target_metadata = Base.metadata
url = config.get_main_option("sqlalchemy.url")


def truncate():
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )
    
    with connectable.connect() as connection:
        Session = sessionmaker(bind=connection)
        session = Session() 
        session.execute(text('SET FOREIGN_KEY_CHECKS = 0;'))
        for table in target_metadata.sorted_tables:
            session.execute(table.delete())
        session.execute(text('SET FOREIGN_KEY_CHECKS = 1;'))
        session.commit()
        
if __name__=="__main__":
    truncate()