468 lines
19 KiB
Python
468 lines
19 KiB
Python
"""Export normalized Ministry fragments for the OpenSearch Bulk API."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
from search.catalog import authority_codes, source_code
|
|
|
|
APP_VERSION = "0.8.1"
|
|
LANGUAGES = {"ru", "ky"}
|
|
DEFAULT_MAPPING = Path(__file__).with_name("minjust-fragments-index.json")
|
|
|
|
|
|
def read_json(path: Path):
|
|
def reject_constant(value: str):
|
|
raise ValueError(f"Invalid JSON constant: {value}")
|
|
|
|
for attempt in range(3):
|
|
try:
|
|
with path.open(encoding="utf-8") as source:
|
|
return json.load(source, parse_constant=reject_constant)
|
|
except (OSError, UnicodeError, json.JSONDecodeError) as error:
|
|
if attempt == 2:
|
|
raise ValueError(f"Cannot read JSON {path}: {error}") from error
|
|
time.sleep(0.1)
|
|
|
|
|
|
def localized(value: object, language: str):
|
|
return value.get(language) if isinstance(value, dict) else None
|
|
|
|
|
|
def paths(document: dict, field: str, language: str) -> list[str]:
|
|
return [" > ".join(item[language]) for item in document.get(field, []) if item.get(language)]
|
|
|
|
|
|
def search_document(document: dict, fragment: dict, expected: tuple[str, str, str, int], current_edition: str) -> dict:
|
|
document_code, edition_code, language, position = expected
|
|
fragment_id = f"document:{document_code}:edition:{edition_code}:lang:{language}:fragment:{position}"
|
|
required = ("id", "document_code", "edition_code", "language", "position", "type", "text", "text_sha256", "source_path", "source_sha256")
|
|
if not isinstance(fragment, dict) or any(key not in fragment for key in required):
|
|
raise ValueError(f"Incomplete fragment {fragment_id}")
|
|
if (fragment["document_code"], fragment["edition_code"], fragment["language"], fragment["position"], fragment["id"]) != (*expected, fragment_id):
|
|
raise ValueError(f"Fragment identity does not match its path: {fragment_id}")
|
|
if language not in LANGUAGES or not isinstance(fragment["text"], str) or not fragment["text"]:
|
|
raise ValueError(f"Invalid fragment content: {fragment_id}")
|
|
for key in ("text_sha256", "source_sha256"):
|
|
value = fragment[key]
|
|
if not isinstance(value, str) or len(value) != 64 or any(char not in "0123456789abcdef" for char in value):
|
|
raise ValueError(f"Invalid {key}: {fragment_id}")
|
|
if hashlib.sha256(fragment["text"].encode()).hexdigest() != fragment["text_sha256"]:
|
|
raise ValueError(f"Text checksum mismatch: {fragment_id}")
|
|
|
|
dates = document.get("dates") or {}
|
|
result = {
|
|
"schema_version": document["schema_version"],
|
|
"document_code": document_code,
|
|
"edition_code": edition_code,
|
|
"is_current_edition": edition_code == current_edition,
|
|
"language": language,
|
|
"position": position,
|
|
"fragment_type": fragment["type"],
|
|
f"text_{language}": fragment["text"],
|
|
"document_name_ru": localized(document.get("name"), "ru"),
|
|
"document_name_ky": localized(document.get("name"), "ky"),
|
|
"document_type_ru": localized(document.get("type"), "ru"),
|
|
"document_type_ky": localized(document.get("type"), "ky"),
|
|
"document_type_code": source_code("document_type", document.get("type")),
|
|
"status_ru": localized(document.get("status"), "ru"),
|
|
"status_ky": localized(document.get("status"), "ky"),
|
|
"status_code": source_code("status", document.get("status")),
|
|
"number": document.get("number"),
|
|
"date_adopted": dates.get("DateAdopted"),
|
|
"authority_paths_ru": paths(document, "authority_paths", "ru"),
|
|
"authority_paths_ky": paths(document, "authority_paths", "ky"),
|
|
"authority_codes": authority_codes(document.get("authority_paths", [])),
|
|
"source_path": fragment["source_path"],
|
|
"source_sha256": fragment["source_sha256"],
|
|
"text_sha256": fragment["text_sha256"],
|
|
}
|
|
return {key: value for key, value in result.items() if value is not None}
|
|
|
|
|
|
def document_codes(input_root: Path, limit: int | None = None, start_at: str | None = None) -> list[str]:
|
|
connection = sqlite3.connect(f"file:{input_root / 'manifest.sqlite3'}?mode=ro", uri=True)
|
|
try:
|
|
query = "SELECT code FROM documents WHERE state = 'success' ORDER BY CAST(code AS INTEGER), code"
|
|
parameters: list[int] = []
|
|
if limit is not None:
|
|
query += " LIMIT ?"
|
|
parameters.append(limit)
|
|
codes = [str(code) for (code,) in connection.execute(query, parameters)]
|
|
finally:
|
|
connection.close()
|
|
if start_at is None:
|
|
return codes
|
|
try:
|
|
return codes[codes.index(start_at):]
|
|
except ValueError as error:
|
|
raise ValueError(f"Resume document not found in selected range: {start_at}") from error
|
|
|
|
|
|
def bulk_pairs(
|
|
input_root: Path,
|
|
index: str,
|
|
limit: int | None = None,
|
|
start_at: str | None = None,
|
|
) -> Iterator[tuple[str, bytes]]:
|
|
document_root = input_root / "documents"
|
|
if not document_root.is_dir():
|
|
raise FileNotFoundError(f"Document directory not found: {document_root}")
|
|
for code in document_codes(input_root, limit, start_at):
|
|
directory = document_root / code
|
|
document = read_json(directory / "document.json")
|
|
if document.get("source_code") != directory.name:
|
|
raise ValueError(f"Document identity does not match its path: {directory}")
|
|
edition_root = directory / "editions"
|
|
editions = [path.name for path in edition_root.iterdir() if path.is_dir()] if edition_root.is_dir() else []
|
|
if not editions:
|
|
continue
|
|
current_edition = max(editions, key=lambda value: (not value.isdigit(), int(value) if value.isdigit() else value))
|
|
for fragment_path in sorted(directory.glob("editions/*/*/fragments.json")):
|
|
edition_code, language = fragment_path.parts[-3:-1]
|
|
values = read_json(fragment_path)
|
|
if not isinstance(values, list):
|
|
raise ValueError(f"Fragments must be a list: {fragment_path}")
|
|
for position, fragment in enumerate(values, 1):
|
|
source = search_document(document, fragment, (directory.name, edition_code, language, position), current_edition)
|
|
action = json.dumps({"index": {"_index": index, "_id": fragment["id"]}}, ensure_ascii=False, allow_nan=False)
|
|
body = json.dumps(source, ensure_ascii=False, allow_nan=False)
|
|
yield directory.name, f"{action}\n{body}\n".encode()
|
|
|
|
|
|
def document_count(input_root: Path, limit: int | None, start_at: str | None = None) -> int:
|
|
return len(document_codes(input_root, limit, start_at))
|
|
|
|
|
|
def bulk_batches(pairs: Iterator[tuple[str, bytes]], maximum_bytes: int) -> Iterator[tuple[str, bytes]]:
|
|
batch = bytearray()
|
|
last_code = ""
|
|
for code, pair in pairs:
|
|
if len(pair) > maximum_bytes:
|
|
raise ValueError(f"One Bulk pair exceeds the {maximum_bytes}-byte batch limit")
|
|
if batch and len(batch) + len(pair) > maximum_bytes:
|
|
yield last_code, bytes(batch)
|
|
batch.clear()
|
|
last_code = code
|
|
batch.extend(pair)
|
|
if batch:
|
|
yield last_code, bytes(batch)
|
|
|
|
|
|
def export_bulk(input_root: Path, output: Path, index: str, limit: int | None = None) -> tuple[int, int]:
|
|
source = input_root.resolve()
|
|
destination = output.resolve()
|
|
if destination == source or destination.is_relative_to(source):
|
|
raise ValueError("--output must not be inside --input")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
fragments = 0
|
|
with tempfile.NamedTemporaryFile("wb", dir=output.parent, delete=False) as target:
|
|
temporary = Path(target.name)
|
|
try:
|
|
for code, pair in bulk_pairs(input_root, index, limit):
|
|
target.write(pair)
|
|
fragments += 1
|
|
target.flush()
|
|
os.fsync(target.fileno())
|
|
os.replace(temporary, output)
|
|
except BaseException:
|
|
temporary.unlink(missing_ok=True)
|
|
raise
|
|
return document_count(input_root, limit), fragments
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def write_json_atomic(path: Path, value: dict) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.NamedTemporaryFile("wb", dir=path.parent, delete=False) as target:
|
|
temporary = Path(target.name)
|
|
try:
|
|
target.write((json.dumps(value, ensure_ascii=False, indent=2) + "\n").encode())
|
|
target.flush()
|
|
os.fsync(target.fileno())
|
|
os.replace(temporary, path)
|
|
except BaseException:
|
|
temporary.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def request_json(
|
|
url: str,
|
|
method: str,
|
|
body: bytes | None,
|
|
content_type: str,
|
|
attempts: int = 5,
|
|
retry_invalid_json: bool = False,
|
|
) -> dict:
|
|
request = urllib.request.Request(url, data=body, method=method, headers={"Content-Type": content_type})
|
|
for attempt in range(attempts):
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
return json.load(response)
|
|
except (UnicodeError, json.JSONDecodeError) as error:
|
|
if not retry_invalid_json or attempt == attempts - 1:
|
|
raise RuntimeError(f"{method} {url} returned invalid JSON: {error}") from error
|
|
delay = 2**attempt
|
|
except urllib.error.HTTPError as error:
|
|
response_body = error.read(500).decode("utf-8", errors="replace")
|
|
error.close()
|
|
retryable = error.code == 429 or 500 <= error.code < 600
|
|
if not retryable or attempt == attempts - 1:
|
|
raise RuntimeError(
|
|
f"{method} {url} failed with HTTP {error.code}: {response_body}"
|
|
) from error
|
|
retry_after = error.headers.get("Retry-After")
|
|
delay = float(retry_after) if retry_after and retry_after.isdigit() else 2**attempt
|
|
except (urllib.error.URLError, TimeoutError, ConnectionResetError) as error:
|
|
if attempt == attempts - 1:
|
|
raise RuntimeError(f"{method} {url} failed after {attempts} attempts: {error}") from error
|
|
delay = 2**attempt
|
|
time.sleep(delay)
|
|
raise AssertionError("unreachable")
|
|
|
|
|
|
def opensearch_identity(base: str, index: str, index_url: str) -> tuple[str, str]:
|
|
cluster = request_json(f"{base}/", "GET", None, "application/json")
|
|
definition = request_json(index_url, "GET", None, "application/json")
|
|
try:
|
|
return cluster["cluster_uuid"], definition[index]["settings"]["index"]["uuid"]
|
|
except (KeyError, TypeError) as error:
|
|
raise RuntimeError("OpenSearch identity response is incomplete") from error
|
|
|
|
|
|
def checkpoint_state(
|
|
input_root: Path,
|
|
index: str,
|
|
checkpoint: Path,
|
|
resume: bool,
|
|
url: str,
|
|
cluster_uuid: str,
|
|
index_uuid: str,
|
|
limit: int | None,
|
|
alias: str | None,
|
|
) -> tuple[dict, str | None]:
|
|
source = input_root.resolve()
|
|
destination = checkpoint.resolve()
|
|
if destination == source or destination.is_relative_to(source):
|
|
raise ValueError("--checkpoint must not be inside --input")
|
|
manifest_sha256 = file_sha256(input_root / "manifest.sqlite3")
|
|
if not resume:
|
|
return {
|
|
"schema_version": 2,
|
|
"url": url,
|
|
"cluster_uuid": cluster_uuid,
|
|
"index": index,
|
|
"index_uuid": index_uuid,
|
|
"input": str(source),
|
|
"manifest_sha256": manifest_sha256,
|
|
"limit": limit,
|
|
"alias": alias,
|
|
"last_document_code": None,
|
|
"complete": False,
|
|
}, None
|
|
|
|
state = read_json(checkpoint)
|
|
expected = {
|
|
"schema_version",
|
|
"url",
|
|
"cluster_uuid",
|
|
"index",
|
|
"index_uuid",
|
|
"input",
|
|
"manifest_sha256",
|
|
"limit",
|
|
"alias",
|
|
"last_document_code",
|
|
"complete",
|
|
}
|
|
legacy_expected = expected - {"alias"}
|
|
if not isinstance(state, dict):
|
|
raise ValueError(f"Invalid checkpoint: {checkpoint}")
|
|
if set(state) == legacy_expected:
|
|
if state["schema_version"] != 1 or alias is not None:
|
|
raise ValueError(f"Legacy checkpoint does not support --alias: {checkpoint}")
|
|
state["schema_version"] = 2
|
|
state["alias"] = None
|
|
elif set(state) != expected:
|
|
raise ValueError(f"Invalid checkpoint: {checkpoint}")
|
|
if (
|
|
state["schema_version"] != 2
|
|
or state["url"] != url
|
|
or state["cluster_uuid"] != cluster_uuid
|
|
or state["index"] != index
|
|
or state["index_uuid"] != index_uuid
|
|
or state["input"] != str(source)
|
|
):
|
|
raise ValueError(f"Checkpoint does not match this load: {checkpoint}")
|
|
if state["limit"] != limit:
|
|
raise ValueError(f"Checkpoint limit does not match --limit: {checkpoint}")
|
|
if state["alias"] != alias:
|
|
raise ValueError(f"Checkpoint alias does not match --alias: {checkpoint}")
|
|
if state["manifest_sha256"] != manifest_sha256:
|
|
raise ValueError("Normalized manifest changed; create a new versioned index")
|
|
if state["complete"] is not False:
|
|
raise ValueError(f"Checkpoint is already complete: {checkpoint}")
|
|
start_at = state["last_document_code"]
|
|
if start_at is not None and not isinstance(start_at, str):
|
|
raise ValueError(f"Invalid checkpoint document code: {checkpoint}")
|
|
return state, start_at
|
|
|
|
|
|
def load_bulk(
|
|
input_root: Path,
|
|
url: str,
|
|
index: str,
|
|
mapping: Path = DEFAULT_MAPPING,
|
|
maximum_bytes: int = 25 * 1024 * 1024,
|
|
limit: int | None = None,
|
|
resume: bool = False,
|
|
checkpoint: Path = Path("data/opensearch/minjust-fragments.checkpoint.json"),
|
|
alias: str | None = None,
|
|
) -> tuple[int, int]:
|
|
base = url.rstrip("/")
|
|
if alias is not None and (not alias or alias == index):
|
|
raise ValueError("--alias must differ from --index")
|
|
index_url = f"{base}/{urllib.parse.quote(index, safe='')}"
|
|
if resume:
|
|
cluster_uuid, index_uuid = opensearch_identity(base, index, index_url)
|
|
else:
|
|
request_json(index_url, "PUT", mapping.read_bytes(), "application/json")
|
|
cluster_uuid, index_uuid = opensearch_identity(base, index, index_url)
|
|
state, start_at = checkpoint_state(
|
|
input_root,
|
|
index,
|
|
checkpoint,
|
|
resume,
|
|
base,
|
|
cluster_uuid,
|
|
index_uuid,
|
|
limit,
|
|
alias,
|
|
)
|
|
documents = document_count(input_root, limit, start_at)
|
|
if not resume:
|
|
write_json_atomic(checkpoint, state)
|
|
fragments = 0
|
|
for last_code, batch in bulk_batches(bulk_pairs(input_root, index, limit, start_at), maximum_bytes):
|
|
result = request_json(
|
|
f"{base}/_bulk",
|
|
"POST",
|
|
batch,
|
|
"application/x-ndjson",
|
|
retry_invalid_json=True,
|
|
)
|
|
expected = batch.count(b"\n") // 2
|
|
items = result.get("items", [])
|
|
if result.get("errors"):
|
|
failures = [item.get("index", {}) for item in items if item.get("index", {}).get("error")]
|
|
details = "; ".join(
|
|
f"{item.get('_id', '<unknown>')}: {item['error'].get('type', 'error')}: "
|
|
f"{item['error'].get('reason', '<no reason>')}"
|
|
for item in failures[:5]
|
|
)
|
|
raise RuntimeError(
|
|
f"OpenSearch Bulk API failed for {len(failures)} item(s); "
|
|
f"resume from {checkpoint}: {details}"
|
|
)
|
|
if len(items) != expected:
|
|
raise RuntimeError(
|
|
f"OpenSearch Bulk API returned {len(items)} of {expected} item result(s); "
|
|
f"resume from {checkpoint}"
|
|
)
|
|
fragments += len(items)
|
|
state["last_document_code"] = last_code
|
|
write_json_atomic(checkpoint, state)
|
|
print(f"checkpoint={last_code} fragments={fragments}", flush=True)
|
|
if alias:
|
|
switch_alias(base, index, alias)
|
|
state["complete"] = True
|
|
write_json_atomic(checkpoint, state)
|
|
return documents, fragments
|
|
|
|
|
|
def switch_alias(base: str, index: str, alias: str) -> None:
|
|
if not alias or alias == index:
|
|
raise ValueError("--alias must differ from --index")
|
|
result = request_json(
|
|
f"{base.rstrip('/')}/_aliases",
|
|
"POST",
|
|
json.dumps({
|
|
"actions": [
|
|
{"remove": {"index": "*", "alias": alias, "must_exist": False}},
|
|
{"add": {"index": index, "alias": alias}},
|
|
]
|
|
}).encode(),
|
|
"application/json",
|
|
)
|
|
if result.get("acknowledged") is not True:
|
|
raise RuntimeError(f"OpenSearch did not acknowledge alias switch: {alias}")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", type=Path, default=Path("data/minjust-normalized"))
|
|
parser.add_argument("--output", type=Path, default=Path("data/opensearch/minjust-fragments.ndjson"))
|
|
parser.add_argument("--index", default="akyldash-fragments-v1")
|
|
parser.add_argument("--limit", type=int)
|
|
parser.add_argument("--url", help="create the index and stream bounded Bulk requests instead of writing a file")
|
|
parser.add_argument("--alias", help="atomically point this alias at --index after a successful load")
|
|
parser.add_argument("--mapping", type=Path, default=DEFAULT_MAPPING)
|
|
parser.add_argument("--batch-mb", type=int, default=25)
|
|
parser.add_argument("--resume", action="store_true", help="load into an existing index")
|
|
parser.add_argument("--checkpoint", type=Path, help="persistent resume checkpoint path")
|
|
parser.add_argument("--version", action="version", version=APP_VERSION)
|
|
arguments = parser.parse_args()
|
|
if arguments.limit is not None and arguments.limit <= 0:
|
|
raise SystemExit("--limit must be greater than zero")
|
|
if arguments.batch_mb <= 0:
|
|
raise SystemExit("--batch-mb must be greater than zero")
|
|
if arguments.resume and not arguments.url:
|
|
raise SystemExit("--resume requires --url")
|
|
if arguments.checkpoint and not arguments.url:
|
|
raise SystemExit("--checkpoint requires --url")
|
|
if arguments.alias == "":
|
|
raise SystemExit("--alias must not be empty")
|
|
if arguments.alias is not None and not arguments.url:
|
|
raise SystemExit("--alias requires --url")
|
|
if arguments.url:
|
|
checkpoint = arguments.checkpoint or Path("data/opensearch") / f"{arguments.index}.checkpoint.json"
|
|
documents, fragments = load_bulk(
|
|
arguments.input,
|
|
arguments.url,
|
|
arguments.index,
|
|
arguments.mapping,
|
|
arguments.batch_mb * 1024 * 1024,
|
|
arguments.limit,
|
|
arguments.resume,
|
|
checkpoint,
|
|
arguments.alias,
|
|
)
|
|
destination = arguments.url
|
|
else:
|
|
documents, fragments = export_bulk(arguments.input, arguments.output, arguments.index, arguments.limit)
|
|
destination = arguments.output
|
|
print(f"documents={documents} fragments={fragments} destination={destination}\nAkyldash Backend v{APP_VERSION} · Frontend — not created")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|