Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/openhound/core/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ def __init__(self, name: str, source_kind: str, help: str = "OpenGraph collector

# Extension metadata is loaded and validated when the extension is loaded in the CollectorManager
self.metadata: Extension | None = None
# The installed distribution version is supplied by the extension entry point.
self.package_version: str | None = None
# Store the collect/convert/preproc methods for this source
self.collector: Callable | None = None
self.converter: Callable | None = None
Expand Down
3 changes: 3 additions & 0 deletions src/openhound/core/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,9 @@ def from_entrypoint(cls, group: str = "openhound.sources") -> "CollectorManager"
is_valid_extension = cls.validate_extension(load_extension, extension.name)
if is_valid_extension:
load_extension.metadata = metadata
load_extension.package_version = (
extension.dist.version if extension.dist is not None else None
)
extensions.append(load_extension)
logger.info(
f"Loaded extension '{extension.name}' from entry point '{group}'",
Expand Down
27 changes: 25 additions & 2 deletions src/openhound/scheduler/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from dataclasses import dataclass
from pathlib import Path

import openhound
import openhound.core.logging as openhound_logging
from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus
from openhound.core.clients.models.jobs import (
Expand Down Expand Up @@ -51,13 +52,35 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result:
ExtensionNotFoundError: If the collector cannot be found via entrypoints.
"""
signal.signal(signal.SIGINT, signal.SIG_IGN)
logger.info(f"Subprocess running collection '{collector_name}' for job {job_id}")
available_collectors = CollectorManager.from_entrypoint()

for collector in available_collectors.collectors:
if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue]
log_fields = {
"collector_extension": collector.name,
"collector_extension_version": (
collector.package_version
or (
str(collector.metadata.version)
if collector.metadata is not None
else "unknown"
)
),
"openhound_version": openhound.__version__,
"job_id": job_id,
}
logger.info(
"Subprocess running collection '%s' for job %s",
collector_name,
job_id,
extra=log_fields,
)
results = dataflow.pipeline(extension=collector)
logger.info(f"Collection for job {job_id} completed successfully.")
logger.info(
"Collection for job %s completed successfully.",
job_id,
extra=log_fields,
)
return Result(results=results, job_id=job_id)

logger.error(f"Collector '{collector_name}' not found in available collectors.")
Expand Down
34 changes: 34 additions & 0 deletions tests/test_bhe_job_scheduling.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@
import gzip
import hashlib
import json
import logging
from concurrent.futures import Future
from concurrent.futures.process import BrokenProcessPool
from pathlib import Path
from types import SimpleNamespace
from urllib.parse import urlsplit

import pytest
Expand Down Expand Up @@ -750,3 +752,35 @@ def unavailable_part_upload(method, path, **kwargs):
"operation_id": operation.id,
"status": ManagementOperationStatus.FAILED.value,
}


def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, caplog):
collector = SimpleNamespace(
name="example",
package_version="2.3.4",
metadata=SimpleNamespace(version="1.2.3"),
)
monkeypatch.setattr(
scheduler_service.CollectorManager,
"from_entrypoint",
lambda: SimpleNamespace(collectors=[collector]),
)
monkeypatch.setattr(
scheduler_service.dataflow, "pipeline", lambda extension: {"collect": []}
)
monkeypatch.setattr(scheduler_service.openhound, "__version__", "4.5.6")

with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"):
_subprocess_collect("example", 42)

collection_records = [
record
for record in caplog.records
if record.getMessage().startswith(("Subprocess running collection", "Collection for job"))
]
assert len(collection_records) == 2
for record in collection_records:
assert record.collector_extension == "example"
assert record.collector_extension_version == "2.3.4"
assert record.openhound_version == "4.5.6"
assert record.job_id == 42
Loading