SQLAlchemy PostgreSQL UNION with ORDER/LIMIT monkey patch fix

sqlalchemy_postgresql_union_fix.py

# Problem: sa.union_all(*qs).order_by('y')  will generate invalid SQL like this:
# SELECT ... ORDER BY x LIMIT 10 UNION ALL SELECT ... ORDER BY x LIMIT 10 ORDER BY y
# Order in inner queries could be required for LIMIT or DISTINCT ON (field).
#
# Solution: sqlalchemy should put all inner queries that contain order by or limit
# clauses into parenthesis, like this:
# SELECT ... UNION ALL (SELECT ... ORDER BY x LIMIT 10) ORDER BY y
#
# sqlalchemy.sql.compiler.SQLCompiler.visit_compound_select
# Fixed to surround subqueries with parens if they contain ORDER BY or LIMIT or OFFSET.
def sa_fixed_visit_compound_select(self, cs, asfrom=False, parens=True, compound_index=1, **kwargs):
    entry = self.stack and self.stack[-1] or {}
    self.stack.append({'from': entry.get('from', None), 'iswrapper': True})
    keyword = self.compound_keywords.get(cs.keyword)

    in_parens = lambda c: len(c._order_by_clause) > 0 or c._limit is not None or c._offset is not None
    text = (" " + keyword + " ").join(
        (c._compiler_dispatch(self, asfrom=asfrom or in_parens(c),
                              parens=in_parens(c),
                              compound_index=i, **kwargs)
         for i, c in enumerate(cs.selects))
    )

    group_by = cs._group_by_clause._compiler_dispatch(self, asfrom=asfrom, **kwargs)
    if group_by:
        text += " GROUP BY " + group_by

    text += self.order_by_clause(cs, **kwargs)
    text += (cs._limit is not None or cs._offset is not None) and \
            self.limit_clause(cs) or ""

    self.stack.pop(-1)
    if asfrom and parens:
        return "(" + text + ")"
    else:
        return text
from sqlalchemy.sql.compiler import SQLCompiler
SQLCompiler.visit_compound_select = sa_fixed_visit_compound_select