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 upsert_file(conn: sqlite3.Connection, path: str) -> Tuple[int, int]:
    loaded = load_posts(path)
    if loaded is None:
        return 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)

    # Remove any existing attachments for posts coming from this source file
    conn.execute(
        """
        DELETE FROM attachments
        WHERE post_id IN (SELECT post_id FROM posts WHERE source_path = ?)
        """,
        (path,),
    )
    conn.execute(
        """
        DELETE FROM attachments_fts
        WHERE post_id IN (SELECT post_id FROM posts WHERE source_path = ?)
        """,
        (path,),
    )

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

    inserted_or_updated = 0
    skipped = 0
    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

        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,
            ),
        )
        conn.execute("DELETE FROM posts_fts WHERE post_id = ?", (post_id,))
        conn.execute("INSERT INTO posts_fts(post_id, message) VALUES (?, ?)", (post_id, message))

        # Index attachments (file names and, for text files, contents)
        if isinstance(files, list) and files:
            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),
                )
        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


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

    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 = upsert_file(conn, path)
            rows_upserted += upserted
            rows_skipped += skipped

        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}, invalid_posts_skipped: {rows_skipped}")
    print(f"Totals: posts={total_posts}, sources={total_sources}")
    return 0


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