from __future__ import annotations

from contextlib import contextmanager
from sqlite3 import Connection
from typing import Iterator


@contextmanager
def unit_of_work(conn: Connection) -> Iterator[Connection]:
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise