massive improve
This commit is contained in:
Сергей Маринкевич
+323
@@ -0,0 +1,323 @@
|
||||
#!/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
|
||||
|
||||
|
||||
def iter_result_files(root: str) -> Iterable[str]:
|
||||
for dirpath, _dirnames, filenames in os.walk(root):
|
||||
for name in filenames:
|
||||
if name.endswith(".json"):
|
||||
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'
|
||||
)
|
||||
"""
|
||||
)
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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())
|
||||
|
||||
Reference in New Issue
Block a user