Search This Blog

04 November 2014

Opimize sqlalchemy queries with a With statement

More often than not, I need SQLAlchemy to issue as little SQL statement as possible. So I gave a try toward joinedload_all.
What I wanted was a nice way to hide this behind a with statement. It's a nice way to move the cumbersome extra options out of the way, but then I realized the inspection of the class could do the job for me.
This is what I got so far.

from sqlalchemy.orm import RelationshipProperty
from sqlalchemy.orm import joinedload_all

class LoadAll(object):
    """ builds a query with all the joined load
    """

    def __init__(self, klass):
        self.klass = klass

    def explore_relations(self):

        def follow(mapper, visited=None):
            """
            builds a tree of relations of the given mapper
            """
            result = {}
            visited = visited or []
            if mapper is not None and mapper not in visited:
                visited.append(mapper)
                for prop in mapper.iterate_properties:
                    if isinstance(prop, RelationshipProperty):
                        # next level
                        prop_arg = prop.argument
                        prop_arg = prop_arg() if callable(prop_arg) else prop_arg
                        prop_map = getattr(prop_arg, '__mapper__', None)
                        if prop_map is not None:
                            result[prop.key] = follow(prop_map, visited)
            return result
        return follow(self.klass.__mapper__)

    def paths(self, tree, cur=[]):
        """
        :param tree: a tree as nested dicts
        :param cur: current path, used to hold the result
        :return: all paths from root to leaves
        """
        if not tree:  # previous was a leaf, path to us is a solution
            yield cur
        else:  # we're a node
            for node, subtree in tree.iteritems():
                for path in self.paths(subtree, cur+[node,]):
                    yield path

    def __enter__(self):
        q = my_session.query(self.klass)
        tree = self.explore_relations()
        all_clauses = self.paths(tree)
        for join_clauses in all_clauses:
            q = q.options(joinedload_all(*join_clauses))
        return q

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

If given a setup like this

Employee N--1 Department N--1 Office N--1 Currency
                                    +N--1 Country 
in SQLAlchemy

from sqlalchemy import Column, DateTime, String, Integer, ForeignKey, func
from sqlalchemy.orm import relationship, backref
from sqlalchemy.ext.declarative import declarative_base

from sqlalchemy import create_engine
engine = create_engine('sqlite:///:memory:', echo=True)

from sqlalchemy.orm import sessionmaker
session = sessionmaker()
session.configure(bind=engine)
my_session = session()

Base = declarative_base()

class Country(Base):
    __tablename__ = 'country'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)

class Currency(Base):
    __tablename__ = 'currency'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)

class Office(Base):
    __tablename__ = 'office'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)
    country_id = Column(Integer, ForeignKey('country.id'))
    country = relationship("Country", backref="office")
    currency_id = Column(Integer, ForeignKey('currency.id'))
    currency = relationship("Currency", backref="office")

class Department(Base):
    __tablename__ = 'department'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)
    office_id = Column(Integer, ForeignKey('office.id'))
    office = relationship("Office", backref="dpts")


class Employee(Base):
    __tablename__ = 'employee'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    hired_on = Column(DateTime, default=func.now())
    department_id = Column(Integer, ForeignKey('department.id'))
    department = relationship(
        "Department",
        backref=backref('employees',
                         uselist=True,
                         cascade='delete,all'))


Base.metadata.create_all(engine)

c1 = Currency(name="euro")
c2 = Currency(name="sterling")
d1 = Department(name="dpt1", office=Office(name="Dublin", country=Country(name="IE"), currency=c1))
d2 = Department(name="dpt2", office=Office(name="London", country=Country(name="UK"), currency=c2))
e1 = Employee(name="dilbert", department=d1)
e2 = Employee(name="wally", department=d2)
e3 = Employee(name="alice", department=d2)
my_session.add(e1)
my_session.add(e2)
my_session.add(e3)

my_session.commit()

If we then use it like this

with LoadAll(Employee) as myq:
    for e in myq.all():
        print e.name, e.department.name, e.department.office.name, \
              e.department.office.country.name, e.department.office.currency.name

sqlalchemy will issue a single query

SELECT employee.id AS employee_id, employee.name AS employee_name, employee.hired_on AS employee_hired_on, employee.department_id AS employee_department_id, country_1.id AS country_1_id, country_1.name AS country_1_name, currency_1.id AS currency_1_id, currency_1.name AS currency_1_name, office_1.id AS office_1_id, office_1.name AS office_1_name, office_1.country_id AS office_1_country_id, office_1.currency_id AS office_1_currency_id, department_1.id AS department_1_id, department_1.name AS department_1_name, department_1.office_id AS department_1_office_id 
FROM employee 
LEFT OUTER JOIN department AS department_1 ON department_1.id = employee.department_id 
LEFT OUTER JOIN office AS office_1 ON office_1.id = department_1.office_id 
LEFT OUTER JOIN country AS country_1 ON country_1.id = office_1.country_id 
LEFT OUTER JOIN currency AS currency_1 ON currency_1.id = office_1.currency_id

01 November 2014

recursivity elegance

while searching for an elegant solution, I stumbled upon this by Andrew Clark
incoming = {'E': {'D': {'A': {},
                        'O': {'Co': {},
                              'Cu': {}
                             }
                        }
                 }
            }

expected = [
   ('E', 'D', 'A'),
   ('E', 'D', 'O', 'Co'),
   ('E', 'D', 'O', 'Cu'),
]


def paths(tree, cur=()):
    if not tree:  # previous was a leaf, path to us is a solution
        yield cur
    else:  # we're a node
        for node, subtree in tree.iteritems():
            for path in paths(subtree, cur+(node,)):
                yield path

assert list(paths(incoming)) == expected