mattermost-dl/mm_index.py

#!/usr/bin/env python3
import argparse
import datetime as dt
import json
import os
import sqlite3
from typing import Any, Dict, Iterable, List, Optional, Tuple

import mm_users


def iter_result_files(root: str) -> Iterable[str]:
    for dirpath, _dirnames, filenames in os.walk(root):
        for name in filenames:
            if name == "users.json" or not name.endswith(".json"):
                continue
            yield os.path.join(dirpath, name)


def parse_post_created(s: str) -> Optional[str]:
    try:
        created = dt.datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc)
    except ValueError:
        return None
    return created.isoformat()


def parse_exported_at(channel_obj: Dict[str, Any]) -> Optional[str]:
    raw = channel_obj.get("exported_at")
    if not isinstance(raw, str):
        return None
    try:
        exported = dt.datetime.strptime(raw, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc)
    except ValueError:
        return None
    return exported.isoformat()


def connect_db(db_path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS sources (
            path TEXT PRIMARY KEY,
            mtime_ns INTEGER NOT NULL,
            size INTEGER NOT NULL,
            indexed_at TEXT NOT NULL
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS posts (
            post_id TEXT PRIMARY KEY,
            idx INTEGER,
            created TEXT NOT NULL,
            username TEXT NOT NULL,
            message TEXT NOT NULL,
            files_json TEXT,
            channel_id TEXT,
            channel_name TEXT,
            channel_display_name TEXT,
            team TEXT,
            source_path TEXT NOT NULL,
            exported_at TEXT
        )
        """
    )
    conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_created ON posts(created)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_username ON posts(username)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_channel_display ON posts(channel_display_name)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_posts_team ON posts(team)")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS attachments (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            post_id TEXT NOT NULL,
            filename TEXT NOT NULL,
            filepath TEXT NOT NULL,
            is_text INTEGER NOT NULL,
            FOREIGN KEY(post_id) REFERENCES posts(post_id) ON DELETE CASCADE
        )
        """
    )
    conn.execute("CREATE INDEX IF NOT EXISTS idx_attachments_post_id ON attachments(post_id)")
    conn.execute(
        """
        CREATE VIRTUAL TABLE IF NOT EXISTS posts_fts
        USING fts5(
            post_id UNINDEXED,
            message,
            tokenize='unicode61'
        )
        """
    )
    conn.execute(
        """
        CREATE VIRTUAL TABLE IF NOT EXISTS attachments_fts
        USING fts5(
            attachment_id UNINDEXED,
            post_id UNINDEXED,
            filename,
            content,
            tokenize='unicode61'
        )
        """
    )
    mm_users.ensure_users_schema(conn)
    return conn


def get_source_state(conn: sqlite3.Connection, path: str) -> Optional[Tuple[int, int]]:
    row = conn.execute("SELECT mtime_ns, size FROM sources WHERE path = ?", (path,)).fetchone()
    if not row:
        return None
    return int(row["mtime_ns"]), int(row["size"])


def source_changed(path: str, known_state: Optional[Tuple[int, int]]) -> bool:
    st = os.stat(path)
    if known_state is None:
        return True
    return known_state != (st.st_mtime_ns, st.st_size)


def load_posts(path: str) -> Optional[Tuple[Dict[str, Any], List[Dict[str, Any]]]]:
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError):
        return None

    if not isinstance(data, dict):
        return None

    channel_obj = data.get("channel")
    posts = data.get("posts")
    if not isinstance(channel_obj, dict) or not isinstance(posts, list):
        return None
    return channel_obj, posts


def _delete_post_search_rows(conn: sqlite3.Connection, post_id: str) -> None:
    conn.execute("DELETE FROM attachments_fts WHERE post_id = ?", (post_id,))
    conn.execute("DELETE FROM attachments WHERE post_id = ?", (post_id,))
    conn.execute("DELETE FROM posts_fts WHERE post_id = ?", (post_id,))


def _delete_post(conn: sqlite3.Connection, post_id: str) -> None:
    _delete_post_search_rows(conn, post_id)
    conn.execute("DELETE FROM posts WHERE post_id = ?", (post_id,))


def _index_attachments(
    conn: sqlite3.Connection,
    post_id: str,
    files: Any,
    idx: Optional[int],
    base_dir: str,
    max_text_size: int,
) -> None:
    if not isinstance(files, list) or not files:
        return
    for filename in files:
        if not isinstance(filename, str) or not filename:
            continue
        disk_name = f"{idx:03d}_{filename}" if isinstance(idx, int) else filename
        disk_path = os.path.join(base_dir, disk_name)
        is_text = 0
        content: Optional[str] = None
        try:
            st = os.stat(disk_path)
            if st.st_size <= max_text_size:
                with open(disk_path, "r", encoding="utf-8") as f:
                    content = f.read()
                is_text = 1
        except (OSError, UnicodeDecodeError):
            is_text = 0
            content = None

        cur = conn.execute(
            """
            INSERT INTO attachments (post_id, filename, filepath, is_text)
            VALUES (?, ?, ?, ?)
            """,
            (post_id, filename, disk_path, is_text),
        )
        attachment_id = cur.lastrowid
        conn.execute(
            """
            INSERT INTO attachments_fts (attachment_id, post_id, filename, content)
            VALUES (?, ?, ?, ?)
            """,
            (attachment_id, post_id, filename, content if is_text else None),
        )


def _content_changed(prev: sqlite3.Row, message: str, username: str, created: str, files_json: Optional[str]) -> bool:
    return (
        prev["message"] != message
        or prev["username"] != username
        or prev["created"] != created
        or prev["files_json"] != files_json
    )


def _meta_changed(
    prev: sqlite3.Row,
    idx: Optional[int],
    channel_id: str,
    channel_name: str,
    channel_display_name: str,
    team: str,
    path: str,
    exported_at: Optional[str],
) -> bool:
    return (
        prev["idx"] != idx
        or prev["channel_id"] != channel_id
        or prev["channel_name"] != channel_name
        or prev["channel_display_name"] != channel_display_name
        or prev["team"] != team
        or prev["source_path"] != path
        or prev["exported_at"] != exported_at
    )


def upsert_file(conn: sqlite3.Connection, path: str, *, force: bool = False) -> Tuple[int, int, int]:
    """Index posts from one export JSON.

    Returns (upserted_or_deleted, invalid_skipped, unchanged).
    When force=False, only new/changed posts rebuild FTS/attachments; idx-only
    renumbers get a cheap posts-row UPDATE. force=True rebuilds every post
    (--reindex-all).
    """
    loaded = load_posts(path)
    if loaded is None:
        return 0, 0, 0

    channel_obj, posts = loaded
    channel_id = str(channel_obj.get("id") or "")
    channel_name = str(channel_obj.get("name") or "")
    channel_display_name = str(channel_obj.get("display_name") or "")
    team = str(channel_obj.get("team") or "")
    exported_at = parse_exported_at(channel_obj)

    existing = {
        row["post_id"]: row
        for row in conn.execute(
            """
            SELECT post_id, idx, created, username, message, files_json,
                   channel_id, channel_name, channel_display_name, team,
                   source_path, exported_at
            FROM posts
            WHERE source_path = ?
            """,
            (path,),
        )
    }

    base_dir = os.path.dirname(path)
    max_text_size = 2 * 1024 * 1024  # 2 MiB

    inserted_or_updated = 0
    unchanged = 0
    skipped = 0
    seen_ids = set()

    for p in posts:
        post_id = p.get("id")
        created_raw = p.get("created")
        username = p.get("username")
        message = p.get("message")
        if not isinstance(post_id, str) or not post_id:
            skipped += 1
            continue
        if not isinstance(created_raw, str) or not isinstance(username, str):
            skipped += 1
            continue
        if not isinstance(message, str):
            message = ""

        created = parse_post_created(created_raw)
        if created is None:
            skipped += 1
            continue

        files = p.get("files")
        files_json = json.dumps(files, ensure_ascii=False) if isinstance(files, list) else None
        idx = p.get("idx") if isinstance(p.get("idx"), int) else None
        seen_ids.add(post_id)
        prev = existing.get(post_id)

        if (
            not force
            and prev is not None
            and not _content_changed(prev, message, username, created, files_json)
        ):
            if _meta_changed(
                prev,
                idx,
                channel_id,
                channel_name,
                channel_display_name,
                team,
                path,
                exported_at,
            ):
                conn.execute(
                    """
                    UPDATE posts SET
                        idx = ?,
                        channel_id = ?,
                        channel_name = ?,
                        channel_display_name = ?,
                        team = ?,
                        source_path = ?,
                        exported_at = ?
                    WHERE post_id = ?
                    """,
                    (
                        idx,
                        channel_id,
                        channel_name,
                        channel_display_name,
                        team,
                        path,
                        exported_at,
                        post_id,
                    ),
                )
            unchanged += 1
            continue

        conn.execute(
            """
            INSERT INTO posts (
                post_id, idx, created, username, message, files_json,
                channel_id, channel_name, channel_display_name, team,
                source_path, exported_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(post_id) DO UPDATE SET
                idx = excluded.idx,
                created = excluded.created,
                username = excluded.username,
                message = excluded.message,
                files_json = excluded.files_json,
                channel_id = excluded.channel_id,
                channel_name = excluded.channel_name,
                channel_display_name = excluded.channel_display_name,
                team = excluded.team,
                source_path = excluded.source_path,
                exported_at = excluded.exported_at
            """,
            (
                post_id,
                idx,
                created,
                username,
                message,
                files_json,
                channel_id,
                channel_name,
                channel_display_name,
                team,
                path,
                exported_at,
            ),
        )
        _delete_post_search_rows(conn, post_id)
        conn.execute("INSERT INTO posts_fts(post_id, message) VALUES (?, ?)", (post_id, message))
        _index_attachments(conn, post_id, files, idx, base_dir, max_text_size)
        inserted_or_updated += 1

    for post_id in existing:
        if post_id not in seen_ids:
            _delete_post(conn, post_id)
            inserted_or_updated += 1

    st = os.stat(path)
    conn.execute(
        """
        INSERT INTO sources(path, mtime_ns, size, indexed_at)
        VALUES (?, ?, ?, ?)
        ON CONFLICT(path) DO UPDATE SET
            mtime_ns = excluded.mtime_ns,
            size = excluded.size,
            indexed_at = excluded.indexed_at
        """,
        (path, st.st_mtime_ns, st.st_size, dt.datetime.now(dt.timezone.utc).isoformat()),
    )
    return inserted_or_updated, skipped, unchanged


def main() -> int:
    ap = argparse.ArgumentParser(description="Build/update sqlite index for Mattermost export JSON dumps.")
    ap.add_argument("-r", "--root", default="results", help="Root folder containing dump directories.")
    ap.add_argument("--db", default="mattermost_index.sqlite", help="SQLite index database file.")
    ap.add_argument("--reindex-all", action="store_true", help="Re-index all files, ignoring sources cache.")
    args = ap.parse_args()

    conn = connect_db(args.db)
    files = list(iter_result_files(args.root))

    processed_files = 0
    skipped_files = 0
    rows_upserted = 0
    rows_skipped = 0
    rows_unchanged = 0

    with conn:
        for path in files:
            known_state = None if args.reindex_all else get_source_state(conn, path)
            if not args.reindex_all and not source_changed(path, known_state):
                skipped_files += 1
                continue
            processed_files += 1
            upserted, skipped, unchanged = upsert_file(conn, path, force=args.reindex_all)
            rows_upserted += upserted
            rows_skipped += skipped
            rows_unchanged += unchanged

        users_json = mm_users.find_users_json(args.root)
        if users_json:
            mm_users.upsert_users(conn, mm_users.load_users_json(users_json))

    total_posts = conn.execute("SELECT COUNT(*) AS c FROM posts").fetchone()["c"]
    total_sources = conn.execute("SELECT COUNT(*) AS c FROM sources").fetchone()["c"]
    conn.close()

    print(f"DB: {args.db}")
    print(f"Files scanned: {len(files)} (processed={processed_files}, unchanged_skipped={skipped_files})")
    print(
        f"Rows upserted: {rows_upserted}, unchanged_posts: {rows_unchanged}, "
        f"invalid_posts_skipped: {rows_skipped}"
    )
    print(f"Totals: posts={total_posts}, sources={total_sources}")
    return 0


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