massive improve
This commit is contained in:
Сергей Маринкевич
+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