dirigent-blocks 0.9.0__py3-none-any.whl

dirigent_blocks/capture.py

@@ -0,0 +1,289 @@
+"""Capturing what a local process printed: decode it, log it, store it, quote it."""
+
+import asyncio
+from collections.abc import Sequence
+from typing import NamedTuple
+
+from dirigent_plugin import ByteSink, StepContext
+
+LOGGED_LINES = 40
+
+FAILURE_TAIL_LINES = 5
+
+#: What a scrubbed secret is replaced by, matching the marker the API redacts a connection with.
+REDACTED = "***"
+
+
+class Captured(NamedTuple):
+    """One stream as an output carries it: the head that inlines, and the truth about the rest."""
+
+    text: str
+    """The head of the stream, as much of it as the instance lets a block inline."""
+
+    total_bytes: int
+    """How much the process printed, whether it inlined or not."""
+
+    truncated: bool
+    """Whether the text above is short of the stream; the artifact URI always holds all of it."""
+
+
+def captured(payload: bytes, limit: int) -> Captured:
+    """Cut one captured stream down to what may be inlined, and say what was cut."""
+    return Captured(text=decode(payload[:limit]), total_bytes=len(payload), truncated=len(payload) > limit)
+
+
+#: How much of the end of a stream is kept, for the lines a failure message quotes. The head
+#: is what a block may inline; this is the other end, and together they bound what is held.
+TAIL_BYTES = 8 * 1024
+
+#: How much is read from a pipe at a time.
+CHUNK_BYTES = 64 * 1024
+
+
+class Logged(NamedTuple):
+    """What a stream put into the run log while it was still being read."""
+
+    lines: int
+    """Non-blank lines the stream carried."""
+
+    live: int
+    """How many of them reached the run log as they arrived."""
+
+
+class Drained(NamedTuple):
+    """What was kept of a stream that was written through to storage as it arrived."""
+
+    head: bytes
+    """The beginning, as much as may be inlined."""
+
+    tail: bytes
+    """The end, for the lines a failure quotes."""
+
+    total_bytes: int
+    """How much the process printed, all of which reached storage."""
+
+    logged: Logged | None = None
+    """What was already logged live, or ``None`` when the stream was not logged as it arrived."""
+
+    def captured(self, limit: int) -> Captured:
+        """Present the head as an output carries it, against what the whole stream was."""
+        return Captured(
+            text=decode(self.head[:limit]),
+            total_bytes=self.total_bytes,
+            truncated=self.total_bytes > limit,
+        )
+
+
+class Ends:
+    """Both ends of a stream and its size, holding nothing of the middle.
+
+    A stream is fed in as it arrives. What is kept is the head, which is what may be inlined,
+    and the last few kilobytes, which is what a failure message quotes.
+    """
+
+    def __init__(self, head_limit: int) -> None:
+        """Keep at most this much of the beginning."""
+        self._head_limit = head_limit
+        self._head = bytearray()
+        self._tail = bytearray()
+        self.total_bytes = 0
+
+    def add(self, chunk: bytes) -> None:
+        """Take one piece of the stream."""
+        self.total_bytes += len(chunk)
+        if len(self._head) < self._head_limit:
+            self._head.extend(chunk[: self._head_limit - len(self._head)])
+        self._tail.extend(chunk)
+        if len(self._tail) > TAIL_BYTES:
+            del self._tail[: len(self._tail) - TAIL_BYTES]
+
+    def drained(self, live: "LiveLines | None" = None) -> Drained:
+        """Present what was kept, and what a live logger already said about it."""
+        return Drained(
+            head=bytes(self._head),
+            tail=bytes(self._tail),
+            total_bytes=self.total_bytes,
+            logged=live.logged() if live is not None else None,
+        )
+
+
+class LiveLines:
+    """Puts the first lines of a stream into the run log as the process prints them.
+
+    Bytes are fed in as they are read. Every complete non-blank line is counted, and the
+    first ``budget`` of them are logged where they happen, so a run screen shows a long
+    command working rather than everything it printed at the moment it exited. Past the
+    budget nothing is held: the line is counted and dropped, and the other end of the stream
+    reaches the log when the stream closes.
+    """
+
+    def __init__(
+        self,
+        ctx: StepContext,
+        stream: str,
+        *,
+        budget: int = LOGGED_LINES // 2,
+        redact: Sequence[str] = (),
+    ) -> None:
+        """Log to this stream's level, redacting the strings the block handed the process."""
+        self._log = ctx.log.info if stream == "stdout" else ctx.log.warning
+        self._stream = stream
+        self._budget = budget
+        self._redact = redact
+        self._pending = bytearray()
+        self._nonblank = False
+        self._lines = 0
+        self._live = 0
+
+    def feed(self, chunk: bytes) -> None:
+        """Take one piece of the stream, logging every complete line it completes."""
+        start = 0
+        while (cut := chunk.find(b"\n", start)) >= 0:
+            self._take(chunk[start:cut])
+            self._finish()
+            start = cut + 1
+        self._take(chunk[start:])
+
+    def close(self) -> None:
+        """Log a last line the process left without a newline."""
+        self._finish()
+
+    def logged(self) -> Logged:
+        """How many lines the stream carried and how many of them are in the log already."""
+        return Logged(lines=self._lines, live=self._live)
+
+    def _take(self, piece: bytes) -> None:
+        """Hold a piece of the line being read, or only whether it has anything in it."""
+        if not piece:
+            return
+        self._nonblank = self._nonblank or bool(piece.strip())
+        if self._live < self._budget:
+            self._pending.extend(piece)
+
+    def _finish(self) -> None:
+        """End the line being read: count it, and log it while the budget lasts."""
+        if self._nonblank:
+            self._lines += 1
+            if self._live < self._budget:
+                self._live += 1
+                self._log(scrub(decode(bytes(self._pending)).rstrip("\r"), self._redact), stream=self._stream)
+        self._pending.clear()
+        self._nonblank = False
+
+
+async def drain(
+    reader: "asyncio.StreamReader | None",
+    sink: ByteSink,
+    *,
+    head_limit: int,
+    live: LiveLines | None = None,
+) -> Drained:
+    """Read one pipe to its end, writing it through to storage and keeping both ends.
+
+    Nothing holds the whole stream. A command that prints a gigabyte is a file in storage and
+    two bounded buffers here, rather than a gigabyte in the worker.
+
+    With a ``live`` logger the lines go into the run log as they are read, up to its budget.
+    """
+    ends = Ends(head_limit)
+    while reader is not None:
+        chunk = await reader.read(CHUNK_BYTES)
+        if not chunk:
+            break
+        ends.add(chunk)
+        if live is not None:
+            live.feed(chunk)
+        await sink.write(chunk)
+    if live is not None:
+        live.close()
+    return ends.drained(live)
+
+
+def decode(payload: bytes) -> str:
+    """Render captured bytes as text without ever failing on what a process printed."""
+    return payload.decode("utf-8", errors="replace")
+
+
+def lines_of(payload: bytes) -> list[str]:
+    """Split a captured stream into its non-blank lines."""
+    return [line for line in decode(payload).splitlines() if line.strip()]
+
+
+def scrub(text: str, secrets: Sequence[str] = ()) -> str:
+    """Replace every occurrence of a secret with the redaction marker.
+
+    A tool prints what it was given. ``git`` echoes a remote URL, an ssh command line and,
+    when a credential is wrong, sometimes the credential itself, so a block that hands a
+    process a secret hands this the same strings and nothing reaches a log or a failure
+    message with them still in it. Longest first, so a secret that contains another is
+    replaced whole rather than left in pieces.
+    """
+    for secret in sorted((one for one in secrets if one), key=len, reverse=True):
+        text = text.replace(secret, REDACTED)
+    return text
+
+
+def tail(
+    payload: bytes,
+    *,
+    count: int = FAILURE_TAIL_LINES,
+    separator: str = " / ",
+    redact: Sequence[str] = (),
+) -> str:
+    """Take the last few lines of a stream, which is what a failure message wants."""
+    return scrub(separator.join(lines_of(payload)[-count:]), redact)
+
+
+async def store(ctx: StepContext, uri: str, payload: bytes) -> None:
+    """Write one captured stream to storage, so it outlives the worker that produced it."""
+    async with ctx.storage.open_write(uri) as sink:
+        await sink.write(payload)
+
+
+def log_stream(ctx: StepContext, stream: str, drained: Drained, redact: Sequence[str] = ()) -> None:
+    """Put what the process printed into the run's log, from both ends of the stream.
+
+    A caller that read the stream through to storage holds only its two ends, and identical
+    ends are how one small enough to be held whole says so. The whole of a stream is elided
+    by line count; what is held in two pieces can only be logged as two pieces.
+
+    A stream that was logged live already has its first lines in the log: this adds the last
+    of them, and only the ones that were never said.
+
+    ``redact`` names the strings a block handed the process that must not reach the log.
+    """
+    log = ctx.log.info if stream == "stdout" else ctx.log.warning
+    half = LOGGED_LINES // 2
+    if drained.logged is not None:
+        remaining = drained.logged.lines - drained.logged.live
+        if remaining <= 0:
+            return
+        last = lines_of(drained.tail)[-min(remaining, half) :]
+        if remaining > len(last):
+            log(f"... {remaining - len(last)} more lines, see the {stream} artifact ...", stream=stream)
+        for line in last:
+            log(scrub(line, redact), stream=stream)
+        return
+    if drained.head == drained.tail:
+        lines = lines_of(drained.head)
+        if not lines:
+            return
+        shown = lines
+        if len(lines) > LOGGED_LINES:
+            shown = [
+                *lines[:half],
+                f"... {len(lines) - LOGGED_LINES} more lines, see the {stream} artifact ...",
+                *lines[-half:],
+            ]
+        for line in shown:
+            log(scrub(line, redact), stream=stream)
+        return
+    first = lines_of(drained.head)
+    last = lines_of(drained.tail)
+    if not first and not last:
+        return
+    for line in first[:half]:
+        log(scrub(line, redact), stream=stream)
+    log(f"... the middle is not in the log, see the {stream} artifact ...", stream=stream)
+    for line in last[-half:]:
+        log(scrub(line, redact), stream=stream)

dirigent_blocks/clock.py

@@ -0,0 +1,220 @@
+"""Clock sensors: ``time.window`` waits for a window to open, ``time.sleep`` waits a duration."""
+
+from datetime import UTC, date, datetime, time, timedelta
+from typing import ClassVar, Final, Literal
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from dirigent_common import BlockModel, Duration
+from dirigent_plugin import NotYet, Sensor, SensorSpec, StepContext
+
+type DayName = Literal["mon", "tue", "wed", "thu", "fri", "sat", "sun"]
+
+#: Indexed with ``date.weekday()``, so the order is that function's and not a preference.
+WEEKDAYS: Final[tuple[DayName, ...]] = ("mon", "tue", "wed", "thu", "fri", "sat", "sun")
+
+DEFAULT_TIMEZONE: Final = "UTC"
+
+MAX_WAIT: Final = timedelta(hours=1)
+
+MIN_WAIT: Final = timedelta(seconds=1)
+
+SEARCH_DAYS: Final = 8
+
+#: The shortest park a sleep asks for, so the last sliver of a wait is one poke and not many.
+MIN_SLEEP_POLL: Final = timedelta(milliseconds=100)
+
+
+class TimeWindowConfig(BlockModel):
+    """The window a step may proceed in, in one named timezone."""
+
+    after: time = time(0, 0)
+    """The local time the window opens."""
+
+    before: time = time(23, 59, 59)
+    """The local time it closes; earlier than ``after`` means the window crosses midnight."""
+
+    timezone: str = DEFAULT_TIMEZONE
+    """The IANA zone the window is read in. A window without one is a window in someone's head."""
+
+    days: list[DayName] = Field(default_factory=list["DayName"])
+    """The days the window opens on; empty means every day."""
+
+    @field_validator("timezone")
+    @classmethod
+    def _check_timezone(cls, value: str) -> str:
+        """Reject a zone this machine's database does not have, at validation and not at poke."""
+        try:
+            ZoneInfo(value)
+        except (ZoneInfoNotFoundError, ValueError) as error:
+            raise ValueError(f"{value!r} is not an IANA timezone name, such as 'Europe/Oslo' or 'UTC'") from error
+        return value
+
+    @model_validator(mode="after")
+    def _check_window(self) -> "TimeWindowConfig":
+        """Reject a window of zero width, which no clock is ever inside."""
+        if self.after == self.before:
+            raise ValueError(
+                f"a window from {self.after} to {self.before} is empty; "
+                f"omit the sensor rather than writing a window nothing falls in"
+            )
+        return self
+
+    @property
+    def zone(self) -> ZoneInfo:
+        """The zone object the window is evaluated in."""
+        return ZoneInfo(self.timezone)
+
+    @property
+    def crosses_midnight(self) -> bool:
+        """Report whether this window opens on one day and closes on the next."""
+        return self.after > self.before
+
+    def opens_on(self, day: date) -> bool:
+        """Report whether a window opens on a given date."""
+        return not self.days or WEEKDAYS[day.weekday()] in self.days
+
+    def opening(self, day: date) -> datetime:
+        """The moment the window opens on a given date, in its own zone."""
+        return datetime.combine(day, self.after, tzinfo=self.zone)
+
+
+class TimeWindowOutput(BlockModel):
+    """The observation that the clock is inside the window, passed downstream like any output."""
+
+    entered_at: datetime
+    """When the window this poke fell inside opened, not when the poke happened."""
+
+    timezone: str
+
+
+class TimeWindowSensor(Sensor[TimeWindowConfig, TimeWindowOutput]):
+    """Waits until the local clock in a named timezone is inside a window."""
+
+    spec = SensorSpec(
+        id="time.window",
+        summary="Wait until the local clock is inside a time window.",
+        default_poll=timedelta(minutes=1),
+        default_deadline=timedelta(hours=24),
+    )
+    config_model: ClassVar[type[BaseModel]] = TimeWindowConfig
+    output_model: ClassVar[type[BaseModel]] = TimeWindowOutput
+
+    async def poke(self, config: TimeWindowConfig, ctx: StepContext) -> TimeWindowOutput | NotYet:
+        """Look at one clock once, and either proceed or say how long the wait is."""
+        moment = now_in(config.zone)
+        entered = entered_at(config, moment)
+        if entered is not None:
+            # The stream carries the time of every line, so a block repeating it says nothing
+            # the reader cannot already see; the window it opened is in the output.
+            ctx.log.info("the window is open", timezone=config.timezone)
+            return TimeWindowOutput(entered_at=entered, timezone=config.timezone)
+        wait = wait_for(config, moment)
+        ctx.log.debug("the window is closed", wait_seconds=round(wait.total_seconds(), 1), timezone=config.timezone)
+        return NotYet(next_poll_in=wait)
+
+
+def now_in(zone: ZoneInfo) -> datetime:
+    """Read the current moment in one zone."""
+    return datetime.now(zone)
+
+
+def entered_at(config: TimeWindowConfig, moment: datetime) -> datetime | None:
+    """Return when the window the moment falls inside opened, or None when it is closed."""
+    local = moment.astimezone(config.zone)
+    today = local.date()
+    clock = local.timetz().replace(tzinfo=None)
+    if not config.crosses_midnight:
+        if config.after <= clock < config.before and config.opens_on(today):
+            return config.opening(today)
+        return None
+    if clock >= config.after and config.opens_on(today):
+        return config.opening(today)
+    yesterday = today - timedelta(days=1)
+    if clock < config.before and config.opens_on(yesterday):
+        return config.opening(yesterday)
+    return None
+
+
+def next_opening(config: TimeWindowConfig, moment: datetime) -> datetime | None:
+    """Find the next moment the window opens, searching a little over a week ahead."""
+    local = moment.astimezone(config.zone)
+    for offset in range(SEARCH_DAYS):
+        day = local.date() + timedelta(days=offset)
+        if not config.opens_on(day):
+            continue
+        opening = config.opening(day)
+        if opening > local:
+            return opening
+    return None  # pragma: no cover - a non-empty day set always opens within eight days
+
+
+def wait_for(config: TimeWindowConfig, moment: datetime) -> timedelta:
+    """Say how long to park a closed window, bounded at both ends."""
+    opening = next_opening(config, moment)
+    if opening is None:  # pragma: no cover - next_opening always finds one
+        return MAX_WAIT
+    return max(MIN_WAIT, min(opening - moment.astimezone(config.zone), MAX_WAIT))
+
+
+class TimeSleepConfig(BlockModel):
+    """How long the wait lasts."""
+
+    wait_for: Duration = Field(alias="for")
+    """How long to wait, measured from when the attempt started.
+
+    The step's deadline still ends the wait, so a `for` longer than a day needs a `deadline`
+    beside it: the sensor default is 24 hours."""
+
+
+class TimeSleepOutput(BlockModel):
+    """What the wait amounted to, passed downstream like any other output."""
+
+    started_at: datetime
+    """When the wait began, which is when the attempt started and not when a poke ran."""
+
+    waited_ms: int
+    """How long the wait actually lasted, which is the configured duration plus poll latency."""
+
+
+class TimeSleepSensor(Sensor[TimeSleepConfig, TimeSleepOutput]):
+    """Waits a fixed duration without occupying a worker.
+
+    Each poke reads the clock and returns, so an hour's wait is a parked row rather than an
+    hour of a worker's concurrency. This is the block to reach for instead of `shell.run`
+    with `sleep`, which asks an instance to allowlist arbitrary code execution in order to do
+    something harmless.
+
+    The deadline is anchored on the attempt's start, so a worker restarting mid-wait resumes
+    the same wait instead of starting the clock again.
+    """
+
+    spec = SensorSpec(
+        id="time.sleep",
+        summary="Wait a fixed duration.",
+        default_poll=timedelta(seconds=1),
+        default_deadline=timedelta(hours=24),
+    )
+    config_model: ClassVar[type[BaseModel]] = TimeSleepConfig
+    output_model: ClassVar[type[BaseModel]] = TimeSleepOutput
+
+    async def poke(self, config: TimeSleepConfig, ctx: StepContext) -> TimeSleepOutput | NotYet:
+        """Read the clock once: finish, or say how much of the wait is left."""
+        moment = datetime.now(UTC)
+        remaining = ctx.started_at + config.wait_for - moment
+        if remaining > timedelta(0):
+            ctx.log.debug(
+                "still waiting",
+                waited_ms=milliseconds(moment - ctx.started_at),
+                remaining_ms=milliseconds(remaining),
+            )
+            return NotYet(next_poll_in=max(remaining, MIN_SLEEP_POLL))
+        waited = moment - ctx.started_at
+        ctx.log.info("the wait is over", waited_ms=milliseconds(waited))
+        return TimeSleepOutput(started_at=ctx.started_at, waited_ms=milliseconds(waited))
+
+
+def milliseconds(value: timedelta) -> int:
+    """Render a duration as the whole milliseconds emitted data measures timings in."""
+    return round(value.total_seconds() * 1000)