Add insert_multiple_mode="executemany"

Author: bdowning

sql_athame/dataclasses.py

@@ -157,7 +157,7 @@ class ModelBase:
     _cache: dict[tuple, Any]
     table_name: str
     primary_key_names: tuple[str, ...]
-    array_safe_insert: bool
+    insert_multiple_mode: str
 
     def __init_subclass__(
         cls,
@@ -169,12 +169,9 @@ def __init_subclass__(
     ):
         cls._cache = {}
         cls.table_name = table_name
-        if insert_multiple_mode == "array_safe":
-            cls.array_safe_insert = True
-        elif insert_multiple_mode == "unnest":
-            cls.array_safe_insert = False
-        else:
+        if insert_multiple_mode not in ("array_safe", "unnest", "executemany"):
             raise ValueError("Unknown `insert_multiple_mode`")
+        cls.insert_multiple_mode = insert_multiple_mode
         if isinstance(primary_key, str):
             cls.primary_key_names = (primary_key,)
         else:
@@ -506,6 +503,37 @@ def insert_multiple_array_safe_sql(cls: type[T], rows: Iterable[T]) -> Fragment:
             ),
         )
 
+    @classmethod
+    def insert_multiple_executemany_chunk_sql(
+        cls: type[T], chunk_size: int
+    ) -> Fragment:
+        def generate() -> Fragment:
+            columns = len(cls.column_info())
+            values = ", ".join(
+                f"({', '.join(f'${i}' for i in chunk)})"
+                for chunk in chunked(range(1, columns * chunk_size + 1), columns)
+            )
+            return sql(
+                "INSERT INTO {table} ({fields}) VALUES {values}",
+                table=cls.table_name_sql(),
+                fields=sql.list(cls.field_names_sql()),
+                values=sql.literal(values),
+            ).flatten()
+
+        return cls._cached(
+            ("insert_multiple_executemany_chunk", chunk_size),
+            generate,
+        )
+
+    @classmethod
+    async def insert_multiple_executemany(
+        cls: type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+    ) -> None:
+        args = [r.field_values() for r in rows]
+        query = cls.insert_multiple_executemany_chunk_sql(1).query()[0]
+        if args:
+            await connection_or_pool.executemany(query, args)
+
     @classmethod
     async def insert_multiple_unnest(
         cls: type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
@@ -527,11 +555,28 @@ async def insert_multiple_array_safe(
     async def insert_multiple(
         cls: type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
     ) -> str:
-        if cls.array_safe_insert:
+        if cls.insert_multiple_mode == "executemany":
+            await cls.insert_multiple_executemany(connection_or_pool, rows)
+            return "INSERT"
+        elif cls.insert_multiple_mode == "array_safe":
             return await cls.insert_multiple_array_safe(connection_or_pool, rows)
         else:
             return await cls.insert_multiple_unnest(connection_or_pool, rows)
 
+    @classmethod
+    async def upsert_multiple_executemany(
+        cls: type[T],
+        connection_or_pool: Union[Connection, Pool],
+        rows: Iterable[T],
+        insert_only: FieldNamesSet = (),
+    ) -> None:
+        args = [r.field_values() for r in rows]
+        query = cls.upsert_sql(
+            cls.insert_multiple_executemany_chunk_sql(1), exclude=insert_only
+        ).query()[0]
+        if args:
+            await connection_or_pool.executemany(query, args)
+
     @classmethod
     async def upsert_multiple_unnest(
         cls: type[T],
@@ -566,7 +611,12 @@ async def upsert_multiple(
         rows: Iterable[T],
         insert_only: FieldNamesSet = (),
     ) -> str:
-        if cls.array_safe_insert:
+        if cls.insert_multiple_mode == "executemany":
+            await cls.upsert_multiple_executemany(
+                connection_or_pool, rows, insert_only=insert_only
+            )
+            return "INSERT"
+        elif cls.insert_multiple_mode == "array_safe":
             return await cls.upsert_multiple_array_safe(
                 connection_or_pool, rows, insert_only=insert_only
             )

tests/test_asyncpg.py

@@ -133,13 +133,14 @@ class Test(ModelBase, table_name="test", primary_key="id"):
     assert db.updated == new.updated
 
 
-async def test_replace_multiple_arrays(conn):
+@pytest.mark.parametrize("insert_multiple_mode", ["array_safe", "executemany"])
+async def test_replace_multiple_arrays(conn, insert_multiple_mode):
     @dataclass(order=True)
     class Test(
         ModelBase,
         table_name="test",
         primary_key="id",
-        insert_multiple_mode="array_safe",
+        insert_multiple_mode=insert_multiple_mode,
     ):
         id: int
         a: Annotated[list[int], ColumnInfo(type="INT[]")]

tests/test_dataclasses.py

@@ -57,6 +57,14 @@ class Test(ModelBase, table_name="table"):
         "hi",
     ]
 
+    assert list(Test.insert_multiple_executemany_chunk_sql(1)) == [
+        'INSERT INTO "table" ("foo", "bar") VALUES ($1, $2)'
+    ]
+
+    assert list(Test.insert_multiple_executemany_chunk_sql(3)) == [
+        'INSERT INTO "table" ("foo", "bar") VALUES ($1, $2), ($3, $4), ($5, $6)'
+    ]
+
     assert sql(
         "INSERT INTO table ({}) VALUES ({})",
         sql(",").join(t.field_names_sql()),