dycw-utilities 0.129.10__py3-none-any.whl → 0.175.17__py3-none-any.whl

utilities/sqlalchemy.py

@@ -12,21 +12,31 @@ 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
 from re import search
-from typing import Any, Literal, TypeGuard, TypeVar, assert_never, cast, override
+from socket import gaierror
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    TypeGuard,
+    assert_never,
+    cast,
+    overload,
+    override,
+)
 
+import sqlalchemy
 from sqlalchemy import (
     URL,
     Column,
     Connection,
     Engine,
     Insert,
-    PrimaryKeyConstraint,
     Selectable,
     Table,
     and_,
@@ -38,16 +48,14 @@ 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.sqlite import Insert as sqlite_Insert
+from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
 from sqlalchemy.exc import ArgumentError, DatabaseError
-from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
-from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 from sqlalchemy.orm import (
     DeclarativeBase,
     InstrumentedAttribute,
@@ -57,17 +65,8 @@ from sqlalchemy.orm import (
 from sqlalchemy.orm.exc import UnmappedClassError
 from sqlalchemy.pool import NullPool, Pool
 
-from utilities.asyncio import Looper, timeout_dur
-from utilities.contextlib import suppress_super_object_attribute_error
-from utilities.datetime import SECOND
-from utilities.functions import (
-    ensure_str,
-    get_class_name,
-    is_sequence_of_tuple_or_str_mapping,
-    is_string_mapping,
-    is_tuple,
-    is_tuple_or_str_mapping,
-)
+from utilities.asyncio import timeout_td
+from utilities.functions import ensure_str, get_class_name, yield_object_attributes
 from utilities.iterables import (
     CheckLengthError,
     CheckSubSetError,
@@ -80,16 +79,63 @@ from utilities.iterables import (
     merge_str_mappings,
     one,
 )
+from utilities.os import is_pytest
 from utilities.reprlib import get_repr
-from utilities.text import snake_case
-from utilities.types import Duration, MaybeIterable, StrMapping, TupleOrStrMapping
+from utilities.text import secret_str, snake_case
+from utilities.types import (
+    Delta,
+    MaybeIterable,
+    MaybeType,
+    StrMapping,
+    TupleOrStrMapping,
+)
+from utilities.typing import (
+    is_sequence_of_tuple_or_str_mapping,
+    is_string_mapping,
+    is_tuple,
+    is_tuple_or_str_mapping,
+)
+
+if TYPE_CHECKING:
+    from enum import Enum, StrEnum
+
 
-_T = TypeVar("_T")
-type _EngineOrConnectionOrAsync = Engine | Connection | AsyncEngine | AsyncConnection
+type EngineOrConnectionOrAsync = Engine | Connection | AsyncEngine | AsyncConnection
 type Dialect = Literal["mssql", "mysql", "oracle", "postgresql", "sqlite"]
+type DialectOrEngineOrConnectionOrAsync = Dialect | EngineOrConnectionOrAsync
 type ORMInstOrClass = DeclarativeBase | type[DeclarativeBase]
 type TableOrORMInstOrClass = Table | ORMInstOrClass
-CHUNK_SIZE_FRAC = 0.95
+CHUNK_SIZE_FRAC = 0.8
+
+
+##
+
+
+_SELECT = text("SELECT 1")
+
+
+def check_connect(engine: Engine, /) -> bool:
+    """Check if an engine can connect."""
+    try:
+        with engine.connect() as conn:
+            return bool(conn.execute(_SELECT).scalar_one())
+    except (gaierror, ConnectionRefusedError, DatabaseError):  # pragma: no cover
+        return False
+
+
+async def check_connect_async(
+    engine: AsyncEngine,
+    /,
+    *,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
+) -> bool:
+    """Check if an engine can connect."""
+    try:
+        async with timeout_td(timeout, error=error), engine.connect() as conn:
+            return bool((await conn.execute(_SELECT)).scalar_one())
+    except (gaierror, ConnectionRefusedError, DatabaseError, TimeoutError):
+        return False
 
 
 ##
@@ -99,8 +145,8 @@ async def check_engine(
     engine: AsyncEngine,
     /,
     *,
-    timeout: Duration | None = None,
-    error: type[Exception] = TimeoutError,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     num_tables: int | tuple[int, float] | None = None,
 ) -> None:
     """Check that an engine can connect.
@@ -115,7 +161,7 @@ async def check_engine(
             query = "select * from all_objects"
         case "sqlite":
             query = "select * from sqlite_master where type='table'"
-        case _ as never:
+        case never:
             assert_never(never)
     statement = text(query)
     async with yield_connection(engine, timeout=timeout, error=error) as conn:
@@ -183,7 +229,36 @@ def _columnwise_minmax(*columns: Any, op: Callable[[Any, Any], Any]) -> Any:
 ##
 
 
-def create_async_engine(
+@overload
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: Literal[True],
+) -> AsyncEngine: ...
+@overload
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: Literal[False] = False,
+) -> Engine: ...
+@overload
+def create_engine(
     drivername: str,
     /,
     *,
@@ -194,7 +269,21 @@ def create_async_engine(
     database: str | None = None,
     query: StrMapping | None = None,
     poolclass: type[Pool] | None = NullPool,
-) -> AsyncEngine:
+    async_: bool = False,
+) -> Engine | AsyncEngine: ...
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: bool = False,
+) -> Engine | AsyncEngine:
     """Create a SQLAlchemy engine."""
     if query is None:
         kwargs = {}
@@ -213,7 +302,47 @@ def create_async_engine(
         database=database,
         **kwargs,
     )
-    return _create_async_engine(url, poolclass=poolclass)
+    match async_:
+        case False:
+            return sqlalchemy.create_engine(url, poolclass=poolclass)
+        case True:
+            return create_async_engine(url, poolclass=poolclass)
+        case never:
+            assert_never(never)
+
+
+##
+
+
+async def ensure_database_created(super_: URL, database: str, /) -> None:
+    """Ensure a database is created."""
+    engine = create_async_engine(super_, isolation_level="AUTOCOMMIT")
+    async with engine.begin() as conn:
+        try:
+            _ = await conn.execute(text(f"CREATE DATABASE {database}"))
+        except DatabaseError as error:
+            _ensure_tables_maybe_reraise(error, 'database ".*" already exists')
+
+
+async def ensure_database_dropped(super_: URL, database: str, /) -> None:
+    """Ensure a database is dropped."""
+    engine = create_async_engine(super_, isolation_level="AUTOCOMMIT")
+    async with engine.begin() as conn:
+        _ = await conn.execute(text(f"DROP DATABASE IF EXISTS {database}"))
+
+
+async def ensure_database_users_disconnected(super_: URL, database: str, /) -> None:
+    """Ensure a databases' users are disconnected."""
+    engine = create_async_engine(super_, isolation_level="AUTOCOMMIT")
+    match dialect := _get_dialect(engine):
+        case "postgresql":  # skipif-ci-and-not-linux
+            query = f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = {database!r} AND pid <> pg_backend_pid()"  # noqa: S608
+        case "mssql" | "mysql" | "oracle" | "sqlite":  # pragma: no cover
+            raise NotImplementedError(dialect)
+        case never:
+            assert_never(never)
+    async with engine.begin() as conn:
+        _ = await conn.execute(text(query))
 
 
 ##
@@ -223,8 +352,8 @@ async def ensure_tables_created(
     engine: AsyncEngine,
     /,
     *tables_or_orms: TableOrORMInstOrClass,
-    timeout: Duration | None = None,
-    error: type[Exception] = TimeoutError,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> None:
     """Ensure a table/set of tables is/are created."""
     tables = set(map(get_table, tables_or_orms))
@@ -239,7 +368,7 @@ async def ensure_tables_created(
             match = "ORA-00955: name is already used by an existing object"
         case "sqlite":
             match = "table .* already exists"
-        case _ as never:
+        case never:
             assert_never(never)
     async with yield_connection(engine, timeout=timeout, error=error) as conn:
         for table in tables:
@@ -252,8 +381,8 @@ async def ensure_tables_created(
 async def ensure_tables_dropped(
     engine: AsyncEngine,
     *tables_or_orms: TableOrORMInstOrClass,
-    timeout: Duration | None = None,
-    error: type[Exception] = TimeoutError,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> None:
     """Ensure a table/set of tables is/are dropped."""
     tables = set(map(get_table, tables_or_orms))
@@ -268,7 +397,7 @@ async def ensure_tables_dropped(
             match = "ORA-00942: table or view does not exist"
         case "sqlite":
             match = "no such table"
-        case _ as never:
+        case never:
             assert_never(never)
     async with yield_connection(engine, timeout=timeout, error=error) as conn:
         for table in tables:
@@ -281,16 +410,120 @@ async def ensure_tables_dropped(
 ##
 
 
+def enum_name(enum: type[Enum], /) -> str:
+    """Get the name of an Enum."""
+    return f"{snake_case(get_class_name(enum))}_enum"
+
+
+##
+
+
+def enum_values(enum: type[StrEnum], /) -> list[str]:
+    """Get the values of a StrEnum."""
+    return [e.value for e in enum]
+
+
+##
+
+
+@dataclass(kw_only=True, slots=True)
+class ExtractURLOutput:
+    username: str
+    password: secret_str
+    host: str
+    port: int
+    database: str
+
+
+def extract_url(url: URL, /) -> ExtractURLOutput:
+    """Extract the database, host & port from a URL."""
+    if url.username is None:
+        raise _ExtractURLUsernameError(url=url)
+    if url.password is None:
+        raise _ExtractURLPasswordError(url=url)
+    if url.host is None:
+        raise _ExtractURLHostError(url=url)
+    if url.port is None:
+        raise _ExtractURLPortError(url=url)
+    if url.database is None:
+        raise _ExtractURLDatabaseError(url=url)
+    return ExtractURLOutput(
+        username=url.username,
+        password=secret_str(url.password),
+        host=url.host,
+        port=url.port,
+        database=url.database,
+    )
+
+
+@dataclass(kw_only=True, slots=True)
+class ExtractURLError(Exception):
+    url: URL
+
+
+@dataclass(kw_only=True, slots=True)
+class _ExtractURLUsernameError(ExtractURLError):
+    @override
+    def __str__(self) -> str:
+        return f"Expected URL to contain a user name; got {self.url}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _ExtractURLPasswordError(ExtractURLError):
+    @override
+    def __str__(self) -> str:
+        return f"Expected URL to contain a password; got {self.url}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _ExtractURLHostError(ExtractURLError):
+    @override
+    def __str__(self) -> str:
+        return f"Expected URL to contain a host; got {self.url}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _ExtractURLPortError(ExtractURLError):
+    @override
+    def __str__(self) -> str:
+        return f"Expected URL to contain a port; got {self.url}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _ExtractURLDatabaseError(ExtractURLError):
+    @override
+    def __str__(self) -> str:
+        return f"Expected URL to contain a database; got {self.url}"
+
+
+##
+
+
 def get_chunk_size(
-    engine_or_conn: _EngineOrConnectionOrAsync,
+    dialect_or_engine_or_conn: DialectOrEngineOrConnectionOrAsync,
+    table_or_orm_or_num_cols: TableOrORMInstOrClass | Sized | int,
     /,
     *,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    scaling: float = 1.0,
 ) -> int:
     """Get the maximum chunk size for an engine."""
-    max_params = _get_dialect_max_params(engine_or_conn)
-    return max(floor(chunk_size_frac * max_params / scaling), 1)
+    max_params = _get_dialect_max_params(dialect_or_engine_or_conn)
+    match table_or_orm_or_num_cols:
+        case Table() | DeclarativeBase() | type() as table_or_orm:
+            return get_chunk_size(
+                dialect_or_engine_or_conn,
+                get_columns(table_or_orm),
+                chunk_size_frac=chunk_size_frac,
+            )
+        case Sized() as sized:
+            return get_chunk_size(
+                dialect_or_engine_or_conn, len(sized), chunk_size_frac=chunk_size_frac
+            )
+        case int() as num_cols:
+            size = floor(chunk_size_frac * max_params / num_cols)
+            return max(size, 1)
+        case never:
+            assert_never(never)
 
 
 ##
@@ -370,18 +603,22 @@ 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: Duration | None = None,
-    error_create: type[Exception] = TimeoutError,
-    timeout_insert: Duration | None = None,
-    error_insert: type[Exception] = TimeoutError,
+    timeout_create: Delta | None = None,
+    error_create: MaybeType[BaseException] = TimeoutError,
+    timeout_insert: Delta | None = None,
+    error_insert: MaybeType[BaseException] = TimeoutError,
 ) -> None:
     """Insert a set of items into a database.
 
@@ -405,37 +642,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
@@ -479,10 +860,10 @@ async def migrate_data(
     table_or_orm_to: TableOrORMInstOrClass | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     assume_tables_exist: bool = False,
-    timeout_create: Duration | None = None,
-    error_create: type[Exception] = TimeoutError,
-    timeout_insert: Duration | None = None,
-    error_insert: type[Exception] = TimeoutError,
+    timeout_create: Delta | None = None,
+    error_create: MaybeType[BaseException] = TimeoutError,
+    timeout_insert: Delta | None = None,
+    error_insert: MaybeType[BaseException] = TimeoutError,
 ) -> None:
     """Migrate the contents of a table from one database to another."""
     table_from = get_table(table_or_orm_from)
@@ -506,82 +887,8 @@ 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 _ as never:
-            assert_never(never)
-
-
-##
-
-
 def selectable_to_string(
-    selectable: Selectable[Any], engine_or_conn: _EngineOrConnectionOrAsync, /
+    selectable: Selectable[Any], engine_or_conn: EngineOrConnectionOrAsync, /
 ) -> str:
     """Convert a selectable into a string."""
     com = selectable.compile(
@@ -604,237 +911,22 @@ class TablenameMixin:
 ##
 
 
-@dataclass(kw_only=True)
-class UpsertService(Looper[_InsertItem]):
-    """Service to upsert items to a database."""
-
-    # base
-    freq: Duration = field(default=SECOND, repr=False)
-    backoff: Duration = field(default=SECOND, repr=False)
-    empty_upon_exit: bool = field(default=True, repr=False)
-    # self
-    engine: AsyncEngine
-    snake: bool = False
-    selected_or_all: _SelectedOrAll = "selected"
-    chunk_size_frac: float = CHUNK_SIZE_FRAC
-    assume_tables_exist: bool = False
-    timeout_create: Duration | None = None
-    error_create: type[Exception] = TimeoutError
-    timeout_insert: Duration | None = None
-    error_insert: type[Exception] = TimeoutError
-
-    @override
-    async def core(self) -> None:
-        await super().core()
-        await upsert_items(
-            self.engine,
-            *self.get_all_nowait(),
-            snake=self.snake,
-            selected_or_all=self.selected_or_all,
-            chunk_size_frac=self.chunk_size_frac,
-            assume_tables_exist=self.assume_tables_exist,
-            timeout_create=self.timeout_create,
-            error_create=self.error_create,
-            timeout_insert=self.timeout_insert,
-            error_insert=self.error_insert,
-        )
-
-
-@dataclass(kw_only=True)
-class UpsertServiceMixin:
-    """Mix-in for the upsert service."""
-
-    # base - looper
-    upsert_service_freq: Duration = field(default=SECOND, repr=False)
-    upsert_service_backoff: Duration = field(default=SECOND, repr=False)
-    upsert_service_empty_upon_exit: bool = field(default=False, repr=False)
-    upsert_service_logger: str | None = field(default=None, repr=False)
-    upsert_service_timeout: Duration | None = field(default=None, repr=False)
-    upsert_service_debug: bool = field(default=False, repr=False)
-    # base - upsert service
-    upsert_service_database: AsyncEngine
-    upsert_service_snake: bool = False
-    upsert_service_selected_or_all: _SelectedOrAll = "selected"
-    upsert_service_chunk_size_frac: float = CHUNK_SIZE_FRAC
-    upsert_service_assume_tables_exist: bool = False
-    upsert_service_timeout_create: Duration | None = None
-    upsert_service_error_create: type[Exception] = TimeoutError
-    upsert_service_timeout_insert: Duration | None = None
-    upsert_service_error_insert: type[Exception] = TimeoutError
-    # self
-    _upsert_service: UpsertService = field(init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        with suppress_super_object_attribute_error():
-            super().__post_init__()  # pyright: ignore[reportAttributeAccessIssue]
-        self._upsert_service = UpsertService(
-            # looper
-            freq=self.upsert_service_freq,
-            backoff=self.upsert_service_backoff,
-            empty_upon_exit=self.upsert_service_empty_upon_exit,
-            logger=self.upsert_service_logger,
-            timeout=self.upsert_service_timeout,
-            _debug=self.upsert_service_debug,
-            # upsert service
-            engine=self.upsert_service_database,
-            snake=self.upsert_service_snake,
-            selected_or_all=self.upsert_service_selected_or_all,
-            chunk_size_frac=self.upsert_service_chunk_size_frac,
-            assume_tables_exist=self.upsert_service_assume_tables_exist,
-            timeout_create=self.upsert_service_timeout_create,
-            error_create=self.upsert_service_error_create,
-            timeout_insert=self.upsert_service_timeout_insert,
-            error_insert=self.upsert_service_error_insert,
-        )
-
-    def _yield_sub_loopers(self) -> Iterator[Looper[Any]]:
-        with suppress_super_object_attribute_error():
-            yield from super()._yield_sub_loopers()  # pyright: ignore[reportAttributeAccessIssue]
-        yield self._upsert_service
-
-
-##
-
-
-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: Duration | None = None,
-    error_create: type[Exception] = TimeoutError,
-    timeout_insert: Duration | 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 _ as 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 _ as 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 _ as 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,
     /,
     *,
-    timeout: Duration | None = None,
-    error: type[Exception] = TimeoutError,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
 ) -> AsyncIterator[AsyncConnection]:
     """Yield an async connection."""
-    async with timeout_dur(duration=timeout, error=error), engine.begin() as conn:
-        yield conn
+    try:
+        async with timeout_td(timeout, error=error), engine.begin() as conn:
+            yield conn
+    except GeneratorExit:  # pragma: no cover
+        if not is_pytest():
+            raise
+        return
 
 
 ##
@@ -865,7 +957,7 @@ def _ensure_tables_maybe_reraise(error: DatabaseError, match: str, /) -> None:
 ##
 
 
-def _get_dialect(engine_or_conn: _EngineOrConnectionOrAsync, /) -> Dialect:
+def _get_dialect(engine_or_conn: EngineOrConnectionOrAsync, /) -> Dialect:
     """Get the dialect of a database."""
     dialect = engine_or_conn.dialect
     if isinstance(dialect, mssql_dialect):  # pragma: no cover
@@ -875,7 +967,7 @@ def _get_dialect(engine_or_conn: _EngineOrConnectionOrAsync, /) -> Dialect:
     if isinstance(dialect, oracle_dialect):  # pragma: no cover
         return "oracle"
     if isinstance(  # skipif-ci-and-not-linux
-        dialect, postgresql_dialect | PGDialect_asyncpg
+        dialect, (postgresql_dialect, PGDialect_asyncpg, PGDialect_psycopg)
     ):
         return "postgresql"
     if isinstance(dialect, sqlite_dialect):
@@ -888,7 +980,7 @@ def _get_dialect(engine_or_conn: _EngineOrConnectionOrAsync, /) -> Dialect:
 
 
 def _get_dialect_max_params(
-    dialect_or_engine_or_conn: Dialect | _EngineOrConnectionOrAsync, /
+    dialect_or_engine_or_conn: DialectOrEngineOrConnectionOrAsync, /
 ) -> int:
     """Get the max number of parameters of a dialect."""
     match dialect_or_engine_or_conn:
@@ -910,7 +1002,7 @@ def _get_dialect_max_params(
         ):
             dialect = _get_dialect(engine_or_conn)
             return _get_dialect_max_params(dialect)
-        case _ as never:
+        case never:
             assert_never(never)
 
 
@@ -943,9 +1035,9 @@ def _is_pair_of_tuple_or_str_mapping_and_table(
     return _is_pair_with_predicate_and_table(obj, is_tuple_or_str_mapping)
 
 
-def _is_pair_with_predicate_and_table(
-    obj: Any, predicate: Callable[[Any], TypeGuard[_T]], /
-) -> TypeGuard[tuple[_T, TableOrORMInstOrClass]]:
+def _is_pair_with_predicate_and_table[T](
+    obj: Any, predicate: Callable[[Any], TypeGuard[T]], /
+) -> TypeGuard[tuple[T, TableOrORMInstOrClass]]:
     """Check if an object is pair and a table."""
     return (
         isinstance(obj, tuple)
@@ -995,7 +1087,7 @@ def _map_mapping_to_table(
 @dataclass(kw_only=True, slots=True)
 class _MapMappingToTableError(Exception):
     mapping: StrMapping
-    columns: Sequence[str]
+    columns: list[str]
 
 
 @dataclass(kw_only=True, slots=True)
@@ -1032,102 +1124,31 @@ class _MapMappingToTableSnakeMapNonUniqueError(_MapMappingToTableError):
 
 def _orm_inst_to_dict(obj: DeclarativeBase, /) -> StrMapping:
     """Map an ORM instance to a dictionary."""
-    cls = type(obj)
-
-    def is_attr(attr: str, key: str, /) -> str | None:
-        if isinstance(value := getattr(cls, attr), InstrumentedAttribute) and (
-            value.name == key
-        ):
-            return attr
-        return None
-
-    def yield_items() -> Iterator[tuple[str, Any]]:
-        for key in get_column_names(cls):
-            attr = one(attr for attr in dir(cls) if is_attr(attr, key) is not None)
-            yield key, getattr(obj, attr)
-
-    return dict(yield_items())
-
-
-##
-
-
-@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 = {
-        table: _prepare_insert_or_upsert_items_merge_items(table, values)
-        for table, values in mapping.items()
+    attrs = {
+        k for k, _ in yield_object_attributes(obj, static_type=InstrumentedAttribute)
+    }
+    return {
+        name: _orm_inst_to_dict_one(obj, attrs, name) for name in get_column_names(obj)
     }
-    max_length = max(lengths, default=1)
-    chunk_size = get_chunk_size(
-        engine, chunk_size_frac=chunk_size_frac, scaling=max_length
-    )
-
-    def yield_pairs() -> Iterator[tuple[Insert, None]]:
-        for table, values in merged.items():
-            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 _orm_inst_to_dict_one(
+    obj: DeclarativeBase, attrs: AbstractSet[str], name: str, /
+) -> Any:
+    attr = one(
+        attr for attr in attrs if _orm_inst_to_dict_predicate(type(obj), attr, name)
+    )
+    return getattr(obj, attr)
 
 
-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 _orm_inst_to_dict_predicate(
+    cls: type[DeclarativeBase], attr: str, name: str, /
+) -> bool:
+    cls_attr = getattr(cls, attr)
+    try:
+        return cls_attr.name == name
+    except AttributeError:
+        return False
 
 
 ##
@@ -1144,18 +1165,27 @@ def _tuple_to_mapping(
 __all__ = [
     "CHUNK_SIZE_FRAC",
     "CheckEngineError",
+    "DialectOrEngineOrConnectionOrAsync",
+    "EngineOrConnectionOrAsync",
+    "ExtractURLError",
+    "ExtractURLOutput",
     "GetTableError",
     "InsertItemsError",
     "TablenameMixin",
-    "UpsertItemsError",
-    "UpsertService",
-    "UpsertServiceMixin",
+    "check_connect",
+    "check_connect_async",
     "check_engine",
     "columnwise_max",
     "columnwise_min",
-    "create_async_engine",
+    "create_engine",
+    "ensure_database_created",
+    "ensure_database_dropped",
+    "ensure_database_users_disconnected",
     "ensure_tables_created",
     "ensure_tables_dropped",
+    "enum_name",
+    "enum_values",
+    "extract_url",
     "get_chunk_size",
     "get_column_names",
     "get_columns",
@@ -1168,7 +1198,6 @@ __all__ = [
     "is_table_or_orm",
     "migrate_data",
     "selectable_to_string",
-    "upsert_items",
     "yield_connection",
     "yield_primary_key_columns",
 ]