Add fast execution helper support.

Added a new flag ``use_batch_mode`` to the psycopg2 dialect.  This flag
enables the use of psycopg2's ``psycopg2.extras.execute_batch``
extension when the :class:`.Engine` calls upon ``cursor.executemany()``.
This extension provides a critical performance increase by over an order of magnitude
when running INSERT statements in batch.  The flag is False by default
as it is considered to be experimental for now.

Change-Id: Ib88d28bc792958d47109f644ff1d08c897db4ff7
Fixes: #4109

1 files changed, 45 insertions, 1 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 3e2968d91..8ac39c201 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -53,6 +53,15 @@ psycopg2-specific keyword arguments which are accepted by
 
     :ref:`psycopg2_unicode`
 
+* ``use_batch_mode``: This flag allows ``psycopg2.extras.execute_batch``
+  for ``cursor.executemany()`` calls performed by the :class:`.Engine`.
+  It is currently experimental but
+  may well become True by default as it is critical for executemany performance.
+
+  .. seealso::
+
+    :ref:`psycopg2_batch_mode`
+
 Unix Domain Connections
 ------------------------
 
@@ -101,6 +110,31 @@ The following DBAPI-specific options are respected when used with
 
   .. versionadded:: 1.0.6
 
+.. _psycopg2_batch_mode:
+
+Psycopg2 Batch Mode (Fast Execution)
+------------------------------------
+
+Modern versions of psycopg2 include a feature known as
+`Fast Execution Helpers <http://initd.org/psycopg/docs/extras.html#fast-execution-helpers>`_,
+which have been shown in benchmarking to improve psycopg2's executemany()
+performance with INSERTS by multiple orders of magnitude.   SQLAlchemy
+allows this extension to be used for all ``executemany()`` style calls
+invoked by an :class:`.Engine` when used with multiple parameter sets,
+by adding the ``use_batch_mode`` flag to :func:`.create_engine`::
+
+    engine = create_engine(
+        "postgresql+psycopg2://scott:tiger@host/dbname",
+        use_batch_mode=True)
+
+Batch mode is considered to be **experimental** at this time, however may
+be enabled by default in a future release.
+
+
+.. versionadded:: 1.2.0
+
+
+
 .. _psycopg2_unicode:
 
 Unicode with Psycopg2
@@ -510,6 +544,7 @@ class PGDialect_psycopg2(PGDialect):
     def __init__(self, server_side_cursors=False, use_native_unicode=True,
                  client_encoding=None,
                  use_native_hstore=True, use_native_uuid=True,
+                 use_batch_mode=False,
                  **kwargs):
         PGDialect.__init__(self, **kwargs)
         self.server_side_cursors = server_side_cursors
@@ -518,6 +553,7 @@ class PGDialect_psycopg2(PGDialect):
         self.use_native_uuid = use_native_uuid
         self.supports_unicode_binds = use_native_unicode
         self.client_encoding = client_encoding
+        self.psycopg2_batch_mode = use_batch_mode
         if self.dbapi and hasattr(self.dbapi, '__version__'):
             m = re.match(r'(\d+)\.(\d+)(?:\.(\d+))?',
                          self.dbapi.__version__)
@@ -540,7 +576,8 @@ class PGDialect_psycopg2(PGDialect):
         # http://initd.org/psycopg/docs/news.html#what-s-new-in-psycopg-2-0-9
         self.supports_sane_multi_rowcount = \
             self.psycopg2_version >= \
-            self.FEATURE_VERSION_MAP['sane_multi_rowcount']
+            self.FEATURE_VERSION_MAP['sane_multi_rowcount'] and \
+            not self.psycopg2_batch_mode
 
     @classmethod
     def dbapi(cls):
@@ -638,6 +675,13 @@ class PGDialect_psycopg2(PGDialect):
         else:
             return None
 
+    def do_executemany(self, cursor, statement, parameters, context=None):
+        if self.psycopg2_batch_mode:
+            extras = self._psycopg2_extras()
+            extras.execute_batch(cursor, statement, parameters)
+        else:
+            cursor.executemany(statement, parameters)
+
     @util.memoized_instancemethod
     def _hstore_oids(self, conn):
         if self.psycopg2_version >= self.FEATURE_VERSION_MAP['hstore_adapter']: