SQLAlchemy SelectInto ORM query support thingie idea test

selectinto.py

from sqlalchemy.sql.selectable import HasPrefixes
from sqlalchemy.sql.base import Executable
from sqlalchemy.sql.elements import ClauseElement
from sqlalchemy.ext.compiler import compiles
from sqlalchemy import Table, MetaData


class SelectInto(HasPrefixes, Executable, ClauseElement):

    def __init__(self,
                 select=None,
                 into=None,
                 prefixes=None,
                 **kwargs):
        self.select = select
        self.into = into
        self._setup_prefixes(
            prefixes if prefixes is not None else ['TEMPORARY'])


@compiles(SelectInto)
def visit_selectinto(element, compiler, **kwargs):
    text = compiler.visit_select(element.select)
    text = text.replace(
        'FROM',
        'INTO %s TABLE %s FROM' % (
            compiler._generate_prefixes(element, element._prefixes, **kwargs),
            element.into
        )
    )
    return text


def query_into(query, name, **kwargs):
    statement = SelectInto(query.selectable, name, **kwargs)
    conn = query._connection_from_session(
        mapper=query._bind_mapper(),
        clause=statement)
    conn.execute(statement, query._params)
    conn.close()
    return Table(
        name,
        MetaData(bind=query.session.bind),
        *map(lambda col: col.copy(), statement.select.columns)
    )