diff --git a/sqlaload/__init__.py b/sqlaload/__init__.py index c040e42..c738d4f 100644 --- a/sqlaload/__init__.py +++ b/sqlaload/__init__.py @@ -4,6 +4,6 @@ from sqlaload.schema import create_table, load_table, get_table, drop_table from sqlaload.schema import create_column from sqlaload.write import add_row, update_row from sqlaload.write import upsert, update, delete -from sqlaload.query import distinct, resultiter, all, find_one, find +from sqlaload.query import distinct, resultiter, all, find_one, find, query from sqlaload.util import dump_csv diff --git a/sqlaload/query.py b/sqlaload/query.py index f22c47d..e3d7d61 100644 --- a/sqlaload/query.py +++ b/sqlaload/query.py @@ -54,6 +54,10 @@ def find(engine, table, _limit=None, _step=5000, _offset=0, for row in rows: yield row +def query(engine, query): + for res in resultiter(engine.execute(query)): + yield res + def distinct(engine, table, *columns): columns = [table.c[c] for c in columns] q = expression.select(columns, distinct=True, diff --git a/sqlaload/schema.py b/sqlaload/schema.py index 44a4ba6..693d1c5 100644 --- a/sqlaload/schema.py +++ b/sqlaload/schema.py @@ -18,8 +18,10 @@ INDEXES = dict() def connect(url): """ Create an engine for the given database URL. """ kw = {} - #if url.startswith('postgres'): - # kw['pool_size'] = 10 + if url.startswith('postgres'): + #kw['pool_size'] = 5 + from sqlalchemy.pool import NullPool + kw['poolclass'] = NullPool engine = create_engine(url, **kw) engine = construct_engine(engine) meta = MetaData()