database/connection.py

"""Shared SQLite connection helper.

Centralizes what psycopg used to provide automatically: dict-shaped rows,
`with`-block commit/rollback/close semantics (including on `conn.cursor()`,
which stock sqlite3.Cursor does not support as a context manager), and
date/datetime round-tripping.
"""

from __future__ import annotations

import sqlite3
from datetime import date, datetime, timezone
from pathlib import Path

sqlite3.register_adapter(date, lambda d: d.isoformat())
sqlite3.register_adapter(
    datetime,
    lambda dt: (dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)).astimezone(timezone.utc).isoformat(),
)
sqlite3.register_converter("DATE", lambda v: date.fromisoformat(v.decode()))
sqlite3.register_converter("TIMESTAMP", lambda v: datetime.fromisoformat(v.decode().replace("Z", "+00:00")))


def _dict_row(cursor: sqlite3.Cursor, row: tuple) -> dict:
    return dict(zip((column[0] for column in cursor.description), row))


class Cursor(sqlite3.Cursor):
    def __enter__(self) -> "Cursor":
        return self

    def __exit__(self, *exc) -> None:
        self.close()


class Connection(sqlite3.Connection):
    def cursor(self, factory=None) -> Cursor:
        return super().cursor(factory or Cursor)

    def __exit__(self, *exc) -> None:
        super().__exit__(*exc)
        self.close()


def _resolve_path(database_url: str) -> str:
    if "sqlite:///" in database_url:
        return database_url.split("sqlite:///", 1)[-1]
    return database_url


def connect(database_url: str) -> Connection:
    path = _resolve_path(database_url)
    if path != ":memory:":
        Path(path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(
        path,
        factory=Connection,
        detect_types=sqlite3.PARSE_DECLTYPES,
        timeout=30,
    )
    conn.row_factory = _dict_row
    conn.execute("PRAGMA foreign_keys = ON")
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA busy_timeout = 5000")
    return conn