Merge pull request #12 from bdowning/safe-array-insert

Add insert_multiple_mode and array safety

Author: bdowning

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.4.0-alpha-1
+current_version = 0.4.0-alpha-2
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(-(?P<release>.*)-(?P<build>\d+))?
 serialize = 
 	{major}.{minor}.{patch}-{release}-{build}

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sql-athame"
-version = "0.4.0-alpha-1"
+version = "0.4.0-alpha-2"
 description = "Python tool for slicing and dicing SQL"
 authors = ["Brian Downing <bdowning@lavos.net>"]
 license = "MIT"

run

@@ -4,7 +4,7 @@ usage=()
 
 usage+=("  $0 tests - run tests")
 tests() {
-    poetry run pytest
+    poetry run pytest "$@"
     lint
 }
 

setup.cfg

@@ -35,6 +35,8 @@ ignore_missing_imports = True
 ignore =
     # max line length
     E501,
+    # E203 whitespace before ':'; fights with black
+    E203,
     # multiple statements on one line (def); fights with black
     E704,
     # line break before binary operator; fights with black

sql_athame/dataclasses.py

@@ -100,12 +100,24 @@ class ModelBase(Mapping[str, Any]):
     _cache: Dict[tuple, Any]
     table_name: str
     primary_key_names: Tuple[str, ...]
+    array_safe_insert: bool
 
     def __init_subclass__(
-        cls, *, table_name: str, primary_key: Union[FieldNames, str] = (), **kwargs: Any
+        cls,
+        *,
+        table_name: str,
+        primary_key: Union[FieldNames, str] = (),
+        insert_multiple_mode: str = "unnest",
+        **kwargs: Any,
     ):
         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:
+            raise ValueError("Unknown `insert_multiple_mode`")
         if isinstance(primary_key, str):
             cls.primary_key_names = (primary_key,)
         else:
@@ -407,19 +419,69 @@ def insert_multiple_sql(cls: Type[T], rows: Iterable[T]) -> Fragment:
         )
 
     @classmethod
-    async def insert_multiple(
+    def insert_multiple_array_safe_sql(cls: Type[T], rows: Iterable[T]) -> Fragment:
+        return sql(
+            "INSERT INTO {table} ({fields}) VALUES {values}",
+            table=cls.table_name_sql(),
+            fields=sql.list(cls.field_names_sql()),
+            values=sql.list(
+                sql("({})", sql.list(row.field_values_sql(default_none=True)))
+                for row in rows
+            ),
+        )
+
+    @classmethod
+    async def insert_multiple_unnest(
         cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
     ) -> str:
         return await connection_or_pool.execute(*cls.insert_multiple_sql(rows))
 
     @classmethod
-    async def upsert_multiple(
+    async def insert_multiple_array_safe(
+        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+    ) -> str:
+        for chunk in chunked(rows, 100):
+            last = await connection_or_pool.execute(
+                *cls.insert_multiple_array_safe_sql(chunk)
+            )
+        return last
+
+    @classmethod
+    async def insert_multiple(
+        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+    ) -> str:
+        if cls.array_safe_insert:
+            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_unnest(
         cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
     ) -> str:
         return await connection_or_pool.execute(
             *cls.upsert_sql(cls.insert_multiple_sql(rows))
         )
 
+    @classmethod
+    async def upsert_multiple_array_safe(
+        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+    ) -> str:
+        for chunk in chunked(rows, 100):
+            last = await connection_or_pool.execute(
+                *cls.upsert_sql(cls.insert_multiple_array_safe_sql(chunk))
+            )
+        return last
+
+    @classmethod
+    async def upsert_multiple(
+        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+    ) -> str:
+        if cls.array_safe_insert:
+            return await cls.upsert_multiple_array_safe(connection_or_pool, rows)
+        else:
+            return await cls.upsert_multiple_unnest(connection_or_pool, rows)
+
     @classmethod
     def _get_equal_ignoring_fn(
         cls: Type[T], ignore: FieldNamesSet = ()
@@ -530,3 +592,8 @@ async def replace_multiple_reporting_differences(
             await cls.delete_multiple(connection, deleted)
 
         return created, updated_triples, deleted
+
+
+def chunked(lst, n):
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]

tests/test_asyncpg.py

@@ -91,6 +91,57 @@ class Test(ModelBase, table_name="test", primary_key="id"):
     ]
 
 
+async def test_replace_multiple_arrays(conn):
+    @dataclass(order=True)
+    class Test(
+        ModelBase,
+        table_name="test",
+        primary_key="id",
+        insert_multiple_mode="array_safe",
+    ):
+        id: int
+        a: list[int] = field(
+            metadata=model_field_metadata(type="INT[]", constraints="NOT NULL")
+        )
+        b: str
+
+    await conn.execute(*Test.create_table_sql())
+
+    data = [
+        Test(1, [1], "foo"),
+        Test(2, [1, 3, 5], "bar"),
+        Test(3, [], "quux"),
+    ]
+    await Test.insert_multiple(conn, data)
+
+    c, u, d = await Test.replace_multiple(conn, [], where=[])
+    assert not c and not u
+    assert len(d) == 3
+    assert await Test.select(conn) == []
+
+    await Test.insert_multiple(conn, data)
+
+    c, u, d = await Test.replace_multiple(conn, [], where=sql("a @> ARRAY[1]"))
+    assert not c and not u
+    assert len(d) == 2
+    assert [x.id for x in await Test.select(conn)] == [3]
+
+    await conn.execute("DELETE FROM test")
+    await Test.insert_multiple(conn, data)
+
+    c, u, d = await Test.replace_multiple(
+        conn, [Test(1, [5], "apples"), Test(4, [6], "fred")], where=sql("a @> ARRAY[1]")
+    )
+    assert len(c) == 1
+    assert len(u) == 1
+    assert len(d) == 1
+    assert list(sorted(await Test.select(conn))) == [
+        Test(1, [5], "apples"),
+        Test(3, [], "quux"),
+        Test(4, [6], "fred"),
+    ]
+
+
 async def test_replace_multiple_reporting_differences(conn):
     @dataclass(order=True)
     class Test(ModelBase, table_name="test", primary_key="id"):