From 1edc36db5d3fc8fc6adab8c9aa0e105a820a744d Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Fri, 28 Aug 2026 14:28:57 +0200 Subject: [PATCH 1/3] Make source_kind optional, remove it from the OG metadata and append it to the node's list of kinds instead --- src/openhound/core/app.py | 185 +++++++++--------- src/openhound/core/app.pyi | 7 +- src/openhound/core/convert.py | 9 +- .../bloodhound_enterprise/destination.py | 11 +- .../destinations/opengraph/destination.py | 24 +-- src/openhound/sources/opengraph/source.py | 35 ++-- tests/test_opengraph_batching.py | 19 ++ tests/test_opengraph_destination.py | 3 +- tests/test_opengraph_destination_retry.py | 4 +- tests/test_opengraph_dlt_integration.py | 6 +- 10 files changed, 154 insertions(+), 149 deletions(-) diff --git a/src/openhound/core/app.py b/src/openhound/core/app.py index 526a49ce..c028548d 100644 --- a/src/openhound/core/app.py +++ b/src/openhound/core/app.py @@ -1,8 +1,3 @@ -import logging -from enum import Enum -from pathlib import Path -from typing import Annotated, Callable, List - import dlt import duckdb import typer @@ -10,7 +5,11 @@ from dlt.common.pipeline import LoadInfo from dlt.extract.resource import DltResource from dlt.extract.source import DltSource +from enum import Enum +from pathlib import Path +from typing import Annotated, Callable, List +import logging from openhound.cli.collect import collect from openhound.cli.convert import convert from openhound.cli.preproc import preprocess @@ -52,7 +51,7 @@ class Contract(str, Enum): class OpenHound: - def __init__(self, name: str, source_kind: str, help: str = "OpenGraph collector"): + def __init__(self, name: str, source_kind: str | None = None, help: str = "OpenGraph collector"): dlt_pydantic.create_list_model = validate.create_list_model dlt_pydantic._classify_validation_errors = validate._classify_validation_errors @@ -78,9 +77,9 @@ def __init__(self, name: str, source_kind: str, help: str = "OpenGraph collector self.edges: list[EdgeDef] = [] def collect( - self, - help: str = "OpenGraph collect pipeline", - **kwargs, + self, + help: str = "OpenGraph collect pipeline", + **kwargs, ): """Register a Typer CLI command that collects resources and stores them (filtered) on disk. @@ -92,29 +91,29 @@ def collect( def decorator(func: Callable): def wrapper( - output_path: OutputPath, - resources: List[str] = typer.Argument(None), - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" - ), - tables_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when data contains newly seen resources/tables previously not collected", + output_path: OutputPath, + resources: List[str] = typer.Argument(None), + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" ), - ] = Contract.evolve, - columns_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when data contains values/keys not found in the Pydantic model", - ), - ] = Contract.evolve, - data_type_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when fields do not match the data types defined in the Pydantic model", - ), - ] = Contract.discard_row, + tables_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when data contains newly seen resources/tables previously not collected", + ), + ] = Contract.evolve, + columns_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when data contains values/keys not found in the Pydantic model", + ), + ] = Contract.evolve, + data_type_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when fields do not match the data types defined in the Pydantic model", + ), + ] = Contract.discard_row, ) -> LoadInfo | None: schema_contract = { "tables": tables_contract, @@ -142,10 +141,10 @@ def wrapper( return decorator def convert( - self, - lookup: Callable | None = None, - help: str = "OpenGraph convert pipeline", - **typer_kwargs, + self, + lookup: Callable | None = None, + help: str = "OpenGraph convert pipeline", + **typer_kwargs, ): """Register a Typer CLI command that converts collected resources to OpenGraph nodes and edges. @@ -158,11 +157,11 @@ def convert( def decorator(func: Callable): def run_convert( - input_path: InputPath, - output_path: Path = Path("/tmp/openhound"), - lookup_file: Path = DEFAULT_LOOKUP_FILE, - progress: Progress = Progress.tqdm, - method: Method = Method.write, + input_path: InputPath, + output_path: Path = Path("/tmp/openhound"), + lookup_file: Path = DEFAULT_LOOKUP_FILE, + progress: Progress = Progress.tqdm, + method: Method = Method.write, ) -> LoadInfo: lookup_session = None if lookup: @@ -197,31 +196,31 @@ def run_convert( ) def wrapper( - input_path: InputPath, - output_path: Annotated[ - Path, - typer.Argument( - exists=False, - file_okay=False, - dir_okay=True, - resolve_path=True, - help="Output path to write OpenGraph JSON files", - ), - ], - # resources: List[str] = typer.Argument(None), - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" - ), - lookup_file: Annotated[ - Path, - typer.Option( - file_okay=True, - dir_okay=False, - readable=True, - resolve_path=True, - help="DuckDB lookup file path", + input_path: InputPath, + output_path: Annotated[ + Path, + typer.Argument( + exists=False, + file_okay=False, + dir_okay=True, + resolve_path=True, + help="Output path to write OpenGraph JSON files", + ), + ], + # resources: List[str] = typer.Argument(None), + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" ), - ] = DEFAULT_LOOKUP_FILE, + lookup_file: Annotated[ + Path, + typer.Option( + file_okay=True, + dir_okay=False, + readable=True, + resolve_path=True, + help="DuckDB lookup file path", + ), + ] = DEFAULT_LOOKUP_FILE, ) -> LoadInfo: return run_convert( input_path=input_path, @@ -241,10 +240,10 @@ def wrapper( return decorator def preproc( - self, - transformer: Callable[[any], None] | None = None, - help: str = "OpenGraph preprocessing pipeline", - **typer_kwargs, + self, + transformer: Callable[[any], None] | None = None, + help: str = "OpenGraph preprocessing pipeline", + **typer_kwargs, ): """Register a Typer CLI command that performs optional preprocessing and builds lookup data for a source. @@ -258,19 +257,19 @@ def decorator(func: Callable): self.preprocessor = func def wrapper( - input_path: InputPath, - output_file: Annotated[ - Path, - typer.Argument( - file_okay=True, - dir_okay=False, - readable=True, - resolve_path=True, + input_path: InputPath, + output_file: Annotated[ + Path, + typer.Argument( + file_okay=True, + dir_okay=False, + readable=True, + resolve_path=True, + ), + ] = DEFAULT_LOOKUP_FILE, + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" ), - ] = DEFAULT_LOOKUP_FILE, - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" - ), ) -> LoadInfo: preprocessor = PreProcessor( name=self.name, @@ -299,9 +298,9 @@ def defer(self, func: Callable) -> Callable: return dlt.defer(safe_func) def transformer( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT transformer with added exception handling.""" @@ -316,9 +315,9 @@ def decorator(func: Callable) -> DltResource: return decorator def resource( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT resource with added exception handling.""" @@ -333,9 +332,9 @@ def decorator(func: Callable) -> DltResource: return decorator def source( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT source with added exception handling. @@ -353,10 +352,10 @@ def decorator(func: Callable) -> Callable: return decorator def asset( - self, - node: NodeDef | None = None, - edges: list[EdgeDef] | None = None, - description: str = "Resource model for OpenGraph", + self, + node: NodeDef | None = None, + edges: list[EdgeDef] | None = None, + description: str = "Resource model for OpenGraph", ): """Decorator to register a resource class and its graph definitions (nodes/edges). This is used to automatically generate documentation for each unique resource and implement rules/warnings when nodes/edges are returned diff --git a/src/openhound/core/app.pyi b/src/openhound/core/app.pyi index 1720a40d..9e45a21f 100644 --- a/src/openhound/core/app.pyi +++ b/src/openhound/core/app.pyi @@ -48,7 +48,7 @@ class Asset: class OpenHound: name: str - source_kind: str + source_kind: str | None = None help: str metadata: Extension | None collector: Callable | None @@ -65,7 +65,10 @@ class OpenHound: edges: list[EdgeDef] def __init__( - self, name: str, source_kind: str, help: str = "OpenGraph collector" + self, + name: str, + source_kind: str | None = None, + help: str = "OpenGraph collector", ): ... def icons(self, color: str, help: str = "BloodHound icons sync"): ... def queries(self, help: str = "BloodHound icons sync"): ... diff --git a/src/openhound/core/convert.py b/src/openhound/core/convert.py index 5af52e12..5df942d5 100644 --- a/src/openhound/core/convert.py +++ b/src/openhound/core/convert.py @@ -33,7 +33,7 @@ def __init__( input_path: Path, lookup: LookupManager, output_path: Path, - source_kind: str, + source_kind: str | None = None, progress: Progress = Progress.tqdm, method: Method = Method.write, ): @@ -51,13 +51,11 @@ def pipeline(self) -> Pipeline: logger.debug( "Initializing BloodHound Enterprise client for converter ingest method" ) - dest = ingest(source_kind=self.source_kind) + dest = ingest() else: logger.debug("Using file output method for converter") - dest = opengraph_file( - output_path=str(self.output_path), source_kind=self.source_kind - ) + dest = opengraph_file(output_path=str(self.output_path)) pipeline = dlt.pipeline( pipeline_name=f"{self.name}_convert", @@ -99,6 +97,7 @@ def run( lookup=self.lookup, bucket_url=str(self.input_path), extras=extra_context, + source_kind=self.source_kind, ) ) diff --git a/src/openhound/destinations/bloodhound_enterprise/destination.py b/src/openhound/destinations/bloodhound_enterprise/destination.py index 725f4d7a..1ee952a0 100644 --- a/src/openhound/destinations/bloodhound_enterprise/destination.py +++ b/src/openhound/destinations/bloodhound_enterprise/destination.py @@ -19,9 +19,7 @@ def ingest( url: str = dlt.config.value, token_id: str = dlt.secrets.value, token_key: str = dlt.secrets.value, - source_kind: str = dlt.config.value, ): - client = BloodHoundEnterprise(token_key=token_key, token_id=token_id, bhe_uri=url) nodes = [] @@ -34,12 +32,5 @@ def ingest( if item["graph"]["entity_type"] == "edge": edges.extend(item["graph"]["content"]) - client.ingest( - json.dumps( - { - "graph": {"nodes": nodes, "edges": edges}, - "metadata": {"source_kind": source_kind}, - } - ) - ) + client.ingest(json.dumps({"graph": {"nodes": nodes, "edges": edges}})) logger.info("Graph ingested to BloodHound Enterprise") diff --git a/src/openhound/destinations/opengraph/destination.py b/src/openhound/destinations/opengraph/destination.py index 87e2decf..0e5d5630 100644 --- a/src/openhound/destinations/opengraph/destination.py +++ b/src/openhound/destinations/opengraph/destination.py @@ -28,12 +28,11 @@ def _load_items(file_path: str) -> Iterable[dict]: 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, + items: list[dict], + table_name: str, + output_path: str, + part_number: int | None = None, + job_file_id: str | None = None, ) -> None: if part_number is None: DEST_PART[table_name] += 1 @@ -60,8 +59,7 @@ def _write_part( fh.write( json.dumps( { - "graph": {"nodes": nodes, "edges": edges}, - "metadata": {"source_kind": source_kind}, + "graph": {"nodes": nodes, "edges": edges} } ), ) @@ -98,12 +96,10 @@ def _publish_parts(staging: Path, committed: Path, output_path: str) -> None: @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, + items: str, + table: TTableSchema, + output_path: str = dlt.config.value, ): - table_name = table.get("name") or "opengraph" staging, committed, file_id = _job_paths(items, output_path) if committed.exists(): @@ -124,7 +120,6 @@ def opengraph_file( batch, table_name, str(staging), - source_kind, part_number, file_id, ) @@ -135,7 +130,6 @@ def opengraph_file( batch, table_name, str(staging), - source_kind, part_number, file_id, ) diff --git a/src/openhound/sources/opengraph/source.py b/src/openhound/sources/opengraph/source.py index e608d41d..dd3803a0 100644 --- a/src/openhound/sources/opengraph/source.py +++ b/src/openhound/sources/opengraph/source.py @@ -7,7 +7,6 @@ from openhound.core.asset import BaseAsset from openhound.core.lookup import LookupManager - from .entries import GraphContent # DLT page boundary; partial pages flush per input file. @@ -21,10 +20,11 @@ class GraphResource: def _generate_graph_content( - resources: Iterable[dict], - model: type[BaseAsset], - batch_size: int, - apply_context: Callable | None = None, + resources: Iterable[dict], + model: type[BaseAsset], + batch_size: int, + apply_context: Callable | None = None, + source_kind: str | None = None ): """Convert one DLT page into bounded OpenGraph batches.""" edge_parts = [] @@ -41,6 +41,8 @@ def serialize(content): as_node = parsed_resource.as_node if as_node: + if source_kind is not None and source_kind not in as_node.kinds: + as_node.kinds.append(source_kind) yield { "graph": { "content": serialize(as_node), @@ -60,11 +62,12 @@ def serialize(content): @dlt.source(name="opengraph", max_table_nesting=0) def opengraph( - graph_resources: list[GraphResource], - bucket_url: str, - lookup: LookupManager, - extras: dict | None = None, - batch_size: int = 150, + graph_resources: list[GraphResource], + bucket_url: str, + lookup: LookupManager, + extras: dict | None = None, + batch_size: int = 150, + source_kind: str | None = None ): if batch_size <= 0: raise ValueError("batch_size must be greater than zero") @@ -76,17 +79,17 @@ def apply_context(obj): for graph_resource in graph_resources: table_name = f"{graph_resource.model.__name__.lower()}_fs" reader = ( - filesystemsource( - bucket_url=bucket_url, - file_glob=f"{graph_resource.table}/**/*.jsonl.gz", - ) - | read_jsonl(chunksize=READ_JSONL_PAGE_SIZE) + filesystemsource( + bucket_url=bucket_url, + file_glob=f"{graph_resource.table}/**/*.jsonl.gz", + ) + | 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): yield from _generate_graph_content( - resources, model, batch_size, apply_context + resources, model, batch_size, apply_context, source_kind=source_kind ) yield reader | generate_graph( diff --git a/tests/test_opengraph_batching.py b/tests/test_opengraph_batching.py index 2672cccd..c07ed14b 100644 --- a/tests/test_opengraph_batching.py +++ b/tests/test_opengraph_batching.py @@ -64,6 +64,25 @@ def _edges(content): ] +def test_source_kind_is_appended_to_emitted_node_kinds(): + content = list( + _generate_graph_content( + [{"value": 1, "edge_count": 0, "has_node": True}], + _Asset, + 1, + source_kind="Test_Source", + ) + ) + + node = next( + item["graph"]["content"] + for item in content + if item["graph"]["entity_type"] == "node" + ) + + assert node["kinds"] == ["Test", "Test_Source"] + + 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"] diff --git a/tests/test_opengraph_destination.py b/tests/test_opengraph_destination.py index b3ec29bd..42340297 100644 --- a/tests/test_opengraph_destination.py +++ b/tests/test_opengraph_destination.py @@ -30,7 +30,7 @@ def test_load_file_streaming_reads_each_non_aligned_jsonl_item_once(tmp_path): 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") + _write_part([_item(1), _item(2)], "test_fs", str(tmp_path)) document = json.loads((tmp_path / "test_fs-1.json").read_text(encoding="utf-8")) assert document["graph"]["nodes"] == [] @@ -38,3 +38,4 @@ def test_write_part_flattens_only_the_items_provided(tmp_path): {"kind": "Test", "start": 1, "end": 2}, {"kind": "Test", "start": 2, "end": 3}, ] + assert "metadata" not in document diff --git a/tests/test_opengraph_destination_retry.py b/tests/test_opengraph_destination_retry.py index 391e6058..1de90380 100644 --- a/tests/test_opengraph_destination_retry.py +++ b/tests/test_opengraph_destination_retry.py @@ -54,9 +54,7 @@ def fail_after_first_part(*args, **kwargs): 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" - ), + destination=destination_module.opengraph_file(output_path=str(output_dir)), ) # DLT retries the transient destination job in this call. The destination diff --git a/tests/test_opengraph_dlt_integration.py b/tests/test_opengraph_dlt_integration.py index 6d24b3c7..ae11e601 100644 --- a/tests/test_opengraph_dlt_integration.py +++ b/tests/test_opengraph_dlt_integration.py @@ -304,9 +304,7 @@ def fail_after_staging(staging, committed, output_path): 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" - ), + destination=file_destination.opengraph_file(output_path=str(output_dir)), pipelines_dir=str(tmp_path / "pipelines"), ) source = opengraph( @@ -376,7 +374,7 @@ def fail_publish(staging, committed, output_path): pipeline_name="opengraph_cold_restart", dataset_name="opengraph_cold_restart", destination=file_destination.opengraph_file( - output_path=os.environ["GRAPH_OUTPUT"], source_kind="test" + output_path=os.environ["GRAPH_OUTPUT"] ), pipelines_dir=os.environ["PIPELINES_DIR"], ) From 7f3689eb0e151df0df99f576a5dcdee8aecf43f0 Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Fri, 28 Aug 2026 14:34:28 +0200 Subject: [PATCH 2/3] Formatting --- src/openhound/core/app.py | 190 ++++++++++++++++++++------------------ 1 file changed, 98 insertions(+), 92 deletions(-) diff --git a/src/openhound/core/app.py b/src/openhound/core/app.py index c028548d..b44e6754 100644 --- a/src/openhound/core/app.py +++ b/src/openhound/core/app.py @@ -1,3 +1,8 @@ +import logging +from enum import Enum +from pathlib import Path +from typing import Annotated, Callable, List + import dlt import duckdb import typer @@ -5,11 +10,7 @@ from dlt.common.pipeline import LoadInfo from dlt.extract.resource import DltResource from dlt.extract.source import DltSource -from enum import Enum -from pathlib import Path -from typing import Annotated, Callable, List -import logging from openhound.cli.collect import collect from openhound.cli.convert import convert from openhound.cli.preproc import preprocess @@ -51,7 +52,12 @@ class Contract(str, Enum): class OpenHound: - def __init__(self, name: str, source_kind: str | None = None, help: str = "OpenGraph collector"): + def __init__( + self, + name: str, + source_kind: str | None = None, + help: str = "OpenGraph collector", + ): dlt_pydantic.create_list_model = validate.create_list_model dlt_pydantic._classify_validation_errors = validate._classify_validation_errors @@ -77,9 +83,9 @@ def __init__(self, name: str, source_kind: str | None = None, help: str = "OpenG self.edges: list[EdgeDef] = [] def collect( - self, - help: str = "OpenGraph collect pipeline", - **kwargs, + self, + help: str = "OpenGraph collect pipeline", + **kwargs, ): """Register a Typer CLI command that collects resources and stores them (filtered) on disk. @@ -91,29 +97,29 @@ def collect( def decorator(func: Callable): def wrapper( - output_path: OutputPath, - resources: List[str] = typer.Argument(None), - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" + output_path: OutputPath, + resources: List[str] = typer.Argument(None), + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" + ), + tables_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when data contains newly seen resources/tables previously not collected", ), - tables_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when data contains newly seen resources/tables previously not collected", - ), - ] = Contract.evolve, - columns_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when data contains values/keys not found in the Pydantic model", - ), - ] = Contract.evolve, - data_type_contract: Annotated[ - Contract, - typer.Option( - help="DLT contract applied when fields do not match the data types defined in the Pydantic model", - ), - ] = Contract.discard_row, + ] = Contract.evolve, + columns_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when data contains values/keys not found in the Pydantic model", + ), + ] = Contract.evolve, + data_type_contract: Annotated[ + Contract, + typer.Option( + help="DLT contract applied when fields do not match the data types defined in the Pydantic model", + ), + ] = Contract.discard_row, ) -> LoadInfo | None: schema_contract = { "tables": tables_contract, @@ -141,10 +147,10 @@ def wrapper( return decorator def convert( - self, - lookup: Callable | None = None, - help: str = "OpenGraph convert pipeline", - **typer_kwargs, + self, + lookup: Callable | None = None, + help: str = "OpenGraph convert pipeline", + **typer_kwargs, ): """Register a Typer CLI command that converts collected resources to OpenGraph nodes and edges. @@ -157,11 +163,11 @@ def convert( def decorator(func: Callable): def run_convert( - input_path: InputPath, - output_path: Path = Path("/tmp/openhound"), - lookup_file: Path = DEFAULT_LOOKUP_FILE, - progress: Progress = Progress.tqdm, - method: Method = Method.write, + input_path: InputPath, + output_path: Path = Path("/tmp/openhound"), + lookup_file: Path = DEFAULT_LOOKUP_FILE, + progress: Progress = Progress.tqdm, + method: Method = Method.write, ) -> LoadInfo: lookup_session = None if lookup: @@ -196,31 +202,31 @@ def run_convert( ) def wrapper( - input_path: InputPath, - output_path: Annotated[ - Path, - typer.Argument( - exists=False, - file_okay=False, - dir_okay=True, - resolve_path=True, - help="Output path to write OpenGraph JSON files", - ), - ], - # resources: List[str] = typer.Argument(None), - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" + input_path: InputPath, + output_path: Annotated[ + Path, + typer.Argument( + exists=False, + file_okay=False, + dir_okay=True, + resolve_path=True, + help="Output path to write OpenGraph JSON files", + ), + ], + # resources: List[str] = typer.Argument(None), + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" + ), + lookup_file: Annotated[ + Path, + typer.Option( + file_okay=True, + dir_okay=False, + readable=True, + resolve_path=True, + help="DuckDB lookup file path", ), - lookup_file: Annotated[ - Path, - typer.Option( - file_okay=True, - dir_okay=False, - readable=True, - resolve_path=True, - help="DuckDB lookup file path", - ), - ] = DEFAULT_LOOKUP_FILE, + ] = DEFAULT_LOOKUP_FILE, ) -> LoadInfo: return run_convert( input_path=input_path, @@ -240,10 +246,10 @@ def wrapper( return decorator def preproc( - self, - transformer: Callable[[any], None] | None = None, - help: str = "OpenGraph preprocessing pipeline", - **typer_kwargs, + self, + transformer: Callable[[any], None] | None = None, + help: str = "OpenGraph preprocessing pipeline", + **typer_kwargs, ): """Register a Typer CLI command that performs optional preprocessing and builds lookup data for a source. @@ -257,19 +263,19 @@ def decorator(func: Callable): self.preprocessor = func def wrapper( - input_path: InputPath, - output_file: Annotated[ - Path, - typer.Argument( - file_okay=True, - dir_okay=False, - readable=True, - resolve_path=True, - ), - ] = DEFAULT_LOOKUP_FILE, - progress: Progress = typer.Option( - Progress.tqdm, help="Select progress tracker option" + input_path: InputPath, + output_file: Annotated[ + Path, + typer.Argument( + file_okay=True, + dir_okay=False, + readable=True, + resolve_path=True, ), + ] = DEFAULT_LOOKUP_FILE, + progress: Progress = typer.Option( + Progress.tqdm, help="Select progress tracker option" + ), ) -> LoadInfo: preprocessor = PreProcessor( name=self.name, @@ -298,9 +304,9 @@ def defer(self, func: Callable) -> Callable: return dlt.defer(safe_func) def transformer( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT transformer with added exception handling.""" @@ -315,9 +321,9 @@ def decorator(func: Callable) -> DltResource: return decorator def resource( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT resource with added exception handling.""" @@ -332,9 +338,9 @@ def decorator(func: Callable) -> DltResource: return decorator def source( - self, - *dlt_args, - **dlt_kwargs, + self, + *dlt_args, + **dlt_kwargs, ): """Decorator to register a DLT source with added exception handling. @@ -352,10 +358,10 @@ def decorator(func: Callable) -> Callable: return decorator def asset( - self, - node: NodeDef | None = None, - edges: list[EdgeDef] | None = None, - description: str = "Resource model for OpenGraph", + self, + node: NodeDef | None = None, + edges: list[EdgeDef] | None = None, + description: str = "Resource model for OpenGraph", ): """Decorator to register a resource class and its graph definitions (nodes/edges). This is used to automatically generate documentation for each unique resource and implement rules/warnings when nodes/edges are returned From f548a7941f904e2b154f6a6a6b2534d773879639 Mon Sep 17 00:00:00 2001 From: Joey Dreijer Date: Fri, 28 Aug 2026 14:35:30 +0200 Subject: [PATCH 3/3] Formatting --- src/openhound/sources/opengraph/source.py | 34 +++++++++++------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/openhound/sources/opengraph/source.py b/src/openhound/sources/opengraph/source.py index dd3803a0..88774437 100644 --- a/src/openhound/sources/opengraph/source.py +++ b/src/openhound/sources/opengraph/source.py @@ -7,6 +7,7 @@ from openhound.core.asset import BaseAsset from openhound.core.lookup import LookupManager + from .entries import GraphContent # DLT page boundary; partial pages flush per input file. @@ -20,11 +21,11 @@ class GraphResource: def _generate_graph_content( - resources: Iterable[dict], - model: type[BaseAsset], - batch_size: int, - apply_context: Callable | None = None, - source_kind: str | None = None + resources: Iterable[dict], + model: type[BaseAsset], + batch_size: int, + apply_context: Callable | None = None, + source_kind: str | None = None, ): """Convert one DLT page into bounded OpenGraph batches.""" edge_parts = [] @@ -62,12 +63,12 @@ def serialize(content): @dlt.source(name="opengraph", max_table_nesting=0) def opengraph( - graph_resources: list[GraphResource], - bucket_url: str, - lookup: LookupManager, - extras: dict | None = None, - batch_size: int = 150, - source_kind: str | None = None + graph_resources: list[GraphResource], + bucket_url: str, + lookup: LookupManager, + extras: dict | None = None, + batch_size: int = 150, + source_kind: str | None = None, ): if batch_size <= 0: raise ValueError("batch_size must be greater than zero") @@ -78,13 +79,10 @@ def apply_context(obj): for graph_resource in graph_resources: table_name = f"{graph_resource.model.__name__.lower()}_fs" - reader = ( - filesystemsource( - bucket_url=bucket_url, - file_glob=f"{graph_resource.table}/**/*.jsonl.gz", - ) - | read_jsonl(chunksize=READ_JSONL_PAGE_SIZE) - ) + reader = filesystemsource( + bucket_url=bucket_url, + file_glob=f"{graph_resource.table}/**/*.jsonl.gz", + ) | 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):