diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 3b8c536e..d704c3cf 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -53,6 +53,11 @@ jobs: run: | pip install -r requirements.txt + - if: ${{ matrix.python-version == '3.12' || matrix.mariadb }} + name: Install SQLAlchemy + run: | + pip install "SQLAlchemy>=2,<3" + - name: Run tests env: TESTDB: actions.cnf @@ -99,10 +104,11 @@ jobs: wget https://github.com/django/django/archive/${DJANGO_VERSION}.tar.gz tar xf ${DJANGO_VERSION}.tar.gz cp ci/test_mysql.py django-${DJANGO_VERSION}/tests/ + cp ci/test_mysql_executemany_multi.py django-${DJANGO_VERSION}/tests/ cd django-${DJANGO_VERSION} pip install . -r tests/requirements/py3.txt - name: Run Django test run: | cd django-${DJANGO_VERSION}/tests/ - PYTHONPATH=.. python3 ./runtests.py --settings=test_mysql + PYTHONPATH=.. python3 ./runtests.py --settings=test_mysql_executemany_multi diff --git a/ci/test_mysql_executemany_multi.py b/ci/test_mysql_executemany_multi.py new file mode 100644 index 00000000..4c585ce7 --- /dev/null +++ b/ci/test_mysql_executemany_multi.py @@ -0,0 +1,5 @@ +from test_mysql import * # noqa: F403 + + +for database in DATABASES.values(): # noqa: F405 + database.setdefault("OPTIONS", {})["executemany_fallback"] = "multi" diff --git a/doc/user_guide.rst b/doc/user_guide.rst index 391b162f..82080521 100644 --- a/doc/user_guide.rst +++ b/doc/user_guide.rst @@ -72,6 +72,7 @@ MySQL C API function mapping ``mysql_get_server_info()`` ``conn.get_server_info()`` ``mysql_info()`` ``conn.info()`` ``mysql_insert_id()`` ``conn.insert_id()`` + ``mysql_more_results()`` ``conn.more_results()`` ``mysql_num_fields()`` ``result.num_fields()`` ``mysql_num_rows()`` ``result.num_rows()`` ``mysql_options()`` various options to ``_mysql.connect()`` @@ -321,6 +322,27 @@ connect(parameters...) overridden. Default: ``MySQLdb.cursors.Cursor``. *This must be a keyword parameter.* + executemany_fallback + Controls how ``Cursor.executemany()`` handles statements that + cannot use the multi-row INSERT/REPLACE optimization. ``"loop"`` + executes the statements one at a time and is the default. + ``"multi"`` may combine safe data manipulation statements into a + multi-statement query. If multi-statements are disabled or a query + is not eligible, execution silently falls back to ``"loop"``. + + This is a connection option so it can be passed through, for + example, SQLAlchemy's ``connect_args`` or Django's database + ``OPTIONS``:: + + create_engine( + "mysql+mysqldb://user:password@host/database", + connect_args={"executemany_fallback": "multi"}, + ) + + DATABASES["default"]["OPTIONS"]["executemany_fallback"] = "multi" + + See ``executemany()`` below for batching and transaction details. + use_unicode If True, CHAR and VARCHAR and TEXT columns are returned as Unicode strings, using the configured character set. It is @@ -562,6 +584,31 @@ close() close the cursor when you are done with it and before creating a new one. +executemany(operation, seq_of_params) + Executes an operation for every parameter set and returns the total + number of affected rows. Multi-row INSERT and REPLACE statements use + MySQLdb's existing single-statement ``VALUES`` rewrite whenever it + applies, independently of ``executemany_fallback``. + + With ``executemany_fallback="multi"``, statements that do not match that + rewrite may instead be sent in multi-statement batches. This applies only + to SQL templates which begin with INSERT, REPLACE, UPDATE, or DELETE + (ignoring leading whitespace), contain no semicolon, and contain no + ``RETURNING`` clause. Other statements, including statements beginning + with a comment or ``WITH``, use the normal loop. The loop is also used + silently when the connection does not have multi-statements enabled. + + Each batch is limited to 16000 encoded bytes, including separators, and + 200 statements. A single rendered statement exceeding the byte limit is + executed alone. On successful completion, ``rowcount`` and the return + value are the sum of the affected-row counts for all statements. + + Batching does not create an implicit transaction and is not atomic. If a + statement fails, statements before it may already have executed, while + statements after it do not execute. Applications needing all-or-nothing + behavior must manage a transaction explicitly; with autocommit enabled, + each statement may be committed independently. + info() Returns some information about the last query. Normally you don't need to check this. If there are any MySQL diff --git a/src/MySQLdb/_mysql.c b/src/MySQLdb/_mysql.c index 0fbb08d9..97914820 100644 --- a/src/MySQLdb/_mysql.c +++ b/src/MySQLdb/_mysql.c @@ -982,6 +982,28 @@ Returns 0 if there are more results; -1 if there are no more results\n\ \n\ Non-standard.\n\ "; + +static const char _mysql_ConnectionObject_more_results__doc__[] = +"Returns True if one or more results follow the current result of a\n\ +multi-statement query. This check does not advance to the next result.\n\ +\n\ +Non-standard.\n\ +"; + +static PyObject * +_mysql_ConnectionObject_more_results( + _mysql_ConnectionObject *self, + PyObject *noargs) +{ + int ret; + BEGIN_CONNECTION_OPERATION(self, return _mysql_Exception(self)); + ret = mysql_more_results(&(self->connection)); + END_CONNECTION_LOCK(self); + if (ret) + Py_RETURN_TRUE; + Py_RETURN_FALSE; +} + static PyObject * _mysql_ConnectionObject_next_result( _mysql_ConnectionObject *self, @@ -2618,6 +2640,12 @@ static PyMethodDef _mysql_ConnectionObject_methods[] = { METH_NOARGS, _mysql_ConnectionObject_rollback__doc__ }, + { + "more_results", + (PyCFunction)_mysql_ConnectionObject_more_results, + METH_NOARGS, + _mysql_ConnectionObject_more_results__doc__ + }, { "next_result", (PyCFunction)_mysql_ConnectionObject_next_result, diff --git a/src/MySQLdb/connections.py b/src/MySQLdb/connections.py index 08dd696e..98847a65 100644 --- a/src/MySQLdb/connections.py +++ b/src/MySQLdb/connections.py @@ -53,6 +53,7 @@ class Connection(_mysql.connection): """MySQL Database Connection Object""" default_cursor = cursors.Cursor + executemany_fallback = "loop" def __init__(self, *args, **kwargs): """ @@ -123,6 +124,13 @@ class object, used to create cursors (keyword only) If True, enable multi statements for clients >= 4.1. Defaults to True. + :param str executemany_fallback: + Controls how ``Cursor.executemany()`` executes statements which + cannot use the multi-row INSERT/REPLACE optimization. ``"loop"`` + executes each statement separately (the default), while + ``"multi"`` batches safe data manipulation statements into a + multi-statement query when multi statements are enabled. + :param str ssl_mode: specify the security settings for connection to the server; see the MySQL documentation for more details @@ -188,6 +196,13 @@ class object, used to create cursors (keyword only) use_unicode = kwargs2.pop("use_unicode", True) sql_mode = kwargs2.pop("sql_mode", "") self._binary_prefix = kwargs2.pop("binary_prefix", False) + executemany_fallback = kwargs2.pop( + "executemany_fallback", self.executemany_fallback + ) + if executemany_fallback not in ("loop", "multi"): + raise ValueError( + "executemany_fallback must be either 'loop' or 'multi'" + ) client_flag = kwargs.get("client_flag", 0) client_flag |= CLIENT.MULTI_RESULTS @@ -203,6 +218,7 @@ class object, used to create cursors (keyword only) super().__init__(*args, **kwargs2) self.cursorclass = cursorclass + self.executemany_fallback = executemany_fallback self.encoders = {k: v for k, v in conv.items() if type(k) is not int} self._server_version = tuple( [numeric_part(n) for n in self.get_server_info().split(".")[:2]] diff --git a/src/MySQLdb/cursors.py b/src/MySQLdb/cursors.py index 7c6eef28..cf8fe4e5 100644 --- a/src/MySQLdb/cursors.py +++ b/src/MySQLdb/cursors.py @@ -6,7 +6,10 @@ import re -from ._exceptions import ProgrammingError +from ._exceptions import InternalError, MySQLError, ProgrammingError +from .constants import CLIENT, CR + +_EXECUTEMANY_MULTI_SEPARATOR = b"\n;\n" #: Regular expression for ``Cursor.executemany```. #: executemany only supports simple bulk insert. @@ -22,6 +25,41 @@ re.IGNORECASE | re.DOTALL, ) +_RE_INSERT_VALUES_BYTES = re.compile( + RE_INSERT_VALUES.pattern.encode("ascii"), re.IGNORECASE | re.DOTALL +) +_RE_EXECUTEMANY_DML = re.compile( + r"\s*(?:INSERT|REPLACE|UPDATE|DELETE)\b", re.IGNORECASE +) +_RE_EXECUTEMANY_DML_BYTES = re.compile( + _RE_EXECUTEMANY_DML.pattern.encode("ascii"), re.IGNORECASE +) +_RE_RETURNING = re.compile(r"\bRETURNING\b", re.IGNORECASE) +_RE_RETURNING_BYTES = re.compile( + _RE_RETURNING.pattern.encode("ascii"), re.IGNORECASE +) + + +def _match_insert_values(query): + if isinstance(query, (bytes, bytearray)): + return _RE_INSERT_VALUES_BYTES.match(query) + return RE_INSERT_VALUES.match(query) + + +def _is_executemany_dml(query): + """Return whether query is safe for client-side multi-statement batching.""" + if isinstance(query, (bytes, bytearray)): + return ( + b";" not in query + and _RE_EXECUTEMANY_DML_BYTES.match(query) is not None + and _RE_RETURNING_BYTES.search(query) is None + ) + return ( + ";" not in query + and _RE_EXECUTEMANY_DML.match(query) is not None + and _RE_RETURNING.search(query) is None + ) + def _backquote_escape(s): return s.replace(b"`", b"``") @@ -50,6 +88,16 @@ class BaseCursor: #: Default value of max_allowed_packet is 1048576. max_stmt_length = 64 * 1024 + #: Maximum encoded size and statement count for multi-statement + #: ``executemany`` fallback batches. The size includes separators and is + #: measured after argument conversion. Subclasses may override them. + max_multi_stmt_length = 16_000 + max_multi_stmt_count = 200 + + #: Override with ``"loop"`` or ``"multi"`` on a cursor subclass or + #: instance. ``None`` inherits the policy from the connection. + executemany_fallback = None + connection = None def __init__(self, connection): @@ -217,19 +265,21 @@ def executemany(self, query, args): :param args: Sequence of sequences or mappings. It is used as parameter. :return: Number of rows affected, if any. - This method improves performance on multiple-row INSERT and - REPLACE. Otherwise it is equivalent to looping over args with - execute(). + This method improves performance on multiple-row INSERT and REPLACE. + When ``executemany_fallback`` is ``"multi"``, it also batches safe DML + statements if the connection has multi statements enabled. Otherwise, + it is equivalent to looping over args with execute(). """ if not args: return - m = RE_INSERT_VALUES.match(query) + m = _match_insert_values(query) if m: q_prefix = m.group(1) % () q_values = m.group(2).rstrip() q_postfix = m.group(3) or "" - assert q_values[0] == "(" and q_values[-1] == ")" + assert q_values[:1] in ("(", b"(") + assert q_values[-1:] in (")", b")") return self._do_execute_many( q_prefix, q_values, @@ -239,9 +289,129 @@ def executemany(self, query, args): self._get_db().encoding, ) - self.rowcount = sum(self.execute(query, arg) for arg in args) + fallback = self.executemany_fallback + db = self._get_db() + if fallback is None: + fallback = getattr(db, "executemany_fallback", "loop") + if fallback not in ("loop", "multi"): + raise ValueError("executemany_fallback must be either 'loop' or 'multi'") + + if ( + fallback == "multi" + and db.client_flag & CLIENT.MULTI_STATEMENTS + and _is_executemany_dml(query) + ): + return self._do_execute_many_multi(query, args) + + rows = None + for arg in args: + result = self.execute(query, arg) + rows = result if rows is None else rows + result + if rows is None: + return + self.rowcount = rows return self.rowcount + def _do_execute_many_multi(self, query, args): + rows = 0 + statement_count = 0 + sql = bytearray() + + for arg in args: + statement = self._mogrify(query, arg) + if statement_count and ( + statement_count >= self.max_multi_stmt_count + or len(sql) + len(_EXECUTEMANY_MULTI_SEPARATOR) + len(statement) + > self.max_multi_stmt_length + ): + rows += self._execute_multi_statement_batch( + bytes(sql), statement_count + ) + sql.clear() + statement_count = 0 + if statement_count: + sql += _EXECUTEMANY_MULTI_SEPARATOR + sql += statement + statement_count += 1 + if not statement_count: + return + rows += self._execute_multi_statement_batch(bytes(sql), statement_count) + self.rowcount = rows + return rows + + def _execute_multi_statement_batch(self, query, statement_count): + """Execute and fully consume one generated multi-statement query.""" + if statement_count == 1: + return self.execute(query) + + db = self._get_db() + query_started = False + try: + query_started = True + self.execute(query) + if self.description is not None: + self._raise_multi_statement_result_mismatch(db) + rows = self.rowcount + for _ in range(statement_count - 1): + if not db.more_results(): + self._raise_multi_statement_result_mismatch(db) + if db.next_result() != 0: + self._raise_multi_statement_result_mismatch(db) + self._do_get_result(db) + if self.description is not None: + self._raise_multi_statement_result_mismatch(db) + self._post_get_result() + rows += self.rowcount + if db.more_results(): + self._raise_multi_statement_result_mismatch(db) + return rows + except BaseException as exc: + # A server-side SQL error from next_result() terminates the rest of + # the multi-statement query and leaves the protocol synchronized. + # Interruptions and client/protocol failures can leave unread + # results, so discard the connection instead of risking reuse. + self.description = None + self.description_flags = None + self.warning_count = 0 + self.rowcount = None + self.lastrowid = None + self._result = None + self._rows = None + self.rownumber = None + if query_started and self._multi_statement_error_needs_close(exc): + self._close_connection(db) + raise + + def _raise_multi_statement_result_mismatch(self, db): + if self._result is not None: + try: + self._result.discard() + except BaseException: # noqa: S110 + pass + self._result = None + self._close_connection(db) + raise InternalError("multi-statement executemany result count mismatch") + + @staticmethod + def _close_connection(db): + try: + db.close() + except BaseException: # noqa: S110 + pass + + def _multi_statement_error_needs_close(self, exc): + if not isinstance(exc, MySQLError): + return True + if not exc.args or not isinstance(exc.args[0], int): + return True + errno = exc.args[0] + return ( + CR.MIN_ERROR <= errno <= CR.MAX_ERROR + or errno == 1153 # ER_NET_PACKET_TOO_LARGE + or errno == 1927 # ER_CONNECTION_KILLED (MariaDB) + or errno == 4031 # ER_CLIENT_INTERACTION_TIMEOUT + ) + def _do_execute_many( self, prefix, values, postfix, args, max_stmt_length, encoding ): @@ -253,7 +423,11 @@ def _do_execute_many( postfix = postfix.encode(encoding) sql = bytearray(prefix) args = iter(args) - v = self._mogrify(values, next(args)) + try: + first_arg = next(args) + except StopIteration: + return + v = self._mogrify(values, first_arg) sql += v rows = 0 for arg in args: diff --git a/tests/test_connection.py b/tests/test_connection.py index 550750f6..fecb07dd 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -2,9 +2,11 @@ import time import pytest -from configdb import connection_factory +from configdb import connection_factory, connection_kwargs +from MySQLdb.connections import Connection from MySQLdb._exceptions import ProgrammingError +from MySQLdb.constants import CLIENT @pytest.fixture @@ -20,14 +22,18 @@ def test_multi_statements_default_true(conn): cursor = conn.cursor() cursor.execute("select 17; select 2") + assert conn.more_results() is True rows = cursor.fetchall() assert rows == ((17,),) + assert cursor.nextset() == 1 + assert conn.more_results() is False def test_multi_statements_false(): conn = connection_factory(multi_statements=False) try: cursor = conn.cursor() + assert not conn.client_flag & CLIENT.MULTI_STATEMENTS with pytest.raises(ProgrammingError): cursor.execute("select 17; select 2") @@ -39,6 +45,30 @@ def test_multi_statements_false(): conn.close() +def test_executemany_fallback_option(): + with connection_factory() as conn: + assert conn.executemany_fallback == "loop" + + with connection_factory(executemany_fallback="multi") as conn: + assert conn.executemany_fallback == "multi" + + with pytest.raises(ValueError, match="executemany_fallback"): + connection_factory(executemany_fallback="invalid") + + +def test_executemany_fallback_connection_subclass_default(): + class MultiConnection(Connection): + executemany_fallback = "multi" + + with MultiConnection(**connection_kwargs({})) as conn: + assert conn.executemany_fallback == "multi" + + with MultiConnection( + **connection_kwargs({"executemany_fallback": "loop"}) + ) as conn: + assert conn.executemany_fallback == "loop" + + def test_ping_false_warns(conn): with pytest.warns(DeprecationWarning, match="reconnect parameter"): conn.ping(False) diff --git a/tests/test_cursor.py b/tests/test_cursor.py index 3c981930..3a9b2af5 100644 --- a/tests/test_cursor.py +++ b/tests/test_cursor.py @@ -1,10 +1,13 @@ from textwrap import dedent +from types import SimpleNamespace import pytest from configdb import connection_factory import MySQLdb.cursors -from MySQLdb.constants import ER +from MySQLdb._exceptions import IntegrityError, InternalError, OperationalError +from MySQLdb.constants import CLIENT, ER +from MySQLdb.converters import conversions _conns = [] _tables = [] @@ -81,6 +84,14 @@ def test_executemany(): "execute many with %s not in one query" ) + # bytes and bytearray queries use the same INSERT/REPLACE fast path. + cursor.executemany(b"insert into test (data) values (%s)", [(10,), (11,)]) + assert cursor._executed.endswith(b"(10),(11)") + cursor.executemany( + bytearray(b"insert into test (data) values (%s)"), [(12,), (13,)] + ) + assert cursor._executed.endswith(b"(12),(13)") + # dict args data_dict = [{"data": i} for i in range(10)] cursor.executemany("insert into test (data) values (%(data)s)", data_dict) @@ -106,6 +117,463 @@ def test_executemany(): cursor.execute("DROP TABLE IF EXISTS percent_test") +@pytest.mark.parametrize( + "Cursor", [MySQLdb.cursors.Cursor, MySQLdb.cursors.SSCursor] +) +def test_executemany_multi_update(Cursor): + conn = connect(executemany_fallback="multi") + cursor = conn.cursor(Cursor) + cursor.execute( + "CREATE TABLE executemany_multi_update " + "(id int primary key, data varchar(100))" + ) + _tables.append("executemany_multi_update") + cursor.executemany( + "INSERT INTO executemany_multi_update (id, data) VALUES (%s, %s)", + [(1, 0), (2, 0), (3, 0)], + ) + assert MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR not in cursor._executed + assert b"),(" in cursor._executed + + rows = cursor.executemany( + "UPDATE executemany_multi_update " + "SET data=%(data)s WHERE id=%(id)s", + [ + {"id": 1, "data": "ten;still-a-value"}, + {"id": 2, "data": "twenty"}, + {"id": 3, "data": "thirty"}, + ], + ) + + assert rows == 3 + assert cursor.rowcount == 3 + assert cursor.description is None + assert cursor._executed.count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) == 2 + assert conn.affected_rows() == 1 + assert conn.warning_count() == 0 + assert conn.more_results() is False + + cursor.execute("SELECT id, data FROM executemany_multi_update ORDER BY id") + assert cursor.fetchall() == ( + (1, "ten;still-a-value"), + (2, "twenty"), + (3, "thirty"), + ) + + +def test_executemany_multi_delete(): + conn = connect(executemany_fallback="multi") + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE executemany_multi_delete (id int primary key, data int)" + ) + _tables.append("executemany_multi_delete") + cursor.executemany( + "INSERT INTO executemany_multi_delete (id, data) VALUES (%s, %s)", + [(1, 10), (2, 20), (3, 30)], + ) + + assert ( + cursor.executemany( + "DELETE FROM executemany_multi_delete WHERE id=%s", [(1,), (3,)] + ) + == 2 + ) + assert cursor.rowcount == 2 + assert cursor._executed.count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) == 1 + assert conn.affected_rows() == 1 + cursor.execute("SELECT id FROM executemany_multi_delete") + assert cursor.fetchall() == ((2,),) + + +def test_executemany_multi_keeps_last_statement_metadata(): + conn = connect(executemany_fallback="multi") + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE executemany_multi_metadata " + "(id int primary key auto_increment, data varchar(1))" + ) + _tables.append("executemany_multi_metadata") + + assert ( + cursor.executemany( + "INSERT IGNORE INTO executemany_multi_metadata SET data=%s", + [("a",), ("b",), ("too long",)], + ) + == 3 + ) + assert cursor.rowcount == 3 + assert cursor.lastrowid == 3 + assert conn.insert_id() == 3 + assert conn.affected_rows() == 1 + assert conn.warning_count() > 0 + assert conn.more_results() is False + + +def test_executemany_multi_policy_and_capability(): + conn = connect() + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE executemany_multi_policy (id int primary key, data int)" + ) + _tables.append("executemany_multi_policy") + cursor.executemany( + "INSERT INTO executemany_multi_policy (id, data) VALUES (%s, %s)", + [(1, 0), (2, 0)], + ) + + query = "UPDATE executemany_multi_policy SET data=%s WHERE id=%s" + cursor.executemany(query, [(10, 1), (20, 2)]) + assert MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR not in cursor._executed + + class MultiCursor(MySQLdb.cursors.Cursor): + executemany_fallback = "multi" + + subclass_cursor = conn.cursor(MultiCursor) + subclass_cursor.executemany(query, [(11, 1), (21, 2)]) + assert ( + subclass_cursor._executed.count( + MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR + ) + == 1 + ) + + cursor.executemany_fallback = "multi" + cursor.executemany(query, [(12, 1), (22, 2)]) + assert cursor._executed.count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) == 1 + assert cursor.executemany( + query + " -- trailing comment", [(13, 1), (23, 2)] + ) == 2 + assert cursor._executed.count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) == 1 + + cursor.executemany_fallback = "loop" + cursor.executemany(query, [(14, 1), (24, 2)]) + assert MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR not in cursor._executed + + with pytest.raises(ValueError, match="executemany_fallback"): + cursor.executemany_fallback = "invalid" + cursor.executemany(query, [(15, 1), (25, 2)]) + + conn.commit() + no_multi_conn = connect( + executemany_fallback="multi", multi_statements=False + ) + no_multi_cursor = no_multi_conn.cursor() + no_multi_cursor.executemany(query, [(16, 1), (26, 2)]) + assert ( + MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR + not in no_multi_cursor._executed + ) + no_multi_conn.rollback() + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("UPDATE t SET value=%s", True), + (b"DELETE FROM t WHERE id=%s", True), + (bytearray(b"UPDATE t SET value=%s"), True), + ("INSERT INTO t SET value=%s", True), + ("REPLACE INTO t SET value=%s", True), + ("WITH values_ AS (SELECT 1) UPDATE t SET value=%s", False), + ("UPDATE t SET value=%s;", False), + ("UPDATE t SET value=%s RETURNING id", False), + ("SELECT %s", False), + ("/* comment */ UPDATE t SET value=%s", False), + ], +) +def test_is_executemany_dml(query, expected): + assert MySQLdb.cursors._is_executemany_dml(query) is expected + + +def test_executemany_multi_batch_limits_and_single_arg(): + class RecordingCursor(MySQLdb.cursors.Cursor): + max_multi_stmt_length = 1_000_000 + max_multi_stmt_count = 2 + + def __init__(self, connection): + super().__init__(connection) + self.execute_calls = [] + + def execute(self, query, args=None): + self.execute_calls.append((query, args)) + return super().execute(query, args) + + conn = connect(executemany_fallback="multi") + cursor = conn.cursor(RecordingCursor) + cursor.execute( + "CREATE TABLE executemany_multi_limits " + "(id int primary key, data text)" + ) + _tables.append("executemany_multi_limits") + cursor.executemany( + "INSERT INTO executemany_multi_limits (id, data) VALUES (%s, %s)", + [(i, 0) for i in range(1, 7)], + ) + + query = "UPDATE executemany_multi_limits SET data=%s WHERE id=%s" + cursor.execute_calls.clear() + assert cursor.executemany(query, [(i * 10, i) for i in range(1, 6)]) == 5 + assert len(cursor.execute_calls) == 3 + assert [ + bytes(q).count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) + for q, args in cursor.execute_calls + ] == [1, 1, 0] + assert all(args is None for query, args in cursor.execute_calls) + assert cursor.rowcount == 5 + assert conn.affected_rows() == 1 + + first_arg = ("a", 1) + second_base_arg = ("", 2) + first_statement = cursor._mogrify(query, first_arg) + second_base_statement = cursor._mogrify(query, second_base_arg) + filler_length = ( + 16_000 + - len(first_statement) + - len(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) + - len(second_base_statement) + ) + boundary_args = [first_arg, ("x" * filler_length, 2), ("c", 3)] + cursor.max_multi_stmt_count = 200 + cursor.max_multi_stmt_length = 16_000 + cursor.execute_calls.clear() + assert cursor.executemany(query, boundary_args) == 3 + assert len(cursor.execute_calls) == 2 + assert len(cursor.execute_calls[0][0]) == 16_000 + + cursor.max_multi_stmt_length = 1_000_000 + cursor.max_multi_stmt_count = 200 + cursor.execute_calls.clear() + assert ( + cursor.executemany( + "DELETE FROM executemany_multi_limits WHERE id=%s", + [(1000 + i,) for i in range(201)], + ) + == 0 + ) + assert len(cursor.execute_calls) == 2 + assert [ + bytes(q).count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) + for q, args in cursor.execute_calls + ] == [199, 0] + + cursor.execute_calls.clear() + arg = (60, 6) + assert cursor.executemany(query, [arg]) == 1 + assert cursor.execute_calls == [(cursor._mogrify(query, arg), None)] + + assert MySQLdb.cursors.BaseCursor.max_multi_stmt_length == 16_000 + assert MySQLdb.cursors.BaseCursor.max_multi_stmt_count == 200 + + +def test_executemany_multi_streams_arguments(): + class RecordingCursor(MySQLdb.cursors.Cursor): + max_multi_stmt_count = 2 + + def _mogrify(self, query, args): + return (query % args).encode() + + def execute(self, query, args=None): + calls.append(query) + self.rowcount = 1 + return 1 + + def _execute_multi_statement_batch(self, query, statement_count): + if statement_count == 1: + return super()._execute_multi_statement_batch(query, statement_count) + calls.append(query) + return statement_count + + calls = [] + cursor = RecordingCursor( + SimpleNamespace( + executemany_fallback="multi", client_flag=CLIENT.MULTI_STATEMENTS + ) + ) + + def params(): + yield (1,) + yield (2,) + yield (3,) + assert len(calls) == 1 + yield (4,) + yield (5,) + + query = "UPDATE t SET value=%s" + assert cursor.executemany(query, params()) == 5 + assert len(calls) == 3 + assert calls[-1] == b"UPDATE t SET value=5" + assert cursor.rowcount == 5 + assert cursor.executemany(query, iter(())) is None + assert len(calls) == 3 + assert cursor.rowcount == 5 + + calls.clear() + cursor.max_multi_stmt_length = 1 + assert cursor.executemany(query, iter([(6,), (7,)])) == 2 + assert calls == [b"UPDATE t SET value=6", b"UPDATE t SET value=7"] + + +def test_executemany_multi_oversized_statement_runs_alone(): + class TinyBatchCursor(MySQLdb.cursors.Cursor): + max_multi_stmt_length = 1 + + def __init__(self, connection): + super().__init__(connection) + self.execute_calls = [] + + def execute(self, query, args=None): + self.execute_calls.append((query, args)) + return super().execute(query, args) + + conn = connect(executemany_fallback="multi") + cursor = conn.cursor(TinyBatchCursor) + cursor.execute( + "CREATE TABLE executemany_multi_oversized (id int primary key, data int)" + ) + _tables.append("executemany_multi_oversized") + cursor.execute_calls.clear() + + query = "UPDATE executemany_multi_oversized SET data=%s WHERE id=%s" + cursor.executemany(query, [(10, 1), (20, 2)]) + assert len(cursor.execute_calls) == 2 + assert all( + MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR not in bytes(q) + for q, args in cursor.execute_calls + ) + + +@pytest.mark.parametrize("fallback", ["loop", "multi"]) +def test_executemany_multi_generator_and_empty_generator(fallback): + conn = connect(executemany_fallback=fallback) + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE executemany_multi_generator (id int primary key, data int)" + ) + _tables.append("executemany_multi_generator") + insert = "INSERT INTO executemany_multi_generator (id, data) VALUES (%s, %s)" + assert cursor.executemany(insert, ((i, 0) for i in range(1, 4))) == 3 + assert cursor.executemany(insert, iter(())) is None + assert cursor.rowcount == 3 + + query = "UPDATE executemany_multi_generator SET data=%s WHERE id=%s" + params = ((i * 10, i) for i in range(1, 4)) + assert cursor.executemany(query, params) == 3 + if fallback == "multi": + assert cursor._executed.count(MySQLdb.cursors._EXECUTEMANY_MULTI_SEPARATOR) == 2 + assert cursor.executemany(query, iter(())) is None + assert cursor.rowcount == 3 + + cursor.execute("SELECT id, data FROM executemany_multi_generator ORDER BY id") + assert cursor.fetchall() == ((1, 10), (2, 20), (3, 30)) + + +@pytest.mark.parametrize( + ("args", "expected_ids"), + [ + ([(99, 10), (1, 10), (2, 20)], (99,)), + ([(1, 10), (99, 10), (2, 20)], (1, 99)), + ([(1, 10), (2, 20), (99, 10)], (1, 2, 99)), + ], +) +def test_executemany_multi_sql_error(args, expected_ids): + conn = connect(executemany_fallback="multi") + cursor = conn.cursor() + cursor.execute( + "CREATE TABLE executemany_multi_error (id int primary key, data int)" + ) + _tables.append("executemany_multi_error") + cursor.execute("INSERT INTO executemany_multi_error VALUES (99, 0)") + + with pytest.raises(IntegrityError): + cursor.executemany( + "INSERT INTO executemany_multi_error SET id=%s, data=%s", args + ) + + assert cursor.rowcount is None + assert conn.open + assert conn.more_results() is False + cursor.execute("SELECT id FROM executemany_multi_error ORDER BY id") + assert tuple(row[0] for row in cursor.fetchall()) == expected_ids + + +@pytest.mark.parametrize( + "Cursor", [MySQLdb.cursors.Cursor, MySQLdb.cursors.SSCursor] +) +def test_executemany_multi_rejects_unexpected_result_count(Cursor): + class RawSQL: + pass + + def raw_sql_literal(value, conv): + return b"1; SELECT 1" + + cleanup_conn = connect() + cleanup_cursor = cleanup_conn.cursor() + cleanup_cursor.execute( + "CREATE TABLE executemany_multi_result_count (id int primary key, data int)" + ) + _tables.append("executemany_multi_result_count") + cleanup_cursor.execute( + "INSERT INTO executemany_multi_result_count VALUES (1, 0)" + ) + cleanup_conn.commit() + + custom_conversions = conversions.copy() + custom_conversions[RawSQL] = raw_sql_literal + conn = connect(executemany_fallback="multi", conv=custom_conversions) + cursor = conn.cursor(Cursor) + + with pytest.raises(InternalError, match="multi-statement executemany"): + cursor.executemany( + "UPDATE executemany_multi_result_count SET data=%s", + [(RawSQL(),), (RawSQL(),)], + ) + + assert not conn.open + _conns.remove(conn) + + +@pytest.mark.parametrize( + "failure", [KeyboardInterrupt(), OperationalError(2013, "server lost")] +) +def test_executemany_multi_drain_failure_closes_connection(failure): + class FailingCursor(MySQLdb.cursors.Cursor): + armed = False + result_number = 0 + + def _do_get_result(self, db): + super()._do_get_result(db) + if self.armed: + self.result_number += 1 + if self.result_number == 2: + raise failure + + cleanup_conn = connect() + cleanup_cursor = cleanup_conn.cursor() + cleanup_cursor.execute( + "CREATE TABLE executemany_multi_drain_failure " + "(id int primary key, data int)" + ) + _tables.append("executemany_multi_drain_failure") + cleanup_cursor.execute( + "INSERT INTO executemany_multi_drain_failure VALUES (1, 0), (2, 0)" + ) + cleanup_conn.commit() + + conn = connect(executemany_fallback="multi") + cursor = conn.cursor(FailingCursor) + cursor.armed = True + with pytest.raises(type(failure)) as exc_info: + cursor.executemany( + "UPDATE executemany_multi_drain_failure SET data=%s WHERE id=%s", + [(10, 1), (20, 2)], + ) + + assert exc_info.value is failure + assert not conn.open + _conns.remove(conn) + + def test_pyparam(): conn = connect() cursor = conn.cursor() diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py new file mode 100644 index 00000000..fb99a174 --- /dev/null +++ b/tests/test_sqlalchemy.py @@ -0,0 +1,151 @@ +from contextlib import contextmanager + +import pytest + + +pytest.importorskip("sqlalchemy", minversion="2.0") + +from sqlalchemy import ( + Integer, + bindparam, + create_engine, + delete, + event, + insert, + select, + update, +) +from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column +from sqlalchemy.pool import NullPool + +from MySQLdb.constants import CLIENT +from configdb import connection_kwargs + + +class Base(DeclarativeBase): + pass + + +class BulkRow(Base): + __tablename__ = "test_sqlalchemy_executemany" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + value: Mapped[int] = mapped_column(Integer, nullable=False) + + +@pytest.fixture(scope="module") +def engine(): + engine = create_engine( + "mysql+mysqldb://", + connect_args=connection_kwargs({"executemany_fallback": "multi"}), + poolclass=NullPool, + ) + Base.metadata.drop_all(engine) + Base.metadata.create_all(engine) + yield engine + Base.metadata.drop_all(engine) + engine.dispose() + + +def reset_rows(engine): + with engine.begin() as connection: + connection.execute(delete(BulkRow)) + connection.execute( + insert(BulkRow), + [ + {"id": 1, "value": 10}, + {"id": 2, "value": 20}, + {"id": 3, "value": 30}, + ], + ) + + +@contextmanager +def capture_executemany(engine): + calls = [] + + def after_cursor_execute( + connection, cursor, statement, parameters, context, executemany + ): + calls.append( + { + "statement": statement, + "executemany": executemany, + "rowcount": cursor.rowcount, + "executed": cursor._executed, + } + ) + + event.listen(engine, "after_cursor_execute", after_cursor_execute) + try: + yield calls + finally: + event.remove(engine, "after_cursor_execute", after_cursor_execute) + + +def assert_executemany_call(calls, operation, rowcount): + calls = [ + call + for call in calls + if call["statement"].lstrip().upper().startswith(operation) + ] + assert len(calls) == 1 + assert calls[0]["executemany"] is True + assert calls[0]["rowcount"] == rowcount + # The ORM passed one statement template to DB-API executemany(), while + # mysqlclient sent the rendered statements in one multi-statement query. + assert b";" in calls[0]["executed"] + + +def test_connect_args_enable_multi_fallback_and_found_rows(engine): + with engine.connect() as connection: + driver_connection = connection.connection.driver_connection + assert driver_connection.executemany_fallback == "multi" + assert driver_connection.client_flag & CLIENT.FOUND_ROWS + + +def test_bulk_update_mappings_uses_executemany(engine): + reset_rows(engine) + + with capture_executemany(engine) as calls, Session(engine) as session: + session.bulk_update_mappings( + BulkRow, + [ + {"id": 1, "value": 10}, # no-op; FOUND_ROWS still counts it + {"id": 2, "value": 21}, + ], + ) + session.commit() + + assert_executemany_call(calls, "UPDATE ", 2) + + +def test_orm_bulk_update_by_primary_key_uses_executemany(engine): + reset_rows(engine) + + with capture_executemany(engine) as calls, Session(engine) as session: + session.execute( + update(BulkRow), + [ + {"id": 1, "value": 11}, + {"id": 2, "value": 22}, + ], + ) + session.commit() + + assert_executemany_call(calls, "UPDATE ", 2) + + +def test_core_executemany_delete_rowcount(engine): + reset_rows(engine) + + with capture_executemany(engine) as calls, engine.begin() as connection: + result = connection.execute( + delete(BulkRow).where(BulkRow.id == bindparam("target_id")), + [{"target_id": 1}, {"target_id": 99}, {"target_id": 3}], + ) + assert result.rowcount == 2 + + assert_executemany_call(calls, "DELETE ", 2) + with engine.connect() as connection: + assert connection.scalars(select(BulkRow.id)).all() == [2]