diff --git a/.github/workflows/build-kotlin-docs.yaml b/.github/workflows/build-kotlin-docs.yaml index fa981350..c6341556 100644 --- a/.github/workflows/build-kotlin-docs.yaml +++ b/.github/workflows/build-kotlin-docs.yaml @@ -143,7 +143,10 @@ jobs: - name: Install system dependencies run: | sudo apt-get update -y - sudo apt-get install -y pngquant unzip zip sqlite3 + # brotli: the CLI, not the Python package. populate_db.py's + # DictionaryCompressor and sync_kdoc_json_to_db.py shell out to it because + # no Python binding exposes a custom dictionary (ADFA-5153). + sudo apt-get install -y pngquant unzip zip sqlite3 brotli - name: Install Python dependencies run: | diff --git a/.github/workflows/docdb-regression-test.yaml b/.github/workflows/docdb-regression-test.yaml index fa5346e0..d9222dfa 100644 --- a/.github/workflows/docdb-regression-test.yaml +++ b/.github/workflows/docdb-regression-test.yaml @@ -98,7 +98,9 @@ jobs: echo "Extracting database from zip file..." # Install unzip if not available - sudo apt-get update -qq && sudo apt-get install -y unzip + # brotli: docdb_studio reads dictionary-compressed Content rows through the + # CLI (ADFA-5153); the downloaded production database is one of those. + sudo apt-get update -qq && sudo apt-get install -y unzip brotli # Extract the zip file if ! unzip -o documentation.zip; then diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/README.md b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/README.md index 8a8bed28..47008e27 100644 --- a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/README.md +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/README.md @@ -13,6 +13,10 @@ and its media straight into a `documentation.db`-schema SQLite database. | [`build_nav.py`](build_nav.py) | Builds `nav.json`/`nav.html` sidebar navigation from `kr.tree`, resolving each `` against `md_to_json.py`'s output. | | [`find_missing_assets.py`](find_missing_assets.py) | QA pass: reports cross-page links, images, and `` targets in the source tree that don't resolve to anything. Reuses `md_to_json.py`'s own resolution logic, so it flags exactly what would end up broken on the rendered site. | | [`populate_db.py`](populate_db.py) | The database path: converts the docs tree the same way `md_to_json.py` does, builds nav the same way `build_nav.py` does, and inserts pages + nav + images + CSS/JS directly into `documentation.db` (replacing everything under `k/html/` and `assets/`). Supports pruning whole `kr.tree` subtrees via `--blacklisted-element-titles`. | +| [`migrate_content_to_dictionary_brotli.py`](migrate_content_to_dictionary_brotli.py) | One-off, resumable: recompresses every `brotli` Content row against the database's shared `CompressionDictionary`, training one first if there is none (ADFA-5153). Covers the rows `populate_db.py` never touches. | +| [`renumber_misnumbered_fragments.py`](renumber_misnumbered_fragments.py) | One-off repair: chunked rows whose continuations start at `-2` (or `-0`) instead of `-1`, which `WebServer.kt` reassembles truncated (ADFA-5171). Moves paths only, never content. | +| [`remint_dictionary.py`](remint_dictionary.py) | One-off, **destructive**: trains a *new* shared dictionary and recompresses every row against it, in one transaction. The only safe way to change a dictionary, since the stored one is otherwise permanent for that database's content. Pair with `verify_remint_dictionary.py` before putting the result in place. | +| [`verify_remint_dictionary.py`](verify_remint_dictionary.py) | Read-only gate for the above: decodes every row out of both databases and requires the plaintexts to match, exiting non-zero otherwise. A mismatched dictionary decodes into wrong bytes without erroring, so this is what makes re-minting safe. | | [`optimize_media.py`](optimize_media.py) | Standalone media optimizer: downscales/recompresses a directory of images (pngquant, Pillow, Scour/cairosvg for SVG) into a mirrored output directory. | | [`insert_optimized_media.py`](insert_optimized_media.py) | Runs `optimize_media.py`'s pipeline over a directory of raw media, then replaces the corresponding `k/html/images/*` rows in an existing database, rewriting any page that referenced a renamed file and deleting anything left unreferenced. | @@ -22,6 +26,7 @@ and its media straight into a `documentation.db`-schema SQLite database. - `pip install markdown-it-py Pillow scour brotli` - `cairosvg` (only needed if an optimized SVG exceeds `--svg-rasterize-threshold`): `pip install cairosvg` - `pngquant` on `PATH` (e.g. `apt install pngquant`) — required by `optimize_media.py`/`insert_optimized_media.py`, and by `populate_db.py` for the images it inserts directly from the Writerside export. +- `brotli` on `PATH` (e.g. `apt install brotli`) — the **command-line tool**, which is a different artifact from the `brotli` Python package listed above. `populate_db.py`, `insert_optimized_media.py`, `migrate_content_to_dictionary_brotli.py` and `remint_dictionary.py` compress against the shared dictionary in `CompressionDictionary` (ADFA-5153), and no Python binding exposes a custom dictionary, so they shell out to this binary. Without it they fail at startup. `populate_db.py` also expects, relative to its own location, and already included in this directory: diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py index 3b08078e..9da4fc8e 100644 --- a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/insert_optimized_media.py @@ -62,15 +62,14 @@ 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, + CHUNK_SIZE, DictionaryCompressor, EXTENSION_TO_CONTENT_TYPE, IMAGES_DB_PATH_PREFIX, IMAGES_URL_PREFIX, + LANGUAGE, PAGE_CONTENT_TYPE, backup_database, fragment_chain, get_content_type, get_id, + insert_chunked_content, load_dictionary, ) WEBP_CONTENT_TYPE = "image/webp" @@ -104,12 +103,22 @@ def delete_content(conn, path: str) -> None: """Deletes a Content row and any chunked continuation fragments for it (see insert_chunked_content/CHUNK_SIZE) - safe to call even if nothing exists yet at that path. 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}-%")) + before any re-insert at the same path. + + Deletes by exact path rather than by a LIKE pattern. `_` is a single- + character wildcard in LIKE and the `-%` suffix does not restrict the tail to + digits, so "DELETE ... WHERE path LIKE '-%'" also removes rows that + merely resemble a continuation - and those are never re-inserted, so the + loss is permanent. populate_db.fragment_chain does the over-matching query + once and re-checks every candidate's suffix, which is what makes the result + exact.""" + conn.execute("DELETE FROM Content WHERE path = ?", (path,)) + for _number, fragment_path in fragment_chain(conn, path): + conn.execute("DELETE FROM Content WHERE path = ?", (fragment_path,)) 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 +134,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 +184,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 +234,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 +255,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 +268,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 +296,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 +306,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 +394,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..129929ff --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/migrate_content_to_dictionary_brotli.py @@ -0,0 +1,409 @@ +#!/usr/bin/env python3 +""" +migrate_content_to_dictionary_brotli.py + +One-time, resumable, 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. Dictionary +compression pays off best when it covers the whole corpus, so this script is +what converts the rows outside populate_db.py's reach. + +Note what it does NOT establish: an invariant that every 'brotli' row in a +shipped database uses the dictionary. That is unachievable by construction - a +plugin installed on-device contributes plain-Brotli rows at any time (see +PluginDocumentationManager/BrotliCompressor in the app). WebServer.kt therefore +tries a dictionary-attached decode and falls back to a plain one, and that +fallback is load-bearing rather than defensive. Any other reader of this +database needs the same fallback. + +Classification, per row, is by *decoding* rather than by assumption, because +"plain decode failed" alone means very little: + + * decodes plainly AND with the dictionary, to identical bytes -> the encoder + never referenced the dictionary (small or already-compressed payloads, + ~0.5% of the real corpus). Nothing to gain; left untouched, so re-runs do + not churn it. This is the case that makes a naive "plain decode succeeded, + so it needs migrating" test re-migrate the same rows on every run. + * decodes plainly only -> not yet migrated. Recompress. + * decodes with the dictionary only -> already migrated. + * decodes neither way -> reported as an ERROR, never counted as success. A + truncated ADFA-5171 chain lands here, and silently counting it as + "already migrated" is exactly how such a row stays plain-Brotli while the + run reports a clean finish. + +Chunked rows are reassembled via populate_db.fragment_chain, which finds a +chain by LIKE plus parsed suffix rather than by probing "-1" - an +ADFA-5171 chain numbered from -2 would otherwise reassemble truncated. Run +renumber_misnumbered_fragments.py first if the database still has those; this +script reports them rather than repairing them. + +Writes are in place: UPDATE on the base row, then the continuation rows are +reconciled by exact path. The base row is never DELETEd and re-INSERTed, +because Content carries AddBook/DeleteBook triggers on paths matching +'%.pdf' - a delete/insert cycle drops the curated Bookshelf entry (title, +description, bookCategoryID) and replaces it with 'CURRENT_TIMESTAMP || id' +under a fresh Content.id. 15 brotli-typed .pdf rows and all 7 Bookshelf rows +are in scope on the real database. + +Dictionary training (only when CompressionDictionary does not exist yet) draws +a sample stratified across doc sets in proportion to their stored bytes, and is +bounded by a plaintext byte budget rather than a row count. Both halves matter, +measured on the real corpus with only the sampling varied: + + first 300 rows by path (all under "a/") 36.2% smaller than plain + 300 rows stratified across doc sets 33.2% <- worse + stratified, 32 MiB plaintext budget 48.3% <- best + first-by-path, same 32 MiB budget 36.4% <- volume alone: nil + +Stratifying at a fixed row count draws quotas from smaller doc sets, so total +material falls and the trainer cannot even fill a 256 KiB dictionary. The +spread is the win; the byte budget is what makes the spread affordable. + +Safety: backs up the database first (VACUUM INTO, same as populate_db.py), +commits in batches (a single transaction spanning the whole run holds a write +lock that the readers below cannot work around under rollback-journal mode), +verifies every recompressed row round-trips before writing it, and VACUUMs +afterward on a separate connection (SQLite refuses VACUUM inside a +transaction). Interrupting it is safe: finished batches stand, and re-running +resumes. + +Performance: all database access happens on the calling thread; the worker pool +only ever receives bytes. Each recompress spawns a `brotli` subprocess, so the +work parallelizes well, and keeping SQLite single-threaded avoids the +"database is locked" failure that worker-owned read connections hit under +journal_mode=delete (documentation.db's actual mode) once the write +transaction's page cache spills. + +Usage: + python3 migrate_content_to_dictionary_brotli.py [--sample-size N] + [--dict-size BYTES] [--max-workers N] [--training-bytes BYTES] [--sample-seed N] +""" +import argparse +import random +import sqlite3 +import sys +import threading +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import brotli + +from populate_db import ( + CHUNK_SIZE, DEFAULT_DICT_SIZE, DictionaryCompressor, backup_database, fragment_chain, + load_or_create_dictionary, +) + +DEFAULT_SAMPLE_SIZE = 300 +# ~128x the 256 KiB dictionary. zstd's cover trainers want roughly two orders of +# magnitude more material than the dictionary they produce; below that they +# return a dictionary smaller than the cap, which measurably compresses worse. +DEFAULT_TRAINING_BYTES = 32 * 1024 * 1024 +DEFAULT_SAMPLE_SEED = 0x5153 +# Rows per write transaction. Small enough that the write lock is never held +# across a long stretch of compression work, large enough that commit overhead +# stays negligible against a q11 recompress. +BATCH_ROWS = 200 + +_thread_local = threading.local() + + +def read_item(conn, path: str) -> bytes: + """The full stored bytes of one logical item: its base row plus every + continuation row in its chain, in suffix order. Suffix-agnostic (see + populate_db.fragment_chain), so an ADFA-5171 chain numbered from -2 + reassembles correctly rather than truncating.""" + row = conn.execute("SELECT content FROM Content WHERE path = ?", (path,)).fetchone() + if row is None: + return b"" + parts = [row[0]] + if len(row[0]) < CHUNK_SIZE: + return parts[0] + for _n, fragment_path in fragment_chain(conn, path): + fragment = conn.execute("SELECT content FROM Content WHERE path = ?", (fragment_path,)).fetchone() + if fragment is not None: + parts.append(fragment[0]) + return b"".join(parts) + + +def write_item(conn, path: str, language_id: int, content_type_id: int, template_id: int, data: bytes) -> None: + """Replaces one item's stored bytes in place: UPDATE on the base row, then + the continuation rows reconciled by exact path (updated, inserted, or + deleted as the new chunk count requires). + + Deliberately never DELETEs the base row: Content's AddBook/DeleteBook + triggers fire on '%.pdf' paths and a delete/insert cycle silently replaces + the curated Bookshelf entry with a timestamp title under a new + Content.id. Continuation paths end in "-", so they never match those + triggers and are safe to delete. Nothing here goes through LIKE, so no + unrelated row can be caught by a `_` wildcard in a path.""" + conn.execute("UPDATE Content SET content = ? WHERE path = ?", (data[:CHUNK_SIZE], path)) + + wanted = {} + for number, offset in enumerate(range(CHUNK_SIZE, len(data), CHUNK_SIZE), start=1): + wanted[f"{path}-{number}"] = data[offset:offset + CHUNK_SIZE] + + existing = {fragment_path for _n, fragment_path in fragment_chain(conn, path)} + for fragment_path, blob in wanted.items(): + if fragment_path in existing: + conn.execute("UPDATE Content SET content = ? WHERE path = ?", (blob, fragment_path)) + else: + conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) " + "VALUES (?, ?, ?, ?, ?)", + (fragment_path, language_id, blob, content_type_id, template_id), + ) + for surplus in sorted(existing - set(wanted)): + conn.execute("DELETE FROM Content WHERE path = ?", (surplus,)) + + +def _thread_compressor(dictionary_data: bytes) -> DictionaryCompressor: + """One DictionaryCompressor per worker thread, reused across every row that + thread processes - creating one per row would re-write the same dictionary + bytes to a fresh temp file on every 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 classify(compressor: DictionaryCompressor, stored: bytes) -> tuple: + """Returns (state, plaintext) for one item's stored bytes, by decoding it + both ways. See the module docstring for why each state means what it does. + state is one of "both", "plain", "dictionary", "undecodable".""" + try: + plain = brotli.decompress(stored) + except brotli.error: + plain = None + try: + via_dictionary = compressor.decompress(stored) + except RuntimeError: + via_dictionary = None + + if plain is not None and via_dictionary is not None and plain == via_dictionary: + return "both", plain + if plain is not None: + return "plain", plain + if via_dictionary is not None: + return "dictionary", via_dictionary + return "undecodable", None + + +def recompress_item(dictionary_data: bytes, path: str, stored: bytes) -> dict: + """Runs on a worker thread: pure bytes in, pure bytes out, no database + access. Returns a result dict the caller writes back (or reports).""" + compressor = _thread_compressor(dictionary_data) + state, plain = classify(compressor, stored) + if state == "undecodable": + return {"path": path, "state": "error", + "detail": "decodes neither plainly nor with the dictionary; " + "run renumber_misnumbered_fragments.py if its chain is numbered from -2"} + if state in ("dictionary", "both"): + return {"path": path, "state": "already"} + + recompressed = compressor.compress(plain) + # A migration that does not round-trip is worse than no migration: the row + # would only fail later, on-device, at read time. + try: + if compressor.decompress(recompressed) != plain: + return {"path": path, "state": "error", "detail": "recompressed bytes do not decode back to the original"} + except RuntimeError as exc: + return {"path": path, "state": "error", "detail": f"recompressed bytes failed to decode: {exc}"} + + return {"path": path, "state": "migrated", "data": recompressed, + "before": len(stored), "after": len(recompressed)} + + +def doc_set(path: str) -> str: + return path.split("/", 1)[0] if "/" in path else "(root)" + + +def collect_training_samples(conn, base_rows: list, sample_size: int, + byte_budget: int = DEFAULT_TRAINING_BYTES, + seed: int = DEFAULT_SAMPLE_SEED, decode=None) -> list: + """Plaintext samples for training a new dictionary, stratified across doc + sets in proportion to their stored bytes and bounded by `byte_budget` of + plaintext. See the module docstring for the measurements behind both + choices. Deterministic for a given seed, because a dictionary is never + retrained once stored - being able to reproduce the training set later is + the only way to explain the bytes you are then stuck with. + + `decode` reads one item's stored bytes back to plaintext, defaulting to plain + Brotli because this only ever runs before CompressionDictionary exists. A + re-mint (see remint_dictionary.py) passes one that decodes against the + outgoing dictionary instead.""" + decode = decode or brotli.decompress + by_set = defaultdict(list) + for row in base_rows: + by_set[doc_set(row[0])].append(row) + + weight = {name: sum(row[4] for row in rows) for name, rows in by_set.items()} + total_weight = sum(weight.values()) or 1 + rng = random.Random(seed) + + drawn = [] + for name, rows in by_set.items(): + shuffled = list(rows) + rng.shuffle(shuffled) + # Oversample per set: the byte budget below is the real limit, and a + # short set should not strand budget that another set could use. + quota = max(1, round(sample_size * weight[name] / total_weight)) + drawn.append((name, shuffled, quota)) + + ordered = [] + for name, shuffled, quota in drawn: + ordered.extend(shuffled[:quota * 3]) + rng.shuffle(ordered) + + samples = [] + used = 0 + for row in ordered: + if used >= byte_budget or len(samples) >= sample_size * 3: + break + try: + plain = decode(read_item(conn, row[0])) + except brotli.error as exc: + print(f"warning: could not decompress {row[0]!r} for training sample: {exc}", file=sys.stderr) + continue + samples.append(plain) + used += len(plain) + + histogram = defaultdict(int) + for row in ordered[:len(samples)]: + histogram[doc_set(row[0])] += 1 + print(f"Training dictionary on {len(samples)} rows, {used / 1048576:.1f} MiB of plaintext " + f"across {len(histogram)} doc set(s): {dict(sorted(histogram.items(), key=lambda kv: -kv[1]))}", + file=sys.stderr) + return samples + + +def load_base_rows(conn) -> list: + """Every brotli-typed base row as (path, language_id, content_type_id, + template_id, stored_len). Blobs are deliberately not selected here - the + real table is ~130 MB of compressed content, and each row's bytes are read + only when its turn comes.""" + rows = conn.execute( + "SELECT C.path, C.languageID, C.contentTypeID, C.templateId, LENGTH(C.content) " + "FROM Content C, ContentTypes CT " + "WHERE C.contentTypeID = CT.id AND CT.compression = 'brotli' " + "ORDER BY C.path" + ).fetchall() + all_paths = {row[0] for row in conn.execute("SELECT path FROM Content")} + base_rows = [] + for row in rows: + prefix, sep, suffix = row[0].rpartition("-") + if sep == "-" and suffix.isdigit() and prefix in all_paths: + continue # a continuation row; handled with its base + base_rows.append(row) + return base_rows + + +def migrate(conn, sample_size: int, dict_size: int, max_workers: int | None = None, + training_bytes: int = DEFAULT_TRAINING_BYTES, sample_seed: int = DEFAULT_SAMPLE_SEED) -> dict: + base_rows = load_base_rows(conn) + + has_dictionary = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'CompressionDictionary'" + ).fetchone() is not None + training_samples = [] if has_dictionary else collect_training_samples( + conn, base_rows, sample_size, training_bytes, sample_seed + ) + dictionary_data = load_or_create_dictionary(conn, training_samples, dict_size) + conn.commit() + + stats = {"scanned": len(base_rows), "migrated": 0, "already": 0, "errors": 0, + "bytes_before": 0, "bytes_after": 0} + problems = [] + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + for start in range(0, len(base_rows), BATCH_ROWS): + batch = base_rows[start:start + BATCH_ROWS] + payloads = [(row, read_item(conn, row[0])) for row in batch] + results = executor.map( + lambda item: recompress_item(dictionary_data, item[0][0], item[1]), payloads + ) + for row, result in zip(batch, results): + if result["state"] == "already": + stats["already"] += 1 + continue + if result["state"] == "error": + stats["errors"] += 1 + problems.append((result["path"], result["detail"])) + continue + stats["migrated"] += 1 + stats["bytes_before"] += result["before"] + stats["bytes_after"] += result["after"] + write_item(conn, row[0], row[1], row[2], row[3], result["data"]) + conn.commit() + + stats["problems"] = problems + 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 aim for when training a dictionary, if none exists yet " + f"(default: {DEFAULT_SAMPLE_SIZE}); --training-bytes is the real limit") + parser.add_argument("--training-bytes", type=int, default=DEFAULT_TRAINING_BYTES, + help=f"Plaintext byte budget for dictionary training " + f"(default: {DEFAULT_TRAINING_BYTES:,})") + parser.add_argument("--sample-seed", type=int, default=DEFAULT_SAMPLE_SEED, + help="Seed for the stratified training sample, so a dictionary's training set " + "stays reproducible (default: %(default)s)") + 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 recompress phase (default: ThreadPoolExecutor's own " + "min(32, cpu_count+4)); database access stays on the calling thread") + 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: + stats = migrate(conn, args.sample_size, args.dict_size, args.max_workers, + args.training_bytes, args.sample_seed) + conn.commit() + 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']}, errors {stats['errors']}." + ) + 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)") + for path, detail in stats["problems"][:20]: + print(f"error: {path}: {detail}", file=sys.stderr) + if len(stats["problems"]) > 20: + print(f"error: ... and {len(stats['problems']) - 20} more", file=sys.stderr) + if stats["errors"]: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/populate_db.py index c338b9dc..ce4b28e3 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,19 @@ chunked is logged by name at the end of the run. """ import argparse +import atexit import json +import re 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 +183,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 +231,151 @@ 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, but the two directions + of that mismatch behave differently, and the difference matters - measured + against the real documentation.db, not assumed: + + * Decoding a dictionary-compressed row with NO dictionary is loud: 398 + of 400 sampled rows raised outright, and the other 2 returned byte- + identical output because the encoder never referenced the dictionary + for them. Zero returned wrong bytes. `migrate_content_to_dictionary_ + brotli.py` relies on this direction, and WebServer.kt's plain-decode + fallback is safe for the same reason. + + * Decoding with the WRONG dictionary is the silent case. Perturbing one + 16 KiB region of the real dictionary and decoding real rows gave 50% + outright failures, 38% that decoded with no error into *different + bytes*, and 12% byte-identical (the perturbed region was never + referenced). Nothing at runtime catches that 38%. + + 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}") @@ -250,6 +417,32 @@ def get_content_type(conn, value: str) -> tuple: return row[0], row[1] == "brotli" +FRAGMENT_SUFFIX_RE = re.compile(r"^(.*)-(\d+)$") + + +def fragment_chain(conn, base_path: str) -> list: + """Every "-" continuation row present, as (n, path) sorted by + n - found by LIKE query and parsed suffix rather than by probing + constructed paths, so it does not matter what N the chain starts at. + + Probing "-1" first (what reassembly used to do) silently returns + a truncated stream for an ADFA-5171 chain numbered from -2, which then + fails to decompress and looks indistinguishable from an already-migrated + row. The LIKE pattern deliberately over-matches - `_` and `%` in a path are + wildcards, and the suffix is not constrained to digits - so the regex + re-check below is what makes the result exact. Never build a DELETE or + UPDATE straight off that pattern. + """ + rows = conn.execute("SELECT path FROM Content WHERE path LIKE ?", (f"{base_path}-%",)).fetchall() + chain = [] + for (path,) in rows: + match = FRAGMENT_SUFFIX_RE.match(path) + if match and match.group(1) == base_path: + chain.append((int(match.group(2)), path)) + chain.sort(key=lambda item: item[0]) + return chain + + def insert_chunked_content(conn, path: str, language_id: int, content_type_id: int, template_id: int, data: bytes, chunked_log: list) -> None: """Inserts `data` (already fully compressed, if applicable - chunking @@ -282,13 +475,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 +493,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 +737,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/remint_dictionary.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/remint_dictionary.py new file mode 100755 index 00000000..5a910d62 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/remint_dictionary.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +""" +remint_dictionary.py + +Trains a NEW shared Brotli dictionary for an already-migrated database and +recompresses every 'brotli' Content row against it, replacing the +CompressionDictionary row. + +This deliberately does what load_or_create_dictionary refuses to do. That +refusal is right for the pipeline: a dictionary-compressed row is only decodable +with the exact dictionary it was compressed against, so replacing the stored +dictionary without recompressing the content orphans every row -- the dictionary +decode fails and the plain fallback fails too. The only safe way to change a +dictionary is to change the content with it, in one operation, which is what this +script is for. Everything runs in a single transaction: either every row is +converted and the dictionary replaced, or nothing is written. + +Why bother re-minting at all: the dictionary a database is first minted with is +permanent for its content, so a poorly-sampled one is permanently expensive. +Measured on the real corpus with only the training sample varied: + + first 300 rows by path (all under "a/") 36.2% smaller than plain + 300 rows stratified across doc sets 33.2% <- worse + stratified, 32 MiB plaintext budget 48.3% <- best + first-by-path, same 32 MiB budget 36.4% <- volume alone: nil + +Re-minting the 21-Aug database took its brotli content from 83.4 MiB to 65.6 +MiB and the vacuumed file from 268 MB to 249 MB, with all 29,677 items verified +byte-identical afterwards. + +Rows are read with the outgoing dictionary, falling back to a plain decode -- +which is not optional, because a dictionary database always holds some plain +rows: anything a plugin contributed on-device, anything written by a script +outside populate_db.py's reach, and any row whose payload the encoder never +needed the dictionary for. + +Every row is proved to round-trip against the new dictionary before it is +written. Afterwards, run verify_remint_dictionary.py against a copy of the +original database: it re-reads every row through both dictionaries and requires +the plaintexts to match. That check is not paranoia -- a row recompressed +against a mismatched dictionary decodes without error into *different* bytes +(measured: 38% of the time), so nothing at runtime would catch it. + +Usage: + cp documentation.db remint.db + ./remint_dictionary.py remint.db + ./verify_remint_dictionary.py documentation.db remint.db && mv remint.db documentation.db +""" +import argparse +import sqlite3 +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import brotli + +from migrate_content_to_dictionary_brotli import ( + BATCH_ROWS, DEFAULT_SAMPLE_SEED, DEFAULT_SAMPLE_SIZE, DEFAULT_TRAINING_BYTES, + collect_training_samples, load_base_rows, read_item, write_item, +) +from populate_db import DEFAULT_DICT_SIZE, DictionaryCompressor, backup_database, train_dictionary + +_thread_local = threading.local() + + +def _compressors(old_dictionary: bytes, new_dictionary: bytes) -> tuple: + """One pair of DictionaryCompressors per worker thread; each writes its + dictionary to a temp file once, so this avoids doing that per row.""" + pair = getattr(_thread_local, "pair", None) + if pair is None: + pair = (DictionaryCompressor(old_dictionary), DictionaryCompressor(new_dictionary)) + _thread_local.pair = pair + return pair + + +def decode_outgoing(old: DictionaryCompressor, stored: bytes) -> bytes | None: + """Plaintext of one item as stored today: against the outgoing dictionary, + or plainly for the rows that never used it.""" + try: + return old.decompress(stored) + except RuntimeError: + pass + try: + return brotli.decompress(stored) + except brotli.error: + return None + + +def convert(old_dictionary: bytes, new_dictionary: bytes, path: str, stored: bytes) -> dict: + """Runs on a worker thread: bytes in, bytes out, no database access.""" + old, new = _compressors(old_dictionary, new_dictionary) + plain = decode_outgoing(old, stored) + if plain is None: + return {"path": path, "error": "decodes with neither the outgoing dictionary nor plainly"} + recompressed = new.compress(plain) + try: + if new.decompress(recompressed) != plain: + return {"path": path, "error": "recompressed bytes do not decode back to the same plaintext"} + except RuntimeError as exc: + return {"path": path, "error": f"recompressed bytes failed to decode: {exc}"} + return {"path": path, "data": recompressed, "before": len(stored), "after": len(recompressed)} + + +def remint(conn, sample_size: int = DEFAULT_SAMPLE_SIZE, dict_size: int = DEFAULT_DICT_SIZE, + training_bytes: int = DEFAULT_TRAINING_BYTES, seed: int = DEFAULT_SAMPLE_SEED, + max_workers: int | None = None) -> dict: + """Re-mints in one transaction. Raises RuntimeError, having written nothing, + if any row fails to convert.""" + row = conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone() + if row is None or not row[0]: + raise RuntimeError("this database has no CompressionDictionary to re-mint; " + "run migrate_content_to_dictionary_brotli.py instead") + old_dictionary = row[0] + base_rows = load_base_rows(conn) + + with DictionaryCompressor(old_dictionary) as old: + samples = collect_training_samples( + conn, base_rows, sample_size, training_bytes, seed, + decode=lambda stored: decode_outgoing(old, stored) or b"", + ) + if not samples: + raise RuntimeError("no training samples could be decoded; refusing to train on nothing") + new_dictionary = train_dictionary(samples, dict_size) + + stats = {"items": len(base_rows), "before": 0, "after": 0} + problems = [] + conn.execute("BEGIN") + with ThreadPoolExecutor(max_workers=max_workers) as executor: + for start in range(0, len(base_rows), BATCH_ROWS): + batch = base_rows[start:start + BATCH_ROWS] + payloads = [(r[0], read_item(conn, r[0])) for r in batch] + results = executor.map( + lambda item: convert(old_dictionary, new_dictionary, item[0], item[1]), payloads + ) + for r, result in zip(batch, results): + if "error" in result: + problems.append((result["path"], result["error"])) + continue + stats["before"] += result["before"] + stats["after"] += result["after"] + write_item(conn, r[0], r[1], r[2], r[3], result["data"]) + + if problems: + conn.rollback() + for path, why in problems[:10]: + print(f"error: {path}: {why}", file=sys.stderr) + raise RuntimeError(f"{len(problems)} row(s) failed to convert; nothing written") + + conn.execute("UPDATE CompressionDictionary SET data = ? WHERE id = 1", (new_dictionary,)) + conn.commit() + stats["dictionary_bytes"] = len(new_dictionary) + return stats + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("db_path", type=Path, help="database to re-mint; work on a copy") + parser.add_argument("--sample-size", type=int, default=DEFAULT_SAMPLE_SIZE, + help=f"rows to aim for when training (default: {DEFAULT_SAMPLE_SIZE}); " + f"--training-bytes is the real limit") + parser.add_argument("--training-bytes", type=int, default=DEFAULT_TRAINING_BYTES, + help=f"plaintext byte budget for training (default: {DEFAULT_TRAINING_BYTES:,})") + parser.add_argument("--sample-seed", type=int, default=DEFAULT_SAMPLE_SEED, + help="seed for the stratified sample, so a dictionary's training set stays " + "reproducible (default: %(default)s)") + parser.add_argument("--dict-size", type=int, default=DEFAULT_DICT_SIZE, + help=f"new dictionary size in bytes (default: {DEFAULT_DICT_SIZE})") + parser.add_argument("--max-workers", type=int, default=None, help="worker threads for recompression") + parser.add_argument("--no-backup", action="store_true", + help="skip the VACUUM INTO backup (for a copy you already treat as disposable)") + args = parser.parse_args() + + if not args.db_path.is_file(): + sys.exit(f"error: {args.db_path} does not exist") + + if not args.no_backup: + print(f"Backing up {args.db_path}...", file=sys.stderr) + print(f"Backup written to {backup_database(args.db_path)}", file=sys.stderr) + + conn = sqlite3.connect(args.db_path) + try: + stats = remint(conn, args.sample_size, args.dict_size, args.training_bytes, + args.sample_seed, args.max_workers) + 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() + + before, after = stats["before"], stats["after"] + print(f"Re-minted {stats['items']:,} item(s) against a {stats['dictionary_bytes']:,}-byte dictionary: " + f"{before:,} -> {after:,} bytes ({100 * (before - after) / before:.1f}% smaller)") + print("Now run verify_remint_dictionary.py before putting it in place.") + + +if __name__ == "__main__": + main() diff --git a/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py new file mode 100644 index 00000000..e945cdd2 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/renumber_misnumbered_fragments.py @@ -0,0 +1,177 @@ +#!/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_chain + +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. Delegates + to populate_db.fragment_chain so this script and the migration cannot drift + apart on how a chain is found - two conventions is what let ADFA-5171's + -2-based chains be silently skipped by the migration.""" + return fragment_chain(conn, base_path) + + +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. + + Two passes, via a parking name no Content row can already hold. A single + ascending pass is collision-free only when the chain shifts *down* (each + target was just vacated by the previous rename); a chain numbered from 0 + shifts up, where ascending order renames onto a slot still occupied and + trips UNIQUE(path) - which rolls back the whole repair run, so one such + chain would block every other fix in the same pass. Parking first is + correct regardless of direction.""" + parked = [] + for _n, path in fragments: + parking_path = f"{path}.renumbering" + conn.execute("UPDATE Content SET path = ? WHERE path = ?", (parking_path, path)) + parked.append(parking_path) + for i, parking_path in enumerate(parked, start=1): + conn.execute("UPDATE Content SET path = ? WHERE path = ?", (f"{base_path}-{i}", parking_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 + # A 0-based chain is repaired, not skipped: WebServer probes "-1" first, + # finds it, and serves the chain with "-0" silently dropped. Shifting it + # up is what renumber_chain's parking pass exists to make safe. + 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..d3b2da30 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_migrate_content_to_dictionary_brotli.py @@ -0,0 +1,342 @@ +#!/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 collect_training_samples, migrate, read_item +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: - the tests below reopen the database to + # check what was actually committed, and the trigger fixture needs a + # schema an in-memory connection would not outlive. + 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, sample_size=20, dict_size=16384) + self.assertEqual(stats["scanned"], 20) + self.assertEqual(stats["migrated"], 20) + self.assertEqual(stats["already"], 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)) + + # High-entropy bytes to push the compressed form over CHUNK_SIZE (text + # from make_text's small vocabulary compresses far too well to cross + # that boundary at any realistic size), plus a text tail so the encoder + # actually references the shared dictionary. Noise alone references + # nothing, and migrate() then correctly reports the row as having + # nothing to gain rather than migrating it. + plain = random.Random(777).randbytes(int(CHUNK_SIZE * 1.2)) + make_text(4000, seed=777) + 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, 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, 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, sample_size=15, dict_size=16384) + self.assertEqual(second_stats["migrated"], 0) + self.assertEqual(second_stats["already"], 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) + + def test_chain_numbered_from_minus_two_is_migrated_not_miscounted(self): + """ADFA-5171 chains must not read as 'already dictionary-compressed'. + + Probing "-1" returns a truncated stream for a chain numbered from + -2; the decode then fails, and counting that as already-migrated leaves + the row plain-Brotli while the run reports a clean finish.""" + 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)) + plain = random.Random(4242).randbytes(int(CHUNK_SIZE * 1.4)) + make_text(4000, seed=4242) + compressed = brotli.compress(plain) + self.assertGreater(len(compressed), CHUNK_SIZE, "fixture must be chunked; adjust its size") + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + ("k/html/misnumbered.html", self.language_id, compressed[:CHUNK_SIZE], self.html_type_id, 0), + ) + # the continuation numbered from -2, with no -1 at all + offset, number = CHUNK_SIZE, 2 + while offset < len(compressed): + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + (f"k/html/misnumbered.html-{number}", self.language_id, + compressed[offset:offset + CHUNK_SIZE], self.html_type_id, 0), + ) + offset += CHUNK_SIZE + number += 1 + self.conn.commit() + + stats = migrate(self.conn, sample_size=21, dict_size=16384) + self.assertEqual(stats["errors"], 0) + self.assertEqual(stats["already"], 0, "a -2 chain was miscounted as already migrated") + self.assertEqual(stats["migrated"], 21) + + # Content survives, and the rewritten chain is -1-based, which is what + # WebServer.kt's reassembly loop actually probes. + with DictionaryCompressor(load_dictionary(self.conn)) as compressor: + self.assertEqual(compressor.decompress(read_item(self.conn, "k/html/misnumbered.html")), plain) + self.assertIsNotNone( + self.conn.execute("SELECT 1 FROM Content WHERE path = 'k/html/misnumbered.html-1'").fetchone()) + + def test_undecodable_row_is_an_error_not_a_success(self): + for i in range(20): + insert_plain_chunked(self.conn, f"k/html/ok{i}.html", self.language_id, self.html_type_id, 0, + make_text(150, seed=i)) + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + ("k/html/corrupt.html", self.language_id, b"not brotli at all", self.html_type_id, 0), + ) + self.conn.commit() + + stats = migrate(self.conn, sample_size=20, dict_size=16384) + self.assertEqual(stats["errors"], 1) + self.assertEqual(stats["already"], 0) + self.assertEqual([path for path, _detail in stats["problems"]], ["k/html/corrupt.html"]) + # left exactly as it was, not half-written + self.assertEqual( + self.conn.execute("SELECT content FROM Content WHERE path = 'k/html/corrupt.html'").fetchone()[0], + b"not brotli at all") + + def test_bookshelf_entry_survives_migration(self): + """Content's AddBook/DeleteBook triggers make a delete+insert cycle on a + '%.pdf' path silently replace the curated Bookshelf row.""" + self.conn.executescript(""" + CREATE TABLE BookCategories (id INTEGER PRIMARY KEY AUTOINCREMENT, category TEXT); + CREATE TABLE Bookshelf ( + contentID INTEGER NOT NULL, title STRING DEFAULT '', description STRING DEFAULT '', + bookCategoryID INTEGER, + FOREIGN KEY (bookCategoryID) REFERENCES BookCategories(id), UNIQUE(title, bookCategoryId)); + CREATE TRIGGER DeleteBook AFTER DELETE ON Content WHEN OLD.path LIKE '%.pdf' + BEGIN DELETE FROM Bookshelf WHERE contentID = OLD.id; END; + CREATE TRIGGER AddBook AFTER INSERT ON Content WHEN NEW.path LIKE '%.pdf' + BEGIN INSERT INTO Bookshelf (contentID, title) VALUES (NEW.id, CURRENT_TIMESTAMP || NEW.id); END; + """) + self.conn.execute("INSERT INTO ContentTypes (value, compression) VALUES ('application/pdf', 'brotli')") + pdf_type_id = self.conn.execute( + "SELECT id FROM ContentTypes WHERE value = 'application/pdf'").fetchone()[0] + self.conn.execute("INSERT INTO BookCategories (category) VALUES ('Programming')") + + for i in range(20): + insert_plain_chunked(self.conn, f"k/html/f{i}.html", self.language_id, self.html_type_id, 0, + make_text(150, seed=i)) + insert_plain_chunked(self.conn, "bookshelfplugin/Notes.pdf", self.language_id, pdf_type_id, 0, + make_text(400, seed=9)) + pdf_id = self.conn.execute( + "SELECT id FROM Content WHERE path = 'bookshelfplugin/Notes.pdf'").fetchone()[0] + self.conn.execute("DELETE FROM Bookshelf") # drop what AddBook just generated + self.conn.execute( + "INSERT INTO Bookshelf (contentID, bookCategoryID, title, description) VALUES (?, 1, ?, ?)", + (pdf_id, "Notes for Professionals", "A curated description")) + self.conn.commit() + + stats = migrate(self.conn, sample_size=21, dict_size=16384) + self.assertEqual(stats["migrated"], 21) + + shelf = self.conn.execute( + "SELECT contentID, bookCategoryID, title, description FROM Bookshelf").fetchall() + self.assertEqual(shelf, [(pdf_id, 1, "Notes for Professionals", "A curated description")], + "the curated Bookshelf entry was replaced") + + def test_second_run_leaves_rows_whose_dictionary_was_never_referenced(self): + """A payload the encoder never needed the dictionary for decodes both + ways, so 'a plain decode succeeded' does not mean 'not yet migrated'.""" + for i in range(20): + insert_plain_chunked(self.conn, f"k/html/t{i}.html", self.language_id, self.html_type_id, 0, + make_text(150, seed=i)) + # 24 bytes of high-entropy noise: nothing in any dictionary can help it. + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, ?, ?, ?, ?)", + ("k/html/tiny.bin", self.language_id, brotli.compress(random.Random(5).randbytes(24)), + self.html_type_id, 0), + ) + self.conn.commit() + + migrate(self.conn, sample_size=20, dict_size=16384) + snapshot = self.conn.execute("SELECT path, content FROM Content").fetchall() + second = migrate(self.conn, sample_size=20, dict_size=16384) + self.assertEqual(second["migrated"], 0, "a second run re-migrated rows it should have left alone") + self.assertEqual(second["errors"], 0) + self.assertEqual(self.conn.execute("SELECT path, content FROM Content").fetchall(), snapshot) + + def test_training_sample_is_stratified_and_reproducible(self): + for doc_set, count in (("a", 40), ("j", 30), ("k", 10)): + for i in range(count): + insert_plain_chunked(self.conn, f"{doc_set}/page{i}.html", self.language_id, + self.html_type_id, 0, make_text(200, seed=hash((doc_set, i)) % 10_000)) + self.conn.commit() + base_rows = self.conn.execute( + "SELECT C.path, C.languageID, C.contentTypeID, C.templateId, LENGTH(C.content) " + "FROM Content C, ContentTypes CT WHERE C.contentTypeID = CT.id AND CT.compression = 'brotli' " + "ORDER BY C.path").fetchall() + + first = collect_training_samples(self.conn, base_rows, sample_size=30, byte_budget=1 << 20, seed=7) + again = collect_training_samples(self.conn, base_rows, sample_size=30, byte_budget=1 << 20, seed=7) + self.assertEqual(first, again, "same seed must produce the same training set") + + # Every doc set should be represented -- the whole point of stratifying. + # Sample paths back out by matching the plaintext we know we inserted. + sampled_sets = set() + for sample in first: + row = self.conn.execute( + "SELECT path FROM Content WHERE content = ?", (brotli.compress(sample),)).fetchone() + if row: + sampled_sets.add(row[0].split("/", 1)[0]) + self.assertGreaterEqual(len(sampled_sets), 2, f"sample covered only {sampled_sets}") + + +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..de128250 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_populate_db_dictionary.py @@ -0,0 +1,132 @@ +#!/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_never_returns_the_original_bytes(self): + # A mismatched dictionary is not guaranteed to fail loudly. Measured on + # the real corpus (perturbing one 16 KiB region of the real dictionary, + # decoding real rows): 50% raised, 38% decoded with no error into + # *different* bytes, 12% decoded identically because the perturbed + # region was never referenced. + # + # So asserting that the decode does not raise would be asserting a coin + # flip, brittle across brotli versions and payloads. The invariant that + # actually holds is the one that matters: a wrong dictionary never + # yields the original bytes *and* reports success. That is why + # load_or_create_dictionary must never retrain over a stored dictionary + # - no runtime check can catch the mismatch afterwards. + 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: + try: + result = wrong_compressor.decompress(compressed) + except RuntimeError: + return # failed loudly, which is the other acceptable outcome + self.assertNotEqual(result, payload, "a wrong dictionary must not appear to succeed") + + 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_remint_dictionary.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_remint_dictionary.py new file mode 100644 index 00000000..143991fb --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_remint_dictionary.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Tests for remint_dictionary.py / verify_remint_dictionary.py (ADFA-5153). + +Run directly: python3 test_remint_dictionary.py +""" +import random +import shutil +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 +from remint_dictionary import remint +from verify_remint_dictionary import verify +from test_migrate_content_to_dictionary_brotli import SCHEMA_SQL, insert_plain_chunked, make_text + + +class RemintDictionaryTest(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 ('text/html', 'brotli')") + self.html_type_id = 1 + self.originals = {} + # Spread across doc sets, so the stratified sampler has something to stratify. + for doc_set in ("a", "j", "k"): + for i in range(12): + path_ = f"{doc_set}/page{i}.html" + plain = make_text(220, seed=hash((doc_set, i)) % 9973) + insert_plain_chunked(self.conn, path_, 1, self.html_type_id, 0, plain) + self.originals[path_] = plain + self.conn.commit() + migrate(self.conn, sample_size=30, dict_size=16384) + self.conn.commit() + self.first_dictionary = load_dictionary(self.conn) + + def tearDown(self): + self.conn.close() + self.db_path.unlink(missing_ok=True) + + def snapshot(self) -> Path: + """A copy of the database as it stands, to verify a re-mint against.""" + fd, path = tempfile.mkstemp(suffix=".db") + Path(path).unlink(missing_ok=True) + copy = Path(path) + shutil.copy2(self.db_path, copy) + self.addCleanup(copy.unlink, True) + return copy + + def plaintext(self, path: str) -> bytes: + with DictionaryCompressor(load_dictionary(self.conn)) as compressor: + row = self.conn.execute("SELECT content FROM Content WHERE path = ?", (path,)).fetchone() + return compressor.decompress(row[0]) + + def test_replaces_the_dictionary_and_preserves_every_payload(self): + before = self.snapshot() + # A different seed is what makes this a real re-mint: the sampler is + # deterministic, so re-training on the same corpus with the same seed + # reproduces the stored dictionary byte for byte (see the sibling test). + stats = remint(self.conn, sample_size=30, dict_size=16384, training_bytes=1 << 20, seed=999) + self.assertEqual(stats["items"], len(self.originals)) + + self.assertNotEqual(load_dictionary(self.conn), self.first_dictionary, + "the dictionary was not actually re-minted") + for path, plain in self.originals.items(): + self.assertEqual(self.plaintext(path), plain, f"{path} lost its content") + self.assertEqual(verify(before, self.db_path), []) + + def test_same_seed_reproduces_the_stored_dictionary(self): + """The sampler is seeded so a dictionary's training set can be + reproduced later -- which also means re-minting with the same seed and + corpus is a no-op, and the verifier says so rather than pretending + something changed.""" + before = self.snapshot() + remint(self.conn, sample_size=30, dict_size=16384, training_bytes=1 << 20) + self.assertEqual(load_dictionary(self.conn), self.first_dictionary) + self.assertEqual(verify(before, self.db_path), []) + + def test_undecodable_row_aborts_with_nothing_written(self): + self.conn.execute( + "INSERT INTO Content (path, languageID, content, contentTypeID, templateId) VALUES (?, 1, ?, ?, 0)", + ("a/broken.html", b"not brotli at all", self.html_type_id)) + self.conn.commit() + snapshot = {path: content for path, content in + self.conn.execute("SELECT path, content FROM Content")} + + with self.assertRaises(RuntimeError): + remint(self.conn, sample_size=30, dict_size=16384, training_bytes=1 << 20) + + self.assertEqual(load_dictionary(self.conn), self.first_dictionary, "dictionary was replaced anyway") + self.assertEqual({path: content for path, content in + self.conn.execute("SELECT path, content FROM Content")}, snapshot, + "rows were rewritten despite the abort") + + def test_verifier_catches_a_corrupted_remint(self): + """Guards against a vacuous verifier: hand it a database whose content + does not match and it must object.""" + before = self.snapshot() + remint(self.conn, sample_size=30, dict_size=16384, training_bytes=1 << 20, seed=4242) + with DictionaryCompressor(load_dictionary(self.conn)) as compressor: + wrong = compressor.compress(b"replaced with something else entirely") + self.conn.execute("UPDATE Content SET content = ? WHERE path = 'j/page3.html'", (wrong,)) + self.conn.commit() + + problems = verify(before, self.db_path) + self.assertEqual([path for path, _why in problems], ["j/page3.html"]) + + +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..dba7f948 --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/test_renumber_misnumbered_fragments.py @@ -0,0 +1,195 @@ +#!/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_renumbers_zero_based_chain(self): + """A chain numbered from -0 shifts *up*, where renaming in ascending + order would land on a slot still occupied and trip UNIQUE(path) - + rolling back every other repair in the same pass. It is as broken as a + -2 chain: WebServer.kt probes "-1", finds it, and serves the chain with + "-0" silently dropped.""" + base = "a/devsite/media/zero-based.gif" + self.insert(base, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{base}-0", chunk_bytes(CHUNK_SIZE, b"B")) + self.insert(f"{base}-1", chunk_bytes(CHUNK_SIZE, b"C")) + self.insert(f"{base}-2", b"D" * 100) + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 1) + self.assertEqual(stats["chains_gapped"], 0) + self.assertEqual(self.all_paths(), {base, f"{base}-1", f"{base}-2", f"{base}-3"}) + # order preserved: -0 -> -1, -1 -> -2, -2 -> -3 + 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"), b"D" * 100) + + def test_zero_based_chain_does_not_block_other_repairs(self): + """One chain tripping UNIQUE(path) used to roll back the whole run.""" + zero_based = "a/devsite/media/zero.gif" + self.insert(zero_based, chunk_bytes(CHUNK_SIZE, b"A")) + self.insert(f"{zero_based}-0", b"B" * 100) + two_based = "a/devsite/media/two.gif" + self.insert(two_based, chunk_bytes(CHUNK_SIZE, b"C")) + self.insert(f"{two_based}-2", b"D" * 100) + self.conn.commit() + + stats = repair(self.conn) + self.conn.commit() + + self.assertEqual(stats["chains_renumbered"], 2) + self.assertEqual(self.all_paths(), + {zero_based, f"{zero_based}-1", two_based, f"{two_based}-1"}) + + 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/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/verify_remint_dictionary.py b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/verify_remint_dictionary.py new file mode 100755 index 00000000..ed3a4edc --- /dev/null +++ b/ProcessDocs/ProcessKotlinDocs/ProcessKotlinWebsiteJSON/verify_remint_dictionary.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +verify_remint_dictionary.py + +Proves a re-minted database still holds exactly the content the original did. + +For every 'brotli' item: decode it out of the new database with the new +dictionary, decode the same path out of the original with the old dictionary +(falling back to plain), and require the two plaintexts to be byte-identical. +Exits non-zero on any mismatch, so it can gate the swap: + + ./verify_remint_dictionary.py documentation.db remint.db && mv remint.db documentation.db + +This is the check that makes re-minting safe to do at all. A row recompressed +against a mismatched dictionary does not fail loudly -- measured on the real +corpus, a wrong dictionary decodes without error into *different* bytes 38% of +the time (50% raises, 12% is identical because the perturbed region was never +referenced). Nothing at runtime detects that, so it has to be caught here, +against the original, before the file is put in place. + +Reads only: neither database is modified. +""" +import argparse +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import brotli + +from populate_db import DictionaryCompressor +from migrate_content_to_dictionary_brotli import load_base_rows, read_item + +_thread_local = threading.local() + + +def _state(old_dictionary: bytes, new_dictionary: bytes, old_db: Path, new_db: Path): + """Per-thread connections and compressors: a sqlite3.Connection cannot be + shared across threads, and a DictionaryCompressor holds a temp file.""" + import sqlite3 + state = getattr(_thread_local, "state", None) + if state is None: + state = { + "old_conn": sqlite3.connect(f"file:{old_db}?mode=ro", uri=True), + "new_conn": sqlite3.connect(f"file:{new_db}?mode=ro", uri=True), + "old": DictionaryCompressor(old_dictionary), + "new": DictionaryCompressor(new_dictionary), + } + _thread_local.state = state + return state + + +def check(old_dictionary: bytes, new_dictionary: bytes, old_db: Path, new_db: Path, path: str): + """Returns None when the item matches, else a description of how it differs.""" + state = _state(old_dictionary, new_dictionary, old_db, new_db) + was = read_item(state["old_conn"], path) + now = read_item(state["new_conn"], path) + if not now: + return "row missing from the re-minted database" + + try: + old_plain = state["old"].decompress(was) + except RuntimeError: + try: + old_plain = brotli.decompress(was) + except brotli.error: + return "original row could not be decoded at all" + try: + new_plain = state["new"].decompress(now) + except RuntimeError: + return "re-minted row does not decode with the new dictionary" + + if old_plain != new_plain: + return f"plaintext differs ({len(old_plain):,} vs {len(new_plain):,} bytes)" + return None + + +def verify(old_db: Path, new_db: Path, max_workers: int | None = None) -> list: + import sqlite3 + old_conn = sqlite3.connect(f"file:{old_db}?mode=ro", uri=True) + new_conn = sqlite3.connect(f"file:{new_db}?mode=ro", uri=True) + old_dictionary = old_conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone()[0] + new_dictionary = new_conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone()[0] + if old_dictionary == new_dictionary: + print("warning: both databases carry the same dictionary; nothing was re-minted", file=sys.stderr) + + items = [row[0] for row in load_base_rows(new_conn)] + print(f"Verifying {len(items):,} item(s)...", file=sys.stderr) + + bad = [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + for i, (path, problem) in enumerate( + zip(items, executor.map( + lambda p: check(old_dictionary, new_dictionary, old_db, new_db, p), items)), start=1 + ): + if problem: + bad.append((path, problem)) + if i % 6000 == 0: + print(f" {i:,}/{len(items):,} checked, {len(bad)} mismatched", file=sys.stderr) + print(f"Byte-identical: {len(items) - len(bad):,}/{len(items):,}") + return bad + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("original_db", type=Path, help="the database as it was before re-minting") + parser.add_argument("reminted_db", type=Path, help="the re-minted database to check") + parser.add_argument("--max-workers", type=int, default=None, help="worker threads") + args = parser.parse_args() + + for path in (args.original_db, args.reminted_db): + if not path.is_file(): + sys.exit(f"error: {path} does not exist") + + bad = verify(args.original_db, args.reminted_db, args.max_workers) + for path, problem in bad[:20]: + print(f"MISMATCH {path}: {problem}", file=sys.stderr) + if len(bad) > 20: + print(f"... and {len(bad) - 20} more", file=sys.stderr) + sys.exit(1 if bad else 0) + + +if __name__ == "__main__": + main() diff --git a/docdb-studio/README.md b/docdb-studio/README.md index 7d174003..d8c654dd 100644 --- a/docdb-studio/README.md +++ b/docdb-studio/README.md @@ -79,6 +79,68 @@ These same commands work on macOS, Linux, and Windows (in PowerShell, Command Pr `uv sync` creates a virtual environment in `.venv/` and installs every dependency from `uv.lock`. There is no need to activate the venv — `uv run` does that for you. +## Installing the `brotli` command-line tool + +One dependency `uv sync` cannot install for you. Databases built since ADFA-5153 compress their `Content` rows against a shared dictionary stored inside the database itself, and no Python library exposes a custom dictionary, so docdb-studio runs the `brotli` **command-line program** to read and write those rows. If it is missing, content rows come up blank and docdb-studio prints an explanatory error in the terminal you launched it from — so a blank preview plus that message means "install this tool", not "this page is empty". A database with no `CompressionDictionary` table needs nothing extra. + +> **The `brotli` in `uv sync` is not this.** The Python package named `brotli` is already installed for you, and it is a *different thing* from the `brotli` program. Installing the Python package again will not help; you need the program, which puts a `brotli` (or `brotli.exe`) command on your `PATH`. + +### macOS + +```bash +brew install brotli +``` + +### Linux + +```bash +sudo apt install brotli # Debian, Ubuntu, Mint +sudo dnf install brotli # Fedora, RHEL +``` + +### Windows + +Pick whichever package manager you already have. Each line searches first, because the package's exact name can change; copy the name from the search results into the install command. + +```powershell +winget search brotli +winget install --id= -e +``` + +```powershell +scoop search brotli +scoop install brotli +``` + +```powershell +choco search brotli +choco install brotli +``` + +`choco` needs an Administrator terminal; `winget` and `scoop` do not. + +If you already use **Git for Windows** or **MSYS2**, its package manager has it too — from the MSYS2 terminal: + +```bash +pacman -S mingw-w64-ucrt-x86_64-brotli +``` + +**Then close your terminal and open a new one.** Windows only picks up a changed `PATH` in newly-opened terminals, so a fresh window is what makes the next check meaningful: + +```powershell +brotli --version +``` + +A version number means you are done. If instead you get "not recognized", the program is installed somewhere that is not on your `PATH`. Find it and add that folder: + +```powershell +where.exe brotli +``` + +If `where.exe` finds nothing, look in the install location your package manager reports (Scoop uses `%USERPROFILE%\scoop\shims`, Chocolatey uses `C:\ProgramData\chocolatey\bin`), then add that folder to `PATH` under **Settings → System → About → Advanced system settings → Environment Variables**, and open a new terminal again. + +To confirm docdb-studio itself can see it, open a database that has a `CompressionDictionary` table and click any content row. If the preview shows the page's text, the tool is wired up correctly. + ## Updating to the latest version When a new version of docdb-studio is released, you can refresh your local copy without re-cloning. Open a terminal, move into the project folder, and pull down the latest changes: diff --git a/docdb-studio/docdb_studio.py b/docdb-studio/docdb_studio.py index 30f14730..25be0c3f 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,10 @@ 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 BrotliCliMissing as exc: + print(f"error: cannot read {base_path!r}: {exc}", file=sys.stderr) + return [] except brotli.error: return [] return extract_html_anchors(full) @@ -708,7 +714,10 @@ def fetch_content_for_path( full = b"".join(parts) if compression == "brotli": try: - full = brotli.decompress(full) + full = decompress_brotli(full, db_path) + except BrotliCliMissing as exc: + print(f"error: cannot read {base_path!r}: {exc}", file=sys.stderr) + return None except brotli.error: return None return full, mime @@ -1220,11 +1229,150 @@ 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] = {} + + +class BrotliCliMissing(brotli.error): + """The `brotli` binary a dictionary database needs is not installed. + + Subclasses brotli.error so it cannot escape as an unhandled RuntimeError out + of content preview or anchor validation -- but the call sites catch it + *separately* from a plain brotli.error and log it, because the two mean + different things to whoever is looking at the screen. A corrupt row is one + bad row; a missing binary means nothing in this database will ever decode, + and it is fixable in one command. Silently returning an empty preview for + that is indistinguishable from an empty page. + + See "Installing the `brotli` command-line tool" in README.md.""" + + +def _find_brotli_cli() -> str: + path = shutil.which("brotli") + if path is None: + raise BrotliCliMissing( + "this database uses a shared Brotli dictionary (ADFA-5153), which needs the `brotli` " + "command-line program -- not the Python package of the same name that is already " + "installed. brew install brotli (macOS), apt install brotli (Linux), or on Windows " + "winget/scoop/choco; see \"Installing the brotli command-line tool\" in README.md. " + "Then reopen the database." + ) + 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. + + Only a *definitive* answer is cached. `sqlite3.OperationalError` also covers + "database is locked", which is entirely plausible for a GUI while another tool + writes the same file; caching that as None would downgrade the whole session to + plain Brotli, so imports would write plain rows into a dictionary database and + reads of existing rows would fail. On any such error this returns None for this + one call and retries on the next.""" + 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 as exc: + print(f"warning: could not read {db_path}'s compression dictionary ({exc}); " + f"not caching that, will retry", file=sys.stderr) + return 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) + try: + result = subprocess.run( + [_find_brotli_cli(), "-D", str(dict_path), "-c"], + input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + except OSError as exc: + raise RuntimeError(f"could not run the brotli tool needed to compress against this " + f"database's shared dictionary: {exc}") from exc + 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. + + Falls back to a plain decode when the dictionary-attached one fails, which + WebServer.kt also does and for the same reason: a dictionary database still + contains plain rows. Any row a plugin contributes on-device is plain (see + PluginDocumentationManager/BrotliCompressor), scripts outside populate_db.py's + reach may not have been migrated yet, and a partially-completed migration leaves + a mixture by design. Without the fallback those rows are unreadable here even + though the app serves them fine. The fallback is safe in this direction: + measured over 400 real dictionary-compressed rows, decoding one without its + dictionary raised 398 times and returned identical bytes twice - never wrong + bytes. (Decoding with the *wrong* dictionary is the silent case, and no fallback + can detect it.) + + 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) + try: + result = subprocess.run( + [_find_brotli_cli(), "-d", "-D", str(dict_path), "-c"], + input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + ) + except OSError as exc: + raise brotli.error(f"could not run the brotli tool: {exc}") from exc + if result.returncode == 0: + return result.stdout + # A plain row in a dictionary database, most likely. brotli.decompress raises + # brotli.error of its own if the blob is genuinely corrupt, so a real failure + # still reaches the caller with the expected exception type. + return brotli.decompress(data) def fragment_blob(blob: bytes, chunk_size: int = CONTENT_CHUNK_SIZE) -> list[bytes]: @@ -1430,7 +1578,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..2ed41d42 --- /dev/null +++ b/docdb-studio/tests/test_compression_dictionary.py @@ -0,0 +1,282 @@ +"""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 + # A *definitive* answer is cached: docdb_studio never writes a dictionary + # itself, so one does not appear mid-session under normal use. (An error + # answer is not cached -- see the locked-database test below.) + 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) + + +def test_locked_database_is_not_cached_as_having_no_dictionary() -> None: + """`sqlite3.OperationalError` also covers "database is locked", which a GUI + hits whenever another tool writes the same file. Caching that as "no + dictionary" would downgrade the whole session to plain Brotli: imports would + write plain rows into a dictionary database and existing rows would fail to + read.""" + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + real_connect = docdb_studio.sqlite3.connect + calls = {"n": 0} + + def flaky_connect(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise docdb_studio.sqlite3.OperationalError("database is locked") + return real_connect(*args, **kwargs) + + try: + docdb_studio.sqlite3.connect = flaky_connect + assert get_compression_dictionary(db) is None # the transient failure + assert get_compression_dictionary(db) == dictionary_data # retried, not cached + finally: + docdb_studio.sqlite3.connect = real_connect + db.unlink(missing_ok=True) + + +def test_plain_row_in_a_dictionary_database_still_decodes() -> None: + """A dictionary database still contains plain rows: anything a plugin + contributes on-device, anything a script outside populate_db.py wrote, and + everything not yet converted during a partial migration. WebServer.kt falls + back to a plain decode for exactly this, and docdb-studio has to agree or it + cannot read rows the app serves fine.""" + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + try: + payload = _make_text(400, seed=4242) + plain_row = brotli.compress(payload) + assert docdb_studio.decompress_brotli(plain_row, db) == payload + finally: + db.unlink(missing_ok=True) + + +def test_missing_brotli_cli_surfaces_as_brotli_error() -> None: + """Every existing call site guards decoding with `except brotli.error`, so a + missing binary has to arrive as one rather than as an unhandled RuntimeError + out of content preview.""" + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + real_which = docdb_studio.shutil.which + try: + docdb_studio.shutil.which = lambda name: None + with pytest.raises(brotli.error) as excinfo: + docdb_studio.decompress_brotli(b"anything", db) + assert "brotli" in str(excinfo.value) + finally: + docdb_studio.shutil.which = real_which + db.unlink(missing_ok=True) + + +def test_missing_brotli_cli_is_reported_not_swallowed(capsys) -> None: + """A missing binary and a corrupt row both leave the preview empty, but they + mean different things: one bad row versus nothing in this database will ever + decode, fixable in one command. The call sites log the second.""" + dictionary_data = _train_dictionary([_make_text(150, i) for i in range(60)]) + db = _make_db(with_dictionary=dictionary_data) + real_which = docdb_studio.shutil.which + try: + html = b"

needs the CLI

" + _insert_content(db, "docs/page.html", compress_for_storage(html, "brotli", db)) + docdb_studio.shutil.which = lambda name: None + + assert fetch_content_for_path(db, "docs/page.html") is None + assert docdb_studio.get_html_anchors_for_path(db, "docs/page.html") == [] + + message = capsys.readouterr().err + assert "docs/page.html" in message + assert "README" in message # tells the reader where to go + assert message.count("error:") == 2 # once per call site, not swallowed + finally: + docdb_studio.shutil.which = real_which + 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 ---------- diff --git a/scripts/sync_kotlin_stdlib_docs/sync_kdoc_json_to_db.py b/scripts/sync_kotlin_stdlib_docs/sync_kdoc_json_to_db.py index fb614926..d1585de6 100755 --- a/scripts/sync_kotlin_stdlib_docs/sync_kdoc_json_to_db.py +++ b/scripts/sync_kotlin_stdlib_docs/sync_kdoc_json_to_db.py @@ -22,16 +22,25 @@ A timestamped backup of the database is made before anything is modified. """ import argparse +import atexit import os import shutil import sqlite3 +import subprocess import sys +import tempfile from datetime import datetime, timezone +from pathlib import Path import brotli PREFIXES = ["k/kotlin-stdlib", "k/kotlin-reflect", "k/kotlin-test"] +# Refuse to run if this fraction or more of the matched rows resolve to no source +# file. A Dokka upgrade that changes the emitted layout makes *every* lookup miss, +# and the only signal would be "Done: updated 0, deleted N" on a gutted database. +MAX_DELETE_FRACTION = 0.5 + def backup_database(db_path): timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") @@ -49,9 +58,54 @@ def relative_target_path(content_path): return without_prefix -def compress_for(compression, raw_bytes, path): +def load_compression_dictionary(conn): + """This database's shared Brotli dictionary (ADFA-5153), or None if it has + none. Rows written here must be compressed the same way populate_db.py + compresses the rest of the database: a plain-Brotli row inside a dictionary + database forfeits the dictionary's compression entirely (readers cope with + it -- WebServer.kt and docdb-studio both fall back to a plain decode -- but + the bytes stay large for no reason).""" + table = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'CompressionDictionary'" + ).fetchone() + if table is None: + return None + row = conn.execute("SELECT data FROM CompressionDictionary WHERE id = 1").fetchone() + return row[0] if row and row[0] else None + + +class DictionaryBrotli: + """Compresses against a raw Brotli dictionary by shelling out to the + `brotli` CLI - the Python `brotli` package exposes no dictionary parameter. + Mirrors populate_db.DictionaryCompressor, kept local because that module + lives in a different tree and this script is standalone.""" + + def __init__(self, dictionary_data): + path = shutil.which("brotli") + if path is None: + raise RuntimeError( + "this database uses a shared Brotli dictionary (ADFA-5153), which needs the " + "`brotli` command-line tool; install it (apt install brotli) and retry" + ) + self._brotli = path + self._dir = Path(tempfile.mkdtemp(prefix="sync-kdoc-brotli-dict-")) + self._dict_path = self._dir / "dictionary.bin" + self._dict_path.write_bytes(dictionary_data) + atexit.register(lambda: shutil.rmtree(self._dir, ignore_errors=True)) + + def compress(self, data): + result = subprocess.run( + [self._brotli, "-D", str(self._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 compress_for(compression, raw_bytes, path, compressor=None): if compression == "brotli": - return brotli.compress(raw_bytes) + return compressor.compress(raw_bytes) if compressor is not None else brotli.compress(raw_bytes) if compression == "none": return raw_bytes raise ValueError(f"Unknown compression '{compression}' needed for {path}") @@ -141,6 +195,30 @@ def main(): deleted_paths = [] unknown_types = set() + dictionary_data = load_compression_dictionary(conn) + compressor = DictionaryBrotli(dictionary_data) if dictionary_data else None + print( + "Compressing brotli rows against this database's shared dictionary." + if compressor else + "This database has no CompressionDictionary; writing plain Brotli.", + file=sys.stderr, + ) + + # Resolve every source file before touching anything, so a wholesale miss + # aborts instead of deleting the rows one at a time (see MAX_DELETE_FRACTION). + missing = [path for _id, path, _type in rows + if not os.path.isfile(os.path.join(args.plugin_output_root, relative_target_path(path)))] + if rows and len(missing) >= max(1, int(len(rows) * MAX_DELETE_FRACTION)): + print( + f"error: {len(missing)} of {len(rows)} matched Content rows resolve to no file under " + f"{args.plugin_output_root!r}. That is a layout mismatch, not {len(missing)} deletions - " + f"refusing to delete them. Check the Dokka output tree, then re-run. Examples: " + f"{', '.join(missing[:3])}", + file=sys.stderr, + ) + conn.close() + sys.exit(1) + try: conn.execute("BEGIN") for content_id, path, content_type_id in rows: @@ -156,7 +234,7 @@ def main(): unknown_types.add(content_type_id) compression = "none" - new_blob = compress_for(compression, raw_bytes, path) + new_blob = compress_for(compression, raw_bytes, path, compressor) if args.dry_run: print(f" [UPDATE] {path} <- {rel_target}")