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

import sqlite3
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