LDAPWeb/server-connection/scripts/geonet-indexer/index_project_metadata.py

#!/usr/bin/env python3
"""Crawl NAS project share metadata into PostgreSQL project_index (Fase 1 POC)."""

from __future__ import annotations

import argparse
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

import psycopg2
from psycopg2.extras import execute_values

SMB_HOST = os.environ.get("NAS_SMB_HOST", "10.100.1.10")
SMB_SHARE = os.environ.get("NAS_SMB_SHARE", "project")
SOURCE_ID = os.environ.get("PROJECT_SOURCE_ID", "project")
MOUNT_BASE = Path(os.environ.get("NAS_MOUNT_BASE", "/mnt/qnap"))
EXCLUDE_DIRS = {".@__thumb", "@Recycle", "@Recently-Snapshot", ".snapshot"}

DB = {
    "host": os.environ.get("PROJECT_SEARCH_DB_HOST", "10.100.1.25"),
    "port": int(os.environ.get("PROJECT_SEARCH_DB_PORT", "5432")),
    "dbname": os.environ.get("PROJECT_SEARCH_DB_DATABASE", "geonet_project_search"),
    "user": os.environ.get("PROJECT_SEARCH_DB_USERNAME", "geonet_project"),
    "password": os.environ.get("PROJECT_SEARCH_DB_PASSWORD", ""),
}

YEAR_RE = re.compile(r"^20\d{2}$")


def parse_years(raw: str) -> set[int] | None:
    raw = raw.strip()
    if not raw:
        return None
    return {int(part) for part in raw.split(",") if part.strip().isdigit()}


def to_smb_path(relative: str) -> str:
    rel = relative.replace("/", "\\")
    return f"\\\\{SMB_HOST}\\{SMB_SHARE}\\{rel}"


def client_name_from_parts(parts: list[str]) -> str | None:
    if len(parts) < 2:
        return None
    return parts[1][:255]


def year_from_parts(parts: list[str]) -> int | None:
    if not parts:
        return None
    if YEAR_RE.match(parts[0]):
        return int(parts[0])
    return None


def collect_rows(root: Path, years: set[int] | None, max_depth: int) -> list[tuple]:
    rows: list[tuple] = []
    if not root.is_dir():
        raise FileNotFoundError(f"Mount tidak ditemukan: {root}")

    for dirpath, dirnames, filenames in os.walk(root, topdown=True):
        dirnames[:] = [
            name
            for name in dirnames
            if name not in EXCLUDE_DIRS and not name.startswith(".")
        ]

        current = Path(dirpath)
        rel_dir = current.relative_to(root).as_posix()
        depth = 0 if rel_dir == "." else len(rel_dir.split("/"))

        if depth > max_depth:
            dirnames.clear()
            continue

        parts = [] if rel_dir == "." else rel_dir.split("/")
        year_folder = year_from_parts(parts)

        if years is not None and parts and YEAR_RE.match(parts[0]) and int(parts[0]) not in years:
            dirnames.clear()
            continue

        if rel_dir != ".":
            stat = current.stat()
            rows.append(
                (
                    SOURCE_ID,
                    rel_dir,
                    to_smb_path(rel_dir),
                    client_name_from_parts(parts),
                    year_folder,
                    True,
                    None,
                    None,
                    datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc),
                )
            )

        for filename in filenames:
            if filename.startswith(".") or filename.startswith("~$"):
                continue
            file_path = current / filename
            try:
                stat = file_path.stat()
            except OSError:
                continue
            rel_file = file_path.relative_to(root).as_posix()
            file_parts = rel_file.split("/")
            ext = file_path.suffix.lower().lstrip(".")[:16] or None
            rows.append(
                (
                    SOURCE_ID,
                    rel_file,
                    to_smb_path(rel_file),
                    client_name_from_parts(file_parts),
                    year_from_parts(file_parts),
                    False,
                    ext,
                    stat.st_size,
                    datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc),
                )
            )

    return rows


def upsert_rows(conn, rows: list[tuple]) -> int:
    if not rows:
        return 0

    sql = """
        INSERT INTO project_index (
            source_id, relative_path, full_smb_path, client_name, year_folder,
            is_directory, file_ext, size_bytes, mtime, indexed_at
        ) VALUES %s
        ON CONFLICT (source_id, relative_path) DO UPDATE SET
            full_smb_path = EXCLUDED.full_smb_path,
            client_name = EXCLUDED.client_name,
            year_folder = EXCLUDED.year_folder,
            is_directory = EXCLUDED.is_directory,
            file_ext = EXCLUDED.file_ext,
            size_bytes = EXCLUDED.size_bytes,
            mtime = EXCLUDED.mtime,
            indexed_at = NOW()
    """
    template = "(%s,%s,%s,%s,%s,%s,%s,%s,%s,NOW())"
    with conn.cursor() as cur:
        batch_size = 1000
        total = 0
        for i in range(0, len(rows), batch_size):
            chunk = rows[i : i + batch_size]
            execute_values(cur, sql, chunk, template=template, page_size=batch_size)
            total += len(chunk)
        conn.commit()
    return total


def main() -> int:
    parser = argparse.ArgumentParser(description="Index NAS project metadata")
    parser.add_argument("--years", default=os.environ.get("PROJECT_INDEX_YEARS", "2020,2021,2022,2023,2024,2025,2026"))
    parser.add_argument("--max-depth", type=int, default=int(os.environ.get("PROJECT_INDEX_MAX_DEPTH", "4")))
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if not DB["password"]:
        print("PROJECT_SEARCH_DB_PASSWORD wajib di-set", file=sys.stderr)
        return 1

    root = MOUNT_BASE / SMB_SHARE
    years = parse_years(args.years)
    rows = collect_rows(root, years, args.max_depth)
    print(f"Collected {len(rows)} rows from {root}")

    if args.dry_run:
        return 0

    conn = psycopg2.connect(**DB)
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO project_index_runs (source_id, status)
                VALUES (%s, 'running')
                RETURNING id
                """,
                (SOURCE_ID,),
            )
            run_id = cur.fetchone()[0]
            conn.commit()

        count = upsert_rows(conn, rows)

        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE project_index_runs
                SET finished_at = NOW(), status = 'completed', rows_upserted = %s
                WHERE id = %s
                """,
                (count, run_id),
            )
            conn.commit()
        print(f"Upserted {count} rows (run_id={run_id})")
    except Exception as exc:
        conn.rollback()
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE project_index_runs
                SET finished_at = NOW(), status = 'failed', error_message = %s
                WHERE id = (SELECT MAX(id) FROM project_index_runs WHERE source_id = %s)
                """,
                (str(exc)[:2000], SOURCE_ID),
            )
            conn.commit()
        raise
    finally:
        conn.close()

    return 0


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