publish

Publisher API.

Publisher

Publisher(publisher_id, url=None, replay_id=None, *, token=None)

Publish timeseries data to Arrakis.

Parameters:

Name	Type	Description	Default
publisher_id	str	Publisher ID string.	required
url	str | None	Initial Flight URL to connect to. Will be automatically determined if not specified.	None
replay_id	str | None	If set, publish under this replay context. Channels will be written to replay-namespaced topics.	None
token	str, bool, or None	Controls authentication. None (default) auto-discovers a token via igwn-auth-utils, falling back to unauthenticated. True auto-discovers but raises if no token is found. False disables authentication. A string is used as the raw JWT directly.	None

Source code in arrakis/publish.py

def __init__(
    self,
    publisher_id: str,
    url: str | None = None,
    replay_id: str | None = None,
    *,
    token: str | bool | None = None,
):
    """Initialize Publisher.

    Parameters
    ----------
    publisher_id : str
        Publisher ID string.
    url : str | None
        Initial Flight URL to connect to.  Will be automatically
        determined if not specified.
    replay_id : str | None
        If set, publish under this replay context.  Channels will
        be written to replay-namespaced topics.
    token : str, bool, or None
        Controls authentication.  ``None`` (default) auto-discovers
        a token via igwn-auth-utils, falling back to
        unauthenticated.  ``True`` auto-discovers but raises if no
        token is found.  ``False`` disables authentication.  A
        string is used as the raw JWT directly.

    """
    if not HAS_KAFKA:
        msg = (
            "Publishing requires confluent-kafka to be installed. "
            "It can be installed through pip or conda."
        )
        raise ImportError(msg)

    self.publisher_id = publisher_id
    self.replay_id = replay_id
    self.initial_url = parse_arrakis_url(url).geturl()
    # discovery must request the publisher-scoped path: a bare
    # authz means path "/", which a token restricted to
    # arrakis.create:/<publisher_id> does not satisfy
    self._token = resolve_token(
        token, self.initial_url, scope=f"{CREATE_SCOPE}:/{publisher_id}"
    )
    self._middleware = build_auth_middleware(self._token)

    self.channels: dict[str, Channel] = {}
    self.stride: int | None = None

    self._producer: Producer
    self._validator = RequestValidator()
    self._last_published_ns: int | None = None

close

close()

Exit publication context manager.

Source code in arrakis/publish.py

def close(self) -> None:
    """Exit publication context manager."""
    logger.info("closing kafka producer...")
    with contextlib.suppress(Exception):
        # bounded: an unreachable broker must not hang shutdown
        pending = self._producer.flush(10)
        if pending:
            logger.warning("abandoning %d undelivered message(s) at close", pending)

enter

enter()

Enter publication context manager

Source code in arrakis/publish.py

def enter(self) -> None:
    """Enter publication context manager"""
    # get connection properties
    producer_info: dict[str, str] = {}
    descriptor = create_descriptor(
        RequestType.Publish,
        publisher_id=self.publisher_id,
        replay_id=self.replay_id,
        validator=self._validator,
    )
    with connect(self.initial_url, middleware=self._middleware) as client:
        flight_info = client.get_flight_info(descriptor)
        with MultiFlightReader(
            flight_info.endpoints, client, middleware=self._middleware
        ) as stream:
            for data in stream.unpack():
                kv_pairs = data["properties"]
                producer_info.update(dict(kv_pairs))
    logger.info(
        "publishing %d channel(s) as publisher %s%s",
        len(self.channels),
        self.publisher_id,
        f" (replay {self.replay_id})" if self.replay_id else "",
    )
    logger.debug("producer info: %s", producer_info)

    # set up producer; librdkafka's own logs (broker failures,
    # authentication errors, ...) are routed into this logger
    # rather than written straight to stderr
    self._producer = Producer(
        {
            "message.max.bytes": 10_000_000,  # 10 MB
            "enable.idempotence": True,
            "logger": logger,
            **producer_info,
        }
    )

publish

publish(block, timeout=None)

Publish timeseries data

Parameters:

Name	Type	Description	Default
block	SeriesBlock	A data block with all channels to publish.	required
timeout	timedelta	The maximum time to wait for delivery confirmation before timing out. By default there is no bound of our own: publishing blocks until every message definitively resolves -- delivered, or failed by the producer once its per-message timeout (message.timeout.ms, which caps its internal retries) expires. Transient broker unavailability thus stalls publishing rather than failing it, and only definitive failures raise.	None

Raises:

Type	Description
PublishError	If the broker rejected or failed delivery of any of the block's messages (e.g. an unwritable partition).
TimeoutError	If a timeout was given and delivery was not confirmed within it. The producer keeps retrying unconfirmed messages in the background (until their message timeout), so the block may still be delivered after this is raised.

Notes

The block's timestamp is recorded for the monotonicity check only once every message's delivery is confirmed, so a block whose publish raised may be retried as-is. A retry after TimeoutError can duplicate messages that were still in flight; producer idempotence does not deduplicate application-level re-produces.

Source code in arrakis/publish.py

def publish(
    self,
    block: SeriesBlock,
    timeout: timedelta | None = None,
) -> None:
    """Publish timeseries data

    Parameters
    ----------
    block : SeriesBlock
        A data block with all channels to publish.
    timeout : timedelta, optional
        The maximum time to wait for delivery confirmation before
        timing out.  By default there is no bound of our own:
        publishing blocks until every message definitively
        resolves -- delivered, or failed by the producer once its
        per-message timeout (``message.timeout.ms``, which caps
        its internal retries) expires.  Transient broker
        unavailability thus stalls publishing rather than failing
        it, and only definitive failures raise.

    Raises
    ------
    PublishError
        If the broker rejected or failed delivery of any of the
        block's messages (e.g. an unwritable partition).
    TimeoutError
        If a *timeout* was given and delivery was not confirmed
        within it.  The producer keeps retrying unconfirmed
        messages in the background (until their message timeout),
        so the block may still be delivered after this is raised.

    Notes
    -----
    The block's timestamp is recorded for the monotonicity check
    only once every message's delivery is confirmed, so a block
    whose publish raised may be retried as-is.  A retry after
    ``TimeoutError`` can duplicate messages that were still in
    flight; producer idempotence does not deduplicate
    application-level re-produces.

    """
    if not hasattr(self, "_producer") or not self._producer:
        msg = (
            "publication interface not initialized, "
            "please use context manager when publishing."
        )
        raise RuntimeError(msg)

    # check for attempt to publish invalid channels
    # FIXME: warning for missing channels
    for name, channel in block.channels.items():
        if channel != self.channels[name]:
            msg = (
                f"channel metadata mismatch for '{name}': "
                f"got {channel!r}, expected {self.channels[name]!r}"
            )
            raise ValueError(msg)

    # check for block duration not matching expected stride
    if block.duration_ns != self.stride:
        msg = (
            "block to publish does not match expected publisher stride "
            f"got {block.duration_ns:_} ns, expected {self.stride:_} ns."
        )
        raise ValueError(msg)

    # check for non-monotonic or duplicate timestamps
    if (
        self._last_published_ns is not None
        and block.time_ns <= self._last_published_ns
    ):
        msg = (
            "block timestamp is not monotonically increasing: got "
            f"{block.time_ns:_} ns, last published {self._last_published_ns:_} ns."
        )
        raise ValueError(msg)

    # publish data for each partition, collecting delivery
    # reports: a delivery failure is only ever surfaced through
    # the report callback -- without one, a failed block would be
    # dropped silently once the producer's message timeout expired
    errors: list[KafkaError] = []

    def on_delivery(error: KafkaError | None, _message: object) -> None:
        if error is not None:
            errors.append(error)

    for partition_id, batch in block.to_row_batches(self.channels):
        topic = topic_name(partition_id, self.replay_id)
        logger.debug("publishing to topic %s: %s", topic, batch)
        self._producer.produce(
            topic=topic,
            value=serialize_batch(batch),
            on_delivery=on_delivery,
        )
        # serve delivery reports of earlier batches while producing
        self._producer.poll(0)

    if timeout is not None:
        pending = self._producer.flush(timeout.total_seconds())
    else:
        pending = self._producer.flush()
    if errors:
        msg = (
            f"delivery failed for {len(errors)} of the block's "
            f"messages: {errors[0]}"
        )
        raise PublishError(msg)
    if pending:
        assert timeout is not None
        msg = (
            f"delivery of {pending} message(s) not confirmed within "
            f"{timeout.total_seconds()} s (the producer will keep "
            "retrying them in the background)"
        )
        raise TimeoutError(msg)
    self._last_published_ns = block.time_ns

register

register(channels=None)

register channels for publication

For most publishers, channels are not specified when registering and this method will query the server for the allowable channels for this publisher and register them internally.

For publishers allowed to register their own channels ("dynamic" publishers), all channels they expect to publish should be provided as argument. Any channel that is new, or whose core metadata (sample rate or data type) has changed relative to the server's current registration for this publisher, will be partitioned dynamically via the server, which will respond with the required Kafka partition information. If the server does not permit dynamic partitioning for this publisher, a :class:pyarrow.flight.FlightUnauthorizedError is raised with a server-provided explanatory message.

Source code in arrakis/publish.py

def register(self, channels: Iterable[Channel] | None = None):
    """register channels for publication

    For most publishers, channels are not specified when
    registering and this method will query the server for the
    allowable channels for this publisher and register them
    internally.

    For publishers allowed to register their own channels
    ("dynamic" publishers), all channels they expect to publish
    should be provided as argument.  Any channel that is new, or
    whose core metadata (sample rate or data type) has changed
    relative to the server's current registration for this
    publisher, will be partitioned dynamically via the server,
    which will respond with the required Kafka partition
    information.  If the server does not permit dynamic
    partitioning for this publisher, a
    :class:`pyarrow.flight.FlightUnauthorizedError` is raised
    with a server-provided explanatory message.

    """
    requested = list(channels) if channels else []

    self._load_registered_channels()

    if requested:
        to_partition = [
            channel
            for channel in requested
            if channel.name not in self.channels
            or channel != self.channels[channel.name]
        ]
        if to_partition:
            # exhaust the generator so the server-side exchange
            # completes; any FlightUnauthorizedError (e.g.
            # PublisherNotDynamicError / PublisherUnknownError)
            # will propagate with the server's message.
            list(self._partition_channels(to_partition))
            # re-query to pick up the newly-registered channels
            # along with any server-supplied fields (stride,
            # max_latency) that are not returned by the partition
            # exchange itself.
            self._load_registered_channels()

        for channel in requested:
            if channel.name not in self.channels:
                msg = f"channel {channel.name} was not properly registered."
                raise ValueError(msg)

    if not self.channels:
        msg = f"no channels registered for publisher ID '{self.publisher_id}'"
        if self.replay_id:
            msg += f" under replay '{self.replay_id}'"
        raise ValueError(msg)

    for channel in self.channels.values():
        assert channel.partition_id is not None, (
            f"Channel {channel} is missing partition_id."
        )
        assert channel.partition_index is not None, (
            f"Channel {channel} is missing partition_index."
        )

    return self

serialize_batch

serialize_batch(batch)

Serialize a record batch to bytes.

Parameters:

Name	Type	Description	Default
batch	RecordBatch	The batch to serialize.	required

Returns:

Type	Description
bytes	The serialized buffer.

Source code in arrakis/publish.py

def serialize_batch(batch: pyarrow.RecordBatch):
    """Serialize a record batch to bytes.

    Parameters
    ----------
    batch : pyarrow.RecordBatch
        The batch to serialize.

    Returns
    -------
    bytes
        The serialized buffer.

    """
    sink = pyarrow.BufferOutputStream()
    with pyarrow.ipc.new_stream(sink, batch.schema) as writer:
        writer.write_batch(batch)
    return sink.getvalue()