dp-zp-agent/scripts/rebuild_index.py

from __future__ import annotations

import argparse
import errno
import fcntl
import json
import os
import subprocess
import sys
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, TextIO

from rich import print


PROJECT_ROOT = Path(__file__).resolve().parents[1]

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


from scripts.build_chunks import build_chunks
from scripts.build_sqlite_index import build_database
from scripts.common import DATA_DIR, DB_FILE, ZPWIKI_ROOT
from scripts.scan_zpwiki import scan_pages


REINDEX_LOCK_FILE = DATA_DIR / ".reindex.lock"
GIT_PULL_TIMEOUT_SECONDS = 120


class ReindexInProgressError(RuntimeError):
    """Iný proces už drží zámok reindexovania."""


@contextmanager
def acquire_reindex_lock(
    lock_file: Path = REINDEX_LOCK_FILE,
) -> Iterator[TextIO]:
    """Získa neblokujúci procesový zámok nad spoločným data volume."""
    lock_file.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    handle = lock_file.open(
        "a+",
        encoding="utf-8",
    )

    try:
        try:
            fcntl.flock(
                handle.fileno(),
                fcntl.LOCK_EX | fcntl.LOCK_NB,
            )

        except OSError as error:
            if error.errno in {
                errno.EACCES,
                errno.EAGAIN,
            }:
                raise ReindexInProgressError(
                    "Reindexovanie už prebieha"
                ) from error

            raise

        handle.seek(0)
        handle.truncate()

        json.dump(
            {
                "pid": os.getpid(),
                "started_at_unix": time.time(),
            },
            handle,
        )

        handle.flush()

        yield handle

    finally:
        try:
            fcntl.flock(
                handle.fileno(),
                fcntl.LOCK_UN,
            )
        except OSError:
            pass

        handle.close()


def git_pull(
    repo_path: Path = ZPWIKI_ROOT,
) -> None:
    if not repo_path.exists():
        raise RuntimeError(
            f"ZPWIKI_ROOT neexistuje: {repo_path}"
        )

    if not (repo_path / ".git").exists():
        raise RuntimeError(
            f"Nie je to git repozitár: {repo_path}"
        )

    try:
        result = subprocess.run(
            [
                "git",
                "pull",
                "--ff-only",
            ],
            cwd=repo_path,
            text=True,
            capture_output=True,
            timeout=GIT_PULL_TIMEOUT_SECONDS,
            check=False,
        )

    except subprocess.TimeoutExpired as error:
        raise RuntimeError(
            "Git pull prekročil limit "
            f"{GIT_PULL_TIMEOUT_SECONDS} sekúnd"
        ) from error

    if result.stdout:
        print(result.stdout.strip())

    if result.stderr:
        print(result.stderr.strip())

    if result.returncode != 0:
        raise RuntimeError(
            "Git pull zlyhal s návratovým kódom "
            f"{result.returncode}"
        )


def rebuild_index(
    pull_git: bool = False,
    *,
    lock_file: Path = REINDEX_LOCK_FILE,
) -> dict:
    with acquire_reindex_lock(lock_file):
        start = time.monotonic()

        print(
            f"[green]ZPWIKI_ROOT:[/green] "
            f"{ZPWIKI_ROOT}"
        )

        if pull_git:
            git_pull()

        documents = scan_pages()
        chunks = build_chunks()
        counts = build_database()

        duration = round(
            time.monotonic() - start,
            2,
        )

        return {
            "duration_seconds": duration,
            "documents_scanned": len(documents),
            "chunks_created": len(chunks),
            "counts": counts,
            "database_path": str(DB_FILE),
        }


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Obnoví JSON súbory "
            "a SQLite FTS5 index."
        )
    )

    parser.add_argument(
        "--pull",
        action="store_true",
        help=(
            "Pred reindexovaním spustí "
            "git pull --ff-only."
        ),
    )

    args = parser.parse_args()

    try:
        result = rebuild_index(
            pull_git=args.pull
        )

    except ReindexInProgressError as error:
        raise SystemExit(
            str(error)
        ) from error

    counts = result["counts"]

    print("[green]Reindex hotový.[/green]")
    print(
        f"Trvanie: "
        f"{result['duration_seconds']} s"
    )
    print(
        f"Dokumentov: "
        f"{counts['documents']}"
    )
    print(
        f"Chunkov: "
        f"{counts['chunks']}"
    )

    if "fts_chunks" in counts:
        print(
            f"FTS5 chunkov: "
            f"{counts['fts_chunks']}"
        )

    print(
        f"Tag záznamov: "
        f"{counts['tags']}"
    )
    print(
        f"Kategória záznamov: "
        f"{counts['categories']}"
    )


if __name__ == "__main__":
    main()