Newer
Older
TIASDataHub / src / tiasdata / catalog / connection.py
"""SQLite の接続."""

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


def connect(path: Path, *, create: bool = False) -> sqlite3.Connection:
    """SQLite の DB を開き,外部キー制約を有効にする.

    Args:
        path: DB ファイルのパス.
        create: True なら DB ファイルが無いときに作る(親フォルダは作らない).
            False なら DB ファイルが無いときに停止する(`tiasdata init` の前に誤って空の目録を
            作らないため).

    Returns:
        `PRAGMA foreign_keys = ON` を有効にした接続.

    Raises:
        ValueError: `create` が False で DB ファイルが無い,または DB を開けない.
    """
    mode = "rwc" if create else "rw"
    uri = f"{path.resolve().as_uri()}?mode={mode}"
    try:
        conn = sqlite3.connect(uri, uri=True)
    except sqlite3.OperationalError as e:
        if not create and not path.exists():
            raise ValueError("目録が無い.先に `tiasdata init` を実行する") from e
        raise ValueError(f"DB を開けない: {e}") from e
    # 外部キー制約は接続ごとの設定で,既定では無効(https://www.sqlite.org/foreignkeys.html)
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


@contextmanager
def write_transaction(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
    """書き込みのトランザクションを始め,成功したら確定,失敗したら取り消す.

    `BEGIN IMMEDIATE` で最初に書き込みのロックを取る(照会してから登録するまでの間に
    別の取り込みが割り込まないようにするため).ATTACH した DB(対応表)への書き込みも
    同じトランザクションに含まれる.SQLite の制約違反は `ValueError` に変換する
    (error-handling.md).

    Args:
        conn: `connect()` で開いた接続(トランザクションの外であること).

    Yields:
        同じ接続.

    Raises:
        ValueError: 制約違反で登録できなかった(変更は取り消される).
    """
    if conn.in_transaction:
        raise RuntimeError("既にトランザクションの中で write_transaction を呼んだ")
    conn.execute("BEGIN IMMEDIATE")
    try:
        yield conn
    except sqlite3.IntegrityError as e:
        conn.rollback()
        # SQLite の制約違反のメッセージは制約・列の名前だけで,値を含まない
        raise ValueError(f"目録の制約に反するため登録を取り消した: {e}") from e
    except BaseException:
        conn.rollback()
        raise
    else:
        conn.commit()