diff options
Diffstat (limited to 'test/dialect/postgresql/test_dialect.py')
| -rw-r--r-- | test/dialect/postgresql/test_dialect.py | 440 |
1 files changed, 15 insertions, 425 deletions
diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py index fdb114d57..27d4a4cf9 100644 --- a/test/dialect/postgresql/test_dialect.py +++ b/test/dialect/postgresql/test_dialect.py @@ -1,7 +1,6 @@ # coding: utf-8 import dataclasses import datetime -import itertools import logging import logging.handlers @@ -18,7 +17,6 @@ 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 @@ -32,15 +30,14 @@ from sqlalchemy import text from sqlalchemy import TypeDecorator from sqlalchemy.dialects.postgresql import base as postgresql from sqlalchemy.dialects.postgresql import HSTORE -from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.dialects.postgresql import psycopg as psycopg_dialect from sqlalchemy.dialects.postgresql import psycopg2 as psycopg2_dialect from sqlalchemy.dialects.postgresql import Range -from sqlalchemy.dialects.postgresql.psycopg2 import EXECUTEMANY_BATCH -from sqlalchemy.dialects.postgresql.psycopg2 import EXECUTEMANY_PLAIN from sqlalchemy.dialects.postgresql.psycopg2 import EXECUTEMANY_VALUES -from sqlalchemy.engine import cursor as _cursor +from sqlalchemy.dialects.postgresql.psycopg2 import ( + EXECUTEMANY_VALUES_PLUS_BATCH, +) from sqlalchemy.engine import url from sqlalchemy.sql.selectable import LABEL_STYLE_TABLENAME_PLUS_COL from sqlalchemy.testing import config @@ -60,11 +57,6 @@ from sqlalchemy.testing.assertions import eq_regex from sqlalchemy.testing.assertions import expect_raises from sqlalchemy.testing.assertions import ne_ -if True: - from sqlalchemy.dialects.postgresql.psycopg2 import ( - EXECUTEMANY_VALUES_PLUS_BATCH, - ) - class DialectTest(fixtures.TestBase): """python-side dialect tests.""" @@ -451,179 +443,6 @@ class ExecuteManyMode: Column("\u6e2c\u8a66", Integer), ) - @testing.combinations( - "insert", "pg_insert", "pg_insert_on_conflict", argnames="insert_type" - ) - def test_insert(self, connection, insert_type): - from psycopg2 import extras - - values_page_size = connection.dialect.executemany_values_page_size - batch_page_size = connection.dialect.executemany_batch_page_size - if connection.dialect.executemany_mode & EXECUTEMANY_VALUES: - meth = extras.execute_values - stmt = "INSERT INTO data (x, y) VALUES %s" - expected_kwargs = { - "template": "(%(x)s, %(y)s)", - "page_size": values_page_size, - "fetch": False, - } - elif connection.dialect.executemany_mode & EXECUTEMANY_BATCH: - meth = extras.execute_batch - stmt = "INSERT INTO data (x, y) VALUES (%(x)s, %(y)s)" - expected_kwargs = {"page_size": batch_page_size} - else: - assert False - - if insert_type == "pg_insert_on_conflict": - stmt += " ON CONFLICT DO NOTHING" - - with mock.patch.object( - extras, meth.__name__, side_effect=meth - ) as mock_exec: - if insert_type == "insert": - ins_stmt = self.tables.data.insert() - elif insert_type == "pg_insert": - ins_stmt = pg_insert(self.tables.data) - elif insert_type == "pg_insert_on_conflict": - ins_stmt = pg_insert(self.tables.data).on_conflict_do_nothing() - else: - assert False - - connection.execute( - ins_stmt, - [ - {"x": "x1", "y": "y1"}, - {"x": "x2", "y": "y2"}, - {"x": "x3", "y": "y3"}, - ], - ) - - eq_( - connection.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, connection): - from psycopg2 import extras - - values_page_size = connection.dialect.executemany_values_page_size - batch_page_size = connection.dialect.executemany_batch_page_size - - if connection.dialect.executemany_mode & EXECUTEMANY_VALUES: - meth = extras.execute_values - stmt = "INSERT INTO data (x, y) VALUES %s" - expected_kwargs = { - "template": "(%(x)s, %(y)s)", - "page_size": values_page_size, - "fetch": False, - } - elif connection.dialect.executemany_mode & EXECUTEMANY_BATCH: - meth = extras.execute_batch - stmt = "INSERT INTO data (x, y) VALUES (%(x)s, %(y)s)" - expected_kwargs = {"page_size": batch_page_size} - else: - assert False - - with mock.patch.object( - extras, meth.__name__, side_effect=meth - ) as mock_exec: - connection.execute( - self.tables.data.insert(), - [ - {"x": "x1", "y": "y1"}, - {"x": "x2", "y": "y2"}, - {"x": "x3", "y": "y3"}, - ], - ) - - 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_page_size(self): - from psycopg2 import extras - - opts = self.options.copy() - opts["executemany_batch_page_size"] = 500 - opts["executemany_values_page_size"] = 1000 - - eng = engines.testing_engine(options=opts) - - if eng.dialect.executemany_mode & EXECUTEMANY_VALUES: - meth = extras.execute_values - stmt = "INSERT INTO data (x, y) VALUES %s" - expected_kwargs = { - "fetch": False, - "page_size": 1000, - "template": "(%(x)s, %(y)s)", - } - elif eng.dialect.executemany_mode & EXECUTEMANY_BATCH: - meth = extras.execute_batch - stmt = "INSERT INTO data (x, y) VALUES (%(x)s, %(y)s)" - expected_kwargs = {"page_size": 500} - else: - assert False - - with mock.patch.object( - extras, meth.__name__, side_effect=meth - ) as mock_exec: - with eng.begin() as conn: - conn.execute( - self.tables.data.insert(), - [ - {"x": "x1", "y": "y1"}, - {"x": "x2", "y": "y2"}, - {"x": "x3", "y": "y3"}, - ], - ) - - 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_unicode_keys(self, connection): table = self.tables["Unitéble2"] @@ -661,7 +480,10 @@ class ExecuteManyMode: ], ) - if connection.dialect.executemany_mode & EXECUTEMANY_BATCH: + if ( + connection.dialect.executemany_mode + is EXECUTEMANY_VALUES_PLUS_BATCH + ): eq_( mock_exec.mock_calls, [ @@ -680,7 +502,10 @@ class ExecuteManyMode: eq_(mock_exec.mock_calls, []) def test_not_sane_rowcount(self, connection): - if connection.dialect.executemany_mode & EXECUTEMANY_BATCH: + if ( + connection.dialect.executemany_mode + is EXECUTEMANY_VALUES_PLUS_BATCH + ): assert not connection.dialect.supports_sane_multi_rowcount else: assert connection.dialect.supports_sane_multi_rowcount @@ -709,257 +534,22 @@ class ExecuteManyMode: ) -class ExecutemanyBatchModeTest(ExecuteManyMode, fixtures.TablesTest): - options = {"executemany_mode": "batch"} - - -class ExecutemanyValuesInsertsTest(ExecuteManyMode, fixtures.TablesTest): - options = {"executemany_mode": "values_only"} - - def test_insert_returning_values(self, connection): - """the psycopg2 dialect needs to assemble a fully buffered result - with the return value of execute_values(). - - """ - t = self.tables.data - - conn = connection - page_size = conn.dialect.executemany_values_page_size or 100 - data = [ - {"x": "x%d" % i, "y": "y%d" % i} - for i in range(1, page_size * 5 + 27) - ] - result = conn.execute(t.insert().returning(t.c.x, t.c.y), data) - - eq_([tup[0] for tup in result.cursor.description], ["x", "y"]) - eq_(result.keys(), ["x", "y"]) - assert t.c.x in result.keys() - assert t.c.id not in result.keys() - assert not result._soft_closed - assert isinstance( - result.cursor_strategy, - _cursor.FullyBufferedCursorFetchStrategy, - ) - assert not result.cursor.closed - assert not result.closed - eq_(result.mappings().all(), data) - - assert result._soft_closed - # assert result.closed - assert result.cursor is None - - def test_insert_returning_preexecute_pk(self, metadata, connection): - counter = itertools.count(1) - - t = Table( - "t", - self.metadata, - Column( - "id", - Integer, - primary_key=True, - default=lambda: next(counter), - ), - Column("data", Integer), - ) - metadata.create_all(connection) - - result = connection.execute( - t.insert().return_defaults(), - [{"data": 1}, {"data": 2}, {"data": 3}], - ) - - eq_(result.inserted_primary_key_rows, [(1,), (2,), (3,)]) - - def test_insert_returning_defaults(self, connection): - t = self.tables.data - - conn = connection - - result = conn.execute(t.insert(), {"x": "x0", "y": "y0"}) - first_pk = result.inserted_primary_key[0] - - page_size = conn.dialect.executemany_values_page_size or 100 - total_rows = page_size * 5 + 27 - data = [{"x": "x%d" % i, "y": "y%d" % i} for i in range(1, total_rows)] - result = conn.execute(t.insert().returning(t.c.id, t.c.z), data) - - eq_( - result.all(), - [(pk, 5) for pk in range(1 + first_pk, total_rows + first_pk)], - ) - - def test_insert_return_pks_default_values(self, connection): - """test sending multiple, empty rows into an INSERT and getting primary - key values back. - - This has to use a format that indicates at least one DEFAULT in - multiple parameter sets, i.e. "INSERT INTO table (anycol) VALUES - (DEFAULT) (DEFAULT) (DEFAULT) ... RETURNING col" - - """ - t = self.tables.data - - conn = connection - - result = conn.execute(t.insert(), {"x": "x0", "y": "y0"}) - first_pk = result.inserted_primary_key[0] - - page_size = conn.dialect.executemany_values_page_size or 100 - total_rows = page_size * 5 + 27 - data = [{} for i in range(1, total_rows)] - result = conn.execute(t.insert().returning(t.c.id), data) - - eq_( - result.all(), - [(pk,) for pk in range(1 + first_pk, total_rows + first_pk)], - ) - - def test_insert_w_newlines(self, connection): - from psycopg2 import extras - - t = self.tables.data - - ins = ( - t.insert() - .inline() - .values( - id=bindparam("id"), - x=select(literal_column("5")) - .select_from(self.tables.data) - .scalar_subquery(), - 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: - - connection.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)", - fetch=False, - page_size=connection.dialect.executemany_values_page_size, - ) - ], - ) - - def test_insert_modified_by_event(self, connection): - from psycopg2 import extras - - t = self.tables.data - - ins = ( - t.insert() - .inline() - .values( - id=bindparam("id"), - x=select(literal_column("5")) - .select_from(self.tables.data) - .scalar_subquery(), - 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: - - # 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( - connection, "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 - - connection.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, []) - - if connection.dialect.executemany_mode & EXECUTEMANY_BATCH: - 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}, - ), - ) - ], - ) - else: - eq_(mock_batch.mock_calls, []) - - class ExecutemanyValuesPlusBatchInsertsTest( ExecuteManyMode, fixtures.TablesTest ): options = {"executemany_mode": "values_plus_batch"} +class ExecutemanyValuesInsertsTest(ExecuteManyMode, fixtures.TablesTest): + options = {"executemany_mode": "values_only"} + + class ExecutemanyFlagOptionsTest(fixtures.TablesTest): __only_on__ = "postgresql+psycopg2" __backend__ = True def test_executemany_correct_flag_options(self): for opt, expected in [ - (None, EXECUTEMANY_PLAIN), - ("batch", EXECUTEMANY_BATCH), ("values_only", EXECUTEMANY_VALUES), ("values_plus_batch", EXECUTEMANY_VALUES_PLUS_BATCH), ]: |