feat: add txt path import and prune command
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
"""Telegram video damage detector package."""
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
HARD_ERROR_KEYWORDS = (
|
||||
"moov atom not found",
|
||||
"error reading header",
|
||||
"ebml header parsing failed",
|
||||
"invalid data found when processing input",
|
||||
)
|
||||
|
||||
SOFT_WARNING_KEYWORDS = (
|
||||
"partial file",
|
||||
"estimating duration from bitrate",
|
||||
"missing key frame",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Stage1Classification:
|
||||
status: str
|
||||
reason_code: str
|
||||
warning_flags: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def classify_stage1_result(
|
||||
*,
|
||||
exit_code: int | None,
|
||||
stderr_text: str,
|
||||
probe_payload: dict[str, object] | None,
|
||||
extension: str,
|
||||
high_risk_extensions: tuple[str, ...],
|
||||
) -> Stage1Classification:
|
||||
normalized_stderr = stderr_text.lower()
|
||||
|
||||
if exit_code not in (0, None):
|
||||
return Stage1Classification(
|
||||
status="confirmed_broken",
|
||||
reason_code="ffprobe_exit_nonzero",
|
||||
)
|
||||
|
||||
for keyword in HARD_ERROR_KEYWORDS:
|
||||
if keyword in normalized_stderr:
|
||||
return Stage1Classification(
|
||||
status="confirmed_broken",
|
||||
reason_code="hard_error_keyword",
|
||||
warning_flags=(keyword,),
|
||||
)
|
||||
|
||||
format_info = _as_dict(probe_payload.get("format") if probe_payload else None)
|
||||
streams = probe_payload.get("streams") if probe_payload else None
|
||||
if not format_info.get("format_name"):
|
||||
return Stage1Classification(
|
||||
status="confirmed_broken",
|
||||
reason_code="missing_format_name",
|
||||
)
|
||||
|
||||
if not _has_video_stream(streams):
|
||||
return Stage1Classification(
|
||||
status="confirmed_broken",
|
||||
reason_code="missing_video_stream",
|
||||
)
|
||||
|
||||
soft_hits = tuple(
|
||||
keyword for keyword in SOFT_WARNING_KEYWORDS if keyword in normalized_stderr
|
||||
)
|
||||
if soft_hits:
|
||||
return Stage1Classification(
|
||||
status="pending_stage2",
|
||||
reason_code="soft_warning_keyword",
|
||||
warning_flags=soft_hits,
|
||||
)
|
||||
|
||||
if not _is_valid_duration(format_info.get("duration")):
|
||||
return Stage1Classification(
|
||||
status="pending_stage2",
|
||||
reason_code="invalid_duration",
|
||||
)
|
||||
|
||||
if extension in high_risk_extensions:
|
||||
return Stage1Classification(
|
||||
status="pending_stage2",
|
||||
reason_code="high_risk_extension",
|
||||
)
|
||||
|
||||
return Stage1Classification(
|
||||
status="passed_stage1",
|
||||
reason_code="passed_stage1",
|
||||
)
|
||||
|
||||
|
||||
def _as_dict(value: object) -> dict[str, object]:
|
||||
if isinstance(value, dict):
|
||||
return value
|
||||
return {}
|
||||
|
||||
|
||||
def _has_video_stream(streams: object) -> bool:
|
||||
if not isinstance(streams, list):
|
||||
return False
|
||||
|
||||
for stream in streams:
|
||||
if isinstance(stream, dict) and stream.get("codec_type") == "video":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_valid_duration(value: object) -> bool:
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
if value in {"N/A", ""}:
|
||||
return False
|
||||
|
||||
try:
|
||||
return float(value) > 0
|
||||
except ValueError:
|
||||
return False
|
||||
@@ -0,0 +1,504 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
import orjson
|
||||
import typer
|
||||
|
||||
from tgvideo_check.classifiers import classify_stage1_result
|
||||
from tgvideo_check.config import ScanConfig, normalize_extension
|
||||
from tgvideo_check.discovery import discover_files
|
||||
from tgvideo_check.ffmpeg_runner import generate_sample_offsets, run_ffmpeg_verification
|
||||
from tgvideo_check.ffprobe_runner import run_ffprobe
|
||||
from tgvideo_check.profiles import resolve_config
|
||||
from tgvideo_check.reporting import ReportRow, export_csv_summary, export_jsonl
|
||||
from tgvideo_check.scheduler import (
|
||||
build_stage_plan,
|
||||
execute_stage_batch,
|
||||
select_pending_stage1_files,
|
||||
select_pending_stage2_files,
|
||||
)
|
||||
from tgvideo_check.store import ScanStore
|
||||
|
||||
|
||||
app = typer.Typer(
|
||||
add_completion=False,
|
||||
help="Batch detector for truncated or damaged Telegram-exported videos.",
|
||||
)
|
||||
|
||||
|
||||
@app.callback()
|
||||
def root() -> None:
|
||||
"""TG video damage detector CLI."""
|
||||
|
||||
|
||||
@app.command()
|
||||
def scan(
|
||||
paths: list[str] | None = typer.Argument(
|
||||
None,
|
||||
help="One or more root directories to scan.",
|
||||
),
|
||||
paths_file: Path | None = typer.Option(
|
||||
None,
|
||||
"--paths-file",
|
||||
exists=True,
|
||||
file_okay=True,
|
||||
dir_okay=False,
|
||||
readable=True,
|
||||
help="Text file with one directory path per line.",
|
||||
),
|
||||
profile: str = typer.Option("tg-hdd", "--profile"),
|
||||
recursive: bool = typer.Option(True, "--recursive/--no-recursive"),
|
||||
extensions: str = typer.Option(
|
||||
"mp4,mov,m4v,mkv,webm,avi",
|
||||
"--extensions",
|
||||
help="Comma-separated video extensions.",
|
||||
),
|
||||
db_path: Path = typer.Option(Path("scan_state.db"), "--db"),
|
||||
resume: bool = typer.Option(True, "--resume/--no-resume"),
|
||||
stage1_workers: int | None = typer.Option(None, "--stage1-workers"),
|
||||
stage2_workers: int | None = typer.Option(None, "--stage2-workers"),
|
||||
queue_limit: int = typer.Option(64, "--queue-limit"),
|
||||
sample_duration: int | None = typer.Option(None, "--sample-duration"),
|
||||
full_decode_confirm: bool = typer.Option(False, "--full-decode-confirm"),
|
||||
export_jsonl_path: Path | None = typer.Option(None, "--export-jsonl"),
|
||||
export_csv_path: Path | None = typer.Option(None, "--export-csv"),
|
||||
) -> None:
|
||||
"""Scan directories for suspicious video files."""
|
||||
|
||||
config = resolve_config(
|
||||
ScanConfig(
|
||||
roots=_resolve_scan_roots(paths or [], paths_file),
|
||||
profile=profile,
|
||||
recursive=recursive,
|
||||
extensions=_parse_extensions(extensions),
|
||||
db_path=db_path,
|
||||
resume=resume,
|
||||
stage1_workers=stage1_workers,
|
||||
stage2_workers=stage2_workers,
|
||||
queue_limit=queue_limit,
|
||||
sample_duration=sample_duration,
|
||||
full_decode_confirm=full_decode_confirm,
|
||||
export_jsonl=export_jsonl_path,
|
||||
export_csv=export_csv_path,
|
||||
)
|
||||
)
|
||||
|
||||
rows = _execute_scan(config)
|
||||
summary = Counter(row.final_status for row in rows)
|
||||
summary_text = " ".join(
|
||||
f"{status}={summary[status]}" for status in sorted(summary.keys())
|
||||
)
|
||||
typer.echo(f"Scanned {len(rows)} file(s). {summary_text}".strip())
|
||||
|
||||
|
||||
@app.command()
|
||||
def prune(
|
||||
db_path: Path = typer.Option(Path("scan_state.db"), "--db"),
|
||||
run_id: int | None = typer.Option(None, "--run-id"),
|
||||
delete_status: str = typer.Option(
|
||||
"confirmed_broken",
|
||||
"--delete-status",
|
||||
help="Comma-separated final statuses eligible for deletion.",
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
True,
|
||||
"--dry-run/--no-dry-run",
|
||||
help="Preview deletion targets without removing files.",
|
||||
),
|
||||
) -> None:
|
||||
"""Delete selected files from the latest scan result and print stats."""
|
||||
|
||||
store = ScanStore(db_path)
|
||||
store.initialize()
|
||||
selected_run_id = run_id or store.get_latest_run_id()
|
||||
if selected_run_id is None:
|
||||
raise typer.BadParameter(f"No scan runs found in database: {db_path}")
|
||||
|
||||
statuses = {
|
||||
item.strip()
|
||||
for item in delete_status.split(",")
|
||||
if item.strip()
|
||||
}
|
||||
if not statuses:
|
||||
raise typer.BadParameter("At least one delete status must be provided.")
|
||||
|
||||
files = store.list_files(selected_run_id)
|
||||
matched_rows = _build_report_rows(files, {})
|
||||
matched_rows = [
|
||||
row for row in matched_rows if row.final_status in statuses
|
||||
]
|
||||
matched_size = sum(row.size for row in matched_rows)
|
||||
deleted = 0
|
||||
deleted_size = 0
|
||||
missing = 0
|
||||
failed = 0
|
||||
|
||||
if not dry_run:
|
||||
for row in matched_rows:
|
||||
try:
|
||||
if not row.path.exists():
|
||||
missing += 1
|
||||
continue
|
||||
if not row.path.is_file():
|
||||
failed += 1
|
||||
continue
|
||||
row.path.unlink()
|
||||
deleted += 1
|
||||
deleted_size += row.size
|
||||
except OSError:
|
||||
failed += 1
|
||||
|
||||
summary_label = "Dry-run" if dry_run else "Prune"
|
||||
typer.echo(
|
||||
f"{summary_label}: run_id={selected_run_id} matched={len(matched_rows)} "
|
||||
f"deleted={deleted} missing={missing} failed={failed} bytes={matched_size}"
|
||||
)
|
||||
if not dry_run:
|
||||
typer.echo(f"Deleted bytes={deleted_size}")
|
||||
|
||||
reason_counts = Counter(row.reason_code for row in matched_rows)
|
||||
if reason_counts:
|
||||
reason_text = " ".join(
|
||||
f"{reason}={reason_counts[reason]}" for reason in sorted(reason_counts)
|
||||
)
|
||||
typer.echo(f"Reasons: {reason_text}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
app()
|
||||
|
||||
|
||||
def _execute_scan(config: ScanConfig) -> list[ReportRow]:
|
||||
store = ScanStore(config.db_path)
|
||||
store.initialize()
|
||||
run_id = store.create_run(profile=config.profile, root_paths=config.roots)
|
||||
reasons: dict[int, str] = {}
|
||||
stage1_payloads: dict[int, dict[str, object]] = {}
|
||||
|
||||
try:
|
||||
candidates = discover_files(
|
||||
config.roots,
|
||||
extensions=config.extensions,
|
||||
recursive=config.recursive,
|
||||
)
|
||||
stored_files = store.upsert_files(
|
||||
run_id,
|
||||
candidates,
|
||||
reuse_previous_results=config.resume,
|
||||
)
|
||||
|
||||
stage1_batches = build_stage_plan(
|
||||
stage1_files=select_pending_stage1_files(stored_files),
|
||||
stage2_files=[],
|
||||
stage1_workers=config.stage1_workers or 1,
|
||||
stage2_workers=config.stage2_workers or 1,
|
||||
queue_limit=config.queue_limit,
|
||||
)
|
||||
for batch in stage1_batches:
|
||||
probe_results = execute_stage_batch(
|
||||
batch,
|
||||
lambda item: run_ffprobe(item.path),
|
||||
)
|
||||
for stored_file, probe_result in zip(batch.files, probe_results):
|
||||
store.record_stage1_result(
|
||||
stored_file.file_id,
|
||||
exit_code=probe_result.exit_code,
|
||||
stdout_text=probe_result.stdout_text,
|
||||
stderr_text=probe_result.stderr_text,
|
||||
elapsed_ms=probe_result.elapsed_ms,
|
||||
)
|
||||
|
||||
if probe_result.status != "completed":
|
||||
store.record_stage1_status(
|
||||
stored_file.file_id,
|
||||
stage1_status="command_error",
|
||||
final_status="command_error",
|
||||
reason_code=probe_result.status,
|
||||
)
|
||||
reasons[stored_file.file_id] = probe_result.status
|
||||
continue
|
||||
|
||||
payload = _parse_probe_payload(probe_result.stdout_text)
|
||||
if probe_result.stdout_text and payload is None:
|
||||
store.record_stage1_status(
|
||||
stored_file.file_id,
|
||||
stage1_status="command_error",
|
||||
final_status="command_error",
|
||||
reason_code="invalid_probe_output",
|
||||
)
|
||||
reasons[stored_file.file_id] = "invalid_probe_output"
|
||||
continue
|
||||
|
||||
classification = classify_stage1_result(
|
||||
exit_code=probe_result.exit_code,
|
||||
stderr_text=probe_result.stderr_text,
|
||||
probe_payload=payload,
|
||||
extension=stored_file.extension,
|
||||
high_risk_extensions=config.high_risk_extensions,
|
||||
)
|
||||
final_status = (
|
||||
classification.status
|
||||
if classification.status in {"confirmed_broken", "passed_stage1"}
|
||||
else None
|
||||
)
|
||||
store.record_stage1_status(
|
||||
stored_file.file_id,
|
||||
stage1_status=classification.status,
|
||||
final_status=final_status,
|
||||
reason_code=classification.reason_code,
|
||||
)
|
||||
reasons[stored_file.file_id] = classification.reason_code
|
||||
if payload is not None:
|
||||
stage1_payloads[stored_file.file_id] = payload
|
||||
|
||||
stored_files = store.list_files(run_id)
|
||||
stage2_batches = build_stage_plan(
|
||||
stage1_files=[],
|
||||
stage2_files=select_pending_stage2_files(stored_files),
|
||||
stage1_workers=config.stage1_workers or 1,
|
||||
stage2_workers=config.stage2_workers or 1,
|
||||
queue_limit=config.queue_limit,
|
||||
)
|
||||
for batch in stage2_batches:
|
||||
stage2_outputs = execute_stage_batch(
|
||||
batch,
|
||||
lambda item: _run_stage2_for_file(
|
||||
item.path,
|
||||
payload=stage1_payloads.get(item.file_id)
|
||||
or _parse_probe_payload(store.get_stage1_stdout_text(item.file_id) or "")
|
||||
or _reload_probe_payload(item.path),
|
||||
sample_duration=config.sample_duration or 3,
|
||||
full_decode_confirm=config.full_decode_confirm,
|
||||
),
|
||||
)
|
||||
for stored_file, (stage2_status, reason_code, stage2_result) in zip(
|
||||
batch.files, stage2_outputs
|
||||
):
|
||||
store.record_stage2_result(
|
||||
stored_file.file_id,
|
||||
mode=stage2_result["mode"],
|
||||
sample_points_json=stage2_result["sample_points_json"],
|
||||
exit_code=stage2_result["exit_code"],
|
||||
stderr_text=stage2_result["stderr_text"],
|
||||
reason_code=reason_code,
|
||||
elapsed_ms=stage2_result["elapsed_ms"],
|
||||
)
|
||||
store.record_stage2_status(
|
||||
stored_file.file_id,
|
||||
stage2_status=stage2_status,
|
||||
final_status=stage2_status,
|
||||
reason_code=reason_code,
|
||||
)
|
||||
reasons[stored_file.file_id] = reason_code
|
||||
|
||||
rows = _build_report_rows(store.list_files(run_id), reasons)
|
||||
if config.export_jsonl is not None:
|
||||
export_jsonl(rows, config.export_jsonl)
|
||||
if config.export_csv is not None:
|
||||
export_csv_summary(rows, config.export_csv)
|
||||
|
||||
store.update_run_status(run_id, status="completed")
|
||||
return rows
|
||||
except Exception:
|
||||
store.update_run_status(run_id, status="failed")
|
||||
raise
|
||||
|
||||
|
||||
def _build_report_rows(
|
||||
stored_files: list,
|
||||
reasons: dict[int, str],
|
||||
) -> list[ReportRow]:
|
||||
rows: list[ReportRow] = []
|
||||
for stored_file in stored_files:
|
||||
final_status = stored_file.final_status or stored_file.stage1_status
|
||||
rows.append(
|
||||
ReportRow(
|
||||
path=stored_file.path,
|
||||
final_status=final_status,
|
||||
reason_code=stored_file.reason_code
|
||||
or reasons.get(stored_file.file_id, final_status or "unknown_reason"),
|
||||
size=stored_file.size,
|
||||
)
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def _run_stage2_for_file(
|
||||
path: Path,
|
||||
*,
|
||||
payload: dict[str, object] | None,
|
||||
sample_duration: int,
|
||||
full_decode_confirm: bool,
|
||||
) -> tuple[str, str, dict[str, int | str | None]]:
|
||||
duration = _extract_duration(payload)
|
||||
if full_decode_confirm or duration is None:
|
||||
sample_points: tuple[float, ...] = ()
|
||||
mode = "full_decode"
|
||||
executions = [
|
||||
run_ffmpeg_verification(
|
||||
path,
|
||||
full_decode=True,
|
||||
)
|
||||
]
|
||||
else:
|
||||
sample_points = generate_sample_offsets(duration)
|
||||
mode = "sample"
|
||||
executions = [
|
||||
run_ffmpeg_verification(
|
||||
path,
|
||||
offset_seconds=offset,
|
||||
sample_duration=sample_duration,
|
||||
)
|
||||
for offset in sample_points
|
||||
]
|
||||
|
||||
total_elapsed = sum(item.elapsed_ms for item in executions)
|
||||
stderr_text = "\n".join(item.stderr_text for item in executions if item.stderr_text)
|
||||
|
||||
for execution in executions:
|
||||
if execution.status != "completed":
|
||||
return (
|
||||
"command_error",
|
||||
execution.reason_code or execution.status,
|
||||
{
|
||||
"mode": mode,
|
||||
"sample_points_json": _dump_json(sample_points),
|
||||
"exit_code": execution.exit_code,
|
||||
"stderr_text": stderr_text,
|
||||
"elapsed_ms": total_elapsed,
|
||||
},
|
||||
)
|
||||
if execution.exit_code not in (0, None):
|
||||
return (
|
||||
"confirmed_broken",
|
||||
execution.reason_code or "ffmpeg_exit_nonzero",
|
||||
{
|
||||
"mode": mode,
|
||||
"sample_points_json": _dump_json(sample_points),
|
||||
"exit_code": execution.exit_code,
|
||||
"stderr_text": stderr_text,
|
||||
"elapsed_ms": total_elapsed,
|
||||
},
|
||||
)
|
||||
|
||||
return (
|
||||
"passed_stage2",
|
||||
"passed_stage2",
|
||||
{
|
||||
"mode": mode,
|
||||
"sample_points_json": _dump_json(sample_points),
|
||||
"exit_code": 0,
|
||||
"stderr_text": stderr_text,
|
||||
"elapsed_ms": total_elapsed,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _reload_probe_payload(path: Path) -> dict[str, object] | None:
|
||||
probe_result = run_ffprobe(path)
|
||||
if probe_result.status != "completed":
|
||||
return None
|
||||
return _parse_probe_payload(probe_result.stdout_text)
|
||||
|
||||
|
||||
def _parse_probe_payload(stdout_text: str) -> dict[str, object] | None:
|
||||
if not stdout_text.strip():
|
||||
return None
|
||||
payload = orjson.loads(stdout_text)
|
||||
if isinstance(payload, dict):
|
||||
return payload
|
||||
return None
|
||||
|
||||
|
||||
def _extract_duration(payload: dict[str, object] | None) -> float | None:
|
||||
if payload is None:
|
||||
return None
|
||||
format_info = payload.get("format")
|
||||
if not isinstance(format_info, dict):
|
||||
return None
|
||||
duration_value = format_info.get("duration")
|
||||
if not isinstance(duration_value, str):
|
||||
return None
|
||||
try:
|
||||
duration = float(duration_value)
|
||||
except ValueError:
|
||||
return None
|
||||
return duration if duration > 0 else None
|
||||
|
||||
|
||||
def _parse_extensions(raw_extensions: str) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
normalize_extension(item)
|
||||
for item in raw_extensions.split(",")
|
||||
if item.strip()
|
||||
)
|
||||
|
||||
|
||||
def _dump_json(values: tuple[float, ...]) -> str:
|
||||
return orjson.dumps(values).decode("utf-8")
|
||||
|
||||
|
||||
def _resolve_scan_roots(
|
||||
positional_paths: list[str],
|
||||
paths_file: Path | None,
|
||||
) -> tuple[Path, ...]:
|
||||
roots: list[Path] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for raw_path in positional_paths:
|
||||
_append_scan_root(
|
||||
Path(raw_path),
|
||||
roots=roots,
|
||||
seen=seen,
|
||||
source="command argument",
|
||||
)
|
||||
|
||||
if paths_file is not None:
|
||||
for line_number, raw_line in enumerate(
|
||||
paths_file.read_text(encoding="utf-8-sig").splitlines(),
|
||||
start=1,
|
||||
):
|
||||
candidate = raw_line.rstrip("\r\n")
|
||||
if not candidate.strip():
|
||||
continue
|
||||
_append_scan_root(
|
||||
Path(_strip_optional_wrapping_quotes(candidate)),
|
||||
roots=roots,
|
||||
seen=seen,
|
||||
source=f"{paths_file}:{line_number}",
|
||||
)
|
||||
|
||||
if not roots:
|
||||
raise typer.BadParameter(
|
||||
"Provide at least one directory path or use --paths-file.",
|
||||
)
|
||||
|
||||
return tuple(roots)
|
||||
|
||||
|
||||
def _append_scan_root(
|
||||
path: Path,
|
||||
*,
|
||||
roots: list[Path],
|
||||
seen: set[str],
|
||||
source: str,
|
||||
) -> None:
|
||||
if not path.exists():
|
||||
raise typer.BadParameter(f"Path from {source} does not exist: {path}")
|
||||
if not path.is_dir():
|
||||
raise typer.BadParameter(f"Path from {source} is not a directory: {path}")
|
||||
|
||||
normalized = str(path.resolve())
|
||||
if normalized in seen:
|
||||
return
|
||||
|
||||
seen.add(normalized)
|
||||
roots.append(path)
|
||||
|
||||
|
||||
def _strip_optional_wrapping_quotes(value: str) -> str:
|
||||
if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
|
||||
return value[1:-1]
|
||||
return value
|
||||
@@ -0,0 +1,35 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def normalize_extension(value: str) -> str:
|
||||
return value.lower().lstrip(".")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScanConfig:
|
||||
roots: tuple[Path, ...]
|
||||
profile: str = "tg-hdd"
|
||||
recursive: bool = True
|
||||
extensions: tuple[str, ...] = ("mp4", "mov", "m4v", "mkv", "webm", "avi")
|
||||
db_path: Path = Path("scan_state.db")
|
||||
resume: bool = True
|
||||
stage1_workers: int | None = None
|
||||
stage2_workers: int | None = None
|
||||
queue_limit: int = 64
|
||||
sample_duration: int | None = None
|
||||
full_decode_confirm: bool = False
|
||||
export_jsonl: Path | None = None
|
||||
export_csv: Path | None = None
|
||||
high_risk_extensions: tuple[str, ...] = ()
|
||||
|
||||
def normalized(self) -> "ScanConfig":
|
||||
return replace(
|
||||
self,
|
||||
extensions=tuple(normalize_extension(item) for item in self.extensions),
|
||||
high_risk_extensions=tuple(
|
||||
normalize_extension(item) for item in self.high_risk_extensions
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from tgvideo_check.config import normalize_extension
|
||||
from tgvideo_check.models import FileCandidate
|
||||
|
||||
|
||||
def discover_files(
|
||||
roots: tuple[Path, ...],
|
||||
*,
|
||||
extensions: tuple[str, ...],
|
||||
recursive: bool = True,
|
||||
) -> list[FileCandidate]:
|
||||
normalized_extensions = {normalize_extension(item) for item in extensions}
|
||||
discovered: list[FileCandidate] = []
|
||||
|
||||
for root in roots:
|
||||
_discover_root(
|
||||
root=root,
|
||||
recursive=recursive,
|
||||
extensions=normalized_extensions,
|
||||
discovered=discovered,
|
||||
)
|
||||
|
||||
return discovered
|
||||
|
||||
|
||||
def _discover_root(
|
||||
*,
|
||||
root: Path,
|
||||
recursive: bool,
|
||||
extensions: set[str],
|
||||
discovered: list[FileCandidate],
|
||||
) -> None:
|
||||
entries = sorted(root.iterdir(), key=lambda item: (not item.is_file(), item.name.lower()))
|
||||
|
||||
for entry in entries:
|
||||
if entry.is_dir():
|
||||
if recursive:
|
||||
_discover_root(
|
||||
root=entry,
|
||||
recursive=recursive,
|
||||
extensions=extensions,
|
||||
discovered=discovered,
|
||||
)
|
||||
continue
|
||||
|
||||
extension = normalize_extension(entry.suffix)
|
||||
if extension not in extensions:
|
||||
continue
|
||||
|
||||
stat = entry.stat()
|
||||
discovered.append(
|
||||
FileCandidate(
|
||||
path=entry,
|
||||
size=stat.st_size,
|
||||
mtime_ns=stat.st_mtime_ns,
|
||||
extension=extension,
|
||||
discovery_order=len(discovered),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,148 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ERROR_REASON_PATTERNS: tuple[tuple[str, str], ...] = (
|
||||
("truncating packet", "truncating_packet"),
|
||||
("invalid nal unit size", "invalid_nal_unit_size"),
|
||||
("decode_slice_header error", "decode_slice_header_error"),
|
||||
("corrupt input packet in stream", "corrupt_input_packet"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FfmpegExecutionResult:
|
||||
status: str
|
||||
command: tuple[str, ...]
|
||||
exit_code: int | None
|
||||
stdout_text: str
|
||||
stderr_text: str
|
||||
elapsed_ms: int
|
||||
reason_code: str | None
|
||||
|
||||
|
||||
def generate_sample_offsets(duration_seconds: float) -> tuple[float, ...]:
|
||||
if duration_seconds <= 60:
|
||||
ratios = (0.10, 0.85)
|
||||
elif duration_seconds <= 600:
|
||||
ratios = (0.10, 0.50, 0.90)
|
||||
else:
|
||||
ratios = (0.10, 0.60, 0.92)
|
||||
|
||||
return tuple(round(duration_seconds * ratio, 3) for ratio in ratios)
|
||||
|
||||
|
||||
def build_ffmpeg_command(
|
||||
path: Path,
|
||||
*,
|
||||
offset_seconds: float | None = None,
|
||||
sample_duration: int = 3,
|
||||
full_decode: bool = False,
|
||||
executable: str = "ffmpeg",
|
||||
launcher_args: tuple[str, ...] = (),
|
||||
) -> list[str]:
|
||||
command = [executable, *launcher_args, "-v", "error", "-xerror"]
|
||||
|
||||
if not full_decode:
|
||||
if offset_seconds is None:
|
||||
raise ValueError("offset_seconds is required for sample verification")
|
||||
command.extend(["-ss", _format_time_value(offset_seconds)])
|
||||
|
||||
command.extend(["-i", str(path)])
|
||||
|
||||
if not full_decode:
|
||||
command.extend(["-t", str(sample_duration), "-map", "0:v:0"])
|
||||
|
||||
command.extend(["-f", "null", "-"])
|
||||
return command
|
||||
|
||||
|
||||
def run_ffmpeg_verification(
|
||||
path: Path,
|
||||
*,
|
||||
offset_seconds: float | None = None,
|
||||
sample_duration: int = 3,
|
||||
full_decode: bool = False,
|
||||
executable: str = "ffmpeg",
|
||||
launcher_args: tuple[str, ...] = (),
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> FfmpegExecutionResult:
|
||||
command = tuple(
|
||||
build_ffmpeg_command(
|
||||
path,
|
||||
offset_seconds=offset_seconds,
|
||||
sample_duration=sample_duration,
|
||||
full_decode=full_decode,
|
||||
executable=executable,
|
||||
launcher_args=launcher_args,
|
||||
)
|
||||
)
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
command,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
timeout=timeout_seconds,
|
||||
check=False,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
return FfmpegExecutionResult(
|
||||
status="missing_command",
|
||||
command=command,
|
||||
exit_code=None,
|
||||
stdout_text="",
|
||||
stderr_text=str(exc),
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
reason_code="missing_command",
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
return FfmpegExecutionResult(
|
||||
status="timeout",
|
||||
command=command,
|
||||
exit_code=None,
|
||||
stdout_text=_coerce_process_output(exc.stdout),
|
||||
stderr_text=_coerce_process_output(exc.stderr),
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
reason_code="timeout",
|
||||
)
|
||||
|
||||
return FfmpegExecutionResult(
|
||||
status="completed",
|
||||
command=command,
|
||||
exit_code=completed.returncode,
|
||||
stdout_text=completed.stdout,
|
||||
stderr_text=completed.stderr,
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
reason_code=classify_ffmpeg_error(completed.stderr),
|
||||
)
|
||||
|
||||
|
||||
def classify_ffmpeg_error(stderr_text: str) -> str | None:
|
||||
normalized = stderr_text.lower()
|
||||
for pattern, reason_code in ERROR_REASON_PATTERNS:
|
||||
if pattern in normalized:
|
||||
return reason_code
|
||||
return None
|
||||
|
||||
|
||||
def _elapsed_ms(start: float) -> int:
|
||||
return int((time.perf_counter() - start) * 1000)
|
||||
|
||||
|
||||
def _coerce_process_output(output: str | bytes | None) -> str:
|
||||
if output is None:
|
||||
return ""
|
||||
if isinstance(output, bytes):
|
||||
return output.decode("utf-8", errors="replace")
|
||||
return output
|
||||
|
||||
|
||||
def _format_time_value(value: float) -> str:
|
||||
return f"{value:.3f}".rstrip("0").rstrip(".")
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
MINIMAL_SHOW_ENTRIES = (
|
||||
"format=format_name,duration,size:"
|
||||
"stream=index,codec_type,codec_name,width,height"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProbeExecutionResult:
|
||||
status: str
|
||||
command: tuple[str, ...]
|
||||
exit_code: int | None
|
||||
stdout_text: str
|
||||
stderr_text: str
|
||||
elapsed_ms: int
|
||||
|
||||
|
||||
def build_ffprobe_command(
|
||||
path: Path,
|
||||
*,
|
||||
executable: str = "ffprobe",
|
||||
launcher_args: tuple[str, ...] = (),
|
||||
) -> list[str]:
|
||||
return [
|
||||
executable,
|
||||
*launcher_args,
|
||||
"-v",
|
||||
"error",
|
||||
"-show_entries",
|
||||
MINIMAL_SHOW_ENTRIES,
|
||||
"-of",
|
||||
"json",
|
||||
str(path),
|
||||
]
|
||||
|
||||
|
||||
def run_ffprobe(
|
||||
path: Path,
|
||||
*,
|
||||
executable: str = "ffprobe",
|
||||
launcher_args: tuple[str, ...] = (),
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> ProbeExecutionResult:
|
||||
command = tuple(
|
||||
build_ffprobe_command(
|
||||
path,
|
||||
executable=executable,
|
||||
launcher_args=launcher_args,
|
||||
)
|
||||
)
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
command,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
timeout=timeout_seconds,
|
||||
check=False,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
return ProbeExecutionResult(
|
||||
status="missing_command",
|
||||
command=command,
|
||||
exit_code=None,
|
||||
stdout_text="",
|
||||
stderr_text=str(exc),
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
return ProbeExecutionResult(
|
||||
status="timeout",
|
||||
command=command,
|
||||
exit_code=None,
|
||||
stdout_text=_coerce_process_output(exc.stdout),
|
||||
stderr_text=_coerce_process_output(exc.stderr),
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
)
|
||||
|
||||
return ProbeExecutionResult(
|
||||
status="completed",
|
||||
command=command,
|
||||
exit_code=completed.returncode,
|
||||
stdout_text=completed.stdout,
|
||||
stderr_text=completed.stderr,
|
||||
elapsed_ms=_elapsed_ms(start),
|
||||
)
|
||||
|
||||
|
||||
def _elapsed_ms(start: float) -> int:
|
||||
return int((time.perf_counter() - start) * 1000)
|
||||
|
||||
|
||||
def _coerce_process_output(output: str | bytes | None) -> str:
|
||||
if output is None:
|
||||
return ""
|
||||
if isinstance(output, bytes):
|
||||
return output.decode("utf-8", errors="replace")
|
||||
return output
|
||||
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FileCandidate:
|
||||
path: Path
|
||||
size: int
|
||||
mtime_ns: int
|
||||
extension: str
|
||||
discovery_order: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StoredFile:
|
||||
file_id: int
|
||||
run_id: int
|
||||
path: Path
|
||||
size: int
|
||||
mtime_ns: int
|
||||
extension: str
|
||||
discovery_order: int
|
||||
stage1_status: str
|
||||
stage2_status: str | None
|
||||
final_status: str | None
|
||||
resumed_from_run_id: int | None
|
||||
reason_code: str | None = None
|
||||
@@ -0,0 +1,58 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
from tgvideo_check.config import ScanConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Profile:
|
||||
name: str
|
||||
stage1_workers: int
|
||||
stage2_workers: int
|
||||
sample_duration: int
|
||||
high_risk_extensions: tuple[str, ...]
|
||||
|
||||
|
||||
PROFILES: dict[str, Profile] = {
|
||||
"tg-hdd": Profile(
|
||||
name="tg-hdd",
|
||||
stage1_workers=1,
|
||||
stage2_workers=1,
|
||||
sample_duration=3,
|
||||
high_risk_extensions=("mp4", "mov", "m4v", "mkv", "webm"),
|
||||
),
|
||||
"tg-fast": Profile(
|
||||
name="tg-fast",
|
||||
stage1_workers=2,
|
||||
stage2_workers=1,
|
||||
sample_duration=3,
|
||||
high_risk_extensions=("mp4", "mov", "m4v", "mkv", "webm"),
|
||||
),
|
||||
"deep-verify": Profile(
|
||||
name="deep-verify",
|
||||
stage1_workers=1,
|
||||
stage2_workers=1,
|
||||
sample_duration=5,
|
||||
high_risk_extensions=("mp4", "mov", "m4v", "mkv", "webm", "avi"),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_profile(name: str) -> Profile:
|
||||
try:
|
||||
return PROFILES[name]
|
||||
except KeyError as exc:
|
||||
raise ValueError(f"Unknown profile: {name}") from exc
|
||||
|
||||
|
||||
def resolve_config(config: ScanConfig) -> ScanConfig:
|
||||
profile = get_profile(config.profile)
|
||||
normalized = config.normalized()
|
||||
return replace(
|
||||
normalized,
|
||||
stage1_workers=normalized.stage1_workers or profile.stage1_workers,
|
||||
stage2_workers=normalized.stage2_workers or profile.stage2_workers,
|
||||
sample_duration=normalized.sample_duration or profile.sample_duration,
|
||||
high_risk_extensions=profile.high_risk_extensions,
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import orjson
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReportRow:
|
||||
path: Path
|
||||
final_status: str
|
||||
reason_code: str
|
||||
size: int
|
||||
|
||||
|
||||
def summarize_reason_counts(rows: list[ReportRow]) -> dict[str, int]:
|
||||
summary: dict[str, int] = {}
|
||||
for row in rows:
|
||||
summary[row.reason_code] = summary.get(row.reason_code, 0) + 1
|
||||
return summary
|
||||
|
||||
|
||||
def collect_redownload_candidates(rows: list[ReportRow]) -> list[Path]:
|
||||
return [
|
||||
row.path
|
||||
for row in rows
|
||||
if row.final_status in {"confirmed_broken", "suspicious"}
|
||||
]
|
||||
|
||||
|
||||
def export_jsonl(rows: list[ReportRow], destination: Path) -> None:
|
||||
lines = [
|
||||
orjson.dumps(
|
||||
{
|
||||
**asdict(row),
|
||||
"path": str(row.path),
|
||||
}
|
||||
).decode("utf-8")
|
||||
for row in rows
|
||||
]
|
||||
destination.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
|
||||
|
||||
|
||||
def export_csv_summary(rows: list[ReportRow], destination: Path) -> None:
|
||||
with destination.open("w", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.DictWriter(
|
||||
handle,
|
||||
fieldnames=["final_status", "reason_code", "path", "size"],
|
||||
)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(
|
||||
{
|
||||
"final_status": row.final_status,
|
||||
"reason_code": row.reason_code,
|
||||
"path": str(row.path),
|
||||
"size": row.size,
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Callable, TypeVar
|
||||
|
||||
from tgvideo_check.models import StoredFile
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StageBatch:
|
||||
stage_name: str
|
||||
worker_count: int
|
||||
files: tuple[StoredFile, ...]
|
||||
|
||||
|
||||
def select_pending_stage1_files(files: list[StoredFile]) -> list[StoredFile]:
|
||||
return sorted(
|
||||
(item for item in files if item.stage1_status == "pending_stage1"),
|
||||
key=lambda item: item.discovery_order,
|
||||
)
|
||||
|
||||
|
||||
def select_pending_stage2_files(files: list[StoredFile]) -> list[StoredFile]:
|
||||
return sorted(
|
||||
(
|
||||
item
|
||||
for item in files
|
||||
if item.stage1_status == "pending_stage2" and item.stage2_status is None
|
||||
),
|
||||
key=lambda item: item.discovery_order,
|
||||
)
|
||||
|
||||
|
||||
def build_stage_plan(
|
||||
*,
|
||||
stage1_files: list[StoredFile],
|
||||
stage2_files: list[StoredFile],
|
||||
stage1_workers: int,
|
||||
stage2_workers: int,
|
||||
queue_limit: int,
|
||||
) -> tuple[StageBatch, ...]:
|
||||
if queue_limit <= 0:
|
||||
raise ValueError("queue_limit must be positive")
|
||||
|
||||
batches: list[StageBatch] = []
|
||||
|
||||
for chunk in _chunk_by_limit(stage1_files, queue_limit):
|
||||
batches.append(
|
||||
StageBatch(
|
||||
stage_name="stage1",
|
||||
worker_count=stage1_workers,
|
||||
files=chunk,
|
||||
)
|
||||
)
|
||||
|
||||
for chunk in _chunk_by_limit(stage2_files, queue_limit):
|
||||
batches.append(
|
||||
StageBatch(
|
||||
stage_name="stage2",
|
||||
worker_count=stage2_workers,
|
||||
files=chunk,
|
||||
)
|
||||
)
|
||||
|
||||
return tuple(batches)
|
||||
|
||||
|
||||
def execute_stage_batch(
|
||||
batch: StageBatch,
|
||||
worker: Callable[[StoredFile], T],
|
||||
) -> tuple[T, ...]:
|
||||
if len(batch.files) <= 1 or batch.worker_count <= 1:
|
||||
return tuple(worker(item) for item in batch.files)
|
||||
|
||||
max_workers = min(batch.worker_count, len(batch.files))
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
return tuple(executor.map(worker, batch.files))
|
||||
|
||||
|
||||
def _chunk_by_limit(
|
||||
files: list[StoredFile], queue_limit: int
|
||||
) -> tuple[tuple[StoredFile, ...], ...]:
|
||||
ordered = sorted(files, key=lambda item: item.discovery_order)
|
||||
return tuple(
|
||||
tuple(ordered[index : index + queue_limit])
|
||||
for index in range(0, len(ordered), queue_limit)
|
||||
)
|
||||
@@ -0,0 +1,476 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import orjson
|
||||
|
||||
from tgvideo_check.models import FileCandidate, StoredFile
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScanRun:
|
||||
run_id: int
|
||||
profile: str
|
||||
root_paths_json: str
|
||||
status: str
|
||||
|
||||
|
||||
class ScanStore:
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = db_path
|
||||
|
||||
def initialize(self) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS scan_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
profile TEXT NOT NULL,
|
||||
root_paths_json TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'running'
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS files (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
run_id INTEGER NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
size INTEGER NOT NULL,
|
||||
mtime_ns INTEGER NOT NULL,
|
||||
extension TEXT NOT NULL,
|
||||
discovery_order INTEGER NOT NULL,
|
||||
stage1_status TEXT NOT NULL,
|
||||
stage2_status TEXT,
|
||||
final_status TEXT,
|
||||
resumed_from_run_id INTEGER,
|
||||
reason_code TEXT,
|
||||
UNIQUE (run_id, path),
|
||||
FOREIGN KEY (run_id) REFERENCES scan_runs(id)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS stage1_results (
|
||||
file_id INTEGER PRIMARY KEY,
|
||||
exit_code INTEGER,
|
||||
stdout_text TEXT,
|
||||
stderr_text TEXT,
|
||||
elapsed_ms INTEGER,
|
||||
FOREIGN KEY (file_id) REFERENCES files(id)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS stage2_results (
|
||||
file_id INTEGER PRIMARY KEY,
|
||||
mode TEXT,
|
||||
sample_points_json TEXT,
|
||||
exit_code INTEGER,
|
||||
stderr_text TEXT,
|
||||
reason_code TEXT,
|
||||
elapsed_ms INTEGER,
|
||||
FOREIGN KEY (file_id) REFERENCES files(id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_files_identity
|
||||
ON files(path, size, mtime_ns);
|
||||
CREATE INDEX IF NOT EXISTS idx_files_run_order
|
||||
ON files(run_id, discovery_order);
|
||||
"""
|
||||
)
|
||||
self._ensure_column(connection, "files", "reason_code", "TEXT")
|
||||
|
||||
def create_run(self, *, profile: str, root_paths: tuple[Path, ...]) -> int:
|
||||
root_paths_json = orjson.dumps([str(path) for path in root_paths]).decode("utf-8")
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"""
|
||||
INSERT INTO scan_runs(profile, root_paths_json, status)
|
||||
VALUES (?, ?, 'running')
|
||||
""",
|
||||
(profile, root_paths_json),
|
||||
)
|
||||
return int(cursor.lastrowid)
|
||||
|
||||
def get_latest_run_id(self, *, prefer_completed: bool = True) -> int | None:
|
||||
with self._connect() as connection:
|
||||
if prefer_completed:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT id
|
||||
FROM scan_runs
|
||||
WHERE status = 'completed'
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
"""
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
return int(row["id"])
|
||||
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT id
|
||||
FROM scan_runs
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
"""
|
||||
).fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
return int(row["id"])
|
||||
|
||||
def upsert_files(
|
||||
self,
|
||||
run_id: int,
|
||||
candidates: list[FileCandidate],
|
||||
*,
|
||||
reuse_previous_results: bool = True,
|
||||
) -> list[StoredFile]:
|
||||
with self._connect() as connection:
|
||||
for candidate in candidates:
|
||||
previous = (
|
||||
self._find_previous_match(connection, run_id, candidate)
|
||||
if reuse_previous_results
|
||||
else None
|
||||
)
|
||||
stage1_status = previous["stage1_status"] if previous else "pending_stage1"
|
||||
stage2_status = previous["stage2_status"] if previous else None
|
||||
final_status = previous["final_status"] if previous else None
|
||||
resumed_from_run_id = previous["run_id"] if previous else None
|
||||
reason_code = previous["reason_code"] if previous else None
|
||||
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO files(
|
||||
run_id,
|
||||
path,
|
||||
size,
|
||||
mtime_ns,
|
||||
extension,
|
||||
discovery_order,
|
||||
stage1_status,
|
||||
stage2_status,
|
||||
final_status,
|
||||
resumed_from_run_id,
|
||||
reason_code
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(run_id, path) DO UPDATE SET
|
||||
size = excluded.size,
|
||||
mtime_ns = excluded.mtime_ns,
|
||||
extension = excluded.extension,
|
||||
discovery_order = excluded.discovery_order,
|
||||
stage1_status = excluded.stage1_status,
|
||||
stage2_status = excluded.stage2_status,
|
||||
final_status = excluded.final_status,
|
||||
resumed_from_run_id = excluded.resumed_from_run_id,
|
||||
reason_code = excluded.reason_code
|
||||
""",
|
||||
(
|
||||
run_id,
|
||||
str(candidate.path),
|
||||
candidate.size,
|
||||
candidate.mtime_ns,
|
||||
candidate.extension,
|
||||
candidate.discovery_order,
|
||||
stage1_status,
|
||||
stage2_status,
|
||||
final_status,
|
||||
resumed_from_run_id,
|
||||
reason_code,
|
||||
),
|
||||
)
|
||||
current_file_id = connection.execute(
|
||||
"""
|
||||
SELECT id
|
||||
FROM files
|
||||
WHERE run_id = ? AND path = ?
|
||||
""",
|
||||
(run_id, str(candidate.path)),
|
||||
).fetchone()["id"]
|
||||
if previous is not None:
|
||||
self._copy_previous_results(
|
||||
connection,
|
||||
previous_file_id=previous["file_id"],
|
||||
current_file_id=current_file_id,
|
||||
)
|
||||
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT
|
||||
id,
|
||||
run_id,
|
||||
path,
|
||||
size,
|
||||
mtime_ns,
|
||||
extension,
|
||||
discovery_order,
|
||||
stage1_status,
|
||||
stage2_status,
|
||||
final_status,
|
||||
resumed_from_run_id,
|
||||
reason_code
|
||||
FROM files
|
||||
WHERE run_id = ?
|
||||
ORDER BY discovery_order ASC
|
||||
""",
|
||||
(run_id,),
|
||||
).fetchall()
|
||||
|
||||
return self._rows_to_stored_files(rows)
|
||||
|
||||
def record_stage1_status(
|
||||
self,
|
||||
file_id: int,
|
||||
*,
|
||||
stage1_status: str,
|
||||
final_status: str | None,
|
||||
reason_code: str | None = None,
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE files
|
||||
SET stage1_status = ?, final_status = ?, reason_code = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(stage1_status, final_status, reason_code, file_id),
|
||||
)
|
||||
|
||||
def record_stage1_result(
|
||||
self,
|
||||
file_id: int,
|
||||
*,
|
||||
exit_code: int | None,
|
||||
stdout_text: str,
|
||||
stderr_text: str,
|
||||
elapsed_ms: int,
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO stage1_results(file_id, exit_code, stdout_text, stderr_text, elapsed_ms)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(file_id) DO UPDATE SET
|
||||
exit_code = excluded.exit_code,
|
||||
stdout_text = excluded.stdout_text,
|
||||
stderr_text = excluded.stderr_text,
|
||||
elapsed_ms = excluded.elapsed_ms
|
||||
""",
|
||||
(file_id, exit_code, stdout_text, stderr_text, elapsed_ms),
|
||||
)
|
||||
|
||||
def record_stage2_status(
|
||||
self,
|
||||
file_id: int,
|
||||
*,
|
||||
stage2_status: str,
|
||||
final_status: str | None,
|
||||
reason_code: str | None = None,
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE files
|
||||
SET stage2_status = ?, final_status = ?, reason_code = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(stage2_status, final_status, reason_code, file_id),
|
||||
)
|
||||
|
||||
def record_stage2_result(
|
||||
self,
|
||||
file_id: int,
|
||||
*,
|
||||
mode: str,
|
||||
sample_points_json: str,
|
||||
exit_code: int | None,
|
||||
stderr_text: str,
|
||||
reason_code: str | None,
|
||||
elapsed_ms: int,
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO stage2_results(
|
||||
file_id,
|
||||
mode,
|
||||
sample_points_json,
|
||||
exit_code,
|
||||
stderr_text,
|
||||
reason_code,
|
||||
elapsed_ms
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(file_id) DO UPDATE SET
|
||||
mode = excluded.mode,
|
||||
sample_points_json = excluded.sample_points_json,
|
||||
exit_code = excluded.exit_code,
|
||||
stderr_text = excluded.stderr_text,
|
||||
reason_code = excluded.reason_code,
|
||||
elapsed_ms = excluded.elapsed_ms
|
||||
""",
|
||||
(
|
||||
file_id,
|
||||
mode,
|
||||
sample_points_json,
|
||||
exit_code,
|
||||
stderr_text,
|
||||
reason_code,
|
||||
elapsed_ms,
|
||||
),
|
||||
)
|
||||
|
||||
def update_run_status(self, run_id: int, *, status: str) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE scan_runs
|
||||
SET status = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(status, run_id),
|
||||
)
|
||||
|
||||
def get_stage1_stdout_text(self, file_id: int) -> str | None:
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT stdout_text
|
||||
FROM stage1_results
|
||||
WHERE file_id = ?
|
||||
""",
|
||||
(file_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return row["stdout_text"]
|
||||
|
||||
def list_files(self, run_id: int) -> list[StoredFile]:
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT
|
||||
id,
|
||||
run_id,
|
||||
path,
|
||||
size,
|
||||
mtime_ns,
|
||||
extension,
|
||||
discovery_order,
|
||||
stage1_status,
|
||||
stage2_status,
|
||||
final_status,
|
||||
resumed_from_run_id,
|
||||
reason_code
|
||||
FROM files
|
||||
WHERE run_id = ?
|
||||
ORDER BY discovery_order ASC
|
||||
""",
|
||||
(run_id,),
|
||||
).fetchall()
|
||||
|
||||
return self._rows_to_stored_files(rows)
|
||||
|
||||
def _rows_to_stored_files(self, rows: list[sqlite3.Row]) -> list[StoredFile]:
|
||||
return [
|
||||
StoredFile(
|
||||
file_id=row["id"],
|
||||
run_id=row["run_id"],
|
||||
path=Path(row["path"]),
|
||||
size=row["size"],
|
||||
mtime_ns=row["mtime_ns"],
|
||||
extension=row["extension"],
|
||||
discovery_order=row["discovery_order"],
|
||||
stage1_status=row["stage1_status"],
|
||||
stage2_status=row["stage2_status"],
|
||||
final_status=row["final_status"],
|
||||
resumed_from_run_id=row["resumed_from_run_id"],
|
||||
reason_code=row["reason_code"],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.db_path)
|
||||
connection.row_factory = sqlite3.Row
|
||||
return connection
|
||||
|
||||
def _find_previous_match(
|
||||
self,
|
||||
connection: sqlite3.Connection,
|
||||
run_id: int,
|
||||
candidate: FileCandidate,
|
||||
) -> sqlite3.Row | None:
|
||||
return connection.execute(
|
||||
"""
|
||||
SELECT id AS file_id, run_id, stage1_status, stage2_status, final_status, reason_code
|
||||
FROM files
|
||||
WHERE path = ? AND size = ? AND mtime_ns = ? AND run_id <> ?
|
||||
ORDER BY run_id DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(str(candidate.path), candidate.size, candidate.mtime_ns, run_id),
|
||||
).fetchone()
|
||||
|
||||
def _copy_previous_results(
|
||||
self,
|
||||
connection: sqlite3.Connection,
|
||||
*,
|
||||
previous_file_id: int,
|
||||
current_file_id: int,
|
||||
) -> None:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO stage1_results(file_id, exit_code, stdout_text, stderr_text, elapsed_ms)
|
||||
SELECT ?, exit_code, stdout_text, stderr_text, elapsed_ms
|
||||
FROM stage1_results
|
||||
WHERE file_id = ?
|
||||
ON CONFLICT(file_id) DO UPDATE SET
|
||||
exit_code = excluded.exit_code,
|
||||
stdout_text = excluded.stdout_text,
|
||||
stderr_text = excluded.stderr_text,
|
||||
elapsed_ms = excluded.elapsed_ms
|
||||
""",
|
||||
(current_file_id, previous_file_id),
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO stage2_results(
|
||||
file_id,
|
||||
mode,
|
||||
sample_points_json,
|
||||
exit_code,
|
||||
stderr_text,
|
||||
reason_code,
|
||||
elapsed_ms
|
||||
)
|
||||
SELECT ?, mode, sample_points_json, exit_code, stderr_text, reason_code, elapsed_ms
|
||||
FROM stage2_results
|
||||
WHERE file_id = ?
|
||||
ON CONFLICT(file_id) DO UPDATE SET
|
||||
mode = excluded.mode,
|
||||
sample_points_json = excluded.sample_points_json,
|
||||
exit_code = excluded.exit_code,
|
||||
stderr_text = excluded.stderr_text,
|
||||
reason_code = excluded.reason_code,
|
||||
elapsed_ms = excluded.elapsed_ms
|
||||
""",
|
||||
(current_file_id, previous_file_id),
|
||||
)
|
||||
|
||||
def _ensure_column(
|
||||
self,
|
||||
connection: sqlite3.Connection,
|
||||
table_name: str,
|
||||
column_name: str,
|
||||
column_definition: str,
|
||||
) -> None:
|
||||
columns = {
|
||||
row["name"]
|
||||
for row in connection.execute(f"PRAGMA table_info({table_name})").fetchall()
|
||||
}
|
||||
if column_name in columns:
|
||||
return
|
||||
connection.execute(
|
||||
f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_definition}"
|
||||
)
|
||||
Reference in New Issue
Block a user