records.py in sqlalchemy

rows.py

import tablib
from sqlalchemy import create_engine
from sqlalchemy.sql import text


class ResultSet(object):
    def __init__(self, result_proxy):
        self._rp = result_proxy
        self.colum_names = self._rp.keys()
        self._all_rows = []
        self._completed = False

    def next(self):
        result = self._rp.fetchone()
        if result:
            return result
        raise StopIteration("ResultSet contains no more rows.")

    def __iter__(self):
        if self._completed:
            for row in self._all_rows:
                yield row

        for row in self._rp:
            self._all_rows.append(row)
            yield row

        self._completed = True

    def all(self):
        if not self._all_rows:
            self._all_rows = list(self._rp)
        return self._all_rows

    @property
    def dataset(self):
        data = tablib.Dataset()
        data.headers = self.colum_names

        for row in self.all():
            row = _reduce_datetimes(row)
            data.append(row)
        return data


class Database(object):
    def __init__(self, dsn):
        self.engine = create_engine(dsn)

    def query(self, query_string, params={}, fetchall=False):
        query = text(query_string)
        result_proxy = self.engine.execute(query, **params)

        return ResultSet(result_proxy)