"""Demultiplex one sample stream into one BDF file per (series, run), with rotation.
BDF's invariant is **one test object, one monotonic timebase, per file** -- but
the world violates it in three directions: one connection can yield many
objects (an account with N cars, a drone with N packs), one object can yield
many runs (a pack flies many flights, each restarting its clock), and many
sources never end at all (a shunt streams forever). :class:`RoutingSink`
resolves all three at the sink layer:
* Samples carrying the reserved routing keys (``series_id`` / ``run_id``, see
:data:`battfeed.RESERVED_KEYS`) are demultiplexed into one child
:class:`~battfeed.BdfCsvSink` per ``(series_id, run_id)``. Samples without
``series_id`` flow to a single default stream, so routing-free sources work
unchanged.
* A new ``run_id`` for a known series closes the previous file and opens the
next one -- runs never merge into a non-monotonic timebase.
* **Rotation** (``rotate_after_s`` / ``rotate_after_rows``, whichever trips
first) does the same *without* a ``run_id`` change: segments *are* runs for
endless streams, turning unbounded telemetry into a sequence of bounded,
valid BDF files (each finalized with its sidecar as soon as it rotates, so a
crash loses at most the open segments -- and even those keep the early,
unfinalised sidecar the child sink writes at open time).
Filenames follow the BDF convention via
:func:`~battfeed.sinks.bdf_csv.dataset_filename`; the ``_XXX`` sequence slot
advances per run/segment for the same series on the same day. Raw
``series_id`` values are device serials, not filenames -- they may contain
``__`` (reserved as the BDF filename separator), path separators, characters
Windows forbids, or be empty -- so they pass through
:func:`sanitize_cell_name`, with deterministic hash-suffix disambiguation when
two distinct ids sanitize to the same name.
**Timebase ownership (invariant I5).** This sink routes and never restamps:
``test_time_second`` is written exactly as the source supplied it. A source
that emits routing keys must therefore supply its own ``test_time_second``,
zero-based per (series, run) -- the harvester's shared elapsed-collection
stamp is wrong for an object that appears mid-run (see
:mod:`battfeed.protocols`).
"""
from __future__ import annotations
import datetime
import hashlib
import logging
import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Iterable, Mapping
from ..protocols import RESERVED_KEYS, SampleValue, Sink
from .bdf_csv import BdfCsvSink, dataset_filename
__all__ = ["RoutingSink", "sanitize_cell_name"]
logger = logging.getLogger(__name__)
#: Routing keys stripped from every row before delegation (defense in depth --
#: the child sink strips them again; see ``protocols.RESERVED_KEYS``).
_RESERVED: frozenset[str] = frozenset(RESERVED_KEYS)
#: Cell name used for samples that carry no ``series_id``. Reserved at
#: construction so a literal series id ``"default"`` can never collide with it.
_DEFAULT_CELL_NAME = "default"
#: Cap on sanitized cell names, keeping full dataset filenames comfortably
#: inside Windows path limits.
_MAX_CELL_NAME_LEN = 60
#: Characters illegal in Windows filenames (plus the path separators).
_ILLEGAL_CHARS = frozenset('<>:"/\\|?*')
[docs]
def sanitize_cell_name(raw: str) -> str:
"""Turn a raw series id (e.g. a device serial) into a safe BDF cell name.
Raw ids are whatever the device reports: they may contain ``__`` (reserved
as the BDF filename separator), path separators, characters illegal on
Windows, control characters, or be empty. The result is deterministic,
non-empty, at most 60 characters, and never contains ``__`` -- nor starts
or ends with ``_`` (which would recreate ``__`` next to the filename
separators).
Distinct raw ids can sanitize to the same name (``"pack/1"`` and
``"pack?1"`` both become ``"pack-1"``, and names differing only by case
collide too -- Windows filesystems are case-insensitive);
:class:`RoutingSink` disambiguates such collisions with a stable hash
suffix derived from the raw id.
"""
cleaned = "".join("-" if ch in _ILLEGAL_CHARS or ord(ch) < 32 else ch for ch in raw)
cleaned = re.sub(r"\s+", "-", cleaned)
cleaned = re.sub(r"_{2,}", "-", cleaned) # "__" is the BDF filename separator
cleaned = re.sub(r"-{2,}", "-", cleaned)
cleaned = cleaned[:_MAX_CELL_NAME_LEN]
cleaned = cleaned.strip("-._ ")
return cleaned or "unknown"
@dataclass
class _SeriesState:
"""Book-keeping for one demultiplexed series (or the default stream)."""
series_id: str | None
cell_name: str
series_metadata: dict[str, Any]
run_id: str | None = None
segment: int = 0
sink: Sink | None = None
opened_at: float = 0.0
rows_in_segment: int = 0
files: list[Path] = field(default_factory=list)
[docs]
class RoutingSink:
"""Route one sample stream into one BDF file per (series, run), rotating.
Implements the :class:`battfeed.Sink` protocol, so it drops in anywhere a
plain :class:`battfeed.BdfCsvSink` does -- the harvester never knows the
difference. Child files open lazily as series/runs appear, so a series that
comes online mid-stream gets its own file and sidecar from its first
sample; closing a segment (run change, rotation, or :meth:`close`)
finalizes that segment's sidecar immediately.
Per-segment sidecar metadata is the shared ``metadata`` mapping, overlaid
with the ``series_info`` metadata for the series, overlaid with
``{"series_id": ..., "run_id": ..., "segment": n}`` (``segment`` counts
from 1 per series).
Args:
directory: Directory the ``.bdf.csv`` files (and their sidecars) are
written into; created on demand.
institution: Institution code for :func:`dataset_filename`.
series_info: Optional ``callable(series_id) -> (cell_name,
metadata_dict)`` hook so lazily-discovered objects get proper
filenames and their own sidecar content. Called once per new
series; without it the sanitized ``series_id`` names the cell.
metadata: Shared base metadata recorded in every segment's sidecar.
rotate_after_s: Close and re-open a series' file once it has been open
this many seconds (measured with ``clock``).
rotate_after_rows: Close and re-open a series' file once it holds this
many rows. With both limits set, whichever trips first rotates.
sink_factory: ``callable(path, *, metadata) -> Sink`` building each
child sink; injectable so tests can capture routed rows in memory.
(Path allocation still reserves each claimed path on disk as an
empty file -- see :meth:`_next_path` -- regardless of the factory.)
clock: Monotonic clock driving time-based rotation; injectable for
tests (see the Harvester's clock/sleep injection).
today: Date provider for filenames; injectable so tests get
deterministic names.
Unlike ``BdfCsvSink``, closing without ever writing produces no files:
there is nothing to route, so nothing is (even emptily) recorded.
"""
def __init__(
self,
directory: str | Path,
*,
institution: str = "LOCAL",
series_info: Callable[[str], tuple[str, Mapping[str, Any]]] | None = None,
metadata: Mapping[str, Any] | None = None,
rotate_after_s: float | None = None,
rotate_after_rows: int | None = None,
sink_factory: Callable[..., Sink] = BdfCsvSink,
clock: Callable[[], float] = time.monotonic,
today: Callable[[], datetime.date] = datetime.date.today,
) -> None:
if not institution:
raise ValueError("institution must be non-empty")
if "__" in institution:
raise ValueError(
"institution must not contain '__' (reserved as the BDF filename "
f"separator): {institution!r}"
)
if rotate_after_s is not None and rotate_after_s <= 0:
raise ValueError(f"rotate_after_s must be positive, got {rotate_after_s}")
if rotate_after_rows is not None and rotate_after_rows < 1:
raise ValueError(f"rotate_after_rows must be >= 1, got {rotate_after_rows}")
self._directory = Path(directory)
self._institution = institution
self._series_info = series_info
self._metadata = dict(metadata or {})
self._rotate_after_s = rotate_after_s
self._rotate_after_rows = rotate_after_rows
self._sink_factory = sink_factory
self._clock = clock
self._today = today
self._series: dict[str | None, _SeriesState] = {}
# casefolded cell_name -> raw series id that owns it ("" owns the
# default stream). Keyed case-insensitively because the dominant
# filesystems are: "CellA" and "cella" would be one file on NTFS.
self._claimed: dict[str, str] = {_DEFAULT_CELL_NAME: ""}
self._seq: dict[tuple[str, str], int] = {}
self._closed = False
@property
def directory(self) -> Path:
return self._directory
@property
def files_by_series(self) -> dict[str | None, list[Path]]:
"""Paths written so far, keyed by ``series_id`` (``None`` = default stream).
Segments appear in the order they were opened; the last entry of a list
may still be open. Returns a copy.
"""
return {series: list(state.files) for series, state in self._series.items()}
[docs]
def write(self, rows: Iterable[Mapping[str, SampleValue]]) -> None:
"""Route a batch of samples to their per-(series, run) child sinks.
Reserved routing keys are stripped from every row before delegation;
everything else -- including any source-supplied ``test_time_second``
-- is passed through untouched (invariant I5: the source owns the
per-(series, run) timebase).
"""
if self._closed:
raise ValueError(f"RoutingSink for {self._directory} is closed")
grouped: dict[str | None, list[tuple[str | None, dict[str, SampleValue]]]] = {}
for row in rows:
series_value = row.get("series_id")
run_value = row.get("run_id")
series = str(series_value) if series_value is not None else None
run = str(run_value) if run_value is not None else None
stripped = {key: value for key, value in row.items() if key not in _RESERVED}
grouped.setdefault(series, []).append((run, stripped))
for series, items in grouped.items():
self._write_series(self._state_for(series), items)
[docs]
def close(self) -> None:
"""Close every open child sink (finalising its sidecar). Idempotent."""
if self._closed:
return
self._closed = True
first_error: Exception | None = None
for state in self._series.values():
if state.sink is None:
continue
try:
self._close_segment(state)
except Exception as exc: # noqa: BLE001 -- keep closing the other children
logger.exception("Failed to close child sink for series %r", state.series_id)
if first_error is None:
first_error = exc
state.sink = None
if first_error is not None:
raise first_error
# -- internals --------------------------------------------------------
def _state_for(self, series: str | None) -> _SeriesState:
state = self._series.get(series)
if state is not None:
return state
if series is None:
cell_name = _DEFAULT_CELL_NAME
series_metadata: dict[str, Any] = {}
else:
if self._series_info is not None:
info_name, info_metadata = self._series_info(series)
base = sanitize_cell_name(info_name)
series_metadata = dict(info_metadata)
else:
base = sanitize_cell_name(series)
series_metadata = {}
cell_name = self._allocate_cell_name(base, series)
state = _SeriesState(series_id=series, cell_name=cell_name, series_metadata=series_metadata)
self._series[series] = state
logger.info("New series %r routed to cell name %r", series, cell_name)
return state
def _allocate_cell_name(self, base: str, series: str) -> str:
"""Claim a unique cell name for ``series``, hash-suffixing collisions.
Claims are matched **case-insensitively** (casefolded), because the
filesystems this targets are: two names differing only by case would
entangle on NTFS/APFS. The returned name keeps the original casing.
Plain (unsuffixed) names are **first-come-first-served**: the first
series to sanitize to ``base`` keeps the plain name, and every later
distinct series gets a stable suffix derived from its raw id. Filename
assignment is therefore **arrival-order dependent across runs** -- the
same series can be plain in one collection and suffixed in the next if
its rival arrived first. Do not join datasets on filenames: the
sidecar's recorded raw ``series_id`` is the stable join key.
"""
if base.casefold() not in self._claimed:
self._claimed[base.casefold()] = series
return base
digest = hashlib.sha256(series.encode("utf-8")).hexdigest()
for length in (8, 16, 32, 64):
candidate = f"{base}-{digest[:length]}"
if candidate.casefold() not in self._claimed:
self._claimed[candidate.casefold()] = series
logger.info(
"Series %r collides with %r on cell name %r; using %r",
series,
self._claimed[base.casefold()],
base,
candidate,
)
return candidate
raise RuntimeError( # pragma: no cover -- would need a sha256 collision
f"Could not allocate a unique cell name for series {series!r}"
)
def _write_series(
self,
state: _SeriesState,
items: list[tuple[str | None, dict[str, SampleValue]]],
) -> None:
index = 0
while index < len(items):
run_id = items[index][0]
if state.sink is not None and (run_id != state.run_id or self._rotation_due(state)):
self._close_segment(state)
if state.sink is None:
self._open_segment(state, run_id)
end = index
while end < len(items) and items[end][0] == run_id:
end += 1
chunk = [row for _, row in items[index:end]]
if self._rotate_after_rows is not None:
chunk = chunk[: self._rotate_after_rows - state.rows_in_segment]
assert state.sink is not None
state.sink.write(chunk)
state.rows_in_segment += len(chunk)
index += len(chunk)
if (
self._rotate_after_rows is not None
and state.rows_in_segment >= self._rotate_after_rows
):
# Rotate eagerly at the row limit so a full segment is
# finalized (sidecar and all) the moment it fills.
self._close_segment(state)
def _rotation_due(self, state: _SeriesState) -> bool:
return (
self._rotate_after_s is not None
and self._clock() - state.opened_at >= self._rotate_after_s
)
def _open_segment(self, state: _SeriesState, run_id: str | None) -> None:
state.segment += 1
state.run_id = run_id
path = self._next_path(state.cell_name)
segment_metadata: dict[str, Any] = {
**self._metadata,
**state.series_metadata,
"series_id": state.series_id,
"run_id": run_id,
"segment": state.segment,
}
state.sink = self._sink_factory(path, metadata=segment_metadata)
state.opened_at = self._clock()
state.rows_in_segment = 0
state.files.append(path)
logger.info(
"Opened segment %d for series %r (run %r): %s",
state.segment,
state.series_id,
run_id,
path,
)
def _next_path(self, cell_name: str) -> Path:
"""Next free dataset path: the ``_XXX`` slot advances per (cell, day).
The claimed path is **reserved atomically** at allocation time by
creating it with ``open(..., "x")`` (``O_CREAT | O_EXCL``): a bare
``exists()`` check would leave a check-then-claim window in which two
concurrent collections into the same directory could pick the same
sequence number and silently clobber each other's file. On
``FileExistsError`` the sequence number advances and the claim is
retried (``dataset_filename`` raises loudly past 999).
The reservation is just an empty file; the child ``BdfCsvSink`` later
reopens it with mode ``"w"``, truncating it before writing the header,
so nothing downstream changes. Segments open lazily, only when a row
is about to be written, so a reservation is never left row-less by
normal operation -- only a crash inside the reserve-to-write window
can leave an empty ``.bdf.csv``, and that file still (correctly)
marks its sequence number as taken.
"""
date = self._today()
key = (cell_name.casefold(), date.isoformat())
seq = self._seq.get(key, 1)
self._directory.mkdir(parents=True, exist_ok=True)
while True:
path = self._directory / dataset_filename(self._institution, cell_name, date, seq)
try:
with open(path, "x", encoding="utf-8"):
pass
except FileExistsError:
seq += 1
continue
break
self._seq[key] = seq + 1
return path
def _close_segment(self, state: _SeriesState) -> None:
assert state.sink is not None
sink = state.sink
state.sink = None
sink.close()
logger.info(
"Closed segment %d of series %r (%d row(s))",
state.segment,
state.series_id,
state.rows_in_segment,
)