Merge branch 'pr182'

6 files changed, 47 insertions, 12 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 35de41fef..36a9d7bf7 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -74,6 +74,8 @@ See also:
 `PQconnectdbParams <http://www.postgresql.org/docs/9.1/static/\
 libpq-connect.html#LIBPQ-PQCONNECTDBPARAMS>`_
 
+.. _psycopg2_execution_options:
+
 Per-Statement/Connection Execution Options
 -------------------------------------------
 
@@ -81,16 +83,23 @@ The following DBAPI-specific options are respected when used with
 :meth:`.Connection.execution_options`, :meth:`.Executable.execution_options`,
 :meth:`.Query.execution_options`, in addition to those not specific to DBAPIs:
 
-* isolation_level - Set the transaction isolation level for the lifespan of a
+* ``isolation_level`` - Set the transaction isolation level for the lifespan of a
   :class:`.Connection` (can only be set on a connection, not a statement
   or query).   See :ref:`psycopg2_isolation_level`.
 
-* stream_results - Enable or disable usage of psycopg2 server side cursors -
+* ``stream_results`` - Enable or disable usage of psycopg2 server side cursors -
   this feature makes use of "named" cursors in combination with special
   result handling methods so that result rows are not fully buffered.
   If ``None`` or not set, the ``server_side_cursors`` option of the
   :class:`.Engine` is used.
 
+* ``max_row_buffer`` - when using ``stream_results``, an integer value that
+  specifies the maximum number of rows to buffer at a time.  This is
+  interpreted by the :class:`.BufferedRowResultProxy`, and if omitted the
+  buffer will grow to ultimately store 1000 rows at a time.
+
+  .. versionadded:: 1.0.6
+
 .. _psycopg2_unicode:
 
 Unicode with Psycopg2
diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
index 6d19cb6d0..b2b78dee8 100644
--- a/lib/sqlalchemy/engine/result.py
+++ b/lib/sqlalchemy/engine/result.py
@@ -1067,10 +1067,27 @@ class BufferedRowResultProxy(ResultProxy):
 
     The pre-fetching behavior fetches only one row initially, and then
     grows its buffer size by a fixed amount with each successive need
-    for additional rows up to a size of 100.
+    for additional rows up to a size of 1000.
+
+    The size argument is configurable using the ``max_row_buffer``
+    execution option::
+
+        with psycopg2_engine.connect() as conn:
+
+            result = conn.execution_options(
+                stream_results=True, max_row_buffer=50
+                ).execute("select * from table")
+
+    .. versionadded:: 1.0.6 Added the ``max_row_buffer`` option.
+
+    .. seealso::
+
+        :ref:`psycopg2_execution_options`
     """
 
     def _init_metadata(self):
+        self._max_row_buffer = self.context.execution_options.get(
+            'max_row_buffer', None)
         self.__buffer_rows()
         super(BufferedRowResultProxy, self)._init_metadata()
 
@@ -1095,6 +1112,8 @@ class BufferedRowResultProxy(ResultProxy):
         size = getattr(self, '_bufsize', 1)
         self.__rowbuffer = collections.deque(self.cursor.fetchmany(size))
         self._bufsize = self.size_growth.get(size, size)
+        if self._max_row_buffer is not None:
+            self._bufsize = min(self._max_row_buffer, self._bufsize)
 
     def _soft_close(self, **kw):
         self.__rowbuffer.clear()
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 4f8c86a14..8b3df08e7 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -756,7 +756,8 @@ class Query(object):
         """
         self._yield_per = count
         self._execution_options = self._execution_options.union(
-            {"stream_results": True})
+            {"stream_results": True,
+             "max_row_buffer": count})
 
     def get(self, ident):
         """Return an instance based on the given primary key identifier,
diff --git a/lib/sqlalchemy/testing/__init__.py b/lib/sqlalchemy/testing/__init__.py
index adfbe85e3..7482e32a1 100644
--- a/lib/sqlalchemy/testing/__init__.py
+++ b/lib/sqlalchemy/testing/__init__.py
@@ -19,7 +19,7 @@ def against(*queries):
     return _against(config._current, *queries)
 
 from .assertions import emits_warning, emits_warning_on, uses_deprecated, \
-    eq_, ne_, is_, is_not_, startswith_, assert_raises, \
+    eq_, ne_, le_, is_, is_not_, startswith_, assert_raises, \
     assert_raises_message, AssertsCompiledSQL, ComparesTables, \
     AssertsExecutionResults, expect_deprecated, expect_warnings
 
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index e5249c296..e0c02c896 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -216,6 +216,11 @@ def ne_(a, b, msg=None):
     assert a != b, msg or "%r == %r" % (a, b)
 
 
+def le_(a, b, msg=None):
+    """Assert a <= b, with repr messaging on failure."""
+    assert a <= b, msg or "%r != %r" % (a, b)
+
+
 def is_(a, b, msg=None):
     """Assert a is b, with repr messaging on failure."""
     assert a is b, msg or "%r is not %r" % (a, b)
diff --git a/lib/sqlalchemy/testing/fixtures.py b/lib/sqlalchemy/testing/fixtures.py
index 7b421952f..e16bc77c0 100644
--- a/lib/sqlalchemy/testing/fixtures.py
+++ b/lib/sqlalchemy/testing/fixtures.py
@@ -134,13 +134,14 @@ class TablesTest(TestBase):
     def _teardown_each_tables(self):
         # no need to run deletes if tables are recreated on setup
         if self.run_define_tables != 'each' and self.run_deletes == 'each':
-            for table in reversed(self.metadata.sorted_tables):
-                try:
-                    table.delete().execute().close()
-                except sa.exc.DBAPIError as ex:
-                    util.print_(
-                        ("Error emptying table %s: %r" % (table, ex)),
-                        file=sys.stderr)
+            with self.bind.connect() as conn:
+                for table in reversed(self.metadata.sorted_tables):
+                    try:
+                        conn.execute(table.delete())
+                    except sa.exc.DBAPIError as ex:
+                        util.print_(
+                            ("Error emptying table %s: %r" % (table, ex)),
+                            file=sys.stderr)
 
     def setup(self):
         self._setup_each_tables()