Add support for server side cursors to mysqldb and pymysql

This allows to skip buffering of the results on the client side, e.g.
the following snippet:

    table = sa.Table(
        'testtbl', sa.MetaData(),
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('a', sa.Integer),
        sa.Column('b', sa.String(512))
    )
    table.create(eng, checkfirst=True)

    with eng.connect() as conn:
        result = conn.execute(table.select().limit(1)).fetchone()
        if result is None:
            for _ in range(1000):
                conn.execute(
                    table.insert(),
                    [{'a': random.randint(1, 100000),
                      'b': ''.join(random.choice(string.ascii_letters) for _ in range(100))}
                      for _ in range(1000)]
                )

    with eng.connect() as conn:
        for row in conn.execution_options(stream_results=True).execute(table.select()):
            pass

now uses ~23 MB of memory instead of ~327 MB on CPython 3.5.2 and
PyMySQL 0.7.9.

psycopg2 implementation and execution options (stream_results,
server_side_cursors) are reused.

Change-Id: I4dc23ce3094f027bdff51b896b050361991c62e2

4 files changed, 65 insertions, 38 deletions

diff --git a/lib/sqlalchemy/dialects/mysql/base.py b/lib/sqlalchemy/dialects/mysql/base.py
index e7e533890..449fffaba 100644
--- a/lib/sqlalchemy/dialects/mysql/base.py
+++ b/lib/sqlalchemy/dialects/mysql/base.py
@@ -177,6 +177,22 @@ multi-column key for some storage engines::
         Column('id', Integer, primary_key=True)
        )
 
+.. _mysql_ss_cursors:
+
+Server Side Cursors
+-------------------
+
+Server-side cursor support is available for the MySQLdb and PyMySQL dialects.
+From a MySQL point of view this means that the ``MySQLdb.cursors.SSCursor`` or
+``pymysql.cursors.SSCursor`` class is used when building up the cursor which
+will receive results.  The most typical way of invoking this feature is via the
+:paramref:`.Connection.execution_options.stream_results` connection execution
+option.   Server side cursors can also be enabled for all SELECT statements
+unconditionally by passing ``server_side_cursors=True`` to
+:func:`.create_engine`.
+
+.. versionadded:: 1.1.4 - added server-side cursor support.
+
 .. _mysql_unicode:
 
 Unicode
@@ -743,6 +759,12 @@ class MySQLExecutionContext(default.DefaultExecutionContext):
     def should_autocommit_text(self, statement):
         return AUTOCOMMIT_RE.match(statement)
 
+    def create_server_side_cursor(self):
+        if self.dialect.supports_server_side_cursors:
+            return self._dbapi_connection.cursor(self.dialect._sscursor)
+        else:
+            raise NotImplementedError()
+
 
 class MySQLCompiler(compiler.SQLCompiler):
 
diff --git a/lib/sqlalchemy/dialects/mysql/mysqldb.py b/lib/sqlalchemy/dialects/mysql/mysqldb.py
index aa8377b27..568c05f62 100644
--- a/lib/sqlalchemy/dialects/mysql/mysqldb.py
+++ b/lib/sqlalchemy/dialects/mysql/mysqldb.py
@@ -38,6 +38,11 @@ using a URL like the following::
 
     mysql+mysqldb://root@/<dbname>?unix_socket=/cloudsql/<projectid>:<instancename>
 
+Server Side Cursors
+-------------------
+
+The mysqldb dialect supports server-side cursors. See :ref:`mysql_ss_cursors`.
+
 """
 
 from .base import (MySQLDialect, MySQLExecutionContext,
@@ -87,6 +92,19 @@ class MySQLDialect_mysqldb(MySQLDialect):
     statement_compiler = MySQLCompiler_mysqldb
     preparer = MySQLIdentifierPreparer_mysqldb
 
+    def __init__(self, server_side_cursors=False, **kwargs):
+        super(MySQLDialect_mysqldb, self).__init__(**kwargs)
+        self.server_side_cursors = server_side_cursors
+
+    @util.langhelpers.memoized_property
+    def supports_server_side_cursors(self):
+        try:
+            cursors = __import__('MySQLdb.cursors').cursors
+            self._sscursor = cursors.SSCursor
+            return True
+        except (ImportError, AttributeError):
+            return False
+
     @classmethod
     def dbapi(cls):
         return __import__('MySQLdb')
diff --git a/lib/sqlalchemy/dialects/mysql/pymysql.py b/lib/sqlalchemy/dialects/mysql/pymysql.py
index 3c493fbfc..e29c17d8b 100644
--- a/lib/sqlalchemy/dialects/mysql/pymysql.py
+++ b/lib/sqlalchemy/dialects/mysql/pymysql.py
@@ -30,7 +30,7 @@ to the pymysql driver as well.
 """
 
 from .mysqldb import MySQLDialect_mysqldb
-from ...util import py3k
+from ...util import langhelpers, py3k
 
 
 class MySQLDialect_pymysql(MySQLDialect_mysqldb):
@@ -44,6 +44,19 @@ class MySQLDialect_pymysql(MySQLDialect_mysqldb):
     supports_unicode_statements = True
     supports_unicode_binds = True
 
+    def __init__(self, server_side_cursors=False, **kwargs):
+        super(MySQLDialect_pymysql, self).__init__(**kwargs)
+        self.server_side_cursors = server_side_cursors
+
+    @langhelpers.memoized_property
+    def supports_server_side_cursors(self):
+        try:
+            cursors = __import__('pymysql.cursors').cursors
+            self._sscursor = cursors.SSCursor
+            return True
+        except (ImportError, AttributeError):
+            return False
+
     @classmethod
     def dbapi(cls):
         return __import__('pymysql')
diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 8488da816..27a1ec099 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -28,7 +28,8 @@ psycopg2-specific keyword arguments which are accepted by
   :class:`~sqlalchemy.engine.ResultProxy` uses special row-buffering
   behavior when this feature is enabled, such that groups of 100 rows at a
   time are fetched over the wire to reduce conversational overhead.
-  Note that the ``stream_results=True`` execution option is a more targeted
+  Note that the :paramref:`.Connection.execution_options.stream_results`
+  execution option is a more targeted
   way of enabling this mode on a per-execution basis.
 * ``use_native_unicode``: Enable the usage of Psycopg2 "native unicode" mode
   per connection.  True by default.
@@ -422,53 +423,24 @@ class _PGUUID(UUID):
                 return value
             return process
 
-# 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)
 
 _server_side_id = util.counter()
 
 
 class PGExecutionContext_psycopg2(PGExecutionContext):
-    def create_cursor(self):
-        # TODO: coverage for server side cursors + select.for_update()
-
-        if self.dialect.server_side_cursors:
-            is_server_side = \
-                self.execution_options.get('stream_results', True) and (
-                    (self.compiled and isinstance(self.compiled.statement,
-                                                  expression.Selectable)
-                     or
-                     (
-                        (not self.compiled or
-                         isinstance(self.compiled.statement,
-                                    expression.TextClause))
-                        and self.statement and SERVER_SIDE_CURSOR_RE.match(
-                            self.statement))
-                     )
-                )
-        else:
-            is_server_side = \
-                self.execution_options.get('stream_results', False)
-
-        self.__is_server_side = is_server_side
-        if is_server_side:
-            # use server-side cursors:
-            # http://lists.initd.org/pipermail/psycopg/2007-January/005251.html
-            ident = "c_%s_%s" % (hex(id(self))[2:],
-                                 hex(_server_side_id())[2:])
-            return self._dbapi_connection.cursor(ident)
-        else:
-            return self._dbapi_connection.cursor()
+    def create_server_side_cursor(self):
+        # use server-side cursors:
+        # http://lists.initd.org/pipermail/psycopg/2007-January/005251.html
+        ident = "c_%s_%s" % (hex(id(self))[2:],
+                             hex(_server_side_id())[2:])
+        return self._dbapi_connection.cursor(ident)
 
     def get_result_proxy(self):
         # TODO: ouch
         if logger.isEnabledFor(logging.INFO):
             self._log_notices(self.cursor)
 
-        if self.__is_server_side:
+        if self._is_server_side:
             return _result.BufferedRowResultProxy(self)
         else:
             return _result.ResultProxy(self)
@@ -502,6 +474,8 @@ class PGDialect_psycopg2(PGDialect):
     if util.py2k:
         supports_unicode_statements = False
 
+    supports_server_side_cursors = True
+
     default_paramstyle = 'pyformat'
     # set to true based on psycopg2 version
     supports_sane_multi_rowcount = False