Newer
Older
TIASDataHub / src / tiasdata / catalog / migrate.py
"""マイグレーション(連番の SQL ファイル)の適用.

適用済みの版を DB の `PRAGMA user_version` に記録し,それより新しいファイルだけを番号順に適用する.
マイグレーションはパッケージ内(`catalog/migrations/<系列>/NNNN_<内容>.sql`)に置き,
`importlib.resources` で読む(editable install でもパッケージとして入れても
同じように見つかるようにするため).
"""

import re
import sqlite3
from importlib import resources
from typing import Literal, NamedTuple

# 目録(catalog.sqlite)と対応表(identity.sqlite)は別の DB なので,版の系列も分ける
Series = Literal["catalog", "identity"]

_FILENAME_PATTERN = re.compile(r"^(\d{4})_[a-z0-9_]+\.sql$")


class Migration(NamedTuple):
    """1 つのマイグレーションファイル."""

    version: int
    name: str
    sql: str


class MigrationResult(NamedTuple):
    """マイグレーションの適用結果."""

    before: int
    after: int

    @property
    def applied(self) -> int:
        """今回適用したファイルの数."""
        return self.after - self.before


def load_migrations(series: Series) -> list[Migration]:
    """系列のマイグレーションを番号順に読む.

    Args:
        series: `catalog`(目録)または `identity`(対応表).

    Returns:
        版の昇順に並べたマイグレーション.

    Raises:
        RuntimeError: ファイル名の形が誤っている,または番号が 1 から連続していない
            (パッケージに同梱するファイルの不備なので内部論理エラーとして扱う).
    """
    directory = resources.files("tiasdata.catalog") / "migrations" / series
    migrations: list[Migration] = []
    for entry in directory.iterdir():
        if not entry.name.endswith(".sql"):
            continue
        match = _FILENAME_PATTERN.match(entry.name)
        if match is None:
            raise RuntimeError(
                f"マイグレーションのファイル名が NNNN_<内容>.sql でない: {entry.name}"
            )
        migrations.append(
            Migration(int(match.group(1)), entry.name, entry.read_text(encoding="utf-8"))
        )
    migrations.sort(key=lambda m: m.version)
    expected = list(range(1, len(migrations) + 1))
    if [m.version for m in migrations] != expected:
        raise RuntimeError(f"マイグレーション({series})の番号が 1 から連続していない")
    return migrations


def latest_version(series: Series) -> int:
    """系列の最新の版(同梱したマイグレーションの最大番号)を返す.

    Args:
        series: `catalog` または `identity`.

    Returns:
        最新の版.マイグレーションが無ければ 0.
    """
    return len(load_migrations(series))


def schema_version(conn: sqlite3.Connection) -> int:
    """DB に適用済みのスキーマの版(`PRAGMA user_version`)を返す.

    Args:
        conn: DB への接続.

    Returns:
        適用済みの版.未適用なら 0.
    """
    return conn.execute("PRAGMA user_version").fetchone()[0]


def require_latest(conn: sqlite3.Connection, series: Series, *, schema: str = "main") -> None:
    """DB のスキーマが最新の版であることを確かめる(取り込み・点検の前に呼ぶ).

    Args:
        conn: DB への接続.
        series: `catalog` または `identity`.
        schema: 接続の中のスキーマ名(`main` または ATTACH した名前).

    Raises:
        ValueError: スキーマが最新でない(`tiasdata init` で更新する),または
            ライブラリより新しい(`tiasdata` を更新する).
    """
    if not schema.isidentifier():
        raise RuntimeError("スキーマ名が識別子でない")
    # PRAGMA はプレースホルダを受け付けないため,コードの定数(識別子)だけを埋め込む
    (current,) = conn.execute(f"PRAGMA {schema}.user_version").fetchone()
    latest = latest_version(series)
    if current < latest:
        raise ValueError(
            f"DB({series})のスキーマの版 {current} が最新 {latest} でない."
            "先に `tiasdata init` を実行する"
        )
    if current > latest:
        raise ValueError(
            f"DB({series})のスキーマの版 {current} が,このライブラリの知る最新の版 "
            f"{latest} より新しい.tiasdata を更新する"
        )


def apply_migrations(conn: sqlite3.Connection, series: Series) -> MigrationResult:
    """未適用のマイグレーションを番号順に適用する.

    1 ファイルを 1 トランザクションで適用し,同じトランザクションの中で `user_version` を
    そのファイルの番号に更新する.途中のファイルで失敗したらそのファイルの変更は残らない
    (それより前のファイルは適用済みのまま).適用済みの DB に再度呼んでも何もしない.

    Args:
        conn: `connect()` で開いた接続.
        series: `catalog`(目録)または `identity`(対応表).

    Returns:
        適用前と適用後の版.

    Raises:
        ValueError: DB の版がこのライブラリの知る最新の版より新しい(ライブラリが古い).
        sqlite3.Error: SQL の実行に失敗した(同梱するマイグレーションの不備).
    """
    migrations = load_migrations(series)
    before = schema_version(conn)
    if before > len(migrations):
        raise ValueError(
            f"DB({series})のスキーマの版 {before} が,このライブラリの知る最新の版 "
            f"{len(migrations)} より新しい.tiasdata を更新する"
        )
    for migration in migrations[before:]:
        # executescript は直前のトランザクションを確定してから,スクリプトを自動コミットで流す.
        # 1 ファイルを 1 トランザクションにするため BEGIN・COMMIT で囲む.
        # PRAGMA はプレースホルダを受け付けないため版を埋め込む(値はファイル名から得た整数のみ)
        script = (
            f"BEGIN;\n{migration.sql}\nPRAGMA user_version = {int(migration.version)};\nCOMMIT;"
        )
        try:
            conn.executescript(script)
        except sqlite3.Error:
            if conn.in_transaction:
                conn.rollback()
            raise
    return MigrationResult(before, schema_version(conn))