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()

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()