"""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()