#!/usr/bin/env python
#-*- coding:utf-8 -*-
import time
from collections import Counter
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import Column, ForeignKey, Integer, String
from sqlalchemy.orm import relationship, backref, joinedload
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
engine = create_engine('sqlite:///:memory:')
Session = sessionmaker(bind=engine)
Base = declarative_base()
session = Session()
class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    fullname = Column(String)
    password = Column(String)
class Address(Base):
    __tablename__ = 'addresses'
    id = Column(Integer, primary_key=True)
    email_address = Column(String, nullable=False)
    user_id = Column(Integer, ForeignKey('users.id'))
    user = relationship("User", backref=backref('addresses', order_by=id))
Base.metadata.create_all(engine)
users = [
    User(name='wendy', fullname='Wendy Williams', password='foobar'),
    User(name='mary', fullname='Mary Contrary', password='xxg527'),
    User(name='fred', fullname='Fred Flinstone', password='blah')
]
for user in users:
    user.addresses = [
        Address(email_address=user.name + '@example.com'),
        Address(email_address=user.name + '@work.corp.com')
    ]
session.add_all(users)
session.commit()
class QMonitor(object):
    def __init__(self, max_queries):
        self.queries = []
        self.max_queries = max_queries
    def assist_user(self):
        """ User requires assistance in query minimization.
        """
        qc = Counter(zip(*self.queries)[0])
        if qc.values().count(1) < len(qc):
            duplicates = [
                '** Repeated {0} times: {1}'.format(c, q)
                for q, c in qc.items()
                if c > 1
            ]
            print (u'Too many queries, try eliminating duplicates:\n' +
                   u'\n'.join(duplicates))
        else:
            print u'Too many queries:\n' + u'\n'.join(zip(*self.queries)[0])
    def check(self):
        """ Check saved database queries.
        """
        if len(self.queries) > self.max_queries:
            self.assist_user()
@contextmanager
def qmonitor(max_queries=15):
    qm = QMonitor(max_queries)
    def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        context._query_start_time = time.time()
    def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        total = time.time() - context._query_start_time
        qm.queries.append((statement, total))
    event.listen(engine, 'before_cursor_execute', before_cursor_execute)
    event.listen(engine, 'after_cursor_execute', after_cursor_execute)
    yield qm
    event.remove(engine, 'before_cursor_execute', before_cursor_execute)
    event.remove(engine, 'after_cursor_execute', after_cursor_execute)
print '== Naive query'
q = session.query(Address)
with qmonitor(max_queries=1) as qm:
    for address in q:
        print address.email_address + ' - ' + address.user.fullname
qm.check()
print
print '== Optimized query'
q = session.query(Address).options(joinedload('user'))
with qmonitor(max_queries=1) as qm:
    for address in q:
        print address.email_address + ' - ' + address.user.fullname
qm.check()