diff --git a/src/openhound/core/app.py b/src/openhound/core/app.py index 526a49ce..b44e6754 100644 --- a/src/openhound/core/app.py +++ b/src/openhound/core/app.py @@ -52,7 +52,12 @@ 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 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..88774437 100644 --- a/src/openhound/sources/opengraph/source.py +++ b/src/openhound/sources/opengraph/source.py @@ -25,6 +25,7 @@ def _generate_graph_content( 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 +42,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), @@ -65,6 +68,7 @@ def opengraph( 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") @@ -75,18 +79,15 @@ 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): 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"], )