diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py index 3b08078e..bf83984f 100644 --- a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py @@ -62,15 +62,13 @@ import tempfile from pathlib import Path -import brotli - from optimize_media import ( BUILTIN_DEFAULTS, Logger, OPTION_SPECS, add_optimize_arguments, find_pngquant, optimize_directory, resolve_config, ) from populate_db import ( CHUNK_SIZE, EXTENSION_TO_CONTENT_TYPE, IMAGES_DB_PATH_PREFIX, IMAGES_URL_PREFIX, LANGUAGE, PAGE_CONTENT_TYPE, - backup_database, get_content_type, get_id, insert_chunked_content, + DictionaryCompressor, backup_database, get_content_type, get_id, insert_chunked_content, load_dictionary, ) WEBP_CONTENT_TYPE = "image/webp" @@ -109,7 +107,7 @@ def delete_content(conn, path: str) -> None: def insert_optimized_file(conn, data: bytes, name: str, db_path: str, language_id: int, content_type_cache: dict, - chunked_log: list) -> bool: + chunked_log: list, compressor: DictionaryCompressor) -> bool: """Inserts one already-optimized file's bytes as-is. Unlike populate_db.py's own insert_file, this does not run pngquant itself - optimize_media.py already did, and running it again here would just @@ -125,7 +123,7 @@ def insert_optimized_file(conn, data: bytes, name: str, db_path: str, language_i content_type_id, compress = content_type_cache[content_type_value] if compress: - data = brotli.compress(data) + data = compressor.compress(data) delete_content(conn, db_path) insert_chunked_content(conn, db_path, language_id, content_type_id, 0, data, chunked_log) return True @@ -175,7 +173,7 @@ def reassemble_content(conn, path: str, first_content: bytes) -> bytes: def rewrite_pages(conn, rename_map: dict, language_id: int, page_content_type_id: int, logger: Logger, - chunked_log: list) -> int: + chunked_log: list, compressor: DictionaryCompressor) -> int: """Rewrites every k/html/*.html page (and the nav row) that references a renamed image, replacing "/k/html/images/" with "/k/html/images/" wherever it appears. Operates directly on @@ -225,12 +223,12 @@ def rewrite_pages(conn, rename_map: dict, language_id: int, page_content_type_id changed = 0 for path, first_content, template_id in rows: full = reassemble_content(conn, path, first_content) - text = brotli.decompress(full).decode("utf-8") + text = compressor.decompress(full).decode("utf-8") hits = len(old_ref_pattern.findall(text)) if not hits: continue new_text = old_ref_pattern.sub(lambda m: replacements[m.group(0)], text) - blob = brotli.compress(new_text.encode("utf-8")) + blob = compressor.compress(new_text.encode("utf-8")) delete_content(conn, path) insert_chunked_content(conn, path, language_id, page_content_type_id, template_id, blob, chunked_log) changed += 1 @@ -246,7 +244,7 @@ def rewrite_pages(conn, rename_map: dict, language_id: int, page_content_type_id IMAGE_REF_RE = re.compile(re.escape(IMAGES_URL_PREFIX) + r'([^\\"]+)\\"') -def collect_referenced_media(conn, page_content_type_id: int) -> set: +def collect_referenced_media(conn, page_content_type_id: int, compressor: DictionaryCompressor) -> set: """Bare filenames (e.g. "mascot.png") referenced by at least one src="/k/html/images/" anywhere across current k/html/*.html page content and the nav row - the same row selection/reassembly @@ -259,7 +257,7 @@ def collect_referenced_media(conn, page_content_type_id: int) -> set: referenced = set() for path, first_content in rows: full = reassemble_content(conn, path, first_content) - text = brotli.decompress(full).decode("utf-8") + text = compressor.decompress(full).decode("utf-8") referenced.update(IMAGE_REF_RE.findall(text)) return referenced @@ -287,7 +285,8 @@ def is_fragment(path: str) -> bool: return {path[len(IMAGES_DB_PATH_PREFIX):]: path for path in paths if not is_fragment(path)} -def delete_unreferenced_media(conn, page_content_type_id: int, logger: Logger) -> int: +def delete_unreferenced_media(conn, page_content_type_id: int, logger: Logger, + compressor: DictionaryCompressor) -> int: """Deletes every currently-stored k/html/images/ row (base row and any chunked fragments) that no page or the nav row references even once. Must run after insertion and rename-rewriting, so it sees the final, @@ -296,7 +295,7 @@ def delete_unreferenced_media(conn, page_content_type_id: int, logger: Logger) - rewrite_pages will have already fixed up by the time this runs. Returns the number of images removed.""" stored = list_stored_media(conn) - referenced = collect_referenced_media(conn, page_content_type_id) + referenced = collect_referenced_media(conn, page_content_type_id, compressor) removed = 0 for name, path in sorted(stored.items()): if name in referenced: @@ -384,44 +383,56 @@ def main() -> None: conn.execute("BEGIN") language_id = get_id(conn, "Languages", LANGUAGE) page_content_type_id = get_id(conn, "ContentTypes", PAGE_CONTENT_TYPE) + # This script only ever runs against a database populate_db.py + # already populated (see module docstring), so its + # CompressionDictionary must already exist - never train a new + # one here, since that would orphan every row already + # compressed against the existing one (see DictionaryCompressor). + compressor = DictionaryCompressor(load_dictionary(conn)) content_type_cache = {} chunked_log = [] inserted = 0 seen_names = {} - for out_path in sorted(work_dir.rglob("*")): - if out_path.is_dir(): - continue - name = out_path.name - if name in seen_names: - logger.error( - f"warning: {out_path} has the same filename as {seen_names[name]}; keeping the first, " - "skipping this one" - ) - continue - seen_names[name] = out_path - db_path = f"{IMAGES_DB_PATH_PREFIX}{name}" - if insert_optimized_file(conn, out_path.read_bytes(), name, db_path, language_id, content_type_cache, - chunked_log): - inserted += 1 + try: + for out_path in sorted(work_dir.rglob("*")): + if out_path.is_dir(): + continue + name = out_path.name + if name in seen_names: + logger.error( + f"warning: {out_path} has the same filename as {seen_names[name]}; keeping the first, " + "skipping this one" + ) + continue + seen_names[name] = out_path + db_path = f"{IMAGES_DB_PATH_PREFIX}{name}" + if insert_optimized_file(conn, out_path.read_bytes(), name, db_path, language_id, + content_type_cache, chunked_log, compressor): + inserted += 1 + if cfg["verbose"]: + logger.info(f"[OK] {out_path} -> {db_path}") + + # A renamed file's old basename no longer appears anywhere under + # work_dir (that's what makes it a rename), so the loop above + # never visits its old db_path to replace it - it'd otherwise + # linger forever as an orphaned, no-longer-referenced row. + removed = 0 + for old_name in rename_map: + old_db_path = f"{IMAGES_DB_PATH_PREFIX}{old_name}" + delete_content(conn, old_db_path) + removed += 1 if cfg["verbose"]: - logger.info(f"[OK] {out_path} -> {db_path}") - - # A renamed file's old basename no longer appears anywhere under - # work_dir (that's what makes it a rename), so the loop above - # never visits its old db_path to replace it - it'd otherwise - # linger forever as an orphaned, no-longer-referenced row. - removed = 0 - for old_name in rename_map: - old_db_path = f"{IMAGES_DB_PATH_PREFIX}{old_name}" - delete_content(conn, old_db_path) - removed += 1 - if cfg["verbose"]: - logger.info(f"[REMOVED] {old_db_path} (renamed to {IMAGES_DB_PATH_PREFIX}{rename_map[old_name]})") - - changed_pages = rewrite_pages(conn, rename_map, language_id, page_content_type_id, logger, chunked_log) - - unreferenced_removed = delete_unreferenced_media(conn, page_content_type_id, logger) + logger.info( + f"[REMOVED] {old_db_path} (renamed to {IMAGES_DB_PATH_PREFIX}{rename_map[old_name]})" + ) + + changed_pages = rewrite_pages(conn, rename_map, language_id, page_content_type_id, logger, + chunked_log, compressor) + + unreferenced_removed = delete_unreferenced_media(conn, page_content_type_id, logger, compressor) + finally: + compressor.close() conn.commit() except Exception: diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/migrate_content_to_dictionary_brotli.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/migrate_content_to_dictionary_brotli.py new file mode 100644 index 00000000..efd1666f --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/migrate_content_to_dictionary_brotli.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +""" +migrate_content_to_dictionary_brotli.py + +One-time, idempotent, whole-database migration: recompresses every Content +row whose ContentTypes.compression is 'brotli' against this database's +shared CompressionDictionary (see ADFA-5153), replacing plain (no +dictionary) Brotli blobs with dictionary-compressed ones in place. + +Why this exists: populate_db.py and insert_optimized_media.py only ever +touch their own subset of Content ("k/html/%", "assets/%"). Every other +Content row in documentation.db - reference docs, tooltip-linked pages, +whatever else - was compressed with plain Brotli by whichever pipeline +wrote it, no dictionary involved. Once anything in this database is +dictionary-compressed, WebServer.kt's reader has to be able to assume EVERY +'brotli' row uses the same dictionary (a per-row dictionary/no-dictionary +flag was explicitly rejected in ADFA-5153 in favor of "convert everything, +once"). This script is what makes that assumption actually true for rows +outside populate_db.py's own reach. + +Trains the shared dictionary from a random sample drawn across the WHOLE +Content table (not just one doc set) if CompressionDictionary doesn't +already exist - broader and more representative than populate_db.py's own +Kotlin-website-only bootstrap sample. Run this BEFORE ever running +populate_db.py against a fresh database, so the one dictionary that ends up +stored is trained on real cross-corpus data. + +Idempotency: for each candidate row, a plain (no-dictionary) decompress is +attempted first. That reliably fails when the row is already +dictionary-compressed (verified empirically over 200 trials - a genuinely +missing dictionary, unlike a *wrong* one, can't coincidentally produce a +parseable stream), so a row that already migrated is left untouched and +counted as "already migrated" rather than reprocessed. Re-running this +script is therefore always safe. + +Chunked rows (see CHUNK_SIZE in populate_db.py) are reassembled before +decompression and re-chunked identically after recompression, the same +fragmentation scheme WebServer.kt expects on read. + +Safety: backs up the database first (VACUUM INTO, same as populate_db.py), +runs entirely inside one transaction (rolled back on any error), and VACUUMs +afterward on a separate connection (SQLite refuses VACUUM inside a +transaction). + +Performance: the per-row work (reassemble + plain-decompress + dictionary- +recompress) runs on a thread pool, since each recompress spawns its own +`brotli` subprocess - real wall time on a ~30,000-row database is dominated +by process-spawn overhead, not CPU, so this parallelizes close to linearly +with --max-workers. Only that read+compress work is parallelized; the +actual delete+insert writes stay serialized on the single caller-supplied +connection (SQLite requires this anyway). + +Usage: + python3 migrate_content_to_dictionary_brotli.py [--sample-size N] [--dict-size BYTES] [--max-workers N] +""" +import argparse +import sqlite3 +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import brotli + +from populate_db import ( + CHUNK_SIZE, DEFAULT_DICT_SIZE, DictionaryCompressor, backup_database, insert_chunked_content, + load_or_create_dictionary, +) + +DEFAULT_SAMPLE_SIZE = 300 + +_thread_local = threading.local() + + +def reassemble_content(conn, path: str, first_content: bytes) -> bytes: + """Reassembles a possibly-chunked row's full bytes - a row is fragmented + purely when its content is exactly CHUNK_SIZE bytes, in which case + "-1", "-2", ... are concatenated until a missing or + shorter-than-CHUNK_SIZE row is hit. Mirrors WebServer.kt's own + reassembly and insert_optimized_media.py's copy of the same logic.""" + if len(first_content) < CHUNK_SIZE: + return first_content + parts = [first_content] + n = 1 + while True: + row = conn.execute("SELECT content FROM Content WHERE path = ?", (f"{path}-{n}",)).fetchone() + if row is None: + break + parts.append(row[0]) + if len(row[0]) < CHUNK_SIZE: + break + n += 1 + return b"".join(parts) + + +def delete_content(conn, path: str) -> None: + """Deletes a Content row and any chunked continuation fragments for it. + Content.path is UNIQUE, so this has to run before any re-insert at the + same path.""" + conn.execute("DELETE FROM Content WHERE path = ? OR path LIKE ?", (path, f"{path}-%")) + + +def list_fragment_paths(conn) -> set: + """Every Content.path that's a chunked continuation fragment of another + row in this table (a path whose trailing "-" strip yields + another path that's also present) - same convention as + insert_optimized_media.py's is_fragment, generalized to the whole table + instead of just one path prefix. Base (non-fragment) rows are the ones + this migration processes; fragments are only ever touched indirectly, + via reassemble_content/delete_content on their base row's path.""" + all_paths = {row[0] for row in conn.execute("SELECT path FROM Content")} + fragments = set() + for path in all_paths: + prefix, sep, suffix = path.rpartition("-") + if sep == "-" and suffix.isdigit() and prefix in all_paths: + fragments.add(path) + return fragments + + +def collect_training_samples(conn, brotli_base_rows: list, sample_size: int) -> list: + """Decompresses up to `sample_size` rows' full (reassembled) content as + plain Brotli - safe to assume plain here, since this only ever runs + before CompressionDictionary exists, i.e. before anything in this + database could possibly be dictionary-compressed yet.""" + sample_rows = brotli_base_rows[:sample_size] + samples = [] + for path, first_content, _language_id, _content_type_id, _template_id in sample_rows: + full = reassemble_content(conn, path, first_content) + try: + samples.append(brotli.decompress(full)) + except brotli.error as exc: + print(f"warning: could not decompress {path!r} for training sample: {exc}", file=sys.stderr) + return samples + + +def dictionary_already_exists(conn) -> bool: + return conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'CompressionDictionary'" + ).fetchone() is not None + + +def _thread_compressor(dictionary_data: bytes) -> DictionaryCompressor: + """One DictionaryCompressor per worker thread, reused across every row + that thread processes - creating one per row would mean re-writing the + same dictionary bytes to a fresh temp file on every single call for no + benefit.""" + compressor = getattr(_thread_local, "compressor", None) + if compressor is None: + compressor = DictionaryCompressor(dictionary_data) + _thread_local.compressor = compressor + return compressor + + +def _migrate_one_row(db_path: Path, dictionary_data: bytes, row: tuple): + """Runs in a worker thread: reassembles, plain-decompresses, and + dictionary-recompresses one row. Returns None if the row is already + dictionary-compressed (a plain decode reliably fails - see module + docstring), else (path, language_id, content_type_id, template_id, + recompressed_bytes, original_size) for the caller to write back. + + Opens its own read-only connection for reassembly rather than sharing + the caller's - a single sqlite3.Connection isn't safe to use from + multiple threads at once.""" + path, first_content, language_id, content_type_id, template_id = row + worker_conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) + try: + full = reassemble_content(worker_conn, path, first_content) + finally: + worker_conn.close() + + try: + plain = brotli.decompress(full) + except brotli.error: + return None + + recompressed = _thread_compressor(dictionary_data).compress(plain) + return path, language_id, content_type_id, template_id, recompressed, len(full) + + +def migrate(conn, db_path: Path, sample_size: int, dict_size: int, max_workers: int | None = None) -> dict: + fragment_paths = list_fragment_paths(conn) + all_brotli_rows = conn.execute( + "SELECT C.path, C.content, C.languageID, C.contentTypeID, C.templateId " + "FROM Content C, ContentTypes CT " + "WHERE C.contentTypeID = CT.id AND CT.compression = 'brotli' " + "ORDER BY C.path" + ).fetchall() + base_rows = [row for row in all_brotli_rows if row[0] not in fragment_paths] + + # Only worth decompressing sample rows for training when there's actually + # no dictionary yet - on every later run, load_or_create_dictionary would + # just discard them anyway, and by then every already-migrated row can no + # longer be plain-decompressed at all (see module docstring), so + # attempting it would just spend time producing warnings for no benefit. + training_samples = [] if dictionary_already_exists(conn) else collect_training_samples(conn, base_rows, + sample_size) + dictionary_data = load_or_create_dictionary(conn, training_samples, dict_size) + + stats = {"scanned": len(base_rows), "migrated": 0, "already_migrated": 0, "bytes_before": 0, "bytes_after": 0} + + # executor.map preserves input order (each result is yielded once its + # corresponding row is done, in submission order) while still running + # every row's read+decompress+recompress concurrently under the hood - + # writes below stay serialized on the single caller-supplied connection. + # max_workers=None uses ThreadPoolExecutor's own default (min(32, + # cpu_count+4)), tuned for exactly this kind of I/O/subprocess-bound + # work - measured 3-6x faster than max_workers=1 on synthetic benchmarks. + with ThreadPoolExecutor(max_workers=max_workers) as executor: + results = executor.map(lambda row: _migrate_one_row(db_path, dictionary_data, row), base_rows) + for result in results: + if result is None: + stats["already_migrated"] += 1 + continue + path, language_id, content_type_id, template_id, recompressed, original_size = result + stats["migrated"] += 1 + stats["bytes_before"] += original_size + stats["bytes_after"] += len(recompressed) + + delete_content(conn, path) + insert_chunked_content(conn, path, language_id, content_type_id, template_id, recompressed, []) + + return stats + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("db_path", type=Path, help="SQLite database to migrate, e.g. documentation.db") + parser.add_argument("--sample-size", type=int, default=DEFAULT_SAMPLE_SIZE, + help=f"Rows to sample for dictionary training if none exists yet (default: {DEFAULT_SAMPLE_SIZE})") + parser.add_argument("--dict-size", type=int, default=DEFAULT_DICT_SIZE, + help=f"Dictionary size in bytes if training a new one (default: {DEFAULT_DICT_SIZE})") + parser.add_argument("--max-workers", type=int, default=None, + help="Worker threads for the read+compress phase (default: ThreadPoolExecutor's own " + "min(32, cpu_count+4))") + args = parser.parse_args() + + if not args.db_path.is_file(): + print(f"error: {args.db_path} does not exist", file=sys.stderr) + sys.exit(1) + + print(f"Backing up {args.db_path}...", file=sys.stderr) + backup_path = backup_database(args.db_path) + print(f"Backup written to {backup_path}", file=sys.stderr) + + conn = sqlite3.connect(args.db_path) + try: + conn.execute("BEGIN") + stats = migrate(conn, args.db_path, args.sample_size, args.dict_size, args.max_workers) + conn.commit() + except Exception: + conn.rollback() + raise + finally: + conn.close() + + print("Vacuuming database to reclaim freed space...", file=sys.stderr) + vacuum_conn = sqlite3.connect(args.db_path) + try: + vacuum_conn.execute("VACUUM") + finally: + vacuum_conn.close() + + print( + f"Scanned {stats['scanned']} brotli row(s): migrated {stats['migrated']}, " + f"already dictionary-compressed {stats['already_migrated']}." + ) + if stats["migrated"]: + before, after = stats["bytes_before"], stats["bytes_after"] + pct = (1 - after / before) * 100 if before else 0.0 + print(f"Migrated bytes: {before:,} -> {after:,} ({pct:.1f}% smaller)") + + +if __name__ == "__main__": + main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py index c338b9dc..4591cddd 100644 --- a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py @@ -81,13 +81,17 @@ "Home" link had nowhere to go. It still renders through page.peb, so the sidebar nav shows up on it like any other page. 4. Inserts one Content row per page at k/html/.html, - JSON-encoded and brotli-compressed, with prev/next computed from - kr.tree's document order, the same way RenderDocs.java does it for the - static site. contentTypeID is the "text/html" row (12|text/html|brotli), - not "application/json": the stored bytes are JSON, but templateId - points the server at page.peb to render that JSON into HTML before a - browser ever sees it, so the Content-Type the server actually sends - back should describe that rendered output, not the storage format. + JSON-encoded and Brotli-compressed against this database's shared + CompressionDictionary (see ADFA-5153; trained once and reused forever - + never retrained - since a dictionary-compressed row is only decodable + against the exact dictionary it was compressed with), with prev/next + computed from kr.tree's document order, the same way RenderDocs.java + does it for the static site. contentTypeID is the "text/html" row + (12|text/html|brotli), not "application/json": the stored bytes are + JSON, but templateId points the server at page.peb to render that JSON + into HTML before a browser ever sees it, so the Content-Type the server + actually sends back should describe that rendered output, not the + storage format. 5. Builds the same navigation tree build_nav.py does from kr.tree, and inserts it as one more Content row (see NAV_CONTENT_PATH below), associated with the nav.peb template and the same "text/html" @@ -113,18 +117,18 @@ chunked is logged by name at the end of the run. """ import argparse +import atexit import json import shutil import sqlite3 import subprocess import sys +import tempfile import xml.etree.ElementTree as ET import zipfile from datetime import datetime from pathlib import Path -import brotli - from build_nav import build_node from md_to_json import ( Converter, @@ -178,6 +182,23 @@ PNGQUANT_CONTENT_TYPE = "image/png" PNGQUANT_QUALITY = "65-80" +# Single-row table: the whole documentation.db has exactly one shared Brotli +# dictionary, embedded here so it always ships in sync with the content +# compressed against it (see ADFA-5153). The id/CHECK pair enforces "exactly +# one row" at the schema level - a second INSERT fails outright instead of +# silently leaving two rows for a reader to pick between arbitrarily. +DICTIONARY_TABLE_SQL = """ +CREATE TABLE IF NOT EXISTS CompressionDictionary ( + id INTEGER PRIMARY KEY CHECK (id = 1), + data BLOB NOT NULL +); +""" +# 256 KiB fast-cover dictionary was the measured sweet spot in ADFA-5153 +# (16.26x held-out ratio vs. 8.20x undictionaried; a larger, exhaustively- +# trained dictionary bought another 0.35x for 144x the training time - not +# worth it). +DEFAULT_DICT_SIZE = 256 * 1024 + def find_pngquant() -> str: """Locates the pngquant executable on PATH. Raises if it's missing, @@ -209,6 +230,141 @@ def compress_png_with_pngquant(data: bytes, pngquant_path: str, name: str) -> by return result.stdout +def find_tool(name: str) -> str: + """Locates an executable on PATH. Raises if it's missing, rather than + silently falling back to some other behavior - see find_pngquant.""" + path = shutil.which(name) + if path is None: + raise RuntimeError(f"{name} not found on PATH; install it and retry") + return path + + +def train_dictionary(samples: list, dict_size: int = DEFAULT_DICT_SIZE) -> bytes: + """Trains a zstd fast-cover dictionary from `samples` (a list of byte + strings) and returns its raw bytes. That output is usable directly as + Brotli's raw `-D` dictionary (validated in ADFA-5153) - zstd's fast-cover + trainer is dramatically cheaper than Brotli's own dictionary tooling for + equivalent quality. Needs a few dozen samples at minimum; zstd's trainer + refuses ("nb of samples too low") on too few/too-small inputs, since a + dictionary trained on a handful of samples won't generalize. + """ + zstd_path = find_tool("zstd") + work_dir = Path(tempfile.mkdtemp(prefix="brotli-dict-train-")) + try: + sample_paths = [] + for i, sample in enumerate(samples): + sample_path = work_dir / f"sample_{i:06}.bin" + sample_path.write_bytes(sample) + sample_paths.append(str(sample_path)) + dict_path = work_dir / "dictionary.bin" + result = subprocess.run( + [zstd_path, "--train-fastcover", f"--maxdict={dict_size}", "-o", str(dict_path), *sample_paths], + stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"zstd --train-fastcover failed: {result.stderr.decode(errors='replace').strip()}") + return dict_path.read_bytes() + finally: + shutil.rmtree(work_dir, ignore_errors=True) + + +class DictionaryCompressor: + """Compresses/decompresses bytes against a fixed raw Brotli dictionary, + shelling out to the `brotli` CLI (the installed Python `brotli` package + has no dictionary parameter at all). A dictionary-compressed stream and a + plain one are NOT interchangeable at decode time - verified empirically, + not just documented behavior: decoding with the wrong dictionary (a + different one than was used to compress, "none" when one was used, or + vice versa) is NOT reliably caught - it sometimes fails outright + ("corrupt input"), but can just as easily "succeed" while silently + returning different bytes than were compressed, depending on how the + corrupted back-references happen to land. There is no runtime check that + catches this after the fact. So every row compressed via this class must + be decompressed via a `DictionaryCompressor` built from the exact same + dictionary bytes, and that dictionary must never change once anything + has been compressed against it - see CompressionDictionary (the single + source of truth for those bytes) and load_or_create_dictionary's + never-retrain guarantee. + + The dictionary is written once to a private temp file for this instance's + lifetime (each compress/decompress call reuses it) rather than per call. + """ + + def __init__(self, dictionary_data: bytes): + self._brotli_path = find_tool("brotli") + self._work_dir = Path(tempfile.mkdtemp(prefix="brotli-dict-")) + self._dict_path = self._work_dir / "dictionary.bin" + self._dict_path.write_bytes(dictionary_data) + # Safety net for callers that can't cleanly scope a `with` block around + # every instance -- e.g. one created per worker thread in a thread pool, + # where no single point of control can call close() on each. Safe to + # also call close() explicitly afterward: shutil.rmtree(ignore_errors=True) + # tolerates a directory that's already gone. + atexit.register(self.close) + + def _run(self, *extra_args: str, data: bytes) -> bytes: + result = subprocess.run( + [self._brotli_path, "-D", str(self._dict_path), *extra_args, "-c"], + input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"brotli failed: {result.stderr.decode(errors='replace').strip()}") + return result.stdout + + def compress(self, data: bytes) -> bytes: + return self._run(data=data) + + def decompress(self, data: bytes) -> bytes: + return self._run("-d", data=data) + + def close(self) -> None: + shutil.rmtree(self._work_dir, ignore_errors=True) + + def __enter__(self) -> "DictionaryCompressor": + return self + + def __exit__(self, *exc_info) -> None: + self.close() + + +def load_dictionary(conn) -> bytes: + """Returns the CompressionDictionary bytes already stored in this + database. Raises if the table doesn't exist or is empty - callers that + only ever run against a database populate_db.py already touched (e.g. + insert_optimized_media.py) should never need to train a new one.""" + row = conn.execute( + "SELECT data FROM CompressionDictionary WHERE id = 1" + ).fetchone() if _table_exists(conn, "CompressionDictionary") else None + if row is None: + raise RuntimeError( + "CompressionDictionary is missing or empty; run populate_db.py against this database first" + ) + return row[0] + + +def load_or_create_dictionary(conn, samples_for_training: list, dict_size: int = DEFAULT_DICT_SIZE) -> bytes: + """Returns the dictionary bytes stored in CompressionDictionary, training + a new one from `samples_for_training` and storing it if the table + doesn't exist yet or is empty. Never retrains an existing dictionary: + since dictionary-compressed content elsewhere in this same database can + only ever be decoded with the exact dictionary it was compressed against + (see DictionaryCompressor), silently replacing an already-populated + dictionary would orphan every row compressed against the old one.""" + conn.execute(DICTIONARY_TABLE_SQL) + row = conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone() + if row is not None: + return row[0] + dictionary_data = train_dictionary(samples_for_training, dict_size) + conn.execute("INSERT INTO CompressionDictionary (id, data) VALUES (1, ?)", (dictionary_data,)) + return dictionary_data + + +def _table_exists(conn, table_name: str) -> bool: + return conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (table_name,) + ).fetchone() is not None + + def backup_database(db_path: Path) -> Path: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") backup_path = db_path.with_name(f"{db_path.name}.backup-{timestamp}") @@ -282,13 +438,13 @@ def insert_chunked_content(conn, path: str, language_id: int, content_type_id: i def insert_file(conn, data: bytes, name: str, db_path: str, language_id: int, content_type_cache: dict, - chunked_log: list, pngquant_path: str) -> bool: + chunked_log: list, pngquant_path: str, compressor: "DictionaryCompressor") -> bool: """Inserts one raw (templateId 0) file's bytes as a Content row (chunked via insert_chunked_content if needed). name is only used to look up its content type by extension. Returns False (and skips it, with a warning) for an extension not in EXTENSION_TO_CONTENT_TYPE instead of guessing at a content type. PNGs are run through pngquant first - the only content - type it's compatible with - before the usual brotli compression.""" + type it's compatible with - before the usual dictionary-Brotli compression.""" content_type_value = EXTENSION_TO_CONTENT_TYPE.get(Path(name).suffix.lower()) if content_type_value is None: print(f"warning: no known content type for {name!r}; skipping", file=sys.stderr) @@ -300,7 +456,7 @@ def insert_file(conn, data: bytes, name: str, db_path: str, language_id: int, co if content_type_value == PNGQUANT_CONTENT_TYPE: data = compress_png_with_pngquant(data, pngquant_path, name) if compress: - data = brotli.compress(data) + data = compressor.compress(data) insert_chunked_content(conn, db_path, language_id, content_type_id, 0, data, chunked_log) return True @@ -544,38 +700,49 @@ def main(): chunked_log = [] - for page in pages: - path = f"{page['id']}.html" - blob = brotli.compress(json.dumps(page, separators=(",", ":"), ensure_ascii=False).encode("utf-8")) - insert_chunked_content(conn, path, language_id, page_content_type_id, page_template_id, blob, chunked_log) - + # Serialized once and reused both as this run's dictionary-training + # samples (only spent if CompressionDictionary doesn't exist yet, see + # load_or_create_dictionary) and as the actual bytes to compress + # below, rather than re-running json.dumps for the same page twice. + page_json_bytes = [ + json.dumps(page, separators=(",", ":"), ensure_ascii=False).encode("utf-8") for page in pages + ] # The server (see layout.pebble) parses each Content row's JSON as an # object and hands its top-level fields to Pebble directly as the # model - a bare JSON array wouldn't parse that way at all, so the # tree goes under a "tree" key here, matching nav.peb's top-level # "{% for node in tree %}". - nav_document = {"tree": nav_tree} - nav_blob = brotli.compress(json.dumps(nav_document, separators=(",", ":"), ensure_ascii=False).encode("utf-8")) - insert_chunked_content(conn, NAV_CONTENT_PATH, language_id, page_content_type_id, nav_template_id, nav_blob, - chunked_log) - - content_type_cache = {} - images_inserted = 0 - with zipfile.ZipFile(args.images_zip) as zf: - for name in image_names: - db_path = f"{IMAGES_DB_PATH_PREFIX}{name}" - if insert_file(conn, zf.read(name), name, db_path, language_id, content_type_cache, chunked_log, - pngquant_path): - images_inserted += 1 - - assets_inserted = 0 - for asset_path in sorted(assets_dir.iterdir()): - if not asset_path.is_file(): - continue - db_path = f"assets/{asset_path.name}" - if insert_file(conn, asset_path.read_bytes(), asset_path.name, db_path, language_id, content_type_cache, - chunked_log, pngquant_path): - assets_inserted += 1 + nav_json_bytes = json.dumps({"tree": nav_tree}, separators=(",", ":"), ensure_ascii=False).encode("utf-8") + + dictionary_data = load_or_create_dictionary(conn, page_json_bytes + [nav_json_bytes]) + with DictionaryCompressor(dictionary_data) as compressor: + for page, json_bytes in zip(pages, page_json_bytes): + path = f"{page['id']}.html" + blob = compressor.compress(json_bytes) + insert_chunked_content(conn, path, language_id, page_content_type_id, page_template_id, blob, + chunked_log) + + nav_blob = compressor.compress(nav_json_bytes) + insert_chunked_content(conn, NAV_CONTENT_PATH, language_id, page_content_type_id, nav_template_id, + nav_blob, chunked_log) + + content_type_cache = {} + images_inserted = 0 + with zipfile.ZipFile(args.images_zip) as zf: + for name in image_names: + db_path = f"{IMAGES_DB_PATH_PREFIX}{name}" + if insert_file(conn, zf.read(name), name, db_path, language_id, content_type_cache, chunked_log, + pngquant_path, compressor): + images_inserted += 1 + + assets_inserted = 0 + for asset_path in sorted(assets_dir.iterdir()): + if not asset_path.is_file(): + continue + db_path = f"assets/{asset_path.name}" + if insert_file(conn, asset_path.read_bytes(), asset_path.name, db_path, language_id, + content_type_cache, chunked_log, pngquant_path, compressor): + assets_inserted += 1 conn.commit() except Exception: diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py new file mode 100644 index 00000000..132220a5 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +""" +renumber_misnumbered_fragments.py + +One-time, idempotent repair for ADFA-5171: some chunked Content rows number +their continuation fragments "-2", "-3", ... with no "-1" +at all. WebServer.kt's reassembly loop always probes "-1" first, so +for these rows it finds nothing and stops after the base CHUNK_SIZE-byte +row - silently truncating (ContentTypes.compression = 'none') or failing to +decompress (compression = 'brotli', since the truncated stream is missing +its tail). + +Every writer in this tool (insert_chunked_content, used by populate_db.py, +insert_optimized_media.py, and migrate_content_to_dictionary_brotli.py) has +always numbered fragments starting at "-1" - none of them produced this, so +it predates this pipeline: inherited data, not something today's code +writes. This script repairs existing databases that still carry it. + +Detects every base row whose content is exactly CHUNK_SIZE bytes, that +isn't itself a fragment of some other chain, and whose own fragment chain +(found by LIKE-querying "-%" and sorting on the numeric suffix, not +by constructed path) doesn't start at 1. A chain with a gap in its +suffixes (a real missing chunk, a different failure than this one) is left +alone and reported rather than guessed at. Renumbers matching chains to a +contiguous "-1", "-2", ... run, lowest original suffix first, so each +rename's target path is always the one just vacated by the previous rename +in the same chain (see renumber_chain). Content bytes are never touched - +only paths move - so this is safe regardless of a row's compression. + +A base row that's exactly CHUNK_SIZE with no continuation fragments at all +is left alone: that's a file that is genuinely exactly 1,048,576 bytes, not +a truncated chain, and WebServer.kt already serves it correctly. + +Idempotent: a chain renumbered by this script starts at -1 afterward, so a +second run finds nothing left to fix. + +Usage: + python3 renumber_misnumbered_fragments.py +""" +import re +import sqlite3 +import sys +from pathlib import Path + +from populate_db import CHUNK_SIZE, backup_database + +FRAGMENT_SUFFIX_RE = re.compile(r"^(.*)-(\d+)$") + + +def find_fragment_paths(conn) -> set: + """Every Content.path that is itself a "-" continuation + fragment of some other row in this table - lets the scan below skip a + fragment that would otherwise also look like a candidate base of its + own (fragments are never themselves further chunked).""" + all_paths = {row[0] for row in conn.execute("SELECT path FROM Content")} + fragments = set() + for path in all_paths: + m = FRAGMENT_SUFFIX_RE.match(path) + if m and m.group(1) in all_paths: + fragments.add(path) + return fragments + + +def chain_fragments(conn, base_path: str) -> list: + """Every "-" row present, as (n, path) sorted by n - found + by LIKE query and parsed suffix, not by constructed path, so it doesn't + matter what N the chain actually starts at or whether it has gaps.""" + rows = conn.execute("SELECT path FROM Content WHERE path LIKE ?", (f"{base_path}-%",)).fetchall() + fragments = [] + for (path,) in rows: + m = FRAGMENT_SUFFIX_RE.match(path) + if m and m.group(1) == base_path: + fragments.append((int(m.group(2)), path)) + fragments.sort(key=lambda item: item[0]) + return fragments + + +def is_contiguous_from_one(fragments: list) -> bool: + return [n for n, _path in fragments] == list(range(1, len(fragments) + 1)) + + +def renumber_chain(conn, base_path: str, fragments: list) -> None: + """Renumbers `fragments` (n, path), sorted ascending by n, to a + contiguous "-1", "-2", ... run. Processed lowest-n first: each target + "-" is either untouched already or was the original path + of the fragment just renamed in the previous iteration, so it's always + free by the time this claims it.""" + for i, (_n, path) in enumerate(fragments, start=1): + new_path = f"{base_path}-{i}" + if path != new_path: + conn.execute("UPDATE Content SET path = ? WHERE path = ?", (new_path, path)) + + +def find_chains(conn, fragment_paths: set) -> tuple: + """Returns (misnumbered, gapped): base paths whose content is exactly + CHUNK_SIZE bytes and aren't themselves a fragment of another chain, + split by whether their fragment chain (if any) is a contiguous run not + starting at 1 (misnumbered - safe to repair) or has an actual gap + (gapped - a real missing chunk, left alone and reported instead of + guessed at).""" + candidates = conn.execute("SELECT path FROM Content WHERE length(content) = ?", (CHUNK_SIZE,)).fetchall() + misnumbered = [] + gapped = [] + for (path,) in candidates: + if path in fragment_paths: + continue + fragments = chain_fragments(conn, path) + if not fragments or fragments[0][0] == 1: + continue + suffixes = [n for n, _path in fragments] + if suffixes == list(range(suffixes[0], suffixes[0] + len(suffixes))): + misnumbered.append((path, fragments)) + else: + gapped.append((path, fragments)) + return misnumbered, gapped + + +def repair(conn) -> dict: + fragment_paths = find_fragment_paths(conn) + misnumbered, gapped = find_chains(conn, fragment_paths) + for base_path, fragments in gapped: + suffixes = [n for n, _path in fragments] + print(f"warning: {base_path!r} has a gapped fragment chain (suffixes {suffixes}); left untouched", + file=sys.stderr) + for base_path, fragments in misnumbered: + renumber_chain(conn, base_path, fragments) + return { + "chains_renumbered": len(misnumbered), + "fragments_moved": sum(len(f) for _, f in misnumbered), + "chains_gapped": len(gapped), + } + + +def main() -> None: + if len(sys.argv) != 2: + print(f"Usage: {sys.argv[0]} ", file=sys.stderr) + sys.exit(1) + db_path = Path(sys.argv[1]) + if not db_path.is_file(): + print(f"error: {db_path} does not exist", file=sys.stderr) + sys.exit(1) + + print(f"Backing up {db_path}...", file=sys.stderr) + backup_path = backup_database(db_path) + print(f"Backup written to {backup_path}", file=sys.stderr) + + conn = sqlite3.connect(db_path) + try: + conn.execute("BEGIN") + stats = repair(conn) + conn.commit() + except Exception: + conn.rollback() + raise + finally: + conn.close() + + print("Vacuuming database to reclaim freed space...", file=sys.stderr) + vacuum_conn = sqlite3.connect(db_path) + try: + vacuum_conn.execute("VACUUM") + finally: + vacuum_conn.close() + + print( + f"Renumbered {stats['chains_renumbered']} chain(s), moved {stats['fragments_moved']} fragment row(s). " + f"{stats['chains_gapped']} chain(s) had a real gap and were left untouched." + ) + + +if __name__ == "__main__": + main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_migrate_content_to_dictionary_brotli.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_migrate_content_to_dictionary_brotli.py new file mode 100644 index 00000000..4c1bfd4f --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_migrate_content_to_dictionary_brotli.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +"""Tests for migrate_content_to_dictionary_brotli.py (ADFA-5153). + +Run directly: python3 test_migrate_content_to_dictionary_brotli.py +""" +import random +import sqlite3 +import tempfile +import unittest +from pathlib import Path + +import brotli + +from migrate_content_to_dictionary_brotli import migrate +from populate_db import CHUNK_SIZE, DictionaryCompressor, load_dictionary + +SCHEMA_SQL = """ +CREATE TABLE Languages (id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT NOT NULL UNIQUE); +CREATE TABLE ContentTypes (id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT NOT NULL UNIQUE, compression TEXT NOT NULL); +CREATE TABLE Content ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + path TEXT NOT NULL, + languageID INTEGER NOT NULL, + content BLOB NOT NULL, + contentTypeID INTEGER NOT NULL, + templateId INTEGER, + UNIQUE(path) +); +""" + +WORDS = [ + "kotlin", "class", "fun", "val", "var", "override", "interface", "object", "companion", + "sidebar", "nav", "template", "docs-sidebar", "toc-element", "page.peb", "Content-Type", +] + + +def make_text(word_count: int, seed: int) -> bytes: + rng = random.Random(seed) + return (" ".join(rng.choice(WORDS) for _ in range(word_count))).encode("utf-8") + + +def insert_plain_chunked(conn, path, language_id, content_type_id, template_id, plain_bytes): + """Mimics populate_db.py's insert_chunked_content, but with plain + (no-dictionary) Brotli - i.e. exactly what every pre-ADFA-5153 pipeline + actually wrote.""" + compressed = brotli.compress(plain_bytes) + conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + (path, language_id, compressed[:CHUNK_SIZE], content_type_id, template_id), + ) + offset = CHUNK_SIZE + n = 1 + while offset < len(compressed): + conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + (f"{path}-{n}", language_id, compressed[offset:offset + CHUNK_SIZE], content_type_id, template_id), + ) + offset += CHUNK_SIZE + n += 1 + return compressed + + +def reassemble(conn, path, first_content): + if len(first_content) < CHUNK_SIZE: + return first_content + parts = [first_content] + n = 1 + while True: + row = conn.execute("SELECT content FROM Content WHERE path = ?", (f"{path}-{n}",)).fetchone() + if row is None: + break + parts.append(row[0]) + if len(row[0]) < CHUNK_SIZE: + break + n += 1 + return b"".join(parts) + + +class MigrateContentToDictionaryBrotliTest(unittest.TestCase): + def setUp(self): + # A real file, not :memory: - migrate() parallelizes the read+compress + # phase across worker threads, each opening its own read-only + # connection to db_path, which an in-memory database has no path for + # (and can't share across connections at all). + fd, path = tempfile.mkstemp(suffix=".db") + Path(path).unlink(missing_ok=True) + self.db_path = Path(path) + self.conn = sqlite3.connect(self.db_path) + self.conn.executescript(SCHEMA_SQL) + self.conn.execute("INSERT INTO Languages (value) VALUES ('en-US')") + self.conn.execute("INSERT INTO ContentTypes (value, compression) VALUES ('text/html', 'brotli')") + self.conn.execute("INSERT INTO ContentTypes (value, compression) VALUES ('image/png', 'none')") + self.language_id = 1 + self.html_type_id = 1 + self.png_type_id = 2 + self.conn.commit() + + def tearDown(self): + self.conn.close() + self.db_path.unlink(missing_ok=True) + + def test_migrates_plain_brotli_rows_preserving_content(self): + originals = {} + for i in range(20): + plain = make_text(200, seed=i) + insert_plain_chunked(self.conn, f"k/html/page{i}.html", self.language_id, self.html_type_id, 5, plain) + originals[f"k/html/page{i}.html"] = plain + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + ("assets/logo.png", self.language_id, b"\x89PNG-not-really-compressed", self.png_type_id, 0), + ) + self.conn.commit() + + stats = migrate(self.conn, self.db_path, sample_size=20, dict_size=16384) + self.assertEqual(stats["scanned"], 20) + self.assertEqual(stats["migrated"], 20) + self.assertEqual(stats["already_migrated"], 0) + + dictionary_data = load_dictionary(self.conn) + with DictionaryCompressor(dictionary_data) as compressor: + for path, plain in originals.items(): + row = self.conn.execute("SELECT content, templateId FROM Content WHERE path = ?", (path,)).fetchone() + first_content, template_id = row + full = reassemble(self.conn, path, first_content) + self.assertEqual(compressor.decompress(full), plain) + self.assertEqual(template_id, 5) + + # untouched: not a 'brotli' content type + png_row = self.conn.execute("SELECT content FROM Content WHERE path = 'assets/logo.png'").fetchone() + self.assertEqual(png_row[0], b"\x89PNG-not-really-compressed") + + def test_preserves_chunked_rows_across_the_1mb_boundary(self): + # A handful of small filler rows so the dictionary trainer has + # more than one sample to work with (zstd's trainer refuses "too + # few samples" on just one row) - a realistic database always has + # many rows, this test's chunked row just happens to be one of them. + for i in range(20): + insert_plain_chunked(self.conn, f"k/html/filler{i}.html", self.language_id, self.html_type_id, 0, + make_text(150, seed=i)) + + # 1.2 MB of high-entropy (incompressible) bytes, so its plain-Brotli + # form still lands comfortably over CHUNK_SIZE - low-entropy text + # (e.g. make_text's small vocabulary) compresses far too well at any + # realistic size to reliably cross that boundary. Exercises the + # multi-row fragment path on both read (reassemble) and write + # (re-chunk) sides. + plain = random.Random(777).randbytes(int(CHUNK_SIZE * 1.2)) + insert_plain_chunked(self.conn, "k/html/big.html", self.language_id, self.html_type_id, 5, plain) + self.conn.commit() + + # Confirm the fixture actually produced a chunked row before relying on it + fragment_exists = self.conn.execute( + "SELECT 1 FROM Content WHERE path = 'k/html/big.html-1'" + ).fetchone() + self.assertIsNotNone(fragment_exists, "test fixture did not produce a chunked row; adjust its size") + + stats = migrate(self.conn, self.db_path, sample_size=21, dict_size=16384) + self.assertEqual(stats["migrated"], 21) + + dictionary_data = load_dictionary(self.conn) + first_content = self.conn.execute( + "SELECT content FROM Content WHERE path = 'k/html/big.html'" + ).fetchone()[0] + with DictionaryCompressor(dictionary_data) as compressor: + full = reassemble(self.conn, "k/html/big.html", first_content) + self.assertEqual(compressor.decompress(full), plain) + + def test_idempotent_second_run_is_a_no_op(self): + originals = {} + for i in range(15): + plain = make_text(150, seed=100 + i) + insert_plain_chunked(self.conn, f"k/html/p{i}.html", self.language_id, self.html_type_id, 0, plain) + originals[f"k/html/p{i}.html"] = plain + self.conn.commit() + + first_stats = migrate(self.conn, self.db_path, sample_size=15, dict_size=16384) + self.assertEqual(first_stats["migrated"], 15) + dictionary_after_first_run = load_dictionary(self.conn) + + snapshot = { + path: self.conn.execute("SELECT content FROM Content WHERE path = ?", (path,)).fetchone()[0] + for path in originals + } + + second_stats = migrate(self.conn, self.db_path, sample_size=15, dict_size=16384) + self.assertEqual(second_stats["migrated"], 0) + self.assertEqual(second_stats["already_migrated"], 15) + + # dictionary must not have been retrained + self.assertEqual(load_dictionary(self.conn), dictionary_after_first_run) + # and no row's bytes changed on the no-op second pass + for path, before in snapshot.items(): + after = self.conn.execute("SELECT content FROM Content WHERE path = ?", (path,)).fetchone()[0] + self.assertEqual(before, after) + + +if __name__ == "__main__": + unittest.main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_populate_db_dictionary.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_populate_db_dictionary.py new file mode 100644 index 00000000..79a96449 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_populate_db_dictionary.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Tests for populate_db.py's shared-dictionary Brotli compression (ADFA-5153). + +Run directly: python3 test_populate_db_dictionary.py +""" +import random +import sqlite3 +import unittest + +from populate_db import DictionaryCompressor, load_dictionary, load_or_create_dictionary, train_dictionary + +WORDS = [ + "kotlin", "class", "fun", "val", "var", "override", "interface", "object", "companion", + "sidebar", "nav", "template", "docs-sidebar", "toc-element", "page.peb", "Content-Type", +] + + +def make_samples(count: int, seed: int = 1) -> list: + rng = random.Random(seed) + return [ + (" ".join(rng.choice(WORDS) for _ in range(150))).encode("utf-8") + for _ in range(count) + ] + + +class TrainDictionaryTest(unittest.TestCase): + def test_produces_nonempty_dictionary(self): + dictionary_data = train_dictionary(make_samples(120)) + self.assertGreater(len(dictionary_data), 0) + + +class DictionaryCompressorTest(unittest.TestCase): + def setUp(self): + self.dictionary_data = train_dictionary(make_samples(120)) + + def test_round_trip(self): + payload = make_samples(1)[0] + with DictionaryCompressor(self.dictionary_data) as compressor: + compressed = compressor.compress(payload) + self.assertNotEqual(compressed, payload) + self.assertEqual(compressor.decompress(compressed), payload) + + def test_compresses_smaller_than_plain_brotli_for_repetitive_corpus(self): + # The whole point of a shared dictionary: content similar to the + # training samples should compress smaller with the dictionary than + # without one. + import brotli + payload = make_samples(1)[0] + with DictionaryCompressor(self.dictionary_data) as compressor: + with_dict = compressor.compress(payload) + without_dict = brotli.compress(payload) + self.assertLess(len(with_dict), len(without_dict)) + + def test_wrong_dictionary_silently_produces_different_bytes(self): + # A mismatched dictionary is NOT guaranteed to fail loudly - it can + # decode "successfully" to silently wrong bytes instead (verified + # empirically: two dictionaries trained on similar-vocabulary + # samples decoded without error but produced garbled output). This + # is exactly why load_or_create_dictionary must never retrain over + # an already-stored dictionary: there is no reliable runtime check + # that would catch the mismatch after the fact. + other_dictionary_data = train_dictionary(make_samples(120, seed=99)) + payload = make_samples(1)[0] + with DictionaryCompressor(self.dictionary_data) as compressor: + compressed = compressor.compress(payload) + with DictionaryCompressor(other_dictionary_data) as wrong_compressor: + result = wrong_compressor.decompress(compressed) + self.assertNotEqual(result, payload) + + def test_no_dictionary_stream_fails_to_decode_with_dictionary_attached(self): + import brotli + payload = make_samples(1)[0] + plain_compressed = brotli.compress(payload) + with DictionaryCompressor(self.dictionary_data) as compressor: + with self.assertRaises(RuntimeError): + compressor.decompress(plain_compressed) + + +class LoadOrCreateDictionaryTest(unittest.TestCase): + def setUp(self): + self.conn = sqlite3.connect(":memory:") + + def tearDown(self): + self.conn.close() + + def test_first_call_trains_and_stores(self): + dictionary_data = load_or_create_dictionary(self.conn, make_samples(120)) + self.assertGreater(len(dictionary_data), 0) + row = self.conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone() + self.assertEqual(row[0], dictionary_data) + + def test_second_call_reuses_stored_dictionary_without_retraining(self): + first = load_or_create_dictionary(self.conn, make_samples(120, seed=1)) + second = load_or_create_dictionary(self.conn, make_samples(120, seed=2)) + self.assertEqual(first, second) + + def test_content_survives_a_reused_dictionary_across_separate_connections(self): + # Mirrors the real cross-repo split: populate_db.py trains/stores the + # dictionary once; a later run (or a different process entirely, + # like WebServer.kt) must be able to decode against the same bytes + # loaded back from the database. + dictionary_data = load_or_create_dictionary(self.conn, make_samples(120)) + payload = make_samples(1)[0] + with DictionaryCompressor(dictionary_data) as compressor: + compressed = compressor.compress(payload) + + reloaded = load_dictionary(self.conn) + self.assertEqual(reloaded, dictionary_data) + with DictionaryCompressor(reloaded) as compressor: + self.assertEqual(compressor.decompress(compressed), payload) + + +class LoadDictionaryTest(unittest.TestCase): + def test_raises_when_missing(self): + conn = sqlite3.connect(":memory:") + try: + with self.assertRaises(RuntimeError): + load_dictionary(conn) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_renumber_misnumbered_fragments.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_renumber_misnumbered_fragments.py new file mode 100644 index 00000000..bc742c32 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_renumber_misnumbered_fragments.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +"""Tests for renumber_misnumbered_fragments.py (ADFA-5171). + +Run directly: python3 test_renumber_misnumbered_fragments.py +""" +import sqlite3 +import tempfile +import unittest +from pathlib import Path + +from populate_db import CHUNK_SIZE +from renumber_misnumbered_fragments import repair + +SCHEMA_SQL = """ +CREATE TABLE Languages (id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT NOT NULL UNIQUE); +CREATE TABLE ContentTypes (id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT NOT NULL UNIQUE, compression TEXT NOT NULL); +CREATE TABLE Content ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + path TEXT NOT NULL, + languageID INTEGER NOT NULL, + content BLOB NOT NULL, + contentTypeID INTEGER NOT NULL, + templateId INTEGER, + UNIQUE(path) +); +""" + + +def chunk_bytes(n: int, fill: bytes) -> bytes: + return (fill * (n // len(fill) + 1))[:n] + + +class RenumberMisnumberedFragmentsTest(unittest.TestCase): + def setUp(self): + fd, path = tempfile.mkstemp(suffix=".db") + Path(path).unlink(missing_ok=True) + self.db_path = Path(path) + self.conn = sqlite3.connect(self.db_path) + self.conn.executescript(SCHEMA_SQL) + self.conn.execute("INSERT INTO Languages (value) VALUES ('en-US')") + self.conn.execute("INSERT INTO ContentTypes (value, compression) VALUES ('image/gif', 'none')") + self.conn.commit() + + def tearDown(self): + self.conn.close() + self.db_path.unlink(missing_ok=True) + + def insert(self, path: str, content: bytes): + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID) VALUES (?, 1, ?, 1)", + (path, content), + ) + + def all_paths(self) -> set: + return {row[0] for row in self.conn.execute("SELECT path FROM Content")} + + def content_at(self, path: str) -> bytes: + return self.conn.execute("SELECT content FROM Content WHERE path = ?", (path,)).fetchone()[0] + + def test_renumbers_chain_starting_at_minus_2(self): + base = "a/devsite/media/size-range.gif" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-2", chunk_bytes(CHUNK_SIZE, b"B")) + self.insert(f"{base}-3", chunk_bytes(CHUNK_SIZE, b"C")) + self.insert(f"{base}-4", chunk_bytes(CHUNK_SIZE, b"D")) + self.insert(f"{base}-5", b"E" * 100) + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 1) + self.assertEqual(stats["fragments_moved"], 4) + self.assertEqual(stats["chains_gapped"], 0) + self.assertEqual( + self.all_paths(), + {base, f"{base}-1", f"{base}-2", f"{base}-3", f"{base}-4"}, + ) + self.assertEqual(self.content_at(f"{base}-1"), chunk_bytes(CHUNK_SIZE, b"B")) + self.assertEqual(self.content_at(f"{base}-2"), chunk_bytes(CHUNK_SIZE, b"C")) + self.assertEqual(self.content_at(f"{base}-3"), chunk_bytes(CHUNK_SIZE, b"D")) + self.assertEqual(self.content_at(f"{base}-4"), b"E" * 100) + + def test_single_orphaned_continuation(self): + base = "j/html/api/index-all.html" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-2", b"tail" * 10) + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 1) + self.assertEqual(stats["fragments_moved"], 1) + self.assertEqual(self.all_paths(), {base, f"{base}-1"}) + self.assertEqual(self.content_at(f"{base}-1"), b"tail" * 10) + + def test_correctly_numbered_chain_untouched(self): + base = "k/html/already-fine.html" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-1", chunk_bytes(CHUNK_SIZE, b"B")) + self.insert(f"{base}-2", b"tail") + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 0) + self.assertEqual(stats["fragments_moved"], 0) + self.assertEqual(self.all_paths(), {base, f"{base}-1", f"{base}-2"}) + + def test_idempotent_second_run(self): + base = "a/devsite/media/size-range.gif" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-2", b"tail") + self.conn.commit() + + repair(self.conn) + self.conn.commit() + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 0) + self.assertEqual(stats["fragments_moved"], 0) + + def test_exact_size_file_with_no_continuation_left_alone(self): + path = "k/html/exactly-one-mb.bin" + self.insert(path, chunk_bytes(CHUNK_SIZE, b"A")) + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 0) + self.assertEqual(stats["chains_gapped"], 0) + self.assertEqual(self.all_paths(), {path}) + + def test_chain_with_real_gap_reported_and_left_untouched(self): + base = "k/html/actually-missing-a-chunk.html" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-2", chunk_bytes(CHUNK_SIZE, b"B")) + self.insert(f"{base}-4", b"tail") # -3 is genuinely missing + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 0) + self.assertEqual(stats["chains_gapped"], 1) + self.assertEqual(self.all_paths(), {base, f"{base}-2", f"{base}-4"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/docdb-studio/docdb_studio.py b/docdb-studio/docdb_studio.py index 30f14730..b6172c45 100644 --- a/docdb-studio/docdb_studio.py +++ b/docdb-studio/docdb_studio.py @@ -9,8 +9,11 @@ import mimetypes import os import platform as _platform +import shutil import sqlite3 +import subprocess import sys +import tempfile import threading import time import unicodedata @@ -334,7 +337,7 @@ def get_html_anchors_for_path(db_path: Path, base_path: str) -> list[str]: full = b"".join(parts) if compression == "brotli": try: - full = brotli.decompress(full) + full = decompress_brotli(full, db_path) except brotli.error: return [] return extract_html_anchors(full) @@ -708,7 +711,7 @@ def fetch_content_for_path( full = b"".join(parts) if compression == "brotli": try: - full = brotli.decompress(full) + full = decompress_brotli(full, db_path) except brotli.error: return None return full, mime @@ -1220,11 +1223,99 @@ def get_languages(db_path: Path) -> list[tuple[int, str]]: return cur.fetchall() -def compress_for_storage(data: bytes, compression: str) -> bytes: - """Apply compression policy. 'brotli' encodes; anything else passes through unchanged.""" - if compression == "brotli": +# db_path -> dictionary bytes, or None if that database has no CompressionDictionary +# (an older/test database predating ADFA-5153). docdb-studio never creates or retrains +# a dictionary itself, so a cached value -- present or None -- can't go stale mid-session. +_dictionary_cache: dict[Path, bytes | None] = {} +# db_path -> temp file holding that database's dictionary bytes, for the brotli CLI's -D +# flag. Written once per db_path and reused, rather than rewriting the same bytes to disk +# on every compress/decompress call. +_dictionary_temp_paths: dict[Path, Path] = {} + + +def _find_brotli_cli() -> str: + path = shutil.which("brotli") + if path is None: + raise RuntimeError("brotli CLI not found on PATH; install it and retry") + return path + + +def get_compression_dictionary(db_path: Path) -> bytes | None: + """Returns db_path's CompressionDictionary bytes (see ADFA-5153), or None if it + doesn't have one yet.""" + if db_path in _dictionary_cache: + return _dictionary_cache[db_path] + dictionary_data: bytes | None = None + try: + with sqlite3.connect(db_path) as conn: + table_row = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'CompressionDictionary'" + ).fetchone() + if table_row is not None: + data_row = conn.execute( + "SELECT data FROM CompressionDictionary WHERE id = 1" + ).fetchone() + if data_row is not None: + dictionary_data = data_row[0] + except sqlite3.OperationalError: + dictionary_data = None + _dictionary_cache[db_path] = dictionary_data + return dictionary_data + + +def _dictionary_temp_path(db_path: Path, dictionary_data: bytes) -> Path: + cached = _dictionary_temp_paths.get(db_path) + if cached is not None and cached.exists(): + return cached + fd, name = tempfile.mkstemp(prefix="docdb-studio-brotli-dict-") + path = Path(name) + with os.fdopen(fd, "wb") as f: + f.write(dictionary_data) + _dictionary_temp_paths[db_path] = path + atexit.register(lambda: path.unlink(missing_ok=True)) + return path + + +def compress_for_storage(data: bytes, compression: str, db_path: Path) -> bytes: + """Apply compression policy. 'brotli' encodes -- against db_path's shared dictionary + if it has one (see ADFA-5153), otherwise plain, matching a database that predates + that migration. Anything else passes through unchanged.""" + if compression != "brotli": + return data + dictionary_data = get_compression_dictionary(db_path) + if dictionary_data is None: return brotli.compress(data) - return data + dict_path = _dictionary_temp_path(db_path, dictionary_data) + result = subprocess.run( + [_find_brotli_cli(), "-D", str(dict_path), "-c"], + input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"brotli failed: {result.stderr.decode(errors='replace').strip()}") + return result.stdout + + +def decompress_brotli(data: bytes, db_path: Path) -> bytes: + """Inverse of compress_for_storage's 'brotli' branch -- decodes against db_path's + shared dictionary if it has one, otherwise plain. A row compressed against the + dictionary is only decodable with that same dictionary (verified empirically to + fail, or silently produce different bytes, otherwise -- see ADFA-5153), so this + must agree with whichever path originally compressed the row. + + Raises brotli.error on failure either way, matching plain brotli.decompress's own + exception type, so existing `except brotli.error:` call sites don't need to change. + """ + dictionary_data = get_compression_dictionary(db_path) + if dictionary_data is None: + return brotli.decompress(data) + dict_path = _dictionary_temp_path(db_path, dictionary_data) + result = subprocess.run( + [_find_brotli_cli(), "-d", "-D", str(dict_path), "-c"], + input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + if result.returncode != 0: + raise brotli.error(result.stderr.decode(errors="replace").strip()) + return result.stdout def fragment_blob(blob: bytes, chunk_size: int = CONTENT_CHUNK_SIZE) -> list[bytes]: @@ -1430,7 +1521,7 @@ def _report(phase: str, current: int, total: int) -> None: _report("add", adds_done, add_total) continue - stored = compress_for_storage(data, item.compression) + stored = compress_for_storage(data, item.compression, db_path) chunks = fragment_blob(stored) paths = target_paths(item.base_path, len(chunks)) diff --git a/docdb-studio/tests/test_compression_dictionary.py b/docdb-studio/tests/test_compression_dictionary.py new file mode 100644 index 00000000..6a5145f7 --- /dev/null +++ b/docdb-studio/tests/test_compression_dictionary.py @@ -0,0 +1,199 @@ +"""Tests for shared-dictionary Brotli compression in docdb_studio.py (ADFA-5153).""" + +import random +import shutil +import sqlite3 +import subprocess +import tempfile +from pathlib import Path + +import brotli +import pytest + +import docdb_studio + +get_compression_dictionary = docdb_studio.get_compression_dictionary +compress_for_storage = docdb_studio.compress_for_storage +decompress_brotli = docdb_studio.decompress_brotli +get_html_anchors_for_path = docdb_studio.get_html_anchors_for_path +fetch_content_for_path = docdb_studio.fetch_content_for_path + +WORDS = [ + "kotlin", "class", "fun", "val", "var", "override", "interface", "object", "companion", + "sidebar", "nav", "template", "docs-sidebar", "toc-element", "page.peb", "Content-Type", +] + + +def _make_text(word_count: int, seed: int) -> bytes: + rng = random.Random(seed) + return (" ".join(rng.choice(WORDS) for _ in range(word_count))).encode("utf-8") + + +def _train_dictionary(samples: list, dict_size: int = 16384) -> bytes: + """Test-only dictionary trainer (docdb_studio.py never trains one itself -- see + its own module docstring/AGENTS.md: it only ever reads a dictionary another tool + already produced).""" + zstd_path = shutil.which("zstd") + assert zstd_path is not None, "zstd must be on PATH to run this test" + work_dir = Path(tempfile.mkdtemp(prefix="docdb-studio-test-dict-")) + try: + sample_paths = [] + for i, sample in enumerate(samples): + p = work_dir / f"sample_{i:03}.bin" + p.write_bytes(sample) + sample_paths.append(str(p)) + dict_path = work_dir / "dictionary.bin" + subprocess.run( + [zstd_path, "--train-fastcover", f"--maxdict={dict_size}", "-o", str(dict_path), *sample_paths], + check=True, capture_output=True, + ) + return dict_path.read_bytes() + finally: + shutil.rmtree(work_dir, ignore_errors=True) + + +def _make_db(with_dictionary: bytes | None = None) -> Path: + """Temp DB with Content/ContentTypes/Languages, and CompressionDictionary + populated iff with_dictionary is given.""" + fd, path = tempfile.mkstemp(suffix=".db") + Path(path).unlink(missing_ok=True) + p = Path(path) + with sqlite3.connect(p) as conn: + conn.executescript( + """ + CREATE TABLE Languages (id INTEGER PRIMARY KEY, value TEXT NOT NULL UNIQUE); + CREATE TABLE ContentTypes ( + id INTEGER PRIMARY KEY, + value TEXT NOT NULL UNIQUE, + compression TEXT NOT NULL + ); + CREATE TABLE "Content" ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + path TEXT NOT NULL, + languageID INTEGER NOT NULL, + content BLOB NOT NULL, + contentTypeID INTEGER NOT NULL, + UNIQUE(path) + ); + INSERT INTO Languages (id, value) VALUES (1, 'en'); + INSERT INTO ContentTypes (id, value, compression) VALUES (1, 'text/html', 'brotli'); + """ + ) + if with_dictionary is not None: + conn.execute( + "CREATE TABLE CompressionDictionary (id INTEGER PRIMARY KEY CHECK (id = 1), data BLOB NOT NULL)" + ) + conn.execute("INSERT INTO CompressionDictionary (id, data) VALUES (1, ?)", (with_dictionary,)) + conn.commit() + return p + + +def _insert_content(db: Path, path: str, blob: bytes, content_type_id: int = 1) -> None: + with sqlite3.connect(db) as conn: + conn.execute( + 'INSERT INTO "Content" (path, languageID, content, contentTypeID) VALUES (?, 1, ?, ?)', + (path, blob, content_type_id), + ) + conn.commit() + + +# ---------- get_compression_dictionary ---------- + + +def test_returns_none_when_table_missing() -> None: + db = _make_db() + try: + assert get_compression_dictionary(db) is None + finally: + db.unlink(missing_ok=True) + + +def test_returns_stored_dictionary_bytes() -> None: + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + try: + assert get_compression_dictionary(db) == dictionary_data + finally: + db.unlink(missing_ok=True) + + +def test_caches_per_db_path() -> None: + db = _make_db() + try: + assert get_compression_dictionary(db) is None + # Mutating the row after the first (cached) lookup must not change the + # cached result -- docdb_studio never expects a dictionary to appear or + # change mid-session, since it never writes one itself. + with sqlite3.connect(db) as conn: + conn.execute( + "CREATE TABLE CompressionDictionary (id INTEGER PRIMARY KEY CHECK (id = 1), data BLOB NOT NULL)" + ) + conn.execute("INSERT INTO CompressionDictionary (id, data) VALUES (1, ?)", (b"late-arriving",)) + conn.commit() + assert get_compression_dictionary(db) is None + finally: + db.unlink(missing_ok=True) + + +# ---------- compress_for_storage / decompress_brotli ---------- + + +def test_round_trip_without_dictionary_matches_plain_brotli() -> None: + db = _make_db() + try: + data = b"hello world " * 500 + compressed = compress_for_storage(data, "brotli", db) + assert brotli.decompress(compressed) == data # plain brotli, no dictionary involved + assert decompress_brotli(compressed, db) == data + finally: + db.unlink(missing_ok=True) + + +def test_round_trip_with_dictionary() -> None: + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + try: + data = _make_text(150, 9999) + compressed = compress_for_storage(data, "brotli", db) + with pytest.raises(brotli.error): + brotli.decompress(compressed) # plain decode of dictionary-compressed data fails + assert decompress_brotli(compressed, db) == data + finally: + db.unlink(missing_ok=True) + + +def test_none_compression_passes_through_unchanged() -> None: + db = _make_db() + try: + data = b"\x00\x01\x02 raw bytes" + assert compress_for_storage(data, "none", db) == data + finally: + db.unlink(missing_ok=True) + + +# ---------- get_html_anchors_for_path / fetch_content_for_path against a real dictionary ---------- + + +def test_get_html_anchors_for_path_decodes_dictionary_compressed_html() -> None: + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + try: + html = b'

Intro

x

' + blob = compress_for_storage(html, "brotli", db) + _insert_content(db, "docs/page.html", blob) + assert get_html_anchors_for_path(db, "docs/page.html") == ["intro", "p1"] + finally: + db.unlink(missing_ok=True) + + +def test_fetch_content_for_path_decodes_dictionary_compressed_content() -> None: + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + try: + html = b"

served via dictionary

" + blob = compress_for_storage(html, "brotli", db) + _insert_content(db, "docs/page.html", blob) + result = fetch_content_for_path(db, "docs/page.html") + assert result == (html, "text/html") + finally: + db.unlink(missing_ok=True) diff --git a/docdb-studio/tests/test_content_import.py b/docdb-studio/tests/test_content_import.py index 092288b2..945188d4 100644 --- a/docdb-studio/tests/test_content_import.py +++ b/docdb-studio/tests/test_content_import.py @@ -122,15 +122,25 @@ def test_target_paths_zero_raises() -> None: def test_compress_for_storage_brotli_round_trip() -> None: - data = b"hello world " * 1000 - compressed = compress_for_storage(data, "brotli") - assert compressed != data - assert brotli.decompress(compressed) == data + # No CompressionDictionary in this fixture -> plain Brotli, same as a database + # that predates ADFA-5153. + db = _make_db_with_content_schema() + try: + data = b"hello world " * 1000 + compressed = compress_for_storage(data, "brotli", db) + assert compressed != data + assert brotli.decompress(compressed) == data + finally: + db.unlink(missing_ok=True) def test_compress_for_storage_none_passthrough() -> None: - data = b"\x00\x01\x02 some bytes" - assert compress_for_storage(data, "none") is data or compress_for_storage(data, "none") == data + db = _make_db_with_content_schema() + try: + data = b"\x00\x01\x02 some bytes" + assert compress_for_storage(data, "none", db) is data or compress_for_storage(data, "none", db) == data + finally: + db.unlink(missing_ok=True) # ---------- mime_for_filename ----------