massive improve
This commit is contained in:
+137
-29
@@ -1,13 +1,14 @@
|
||||
import os
|
||||
import pickle
|
||||
import sqlite3
|
||||
from datetime import datetime, date
|
||||
from datetime import datetime, date, timezone
|
||||
from typing import Tuple, Dict, List, Optional
|
||||
import getpass
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
from mattermostdriver import Driver
|
||||
from mattermostdriver.exceptions import ResourceNotFound
|
||||
import pathlib
|
||||
import json
|
||||
|
||||
@@ -184,6 +185,11 @@ def select_channel(d: Driver, team: str, my_user_id: str, user_id_to_name: Dict[
|
||||
# The channel name consists of two user ids connected by a double underscore
|
||||
user_ids = channel["name"].split("__")
|
||||
other_user_id = user_ids[1] if user_ids[0] == my_user_id else user_ids[0]
|
||||
if other_user_id not in user_id_to_name:
|
||||
try:
|
||||
user_id_to_name[other_user_id] = d.users.get_user(other_user_id)["username"]
|
||||
except Exception:
|
||||
user_id_to_name[other_user_id] = channel.get("display_name") or other_user_id
|
||||
channel["display_name"] = user_id_to_name[other_user_id]
|
||||
# Sort channels by name for easier search
|
||||
channels = sorted(channels, key=lambda x: x["display_name"].lower())
|
||||
@@ -205,29 +211,73 @@ def select_channel(d: Driver, team: str, my_user_id: str, user_id_to_name: Dict[
|
||||
return selected_channels
|
||||
|
||||
|
||||
def _resolve_team_name(d: Driver, channel: dict, fallback_team: Optional[dict] = None) -> str:
|
||||
team_id = channel.get("team_id") or ""
|
||||
if team_id:
|
||||
try:
|
||||
return d.teams.get_team(team_id)["name"]
|
||||
except ResourceNotFound:
|
||||
pass
|
||||
if fallback_team:
|
||||
return fallback_team["name"]
|
||||
return ""
|
||||
|
||||
|
||||
def _parse_filter_timestamp(s: str) -> float:
|
||||
for fmt in ('%Y-%m-%dT%H:%M:%SZ', '%Y-%m-%d'):
|
||||
try:
|
||||
dt = datetime.strptime(s, fmt)
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt.timestamp()
|
||||
except ValueError:
|
||||
continue
|
||||
raise ValueError(f"Invalid date format: {s!r} (expected YYYY-MM-DD or YYYY-MM-DDTHH:MM:SSZ)")
|
||||
|
||||
|
||||
def save_config(config: dict, config_filename: str = "config.json") -> None:
|
||||
to_save = {k: v for k, v in config.items() if k != "password"}
|
||||
with open(config_filename, "w") as f:
|
||||
json.dump(to_save, f, indent=2)
|
||||
|
||||
|
||||
def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], output_base: str,
|
||||
download_files: bool = True, before: str = None, after: str = None):
|
||||
download_files: bool = True, before: str = None, after: str = None,
|
||||
team: Optional[dict] = None, merge: bool = False):
|
||||
# Sanitize channel name
|
||||
channel_name = channel["display_name"].replace("\\", "").replace("/", "")
|
||||
|
||||
print("Exporting channel", channel_name)
|
||||
after_ts = _parse_filter_timestamp(after) if after else None
|
||||
before_ts = _parse_filter_timestamp(before) if before else None
|
||||
if after:
|
||||
after = datetime.strptime(after, '%Y-%m-%d').timestamp()
|
||||
if before:
|
||||
before = datetime.strptime(before, '%Y-%m-%d').timestamp()
|
||||
print(f" fetching posts after {after}")
|
||||
|
||||
# Get all posts for selected channel
|
||||
page = 0
|
||||
all_posts = []
|
||||
while True:
|
||||
print(f"Requesting channel page {page}")
|
||||
posts = d.posts.get_posts_for_channel(channel["id"], params={"per_page": 200, "page": page})
|
||||
params = {"per_page": 200, "page": page}
|
||||
if after_ts is not None:
|
||||
# Ask server only for posts changed since last export.
|
||||
# For incremental updates we typically expect few new posts,
|
||||
# so a single page is enough and avoids walking the whole history.
|
||||
params["since"] = int(after_ts * 1000)
|
||||
posts = d.posts.get_posts_for_channel(channel["id"], params=params)
|
||||
|
||||
if len(posts["posts"]) == 0:
|
||||
# If no posts are returned, we have reached the end
|
||||
break
|
||||
|
||||
all_posts.extend([posts["posts"][post] for post in posts["order"]])
|
||||
|
||||
# In incremental mode (after_ts is set) we intentionally stop
|
||||
# after the first non-empty page to avoid downloading the
|
||||
# entire channel history again. For full dumps, keep paging.
|
||||
if after_ts is not None:
|
||||
break
|
||||
|
||||
page += 1
|
||||
print(f"Found {len(all_posts)} posts")
|
||||
# Create output directory
|
||||
@@ -239,14 +289,16 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out
|
||||
|
||||
# Filter posts by date range
|
||||
created = post["create_at"] / 1000
|
||||
if (before and created > before) or (after and created < after):
|
||||
if (before_ts and created > before_ts) or (after_ts and created <= after_ts):
|
||||
continue
|
||||
|
||||
user_id = post["user_id"]
|
||||
if user_id not in user_id_to_name:
|
||||
user_id_to_name[user_id] = d.users.get_user(user_id)["username"]
|
||||
username = user_id_to_name[user_id]
|
||||
created = datetime.utcfromtimestamp(post["create_at"] / 1000).strftime('%Y-%m-%dT%H:%M:%SZ')
|
||||
created = datetime.fromtimestamp(
|
||||
post["create_at"] / 1000, timezone.utc
|
||||
).strftime('%Y-%m-%dT%H:%M:%SZ')
|
||||
message = post["message"]
|
||||
simple_post = dict(idx=i_post, id=post["id"], created=created, username=username, message=message)
|
||||
|
||||
@@ -270,12 +322,12 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out
|
||||
if download_files:
|
||||
filename = "%03d" % i_post + "_" + file["name"]
|
||||
print("Downloading", file["name"])
|
||||
while True:
|
||||
try:
|
||||
resp = d.files.get_file(file["id"])
|
||||
break
|
||||
except:
|
||||
print("Downloading file failed")
|
||||
except ResourceNotFound:
|
||||
print(f"File not found, skipping: {file['name']}")
|
||||
filenames.append(file["name"])
|
||||
continue
|
||||
# Mattermost Driver unfortunately parses json files to dicts
|
||||
if isinstance(resp, dict):
|
||||
with open(output_base / filename, "w") as f:
|
||||
@@ -288,23 +340,42 @@ def export_channel(d: Driver, channel: str, user_id_to_name: Dict[str, str], out
|
||||
simple_post["files"] = filenames
|
||||
simple_posts.append(simple_post)
|
||||
|
||||
filtered_channel_name = ''.join(filter(lambda ch: ch not in "?!/\\.;:*\"<>|", channel_name))
|
||||
output_filename = filtered_channel_name + ".json"
|
||||
output_filepath = output_base / output_filename
|
||||
|
||||
merged_posts = simple_posts
|
||||
if merge and output_filepath.exists():
|
||||
with open(output_filepath, encoding='utf8') as f:
|
||||
existing = json.load(f)
|
||||
seen_ids = {post["id"] for post in existing.get("posts", [])}
|
||||
merged_posts = list(existing.get("posts", []))
|
||||
new_count = 0
|
||||
for post in simple_posts:
|
||||
if post["id"] not in seen_ids:
|
||||
merged_posts.append(post)
|
||||
seen_ids.add(post["id"])
|
||||
new_count += 1
|
||||
merged_posts.sort(key=lambda p: p["created"])
|
||||
for i, post in enumerate(merged_posts):
|
||||
post["idx"] = i
|
||||
print(f" merged {new_count} new posts ({len(merged_posts)} total)")
|
||||
elif merge:
|
||||
print(f" no existing export at '{output_filepath}', writing fresh")
|
||||
|
||||
output = {
|
||||
"channel": {
|
||||
"name": channel["name"],
|
||||
"display_name": channel["display_name"],
|
||||
"header": channel["header"],
|
||||
"id": channel["id"],
|
||||
"team": d.teams.get_team(channel["team_id"])["name"],
|
||||
"team_id": channel["team_id"],
|
||||
"exported_at": datetime.now().strftime('%Y-%m-%dT%H:%M:%SZ')
|
||||
"team": _resolve_team_name(d, channel, team),
|
||||
"team_id": channel.get("team_id") or "",
|
||||
"exported_at": datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
|
||||
},
|
||||
"posts": simple_posts
|
||||
"posts": merged_posts
|
||||
}
|
||||
|
||||
# Export posts to json file
|
||||
filtered_channel_name = ''.join(filter(lambda ch: ch not in "?!/\\.;:*\"<>|", channel_name))
|
||||
output_filename = filtered_channel_name + ".json"
|
||||
output_filepath = output_base / output_filename
|
||||
with open(output_filepath, "w", encoding='utf8') as f:
|
||||
json.dump(output, f, indent=2, ensure_ascii=False)
|
||||
print(f"Exported channel data to '{output_filepath}'")
|
||||
@@ -350,6 +421,11 @@ def complete_config(config: dict, config_filename: str = "config.json") -> dict:
|
||||
|
||||
password = getpass.getpass("Enter your password (hidden): ")
|
||||
elif config["login_mode"] == "gitlab":
|
||||
# If we already have a token in config, reuse it and do not ask
|
||||
# for GitLab credentials again.
|
||||
if config.get("token"):
|
||||
print("Using existing token from config for GitLab login_mode.")
|
||||
else:
|
||||
if config.get("gitlab_url", False):
|
||||
print(f"Using GitLab URL '{config['gitlab_url']}' from config")
|
||||
else:
|
||||
@@ -465,23 +541,55 @@ def find_mmauthtoken_firefox(host):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
config = get_config_from_json()
|
||||
config = complete_config(config)
|
||||
config_filename = "config.json"
|
||||
config = get_config_from_json(config_filename)
|
||||
config = complete_config(config, config_filename)
|
||||
|
||||
incremental = config.get("incremental", bool(config.get("last_export_at")))
|
||||
if incremental and config.get("last_export_at"):
|
||||
print(f"Incremental update from {config['last_export_at']}")
|
||||
output_base = config.get("output_base") or ("results/" + date.today().strftime("%Y%m%d"))
|
||||
after = config["last_export_at"]
|
||||
before = config.get("before")
|
||||
merge = True
|
||||
else:
|
||||
output_base = "results/" + date.today().strftime("%Y%m%d")
|
||||
after = config.get("after")
|
||||
before = config.get("before")
|
||||
merge = False
|
||||
|
||||
print(f"Storing downloaded data in {output_base}")
|
||||
export_started_at = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
|
||||
|
||||
# Range of posts to be exported as string in format "YYYY-MM-DD". Use None if no filter should be applied
|
||||
after = config.get("after", None)
|
||||
before = config.get("before", None)
|
||||
|
||||
d = connect(config["host"], config.get("token", None),
|
||||
config.get("username", None), config.get("password", None))
|
||||
try:
|
||||
d = connect(
|
||||
config["host"],
|
||||
config.get("token", None),
|
||||
config.get("username", None),
|
||||
config.get("password", None),
|
||||
)
|
||||
except Exception as e:
|
||||
# If token-based authentication failed, drop the token and ask
|
||||
# the user to rerun the script so that fresh credentials can be
|
||||
# obtained or entered.
|
||||
if config.get("token"):
|
||||
print("Token-based authentication failed, removing stored token.")
|
||||
config.pop("token", None)
|
||||
save_config(config, config_filename)
|
||||
print("Please restart the script to reconfigure authentication.")
|
||||
raise SystemExit(1) from e
|
||||
raise
|
||||
user_id_to_name, my_user_id = get_users(d)
|
||||
team = select_team(d, my_user_id)
|
||||
channels = select_channel(d, team, my_user_id, user_id_to_name)
|
||||
for i_channel, channel in enumerate(channels):
|
||||
print(f"Start export of channel {i_channel + 1}/{len(channels)}")
|
||||
export_channel(d, channel, user_id_to_name, output_base, config["download_files"],
|
||||
before, after)
|
||||
before, after, team=team, merge=merge)
|
||||
|
||||
config["last_export_at"] = export_started_at
|
||||
config["output_base"] = output_base
|
||||
config["incremental"] = True
|
||||
save_config(config, config_filename)
|
||||
print(f"Updated last_export_at to {export_started_at}")
|
||||
print("Finished export")
|
||||
|
||||
+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())
|
||||
|
||||
+481
@@ -0,0 +1,481 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
|
||||
def iter_result_files(root: str) -> Iterable[str]:
|
||||
for dirpath, _dirnames, filenames in os.walk(root):
|
||||
for name in filenames:
|
||||
if not name.endswith(".json"):
|
||||
continue
|
||||
yield os.path.join(dirpath, name)
|
||||
|
||||
|
||||
def parse_iso_date(s: str) -> dt.datetime:
|
||||
# Accept YYYY-MM-DD or full ISO with time; normalize to UTC
|
||||
try:
|
||||
if "T" in s:
|
||||
return dt.datetime.fromisoformat(s.replace("Z", "+00:00")).astimezone(dt.timezone.utc)
|
||||
return dt.datetime.strptime(s, "%Y-%m-%d").replace(tzinfo=dt.timezone.utc)
|
||||
except ValueError as e:
|
||||
raise argparse.ArgumentTypeError(f"invalid date '{s}': {e}") from e
|
||||
|
||||
|
||||
def parse_post_created(s: str) -> dt.datetime:
|
||||
# Exporter uses '%Y-%m-%dT%H:%M:%SZ'
|
||||
return dt.datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=dt.timezone.utc)
|
||||
|
||||
|
||||
def make_matcher(pattern: str, use_regex: bool, ignore_case: bool):
|
||||
if use_regex:
|
||||
flags = re.IGNORECASE if ignore_case else 0
|
||||
rx = re.compile(pattern, flags)
|
||||
|
||||
def match(text: str) -> bool:
|
||||
return bool(rx.search(text))
|
||||
|
||||
return match
|
||||
|
||||
needle = pattern.lower() if ignore_case else pattern
|
||||
|
||||
def match(text: str) -> bool:
|
||||
hay = text.lower() if ignore_case else text
|
||||
return needle in hay
|
||||
|
||||
return match
|
||||
|
||||
|
||||
def build_fts_query(query: str) -> Optional[str]:
|
||||
# Tokenize user query and run prefix search by tokens.
|
||||
tokens = re.findall(r"\w+", query, flags=re.UNICODE)
|
||||
if not tokens:
|
||||
return None
|
||||
return " AND ".join(f"{tok}*" for tok in tokens)
|
||||
|
||||
|
||||
def search_index(
|
||||
db_path: str,
|
||||
query: str,
|
||||
user: Optional[str],
|
||||
channel: Optional[str],
|
||||
team: Optional[str],
|
||||
after: Optional[dt.datetime],
|
||||
before: Optional[dt.datetime],
|
||||
limit: Optional[int],
|
||||
ignore_case: bool,
|
||||
) -> Iterable[Dict[str, Any]]:
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
filters = []
|
||||
params: List[Any] = []
|
||||
|
||||
if user:
|
||||
filters.append("p.username = ?")
|
||||
params.append(user)
|
||||
if channel:
|
||||
filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
|
||||
params.extend([channel, channel])
|
||||
if team:
|
||||
filters.append("p.team = ?")
|
||||
params.append(team)
|
||||
if after:
|
||||
filters.append("p.created >= ?")
|
||||
params.append(after.isoformat())
|
||||
if before:
|
||||
filters.append("p.created <= ?")
|
||||
params.append(before.isoformat())
|
||||
|
||||
where_filters = f"AND {' AND '.join(filters)}" if filters else ""
|
||||
fts_q = build_fts_query(query)
|
||||
if fts_q:
|
||||
sql = f"""
|
||||
SELECT
|
||||
p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
|
||||
p.channel_display_name, p.channel_name, p.team, p.source_path
|
||||
FROM posts p
|
||||
JOIN posts_fts fts ON fts.post_id = p.post_id
|
||||
WHERE fts.message MATCH ?
|
||||
{where_filters}
|
||||
ORDER BY p.created ASC
|
||||
"""
|
||||
params = [fts_q] + params
|
||||
else:
|
||||
cmp_col = "LOWER(p.message)" if ignore_case else "p.message"
|
||||
needle = query.lower() if ignore_case else query
|
||||
sql = f"""
|
||||
SELECT
|
||||
p.post_id, p.idx, p.created, p.username, p.message, p.files_json,
|
||||
p.channel_display_name, p.channel_name, p.team, p.source_path
|
||||
FROM posts p
|
||||
WHERE {cmp_col} LIKE ?
|
||||
{where_filters}
|
||||
ORDER BY p.created ASC
|
||||
"""
|
||||
params = [f"%{needle}%"] + params
|
||||
|
||||
if limit is not None:
|
||||
sql += " LIMIT ?"
|
||||
params.append(limit)
|
||||
|
||||
rows = conn.execute(sql, params).fetchall()
|
||||
for row in rows:
|
||||
files: Optional[List[str]] = None
|
||||
if row["files_json"]:
|
||||
try:
|
||||
maybe_files = json.loads(row["files_json"])
|
||||
if isinstance(maybe_files, list):
|
||||
files = maybe_files
|
||||
except json.JSONDecodeError:
|
||||
files = None
|
||||
post: Dict[str, Any] = {
|
||||
"id": row["post_id"],
|
||||
"idx": row["idx"],
|
||||
"created": row["created"].replace("+00:00", "Z"),
|
||||
"username": row["username"],
|
||||
"message": row["message"],
|
||||
}
|
||||
if files:
|
||||
post["files"] = files
|
||||
|
||||
yield {
|
||||
"file": row["source_path"],
|
||||
"channel": row["channel_display_name"] or row["channel_name"],
|
||||
"team": row["team"],
|
||||
"post": post,
|
||||
}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def search_attachments_index(
|
||||
db_path: str,
|
||||
query: str,
|
||||
user: Optional[str],
|
||||
channel: Optional[str],
|
||||
team: Optional[str],
|
||||
after: Optional[dt.datetime],
|
||||
before: Optional[dt.datetime],
|
||||
limit: Optional[int],
|
||||
mode: str,
|
||||
) -> Iterable[Dict[str, Any]]:
|
||||
"""
|
||||
Search in attachments index.
|
||||
|
||||
mode: "name" to search by file name, "content" to search by text content.
|
||||
"""
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
filters = []
|
||||
params: List[Any] = []
|
||||
|
||||
if user:
|
||||
filters.append("p.username = ?")
|
||||
params.append(user)
|
||||
if channel:
|
||||
filters.append("(p.channel_display_name = ? OR p.channel_name = ?)")
|
||||
params.extend([channel, channel])
|
||||
if team:
|
||||
filters.append("p.team = ?")
|
||||
params.append(team)
|
||||
if after:
|
||||
filters.append("p.created >= ?")
|
||||
params.append(after.isoformat())
|
||||
if before:
|
||||
filters.append("p.created <= ?")
|
||||
params.append(before.isoformat())
|
||||
|
||||
where_filters = f"AND {' AND '.join(filters)}" if filters else ""
|
||||
|
||||
fts_column = "filename" if mode == "name" else "content"
|
||||
fts_q = build_fts_query(query)
|
||||
if not fts_q:
|
||||
return []
|
||||
|
||||
sql = f"""
|
||||
SELECT
|
||||
p.post_id,
|
||||
p.idx,
|
||||
p.created,
|
||||
p.username,
|
||||
p.message,
|
||||
p.files_json,
|
||||
p.channel_display_name,
|
||||
p.channel_name,
|
||||
p.team,
|
||||
p.source_path,
|
||||
a.filename AS attachment_filename,
|
||||
a.filepath AS attachment_path,
|
||||
a.is_text AS attachment_is_text
|
||||
FROM attachments a
|
||||
JOIN attachments_fts aft ON aft.attachment_id = a.id
|
||||
JOIN posts p ON p.post_id = a.post_id
|
||||
WHERE aft.{fts_column} MATCH ?
|
||||
{where_filters}
|
||||
ORDER BY p.created ASC
|
||||
"""
|
||||
|
||||
params = [fts_q] + params
|
||||
if limit is not None:
|
||||
sql += " LIMIT ?"
|
||||
params.append(limit)
|
||||
|
||||
rows = conn.execute(sql, params).fetchall()
|
||||
for row in rows:
|
||||
files: Optional[List[str]] = None
|
||||
if row["files_json"]:
|
||||
try:
|
||||
maybe_files = json.loads(row["files_json"])
|
||||
if isinstance(maybe_files, list):
|
||||
files = maybe_files
|
||||
except json.JSONDecodeError:
|
||||
files = None
|
||||
post: Dict[str, Any] = {
|
||||
"id": row["post_id"],
|
||||
"idx": row["idx"],
|
||||
"created": row["created"].replace("+00:00", "Z"),
|
||||
"username": row["username"],
|
||||
"message": row["message"],
|
||||
}
|
||||
if files:
|
||||
post["files"] = files
|
||||
|
||||
attachment = {
|
||||
"filename": row["attachment_filename"],
|
||||
"path": row["attachment_path"],
|
||||
"is_text": bool(row["attachment_is_text"]),
|
||||
}
|
||||
|
||||
yield {
|
||||
"file": row["source_path"],
|
||||
"channel": row["channel_display_name"] or row["channel_name"],
|
||||
"team": row["team"],
|
||||
"post": post,
|
||||
"attachment": attachment,
|
||||
}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def search_file(
|
||||
path: str,
|
||||
matcher,
|
||||
user: Optional[str],
|
||||
channel: Optional[str],
|
||||
team: Optional[str],
|
||||
after: Optional[dt.datetime],
|
||||
before: Optional[dt.datetime],
|
||||
) -> Iterable[Dict[str, Any]]:
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return []
|
||||
|
||||
ch = data.get("channel", {}) or {}
|
||||
posts: List[Dict[str, Any]] = data.get("posts") or []
|
||||
|
||||
channel_name = ch.get("display_name") or ch.get("name") or ""
|
||||
team_name = ch.get("team") or ""
|
||||
|
||||
if channel and channel not in (channel_name, ch.get("name", "")):
|
||||
return []
|
||||
if team and team != team_name:
|
||||
return []
|
||||
|
||||
for p in posts:
|
||||
msg = p.get("message") or ""
|
||||
username = p.get("username") or ""
|
||||
created_raw = p.get("created")
|
||||
if not isinstance(created_raw, str):
|
||||
continue
|
||||
try:
|
||||
created = parse_post_created(created_raw)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if user and username != user:
|
||||
continue
|
||||
if after and created < after:
|
||||
continue
|
||||
if before and created > before:
|
||||
continue
|
||||
if not matcher(msg):
|
||||
continue
|
||||
|
||||
yield {
|
||||
"file": path,
|
||||
"channel": channel_name,
|
||||
"team": team_name,
|
||||
"post": p,
|
||||
}
|
||||
|
||||
|
||||
def print_human(matches: Iterable[Dict[str, Any]], limit: Optional[int], show_file: bool) -> None:
|
||||
count = 0
|
||||
for m in matches:
|
||||
if limit is not None and count >= limit:
|
||||
break
|
||||
p = m["post"]
|
||||
created = p.get("created", "")
|
||||
username = p.get("username", "")
|
||||
msg = (p.get("message") or "").replace("\n", "\\n")
|
||||
attachment = m.get("attachment")
|
||||
attachment_suffix = ""
|
||||
if isinstance(attachment, dict) and attachment.get("filename"):
|
||||
attachment_suffix = f" [file: {attachment['filename']}]"
|
||||
if show_file:
|
||||
file_rel = os.path.relpath(m["file"])
|
||||
idx = p.get("idx", "?")
|
||||
prefix = f"{file_rel}:{idx}: "
|
||||
else:
|
||||
prefix = ""
|
||||
channel = m.get("channel") or "?"
|
||||
team = m.get("team") or ""
|
||||
team_prefix = f"[{team}]" if team else ""
|
||||
print(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
|
||||
count += 1
|
||||
|
||||
|
||||
def print_json(matches: Iterable[Dict[str, Any]], limit: Optional[int]) -> None:
|
||||
out: List[Dict[str, Any]] = []
|
||||
count = 0
|
||||
for m in matches:
|
||||
if limit is not None and count >= limit:
|
||||
break
|
||||
out.append(m)
|
||||
count += 1
|
||||
json.dump(out, fp=os.sys.stdout, ensure_ascii=False, indent=2)
|
||||
print()
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Search through Mattermost export JSON files produced by mattermost-dl.py.",
|
||||
)
|
||||
ap.add_argument("query", help="Search string or regex (see --regex).")
|
||||
ap.add_argument(
|
||||
"-r",
|
||||
"--root",
|
||||
default="results",
|
||||
help="Root directory with export results (default: results).",
|
||||
)
|
||||
ap.add_argument("--user", help="Filter by username (exact match).")
|
||||
ap.add_argument("--channel", help="Filter by channel display_name or name (exact match).")
|
||||
ap.add_argument("--team", help="Filter by team name (exact match).")
|
||||
ap.add_argument("--after", type=parse_iso_date, help="Only messages on/after this date (YYYY-MM-DD or ISO).")
|
||||
ap.add_argument("--before", type=parse_iso_date, help="Only messages on/before this date (YYYY-MM-DD or ISO).")
|
||||
ap.add_argument("--regex", action="store_true", help="Treat query as regular expression.")
|
||||
group = ap.add_mutually_exclusive_group()
|
||||
group.add_argument(
|
||||
"--file-name",
|
||||
action="store_true",
|
||||
help="Search by attachment file name instead of message text (requires sqlite index).",
|
||||
)
|
||||
group.add_argument(
|
||||
"--file-content",
|
||||
action="store_true",
|
||||
help="Search by attachment text content instead of message text (requires sqlite index; only UTF-8 text files are indexed).",
|
||||
)
|
||||
ap.add_argument(
|
||||
"-i",
|
||||
"--ignore-case",
|
||||
action="store_true",
|
||||
help="Case-insensitive search (default for substring search).",
|
||||
)
|
||||
ap.add_argument("--json", action="store_true", help="Output matches as JSON array (for jq, etc.).")
|
||||
ap.add_argument(
|
||||
"-n",
|
||||
"--limit",
|
||||
type=int,
|
||||
help="Maximum number of matches to output.",
|
||||
)
|
||||
ap.add_argument("--db", default="mattermost_index.sqlite", help="SQLite index database path.")
|
||||
ap.add_argument(
|
||||
"--no-index",
|
||||
action="store_true",
|
||||
help="Do not use sqlite index even if DB exists.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--show-file",
|
||||
action="store_true",
|
||||
help="Include source filename and per-file index in human output.",
|
||||
)
|
||||
|
||||
args = ap.parse_args(argv)
|
||||
|
||||
# Default ignore-case for non-regex search unless explicitly disabled
|
||||
ignore_case = args.ignore_case or not args.regex
|
||||
matcher = make_matcher(args.query, args.regex, ignore_case)
|
||||
|
||||
if (args.file_name or args.file_content) and args.regex:
|
||||
raise SystemExit("--file-name/--file-content do not support --regex yet.")
|
||||
|
||||
use_index = (not args.no_index) and os.path.exists(args.db)
|
||||
if args.file_name or args.file_content:
|
||||
if not use_index:
|
||||
raise SystemExit(
|
||||
"Attachment search (--file-name/--file-content) requires sqlite index. "
|
||||
"Create it with mm_index.py and do not use --no-index."
|
||||
)
|
||||
mode = "name" if args.file_name else "content"
|
||||
matches_iter = search_attachments_index(
|
||||
db_path=args.db,
|
||||
query=args.query,
|
||||
user=args.user,
|
||||
channel=args.channel,
|
||||
team=args.team,
|
||||
after=args.after,
|
||||
before=args.before,
|
||||
limit=args.limit,
|
||||
mode=mode,
|
||||
)
|
||||
output_limit = None
|
||||
elif use_index and not args.regex:
|
||||
matches_iter = search_index(
|
||||
db_path=args.db,
|
||||
query=args.query,
|
||||
user=args.user,
|
||||
channel=args.channel,
|
||||
team=args.team,
|
||||
after=args.after,
|
||||
before=args.before,
|
||||
limit=args.limit,
|
||||
ignore_case=ignore_case,
|
||||
)
|
||||
# Limit already applied in SQL.
|
||||
output_limit = None
|
||||
else:
|
||||
files = list(iter_result_files(args.root))
|
||||
matches_iter = (
|
||||
m
|
||||
for path in files
|
||||
for m in search_file(
|
||||
path=path,
|
||||
matcher=matcher,
|
||||
user=args.user,
|
||||
channel=args.channel,
|
||||
team=args.team,
|
||||
after=args.after,
|
||||
before=args.before,
|
||||
)
|
||||
)
|
||||
output_limit = args.limit
|
||||
|
||||
if args.json:
|
||||
print_json(matches_iter, output_limit)
|
||||
else:
|
||||
print_human(matches_iter, output_limit, show_file=args.show_file)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user