flow.record 3.15.dev16__tar.gz → 3.16.dev1__tar.gz

{flow_record-3.15.dev16/flow.record.egg-info → flow_record-3.16.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flow.record
-Version: 3.15.dev16
+Version: 3.16.dev1
 Summary: A library for defining and creating structured data (called records) that can be streamed to disk or piped to other tools that use flow.record
 Author-email: Dissect Team <dissect@fox-it.com>
 License: Affero General Public License v3
@@ -40,6 +40,7 @@ Requires-Dist: httpx; extra == "splunk"
 Provides-Extra: test
 Requires-Dist: flow.record[compression]; extra == "test"
 Requires-Dist: flow.record[avro]; extra == "test"
+Requires-Dist: flow.record[elastic]; extra == "test"
 Requires-Dist: duckdb; (platform_python_implementation != "PyPy" and python_version < "3.12") and extra == "test"
 Requires-Dist: pytz; (platform_python_implementation != "PyPy" and python_version < "3.12") and extra == "test"
 

{flow_record-3.15.dev16 → flow_record-3.16.dev1}/flow/record/adapter/elastic.py

@@ -2,7 +2,7 @@ import hashlib
 import logging
 import queue
 import threading
-from typing import Iterator, Union
+from typing import Iterator, Optional, Union
 
 import elasticsearch
 import elasticsearch.helpers
@@ -22,9 +22,11 @@ Read usage: rdump elastic+[PROTOCOL]://[IP]:[PORT]?index=[INDEX]
 [PROTOCOL]: http or https. Defaults to https when "+[PROTOCOL]" is omitted
 
 Optional arguments:
+  [API_KEY]: base64 encoded api key to authenticate with (default: False)
   [INDEX]: name of the index to use (default: records)
   [VERIFY_CERTS]: verify certs of Elasticsearch instance (default: True)
   [HASH_RECORD]: make record unique by hashing record [slow] (default: False)
+  [_META_*]: record metadata fields (default: None)
 """
 
 log = logging.getLogger(__name__)
@@ -38,6 +40,7 @@ class ElasticWriter(AbstractWriter):
         verify_certs: Union[str, bool] = True,
         http_compress: Union[str, bool] = True,
         hash_record: Union[str, bool] = False,
+        api_key: Optional[str] = None,
         **kwargs,
     ) -> None:
         self.index = index
@@ -45,7 +48,17 @@ class ElasticWriter(AbstractWriter):
         verify_certs = str(verify_certs).lower() in ("1", "true")
         http_compress = str(http_compress).lower() in ("1", "true")
         self.hash_record = str(hash_record).lower() in ("1", "true")
-        self.es = elasticsearch.Elasticsearch(uri, verify_certs=verify_certs, http_compress=http_compress)
+
+        if not uri.lower().startswith(("http://", "https://")):
+            uri = "http://" + uri
+
+        self.es = elasticsearch.Elasticsearch(
+            uri,
+            verify_certs=verify_certs,
+            http_compress=http_compress,
+            api_key=api_key,
+        )
+
         self.json_packer = JsonRecordPacker()
         self.queue: queue.Queue[Union[Record, StopIteration]] = queue.Queue()
         self.event = threading.Event()
@@ -58,25 +71,34 @@ class ElasticWriter(AbstractWriter):
 
             urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
+        self.metadata_fields = {}
+        for arg_key, arg_val in kwargs.items():
+            if arg_key.startswith("_meta_"):
+                self.metadata_fields[arg_key[6:]] = arg_val
+
     def record_to_document(self, record: Record, index: str) -> dict:
         """Convert a record to a Elasticsearch compatible document dictionary"""
         rdict = record._asdict()
 
-        # Store record metadata under `_record_metadata`
+        # Store record metadata under `_record_metadata`.
         rdict_meta = {
             "descriptor": {
                 "name": record._desc.name,
                 "hash": record._desc.descriptor_hash,
             },
         }
+
         # Move all dunder fields to `_record_metadata` to avoid naming clash with ES.
         dunder_keys = [key for key in rdict if key.startswith("_")]
         for key in dunder_keys:
             rdict_meta[key.lstrip("_")] = rdict.pop(key)
-        # remove _generated field from metadata to ensure determinstic documents
+
+        # Remove _generated field from metadata to ensure determinstic documents.
         if self.hash_record:
             rdict_meta.pop("generated", None)
-        rdict["_record_metadata"] = rdict_meta
+
+        rdict["_record_metadata"] = rdict_meta.copy()
+        rdict["_record_metadata"].update(self.metadata_fields)
 
         document = {
             "_index": index,
@@ -106,6 +128,7 @@ class ElasticWriter(AbstractWriter):
         ):
             if not ok:
                 log.error("Failed to insert %r", item)
+
         self.event.set()
 
     def write(self, record: Record) -> None:
@@ -129,6 +152,7 @@ class ElasticReader(AbstractReader):
         verify_certs: Union[str, bool] = True,
         http_compress: Union[str, bool] = True,
         selector: Union[None, Selector, CompiledSelector] = None,
+        api_key: Optional[str] = None,
         **kwargs,
     ) -> None:
         self.index = index
@@ -136,7 +160,16 @@ class ElasticReader(AbstractReader):
         self.selector = selector
         verify_certs = str(verify_certs).lower() in ("1", "true")
         http_compress = str(http_compress).lower() in ("1", "true")
-        self.es = elasticsearch.Elasticsearch(uri, verify_certs=verify_certs, http_compress=http_compress)
+
+        if not uri.lower().startswith(("http://", "https://")):
+            uri = "http://" + uri
+
+        self.es = elasticsearch.Elasticsearch(
+            uri,
+            verify_certs=verify_certs,
+            http_compress=http_compress,
+            api_key=api_key,
+        )
 
         if not verify_certs:
             # Disable InsecureRequestWarning of urllib3, caused by the verify_certs flag.

{flow_record-3.15.dev16 → flow_record-3.16.dev1}/flow/record/version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.15.dev16'
-__version_tuple__ = version_tuple = (3, 15, 'dev16')
+__version__ = version = '3.16.dev1'
+__version_tuple__ = version_tuple = (3, 16, 'dev1')

{flow_record-3.15.dev16 → flow_record-3.16.dev1/flow.record.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flow.record
-Version: 3.15.dev16
+Version: 3.16.dev1
 Summary: A library for defining and creating structured data (called records) that can be streamed to disk or piped to other tools that use flow.record
 Author-email: Dissect Team <dissect@fox-it.com>
 License: Affero General Public License v3
@@ -40,6 +40,7 @@ Requires-Dist: httpx; extra == "splunk"
 Provides-Extra: test
 Requires-Dist: flow.record[compression]; extra == "test"
 Requires-Dist: flow.record[avro]; extra == "test"
+Requires-Dist: flow.record[elastic]; extra == "test"
 Requires-Dist: duckdb; (platform_python_implementation != "PyPy" and python_version < "3.12") and extra == "test"
 Requires-Dist: pytz; (platform_python_implementation != "PyPy" and python_version < "3.12") and extra == "test"
 

{flow_record-3.15.dev16 → flow_record-3.16.dev1}/flow.record.egg-info/SOURCES.txt

@@ -59,6 +59,7 @@ tests/test_avro_adapter.py
 tests/test_compiled_selector.py
 tests/test_csv_adapter.py
 tests/test_deprecations.py
+tests/test_elastic_adapter.py
 tests/test_fieldtype_ip.py
 tests/test_fieldtypes.py
 tests/test_json_packer.py

{flow_record-3.15.dev16 → flow_record-3.16.dev1}/flow.record.egg-info/requires.txt

@@ -29,6 +29,7 @@ httpx
 [test]
 flow.record[compression]
 flow.record[avro]
+flow.record[elastic]
 
 [test:platform_python_implementation != "PyPy" and python_version < "3.12"]
 duckdb

{flow_record-3.15.dev16 → flow_record-3.16.dev1}/pyproject.toml

@@ -59,6 +59,7 @@ splunk = [
 test = [
     "flow.record[compression]",
     "flow.record[avro]",
+    "flow.record[elastic]",
     "duckdb; platform_python_implementation != 'PyPy' and python_version < '3.12'", # duckdb
     "pytz; platform_python_implementation != 'PyPy' and python_version < '3.12'", # duckdb
 ]

flow_record-3.16.dev1/tests/test_elastic_adapter.py

@@ -0,0 +1,53 @@
+import json
+
+import pytest
+
+from flow.record import RecordDescriptor
+from flow.record.adapter.elastic import ElasticWriter
+
+MyRecord = RecordDescriptor(
+    "my/record",
+    [
+        ("string", "field_one"),
+        ("string", "field_two"),
+    ],
+)
+
+
+@pytest.mark.parametrize(
+    "record",
+    [
+        MyRecord("first", "record"),
+        MyRecord("second", "record"),
+    ],
+)
+def test_elastic_writer_metadata(record):
+    options = {
+        "_meta_foo": "some value",
+        "_meta_bar": "another value",
+    }
+
+    with ElasticWriter(uri="elasticsearch:9200", **options) as writer:
+        assert writer.metadata_fields == {"foo": "some value", "bar": "another value"}
+
+        assert writer.record_to_document(record, "some-index") == {
+            "_index": "some-index",
+            "_source": json.dumps(
+                {
+                    "field_one": record.field_one,
+                    "field_two": record.field_two,
+                    "_record_metadata": {
+                        "descriptor": {
+                            "name": "my/record",
+                            "hash": record._desc.descriptor_hash,
+                        },
+                        "source": None,
+                        "classification": None,
+                        "generated": record._generated.isoformat(),
+                        "version": 1,
+                        "foo": "some value",
+                        "bar": "another value",
+                    },
+                }
+            ),
+        }