From: Mike Bayer Date: Sat, 17 Aug 2019 02:38:51 +0000 (-0400) Subject: Revise psycopg2 execute_values approach X-Git-Tag: rel_1_4_0b1~756^2 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=0a7ca00e04f7c1cfdbb8bdfe7da5f62fc9b40930;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git 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 --- diff --git a/doc/build/changelog/unreleased_13/4623.rst b/doc/build/changelog/unreleased_13/4623.rst new file mode 100644 index 0000000000..462e77f92d --- /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 26014dadce..1a4db1108d 100644 --- 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 8bd2249e27..fa7eeaecf4 100644 --- a/lib/sqlalchemy/sql/compiler.py +++ b/lib/sqlalchemy/sql/compiler.py @@ -498,6 +498,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__( @@ -2512,7 +2521,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 798831cd38..cf33f7cbbf 100644 --- 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"