From 452f9d22f09c9a1bb8cff2d5534d0c850d82928c Mon Sep 17 00:00:00 2001 From: Katie Strader Date: Mon, 24 Aug 2026 15:09:21 -0700 Subject: [PATCH 1/2] fix: Fix OpenGraph batching and prevent duplicate DLT delivery --- .gitignore | 1 + benchmarks/README.md | 87 +++ benchmarks/opengraph_batching.py | 523 ++++++++++++++++++ benchmarks/opengraph_batching_results.json | 135 +++++ benchmarks/opengraph_dlt_pipeline_metrics.py | 116 ++++ .../destinations/opengraph/destination.py | 125 ++++- src/openhound/sources/opengraph/source.py | 73 ++- tests/test_opengraph_batching.py | 131 +++++ tests/test_opengraph_destination.py | 40 ++ tests/test_opengraph_destination_retry.py | 70 +++ tests/test_opengraph_dlt_integration.py | 436 +++++++++++++++ 11 files changed, 1696 insertions(+), 41 deletions(-) create mode 100644 benchmarks/README.md create mode 100644 benchmarks/opengraph_batching.py create mode 100644 benchmarks/opengraph_batching_results.json create mode 100644 benchmarks/opengraph_dlt_pipeline_metrics.py create mode 100644 tests/test_opengraph_batching.py create mode 100644 tests/test_opengraph_destination.py create mode 100644 tests/test_opengraph_destination_retry.py create mode 100644 tests/test_opengraph_dlt_integration.py diff --git a/.gitignore b/.gitignore index 0819916b..78d245ea 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ collectors/ notebooks output +.benchmark-results/ graph logs dbt_packages diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..a0416085 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,87 @@ +# OpenGraph batching benchmark + +`opengraph_batching.py` is a manual tool for comparing batching behavior. It is +separate from pytest; generated files go to git-ignored `.benchmark-results/`. + +Run commands from the repository root with the project environment: + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_batching.py --help +``` + +Each run creates a timestamped folder with `metrics.json` and generated output. +Use `--output-dir ` to change the location. + +To measure a real local DLT load, including package-file and callback metrics: + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_dlt_pipeline_metrics.py --rows 100000 +``` + +## Modes + +### Synthetic: `per-row` + +Simulates the old behavior: one conversion per input row. + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_batching.py ` + --rows 100000 --edges-per-row 1 --batch-size 150 --mode per-row +``` + +### Synthetic: `page-batched` + +Simulates the new behavior: relationships are combined across DLT-page rows, +then flushed at `--batch-size` and page end. + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_batching.py ` + --rows 100000 --edges-per-row 1 --batch-size 150 --mode page-batched +``` + +Run both modes with the same arguments for a fair comparison. Saved 100k/1m +results are in `opengraph_batching_results.json`. + +### Graph replay + +Replays final graph JSON from a collection, such as `graph/okta` or +`graph/github`. + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_batching.py ` + --graph-dir graph/okta --graph-glob 'applicationuser_fs-*.json' ` + --batch-size 150 +``` + +`--graph-glob` defaults to `*_fs-*.json`. This cannot exactly compare old/new +source batching because final graph files lack raw row/page boundaries. + +### Raw DLT JSONL replay + +Replays raw DLT JSONL through the real extension model. It can exactly compare +legacy per-row and page-batched output. Supported targets: + +- `--raw-source okta --raw-table application_users` +- `--raw-source github --raw-table org_role_members` + +Pass `--lookup-file` when the model needs the collection's DuckDB lookup file. + +```powershell +.\.venv\Scripts\python.exe benchmarks\opengraph_batching.py ` + --raw-dir output/okta --raw-source okta --raw-table application_users ` + --lookup-file lookup.duckdb --batch-size 150 --compare-source-batching +``` + +`--compare-source-batching` runs both versions, records hashes, and fails if +they differ. Normal metrics are page-batched; +`legacy_per_row_comparison` contains the baseline. + +## Reading `metrics.json` + +Key fields: `wrapper_items`, `inner_relationships`, `destination_parts`, +`maximum_relationships_per_part`, `maximum_part_bytes`, timing, peak memory, +and `flattened_semantics` (result hashes). + +The tool does not create a real DLT package, so `dlt_package_files` is `null` +and callbacks are simulated. Synthetic runs exclude model and DuckDB work; use +the results as local comparisons, not customer performance claims. diff --git a/benchmarks/opengraph_batching.py b/benchmarks/opengraph_batching.py new file mode 100644 index 00000000..60abc5a2 --- /dev/null +++ b/benchmarks/opengraph_batching.py @@ -0,0 +1,523 @@ +"""Manually benchmark OpenGraph source batching and file-destination output. + +This is deliberately not a pytest test. Run one mode at a time, for example: + + .\\.venv\\Scripts\\python.exe benchmarks\\opengraph_batching.py \ + --rows 100000 --batch-size 150 --mode per-row + +Then repeat with ``--mode page-batched``. To replay real saved graph output, use +``--graph-dir graph/okta --graph-glob 'applicationuser_fs-*.json'``. Each +invocation creates a timestamped subdirectory and writes a ``metrics.json`` +report, so prior benchmark outputs are never overwritten. +""" + +from __future__ import annotations + +import argparse +import gzip +import hashlib +import importlib +import json as stdlib_json +import threading +import time +from collections.abc import Iterable, Iterator +from datetime import UTC, datetime +from pathlib import Path +from typing import Literal, Self + +import psutil +from pydantic import PrivateAttr + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries_dataclass import Edge, EdgePath +from openhound.destinations.opengraph.destination import ( + DEST_PART, + DESTINATION_ITEM_BATCH_SIZE, + _load_items, + _write_part, +) +from openhound.sources.opengraph.source import ( + READ_JSONL_PAGE_SIZE, + _generate_graph_content, +) + +Mode = Literal["per-row", "page-batched"] + +RAW_REPLAY_MODELS = { + ("okta", "application_users"): ( + "openhound_okta.models", + "ApplicationUser", + "openhound_okta.lookup", + "OktaLookup", + ), + ("github", "org_role_members"): ( + "openhound_github.models", + "OrgRoleMember", + "openhound_github.lookup", + "GithubLookup", + ), +} + + +class SyntheticAsset(BaseAsset): + """A generic membership-like asset with a configurable edge count.""" + + row: int + edges_per_row: int + _lookup: object = PrivateAttr(default=None) + _extras: dict = PrivateAttr(default_factory=dict) + + @property + def as_node(self): + return None + + @property + def edges(self) -> Iterable[Edge]: + return ( + Edge( + kind="Benchmark_MemberOf", + start=EdgePath(match_by="id", value=f"principal-{self.row}-{edge}"), + end=EdgePath(match_by="id", value=f"group-{self.row}-{edge}"), + ) + for edge in range(self.edges_per_row) + ) + + +class PeakRss: + """Sample process RSS while a benchmark is running.""" + + def __init__(self) -> None: + self._process = psutil.Process() + self._stop = threading.Event() + self._peak = self._process.memory_info().rss + self._thread = threading.Thread(target=self._sample, daemon=True) + + def _sample(self) -> None: + while not self._stop.wait(0.02): + self._peak = max(self._peak, self._process.memory_info().rss) + + def __enter__(self) -> Self: + self._thread.start() + return self + + def __exit__(self, *_: object) -> None: + self._stop.set() + self._thread.join() + self._peak = max(self._peak, self._process.memory_info().rss) + + @property + def bytes(self) -> int: + return self._peak + + +class SemanticDigest: + """Track ordered and order-independent hashes of flattened graph content.""" + + def __init__(self) -> None: + self._sequence = hashlib.sha256() + self._individual: list[str] = [] + self.count = 0 + + def add(self, content: dict) -> None: + canonical = stdlib_json.dumps( + content, default=str, separators=(",", ":"), sort_keys=True + ).encode() + self._sequence.update(canonical + b"\n") + self._individual.append(hashlib.sha256(canonical).hexdigest()) + self.count += 1 + + def metrics(self) -> dict[str, int | str]: + multiset = hashlib.sha256("\n".join(sorted(self._individual)).encode()).hexdigest() + return { + "count": self.count, + "sequence_sha256": self._sequence.hexdigest(), + "multiset_sha256": multiset, + } + + +def _pages(rows: int, edges_per_row: int) -> Iterator[list[dict[str, int]]]: + for start in range(0, rows, READ_JSONL_PAGE_SIZE): + yield [ + {"row": row, "edges_per_row": edges_per_row} + for row in range(start, min(start + READ_JSONL_PAGE_SIZE, rows)) + ] + + +def _content( + rows: int, edges_per_row: int, batch_size: int, mode: Mode +) -> Iterator[dict]: + for page in _pages(rows, edges_per_row): + if mode == "per-row": + for row in page: + yield from _generate_graph_content([row], SyntheticAsset, batch_size) + else: + yield from _generate_graph_content(page, SyntheticAsset, batch_size) + + +def _replay_graph_content(files: Iterable[Path], batch_size: int) -> Iterator[dict]: + """Re-wrap final OpenGraph files for a destination-path replay. + + Final graph files no longer retain raw source-row or DLT-page provenance, so + this measures real relationship shape and file-destination behavior, not + source-model evaluation or an exact before/after source batching comparison. + """ + edge_parts: list[dict] = [] + for file_path in files: + document = stdlib_json.loads(file_path.read_text(encoding="utf-8")) + graph = document.get("graph", {}) + for node in graph.get("nodes", []): + yield {"graph": {"content": node, "entity_type": "node"}} + for edge in graph.get("edges", []): + edge_parts.append(edge) + if len(edge_parts) == batch_size: + yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + edge_parts = [] + if edge_parts: + yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + + +def _replay_raw_content( + files: Iterable[Path], + model: type[BaseAsset], + batch_size: int, + lookup: object | None, + row_counter: dict[str, int], + mode: Mode, +) -> Iterator[dict]: + """Replay raw DLT JSONL files through the real extension asset model.""" + + def apply_context(asset: BaseAsset) -> None: + asset._lookup = lookup + asset._extras = {} + + for file_path in files: + opener = gzip.open if file_path.suffix == ".gz" else open + with opener(file_path, "rt", encoding="utf-8") as fh: + page = [] + for line in fh: + page.append(stdlib_json.loads(line)) + if len(page) == READ_JSONL_PAGE_SIZE: + row_counter["rows"] += len(page) + if mode == "per-row": + for row in page: + yield from _generate_graph_content( + [row], model, batch_size, apply_context + ) + else: + yield from _generate_graph_content( + page, model, batch_size, apply_context + ) + page = [] + if page: + row_counter["rows"] += len(page) + if mode == "per-row": + for row in page: + yield from _generate_graph_content( + [row], model, batch_size, apply_context + ) + else: + yield from _generate_graph_content( + page, model, batch_size, apply_context + ) + + +def _raw_model_and_lookup(args: argparse.Namespace) -> tuple[type[BaseAsset], object | None]: + model_module_name, model_name, lookup_module_name, lookup_name = RAW_REPLAY_MODELS[ + (args.raw_source, args.raw_table) + ] + model = getattr(importlib.import_module(model_module_name), model_name) + if not args.lookup_file: + return model, None + + import duckdb + + connection = duckdb.connect(str(args.lookup_file), read_only=True) + lookup_class = getattr(importlib.import_module(lookup_module_name), lookup_name) + return model, lookup_class(connection) + + +def _write_normalized_load_file( + content: Iterable[dict], load_file: Path +) -> tuple[int, int, dict[str, dict[str, int | str]]]: + wrappers = 0 + relationships = 0 + nodes = SemanticDigest() + edges = SemanticDigest() + with gzip.open(load_file, "wt", encoding="utf-8") as fh: + for item in content: + wrappers += 1 + if item["graph"]["entity_type"] == "edge": + relationships += len(item["graph"]["content"]) + for edge in item["graph"]["content"]: + edges.add(edge) + else: + nodes.add(item["graph"]["content"]) + fh.write(stdlib_json.dumps(item, separators=(",", ":")) + "\n") + return wrappers, relationships, {"nodes": nodes.metrics(), "edges": edges.metrics()} + + +def _write_destination_parts(load_file: Path, output_dir: Path) -> tuple[int, int, int]: + batch: list[dict] = [] + parts = 0 + max_relationships = 0 + max_part_bytes = 0 + DEST_PART.clear() + + def flush() -> None: + nonlocal batch, parts, max_relationships, max_part_bytes + relationships = sum( + len(item["graph"]["content"]) + for item in batch + if item["graph"]["entity_type"] == "edge" + ) + _write_part(batch, "synthetic_fs", str(output_dir), "benchmark") + part_file = output_dir / f"synthetic_fs-{DEST_PART['synthetic_fs']}.json" + parts += 1 + max_relationships = max(max_relationships, relationships) + max_part_bytes = max(max_part_bytes, part_file.stat().st_size) + batch = [] + + for item in _load_items(str(load_file)): + batch.append(item) + if len(batch) == DESTINATION_ITEM_BATCH_SIZE: + flush() + if batch: + flush() + return parts, max_relationships, max_part_bytes + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int) + parser.add_argument("--batch-size", type=int, default=150) + parser.add_argument("--edges-per-row", type=int, default=1) + parser.add_argument("--mode", choices=("per-row", "page-batched")) + parser.add_argument( + "--graph-dir", + type=Path, + help="Directory containing final graph JSON files, such as graph/okta", + ) + parser.add_argument( + "--graph-glob", + default="*_fs-*.json", + help="Glob within --graph-dir to replay (default: *_fs-*.json)", + ) + parser.add_argument( + "--raw-dir", + type=Path, + help="Raw DLT output root, such as output/okta or output/github", + ) + parser.add_argument("--raw-source", choices=("okta", "github")) + parser.add_argument("--raw-table") + parser.add_argument( + "--compare-source-batching", + action="store_true", + help="Compare legacy per-row and page-batched raw conversion semantics", + ) + parser.add_argument( + "--lookup-file", + type=Path, + help="Optional DuckDB lookup file for model-dependent relationship generation", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path(".benchmark-results"), + help="Ignored directory for timestamped artifacts (default: .benchmark-results)", + ) + args = parser.parse_args() + if args.batch_size <= 0 or args.edges_per_row < 0: + parser.error("edges-per-row must be non-negative; batch-size must be positive") + if args.raw_dir: + if args.graph_dir or args.rows is not None or args.mode is not None: + parser.error("--raw-dir cannot be combined with --graph-dir, --rows, or --mode") + if not args.raw_source or not args.raw_table: + parser.error("--raw-dir requires --raw-source and --raw-table") + if (args.raw_source, args.raw_table) not in RAW_REPLAY_MODELS: + supported = ", ".join( + f"{source}/{table}" for source, table in RAW_REPLAY_MODELS + ) + parser.error(f"Unsupported raw replay target; supported: {supported}") + if not args.raw_dir.is_dir(): + parser.error(f"raw directory does not exist: {args.raw_dir}") + if args.lookup_file and not args.lookup_file.is_file(): + parser.error(f"lookup file does not exist: {args.lookup_file}") + elif args.compare_source_batching: + parser.error("--compare-source-batching requires --raw-dir") + elif args.graph_dir: + if args.rows is not None or args.mode is not None: + parser.error("--graph-dir cannot be combined with --rows or --mode") + if not args.graph_dir.is_dir(): + parser.error(f"graph directory does not exist: {args.graph_dir}") + elif args.rows is None or args.mode is None: + parser.error("synthetic runs require both --rows and --mode") + elif args.rows < 0: + parser.error("rows must be non-negative") + return args + + +def main() -> None: + args = _parse_args() + run_name = datetime.now(UTC).strftime("openhound-batching-%Y%m%dT%H%M%SZ") + run_dir = args.output_dir / run_name + suffix = 1 + while run_dir.exists(): + suffix += 1 + run_dir = args.output_dir / f"{run_name}-{suffix}" + run_dir.mkdir(parents=True) + + load_file = run_dir / "synthetic.normalized.jsonl.gz" + graph_files = [] + raw_files = [] + raw_row_counter = {"rows": 0} + if args.graph_dir: + graph_files = sorted(args.graph_dir.glob(args.graph_glob)) + if not graph_files: + raise ValueError( + f"No graph files matched {args.graph_glob!r} in {args.graph_dir}" + ) + if args.raw_dir: + raw_files = sorted((args.raw_dir / args.raw_table).glob("*.jsonl*")) + if not raw_files: + raise ValueError(f"No raw JSONL files found in {args.raw_dir / args.raw_table}") + model, lookup = _raw_model_and_lookup(args) + + legacy_semantics = None + legacy_wrappers = None + if args.compare_source_batching: + legacy_rows = {"rows": 0} + legacy_content = _replay_raw_content( + raw_files, + model, + args.batch_size, + lookup, + legacy_rows, + "per-row", + ) + legacy_wrappers = 0 + legacy_nodes = SemanticDigest() + legacy_edges = SemanticDigest() + for item in legacy_content: + legacy_wrappers += 1 + if item["graph"]["entity_type"] == "edge": + for edge in item["graph"]["content"]: + legacy_edges.add(edge) + else: + legacy_nodes.add(item["graph"]["content"]) + legacy_semantics = { + "rows": legacy_rows["rows"], + "wrapper_items": legacy_wrappers, + "nodes": legacy_nodes.metrics(), + "edges": legacy_edges.metrics(), + } + process = psutil.Process() + cpu_start = process.cpu_times() + wall_start = time.perf_counter() + with PeakRss() as peak_rss: + wrappers, relationships, semantics = _write_normalized_load_file( + _replay_raw_content( + raw_files, + model, + args.batch_size, + lookup, + raw_row_counter, + "page-batched", + ) + if args.raw_dir + else _replay_graph_content(graph_files, args.batch_size) + if args.graph_dir + else _content(args.rows, args.edges_per_row, args.batch_size, args.mode), + load_file, + ) + parts, max_relationships, max_part_bytes = _write_destination_parts( + load_file, run_dir + ) + wall_seconds = time.perf_counter() - wall_start + cpu_end = process.cpu_times() + destination_files = list(run_dir.glob("synthetic_fs-*.json")) + + metrics = { + "mode": "raw-replay" if args.raw_dir else "graph-replay" if args.graph_dir else args.mode, + "rows": raw_row_counter["rows"] if args.raw_dir else args.rows, + "edges_per_row": args.edges_per_row, + "graph_input_directory": str(args.graph_dir) if args.graph_dir else None, + "graph_input_glob": args.graph_glob if args.graph_dir else None, + "graph_input_files": len(graph_files), + "raw_input_directory": str(args.raw_dir) if args.raw_dir else None, + "raw_input_table": args.raw_table if args.raw_dir else None, + "raw_input_files": len(raw_files), + "lookup_file": str(args.lookup_file) if args.lookup_file else None, + "flattened_semantics": semantics, + "legacy_per_row_comparison": legacy_semantics, + "source_batch_size": args.batch_size, + "source_page_size": READ_JSONL_PAGE_SIZE, + "wrapper_items": wrappers, + "inner_relationships": relationships, + "normalized_dlt_items": wrappers, + "destination_callbacks": 1, + "destination_parts": parts, + "dlt_package_files": None, + "post_package_files": len(destination_files), + "maximum_relationships_per_part": max_relationships, + "maximum_part_bytes": max_part_bytes, + "normalized_compressed_bytes": load_file.stat().st_size, + "destination_uncompressed_bytes": sum(path.stat().st_size for path in destination_files), + "wall_seconds": wall_seconds, + "process_cpu_seconds": (cpu_end.user - cpu_start.user) + + (cpu_end.system - cpu_start.system), + "peak_rss_bytes": peak_rss.bytes, + "notes": [ + "Synthetic benchmark; it does not measure model evaluation, DuckDB, or DLT package creation." + if not args.graph_dir and not args.raw_dir + else "Raw replay evaluates extension model conversion but does not create a DLT load package." + if args.raw_dir + else "Graph replay uses final OpenGraph output; raw source-row and DLT-page provenance is unavailable.", + "dlt_package_files is null because this benchmark streams a synthetic normalized load file rather than running a DLT pipeline.", + ], + } + (run_dir / "metrics.json").write_text( + stdlib_json.dumps(metrics, indent=2) + "\n", encoding="utf-8" + ) + if legacy_semantics: + legacy_graph = {key: legacy_semantics[key] for key in ("nodes", "edges")} + if legacy_graph != semantics: + raise AssertionError( + "Legacy per-row and page-batched flattened graph semantics differ; " + f"see {run_dir / 'metrics.json'}" + ) + print("\nOpenGraph batching benchmark") + print(f" Mode: {metrics['mode']}") + if args.graph_dir: + print( + " Graph files / relations: " + f"{metrics['graph_input_files']:,} / {metrics['inner_relationships']:,}" + ) + else: + print( + f" Rows / relationships: {metrics['rows']:,} / " + f"{metrics['inner_relationships']:,}" + ) + print(f" Wrapper items: {metrics['wrapper_items']:,}") + if legacy_semantics: + print( + " Legacy / new wrappers: " + f"{legacy_wrappers:,} / {metrics['wrapper_items']:,} (semantics match)" + ) + print( + " Destination: " + f"{metrics['destination_callbacks']:,} callback, " + f"{metrics['destination_parts']:,} part(s)" + ) + print(f" Wall / CPU: {metrics['wall_seconds']:.3f}s / {metrics['process_cpu_seconds']:.3f}s") + print(f" Peak RSS: {metrics['peak_rss_bytes'] / 1024 / 1024:.1f} MiB") + print( + " Output bytes: " + f"{metrics['normalized_compressed_bytes']:,} compressed, " + f"{metrics['destination_uncompressed_bytes']:,} destination" + ) + print(f" Metrics: {run_dir / 'metrics.json'}") + print(f" Artifacts: {run_dir}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/opengraph_batching_results.json b/benchmarks/opengraph_batching_results.json new file mode 100644 index 00000000..9cd61b23 --- /dev/null +++ b/benchmarks/opengraph_batching_results.json @@ -0,0 +1,135 @@ +{ + "benchmark": "OpenGraph synthetic source batching", + "recorded_at": "2026-08-24", + "commit": "e343a98e525e4ca8aed8cfbf0b316151abe914e3", + "environment": { + "os": "Windows-11-10.0.26200-SP0", + "python": "3.14.2", + "dlt": "1.26.0" + }, + "configuration": { + "edges_per_row": 1, + "source_batch_size": 150, + "source_page_size": 1000, + "destination_item_batch_size": 1000 + }, + "runs": [ + { + "rows": 100000, + "mode": "per-row", + "wrappers": 100000, + "relationships": 100000, + "destination_callbacks": 1, + "destination_parts": 100, + "maximum_relationships_per_part": 1000, + "maximum_part_bytes": 184071, + "normalized_compressed_bytes": 603508, + "destination_uncompressed_bytes": 18384880, + "wall_seconds": 6.7765605, + "process_cpu_seconds": 6.65625, + "peak_rss_bytes": 114798592, + "artifact": ".benchmark-results/openhound-batching-20260824T203354Z/metrics.json" + }, + { + "rows": 100000, + "mode": "page-batched", + "wrappers": 700, + "relationships": 100000, + "destination_callbacks": 1, + "destination_parts": 1, + "maximum_relationships_per_part": 100000, + "maximum_part_bytes": 18377851, + "normalized_compressed_bytes": 602668, + "destination_uncompressed_bytes": 18377851, + "wall_seconds": 5.1618227, + "process_cpu_seconds": 5.015625, + "peak_rss_bytes": 223166464, + "artifact": ".benchmark-results/openhound-batching-20260824T203402Z/metrics.json" + }, + { + "rows": 1000000, + "mode": "per-row", + "wrappers": 1000000, + "relationships": 1000000, + "destination_callbacks": 1, + "destination_parts": 1000, + "maximum_relationships_per_part": 1000, + "maximum_part_bytes": 186071, + "normalized_compressed_bytes": 6032827, + "destination_uncompressed_bytes": 185848780, + "wall_seconds": 92.1498802, + "process_cpu_seconds": 89.796875, + "peak_rss_bytes": 283844608, + "artifact": ".benchmark-results/openhound-batching-20260824T203408Z/metrics.json" + }, + { + "rows": 1000000, + "mode": "page-batched", + "wrappers": 7000, + "relationships": 1000000, + "destination_callbacks": 1, + "destination_parts": 7, + "maximum_relationships_per_part": 142900, + "maximum_part_bytes": 26570171, + "normalized_compressed_bytes": 6023827, + "destination_uncompressed_bytes": 185778277, + "wall_seconds": 58.9026913, + "process_cpu_seconds": 57.390625, + "peak_rss_bytes": 306421760, + "artifact": ".benchmark-results/openhound-batching-20260824T203639Z/metrics.json" + } + ], + "real_dlt_pipeline_runs": [ + { + "rows": 100000, + "dlt_package_files": 2, + "destination_callbacks": 1, + "maximum_items_per_destination_callback": 700, + "maximum_bytes_per_destination_callback": 565874, + "destination_parts": 1, + "artifact": ".benchmark-results/openhound-dlt-2s2earqb/dlt_metrics.json" + }, + { + "rows": 1000000, + "dlt_package_files": 2, + "destination_callbacks": 1, + "maximum_items_per_destination_callback": 7000, + "maximum_bytes_per_destination_callback": 5643496, + "destination_parts": 7, + "artifact": ".benchmark-results/openhound-dlt-tv5d37yc/dlt_metrics.json" + } + ], + "customer_shaped_raw_replay_runs": [ + { + "source": "okta", + "table": "application_users", + "rows": 10489, + "raw_input_files": 1, + "lookup_file": "lookup.duckdb", + "relationships": 8577, + "legacy_per_row_wrappers": 4120, + "page_batched_wrappers": 63, + "destination_parts": 1, + "wall_seconds": 13.68111989996396, + "process_cpu_seconds": 13.203125, + "peak_rss_bytes": 151158784, + "semantic_hashes_match": true, + "artifact": ".benchmark-results/openhound-batching-20260824T205713Z/metrics.json" + } + ], + "manual_bloodhound_ingest_validation": { + "benchmark_artifact": ".benchmark-results/openhound-dlt-k_ausfim/dlt_metrics.json", + "rows": 1000000, + "graph_files": 7, + "bloodhound_ingest_job_id": 8, + "status": "Complete", + "failed_files": 0, + "partial_failed_files": 0, + "ingest_duration_seconds": 144.330721879 + }, + "limitations": [ + "Standalone synthetic runs do not create DLT load packages.", + "Real-DLT runs use one synthetic edge per row and do not include extension-model or DuckDB lookup work.", + "Customer-shaped raw replay evaluates a saved raw collection through the real model and DuckDB lookup, but does not perform collection API calls or create a DLT package." + ] +} diff --git a/benchmarks/opengraph_dlt_pipeline_metrics.py b/benchmarks/opengraph_dlt_pipeline_metrics.py new file mode 100644 index 00000000..66569b9f --- /dev/null +++ b/benchmarks/opengraph_dlt_pipeline_metrics.py @@ -0,0 +1,116 @@ +"""Run a real local DLT/OpenGraph load and report package and callback metrics. + +Example: + .\\.venv\\Scripts\\python.exe benchmarks\\opengraph_dlt_pipeline_metrics.py --rows 100000 +""" + +from __future__ import annotations + +import argparse +import gzip +import json +import tempfile +from collections.abc import Iterable +from pathlib import Path + +import dlt + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries_dataclass import Edge, EdgePath +from openhound.destinations.opengraph import destination as destination_module +from openhound.sources.opengraph.source import GraphResource, opengraph + + +class PipelineBenchmarkAsset(BaseAsset): + """One relationship per raw input row.""" + + row: int + + @property + def as_node(self): + return None + + @property + def edges(self) -> Iterable[Edge]: + yield Edge( + kind="PipelineBenchmarkEdge", + start=EdgePath(match_by="id", value=f"principal-{self.row}"), + end=EdgePath(match_by="id", value=f"group-{self.row}"), + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, required=True) + parser.add_argument("--batch-size", type=int, default=150) + parser.add_argument("--output-dir", type=Path, default=Path(".benchmark-results")) + args = parser.parse_args() + if args.rows < 0 or args.batch_size <= 0: + parser.error("rows must be non-negative and batch-size must be positive") + return args + + +def main() -> None: + args = _parse_args() + run_dir = Path(tempfile.mkdtemp(prefix="openhound-dlt-", dir=args.output_dir)) + raw_dir = run_dir / "raw" / "assets" + raw_dir.mkdir(parents=True) + with gzip.open(raw_dir / "rows.jsonl.gz", "wt", encoding="utf-8") as raw_file: + for row in range(args.rows): + raw_file.write(json.dumps({"row": row}) + "\n") + + callback_items: list[int] = [] + callback_bytes: list[int] = [] + original_load_items = destination_module._load_items + + def measured_load_items(file_path: str): + callback_bytes.append(Path(file_path).stat().st_size) + count = 0 + for item in original_load_items(file_path): + count += 1 + yield item + callback_items.append(count) + + destination_module._load_items = measured_load_items + try: + pipeline = dlt.pipeline( + pipeline_name="opengraph_dlt_pipeline_benchmark", + dataset_name="opengraph_dlt_pipeline_benchmark", + destination=destination_module.opengraph_file( + output_path=str(run_dir / "graph"), source_kind="benchmark" + ), + pipelines_dir=str(run_dir / "pipelines"), + ) + load_info = pipeline.run( + opengraph( + [GraphResource(table="assets", model=PipelineBenchmarkAsset)], + bucket_url=str(run_dir / "raw"), + lookup=None, + batch_size=args.batch_size, + ) + ) + finally: + destination_module._load_items = original_load_items + + completed_jobs = [ + job + for package in load_info.load_packages + for job in package.jobs["completed_jobs"] + ] + graph_parts = list((run_dir / "graph").glob("pipelinebenchmarkasset_fs-*.json")) + metrics = { + "rows": args.rows, + "source_batch_size": args.batch_size, + "dlt_package_files": len(completed_jobs), + "destination_callbacks": len(callback_items), + "maximum_items_per_destination_callback": max(callback_items, default=0), + "maximum_bytes_per_destination_callback": max(callback_bytes, default=0), + "destination_parts": len(graph_parts), + } + (run_dir / "dlt_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + print(json.dumps(metrics, indent=2)) + print(f"Metrics: {run_dir / 'dlt_metrics.json'}") + + +if __name__ == "__main__": + main() diff --git a/src/openhound/destinations/opengraph/destination.py b/src/openhound/destinations/opengraph/destination.py index 2f94e0ce..87e2decf 100644 --- a/src/openhound/destinations/opengraph/destination.py +++ b/src/openhound/destinations/opengraph/destination.py @@ -1,43 +1,61 @@ import logging +import shutil from collections import defaultdict +from collections.abc import Iterable from pathlib import Path import dlt from dlt.common import json from dlt.common.schema import TTableSchema -from dlt.common.typing import TDataItems +from dlt.common.storages.file_storage import FileStorage +from dlt.common.storages.load_package import ParsedLoadJobFileName logger = logging.getLogger(__name__) DEST_PART: defaultdict[str, int] = defaultdict(int) +DESTINATION_ITEM_BATCH_SIZE = 1000 -@dlt.destination(skip_dlt_columns_and_tables=True, batch_size=1000) -def opengraph_file( - items: TDataItems, - table: TTableSchema, - output_path: str = dlt.config.value, - source_kind: str = dlt.config.value, -): +def _load_items(file_path: str) -> Iterable[dict]: + """Read normalized JSONL directly to avoid DLT 1.26.x duplicate batches.""" + with FileStorage.open_zipsafe_ro(file_path) as fh: + for line in fh: + decoded = json.typed_loads(line) + if isinstance(decoded, dict): + yield decoded + else: + yield from decoded - table_name = table.get("name") or "opengraph" - DEST_PART[table_name] += 1 +def _write_part( + items: list[dict], + table_name: str, + output_path: str, + source_kind: str, + part_number: int | None = None, + job_file_id: str | None = None, +) -> None: + if part_number is None: + DEST_PART[table_name] += 1 + part_number = DEST_PART[table_name] + if job_file_id: + file_name = f"{table_name}-{job_file_id}-{part_number:04d}.json" + else: + file_name = f"{table_name}-{part_number}.json" nodes = [] edges = [] logger.debug( - f"Processing {len(items)} items for OpenGraph file output (part {DEST_PART[table_name]})" + "Processing %d items for OpenGraph file output (part %d)", + len(items), + part_number, ) for item in items: if item["graph"]["entity_type"] == "node": nodes.append(item["graph"]["content"]) - if item["graph"]["entity_type"] == "edge": + elif item["graph"]["entity_type"] == "edge": edges.extend(item["graph"]["content"]) - file_name = f"{table_name}-{DEST_PART[table_name]}.json" - output_dir = Path(output_path) - file_path = output_dir / file_name - + file_path = Path(output_path) / file_name with file_path.open("w", encoding="utf-8") as fh: fh.write( json.dumps( @@ -47,3 +65,78 @@ def opengraph_file( } ), ) + + +def _job_paths(file_path: str, output_path: str) -> tuple[Path, Path, str]: + """Return retry-stable paths for one DLT load job.""" + parsed = ParsedLoadJobFileName.parse(file_path) + load_id = Path(file_path).parent.parent.name + job_id = parsed.job_id() + root = Path(output_path) + staging = root / ".openhound-staging" / load_id / job_id + committed = root / ".openhound-commits" / load_id / f"{job_id}.json" + return staging, committed, parsed.file_id + + +def _publish_parts(staging: Path, committed: Path, output_path: str) -> None: + """Publish a complete staged job and record a completion marker.""" + output_dir = Path(output_path) + output_dir.mkdir(parents=True, exist_ok=True) + part_files = sorted(staging.glob("*.json")) + for part_file in part_files: + part_file.replace(output_dir / part_file.name) + + committed.parent.mkdir(parents=True, exist_ok=True) + marker = committed.with_suffix(".tmp") + marker.write_text( + json.dumps({"parts": [path.name for path in part_files]}), + encoding="utf-8", + ) + marker.replace(committed) + shutil.rmtree(staging) + + +@dlt.destination(skip_dlt_columns_and_tables=True, batch_size=0) +def opengraph_file( + items: str, + table: TTableSchema, + output_path: str = dlt.config.value, + source_kind: str = dlt.config.value, +): + + table_name = table.get("name") or "opengraph" + staging, committed, file_id = _job_paths(items, output_path) + if committed.exists(): + logger.debug("OpenGraph destination job %s was already published", file_id) + return + + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + batch = [] + part_number = 0 + for item in _load_items(items): + batch.append(item) + if len(batch) == DESTINATION_ITEM_BATCH_SIZE: + part_number += 1 + _write_part( + batch, + table_name, + str(staging), + source_kind, + part_number, + file_id, + ) + batch = [] + if batch: + part_number += 1 + _write_part( + batch, + table_name, + str(staging), + source_kind, + part_number, + file_id, + ) + _publish_parts(staging, committed, output_path) diff --git a/src/openhound/sources/opengraph/source.py b/src/openhound/sources/opengraph/source.py index eb0c56f6..e608d41d 100644 --- a/src/openhound/sources/opengraph/source.py +++ b/src/openhound/sources/opengraph/source.py @@ -1,5 +1,5 @@ +from collections.abc import Callable, Iterable from dataclasses import asdict, dataclass -from typing import Callable import dlt from dlt.sources.filesystem import filesystem as filesystemsource @@ -10,6 +10,9 @@ from .entries import GraphContent +# DLT page boundary; partial pages flush per input file. +READ_JSONL_PAGE_SIZE = 1000 + @dataclass class GraphResource: @@ -17,6 +20,44 @@ class GraphResource: model: BaseAsset +def _generate_graph_content( + resources: Iterable[dict], + model: type[BaseAsset], + batch_size: int, + apply_context: Callable | None = None, +): + """Convert one DLT page into bounded OpenGraph batches.""" + edge_parts = [] + + def serialize(content): + if hasattr(content, "model_dump"): + return content.model_dump() + return asdict(content) + + for resource in resources: + parsed_resource = model(**resource) + if apply_context: + apply_context(parsed_resource) + + as_node = parsed_resource.as_node + if as_node: + yield { + "graph": { + "content": serialize(as_node), + "entity_type": "node", + }, + } + + for edge in parsed_resource.edges or []: + edge_parts.append(serialize(edge)) + if len(edge_parts) == batch_size: + yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + edge_parts = [] + + if edge_parts: + yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + + @dlt.source(name="opengraph", max_table_nesting=0) def opengraph( graph_resources: list[GraphResource], @@ -25,6 +66,8 @@ def opengraph( extras: dict | None = None, batch_size: int = 150, ): + if batch_size <= 0: + raise ValueError("batch_size must be greater than zero") def apply_context(obj): obj._lookup = lookup @@ -37,34 +80,14 @@ def apply_context(obj): bucket_url=bucket_url, file_glob=f"{graph_resource.table}/**/*.jsonl.gz", ) - | read_jsonl() + | read_jsonl(chunksize=READ_JSONL_PAGE_SIZE) ) @dlt.transformer(parallelized=False, name=table_name, columns=GraphContent) def generate_graph(resources, model, apply_context: Callable | None = None): - for resource in resources: - parsed_resource = model(**resource) - if apply_context: - apply_context(parsed_resource) - - as_node = parsed_resource.as_node - if as_node: - yield { - "graph": { - "content": asdict(as_node), - "entity_type": "node", - }, - } - - edge_parts = [] - for edge in parsed_resource.edges: - edge_parts.append(asdict(edge)) - if len(edge_parts) >= batch_size: - yield {"graph": {"content": edge_parts, "entity_type": "edge"}} - edge_parts = [] - - if edge_parts: - yield {"graph": {"content": edge_parts, "entity_type": "edge"}} + yield from _generate_graph_content( + resources, model, batch_size, apply_context + ) yield reader | generate_graph( model=graph_resource.model, apply_context=apply_context diff --git a/tests/test_opengraph_batching.py b/tests/test_opengraph_batching.py new file mode 100644 index 00000000..2672cccd --- /dev/null +++ b/tests/test_opengraph_batching.py @@ -0,0 +1,131 @@ +import pytest +from pydantic import computed_field + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries import Node, NodeProperties +from openhound.core.models.entries_dataclass import Edge, EdgePath +from openhound.sources.opengraph.source import _generate_graph_content, opengraph + + +class _Node(Node): + value: int + + @classmethod + def guid(cls, name: str, node_type: str, *args: str) -> str: + return name + + @computed_field + @property + def id(self) -> str: + return f"node-{self.value}" + + +class _Asset(BaseAsset): + value: int + edge_count: int = 1 + has_node: bool = False + + @property + def as_node(self): + if not self.has_node: + return None + return _Node( + value=self.value, + kinds=["Test"], + properties=NodeProperties( + name=str(self.value), + displayname=str(self.value), + environmentid="test", + ), + ) + + @property + def edges(self): + return [ + Edge( + kind="TestEdge", + start=EdgePath(match_by="id", value=f"start-{self.value}-{index}"), + end=EdgePath(match_by="id", value=f"end-{self.value}-{index}"), + ) + for index in range(self.edge_count) + ] + + +def _rows(count: int): + return [{"value": index} for index in range(count)] + + +def _edges(content): + return [ + edge + for item in content + if item["graph"]["entity_type"] == "edge" + for edge in item["graph"]["content"] + ] + + +def test_batches_edges_across_successive_rows_and_preserves_order(): + content = list(_generate_graph_content(_rows(1_001), _Asset, 150)) + wrappers = [item for item in content if item["graph"]["entity_type"] == "edge"] + + assert len(wrappers) == 7 + assert [len(wrapper["graph"]["content"]) for wrapper in wrappers] == [150] * 6 + [101] + assert [edge["start"]["value"] for edge in _edges(content)] == [ + f"start-{index}-0" for index in range(1_001) + ] + + +def test_page_boundary_and_file_boundary_have_explicit_count_semantics(): + # DLT calls the transformer once per read_jsonl page; 1,001 rows are 1,000 + 1. + content = [ + *list(_generate_graph_content(_rows(1_000), _Asset, 150)), + *list(_generate_graph_content(_rows(1_001)[1_000:], _Asset, 150)), + ] + wrappers = [item for item in content if item["graph"]["entity_type"] == "edge"] + + assert len(wrappers) == 8 # ceil(1000 / 150) + ceil(1 / 150) + assert len(_edges(content)) == 1_001 + + +def test_page_accumulators_do_not_leak_between_extraction_attempts(): + first_attempt = list(_generate_graph_content(_rows(1), _Asset, 3)) + retry_attempt = list(_generate_graph_content(_rows(1), _Asset, 3)) + + assert _edges(first_attempt) == _edges(retry_attempt) + assert len(first_attempt) == len(retry_attempt) == 1 + + +def test_batch_size_one_preserves_one_wrapper_per_edge(): + content = list(_generate_graph_content(_rows(3), _Asset, 1)) + wrappers = [item for item in content if item["graph"]["entity_type"] == "edge"] + + assert [len(wrapper["graph"]["content"]) for wrapper in wrappers] == [1, 1, 1] + + +def test_mixed_assets_keep_nodes_separate_and_preserve_duplicate_edges(): + rows = [ + {"value": 0, "edge_count": 0, "has_node": True}, + {"value": 1, "edge_count": 2}, + {"value": 1, "edge_count": 2}, + {"value": 2, "edge_count": 1, "has_node": True}, + ] + content = list(_generate_graph_content(rows, _Asset, 3)) + + nodes = [item for item in content if item["graph"]["entity_type"] == "node"] + wrappers = [item for item in content if item["graph"]["entity_type"] == "edge"] + assert len(nodes) == 2 + assert all(not isinstance(item["graph"]["content"], list) for item in nodes) + assert [len(item["graph"]["content"]) for item in wrappers] == [3, 2] + assert [edge["start"]["value"] for edge in _edges(content)] == [ + "start-1-0", + "start-1-1", + "start-1-0", + "start-1-1", + "start-2-0", + ] + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_invalid_batch_sizes_are_rejected(batch_size): + with pytest.raises(ValueError, match="greater than zero"): + opengraph([], "unused", lookup=None, batch_size=batch_size) diff --git a/tests/test_opengraph_destination.py b/tests/test_opengraph_destination.py new file mode 100644 index 00000000..b3ec29bd --- /dev/null +++ b/tests/test_opengraph_destination.py @@ -0,0 +1,40 @@ +import gzip + +from dlt.common import json + +from openhound.destinations.opengraph.destination import ( + DEST_PART, + _load_items, + _write_part, +) + + +def _item(value: int): + return { + "graph": { + "entity_type": "edge", + "content": [{"kind": "Test", "start": value, "end": value + 1}], + } + } + + +def test_load_file_streaming_reads_each_non_aligned_jsonl_item_once(tmp_path): + load_file = tmp_path / "test.jsonl.gz" + expected = [_item(index) for index in range(5)] + with gzip.open(load_file, "wt", encoding="utf-8") as fh: + fh.write(json.dumps(expected[:2]) + "\n") + fh.write(json.dumps(expected[2:]) + "\n") + + assert list(_load_items(str(load_file))) == expected + + +def test_write_part_flattens_only_the_items_provided(tmp_path): + DEST_PART.clear() + _write_part([_item(1), _item(2)], "test_fs", str(tmp_path), "test") + + document = json.loads((tmp_path / "test_fs-1.json").read_text(encoding="utf-8")) + assert document["graph"]["nodes"] == [] + assert document["graph"]["edges"] == [ + {"kind": "Test", "start": 1, "end": 2}, + {"kind": "Test", "start": 2, "end": 3}, + ] diff --git a/tests/test_opengraph_destination_retry.py b/tests/test_opengraph_destination_retry.py new file mode 100644 index 00000000..391e6058 --- /dev/null +++ b/tests/test_opengraph_destination_retry.py @@ -0,0 +1,70 @@ +import json + +import dlt + +import openhound.destinations.opengraph.destination as destination_module + + +def _item(index: int) -> dict: + return { + "graph": { + "entity_type": "edge", + "content": [ + { + "kind": "Validation", + "start": {"match_by": "id", "value": f"start-{index}"}, + "end": {"match_by": "id", "value": f"end-{index}"}, + "properties": {"sequence": index}, + } + ], + } + } + + +def _edge_sequences(output_dir) -> list[int]: + values = [] + for path in sorted(output_dir.glob("validation_fs-*.json")): + document = json.loads(path.read_text(encoding="utf-8")) + values.extend(edge["properties"]["sequence"] for edge in document["graph"]["edges"]) + return values + + +def test_destination_retry_republishes_no_duplicate_parts(tmp_path, monkeypatch): + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + output_dir = tmp_path / "graph" + output_dir.mkdir() + items = [_item(index) for index in range(1_001)] + + @dlt.resource(name="validation_fs", max_table_nesting=0) + def non_aligned_source(): + for start, end in ((0, 2), (2, 155), (155, 304), (304, 1_001)): + yield items[start:end] + + original_write_part = destination_module._write_part + write_attempts = 0 + + def fail_after_first_part(*args, **kwargs): + nonlocal write_attempts + original_write_part(*args, **kwargs) + write_attempts += 1 + if write_attempts == 1: + raise RuntimeError("intentional destination failure") + + monkeypatch.setattr(destination_module, "_write_part", fail_after_first_part) + pipeline = dlt.pipeline( + pipeline_name="destination_retry_validation", + dataset_name="destination_retry_validation", + destination=destination_module.opengraph_file( + output_path=str(output_dir), source_kind="test" + ), + ) + + # DLT retries the transient destination job in this call. The destination + # must discard its staged first attempt and publish one complete part set. + pipeline.run(non_aligned_source()) + + sequences = _edge_sequences(output_dir) + assert sequences == list(range(1_001)) + assert write_attempts == 3 + assert not list((output_dir / ".openhound-staging").rglob("*.json")) + assert len(list((output_dir / ".openhound-commits").rglob("*.json"))) == 1 diff --git a/tests/test_opengraph_dlt_integration.py b/tests/test_opengraph_dlt_integration.py new file mode 100644 index 00000000..6d24b3c7 --- /dev/null +++ b/tests/test_opengraph_dlt_integration.py @@ -0,0 +1,436 @@ +import gzip +import json as stdlib_json +import os +import subprocess +import sys +import textwrap +from collections import defaultdict +from pathlib import Path + +import dlt +import pytest +from dlt.common import json +from dlt.common.storages.file_storage import FileStorage + +from openhound.core.asset import BaseAsset +from openhound.core.models.entries_dataclass import Edge, EdgePath +from openhound.destinations.opengraph import destination as file_destination +from openhound.sources.opengraph.source import GraphResource, opengraph + + +class _RawEdgeAsset(BaseAsset): + row: int + + @property + def as_node(self): + return None + + @property + def edges(self): + yield Edge( + kind="IntegrationEdge", + start=EdgePath(match_by="id", value=f"start-{self.row}"), + end=EdgePath(match_by="id", value=f"end-{self.row}"), + ) + + +class _SecondRawEdgeAsset(_RawEdgeAsset): + @property + def edges(self): + yield Edge( + kind="SecondIntegrationEdge", + start=EdgePath(match_by="id", value=f"second-start-{self.row}"), + end=EdgePath(match_by="id", value=f"second-end-{self.row}"), + ) + + +class _VariableEdgeAsset(BaseAsset): + row: int + edge_count: int = 1 + + @property + def as_node(self): + return None + + @property + def edges(self): + for edge in range(self.edge_count): + yield Edge( + kind="VariableIntegrationEdge", + start=EdgePath( + match_by="id", value=f"variable-start-{self.row}-{edge}" + ), + end=EdgePath(match_by="id", value=f"variable-end-{self.row}-{edge}"), + ) + + +def _write_jsonl(path, rows): + with gzip.open(path, "wt", encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row) + "\n") + + +def test_opengraph_dlt_batches_across_rows_and_flushes_per_input_file( + monkeypatch, tmp_path +): + """Exercise the DLT filesystem reader, not just the batching helper. + + The first file crosses DLT's 1,000-row read_jsonl page boundary and the + second is a separate input file. With a page-scoped batch size of 150, + this must produce ceil(1000 / 150) + ceil(1 / 150) == 8 edge wrappers. + """ + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + raw_dir = tmp_path / "raw" / "assets" + raw_dir.mkdir(parents=True) + _write_jsonl(raw_dir / "one.jsonl.gz", ({"row": row} for row in range(1_000))) + _write_jsonl(raw_dir / "two.jsonl.gz", [{"row": 1_000}]) + + captured = [] + + @dlt.destination(skip_dlt_columns_and_tables=True, batch_size=0) + def capture_normalized_file(items: str, table): + with FileStorage.open_zipsafe_ro(items) as fh: + for line in fh: + decoded = json.typed_loads(line) + captured.extend([decoded] if isinstance(decoded, dict) else decoded) + + pipeline = dlt.pipeline( + pipeline_name="opengraph_dlt_page_boundary", + dataset_name="opengraph_dlt_page_boundary", + destination=capture_normalized_file(), + pipelines_dir=str(tmp_path / "pipelines"), + ) + pipeline.run( + opengraph( + [GraphResource(table="assets", model=_RawEdgeAsset)], + bucket_url=str(tmp_path / "raw"), + lookup=None, + batch_size=150, + ) + ) + + edge_wrappers = [ + item for item in captured if item["graph"]["entity_type"] == "edge" + ] + flattened = [ + edge for wrapper in edge_wrappers for edge in wrapper["graph"]["content"] + ] + assert sorted(len(wrapper["graph"]["content"]) for wrapper in edge_wrappers) == [ + 1, + 100, + 150, + 150, + 150, + 150, + 150, + 150, + ] + assert sorted( + int(edge["start"]["value"].removeprefix("start-")) for edge in flattened + ) == list(range(1_001)) + + +def test_opengraph_dlt_keeps_multiple_graph_resources_isolated(monkeypatch, tmp_path): + """Each GraphResource receives a fresh table-specific accumulator.""" + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + raw_root = tmp_path / "raw" + for table_name, rows in {"first": range(3), "second": range(10, 13)}.items(): + table_dir = raw_root / table_name + table_dir.mkdir(parents=True) + _write_jsonl(table_dir / "rows.jsonl.gz", ({"row": row} for row in rows)) + + captured: defaultdict[str, list[dict]] = defaultdict(list) + + @dlt.destination(skip_dlt_columns_and_tables=True, batch_size=0) + def capture_normalized_file(items: str, table): + with FileStorage.open_zipsafe_ro(items) as fh: + for line in fh: + decoded = json.typed_loads(line) + captured[table["name"]].extend( + [decoded] if isinstance(decoded, dict) else decoded + ) + + pipeline = dlt.pipeline( + pipeline_name="opengraph_dlt_multiple_resources", + dataset_name="opengraph_dlt_multiple_resources", + destination=capture_normalized_file(), + pipelines_dir=str(tmp_path / "pipelines"), + ) + pipeline.run( + opengraph( + [ + GraphResource(table="first", model=_RawEdgeAsset), + GraphResource(table="second", model=_SecondRawEdgeAsset), + ], + bucket_url=str(raw_root), + lookup=None, + batch_size=2, + ) + ) + + first = captured["_rawedgeasset_fs"] + second = captured["_secondrawedgeasset_fs"] + assert [len(item["graph"]["content"]) for item in first] == [2, 1] + assert [len(item["graph"]["content"]) for item in second] == [2, 1] + assert [ + edge["start"]["value"] for item in first for edge in item["graph"]["content"] + ] == ["start-0", "start-1", "start-2"] + assert [ + edge["start"]["value"] for item in second for edge in item["graph"]["content"] + ] == ["second-start-10", "second-start-11", "second-start-12"] + + +@pytest.mark.parametrize( + ("rows", "batch_size", "expected_lengths", "expected_starts"), + [ + (None, 3, [], []), + ( + [{"row": 0, "edge_count": 0}, {"row": 1, "edge_count": 0}], + 3, + [], + [], + ), + ( + [{"row": 0}, {"row": 1}, {"row": 2}], + 3, + [3], + ["variable-start-0-0", "variable-start-1-0", "variable-start-2-0"], + ), + ( + [{"row": 0}, {"row": 1}, {"row": 2}, {"row": 3}], + 3, + [3, 1], + [ + "variable-start-0-0", + "variable-start-1-0", + "variable-start-2-0", + "variable-start-3-0", + ], + ), + ( + [ + {"row": 0, "edge_count": 2}, + {"row": 1, "edge_count": 2}, + {"row": 2, "edge_count": 1}, + ], + 3, + [3, 2], + [ + "variable-start-0-0", + "variable-start-0-1", + "variable-start-1-0", + "variable-start-1-1", + "variable-start-2-0", + ], + ), + ], +) +def test_opengraph_dlt_boundary_cases( + monkeypatch, tmp_path, rows, batch_size, expected_lengths, expected_starts +): + """Exercise boundary cases through the DLT filesystem and load pipeline.""" + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + raw_dir = tmp_path / "raw" / "assets" + raw_dir.mkdir(parents=True) + if rows is not None: + _write_jsonl(raw_dir / "rows.jsonl.gz", rows) + + captured: list[dict] = [] + + @dlt.destination(skip_dlt_columns_and_tables=True, batch_size=0) + def capture_normalized_file(items: str, table): + if table["name"] != "_variableedgeasset_fs": + return + with FileStorage.open_zipsafe_ro(items) as fh: + for line in fh: + decoded = json.typed_loads(line) + captured.extend([decoded] if isinstance(decoded, dict) else decoded) + + pipeline = dlt.pipeline( + pipeline_name="opengraph_dlt_boundary_cases", + dataset_name="opengraph_dlt_boundary_cases", + destination=capture_normalized_file(), + pipelines_dir=str(tmp_path / "pipelines"), + ) + pipeline.run( + opengraph( + [GraphResource(table="assets", model=_VariableEdgeAsset)], + bucket_url=str(tmp_path / "raw"), + lookup=None, + batch_size=batch_size, + ) + ) + + edge_wrappers = [ + item for item in captured if item["graph"]["entity_type"] == "edge" + ] + assert [ + len(wrapper["graph"]["content"]) for wrapper in edge_wrappers + ] == expected_lengths + assert [ + edge["start"]["value"] + for wrapper in edge_wrappers + for edge in wrapper["graph"]["content"] + ] == expected_starts + + +def test_opengraph_file_destination_retries_staged_load_without_duplicates( + monkeypatch, tmp_path +): + """A failed destination publish must resume without source re-extraction. + + The first load attempt writes all destination parts to the job staging area, + then fails before publishing them. DLT retries the pending load job from + the normalized package. The retry must publish one complete file, not lose + pending edges or create a second delivery. + """ + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + raw_dir = tmp_path / "raw" / "assets" + raw_dir.mkdir(parents=True) + _write_jsonl(raw_dir / "assets.jsonl.gz", ({"row": row} for row in range(1_001))) + + output_dir = tmp_path / "graph" + original_publish = file_destination._publish_parts + fail_once = True + + def fail_after_staging(staging, committed, output_path): + nonlocal fail_once + if fail_once: + fail_once = False + raise RuntimeError("simulated destination publish failure") + original_publish(staging, committed, output_path) + + monkeypatch.setattr(file_destination, "_publish_parts", fail_after_staging) + pipeline = dlt.pipeline( + pipeline_name="opengraph_dlt_restart", + dataset_name="opengraph_dlt_restart", + destination=file_destination.opengraph_file( + output_path=str(output_dir), source_kind="test" + ), + pipelines_dir=str(tmp_path / "pipelines"), + ) + source = opengraph( + [GraphResource(table="assets", model=_RawEdgeAsset)], + bucket_url=str(tmp_path / "raw"), + lookup=None, + batch_size=150, + ) + + pipeline.run(source) + + assert fail_once is False + + published_parts = list(output_dir.glob("_rawedgeasset_fs-*.json")) + assert len(published_parts) == 1 + document = stdlib_json.loads(published_parts[0].read_text(encoding="utf-8")) + assert len(document["graph"]["edges"]) == 1_001 + assert {edge["start"]["value"] for edge in document["graph"]["edges"]} == { + f"start-{row}" for row in range(1_001) + } + + +def test_opengraph_file_destination_cold_restart_resumes_pending_job(tmp_path): + """A new Python process resumes a failed staged destination load exactly once.""" + raw_dir = tmp_path / "raw" / "assets" + raw_dir.mkdir(parents=True) + _write_jsonl(raw_dir / "assets.jsonl.gz", ({"row": row} for row in range(5))) + + runner = tmp_path / "cold_restart_runner.py" + runner.write_text( + textwrap.dedent( + """ + import os + import sys + + import dlt + from openhound.core.asset import BaseAsset + from openhound.core.models.entries_dataclass import Edge, EdgePath + from openhound.destinations.opengraph import destination as file_destination + from openhound.sources.opengraph.source import GraphResource, opengraph + + + class Asset(BaseAsset): + row: int + + @property + def as_node(self): + return None + + @property + def edges(self): + yield Edge( + kind="ColdRestartEdge", + start=EdgePath(match_by="id", value=f"start-{self.row}"), + end=EdgePath(match_by="id", value=f"end-{self.row}"), + ) + + + def main(): + if os.environ.get("FAIL_DESTINATION_PUBLISH"): + def fail_publish(staging, committed, output_path): + raise RuntimeError("simulated persistent publish failure") + + file_destination._publish_parts = fail_publish + + pipeline = dlt.pipeline( + pipeline_name="opengraph_cold_restart", + dataset_name="opengraph_cold_restart", + destination=file_destination.opengraph_file( + output_path=os.environ["GRAPH_OUTPUT"], source_kind="test" + ), + pipelines_dir=os.environ["PIPELINES_DIR"], + ) + if sys.argv[1] == "initial": + pipeline.run( + opengraph( + [GraphResource(table="assets", model=Asset)], + bucket_url=os.environ["RAW_ROOT"], + lookup=None, + batch_size=3, + ) + ) + else: + pipeline.run() + + + if __name__ == "__main__": + main() + """ + ), + encoding="utf-8", + ) + project_root = Path(__file__).resolve().parents[1] + environment = { + **os.environ, + "DLT_DATA_DIR": str(tmp_path / ".dlt"), + "GRAPH_OUTPUT": str(tmp_path / "graph"), + "PIPELINES_DIR": str(tmp_path / "pipelines"), + "RAW_ROOT": str(tmp_path / "raw"), + "PYTHONPATH": str(project_root / "src"), + } + failed = subprocess.run( + [sys.executable, str(runner), "initial"], + cwd=project_root, + env={**environment, "FAIL_DESTINATION_PUBLISH": "1"}, + capture_output=True, + check=False, + text=True, + ) + assert failed.returncode != 0, failed.stderr + + resumed = subprocess.run( + [sys.executable, str(runner), "resume"], + cwd=project_root, + env=environment, + capture_output=True, + check=False, + text=True, + ) + assert resumed.returncode == 0, resumed.stderr + + published_parts = list((tmp_path / "graph").glob("asset_fs-*.json")) + assert len(published_parts) == 1 + document = stdlib_json.loads(published_parts[0].read_text(encoding="utf-8")) + assert [edge["start"]["value"] for edge in document["graph"]["edges"]] == [ + f"start-{row}" for row in range(5) + ] From 6b2e08eb946717c68941ad1daecc9ba187b503b1 Mon Sep 17 00:00:00 2001 From: Katie Strader Date: Mon, 24 Aug 2026 15:22:22 -0700 Subject: [PATCH 2/2] fix: Create the output directory before mkdtemp --- benchmarks/opengraph_dlt_pipeline_metrics.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/opengraph_dlt_pipeline_metrics.py b/benchmarks/opengraph_dlt_pipeline_metrics.py index 66569b9f..f2189a6c 100644 --- a/benchmarks/opengraph_dlt_pipeline_metrics.py +++ b/benchmarks/opengraph_dlt_pipeline_metrics.py @@ -52,6 +52,7 @@ def _parse_args() -> argparse.Namespace: def main() -> None: args = _parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) run_dir = Path(tempfile.mkdtemp(prefix="openhound-dlt-", dir=args.output_dir)) raw_dir = run_dir / "raw" / "assets" raw_dir.mkdir(parents=True)