paper_db/paper_db_common/scripts/delete_paper.py

#!/usr/bin/env python3
"""paper_idを指定してpapers.dbから論文を削除する。

PDFファイルそのものは削除しない。
既定はdry-run。実際に削除する場合は--applyを指定する。
"""

from __future__ import annotations

import argparse
import shutil
import sqlite3
import sys
from dataclasses import dataclass
from pathlib import Path

from common import CONFIG


@dataclass(frozen=True)
class Paper:
    paper_id: int
    title: str
    first_author: str | None
    year: int | None
    doi: str | None
    pdf_path: str | None
    file_sha256: str | None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="paper_idを指定してDBレコードを削除する"
    )
    parser.add_argument(
        "paper_id",
        type=int,
        help="削除するpapers.id",
    )
    parser.add_argument(
        "--database",
        type=Path,
        default=CONFIG.path.database,
        help="SQLiteデータベース",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="削除を実際にDBへ反映する",
    )
    parser.add_argument(
        "--no-backup",
        action="store_true",
        help="削除前バックアップを作成しない",
    )
    return parser.parse_args()


def open_database(
    database: Path,
) -> sqlite3.Connection:
    connection = sqlite3.connect(database)
    connection.row_factory = sqlite3.Row
    connection.execute("PRAGMA foreign_keys=ON")
    return connection


def load_paper(
    connection: sqlite3.Connection,
    paper_id: int,
) -> Paper | None:
    row = connection.execute(
        """
        SELECT
            id,
            title,
            first_author,
            year,
            doi,
            pdf_path,
            file_sha256
        FROM papers
        WHERE id = ?
        """,
        (paper_id,),
    ).fetchone()

    if row is None:
        return None

    return Paper(
        paper_id=row["id"],
        title=row["title"],
        first_author=row["first_author"],
        year=row["year"],
        doi=row["doi"],
        pdf_path=row["pdf_path"],
        file_sha256=row["file_sha256"],
    )


def print_paper(paper: Paper) -> None:
    print("=== 削除対象 ===")
    print(f"id:           {paper.paper_id}")
    print(f"title:        {paper.title}")
    print(
        "first_author: "
        f"{paper.first_author or '(none)'}"
    )
    print(f"year:         {paper.year or '(none)'}")
    print(f"doi:          {paper.doi or '(none)'}")
    print(
        f"pdf_path:     {paper.pdf_path or '(none)'}"
    )
    print(
        "file_sha256:  "
        f"{paper.file_sha256 or '(none)'}"
    )


def create_backup(database: Path) -> Path:
    backup = database.with_name(
        database.name + ".before-delete"
    )

    counter = 1

    while backup.exists():
        backup = database.with_name(
            database.name + f".before-delete.{counter}"
        )
        counter += 1

    shutil.copy2(database, backup)
    return backup


def delete_paper(
    connection: sqlite3.Connection,
    paper_id: int,
) -> None:
    """関連レコードを含めて1トランザクションで削除する。"""
    with connection:
        tables = (
            "paper_authors",
            "paper_keywords",
            "processing_log",
        )

        for table in tables:
            if table_exists(connection, table):
                connection.execute(
                    f"DELETE FROM {table} WHERE paper_id = ?",
                    (paper_id,),
                )

        cursor = connection.execute(
            "DELETE FROM papers WHERE id = ?",
            (paper_id,),
        )

        if cursor.rowcount != 1:
            raise RuntimeError(
                f"papers.id={paper_id}を削除できませんでした"
            )


def table_exists(
    connection: sqlite3.Connection,
    table: str,
) -> bool:
    row = connection.execute(
        """
        SELECT 1
        FROM sqlite_master
        WHERE type = 'table'
          AND name = ?
        """,
        (table,),
    ).fetchone()

    return row is not None


def validate_database(
    connection: sqlite3.Connection,
) -> None:
    integrity = connection.execute(
        "PRAGMA integrity_check"
    ).fetchone()[0]

    if integrity != "ok":
        raise RuntimeError(
            f"integrity_check: {integrity}"
        )

    foreign_keys = connection.execute(
        "PRAGMA foreign_key_check"
    ).fetchall()

    if foreign_keys:
        raise RuntimeError(
            "foreign_key_check: "
            f"{len(foreign_keys)}件"
        )


def main() -> int:
    args = parse_args()

    try:
        database = args.database.expanduser().resolve()

        if not database.is_file():
            raise FileNotFoundError(database)

        with open_database(database) as connection:
            paper = load_paper(
                connection,
                args.paper_id,
            )

            if paper is None:
                print(
                    f"paper_id={args.paper_id}は存在しません。",
                    file=sys.stderr,
                )
                return 1

            print_paper(paper)

            if not args.apply:
                print()
                print(
                    "削除するには --apply を指定してください。"
                )
                return 0

            if not args.no_backup:
                backup = create_backup(database)
                print()
                print(f"DBバックアップ: {backup}")

            delete_paper(
                connection,
                args.paper_id,
            )
            validate_database(connection)

        print()
        print(
            f"Deleted paper_id={args.paper_id}"
        )
        print("PDFファイルは削除していません。")
        return 0

    except KeyboardInterrupt:
        return 130

    except Exception as exc:
        print(
            f"delete_paper.py: {exc}",
            file=sys.stderr,
        )
        return 1


if __name__ == "__main__":
    raise SystemExit(main())