import sqlite3
from collections.abc import Iterator
from contextlib import contextmanager

from folkugat_web.config.db import DB_FILE

Connection = sqlite3.Connection


@contextmanager
def get_connection(con: Connection | None = None) -> Iterator[Connection]:
    if con:
        yield con
    else:
        con = sqlite3.connect(DB_FILE)
        _ = con.execute('PRAGMA foreign_keys = ON;')
        try:
            yield con
            con.commit()
        except Exception:
            con.rollback()
            raise
        finally:
            con.close()