Batch Upsert Rows into SQLite Transactionally

from __future__ import annotations

import sqlite3

SCHEMA = """
CREATE TABLE inventory (
    sku TEXT PRIMARY KEY,
    quantity INTEGER NOT NULL CHECK (quantity >= 0),
    price_cents INTEGER NOT NULL CHECK (price_cents >= 0)
)
"""

UPSERT = """
INSERT INTO inventory (sku, quantity, price_cents)
VALUES (?, ?, ?)
ON CONFLICT(sku) DO UPDATE SET
    quantity = excluded.quantity,
    price_cents = excluded.price_cents
"""


def apply_stock_updates(
    connection: sqlite3.Connection,
    rows: list[tuple[str, int, int]],
) -> int:
    for sku, quantity, price_cents in rows:
        if not sku or quantity < 0 or price_cents < 0:
            raise ValueError(f"invalid inventory row: {sku!r}")

    connection.execute("BEGIN IMMEDIATE")
    try:
        cursor = connection.executemany(UPSERT, rows)
        connection.commit()
    except Exception:
        connection.rollback()
        raise
    return cursor.rowcount


def main() -> None:
    connection = sqlite3.connect(":memory:")
    connection.executescript(SCHEMA)

    first = apply_stock_updates(connection, [("A-1", 5, 1200), ("B-2", 3, 800)])
    second = apply_stock_updates(connection, [("A-1", 9, 1150), ("C-3", 1, 450)])
    print("affected rows:", first, second)
    print(connection.execute(
        "SELECT sku, quantity, price_cents FROM inventory ORDER BY sku"
    ).fetchall())

    try:
        apply_stock_updates(connection, [("D-4", -1, 100)])
    except ValueError as exc:
        print(f"rejected: {exc}")
    assert connection.execute("SELECT COUNT(*) FROM inventory").fetchone()[0] == 3


if __name__ == "__main__":
    main()