Some functions to see the execution plan of a Postgres query emitted by SQLAlchemy

Explain SQLAlchemy Postgres queries

# Source: https://github.com/sqlalchemy/sqlalchemy/wiki/Query-Plan-SQL-construct
# This adds the last function, to print the query plan
# Caveats: stmt has to be built using sqlalchemy.select(...). If you use session.query(...) it'll fail.
# This is Postgres-only
# Guide to migrate to SQLAlchemy 2.0-style (from session.query() to select(...) ): https://docs.sqlalchemy.org/en/14/changelog/migration_20.html#migration-orm-usage


from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ClauseElement
from sqlalchemy.sql.expression import Executable



class Explain(Executable, ClauseElement):
    def __init__(self, stmt, analyze=False):
        self.statement = stmt
        self.analyze = analyze


@compiles(Explain, "postgresql")
def pg_explain(element, compiler, **kw):
    text = "EXPLAIN "
    if element.analyze:
        text += "ANALYZE "
    text += compiler.process(element.statement, **kw)

    return text


def explain_query(db: Session, stmt: Select, analyze=False):
    for el in db.execute(explain(stmt, analyze=analyze)).fetchall():
        print(el[0])