- added "statement_options()" to Query, to so options can be

passed to the resulting statement. Currently only
Select-statements have these options, and the only option
used is "stream_results", and the only dialect which knows
"stream_results" is psycopg2.
- Query.yield_per() will set the "stream_results" statement
option automatically.
- Added "statement_options()" to Selects, which set statement
specific options. These enable e.g. dialect specific options
such as whether to enable using server side cursors, etc.
- The psycopg2 now respects the statement option
"stream_results". This option overrides the connection setting
"server_side_cursors". If true, server side cursors will be
used for the statement. If false, they will not be used, even
if "server_side_cursors" is true on the
connection. [ticket:1619]
- added a "frozendict" from http://code.activestate.com/recipes/414283/,
adding more default collections as immutable class vars on
Query, Insert, Select

1 files changed, 28 insertions, 12 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index a46fdbddb..7733aadcd 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -35,6 +35,15 @@ Transactions
 The psycopg2 dialect fully supports SAVEPOINT and two-phase commit operations.
 
 
+Statement options
+-----------------
+
+The following statement options are respected:
+
+* *stream_results* - Enable or disable usage of server side cursors for the SELECT-statement.
+  If *None* or not set, the *server_side_cursors* option of the connection is used. If
+  auto-commit is enabled, the option is ignored.
+
 """
 
 import decimal, random, re
@@ -93,8 +102,9 @@ class _PGArray(ARRAY):
         if isinstance(self.item_type, sqltypes.String) and \
                     self.item_type.convert_unicode:
             self.item_type.convert_unicode = "force"
-    
-# TODO: filter out 'FOR UPDATE' statements
+
+# When we're handed literal SQL, ensure it's a SELECT-query. Since
+# 8.3, combining cursors and "FOR UPDATE" has been fine.
 SERVER_SIDE_CURSOR_RE = re.compile(
     r'\s*SELECT',
     re.I | re.UNICODE)
@@ -102,16 +112,22 @@ SERVER_SIDE_CURSOR_RE = re.compile(
 class PostgreSQL_psycopg2ExecutionContext(PGExecutionContext):
     def create_cursor(self):
         # TODO: coverage for server side cursors + select.for_update()
-        is_server_side = \
-            self.dialect.server_side_cursors and \
-            not self.should_autocommit and \
-            ((self.compiled and isinstance(self.compiled.statement, expression.Selectable) 
-                and not getattr(self.compiled.statement, 'for_update', False)) \
-            or \
-            (
-                (not self.compiled or isinstance(self.compiled.statement, expression._TextClause)) 
-                and self.statement and SERVER_SIDE_CURSOR_RE.match(self.statement))
-            )
+        stream_results_option = self.statement_options.get('stream_results')
+        is_server_side = (
+            # Enabled for this statement ...
+            (stream_results_option or
+             # ... or enabled for all statements
+             (self.dialect.server_side_cursors and
+              # ... and not explicitly disabled for this one.
+              (stream_results_option or stream_results_option is None))
+             ) and (
+                # But don't use SS-cursors when autocommit is on ...
+                (not self.should_autocommit and
+                 self.compiled and isinstance(self.compiled.statement, expression.Selectable))
+                or (
+                    # ... or if it's not even a SELECT.
+                    (not self.compiled or isinstance(self.compiled.statement, expression._TextClause))
+                    and self.statement and SERVER_SIDE_CURSOR_RE.match(self.statement))))
 
         self.__is_server_side = is_server_side
         if is_server_side: