Merge pull request #838 from pyathena-dev/feat/835-insertmanyvalues

Batch SQLAlchemy executemany inserts into multi-row INSERT statements

Author: laughingman7743

docs/sqlalchemy.md

@@ -756,6 +756,34 @@ Athena `FLOAT` and `REAL` are the same 32-bit floating-point type, which keeps a
 Use `Double` (SQLAlchemy 2.0+) for 64-bit values.
 `Float(precision)` does not change the Athena type.
 
+## Bulk inserts
+
+With SQLAlchemy 2.0, an insert executed with a list of parameter sets runs as multi-row `INSERT INTO ... VALUES (...), (...)` statements of up to 100 rows each, instead of one query per row.
+This applies to Core `insert()` with a list of parameters and to ORM flushes that insert several objects.
+`CursorResult.rowcount` is the total number of inserted rows, or -1 if Athena does not report a count.
+If a statement fails, the rows of the earlier statements remain inserted.
+
+Athena limits a query to 262,144 bytes.
+For wide rows, lower the number of rows per statement for the engine or for one execution.
+To run one query per row, disable the batching:
+
+```python
+from sqlalchemy import create_engine
+
+engine = create_engine(
+    "awsathena+rest://:@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://YOUR_S3_BUCKET/path/to/",
+    insertmanyvalues_page_size=20,
+)
+
+with engine.begin() as conn:
+    conn.execution_options(insertmanyvalues_page_size=5).execute(table.insert(), rows)
+
+engine_per_row = create_engine(
+    "awsathena+rest://:@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://YOUR_S3_BUCKET/path/to/",
+    use_insertmanyvalues=False,
+)
+```
+
 ## Complex data types
 
 ### STRUCT type support

pyathena/sqlalchemy/base.py

@@ -13,7 +13,7 @@
 
 from sqlalchemy import exc, schema, types, util
 from sqlalchemy.engine import Engine, reflection
-from sqlalchemy.engine.default import DefaultDialect
+from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
 from sqlalchemy.engine.interfaces import ExecutionContext
 from sqlalchemy.sql.compiler import (
     DDLCompiler,
@@ -153,6 +153,12 @@ class AthenaDialect(DefaultDialect):
     supports_default_values: bool = False
     supports_empty_insert: bool = False
     supports_multivalues_insert: bool = True
+    # Render executemany inserts as multi-row INSERT statements. Athena has no
+    # RETURNING, so batching must also apply to inserts without it. The page
+    # size keeps typical rows well below Athena's 262,144-byte query limit.
+    use_insertmanyvalues: bool = True
+    use_insertmanyvalues_wo_returning: bool = True
+    insertmanyvalues_page_size: int = 100
     supports_sane_rowcount: bool = True
     supports_sane_multi_rowcount: bool = True
     supports_native_decimal: bool = True
@@ -708,6 +714,19 @@ def get_indexes(
         return []  # pragma: no cover
 
     def do_execute(self, cursor, statement, parameters, context=None):
+        """Execute a statement with the DB API cursor.
+
+        SQLAlchemy calls this once per page of an "insertmanyvalues" insert.
+        For those pages, the execution context's row count accumulates the
+        cursor row counts, so that ``CursorResult.rowcount`` reports the total
+        like ``Cursor.executemany``, or -1 if any page has an unknown count.
+
+        Args:
+            cursor: The DB API cursor.
+            statement: The SQL statement.
+            parameters: The statement parameters.
+            context: The SQLAlchemy execution context, if any.
+        """
         on_start_query_execution = None
         if isinstance(context, ExecutionContext):
             execution_options = context.execution_options
@@ -719,6 +738,16 @@ def do_execute(self, cursor, statement, parameters, context=None):
         else:
             cursor.execute(statement, parameters)
 
+        # An executemany context reaches do_execute only for insertmanyvalues
+        # pages; other executemany statements go through do_executemany.
+        if isinstance(context, DefaultExecutionContext) and context.executemany:
+            total = context._rowcount
+            count = cursor.rowcount
+            if total is None:
+                context._rowcount = count
+            else:
+                context._rowcount = total + count if total >= 0 and count >= 0 else -1
+
     def do_rollback(self, dbapi_connection: PoolProxiedConnection) -> None:
         # No transactions for Athena
         pass  # pragma: no cover

tests/pyathena/aio/sqlalchemy/test_base.py

@@ -1,8 +1,12 @@
+from datetime import datetime
+from decimal import Decimal
+
 import pytest
 import sqlalchemy
 from sqlalchemy import cast, literal, select, text, types
-from sqlalchemy.sql.schema import MetaData, Table
+from sqlalchemy.sql.schema import Column, MetaData, Table
 
+from pyathena.sqlalchemy.types import AthenaArray
 from tests import ENV
 from tests.pyathena.util import throttle_metadata_api
 
@@ -229,3 +233,60 @@ async def test_executemany_failure(self, async_engine, executemany_table):
         assert rows == [(1, 11), (2, 21), (3, 30)]
         result = await conn.execute(statement, {"group_id": "2"})
         assert result.rowcount == 1
+
+    async def test_insertmanyvalues(self, async_engine):
+        _, conn = async_engine
+        table_name = "insertmanyvalues_async"
+        table = Table(
+            table_name,
+            MetaData(schema=ENV.schema),
+            Column("id", types.Integer),
+            Column("name", types.String),
+            Column("data", types.LargeBinary),
+            Column("ts", types.DateTime),
+            Column("amount", types.Numeric(10, 3)),
+            Column("tags", AthenaArray(types.Integer)),
+            awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/",
+            awsathena_tblproperties={"table_type": "ICEBERG"},
+        )
+        rows = [
+            {
+                "id": 1,
+                "name": "it's",
+                "data": b"\x00\x01",
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123000),
+                "amount": Decimal("1.5"),
+                "tags": [1, None],
+            },
+            {
+                "id": 2,
+                "name": None,
+                "data": None,
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123456),
+                "amount": Decimal("12.345"),
+                "tags": None,
+            },
+            {"id": 3, "name": "c", "data": b"", "ts": None, "amount": None, "tags": []},
+            {"id": 4, "name": "d", "data": b"d", "ts": None, "amount": None, "tags": [4]},
+            {"id": 5, "name": "e", "data": b"e", "ts": None, "amount": None, "tags": [5, 5]},
+        ]
+
+        statements = []
+
+        def record(conn, cursor, statement, parameters, context, executemany):
+            statements.append(statement)
+
+        await conn.run_sync(table.create)
+        sqlalchemy.event.listen(conn.sync_connection, "before_cursor_execute", record)
+        try:
+            result = await conn.execute(
+                table.insert(), rows, execution_options={"insertmanyvalues_page_size": 2}
+            )
+        finally:
+            sqlalchemy.event.remove(conn.sync_connection, "before_cursor_execute", record)
+
+        # One event per page; a DB API executemany would fire a single event.
+        assert len(statements) == 3
+        assert result.rowcount == 5
+        actual = (await conn.execute(select(table).order_by(table.c.id))).mappings().all()
+        assert [dict(row) for row in actual] == rows

tests/pyathena/sqlalchemy/test_base.py

@@ -22,6 +22,7 @@
 from pyathena.converter import DefaultTypeConverter
 from pyathena.cursor import Cursor
 from pyathena.error import DatabaseError, OperationalError
+from pyathena.formatter import DefaultParameterFormatter
 from pyathena.sqlalchemy.base import AthenaDialect
 from pyathena.sqlalchemy.types import (
     TINYINT,
@@ -60,6 +61,48 @@ def unique_s3tables_table_name(base: str) -> str:
     return f"{base}_{uuid.uuid4().hex[:8]}"
 
 
+def recording_engine(rowcounts=None, **kwargs):
+    """Create an engine whose DB API connection records statements offline.
+
+    Args:
+        rowcounts: Row counts reported by successive cursor calls. By default,
+            a call reports the number of rows it inserted.
+        **kwargs: Additional keyword arguments for ``create_engine``.
+
+    Returns:
+        A tuple of the engine and the list of recorded
+        ``(method, operation, parameters)`` calls.
+    """
+    calls = []
+    counts = iter(rowcounts or ())
+
+    class RecordingCursor:
+        description = None
+        rowcount = -1
+
+        def execute(self, operation, parameters=None, **_):
+            calls.append(("execute", operation, parameters))
+            rows = sum(1 for key in parameters if key.startswith("id"))
+            self.rowcount = next(counts, rows)
+
+        def executemany(self, operation, seq_of_parameters, **_):
+            calls.append(("executemany", operation, seq_of_parameters))
+            self.rowcount = next(counts, len(seq_of_parameters))
+
+        def close(self):
+            pass
+
+    connection = SimpleNamespace(
+        cursor=RecordingCursor, close=lambda: None, commit=lambda: None, rollback=lambda: None
+    )
+    engine = create_engine(
+        "awsathena+rest://athena.us-west-2.amazonaws.com/default",
+        creator=lambda: connection,
+        **kwargs,
+    )
+    return engine, calls
+
+
 class TestAthenaDialect:
     def test_columns_from_information_schema(self):
         # Rows arrive unordered, and Athena reports a missing comment as NULL.
@@ -490,6 +533,108 @@ def list_table_metadata(**kwargs):
             ("pyathena_table_metadata", "other_catalog", "default", table_name): listed
         }
 
+    def test_insertmanyvalues_pages(self):
+        engine, calls = recording_engine()
+        table = Table("t", MetaData(), Column("id", types.Integer), Column("name", types.String))
+
+        with engine.connect() as conn:
+            result = conn.execute(
+                table.insert(), [{"id": i, "name": f"name {i}"} for i in range(250)]
+            )
+
+        # 100 rows per statement by default, and the total row count.
+        assert [(method, len(parameters) // 2) for method, _, parameters in calls] == [
+            ("execute", 100),
+            ("execute", 100),
+            ("execute", 50),
+        ]
+        assert result.rowcount == 250
+        _, operation, parameters = calls[-1]
+        assert (
+            DefaultParameterFormatter()
+            .format(operation, parameters)
+            .startswith("INSERT INTO t (id, name) VALUES (200, 'name 200'), (201, 'name 201'), ")
+        )
+
+    def test_insertmanyvalues_unknown_rowcount(self):
+        engine, _ = recording_engine(rowcounts=[100, -1, 50])
+        table = Table("t", MetaData(), Column("id", types.Integer))
+
+        with engine.connect() as conn:
+            result = conn.execute(table.insert(), [{"id": i} for i in range(250)])
+
+        assert result.rowcount == -1
+
+    @pytest.mark.parametrize("configure", ["engine", "execution_options"])
+    def test_insertmanyvalues_page_size(self, configure):
+        engine_kwargs = {"insertmanyvalues_page_size": 2} if configure == "engine" else {}
+        engine, calls = recording_engine(**engine_kwargs)
+        table = Table("t", MetaData(), Column("id", types.Integer))
+
+        with engine.connect() as conn:
+            if configure == "execution_options":
+                conn = conn.execution_options(insertmanyvalues_page_size=2)
+            result = conn.execute(table.insert(), [{"id": i} for i in range(5)])
+
+        assert [len(parameters) for _, _, parameters in calls] == [2, 2, 1]
+        assert result.rowcount == 5
+
+    def test_insertmanyvalues_disabled(self):
+        engine, calls = recording_engine(use_insertmanyvalues=False)
+        table = Table("t", MetaData(), Column("id", types.Integer))
+
+        with engine.connect() as conn:
+            result = conn.execute(table.insert(), [{"id": i} for i in range(3)])
+
+        ((method, operation, parameters),) = calls
+        assert method == "executemany"
+        assert operation == "INSERT INTO t (id) VALUES (%(id)s)"
+        assert parameters == [{"id": 0}, {"id": 1}, {"id": 2}]
+        assert result.rowcount == 3
+
+    def test_insertmanyvalues_formats_rows(self):
+        engine, calls = recording_engine()
+        table = Table(
+            "t",
+            MetaData(),
+            Column("id", types.Integer),
+            Column("name", types.String),
+            Column("data", types.LargeBinary),
+            Column("ts", types.DateTime),
+            Column("amount", types.Numeric(10, 3)),
+            Column("tags", AthenaArray(types.Integer)),
+        )
+        rows = [
+            {
+                "id": 1,
+                "name": "it's 100%",
+                "data": b"\x00\x01",
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123000),
+                "amount": Decimal("1.5"),
+                "tags": [1, None],
+            },
+            {
+                "id": 2,
+                "name": None,
+                "data": None,
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123456),
+                "amount": None,
+                "tags": None,
+            },
+        ]
+
+        with engine.connect() as conn:
+            conn.execute(table.insert(), rows)
+
+        ((_, operation, parameters),) = calls
+        assert DefaultParameterFormatter().format(operation, parameters) == (
+            "INSERT INTO t (id, name, data, ts, amount, tags) VALUES "
+            "(1, 'it''s 100%', X'0001', TIMESTAMP '2026-01-02 03:04:05.123', DECIMAL '1.5', "
+            "CAST(ARRAY[1, null] AS ARRAY(INTEGER))), "
+            "(2, null, null, TIMESTAMP '2026-01-02 03:04:05.123456', null, "
+            "CAST(null AS ARRAY(INTEGER)))"
+        )
+
 
 class TestSQLAlchemyAthena:
     @pytest.mark.parametrize(
@@ -2951,6 +3096,62 @@ def test_insert_from_select_cte_follows_insert_one(self, engine):
         )
         assert actual == [(2, "bar")]
 
+    def test_insertmanyvalues(self, engine):
+        engine, conn = engine
+        table_name = "insertmanyvalues"
+        table = Table(
+            table_name,
+            MetaData(schema=ENV.schema),
+            Column("id", types.Integer),
+            Column("name", types.String),
+            Column("data", types.LargeBinary),
+            Column("ts", types.DateTime),
+            Column("amount", types.Numeric(10, 3)),
+            Column("tags", AthenaArray(types.Integer)),
+            awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/",
+            awsathena_tblproperties={"table_type": "ICEBERG"},
+        )
+        rows = [
+            {
+                "id": 1,
+                "name": "it's",
+                "data": b"\x00\x01",
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123000),
+                "amount": Decimal("1.5"),
+                "tags": [1, None],
+            },
+            {
+                "id": 2,
+                "name": None,
+                "data": None,
+                "ts": datetime(2026, 1, 2, 3, 4, 5, 123456),
+                "amount": Decimal("12.345"),
+                "tags": None,
+            },
+            {"id": 3, "name": "c", "data": b"", "ts": None, "amount": None, "tags": []},
+            {"id": 4, "name": "d", "data": b"d", "ts": None, "amount": None, "tags": [4]},
+            {"id": 5, "name": "e", "data": b"e", "ts": None, "amount": None, "tags": [5, 5]},
+        ]
+        statements = []
+
+        def record(conn, cursor, statement, parameters, context, executemany):
+            statements.append(statement)
+
+        table.create(bind=conn)
+        sqlalchemy.event.listen(conn, "before_cursor_execute", record)
+        try:
+            result = conn.execution_options(insertmanyvalues_page_size=2).execute(
+                table.insert(), rows
+            )
+        finally:
+            sqlalchemy.event.remove(conn, "before_cursor_execute", record)
+
+        # Rows with different literal precisions share each multi-row statement.
+        assert len(statements) == 3
+        assert result.rowcount == 5
+        actual = conn.execute(sqlalchemy.select(table).order_by(table.c.id)).mappings().all()
+        assert [dict(row) for row in actual] == rows
+
     def test_get_view_definition(self, engine):
         engine, conn = engine
         insp = sqlalchemy.inspect(engine)