feat: add dry run mode for Python SDK

README.md

@@ -24,6 +24,7 @@ pip install netboxlabs-diode-sdk
 * `DIODE_SENTRY_DSN` - Optional Sentry DSN for error reporting
 * `DIODE_CLIENT_ID` - Client ID for OAuth2 authentication
 * `DIODE_CLIENT_SECRET` - Client Secret for OAuth2 authentication
+* `DIODE_DRY_RUN_OUTPUT_FILE` - Path to store JSON Lines output when using `DiodeDryRunClient`
 
 ### Example
 
@@ -75,6 +76,35 @@ if __name__ == "__main__":
 
 ```
 
+### Dry run mode
+
+`DiodeDryRunClient` allows generating ingestion requests without contacting a
+Diode server. Each request is printed to stdout or written to a JSON Lines file
+when `dry_run_output_file` (or `DIODE_DRY_RUN_OUTPUT_FILE`) is set.
+
+```python
+from netboxlabs.diode.sdk import DiodeDryRunClient
+
+with DiodeDryRunClient(dry_run_output_file="dryrun.jsonl") as client:
+    client.ingest([
+        Entity(device="Device A"),
+    ])
+```
+
+The produced file can later be ingested by a real Diode instance using
+`load_dryrun_entities` with a standard `DiodeClient`:
+
+```python
+from netboxlabs.diode.sdk import DiodeClient, load_dryrun_entities
+
+with DiodeClient(
+    target="grpc://localhost:8080/diode",
+    app_name="my-test-app",
+    app_version="0.0.1",
+) as client:
+    client.ingest(entities=load_dryrun_entities("dryrun.jsonl"))
+```
+
 ## Supported entities (object types)
 
 * ASN

netboxlabs/diode/sdk/__init__.py

@@ -2,6 +2,12 @@
 # Copyright 2024 NetBox Labs Inc
 """NetBox Labs, Diode - SDK."""
 
-from netboxlabs.diode.sdk.client import DiodeClient
+from netboxlabs.diode.sdk.client import (
+    DiodeClient,
+    DiodeDryRunClient,
+    load_dryrun_entities,
+)
 
 assert DiodeClient
+assert DiodeDryRunClient
+assert load_dryrun_entities

netboxlabs/diode/sdk/client.py

@@ -9,13 +9,16 @@
 import os
 import platform
 import ssl
+import sys
 import uuid
 from collections.abc import Iterable
+from pathlib import Path
 from urllib.parse import urlencode, urlparse
 
 import certifi
 import grpc
 import sentry_sdk
+from google.protobuf.json_format import MessageToDict, ParseDict
 
 from netboxlabs.diode.sdk.diode.v1 import ingester_pb2, ingester_pb2_grpc
 from netboxlabs.diode.sdk.exceptions import DiodeClientError, DiodeConfigError
@@ -27,11 +30,29 @@
 _DIODE_SENTRY_DSN_ENVVAR_NAME = "DIODE_SENTRY_DSN"
 _CLIENT_ID_ENVVAR_NAME = "DIODE_CLIENT_ID"
 _CLIENT_SECRET_ENVVAR_NAME = "DIODE_CLIENT_SECRET"
+_DRY_RUN_OUTPUT_ENVVAR_NAME = "DIODE_DRY_RUN_OUTPUT_FILE"
 _INGEST_SCOPE = "diode:ingest"
 _DEFAULT_STREAM = "latest"
 _LOGGER = logging.getLogger(__name__)
 
 
+def load_dryrun_entities(file_path: str | Path) -> Iterable[Entity]:
+    """Yield entities stored in a JSON Lines file produced by ``DiodeDryRunClient``."""
+    with open(file_path) as fh:
+        for line in fh:
+            data = json.loads(line)
+            for entity_dict in data.get("entities", []):
+                pb_entity = Entity()
+                ParseDict(entity_dict, pb_entity)
+                yield pb_entity
+
+
+class DiodeClientInterface:
+    """Runtime placeholder for the Diode client interface."""
+
+    pass
+
+
 def _load_certs() -> bytes:
     """Loads cacert.pem."""
     with open(certifi.where(), "rb") as f:
@@ -82,7 +103,7 @@ def _get_optional_config_value(
     return value
 
 
-class DiodeClient:
+class DiodeClient(DiodeClientInterface):
     """Diode Client."""
 
     _name = "diode-sdk-python"
@@ -287,6 +308,65 @@ def _authenticate(self, scope: str):
         ) + [("authorization", f"Bearer {access_token}")]
 
 
+class DiodeDryRunClient(DiodeClientInterface):
+    """Client that outputs ingestion requests instead of sending them."""
+
+    _name = "diode-sdk-python-dry-run"
+    _version = version_semver()
+    _app_name = None
+    _app_version = None
+
+    def __init__(self, dry_run_output_file: str | None = None):
+        """Initiate a new dry run client."""
+        self._dry_run_output_file = os.getenv(
+            _DRY_RUN_OUTPUT_ENVVAR_NAME, dry_run_output_file
+        )
+
+    @property
+    def name(self) -> str:
+        """Retrieve the name."""
+        return self._name
+
+    @property
+    def version(self) -> str:
+        """Retrieve the version."""
+        return self._version
+
+    @property
+    def dry_run_output_file(self) -> str | None:
+        """Retrieve the dry run output file."""
+        return self._dry_run_output_file
+
+    def __enter__(self):
+        """Enters the runtime context related to the channel object."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        """Exits the runtime context related to the channel object."""
+
+    def ingest(
+        self,
+        entities: Iterable[Entity | ingester_pb2.Entity | None],
+        stream: str | None = _DEFAULT_STREAM,
+    ) -> ingester_pb2.IngestResponse:
+        """Ingest entities in dry run mode."""
+        request = ingester_pb2.IngestRequest(
+            stream=stream,
+            id=str(uuid.uuid4()),
+            entities=entities,
+            sdk_name=self.name,
+            sdk_version=self.version,
+        )
+
+        output = json.dumps(MessageToDict(request))
+        if self._dry_run_output_file:
+            with open(self._dry_run_output_file, "a") as fh:
+                fh.write(output + "\n")
+        else:
+            print(output, file=sys.stdout)
+        return ingester_pb2.IngestResponse()
+
+
 class _DiodeAuthentication:
     def __init__(
         self,

netboxlabs/diode/sdk/client.pyi

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Protocol, runtime_checkable
+
+from netboxlabs.diode.sdk.diode.v1 import ingester_pb2
+from netboxlabs.diode.sdk.ingester import Entity
+
+_DEFAULT_STREAM: str
+
+@runtime_checkable
+class DiodeClientInterface(Protocol):
+    """Interface implemented by diode clients."""
+
+    @property
+    def name(self) -> str: ...
+    @property
+    def version(self) -> str: ...
+    def ingest(
+        self,
+        entities: Iterable[Entity | ingester_pb2.Entity | None],
+        stream: str | None = _DEFAULT_STREAM,
+    ) -> ingester_pb2.IngestResponse: ...
+    def __enter__(self) -> DiodeClientInterface: ...
+    def __exit__(self, exc_type, exc_value, exc_traceback) -> None: ...