flow.record 3.18.dev2__tar.gz → 3.19__tar.gz

flow_record-3.19/.git-blame-ignore-revs

@@ -0,0 +1,6 @@
+# Formatting commits. You can ignore them during git-blame with `--ignore-rev` or `--ignore-revs-file`.
+#
+#   $ git config --add 'blame.ignoreRevsFile' '.git-blame-ignore-revs'
+#
+# Change linter to Ruff (#158)
+c67f778c653c295ec26146cf6422d3b06ac640e8

{flow_record-3.18.dev2/flow.record.egg-info → flow_record-3.19}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: flow.record
-Version: 3.18.dev2
+Version: 3.19
 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

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/__init__.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import gzip
-import os
+from pathlib import Path
 
 from flow.record.base import (
     IGNORE_FIELDS_FOR_COMPARISON,
@@ -39,71 +41,61 @@ from flow.record.stream import (
 
 __all__ = [
     "IGNORE_FIELDS_FOR_COMPARISON",
-    "RECORD_VERSION",
     "RECORDSTREAM_MAGIC",
+    "RECORD_VERSION",
+    "DynamicDescriptor",
     "FieldType",
-    "Record",
     "GroupedRecord",
-    "RecordDescriptor",
+    "JsonRecordPacker",
+    "PathTemplateWriter",
+    "Record",
     "RecordAdapter",
+    "RecordArchiver",
+    "RecordDescriptor",
+    "RecordDescriptorError",
     "RecordField",
-    "RecordReader",
-    "RecordWriter",
     "RecordOutput",
-    "RecordPrinter",
     "RecordPacker",
-    "JsonRecordPacker",
-    "RecordStreamWriter",
+    "RecordPrinter",
+    "RecordReader",
     "RecordStreamReader",
-    "open_path_or_stream",
+    "RecordStreamWriter",
+    "RecordWriter",
+    "dynamic_fieldtype",
+    "extend_record",
+    "ignore_fields_for_comparison",
+    "iter_timestamped_records",
     "open_path",
+    "open_path_or_stream",
     "open_stream",
-    "ignore_fields_for_comparison",
+    "record_stream",
     "set_ignored_fields_for_comparison",
     "stream",
-    "dynamic_fieldtype",
-    "DynamicDescriptor",
-    "PathTemplateWriter",
-    "RecordArchiver",
-    "RecordDescriptorError",
-    "record_stream",
-    "extend_record",
-    "iter_timestamped_records",
 ]
 
 
-class View:
-    fields = None
-
-    def __init__(self, fields):
-        self.fields = fields
-
-    def __iter__(self, fields):
-        pass
-
-
 class RecordDateSplitter:
     basepath = None
     out = None
 
-    def __init__(self, basepath):
-        self.basepath = basepath
+    def __init__(self, basepath: str | Path):
+        self.basepath = Path(basepath)
         self.out = {}
 
-    def getstream(self, t):
+    def getstream(self, t: tuple[int, int, int]) -> RecordStreamWriter:
         if t not in self.out:
-            path = os.path.join(self.basepath, "-".join(["{:2d}".format(v) for v in t]) + ".rec.gz")
+            path = self.basepath.joinpath("-".join([f"{v:2d}" for v in t]) + ".rec.gz")
             f = gzip.GzipFile(path, "wb")
             rs = RecordStreamWriter(f)
             self.out[t] = rs
         return self.out[t]
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         t = (r.ts.year, r.ts.month, r.ts.day)
         rs = self.getstream(t)
         rs.write(r)
         rs.fp.flush()
 
-    def close(self):
+    def close(self) -> None:
         for rs in self.out.values():
             rs.close()

flow_record-3.19/flow/record/adapter/__init__.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # make this namespace extensible from other packages
+import abc
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.base import Record
+
+
+class AbstractWriter(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def write(self, rec: Record) -> None:
+        """Write a record."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def flush(self) -> None:
+        """Flush any buffered writes."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Close the Writer, no more writes will be possible."""
+        raise NotImplementedError
+
+    def __del__(self) -> None:
+        self.close()
+
+    def __enter__(self) -> AbstractWriter:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.flush()
+        self.close()
+
+
+class AbstractReader(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def __iter__(self) -> Iterator[Record]:
+        """Return a record iterator."""
+        raise NotImplementedError
+
+    def close(self) -> None:  # noqa: B027
+        """Close the Reader, can be overriden to properly free resources."""
+
+    def __enter__(self) -> AbstractReader:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.close()

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/adapter/archive.py

@@ -1,6 +1,13 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.stream import RecordArchiver
 
+if TYPE_CHECKING:
+    from flow.record.base import Record
+
 __usage__ = """
 Record archiver adapter, writes records to YYYY/mm/dd directories (writer only)
 ---
@@ -12,7 +19,7 @@ Write usage: rdump -w archive://[PATH]
 class ArchiveWriter(AbstractWriter):
     writer = None
 
-    def __init__(self, path, **kwargs):
+    def __init__(self, path: str, **kwargs):
         self.path = path
 
         path_template = kwargs.get("path_template")
@@ -20,19 +27,19 @@ class ArchiveWriter(AbstractWriter):
 
         self.writer = RecordArchiver(self.path, path_template=path_template, name=name)
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         self.writer.write(r)
 
-    def flush(self):
+    def flush(self) -> None:
         # RecordArchiver already flushes after every write
         pass
 
-    def close(self):
+    def close(self) -> None:
         if self.writer:
             self.writer.close()
         self.writer = None
 
 
 class ArchiveReader(AbstractReader):
-    def __init__(self, path, **kwargs):
+    def __init__(self, path: str, **kwargs):
         raise NotImplementedError

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/adapter/avro.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 from datetime import datetime, timedelta, timezone
 from importlib.util import find_spec
-from typing import Any, Iterator
+from typing import TYPE_CHECKING, Any, BinaryIO
 
 import fastavro
 
@@ -12,6 +12,10 @@ from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.selector import make_selector
 from flow.record.utils import is_stdout
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
 __usage__ = """
 Apache AVRO adapter
 ---
@@ -52,7 +56,7 @@ class AvroWriter(AbstractWriter):
     fp = None
     writer = None
 
-    def __init__(self, path, key=None, **kwargs):
+    def __init__(self, path: str | Path | BinaryIO, **kwargs):
         self.fp = record.open_path_or_stream(path, "wb")
 
         self.desc = None
@@ -69,11 +73,11 @@ class AvroWriter(AbstractWriter):
             self.writer = fastavro.write.Writer(self.fp, self.parsed_schema, codec=self.codec)
 
         if self.desc != r._desc:
-            raise Exception("Mixed record types")
+            raise ValueError("Mixed record types")
 
         self.writer.write(r._packdict())
 
-    def flush(self):
+    def flush(self) -> None:
         if not self.writer:
             self.writer = fastavro.write.Writer(
                 self.fp,
@@ -92,21 +96,21 @@ class AvroWriter(AbstractWriter):
 class AvroReader(AbstractReader):
     fp = None
 
-    def __init__(self, path, selector=None, **kwargs):
+    def __init__(self, path: str, selector: str | None = None, **kwargs):
         self.fp = record.open_path_or_stream(path, "rb")
         self.selector = make_selector(selector)
 
         self.reader = fastavro.reader(self.fp)
         self.schema = self.reader.writer_schema
         if not self.schema:
-            raise Exception("Missing Avro schema")
+            raise ValueError("Missing Avro schema")
 
         self.desc = schema_to_descriptor(self.schema)
 
         # Store the fieldnames that are of type "datetime"
-        self.datetime_fields = set(
+        self.datetime_fields = {
             name for name, field in self.desc.get_all_fields().items() if field.typename == "datetime"
-        )
+        }
 
     def __iter__(self) -> Iterator[record.Record]:
         for obj in self.reader:
@@ -149,7 +153,7 @@ def descriptor_to_schema(desc: record.RecordDescriptor) -> dict[str, Any]:
         else:
             avro_type = AVRO_TYPE_MAP.get(field_type)
             if not avro_type:
-                raise Exception("Unsupported Avro type: {}".format(field_type))
+                raise ValueError(f"Unsupported Avro type: {field_type}")
 
             field_schema["type"] = [avro_type, "null"]
 
@@ -190,11 +194,10 @@ def avro_type_to_flow_type(ftype: list) -> str:
         if isinstance(t, dict):
             if t.get("type") == "array":
                 item_type = avro_type_to_flow_type(t.get("items"))
-                return "{}[]".format(item_type)
-            else:
-                logical_type = t.get("logicalType")
-                if logical_type and ("time" in logical_type or "date" in logical_type):
-                    return "datetime"
+                return f"{item_type}[]"
+            logical_type = t.get("logicalType")
+            if logical_type and ("time" in logical_type or "date" in logical_type):
+                return "datetime"
 
         if t == "null":
             continue
@@ -202,4 +205,4 @@ def avro_type_to_flow_type(ftype: list) -> str:
         if t in RECORD_TYPE_MAP:
             return RECORD_TYPE_MAP[t]
 
-    raise TypeError("Can't map avro type to flow type: {}".format(t))
+    raise TypeError(f"Can't map avro type to flow type: {t}")

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/adapter/broker.py

@@ -1,7 +1,15 @@
-from flow.broker import Publisher, Subscriber
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
+from flow.broker import Publisher, Subscriber
 from flow.record.adapter import AbstractReader, AbstractWriter
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.base import Record
+
 __usage__ = """
 PubSub adapter using flow.broker
 ---
@@ -13,23 +21,23 @@ Read usage: rdump broker+tcp://[IP]:[PORT] -s True
 class BrokerWriter(AbstractWriter):
     publisher = None
 
-    def __init__(self, uri, source=None, classification=None, **kwargs):
+    def __init__(self, uri: str, source: str | None = None, classification: str | None = None, **kwargs):
         self.publisher = Publisher(uri, **kwargs)
         self.source = source
         self.classification = classification
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         record = r._replace(
             _source=self.source or r._source,
             _classification=self.classification or r._classification,
         )
         self.publisher.send(record)
 
-    def flush(self):
+    def flush(self) -> None:
         if self.publisher:
             self.publisher.flush()
 
-    def close(self):
+    def close(self) -> None:
         if self.publisher:
             if hasattr(self.publisher, "stop"):
                 # Requires flow.broker >= 1.1.1
@@ -42,14 +50,14 @@ class BrokerWriter(AbstractWriter):
 class BrokerReader(AbstractReader):
     subscriber = None
 
-    def __init__(self, uri, name=None, selector=None, **kwargs):
+    def __init__(self, uri: str, name: str | None = None, selector: str | None = None, **kwargs):
         self.subscriber = Subscriber(uri, **kwargs)
         self.subscription = self.subscriber.select(name, str(selector))
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Record]:
         return iter(self.subscription)
 
-    def close(self):
+    def close(self) -> None:
         if self.subscriber:
             self.subscriber.stop()
         self.subscriber = None

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/adapter/csvfile.py

@@ -1,14 +1,19 @@
-from __future__ import absolute_import
+from __future__ import annotations
 
 import csv
 import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 from flow.record import RecordDescriptor
 from flow.record.adapter import AbstractReader, AbstractWriter
-from flow.record.base import normalize_fieldname
+from flow.record.base import Record, normalize_fieldname
 from flow.record.selector import make_selector
 from flow.record.utils import is_stdout
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 __usage__ = """
 Comma-separated values (CSV) adapter
 ---
@@ -23,13 +28,20 @@ Optional parameters:
 
 
 class CsvfileWriter(AbstractWriter):
-    def __init__(self, path, fields=None, exclude=None, lineterminator=None, **kwargs):
+    def __init__(
+        self,
+        path: str | Path | None,
+        fields: str | list[str] | None = None,
+        exclude: str | list[str] | None = None,
+        lineterminator: str = "\r\n",
+        **kwargs,
+    ):
         self.fp = None
         if path in (None, "", "-"):
             self.fp = sys.stdout
         else:
-            self.fp = open(path, "w", newline="")
-        self.lineterminator = lineterminator or "\r\n"
+            self.fp = Path(path).open("w", newline="")  # noqa: SIM115
+        self.lineterminator = lineterminator
         for r, n in ((r"\r", "\r"), (r"\n", "\n"), (r"\t", "\t")):
             self.lineterminator = self.lineterminator.replace(r, n)
         self.desc = None
@@ -41,7 +53,7 @@ class CsvfileWriter(AbstractWriter):
         if isinstance(self.exclude, str):
             self.exclude = self.exclude.split(",")
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         rdict = r._asdict(fields=self.fields, exclude=self.exclude)
         if not self.desc or self.desc != r._desc:
             self.desc = r._desc
@@ -49,24 +61,26 @@ class CsvfileWriter(AbstractWriter):
             self.writer.writeheader()
         self.writer.writerow(rdict)
 
-    def flush(self):
+    def flush(self) -> None:
         if self.fp:
             self.fp.flush()
 
-    def close(self):
+    def close(self) -> None:
         if self.fp and not is_stdout(self.fp):
             self.fp.close()
         self.fp = None
 
 
 class CsvfileReader(AbstractReader):
-    def __init__(self, path, selector=None, fields=None, **kwargs):
+    def __init__(
+        self, path: str | Path | None, selector: str | None = None, fields: str | list[str] | None = None, **kwargs
+    ):
         self.fp = None
         self.selector = make_selector(selector)
         if path in (None, "", "-"):
             self.fp = sys.stdin
         else:
-            self.fp = open(path, "r", newline="")
+            self.fp = Path(path).open("r", newline="")  # noqa: SIM115
 
         self.dialect = "excel"
         if self.fp.seekable():
@@ -87,12 +101,12 @@ class CsvfileReader(AbstractReader):
         # Create RecordDescriptor from fields, skipping fields starting with "_" (reserved for internal use)
         self.desc = RecordDescriptor("csv/reader", [("string", col) for col in self.fields if not col.startswith("_")])
 
-    def close(self):
+    def close(self) -> None:
         if self.fp:
             self.fp.close()
         self.fp = None
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Record]:
         for row in self.reader:
             rdict = dict(zip(self.fields, row))
             record = self.desc.init_from_dict(rdict)

{flow_record-3.18.dev2 → flow_record-3.19}/flow/record/adapter/elastic.py

@@ -4,16 +4,26 @@ import hashlib
 import logging
 import queue
 import threading
-from typing import Iterator
+from typing import TYPE_CHECKING
 
-import elasticsearch
-import elasticsearch.helpers
+try:
+    import elasticsearch
+    import elasticsearch.helpers
+
+    HAS_ELASTIC = True
+
+except ImportError:
+    HAS_ELASTIC = False
 
 from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.base import Record, RecordDescriptor
 from flow.record.fieldtypes import fieldtype_for_value
 from flow.record.jsonpacker import JsonRecordPacker
-from flow.record.selector import CompiledSelector, Selector
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.selector import CompiledSelector, Selector
 
 __usage__ = """
 ElasticSearch adapter
@@ -25,9 +35,12 @@ Read usage: rdump elastic+[PROTOCOL]://[IP]:[PORT]?index=[INDEX]
 
 Optional arguments:
   [API_KEY]: base64 encoded api key to authenticate with (default: False)
+  [QUEUE_SIZE]: maximum queue size for writing records; limits memory usage (default: 100000)
   [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)
+  [REQUEST_TIMEOUT]: maximum duration in seconds for a request to Elastic (default: 30)
+  [MAX_RETRIES]: maximum retries before a record is marked as failed (default: 3)
   [_META_*]: record metadata fields (default: None)
 """
 
@@ -43,33 +56,49 @@ class ElasticWriter(AbstractWriter):
         http_compress: str | bool = True,
         hash_record: str | bool = False,
         api_key: str | None = None,
+        queue_size: int = 100000,
+        request_timeout: int = 30,
+        max_retries: int = 3,
         **kwargs,
     ) -> None:
+        """Initialize the ElasticWriter.
+
+        Resources:
+            - https://elasticsearch-py.readthedocs.io/en/v8.17.1/api/elasticsearch.html
+        """
+
+        if not HAS_ELASTIC:
+            raise RuntimeError("Required dependency 'elasticsearch' missing")
+
         self.index = index
         self.uri = uri
         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")
+        queue_size = int(queue_size)
 
         if not uri.lower().startswith(("http://", "https://")):
             uri = "http://" + uri
 
-        self.queue: queue.Queue[Record | StopIteration] = queue.Queue()
+        self.queue: queue.Queue[Record | StopIteration] = queue.Queue(maxsize=queue_size)
         self.event = threading.Event()
+        self.exception: Exception | None = None
+        threading.excepthook = self.excepthook
 
         self.es = elasticsearch.Elasticsearch(
             uri,
             verify_certs=verify_certs,
             http_compress=http_compress,
             api_key=api_key,
+            request_timeout=request_timeout,
+            retry_on_timeout=True,
+            max_retries=max_retries,
         )
 
         self.json_packer = JsonRecordPacker()
 
         self.thread = threading.Thread(target=self.streaming_bulk_thread)
         self.thread.start()
-        self.exception: Exception | None = None
-        threading.excepthook = self.excepthook
 
         if not verify_certs:
             # Disable InsecureRequestWarning of urllib3, caused by the verify_certs flag.
@@ -83,8 +112,8 @@ class ElasticWriter(AbstractWriter):
                 self.metadata_fields[arg_key[6:]] = arg_val
 
     def excepthook(self, exc: threading.ExceptHookArgs, *args, **kwargs) -> None:
-        log.error("Exception in thread: %s", exc.exc_value.message)
-        self.exception = exc.exc_value
+        log.error("Exception in thread: %s", exc)
+        self.exception = getattr(exc, "exc_value", exc)
         self.event.set()
         self.close()
 
@@ -128,24 +157,34 @@ class ElasticWriter(AbstractWriter):
             record = self.queue.get()
             if record is StopIteration:
                 break
+            if not record:
+                continue
             yield self.record_to_document(record, index=self.index)
 
     def streaming_bulk_thread(self) -> None:
-        """Thread that streams the documents to ES via the bulk api"""
+        """Thread that streams the documents to ES via the bulk api.
 
-        for ok, item in elasticsearch.helpers.streaming_bulk(
+        Resources:
+            - https://elasticsearch-py.readthedocs.io/en/v8.17.1/helpers.html#elasticsearch.helpers.streaming_bulk
+            - https://github.com/elastic/elasticsearch-py/blob/main/elasticsearch/helpers/actions.py#L362
+        """
+        for _ok, _item in elasticsearch.helpers.streaming_bulk(
             self.es,
             self.document_stream(),
-            raise_on_error=False,
-            raise_on_exception=False,
+            raise_on_error=True,
+            raise_on_exception=True,
+            # Some settings have to be redefined because streaming_bulk does not inherit them from the self.es instance.
+            max_retries=3,
         ):
-            if not ok:
-                log.error("Failed to insert %r", item)
+            pass
 
         self.event.set()
 
     def write(self, record: Record) -> None:
-        self.queue.put_nowait(record)
+        if self.exception:
+            raise self.exception
+
+        self.queue.put(record)
 
     def flush(self) -> None:
         pass
@@ -170,6 +209,8 @@ class ElasticReader(AbstractReader):
         http_compress: str | bool = True,
         selector: None | Selector | CompiledSelector = None,
         api_key: str | None = None,
+        request_timeout: int = 30,
+        max_retries: int = 3,
         **kwargs,
     ) -> None:
         self.index = index
@@ -186,6 +227,9 @@ class ElasticReader(AbstractReader):
             verify_certs=verify_certs,
             http_compress=http_compress,
             api_key=api_key,
+            request_timeout=request_timeout,
+            retry_on_timeout=True,
+            max_retries=max_retries,
         )
 
         if not verify_certs: