Newer
Older
TIASDataHub / src / tiasdata / export.py
"""バンドルの書き出し(`tiasdata export`.SPEC_03「バンドル」).

目録・ストア・設定を使うため,読み込み側(`tiasdata.bundle`)とはモジュールを分ける
(`import tiasdata` で目録の依存を計算サーバに持ち込まないため.SPEC_01「依存の向き」).
形の組み立て・読み込み・照合は `tiasdata.bundle` のものを使う.

書き出しは「一時フォルダ(`.tmp-` で始まる名前)に作る → 開いて形を確かめる → 正式な名前に
変える」の順で行い,書きかけのバンドルを正式な名前で残さない.既にあるバンドルは上書きせず,
版の記録と照合して使う(食い違えば停止する).

出力・例外メッセージにはパスを出さない(SPEC_03「出力の規則」).
"""

import os
import shutil
import sqlite3
import tempfile
from contextlib import closing
from pathlib import Path

from pydantic import BaseModel, ConfigDict

from tiasdata import bundle, catalog, datasets, store
from tiasdata.catalog import versions
from tiasdata.config import Config
from tiasdata.manifest import FILES_DIRNAME, bundle_file_path, manifest_sha256, sha256_file

# 書きかけのバンドルのフォルダ名の接頭辞(ストアの一時ファイルと同じ印)
TEMP_PREFIX = store.TEMP_PREFIX


class ExportResult(BaseModel):
    """書き出しの結果.

    Attributes:
        name: データセット名.
        version: 版.
        files: ファイルの数.
        size_bytes: ファイルのサイズの合計.
        manifest_sha256: manifest.csv の SHA-256.
        reused: 既にあったバンドルを照合して使ったか(False なら今回書き出した).
    """

    model_config = ConfigDict(frozen=True)

    name: str
    version: int
    files: int
    size_bytes: int
    manifest_sha256: str
    reused: bool


class _Planned(BaseModel):
    # 書き出す(または照合する)バンドルの中身.ファイルはストアからコピーする
    model_config = ConfigDict(frozen=True)

    info: bundle.BundleInfo
    manifest: bytes
    dataset_json: bytes
    sha256sums: bytes
    files: list[tuple[str, str]]  # (SHA-256, 拡張子)


def _plan(conn: sqlite3.Connection, name: str, version: int) -> _Planned:
    record = versions.get_version(conn, name, version)
    if record is None:
        raise ValueError(f"データセット '{name}' の版 {version} が無い")
    # 版の構成から作り直し,記録した SHA-256 と照合してから書き出す(目録の書き換え等で
    # 版を作ったときと違う manifest を配らないため)
    data = datasets.rebuild_manifest(conn, name, version)
    digest = manifest_sha256(data)
    if digest != record.manifest_sha256:
        raise ValueError(
            f"データセット '{name}' の版 {version}: 構成から作り直した manifest の SHA-256 が"
            "版の記録と一致しない.書き出さずに停止した.tiasdata check で点検する"
        )
    files = [(m.sha256, m.ext) for m in versions.version_members(conn, name, version)]
    info = bundle.BundleInfo(
        format=bundle.BUNDLE_FORMAT,
        name=name,
        version=version,
        created_at=record.created_at,
        hub_commit=record.hub_commit,
        files=len(files),
        manifest_sha256=digest,
    )
    return _Planned(
        info=info,
        manifest=data,
        dataset_json=bundle.build_dataset_json(info),
        sha256sums=bundle.build_sha256sums(
            (sha256, bundle_file_path(sha256, ext)) for sha256, ext in files
        ),
        files=files,
    )


def export_bundle(
    conn: sqlite3.Connection, *, bundles_dir: Path, store_dir: Path, name: str, version: int
) -> ExportResult:
    """凍結した版のバンドルを `bundles_dir/<名前>/v<版>/` に書き出す.

    既にあれば上書きせず,dataset.json・manifest.csv・SHA256SUMS が版の記録から作ったものと
    バイト単位で同じで,ファイルの照合(`Bundle.verify`)が通ることを確かめて使う.

    Args:
        conn: 目録への接続(スキーマ最新).
        bundles_dir: バンドルの置き場所(`<catalog.root>/bundles`.`tiasdata init` で作成済み).
        store_dir: ストアのフォルダ.
        name: データセット名.
        version: 版.

    Returns:
        書き出しの結果.

    Raises:
        ValueError: 名前・版の形が誤っている,版が無い,作り直した manifest が記録と一致しない,
            バンドルの置き場所が無い,ストアのファイルが無い・照合が合わない,
            または既にあるバンドルが版の記録と食い違う.
    """
    relative = bundle.bundle_relative_dir(name, version)
    if not bundles_dir.is_dir():
        raise ValueError("バンドルの置き場所(bundles)が無い(先に tiasdata init を実行する)")
    planned = _plan(conn, name, version)
    target = bundles_dir / relative
    if target.exists():
        return _reuse_existing(target, planned)

    parent = target.parent
    parent.mkdir(exist_ok=True)
    work = Path(tempfile.mkdtemp(prefix=f"{TEMP_PREFIX}{target.name}-", dir=parent))
    try:
        size_bytes = _write_bundle(work, planned, store_dir)
        # 書いたものを読み込み側で開き直し,使う側と同じ規則で読めることを確かめる
        bundle.open_bundle(work)
        published = _publish(work, target)
    except BaseException:
        shutil.rmtree(work, ignore_errors=True)
        raise
    if not published:
        # 照合の間に同じ版が書き出された(同時に実行した等)
        shutil.rmtree(work, ignore_errors=True)
        return _reuse_existing(target, planned)
    return _result(planned, size_bytes, reused=False)


def _write_bundle(work: Path, planned: _Planned, store_dir: Path) -> int:
    files_dir = work / FILES_DIRNAME
    files_dir.mkdir()
    size_bytes = 0
    for sha256, ext in planned.files:
        source = store.store_path(store_dir, sha256, ext)
        if not source.is_file():
            raise ValueError(
                f"ストアにファイルが無い(SHA-256 {sha256}).tiasdata check で点検する"
            )
        destination = work / bundle_file_path(sha256, ext)
        # copyfile は中身だけを写す(ストアの読み取り専用の属性をバンドルに持ち込まない.
        # バンドルは再生成でき,受け渡し後に消すため)
        shutil.copyfile(source, destination)
        # 書き込んだファイルを読み直して照合する(ストアの破損・書き込みの誤りを検出するため)
        if sha256_file(destination) != sha256:
            raise ValueError(
                f"バンドルにコピーしたファイルの照合が合わない(SHA-256 {sha256})."
                "tiasdata check でストアを点検する"
            )
        size_bytes += destination.stat().st_size
    (work / bundle.MANIFEST_CSV).write_bytes(planned.manifest)
    (work / bundle.SHA256SUMS).write_bytes(planned.sha256sums)
    (work / bundle.DATASET_JSON).write_bytes(planned.dataset_json)
    return size_bytes


def _publish(work: Path, target: Path) -> bool:
    """一時フォルダを正式な名前にする.既に同じ名前があれば何もせず False を返す."""
    if target.exists():
        return False
    try:
        # Windows は移動先があると FileExistsError.POSIX は空でないフォルダへの rename が失敗する
        os.rename(work, target)
    except FileExistsError:
        return False
    except OSError:
        if target.exists():
            return False
        raise
    return True


def _reuse_existing(target: Path, planned: _Planned) -> ExportResult:
    label = f"データセット '{planned.info.name}' の版 {planned.info.version}"
    message = f"{label}: 既にあるバンドルが版の記録と食い違う.上書きせずに停止した"
    if not target.is_dir():
        raise ValueError(f"{message}(フォルダでない)")
    for filename, expected in (
        (bundle.DATASET_JSON, planned.dataset_json),
        (bundle.MANIFEST_CSV, planned.manifest),
        (bundle.SHA256SUMS, planned.sha256sums),
    ):
        path = target / filename
        if not path.is_file() or path.read_bytes() != expected:
            raise ValueError(f"{message}({filename} が無い,または違う).消してから書き出し直す")
    try:
        existing = bundle.open_bundle(target)
        existing.verify()
    except ValueError as e:
        raise ValueError(f"{message}({e}).消してから書き出し直す") from e
    size_bytes = sum(existing.file_path(row).stat().st_size for row in existing.rows)
    return _result(planned, size_bytes, reused=True)


def _result(planned: _Planned, size_bytes: int, *, reused: bool) -> ExportResult:
    return ExportResult(
        name=planned.info.name,
        version=planned.info.version,
        files=planned.info.files,
        size_bytes=size_bytes,
        manifest_sha256=planned.info.manifest_sha256,
        reused=reused,
    )


def run_export(config: Config, name: str, version: int) -> ExportResult:
    """設定の目録・ストアから,凍結した版のバンドルを `bundles/` に書き出す.

    Args:
        config: 設定.
        name: データセット名.
        version: 版.

    Returns:
        書き出しの結果.

    Raises:
        ValueError: 目録が無いか古い,または書き出せない(`export_bundle` を参照).
    """
    # 名前・版の形の誤りを,目録が無い・古い等のエラーより先に知らせる(引数を先に直せるように)
    bundle.bundle_relative_dir(name, version)
    paths = config.catalog
    with closing(catalog.connect_latest(paths.catalog_path)) as conn:
        return export_bundle(
            conn,
            bundles_dir=paths.bundles_dir,
            store_dir=paths.store_dir,
            name=name,
            version=version,
        )