Revise psycopg2 execute_values approach

Revised the approach for the just added support for the psycopg2
"execute_values()" feature added in 1.3.7 for :ticket:`4623`.  The approach
relied upon a regular expression that would fail to match for a more
complex INSERT statement such as one which had subqueries involved.   The
new approach matches exactly the string that was rendered as the VALUES
clause.

Fixes: #4623
Change-Id: Icaae0f7b6bcf87a2cf5c6290a839c8429dd5fac3
(cherry picked from commit 0a7ca00e04f7c1cfdbb8bdfe7da5f62fc9b40930)

diff --git a/doc/build/changelog/unreleased_13/4623.rst b/doc/build/changelog/unreleased_13/4623.rst
new file mode 100644 (file)
index 0000000..462e77f
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/4623.rst
@@ -0,0 +1,10 @@
+.. change::
+    :tags: bug, postgresql
+    :tickets: 4623
+
+    Revised the approach for the just added support for the psycopg2
+    "execute_values()" feature added in 1.3.7 for :ticket:`4623`.  The approach
+    relied upon a regular expression that would fail to match for a more
+    complex INSERT statement such as one which had subqueries involved.   The
+    new approach matches exactly the string that was rendered as the VALUES
+    clause.

diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 26014dadce3e5b6e62665f123946ec24123637af..1a4db1108dcd6d241bb5a5bdfd218020d8923598 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -849,8 +849,6 @@ class PGDialect_psycopg2(PGDialect):
         else:
             return None
 
-    _insert_values_match = re.compile(r".* VALUES (\(.+\))").match
-
     def do_executemany(self, cursor, statement, parameters, context=None):
         if self.executemany_mode is EXECUTEMANY_DEFAULT:
             cursor.executemany(statement, parameters)
@@ -860,8 +858,14 @@ class PGDialect_psycopg2(PGDialect):
             self.executemany_mode is EXECUTEMANY_VALUES
             and context
             and context.isinsert
+            and context.compiled.insert_single_values_expr
         ):
-            executemany_values = self._insert_values_match(statement)
+            executemany_values = (
+                "(%s)" % context.compiled.insert_single_values_expr
+            )
+            # guard for statement that was altered via event hook or similar
+            if executemany_values not in statement:
+                executemany_values = None
         else:
             executemany_values = None
 
@@ -870,7 +874,7 @@ class PGDialect_psycopg2(PGDialect):
             # into executemany(), since no DBAPI has ever supported that
             # until the introduction of psycopg2's executemany_values, so
             # we are not yet using the fetch=True flag.
-            statement = statement.replace(executemany_values.group(1), "%s")
+            statement = statement.replace(executemany_values, "%s")
             if self.executemany_values_page_size:
                 kwargs = {"page_size": self.executemany_values_page_size}
             else:
@@ -879,7 +883,7 @@ class PGDialect_psycopg2(PGDialect):
                 cursor,
                 statement,
                 parameters,
-                template=executemany_values.group(1),
+                template=executemany_values,
                 **kwargs
             )
 

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 7c4a7a518356b42dc6ffd62de1b19cbdbb8b78c0..e395caeaa84c845d9ef6f7ca35df1f19114468ed 100644 (file)
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -495,6 +495,15 @@ class SQLCompiler(Compiled):
 
     """
 
+    insert_single_values_expr = None
+    """When an INSERT is compiled with a single set of parameters inside
+    a VALUES expression, the string is assigned here, where it can be
+    used for insert batching schemes to rewrite the VALUES expression.
+
+    .. versionadded:: 1.3.8
+
+    """
+
     insert_prefetch = update_prefetch = ()
 
     def __init__(
@@ -2468,7 +2477,10 @@ class SQLCompiler(Compiled):
                 )
             )
         else:
-            text += " VALUES (%s)" % ", ".join([c[1] for c in crud_params])
+            insert_single_values_expr = ", ".join([c[1] for c in crud_params])
+            text += " VALUES (%s)" % insert_single_values_expr
+            if toplevel:
+                self.insert_single_values_expr = insert_single_values_expr
 
         if insert_stmt._post_values_clause is not None:
             post_values_clause = self.process(

diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index 798831cd383a6a4824ce136fd6f0a6093d2e2b57..cf33f7cbbf5422e750611fcbbf3f25beb5492d44 100644 (file)
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -10,11 +10,13 @@ from sqlalchemy import cast
 from sqlalchemy import Column
 from sqlalchemy import DateTime
 from sqlalchemy import dialects
+from sqlalchemy import event
 from sqlalchemy import exc
 from sqlalchemy import extract
 from sqlalchemy import func
 from sqlalchemy import Integer
 from sqlalchemy import literal
+from sqlalchemy import literal_column
 from sqlalchemy import MetaData
 from sqlalchemy import Numeric
 from sqlalchemy import schema
@@ -191,20 +193,53 @@ class ExecuteManyMode(object):
         super(ExecuteManyMode, self).teardown()
 
     def test_insert(self):
-        with self.engine.connect() as conn:
-            conn.execute(
-                self.tables.data.insert(),
-                [
-                    {"x": "x1", "y": "y1"},
-                    {"x": "x2", "y": "y2"},
-                    {"x": "x3", "y": "y3"},
-                ],
-            )
+        from psycopg2 import extras
 
-            eq_(
-                conn.execute(select([self.tables.data])).fetchall(),
-                [(1, "x1", "y1", 5), (2, "x2", "y2", 5), (3, "x3", "y3", 5)],
-            )
+        if self.engine.dialect.executemany_mode is EXECUTEMANY_BATCH:
+            meth = extras.execute_batch
+            stmt = "INSERT INTO data (x, y) VALUES (%(x)s, %(y)s)"
+            expected_kwargs = {}
+        else:
+            meth = extras.execute_values
+            stmt = "INSERT INTO data (x, y) VALUES %s"
+            expected_kwargs = {"template": "(%(x)s, %(y)s)"}
+
+        with mock.patch.object(
+            extras, meth.__name__, side_effect=meth
+        ) as mock_exec:
+            with self.engine.connect() as conn:
+                conn.execute(
+                    self.tables.data.insert(),
+                    [
+                        {"x": "x1", "y": "y1"},
+                        {"x": "x2", "y": "y2"},
+                        {"x": "x3", "y": "y3"},
+                    ],
+                )
+
+                eq_(
+                    conn.execute(select([self.tables.data])).fetchall(),
+                    [
+                        (1, "x1", "y1", 5),
+                        (2, "x2", "y2", 5),
+                        (3, "x3", "y3", 5),
+                    ],
+                )
+        eq_(
+            mock_exec.mock_calls,
+            [
+                mock.call(
+                    mock.ANY,
+                    stmt,
+                    (
+                        {"x": "x1", "y": "y1"},
+                        {"x": "x2", "y": "y2"},
+                        {"x": "x3", "y": "y3"},
+                    ),
+                    **expected_kwargs
+                )
+            ],
+        )
 
     def test_insert_no_page_size(self):
         from psycopg2 import extras
@@ -249,6 +284,8 @@ class ExecuteManyMode(object):
         )
 
     def test_insert_page_size(self):
+        from psycopg2 import extras
+
         opts = self.options.copy()
         opts["executemany_batch_page_size"] = 500
         opts["executemany_values_page_size"] = 1000
@@ -256,8 +293,6 @@ class ExecuteManyMode(object):
         with self.expect_deprecated_opts():
             eng = engines.testing_engine(options=opts)
 
-        from psycopg2 import extras
-
         if eng.dialect.executemany_mode is EXECUTEMANY_BATCH:
             meth = extras.execute_batch
             stmt = "INSERT INTO data (x, y) VALUES (%(x)s, %(y)s)"
@@ -379,6 +414,119 @@ class ExecutemanyBatchModeTest(ExecuteManyMode, fixtures.TablesTest):
 class ExecutemanyValuesInsertsTest(ExecuteManyMode, fixtures.TablesTest):
     options = {"executemany_mode": "values"}
 
+    def test_insert_w_newlines(self):
+        from psycopg2 import extras
+
+        t = self.tables.data
+
+        ins = t.insert(inline=True).values(
+            id=bindparam("id"),
+            x=select([literal_column("5")]).select_from(self.tables.data),
+            y=bindparam("y"),
+            z=bindparam("z"),
+        )
+        # compiled SQL has a newline in it
+        eq_(
+            str(ins.compile(testing.db)),
+            "INSERT INTO data (id, x, y, z) VALUES (%(id)s, "
+            "(SELECT 5 \nFROM data), %(y)s, %(z)s)",
+        )
+        meth = extras.execute_values
+        with mock.patch.object(
+            extras, "execute_values", side_effect=meth
+        ) as mock_exec:
+
+            with self.engine.connect() as conn:
+                conn.execute(
+                    ins,
+                    [
+                        {"id": 1, "y": "y1", "z": 1},
+                        {"id": 2, "y": "y2", "z": 2},
+                        {"id": 3, "y": "y3", "z": 3},
+                    ],
+                )
+
+        eq_(
+            mock_exec.mock_calls,
+            [
+                mock.call(
+                    mock.ANY,
+                    "INSERT INTO data (id, x, y, z) VALUES %s",
+                    (
+                        {"id": 1, "y": "y1", "z": 1},
+                        {"id": 2, "y": "y2", "z": 2},
+                        {"id": 3, "y": "y3", "z": 3},
+                    ),
+                    template="(%(id)s, (SELECT 5 \nFROM data), %(y)s, %(z)s)",
+                )
+            ],
+        )
+
+    def test_insert_modified_by_event(self):
+        from psycopg2 import extras
+
+        t = self.tables.data
+
+        ins = t.insert(inline=True).values(
+            id=bindparam("id"),
+            x=select([literal_column("5")]).select_from(self.tables.data),
+            y=bindparam("y"),
+            z=bindparam("z"),
+        )
+        # compiled SQL has a newline in it
+        eq_(
+            str(ins.compile(testing.db)),
+            "INSERT INTO data (id, x, y, z) VALUES (%(id)s, "
+            "(SELECT 5 \nFROM data), %(y)s, %(z)s)",
+        )
+        meth = extras.execute_batch
+        with mock.patch.object(
+            extras, "execute_values"
+        ) as mock_values, mock.patch.object(
+            extras, "execute_batch", side_effect=meth
+        ) as mock_batch:
+
+            with self.engine.connect() as conn:
+
+                # create an event hook that will change the statement to
+                # something else, meaning the dialect has to detect that
+                # insert_single_values_expr is no longer useful
+                @event.listens_for(conn, "before_cursor_execute", retval=True)
+                def before_cursor_execute(
+                    conn, cursor, statement, parameters, context, executemany
+                ):
+                    statement = (
+                        "INSERT INTO data (id, y, z) VALUES "
+                        "(%(id)s, %(y)s, %(z)s)"
+                    )
+                    return statement, parameters
+
+                conn.execute(
+                    ins,
+                    [
+                        {"id": 1, "y": "y1", "z": 1},
+                        {"id": 2, "y": "y2", "z": 2},
+                        {"id": 3, "y": "y3", "z": 3},
+                    ],
+                )
+
+        eq_(mock_values.mock_calls, [])
+        eq_(
+            mock_batch.mock_calls,
+            [
+                mock.call(
+                    mock.ANY,
+                    "INSERT INTO data (id, y, z) VALUES "
+                    "(%(id)s, %(y)s, %(z)s)",
+                    (
+                        {"id": 1, "y": "y1", "z": 1},
+                        {"id": 2, "y": "y2", "z": 2},
+                        {"id": 3, "y": "y3", "z": 3},
+                    ),
+                )
+            ],
+        )
+
 
 class ExecutemanyFlagOptionsTest(fixtures.TablesTest):
     __only_on__ = "postgresql+psycopg2"