dycw-utilities 0.157.1__py3-none-any.whl → 0.158.0__py3-none-any.whl

{dycw_utilities-0.157.1.dist-info → dycw_utilities-0.158.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.157.1
+Version: 0.158.0
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.157.1.dist-info → dycw_utilities-0.158.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=aCEk399BQ2RGpGQ17jdGxjzeW08BgR-YI0pg_VvNlKM,60
+utilities/__init__.py,sha256=dllUsKScS7RBnWFHME-sLeOoDAWms1u25frX9QMEP-4,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=PUedzQ5deqlSECQ33sam9cRzI9TnygHz3FdOqWJWPTM,15288
 utilities/atomicwrites.py,sha256=tPo6r-Rypd9u99u66B9z86YBPpnLrlHtwox_8Z7T34Y,5790
@@ -14,7 +14,7 @@ utilities/dataclasses.py,sha256=MXrvIPSZHlpV4msRdVVDRZZo7MC3gX5C9jDUSoNOdpE,3247
 utilities/enum.py,sha256=5l6pwZD1cjSlVW4ss-zBPspWvrbrYrdtJWcg6f5_J5w,5781
 utilities/errors.py,sha256=mFlDGSM0LI1jZ1pbqwLAH3ttLZ2JVIxyZLojw8tGVZU,1479
 utilities/eventkit.py,sha256=ddoleSwW9zdc2tjX5Ge0pMKtYwV_JMxhHYOxnWX2AGM,12609
-utilities/fastapi.py,sha256=3wpd63Tw9paSyy7STpAD7GGe8fLkLaRC6TPCwIGm1BU,1361
+utilities/fastapi.py,sha256=TqyKvBjiMS594sXPjrz-KRTLMb3l3D3rZ1zAYV7GfOk,1454
 utilities/fpdf2.py,sha256=HgM8JSvoioDXrjC0UR3HVLjnMnnb_mML7nL2EmkTwGI,1854
 utilities/functions.py,sha256=RNVAoLeT_sl-gXaBv2VI_U_EB-d-nSVosYR4gTeeojE,28261
 utilities/functools.py,sha256=I00ru2gQPakZw2SHVeKIKXfTv741655s6HI0lUoE0D4,1552
@@ -22,7 +22,7 @@ utilities/getpass.py,sha256=DfN5UgMAtFCqS3dSfFHUfqIMZX2shXvwphOz_6J6f6A,103
 utilities/gzip.py,sha256=fkGP3KdsBfXlstodT4wtlp-PwNyUsogpbDCVVVGdsm4,781
 utilities/hashlib.py,sha256=SVTgtguur0P4elppvzOBbLEjVM3Pea0eWB61yg2ilxo,309
 utilities/http.py,sha256=TsavEfHlRtlLaeV21Z6KZh0qbPw-kvD1zsQdZ7Kep5Q,977
-utilities/hypothesis.py,sha256=qqV0O2ynV73tT-bAPXxlfWdR1X0iZgrjn7sK-eJOqBM,43812
+utilities/hypothesis.py,sha256=OG8mxN6Y3fSEjRg4NiIjsO_JUHJBzh4g8fvpmKRoRU8,44370
 utilities/importlib.py,sha256=mV1xT_O_zt_GnZZ36tl3xOmMaN_3jErDWY54fX39F6Y,429
 utilities/inflect.py,sha256=v7YkOWSu8NAmVghPcf4F3YBZQoJCS47_DLf9jbfWIs0,581
 utilities/ipython.py,sha256=V2oMYHvEKvlNBzxDXdLvKi48oUq2SclRg5xasjaXStw,763
@@ -62,10 +62,10 @@ utilities/reprlib.py,sha256=ssYTcBW-TeRh3fhCJv57sopTZHF5FrPyyUg9yp5XBlo,3953
 utilities/scipy.py,sha256=wZJM7fEgBAkLSYYvSmsg5ac-QuwAI0BGqHVetw1_Hb0,947
 utilities/sentinel.py,sha256=A_p5jX2K0Yc5XBfoYHyBLqHsEWzE1ByOdDuzzA2pZnE,1434
 utilities/shelve.py,sha256=4OzjQI6kGuUbJciqf535rwnao-_IBv66gsT6tRGiUt0,759
-utilities/slack_sdk.py,sha256=ppFBvKgfg5IRWiIoKPtpTyzBtBF4XmwEvU3I5wLJikM,2140
+utilities/slack_sdk.py,sha256=76-DYtcGiUhEvl-voMamc5OjfF7Y7nCq54Bys1arqzw,2233
 utilities/socket.py,sha256=K77vfREvzoVTrpYKo6MZakol0EYu2q1sWJnnZqL0So0,118
-utilities/sqlalchemy.py,sha256=wtbIp6XDjKwfrvl-wfoY4FQXo_a9vSoHq5K_dYeBBeY,40541
-utilities/sqlalchemy_polars.py,sha256=5Q9HReETYg0qB6E6WQhFh4QAZlKE-IWlogj2BVif_-w,14246
+utilities/sqlalchemy.py,sha256=4rLjecgDe60MnuWxUqZVIQy3dKepNaemoGyrnuCdNZg,36347
+utilities/sqlalchemy_polars.py,sha256=JCGhB37raSR7fqeWV5dTsciRTMVzIdVT9YSqKT0piT0,13370
 utilities/statsmodels.py,sha256=koyiBHvpMcSiBfh99wFUfSggLNx7cuAw3rwyfAhoKpQ,3410
 utilities/string.py,sha256=shmBK87zZwzGyixuNuXCiUbqzfeZ9xlrFwz6JTaRvDk,582
 utilities/tempfile.py,sha256=HxB2BF28CcecDJLQ3Bx2Ej-Pb6RJc6W9ngSpB9CnP4k,2018
@@ -87,8 +87,8 @@ utilities/zoneinfo.py,sha256=FBMcUQ4662Aq8SsuCL1OAhDQiyANmVjtb-C30DRrWoE,1966
 utilities/pytest_plugins/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 utilities/pytest_plugins/pytest_randomly.py,sha256=B1qYVlExGOxTywq2r1SMi5o7btHLk2PNdY_b1p98dkE,409
 utilities/pytest_plugins/pytest_regressions.py,sha256=9v8kAXDM2ycIXJBimoiF4EgrwbUvxTycFWJiGR_GHhM,1466
-dycw_utilities-0.157.1.dist-info/METADATA,sha256=OvC7ucuMA0zBo4CHiaNRLyMYPjJgGjD7meIO_vUXdG4,1643
-dycw_utilities-0.157.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.157.1.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.157.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.157.1.dist-info/RECORD,,
+dycw_utilities-0.158.0.dist-info/METADATA,sha256=-lYD9v1af9rVk_ibo4Z2cLp-i4JgpoR7qiHy1svZeDU,1643
+dycw_utilities-0.158.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.158.0.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.158.0.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.158.0.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.157.1"
+__version__ = "0.158.0"

utilities/fastapi.py

@@ -13,7 +13,7 @@ from utilities.whenever import get_now_local
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
 
-    from utilities.types import Delta
+    from utilities.types import Delta, MaybeType
 
 
 _TASKS: list[Task[None]] = []
@@ -35,14 +35,19 @@ class _PingerReceiverApp(FastAPI):
 
 @enhanced_async_context_manager
 async def yield_ping_receiver(
-    port: int, /, *, host: str = "localhost", timeout: Delta | None = None
+    port: int,
+    /,
+    *,
+    host: str = "localhost",
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> AsyncIterator[None]:
     """Yield the ping receiver."""
     app = _PingerReceiverApp()  # skipif-ci
     server = Server(Config(app, host=host, port=port))  # skipif-ci
     _TASKS.append(create_task(server.serve()))  # skipif-ci
     try:  # skipif-ci
-        async with timeout_td(timeout):
+        async with timeout_td(timeout, error=error):
             yield
     finally:  # skipif-ci
         await server.shutdown()

utilities/hypothesis.py

@@ -897,6 +897,27 @@ def py_datetimes(
 ##
 
 
+def quadruples[T](
+    strategy: SearchStrategy[T],
+    /,
+    *,
+    unique: MaybeSearchStrategy[bool] = False,
+    sorted: MaybeSearchStrategy[bool] = False,  # noqa: A002
+) -> SearchStrategy[tuple[T, T, T, T]]:
+    """Strategy for generating quadruples of elements."""
+    return lists_fixed_length(strategy, 4, unique=unique, sorted=sorted).map(
+        _quadruples_map
+    )
+
+
+def _quadruples_map[T](elements: list[T], /) -> tuple[T, T, T, T]:
+    first, second, third, fourth = elements
+    return first, second, third, fourth
+
+
+##
+
+
 @composite
 def random_states(
     draw: DrawFn, /, *, seed: MaybeSearchStrategy[int | None] = None
@@ -1555,6 +1576,7 @@ __all__ = [
     "paths",
     "plain_date_times",
     "py_datetimes",
+    "quadruples",
     "random_states",
     "sentinels",
     "sets_fixed_length",

utilities/slack_sdk.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:
     from slack_sdk.webhook import WebhookResponse
     from whenever import TimeDelta
 
-    from utilities.types import Delta
+    from utilities.types import Delta, MaybeType
 
 
 _TIMEOUT: Delta = MINUTE
@@ -39,11 +39,16 @@ def _get_client(url: str, /, *, timeout: Delta = _TIMEOUT) -> WebhookClient:
 
 
 async def send_to_slack_async(
-    url: str, text: str, /, *, timeout: TimeDelta = _TIMEOUT
+    url: str,
+    text: str,
+    /,
+    *,
+    timeout: TimeDelta = _TIMEOUT,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> None:
     """Send a message via Slack."""
     client = _get_async_client(url, timeout=timeout)
-    async with timeout_td(timeout):
+    async with timeout_td(timeout, error=error):
         response = await client.send(text=text)
     if response.status_code != HTTPStatus.OK:  # pragma: no cover
         raise SendToSlackError(text=text, response=response)

utilities/sqlalchemy.py

@@ -12,8 +12,8 @@ from collections.abc import (
 )
 from collections.abc import Set as AbstractSet
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
-from functools import partial, reduce
+from dataclasses import dataclass
+from functools import reduce
 from itertools import chain
 from math import floor
 from operator import ge, le
@@ -37,7 +37,6 @@ from sqlalchemy import (
     Connection,
     Engine,
     Insert,
-    PrimaryKeyConstraint,
     Selectable,
     Table,
     and_,
@@ -49,20 +48,13 @@ from sqlalchemy import (
 from sqlalchemy.dialects.mssql import dialect as mssql_dialect
 from sqlalchemy.dialects.mysql import dialect as mysql_dialect
 from sqlalchemy.dialects.oracle import dialect as oracle_dialect
-from sqlalchemy.dialects.postgresql import Insert as postgresql_Insert
 from sqlalchemy.dialects.postgresql import dialect as postgresql_dialect
 from sqlalchemy.dialects.postgresql import insert as postgresql_insert
 from sqlalchemy.dialects.postgresql.asyncpg import PGDialect_asyncpg
 from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
-from sqlalchemy.dialects.sqlite import Insert as sqlite_Insert
 from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-from sqlalchemy.exc import (
-    ArgumentError,
-    DatabaseError,
-    OperationalError,
-    ProgrammingError,
-)
+from sqlalchemy.exc import ArgumentError, DatabaseError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 from sqlalchemy.orm import (
     DeclarativeBase,
@@ -129,24 +121,22 @@ def check_connect(engine: Engine, /) -> bool:
     try:
         with engine.connect() as conn:
             return bool(conn.execute(_SELECT).scalar_one())
-    except (gaierror, ConnectionRefusedError, OperationalError, ProgrammingError):
+    except (gaierror, ConnectionRefusedError, DatabaseError):  # pragma: no cover
         return False
 
 
 async def check_connect_async(
-    engine: AsyncEngine, /, *, timeout: Delta | None = None
+    engine: AsyncEngine,
+    /,
+    *,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> bool:
     """Check if an engine can connect."""
     try:
-        async with timeout_td(timeout), engine.connect() as conn:
+        async with timeout_td(timeout, error=error), engine.connect() as conn:
             return bool((await conn.execute(_SELECT)).scalar_one())
-    except (
-        gaierror,
-        ConnectionRefusedError,
-        OperationalError,
-        ProgrammingError,
-        TimeoutError,
-    ):
+    except (gaierror, ConnectionRefusedError, DatabaseError, TimeoutError):
         return False
 
 
@@ -158,7 +148,7 @@ async def check_engine(
     /,
     *,
     timeout: Delta | None = None,
-    error: type[Exception] = TimeoutError,
+    error: MaybeType[Exception] = TimeoutError,
     num_tables: int | tuple[int, float] | None = None,
 ) -> None:
     """Check that an engine can connect.
@@ -332,9 +322,8 @@ async def ensure_database_created(super_: URL, database: str, /) -> None:
     async with engine.begin() as conn:
         try:
             _ = await conn.execute(text(f"CREATE DATABASE {database}"))
-        except (OperationalError, ProgrammingError) as error:
-            if not search('database ".*" already exists', ensure_str(one(error.args))):
-                raise
+        except DatabaseError as error:
+            _ensure_tables_maybe_reraise(error, 'database ".*" already exists')
 
 
 async def ensure_database_dropped(super_: URL, database: str, /) -> None:
@@ -352,7 +341,7 @@ async def ensure_tables_created(
     /,
     *tables_or_orms: TableOrORMInstOrClass,
     timeout: Delta | None = None,
-    error: type[Exception] = TimeoutError,
+    error: MaybeType[Exception] = TimeoutError,
 ) -> None:
     """Ensure a table/set of tables is/are created."""
     tables = set(map(get_table, tables_or_orms))
@@ -381,7 +370,7 @@ async def ensure_tables_dropped(
     engine: AsyncEngine,
     *tables_or_orms: TableOrORMInstOrClass,
     timeout: Delta | None = None,
-    error: type[Exception] = TimeoutError,
+    error: MaybeType[Exception] = TimeoutError,
 ) -> None:
     """Ensure a table/set of tables is/are dropped."""
     tables = set(map(get_table, tables_or_orms))
@@ -602,12 +591,16 @@ type _InsertItem = (
     | Sequence[_PairOfTupleOrStrMappingAndTable]
     | Sequence[DeclarativeBase]
 )
+type _NormalizedItem = tuple[Table, StrMapping]
+type _InsertPair = tuple[Table, Sequence[StrMapping]]
+type _InsertTriple = tuple[Table, Insert, Sequence[StrMapping] | None]
 
 
 async def insert_items(
     engine: AsyncEngine,
     *items: _InsertItem,
     snake: bool = False,
+    is_upsert: bool = False,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     assume_tables_exist: bool = False,
     timeout_create: Delta | None = None,
@@ -637,37 +630,181 @@ async def insert_items(
                                                Obj(k1=v21, k2=v22, ...),
                                                ...]
     """
-
-    def build_insert(
-        table: Table, values: Iterable[StrMapping], /
-    ) -> tuple[Insert, Any]:
-        match _get_dialect(engine):
-            case "oracle":  # pragma: no cover
-                return insert(table), values
-            case _:
-                return insert(table).values(list(values)), None
-
-    try:
-        prepared = _prepare_insert_or_upsert_items(
-            partial(_normalize_insert_item, snake=snake),
-            engine,
-            build_insert,
-            *items,
-            chunk_size_frac=chunk_size_frac,
-        )
-    except _PrepareInsertOrUpsertItemsError as error:
-        raise InsertItemsError(item=error.item) from None
+    normalized = chain.from_iterable(
+        _insert_items_yield_normalized(i, snake=snake) for i in items
+    )
+    triples = _insert_items_yield_triples(
+        engine, normalized, is_upsert=is_upsert, chunk_size_frac=chunk_size_frac
+    )
     if not assume_tables_exist:
+        triples = list(triples)
+        tables = {table for table, _, _ in triples}
         await ensure_tables_created(
-            engine, *prepared.tables, timeout=timeout_create, error=error_create
+            engine, *tables, timeout=timeout_create, error=error_create
         )
-    for ins, parameters in prepared.yield_pairs():
+    for _, ins, parameters in triples:
         async with yield_connection(
             engine, timeout=timeout_insert, error=error_insert
         ) as conn:
             _ = await conn.execute(ins, parameters=parameters)
 
 
+def _insert_items_yield_normalized(
+    item: _InsertItem, /, *, snake: bool = False
+) -> Iterator[_NormalizedItem]:
+    if _is_pair_of_str_mapping_and_table(item):
+        mapping, table_or_orm = item
+        adjusted = _map_mapping_to_table(mapping, table_or_orm, snake=snake)
+        yield (get_table(table_or_orm), adjusted)
+        return
+    if _is_pair_of_tuple_and_table(item):
+        tuple_, table_or_orm = item
+        mapping = _tuple_to_mapping(tuple_, table_or_orm)
+        yield from _insert_items_yield_normalized((mapping, table_or_orm), snake=snake)
+        return
+    if _is_pair_of_sequence_of_tuple_or_string_mapping_and_table(item):
+        items, table_or_orm = item
+        pairs = [(i, table_or_orm) for i in items]
+        for p in pairs:
+            yield from _insert_items_yield_normalized(p, snake=snake)
+        return
+    if isinstance(item, DeclarativeBase):
+        mapping = _orm_inst_to_dict(item)
+        yield from _insert_items_yield_normalized((mapping, item), snake=snake)
+        return
+    try:
+        _ = iter(item)
+    except TypeError:
+        raise InsertItemsError(item=item) from None
+    if all(map(_is_pair_of_tuple_or_str_mapping_and_table, item)):
+        pairs = cast("Sequence[_PairOfTupleOrStrMappingAndTable]", item)
+        for p in pairs:
+            yield from _insert_items_yield_normalized(p, snake=snake)
+        return
+    if all(map(is_orm, item)):
+        classes = cast("Sequence[DeclarativeBase]", item)
+        for c in classes:
+            yield from _insert_items_yield_normalized(c, snake=snake)
+        return
+    raise InsertItemsError(item=item)
+
+
+def _insert_items_yield_triples(
+    engine: AsyncEngine,
+    items: Iterable[_NormalizedItem],
+    /,
+    *,
+    is_upsert: bool = False,
+    chunk_size_frac: float = CHUNK_SIZE_FRAC,
+) -> Iterable[_InsertTriple]:
+    pairs = _insert_items_yield_chunked_pairs(
+        engine, items, is_upsert=is_upsert, chunk_size_frac=chunk_size_frac
+    )
+    for table, mappings in pairs:
+        match is_upsert, _get_dialect(engine):
+            case False, "oracle":  # pragma: no cover
+                ins = insert(table)
+                parameters = mappings
+            case False, _:
+                ins = insert(table).values(mappings)
+                parameters = None
+            case True, _:
+                ins = _insert_items_build_insert_with_on_conflict_do_update(
+                    engine, table, mappings
+                )
+                parameters = None
+            case never:
+                assert_never(never)
+        yield table, ins, parameters
+
+
+def _insert_items_yield_chunked_pairs(
+    engine: AsyncEngine,
+    items: Iterable[_NormalizedItem],
+    /,
+    *,
+    is_upsert: bool = False,
+    chunk_size_frac: float = CHUNK_SIZE_FRAC,
+) -> Iterable[_InsertPair]:
+    for table, mappings in _insert_items_yield_raw_pairs(items, is_upsert=is_upsert):
+        chunk_size = get_chunk_size(engine, table, chunk_size_frac=chunk_size_frac)
+        for mappings_i in chunked(mappings, chunk_size):
+            yield table, list(mappings_i)
+
+
+def _insert_items_yield_raw_pairs(
+    items: Iterable[_NormalizedItem], /, *, is_upsert: bool = False
+) -> Iterable[_InsertPair]:
+    by_table: defaultdict[Table, list[StrMapping]] = defaultdict(list)
+    for table, mapping in items:
+        by_table[table].append(mapping)
+    for table, mappings in by_table.items():
+        yield from _insert_items_yield_raw_pairs_one(
+            table, mappings, is_upsert=is_upsert
+        )
+
+
+def _insert_items_yield_raw_pairs_one(
+    table: Table, mappings: Iterable[StrMapping], /, *, is_upsert: bool = False
+) -> Iterable[_InsertPair]:
+    merged = _insert_items_yield_merged_mappings(table, mappings)
+    match is_upsert:
+        case True:
+            by_keys: defaultdict[frozenset[str], list[StrMapping]] = defaultdict(list)
+            for mapping in merged:
+                non_null = {k: v for k, v in mapping.items() if v is not None}
+                by_keys[frozenset(non_null)].append(non_null)
+            for mappings_i in by_keys.values():
+                yield table, mappings_i
+        case False:
+            yield table, list(merged)
+        case never:
+            assert_never(never)
+
+
+def _insert_items_yield_merged_mappings(
+    table: Table, mappings: Iterable[StrMapping], /
+) -> Iterable[StrMapping]:
+    columns = list(yield_primary_key_columns(table))
+    col_names = [c.name for c in columns]
+    cols_auto = {c.name for c in columns if c.autoincrement in {True, "auto"}}
+    cols_non_auto = set(col_names) - cols_auto
+    by_key: defaultdict[tuple[Hashable, ...], list[StrMapping]] = defaultdict(list)
+    for mapping in mappings:
+        check_subset(cols_non_auto, mapping)
+        has_all_auto = set(cols_auto).issubset(mapping)
+        if has_all_auto:
+            pkey = tuple(mapping[k] for k in col_names)
+            rest: StrMapping = {k: v for k, v in mapping.items() if k not in col_names}
+            by_key[pkey].append(rest)
+        else:
+            yield mapping
+    for k, v in by_key.items():
+        head = dict(zip(col_names, k, strict=True))
+        yield merge_str_mappings(head, *v)
+
+
+def _insert_items_build_insert_with_on_conflict_do_update(
+    engine: AsyncEngine, table: Table, mappings: Iterable[StrMapping], /
+) -> Insert:
+    primary_key = cast("Any", table.primary_key)
+    mappings = list(mappings)
+    columns = merge_sets(*mappings)
+    match _get_dialect(engine):
+        case "postgresql":  # skipif-ci-and-not-linux
+            ins = postgresql_insert(table).values(mappings)
+            set_ = {c: getattr(ins.excluded, c) for c in columns}
+            return ins.on_conflict_do_update(constraint=primary_key, set_=set_)
+        case "sqlite":
+            ins = sqlite_insert(table).values(mappings)
+            set_ = {c: getattr(ins.excluded, c) for c in columns}
+            return ins.on_conflict_do_update(index_elements=primary_key, set_=set_)
+        case "mssql" | "mysql" | "oracle" as dialect:  # pragma: no cover
+            raise NotImplementedError(dialect)
+        case never:
+            assert_never(never)
+
+
 @dataclass(kw_only=True, slots=True)
 class InsertItemsError(Exception):
     item: _InsertItem
@@ -738,80 +875,6 @@ async def migrate_data(
 ##
 
 
-def _normalize_insert_item(
-    item: _InsertItem, /, *, snake: bool = False
-) -> list[_NormalizedItem]:
-    """Normalize an insertion item."""
-    if _is_pair_of_str_mapping_and_table(item):
-        mapping, table_or_orm = item
-        adjusted = _map_mapping_to_table(mapping, table_or_orm, snake=snake)
-        normalized = _NormalizedItem(mapping=adjusted, table=get_table(table_or_orm))
-        return [normalized]
-    if _is_pair_of_tuple_and_table(item):
-        tuple_, table_or_orm = item
-        mapping = _tuple_to_mapping(tuple_, table_or_orm)
-        return _normalize_insert_item((mapping, table_or_orm), snake=snake)
-    if _is_pair_of_sequence_of_tuple_or_string_mapping_and_table(item):
-        items, table_or_orm = item
-        pairs = [(i, table_or_orm) for i in items]
-        normalized = (_normalize_insert_item(p, snake=snake) for p in pairs)
-        return list(chain.from_iterable(normalized))
-    if isinstance(item, DeclarativeBase):
-        mapping = _orm_inst_to_dict(item)
-        return _normalize_insert_item((mapping, item), snake=snake)
-    try:
-        _ = iter(item)
-    except TypeError:
-        raise _NormalizeInsertItemError(item=item) from None
-    if all(map(_is_pair_of_tuple_or_str_mapping_and_table, item)):
-        seq = cast("Sequence[_PairOfTupleOrStrMappingAndTable]", item)
-        normalized = (_normalize_insert_item(p, snake=snake) for p in seq)
-        return list(chain.from_iterable(normalized))
-    if all(map(is_orm, item)):
-        seq = cast("Sequence[DeclarativeBase]", item)
-        normalized = (_normalize_insert_item(p, snake=snake) for p in seq)
-        return list(chain.from_iterable(normalized))
-    raise _NormalizeInsertItemError(item=item)
-
-
-@dataclass(kw_only=True, slots=True)
-class _NormalizeInsertItemError(Exception):
-    item: _InsertItem
-
-    @override
-    def __str__(self) -> str:
-        return f"Item must be valid; got {self.item}"
-
-
-@dataclass(kw_only=True, slots=True)
-class _NormalizedItem:
-    mapping: StrMapping
-    table: Table
-
-
-def _normalize_upsert_item(
-    item: _InsertItem,
-    /,
-    *,
-    snake: bool = False,
-    selected_or_all: Literal["selected", "all"] = "selected",
-) -> Iterator[_NormalizedItem]:
-    """Normalize an upsert item."""
-    normalized = _normalize_insert_item(item, snake=snake)
-    match selected_or_all:
-        case "selected":
-            for norm in normalized:
-                values = {k: v for k, v in norm.mapping.items() if v is not None}
-                yield _NormalizedItem(mapping=values, table=norm.table)
-        case "all":
-            yield from normalized
-        case never:
-            assert_never(never)
-
-
-##
-
-
 def selectable_to_string(
     selectable: Selectable[Any], engine_or_conn: EngineOrConnectionOrAsync, /
 ) -> str:
@@ -836,134 +899,6 @@ class TablenameMixin:
 ##
 
 
-type _SelectedOrAll = Literal["selected", "all"]
-
-
-async def upsert_items(
-    engine: AsyncEngine,
-    /,
-    *items: _InsertItem,
-    snake: bool = False,
-    selected_or_all: _SelectedOrAll = "selected",
-    chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    assume_tables_exist: bool = False,
-    timeout_create: Delta | None = None,
-    error_create: type[Exception] = TimeoutError,
-    timeout_insert: Delta | None = None,
-    error_insert: type[Exception] = TimeoutError,
-) -> None:
-    """Upsert a set of items into a database.
-
-    These can be one of the following:
-     - pair of dict & table/class:            {k1=v1, k2=v2, ...), table_cls
-     - pair of list of dicts & table/class:   [{k1=v11, k2=v12, ...},
-                                               {k1=v21, k2=v22, ...},
-                                               ...], table/class
-     - list of pairs of dict & table/class:   [({k1=v11, k2=v12, ...}, table_cls1),
-                                               ({k1=v21, k2=v22, ...}, table_cls2),
-                                               ...]
-     - mapped class:                          Obj(k1=v1, k2=v2, ...)
-     - list of mapped classes:                [Obj(k1=v11, k2=v12, ...),
-                                               Obj(k1=v21, k2=v22, ...),
-                                               ...]
-    """
-
-    def build_insert(
-        table: Table, values: Iterable[StrMapping], /
-    ) -> tuple[Insert, None]:
-        ups = _upsert_items_build(
-            engine, table, values, selected_or_all=selected_or_all
-        )
-        return ups, None
-
-    try:
-        prepared = _prepare_insert_or_upsert_items(
-            partial(
-                _normalize_upsert_item, snake=snake, selected_or_all=selected_or_all
-            ),
-            engine,
-            build_insert,
-            *items,
-            chunk_size_frac=chunk_size_frac,
-        )
-    except _PrepareInsertOrUpsertItemsError as error:
-        raise UpsertItemsError(item=error.item) from None
-    if not assume_tables_exist:
-        await ensure_tables_created(
-            engine, *prepared.tables, timeout=timeout_create, error=error_create
-        )
-    for ups, _ in prepared.yield_pairs():
-        async with yield_connection(
-            engine, timeout=timeout_insert, error=error_insert
-        ) as conn:
-            _ = await conn.execute(ups)
-
-
-def _upsert_items_build(
-    engine: AsyncEngine,
-    table: Table,
-    values: Iterable[StrMapping],
-    /,
-    *,
-    selected_or_all: Literal["selected", "all"] = "selected",
-) -> Insert:
-    values = list(values)
-    keys = merge_sets(*values)
-    dict_nones = dict.fromkeys(keys)
-    values = [{**dict_nones, **v} for v in values]
-    match _get_dialect(engine):
-        case "postgresql":  # skipif-ci-and-not-linux
-            insert = postgresql_insert
-        case "sqlite":
-            insert = sqlite_insert
-        case "mssql" | "mysql" | "oracle" as dialect:  # pragma: no cover
-            raise NotImplementedError(dialect)
-        case never:
-            assert_never(never)
-    ins = insert(table).values(values)
-    primary_key = cast("Any", table.primary_key)
-    return _upsert_items_apply_on_conflict_do_update(
-        values, ins, primary_key, selected_or_all=selected_or_all
-    )
-
-
-def _upsert_items_apply_on_conflict_do_update(
-    values: Iterable[StrMapping],
-    insert: postgresql_Insert | sqlite_Insert,
-    primary_key: PrimaryKeyConstraint,
-    /,
-    *,
-    selected_or_all: Literal["selected", "all"] = "selected",
-) -> Insert:
-    match selected_or_all:
-        case "selected":
-            columns = merge_sets(*values)
-        case "all":
-            columns = {c.name for c in insert.excluded}
-        case never:
-            assert_never(never)
-    set_ = {c: getattr(insert.excluded, c) for c in columns}
-    match insert:
-        case postgresql_Insert():  # skipif-ci
-            return insert.on_conflict_do_update(constraint=primary_key, set_=set_)
-        case sqlite_Insert():
-            return insert.on_conflict_do_update(index_elements=primary_key, set_=set_)
-        case never:
-            assert_never(never)
-
-
-@dataclass(kw_only=True, slots=True)
-class UpsertItemsError(Exception):
-    item: _InsertItem
-
-    @override
-    def __str__(self) -> str:
-        return f"Item must be valid; got {self.item}"
-
-
-##
-
-
 @asynccontextmanager
 async def yield_connection(
     engine: AsyncEngine,
@@ -1207,84 +1142,6 @@ def _orm_inst_to_dict_predicate(
 ##
 
 
-@dataclass(kw_only=True, slots=True)
-class _PrepareInsertOrUpsertItems:
-    mapping: dict[Table, list[StrMapping]] = field(default_factory=dict)
-    yield_pairs: Callable[[], Iterator[tuple[Insert, Any]]]
-
-    @property
-    def tables(self) -> Sequence[Table]:
-        return list(self.mapping)
-
-
-def _prepare_insert_or_upsert_items(
-    normalize_item: Callable[[_InsertItem], Iterable[_NormalizedItem]],
-    engine: AsyncEngine,
-    build_insert: Callable[[Table, Iterable[StrMapping]], tuple[Insert, Any]],
-    /,
-    *items: Any,
-    chunk_size_frac: float = CHUNK_SIZE_FRAC,
-) -> _PrepareInsertOrUpsertItems:
-    """Prepare a set of insert/upsert items."""
-    mapping: defaultdict[Table, list[StrMapping]] = defaultdict(list)
-    lengths: set[int] = set()
-    try:
-        for item in items:
-            for normed in normalize_item(item):
-                mapping[normed.table].append(normed.mapping)
-                lengths.add(len(normed.mapping))
-    except _NormalizeInsertItemError as error:
-        raise _PrepareInsertOrUpsertItemsError(item=error.item) from None
-    merged: dict[Table, list[StrMapping]] = {
-        table: _prepare_insert_or_upsert_items_merge_items(table, values)
-        for table, values in mapping.items()
-    }
-
-    def yield_pairs() -> Iterator[tuple[Insert, None]]:
-        for table, values in merged.items():
-            chunk_size = get_chunk_size(engine, table, chunk_size_frac=chunk_size_frac)
-            for chunk in chunked(values, chunk_size):
-                yield build_insert(table, chunk)
-
-    return _PrepareInsertOrUpsertItems(mapping=mapping, yield_pairs=yield_pairs)
-
-
-@dataclass(kw_only=True, slots=True)
-class _PrepareInsertOrUpsertItemsError(Exception):
-    item: Any
-
-    @override
-    def __str__(self) -> str:
-        return f"Item must be valid; got {self.item}"
-
-
-def _prepare_insert_or_upsert_items_merge_items(
-    table: Table, items: Iterable[StrMapping], /
-) -> list[StrMapping]:
-    columns = list(yield_primary_key_columns(table))
-    col_names = [c.name for c in columns]
-    cols_auto = {c.name for c in columns if c.autoincrement in {True, "auto"}}
-    cols_non_auto = set(col_names) - cols_auto
-    mapping: defaultdict[tuple[Hashable, ...], list[StrMapping]] = defaultdict(list)
-    unchanged: list[StrMapping] = []
-    for item in items:
-        check_subset(cols_non_auto, item)
-        has_all_auto = set(cols_auto).issubset(item)
-        if has_all_auto:
-            pkey = tuple(item[k] for k in col_names)
-            rest: StrMapping = {k: v for k, v in item.items() if k not in col_names}
-            mapping[pkey].append(rest)
-        else:
-            unchanged.append(item)
-    merged = {k: merge_str_mappings(*v) for k, v in mapping.items()}
-    return [
-        dict(zip(col_names, k, strict=True)) | dict(v) for k, v in merged.items()
-    ] + unchanged
-
-
-##
-
-
 def _tuple_to_mapping(
     values: tuple[Any, ...], table_or_orm: TableOrORMInstOrClass, /
 ) -> dict[str, Any]:
@@ -1303,7 +1160,6 @@ __all__ = [
     "GetTableError",
     "InsertItemsError",
     "TablenameMixin",
-    "UpsertItemsError",
     "check_connect",
     "check_connect_async",
     "check_engine",
@@ -1329,7 +1185,6 @@ __all__ = [
     "is_table_or_orm",
     "migrate_data",
     "selectable_to_string",
-    "upsert_items",
     "yield_connection",
     "yield_primary_key_columns",
 ]

utilities/sqlalchemy_polars.py

@@ -4,7 +4,7 @@ import datetime as dt
 import decimal
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal, assert_never, cast, overload, override
+from typing import TYPE_CHECKING, Any, cast, overload, override
 from uuid import UUID
 
 import polars as pl
@@ -44,7 +44,6 @@ from utilities.sqlalchemy import (
     get_chunk_size,
     get_columns,
     insert_items,
-    upsert_items,
 )
 from utilities.text import snake_case
 from utilities.typing import is_subclass_gen
@@ -75,9 +74,9 @@ async def insert_dataframe(
     /,
     *,
     snake: bool = False,
+    is_upsert: bool = False,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     assume_tables_exist: bool = False,
-    upsert: Literal["selected", "all"] | None = None,
     timeout_create: TimeDelta | None = None,
     error_create: type[Exception] = TimeoutError,
     timeout_insert: TimeDelta | None = None,
@@ -87,43 +86,25 @@ async def insert_dataframe(
     mapping = _insert_dataframe_map_df_schema_to_table(
         df.schema, table_or_orm, snake=snake
     )
-    items = df.select(*mapping).rename(mapping).to_dicts()
+    items = df.select(*mapping).rename(mapping).rows(named=True)
     if len(items) == 0:
-        if not df.is_empty():
-            raise InsertDataFrameError(df=df)
         if not assume_tables_exist:
             await ensure_tables_created(
                 engine, table_or_orm, timeout=timeout_create, error=error_create
             )
         return
-    match upsert:
-        case None:
-            await insert_items(
-                engine,
-                (items, table_or_orm),
-                snake=snake,
-                chunk_size_frac=chunk_size_frac,
-                assume_tables_exist=assume_tables_exist,
-                timeout_create=timeout_create,
-                error_create=error_create,
-                timeout_insert=timeout_insert,
-                error_insert=error_insert,
-            )
-        case "selected" | "all" as selected_or_all:  # skipif-ci-and-not-linux
-            await upsert_items(
-                engine,
-                (items, table_or_orm),
-                snake=snake,
-                chunk_size_frac=chunk_size_frac,
-                selected_or_all=selected_or_all,
-                assume_tables_exist=assume_tables_exist,
-                timeout_create=timeout_create,
-                error_create=error_create,
-                timeout_insert=timeout_insert,
-                error_insert=error_insert,
-            )
-        case never:
-            assert_never(never)
+    await insert_items(
+        engine,
+        (items, table_or_orm),
+        snake=snake,
+        is_upsert=is_upsert,
+        chunk_size_frac=chunk_size_frac,
+        assume_tables_exist=assume_tables_exist,
+        timeout_create=timeout_create,
+        error_create=error_create,
+        timeout_insert=timeout_insert,
+        error_insert=error_insert,
+    )
 
 
 def _insert_dataframe_map_df_schema_to_table(
@@ -207,15 +188,6 @@ def _insert_dataframe_check_df_and_db_types(
     )
 
 
-@dataclass(kw_only=True, slots=True)
-class InsertDataFrameError(Exception):
-    df: DataFrame
-
-    @override
-    def __str__(self) -> str:
-        return f"Non-empty DataFrame must resolve to at least 1 item\n\n{self.df}"
-
-
 @overload
 async def select_to_dataframe(
     sel: Select[Any],
@@ -229,6 +201,7 @@ async def select_to_dataframe(
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame: ...
 @overload
@@ -244,6 +217,7 @@ async def select_to_dataframe(
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> Iterable[DataFrame]: ...
 @overload
@@ -259,6 +233,7 @@ async def select_to_dataframe(
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> AsyncIterable[DataFrame]: ...
 @overload
@@ -274,6 +249,7 @@ async def select_to_dataframe(
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]: ...
 async def select_to_dataframe(
@@ -439,4 +415,4 @@ def _select_to_dataframe_yield_selects_with_in_clauses(
     return (sel.where(in_col.in_(values)) for values in chunked(in_values, chunk_size))
 
 
-__all__ = ["InsertDataFrameError", "insert_dataframe", "select_to_dataframe"]
+__all__ = ["insert_dataframe", "select_to_dataframe"]