adc-streaming 0.0.0__py2-none-any.whl

adc/__init__.py

@@ -0,0 +1,11 @@
+try:
+    from importlib.metadata import PackageNotFoundError, version
+except ImportError:
+    # NOTE: remove after dropping support for Python < 3.8
+    from importlib_metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("adc-streaming")
+except PackageNotFoundError:
+    # package is not installed
+    pass

adc/auth.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+
+from enum import Enum
+
+import certifi
+
+
+class SASLMethod(Enum):
+    """SASL method to use for authentication.
+    """
+
+    PLAIN = 1
+    SCRAM_SHA_256 = 2
+    SCRAM_SHA_512 = 3
+    OAUTHBEARER = 4
+
+    def __str__(self):
+        return self.name.replace("_", "-")
+
+
+class SASLAuth(object):
+    """Attach SASL-based authentication to a client.
+
+    Returns client-based auth options when called.
+
+    Parameters
+    ----------
+    user : `str`
+        Username to authenticate with.
+    password : `str`
+        Password to authenticate with.
+    ssl : `bool`, optional
+        Whether to enable SSL (enabled by default).
+    method : `SASLMethod`, optional
+        The SASL method to authenticate. The default is SASLMethod.OAUTHBEARER
+        if token_endpoint is provided, or SASLMethod.PLAIN otherwise.
+        See valid SASL methods in SASLMethod.
+    ssl_ca_location : `str`, optional
+        If using SSL via a self-signed cert, a path/location
+        to the certificate.
+    ssl_endpoint_identification_algorithm : `str`, optional
+        If using SSL, the algorithm used to verify that certificate is valid for the endpoint.
+    token_endpoint : `str`, optional
+        The OpenID Connect token endpoint URL.
+        Required for OAUTHBEARER / OpenID Connect, otherwise ignored.
+
+    """
+
+    def __init__(self, user, password, ssl=True, method=None, token_endpoint=None, **kwargs):
+        if method is None:
+            if token_endpoint is None:
+                method = SASLMethod.PLAIN
+            else:
+                method = SASLMethod.OAUTHBEARER
+        self._method = method
+
+        # set up SSL options
+        if ssl:
+            if "ssl_ca_location" in kwargs:
+                ssl_cert = kwargs["ssl_ca_location"]
+            else:
+                ssl_cert = certifi.where()
+
+            self._config = {
+                "security.protocol": "SASL_SSL",
+                "ssl.ca.location": ssl_cert,
+                "https.ca.location": ssl_cert,
+            }
+            if "ssl_endpoint_identification_algorithm" in kwargs:
+                self._config["ssl.endpoint.identification.algorithm"] = \
+                    kwargs["ssl_endpoint_identification_algorithm"]
+        else:
+            self._config = {"security.protocol": "SASL_PLAINTEXT"}
+
+        # set up SASL options
+        self._config["sasl.mechanism"] = str(self._method)
+        if token_endpoint:
+            self._config["sasl.oauthbearer.client.id"] = user
+            self._config["sasl.oauthbearer.client.secret"] = password
+            self._config["sasl.oauthbearer.method"] = "oidc"
+            self._config["sasl.oauthbearer.token.endpoint.url"] = token_endpoint
+        else:
+            self._config["sasl.username"] = user
+            self._config["sasl.password"] = password
+
+    def __call__(self):
+        return self._config

adc/consumer.py

@@ -0,0 +1,362 @@
+import dataclasses
+import enum
+import logging
+import threading
+from collections import defaultdict
+from datetime import datetime, timedelta
+# Imports from typing are deprecated as of Python 3.9 but required for
+# compatibility with earlier versions
+from typing import (Collection, Dict, Iterable, Iterator, List, Optional, Set,
+                    Union)
+
+import confluent_kafka  # type: ignore
+import confluent_kafka.admin  # type: ignore
+
+from .auth import SASLAuth
+from .errors import ErrorCallback, log_client_errors
+
+
+class LogicalOffset(enum.IntEnum):
+    BEGINNING = confluent_kafka.OFFSET_BEGINNING
+    EARLIEST = confluent_kafka.OFFSET_BEGINNING
+
+    END = confluent_kafka.OFFSET_END
+    LATEST = confluent_kafka.OFFSET_END
+
+    STORED = confluent_kafka.OFFSET_STORED
+
+    INVALID = confluent_kafka.OFFSET_INVALID
+
+
+class Consumer:
+    conf: 'ConsumerConfig'
+    _consumer: confluent_kafka.Consumer
+    logger: logging.Logger
+
+    def __init__(self, conf: 'ConsumerConfig') -> None:
+        self.logger = logging.getLogger("adc-streaming.consumer")
+        self.conf = conf
+        self._consumer = confluent_kafka.Consumer(conf._to_confluent_kafka())
+        # Workaround for
+        # https://github.com/confluentinc/librdkafka/issues/3753#issuecomment-1058272987.
+        # FIXME: Remove once fixed upstream, or on removal of oauth_cb.
+        self._consumer.poll(0)
+        self._stop_event = threading.Event()
+
+    def subscribe(self,
+                  topics: Union[str, Iterable],
+                  timeout: timedelta = timedelta(seconds=10)):
+        """Subscribes to topics for consuming. This method doesn't use Kafka's
+        Consumer Groups; it assigns all partitions manually to this
+        process.
+
+        The topics must already exist for the subscription to succeed.
+        """
+        if isinstance(topics, str):
+            topics = [topics]
+
+        assignment = []
+        for topic in topics:
+            self.logger.debug(f"subscribing to topic {topic}")
+
+            try:
+                topic_meta = self.describe_topic(topic, timeout)
+            except KeyError:
+                raise ValueError(f"topic {topic} does not exist on the broker, so can't subscribe")
+
+            for partition_id in topic_meta.partitions.keys():
+                self.logger.debug(f"adding subscription to topic partition={partition_id}")
+                tp = confluent_kafka.TopicPartition(
+                    topic=topic,
+                    partition=partition_id,
+                )
+                assignment.append(tp)
+
+        self.logger.debug("registering topic assignment")
+        self._consumer.assign(assignment)
+
+    def describe_topic(
+            self,
+            topic: str,
+            timeout: timedelta = timedelta(seconds=5.0)) -> confluent_kafka.admin.TopicMetadata:
+        """Fetch confluent_kafka.admin.TopicMetadata describing a topic.
+        """
+        self.logger.debug(f"fetching cluster metadata to describe topic name={topic}")
+        cluster_meta = self._consumer.list_topics(timeout=timeout.total_seconds())
+        self.logger.debug(f"cluster metadata: {cluster_meta.topics}")
+        return cluster_meta.topics[topic]
+
+    def mark_done(self, msg: confluent_kafka.Message, asynchronous: bool = True):
+        """
+        Mark a message as fully-processed. In the background, the client will
+        continuously synchronize this information with Kafka so that the stream can be
+        resumed from this point in the future.
+
+        If asynchronous is set to False, however, the information will be sent
+        to Kafka immediately. This option allows fine-grained reliablity, but
+        can seriously reduce throughput if used on every message.
+        """
+        if asynchronous:
+            self._consumer.store_offsets(msg)
+        else:
+            self._consumer.commit(msg, asynchronous=False)
+
+    def _offsets_for_position(self, partitions: Collection[confluent_kafka.TopicPartition],
+                              position: Union[datetime, LogicalOffset]) \
+            -> List[confluent_kafka.TopicPartition]:
+        if isinstance(position, datetime):
+            offset = int(position.timestamp() * 1000)
+        elif isinstance(position, LogicalOffset):
+            offset = position
+        else:
+            raise TypeError("Only datetime objects and logical offsets supported")
+
+        _partitions = [
+            confluent_kafka.TopicPartition(topic=tp.topic, partition=tp.partition, offset=offset)
+            for tp in partitions
+        ]
+
+        if isinstance(position, datetime):
+            self.logger.debug("looking up offsets for time")
+            return self._consumer.offsets_for_times(_partitions)
+        else:
+            return _partitions
+
+    def stop(self):
+        """Stops the runloop of the consumer. Useful when running the
+        consumer in a different thread.
+        """
+        self._stop_event.set()
+
+    def stream(self,
+               autocommit: bool = True,
+               batch_size: int = 100,
+               batch_timeout: timedelta = timedelta(seconds=1.0),
+               start_at: Union[datetime, LogicalOffset, None] = None
+               ) -> Iterator[confluent_kafka.Message]:
+        """Returns a stream which iterates over the messages in the topics
+        to which the client is subscribed.
+
+        If autocommit is true, then messages are automatically marked as handled
+        when they are yielded. This removes the need to call
+        'mark_done' on each message. Callers using asynchronous message
+        processing or with complex processing needs should disable this.
+
+        batch_size controls the number of messages to request from Kafka per
+        batch. Higher values may be more efficient, but may add latency.
+
+        batch_timeout controls how long the client should wait for Kafka to
+        provide a full batch of batch_size messages. Higher values may be more
+        efficient, but may add latency.
+
+        start_at controls the location in every partition the client reads
+        from, overriding the configured default position or the stored offsets.
+        Either a special logical offset value (END, BEGINNING, STORE, INVALID)
+        or a datetime. Using a datetime will cause reading to start from the
+        first message *after* the specified datetime (or END if none exists).
+        Using this parameter will override any previously assigned but not
+        committed offsets if they are managed from outside adc. Passing None
+        (or leaving it unspecified) avoids reassignment.
+
+        If the consumer's configuration has read_forever set to False, then the
+        stream stops when the client has hit the last message in all partitions.
+        This set of partitions is calculated just once when iterate() is first
+        called; calling subscribe() after iterate() may cause inconsistent
+        behavior in this case.
+
+        """
+
+        if start_at is not None:
+            assignment = self._consumer.assignment()
+            self._consumer.assign(self._offsets_for_position(assignment, start_at))
+
+        if self.conf.read_forever:
+            return self._stream_forever(autocommit, batch_size, batch_timeout)
+        else:
+            return self._stream_until_eof(autocommit, batch_size, batch_timeout)
+
+    def _stream_forever(self,
+                        autocommit: bool = True,
+                        batch_size: int = 100,
+                        batch_timeout: timedelta = timedelta(seconds=1.0),
+                        ) -> Iterator[confluent_kafka.Message]:
+        self._stop_event.clear()
+        while not self._stop_event.is_set():
+            try:
+                messages = self._consumer.consume(batch_size,
+                                                  batch_timeout.total_seconds())
+                for m in messages:
+                    if self._stop_event.is_set():
+                        break
+                    err = m.error()
+                    if err is None:
+                        self.logger.debug(f"read message from partition {m.partition()}")
+                        # Automatically mark message as processed, if desired
+                        if autocommit:
+                            self.mark_done(m, asynchronous=True)
+                        yield m
+                    else:
+                        raise (confluent_kafka.KafkaException(err))
+            finally:
+                if autocommit:
+                    self._consumer.commit(asynchronous=True)
+
+    def _stream_until_eof(self,
+                          autocommit: bool = True,
+                          batch_size: int = 100,
+                          batch_timeout: timedelta = timedelta(seconds=1.0),
+                          ) -> Iterator[confluent_kafka.Message]:
+        assignment = self._consumer.assignment()
+
+        # Make a map of topic-name -> set of partition IDs we're assigned to.
+        # When we hit a partition EOF, remove that partition from the map.
+        active_partitions: defaultdict[str, Set[int]] = defaultdict(set)
+        for tp in assignment:
+            self.logger.debug(f"tracking until eof for topic={tp.topic} partition={tp.partition}")
+            active_partitions[tp.topic].add(tp.partition)
+
+        self._stop_event.clear()
+        while len(active_partitions) > 0 and not self._stop_event.is_set():
+            messages = self._consumer.consume(batch_size, batch_timeout.total_seconds())
+            try:
+                for m in messages:
+                    if self._stop_event.is_set():
+                        raise StopIteration
+                    err = m.error()
+                    # A new message may arrive from a previously removed topic/partition,
+                    # in which case it must be re-added
+                    partition_set = active_partitions[m.topic()]
+                    partition_set.add(m.partition())
+
+                    if err is None:
+                        self.logger.debug(f"read message from partition {m.partition()}")
+                        # Automatically mark message as processed, if desired
+                        if autocommit:
+                            self.mark_done(m, asynchronous=True)
+                        yield m
+                    elif err.code() == confluent_kafka.KafkaError._PARTITION_EOF:
+                        self.logger.debug(f"eof for topic={m.topic()} partition={m.partition()}")
+                        # Done with this partition, remove it
+                        partition_set.remove(m.partition())
+                        if len(partition_set) == 0:
+                            # Done with all partitions for the topic, remove it
+                            del active_partitions[m.topic()]
+                    else:
+                        raise (confluent_kafka.KafkaException(err))
+            finally:
+                if autocommit:
+                    self._consumer.commit(asynchronous=True)
+        self._stop_event.set()
+
+    def close(self):
+        """ Close the consumer, ending its subscriptions. """
+        self._consumer.close()
+
+
+# Used to be called ConsumerStartPosition, though this was confusing because
+# it only affects "auto.offset.reset" not the start position for a call to
+# consume.
+class ConsumerDefaultPosition(enum.Enum):
+    EARLIEST = 1
+    LATEST = 2
+
+    def __str__(self):
+        return self.name.lower()
+
+
+# Alias to the old name
+# TODO: Remove alias on the next breaking release
+ConsumerStartPosition = ConsumerDefaultPosition
+
+
+@dataclasses.dataclass
+class ConsumerConfig:
+    broker_urls: List[str]
+    group_id: str
+
+    # When we have reached the last message on a topic, should we hold the
+    # stream open to wait for more messages?
+    read_forever: bool = True
+
+    # When reading a topic for the first time, where should we start in the
+    # stream? Note that, if the topic has already been consumed under the
+    # provided group_id, then consumption will start after the last message that
+    # was marked done with consumer.mark_done, regardless of this setting. This
+    # is only used when the position in the stream is unknown.
+    #
+    # You can force reading at a logical offset or datetime with the "start_at"
+    # argument to Consumer.stream().
+    #
+    # This is specified via a ConsumerDefaultPosition value.
+    #
+    # TODO: rename on next breaking release
+    start_at: ConsumerDefaultPosition = ConsumerDefaultPosition.EARLIEST
+
+    # Authentication package to pass in to read from Kafka.
+    auth: Optional[SASLAuth] = None
+
+    # Callback to execute whenever an internal Kafka error occurs.
+    error_callback: Optional[ErrorCallback] = log_client_errors
+
+    # How often should we save our progress to Kafka?
+    offset_commit_interval: timedelta = timedelta(seconds=5)
+
+    # Whether ncoming message CRCs should be checked to detect corruption in
+    # transit. Enabling this option has a small CPU use/throughput cost.
+    check_crcs: bool = False
+
+    # reconnect_backoff_time is the time that the backend should initially wait
+    # before attempting to reconnect to Kafka if its connection fails.
+    # Repeated failures will cause the wait time to be increased exponentially,
+    # with a random variation, until reconnect_max_time is reached.
+    reconnect_backoff_time: timedelta = timedelta(milliseconds=100)
+
+    # reconnect_max_time is the longest time that the backend should wait
+    # between attempts to reconnect to Kafka.
+    reconnect_max_time: timedelta = timedelta(seconds=10)
+
+    def _to_confluent_kafka(self) -> Dict:
+        def as_ms(td: timedelta):
+            """Convert a timedelta object to a duration in milliseconds"""
+            return int(td.total_seconds() * 1000.0)
+
+        config = {
+            "bootstrap.servers": ",".join(self.broker_urls),
+            "check.crcs": self.check_crcs,
+            "error_cb": self.error_callback,
+            "group.id": self.group_id,
+            "enable.auto.commit": True,
+            "auto.commit.interval.ms": as_ms(self.offset_commit_interval),
+            "enable.auto.offset.store": False,
+            "queued.min.messages": 1000,
+            "enable.partition.eof": not self.read_forever,
+            "reconnect.backoff.max.ms": as_ms(self.reconnect_max_time),
+            "reconnect.backoff.ms": as_ms(self.reconnect_backoff_time),
+        }
+        if self.start_at is ConsumerDefaultPosition.EARLIEST:
+            default_topic_config = config.get("default.topic.config", {})
+            default_topic_config = {
+                "auto.offset.reset": "EARLIEST",
+            }
+            config["default.topic.config"] = default_topic_config
+        elif self.start_at is ConsumerDefaultPosition.LATEST:
+            # FIXME: librdkafka has a bug in offset handling - it caches
+            # "OFFSET_END", and will repeatedly move to the end of the
+            # topic. See https://github.com/edenhill/librdkafka/pull/2876 -
+            # it should get fixed in v1.5 of librdkafka.
+
+            librdkafka_version = confluent_kafka.libversion()[0]
+            if librdkafka_version < "1.5.0":
+                self.logger.warn(
+                    "In librdkafka before v1.5, LATEST offsets have buggy behavior; you may "
+                    f"not receive data (your librdkafka version is {librdkafka_version}). See "
+                    "https://github.com/confluentinc/confluent-kafka-dotnet/issues/1254.")
+            default_topic_config = config.get("default.topic.config", {})
+            default_topic_config = {
+                "auto.offset.reset": "LATEST",
+            }
+            config["default.topic.config"] = default_topic_config
+
+        if self.auth is not None:
+            config.update(self.auth())
+        return config

adc/errors.py

@@ -0,0 +1,102 @@
+import logging
+from typing import Callable
+
+import confluent_kafka  # type: ignore
+
+logger = logging.getLogger("adc-streaming")
+
+
+ErrorCallback = Callable[[confluent_kafka.KafkaError], None]
+
+
+def log_client_errors(kafka_error: confluent_kafka.KafkaError):
+    if kafka_error.code() == confluent_kafka.KafkaError._ALL_BROKERS_DOWN:
+        # This error occurs very frequently. It's not nearly as fatal as it
+        # sounds: it really indicates that the client's broker metadata has
+        # timed out. It appears to get triggered in races during client
+        # shutdown, too. See https://github.com/edenhill/librdkafka/issues/2543
+        # for more background.
+        logger.warn("client is currently disconnected from all brokers")
+    else:
+        logger.error(f"internal kafka error: {kafka_error}")
+
+
+DeliveryCallback = Callable[[confluent_kafka.KafkaError, confluent_kafka.Message], None]
+
+
+def log_delivery_errors(
+        kafka_error: confluent_kafka.KafkaError,
+        msg: confluent_kafka.Message) -> None:
+    if kafka_error is not None:
+        logger.error(f"delivery error: {kafka_error}")
+
+
+def raise_delivery_errors(kafka_error: confluent_kafka.KafkaError,
+                          msg: confluent_kafka.Message) -> None:
+    if kafka_error is not None:
+        raise KafkaException.from_kafka_error(kafka_error, msg)
+    elif msg.error() is not None:
+        raise KafkaException.from_kafka_error(msg.error(), msg)
+
+
+def _get_topic_related_errors():
+    """Build a set of all Kafka error codes which seem to relate to a specific topic.
+
+    This uses a list extracted from all documented error codes up to confluent_kafka v2.4,
+    but some of these errors did not exist or were not exposed in earlier versions.
+    To maintain backward compatibility, this function checks whether each error exists before
+    attempting to otherwise refer to it.
+    """
+    err_names = [
+        "_UNKNOWN_TOPIC",
+        "_NO_OFFSET",
+        "_LOG_TRUNCATION",
+        "OFFSET_OUT_OF_RANGE",
+        "UNKNOWN_TOPIC_OR_PART",
+        "NOT_LEADER_FOR_PARTITION",
+        "TOPIC_EXCEPTION",
+        "NOT_ENOUGH_REPLICAS",
+        "NOT_ENOUGH_REPLICAS_AFTER_APPEND",
+        "INVALID_COMMIT_OFFSET_SIZE",
+        "TOPIC_AUTHORIZATION_FAILED",
+        "TOPIC_ALREADY_EXISTS",
+        "INVALID_PARTITIONS",
+        "INVALID_REPLICATION_FACTOR",
+        "INVALID_REPLICA_ASSIGNMENT",
+        "REASSIGNMENT_IN_PROGRESS",
+        "TOPIC_DELETION_DISABLED",
+        "OFFSET_NOT_AVAILABLE",
+        "PREFERRED_LEADER_NOT_AVAILABLE",
+        "NO_REASSIGNMENT_IN_PROGRESS",
+        "GROUP_SUBSCRIBED_TO_TOPIC",
+        "UNSTABLE_OFFSET_COMMIT",
+        "UNKNOWN_TOPIC_ID",
+    ]
+    errors = set()
+    for name in err_names:
+        if hasattr(confluent_kafka.KafkaError, name):
+            errors.add(getattr(confluent_kafka.KafkaError, name))
+        else:
+            logger.debug(f"{name} does not exist in confluent_kafka version "
+                         f"{confluent_kafka.__version__} ({confluent_kafka.libversion()})")
+    return errors
+
+
+class KafkaException(Exception):
+    @classmethod
+    def from_kafka_error(cls, error, msg=None):
+        return cls(error, msg)
+
+    topic_related_errors = _get_topic_related_errors()
+
+    def __init__(self, error, msg=None):
+        self.error = error
+        self.name = error.name()
+        self.reason = error.str()
+        self.retriable = error.retriable()
+        self.fatal = error.fatal()
+        self.message = msg
+        ex_msg = f"Error communicating with Kafka: code={self.name} {self.reason}"
+        if msg and error.code() in KafkaException.topic_related_errors:
+            ex_msg += f" on topic {msg.topic()}"
+        super(KafkaException, self).__init__(ex_msg)

adc/io.py

@@ -0,0 +1,74 @@
+import logging
+import warnings
+from contextlib import contextmanager
+from typing import Iterable, List, Optional, Union
+
+import confluent_kafka  # type: ignore
+
+from adc import auth, consumer, kafka, producer
+
+logger = logging.getLogger("adc-streaming")
+
+
+def open(url: str,
+         mode: str = 'r',
+         auth: Optional[auth.SASLAuth] = None,
+         start_at: consumer.ConsumerStartPosition = consumer.ConsumerStartPosition.EARLIEST,  # noqa: E501
+         read_forever: bool = True,
+         ) -> Union[producer.Producer, Iterable[confluent_kafka.Message]]:
+    group_id, broker_addresses, topics = kafka.parse_kafka_url(url)
+    logger.debug("connecting to addresses=%s  group_id=%s  topics=%s",
+                 broker_addresses, group_id, topics)
+    if mode == "r":
+        if group_id is None:
+            raise ValueError("group ID must be set when in reader mode")
+        return _open_consumer(group_id, broker_addresses, topics, auth, start_at, read_forever)
+    elif mode == "w":
+        if len(topics) != 1:
+            raise ValueError("must specify exactly one topic in write mode")
+        if group_id is not None:
+            warnings.warn("group ID has no effect when opening a stream in write mode")
+        if start_at is not consumer.ConsumerStartPosition.EARLIEST:
+            warnings.warn("start_at has no effect when opening a stream in write mode")
+        if read_forever is not True:
+            warnings.warn("read_forever has no effect when opening a stream in write mode")
+        return _open_producer(broker_addresses, topics[0], auth)
+    else:
+        raise ValueError("mode must be either 'w' or 'r'")
+
+
+@contextmanager
+def _open_consumer(
+        group_id: str,
+        broker_addresses: List[str],
+        topics: List[str],
+        auth: Optional[auth.SASLAuth],
+        start_at: consumer.ConsumerStartPosition,
+        read_forever: bool,
+) -> Iterable[confluent_kafka.Message]:
+    client = consumer.Consumer(consumer.ConsumerConfig(
+        broker_urls=broker_addresses,
+        group_id=group_id,
+        auth=auth,
+        start_at=start_at,
+        read_forever=read_forever,
+    ))
+    for t in topics:
+        client.subscribe(t)
+
+    try:
+        yield client.stream()
+    finally:
+        client.close()
+
+
+def _open_producer(
+        broker_addresses: List[str],
+        topic: str,
+        auth: Optional[auth.SASLAuth],
+) -> producer.Producer:
+    return producer.Producer(producer.ProducerConfig(
+        broker_urls=broker_addresses,
+        topic=topic,
+        auth=auth,
+    ))

adc/kafka.py

@@ -0,0 +1,30 @@
+from urllib.parse import urlparse
+
+
+def parse_kafka_url(val):
+    """Extracts the group ID, broker addresses, and topic names from a Kafka URL.
+
+    The URL should be in this form:
+        ``kafka://[groupid@]broker[,broker2[,...]]/topic[,topic2[,...]]``
+
+    The returned group ID and topic may be None if they aren't in the URL.
+
+    """
+    parsed = urlparse(val)
+    if parsed.scheme != "kafka":
+        raise ValueError("invalid kafka URL: must start with 'kafka://'")
+
+    split_netloc = parsed.netloc.split("@", maxsplit=1)
+    if len(split_netloc) == 2:
+        group_id = split_netloc[0]
+        broker_addresses = split_netloc[1].split(",")
+    else:
+        group_id = None
+        broker_addresses = split_netloc[0].split(",")
+
+    topics = parsed.path.lstrip("/")
+    if len(topics) == 0:
+        split_topics = None
+    else:
+        split_topics = topics.split(",")
+    return group_id, broker_addresses, split_topics

adc/oidc.py

@@ -0,0 +1,27 @@
+def set_oauth_cb(config):
+    """Implement client support for KIP-768 OpenID Connect.
+
+    Apache Kafka 3.1.0 supports authentication using OpenID Client Credentials.
+    Native support for Python is still incomplete due to this issue:
+    https://github.com/confluentinc/librdkafka/issues/3751
+
+    Meanwhile, this is a pure Python implementation of the refresh token
+    callback.
+    """
+    if config.pop('sasl.oauthbearer.method', None) != 'oidc':
+        return
+
+    client_id = config.pop('sasl.oauthbearer.client.id')
+    client_secret = config.pop('sasl.oauthbearer.client.secret')
+    scope = config.pop('sasl.oauthbearer.scope', None)
+    token_endpoint = config.pop('sasl.oauthbearer.token.endpoint.url')
+
+    from authlib.integrations.requests_client import OAuth2Session
+    session = OAuth2Session(client_id, client_secret, scope=scope)
+
+    def oauth_cb(*_, **__):
+        token = session.fetch_token(
+            token_endpoint, grant_type='client_credentials')
+        return token['access_token'], token['expires_at']
+
+    config['oauth_cb'] = oauth_cb

adc/producer.py

@@ -0,0 +1,164 @@
+import abc
+import dataclasses
+import logging
+from datetime import timedelta
+from typing import Dict, List, Optional, Union
+
+try:  # this will work only in python >= 3.8
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+import confluent_kafka  # type: ignore
+
+from .auth import SASLAuth
+from .errors import (DeliveryCallback, ErrorCallback, log_client_errors,
+                     log_delivery_errors)
+
+
+class Producer:
+    conf: 'ProducerConfig'
+    _producer: confluent_kafka.Producer
+    logger: logging.Logger
+
+    def __init__(self, conf: 'ProducerConfig') -> None:
+        self.logger = logging.getLogger("adc-streaming.producer")
+        self.conf = conf
+        self.logger.debug(f"connecting to producer with config {conf._to_confluent_kafka()}")
+        self._producer = confluent_kafka.Producer(conf._to_confluent_kafka())
+        # Workaround for
+        # https://github.com/confluentinc/librdkafka/issues/3753#issuecomment-1058272987.
+        # FIXME: Remove once fixed upstream, or on removal of oauth_cb.
+        self._producer.poll(0)
+
+    def write(self,
+              msg: Union[bytes, 'Serializable'],
+              headers: Optional[Union[dict, list]] = None,
+              delivery_callback: Optional[DeliveryCallback] = log_delivery_errors,
+              topic: Optional[str] = None,
+              key: Optional[Union[str, bytes]] = None) -> None:
+        if isinstance(msg, Serializable):
+            msg = msg.serialize()
+        if topic is None:
+            if self.conf.topic is not None:
+                topic = self.conf.topic
+            else:
+                raise Exception("No topic specified for write: "
+                                "Either configure a topic when constructing the Producer, "
+                                "or specify the topic argument to write()")
+        self.logger.debug("writing message to %s", topic)
+        produce_kwargs = {"headers": headers, "key": key}
+        if delivery_callback is not None:
+            produce_kwargs["on_delivery"] = delivery_callback
+        while True:
+            try:
+                self._producer.produce(topic, msg, **produce_kwargs)
+                break
+            except BufferError:
+                # It's hard to know what the size limit on the buffer is, so we will try to
+                # wait until the number of items in it decreases (or it is empty, in case all
+                # messages were successfully sent in the time it took to handle the error).
+                buffer_len = len(self._producer)
+                self.logger.debug(f"Blocking due to BufferError, buffer size: {buffer_len}")
+                while buffer_len > 0 and len(self._producer) == buffer_len:
+                    self._producer.poll(0.01)
+
+    def queued_message_count(self):
+        """Get the number of messages waiting to be sent to the broker.
+        """
+        return len(self._producer)
+
+    def flush(self, timeout: timedelta = timedelta(seconds=10)) -> int:
+        """Attempt to flush enqueued messages. Return the number of messages still
+        enqueued after the attempt.
+
+        """
+        n = self._producer.flush(timeout.total_seconds())
+        if n > 0:
+            self.logger.debug("flushed messages, %d still enqueued", n)
+        else:
+            self.logger.debug("flushed all messages")
+        return n
+
+    def close(self, timeout: timedelta = timedelta(seconds=10)) -> int:
+        self.logger.debug("shutting down producer")
+        return self.flush(timeout)
+
+    def __enter__(self) -> 'Producer':
+        return self
+
+    def __exit__(self, type, value, traceback) -> bool:
+        if type == KeyboardInterrupt:
+            print("Aborted (CTRL-C).")
+            return True
+        if type is None and value is None and traceback is None:
+            n_unsent = self.close()
+            if n_unsent > 0:
+                raise Exception(f"{n_unsent} messages remain unsent, some data may have been lost!")
+            return False
+        return False
+
+
+@dataclasses.dataclass
+class ProducerConfig:
+    broker_urls: List[str]
+    topic: Optional[str]
+    auth: Optional[SASLAuth] = None
+    error_callback: Optional[ErrorCallback] = log_client_errors
+
+    # produce_timeout sets the maximum amount of time that the backend can take
+    # to send a message to Kafka. Use a value of 0 to never timeout.
+    produce_timeout: timedelta = timedelta(seconds=10)
+
+    # produce_backoff_time sets the time the backend will wait before retrying
+    # to send a message to Kafka. May not be less than one millisecond.
+    produce_backoff_time: timedelta = timedelta(milliseconds=100)
+
+    # use_idempotence instructs the backend whether to ensure that messages are
+    # recorded by the broker exactly once and in the order of production.
+    use_idempotence: bool = False
+
+    # reconnect_backoff_time is the time that the backend should initially wait
+    # before attempting to reconnect to Kafka if its connection fails.
+    # Repeated failures will cause the wait time to be increased exponentially,
+    # with a random variation, until reconnect_max_time is reached.
+    reconnect_backoff_time: timedelta = timedelta(milliseconds=100)
+
+    # reconnect_max_time is the longest time that the backend should wait
+    # between attempts to reconnect to Kafka.
+    reconnect_max_time: timedelta = timedelta(seconds=10)
+
+    compression_type: Optional[Union[Literal['gzip'], Literal['snappy'],
+                                     Literal['lz4'], Literal['zstd']]] = None
+
+    # maximum message size, before compression
+    message_max_bytes: Optional[int] = None
+
+    def _to_confluent_kafka(self) -> Dict:
+        def as_ms(td: timedelta):
+            """Convert a timedelta object to a duration in milliseconds"""
+            return int(td.total_seconds() * 1000.0)
+
+        if self.produce_backoff_time < timedelta(milliseconds=1):
+            raise ValueError("produce_backoff_time may not be less than one millisecond")
+        config = {
+            "bootstrap.servers": ",".join(self.broker_urls),
+            "enable.idempotence": self.use_idempotence,
+            "message.timeout.ms": as_ms(self.produce_timeout),
+            "reconnect.backoff.max.ms": as_ms(self.reconnect_max_time),
+            "reconnect.backoff.ms": as_ms(self.reconnect_backoff_time),
+            "retry.backoff.ms": as_ms(self.produce_backoff_time),
+            "compression.type": self.compression_type or 'none',
+        }
+        if self.message_max_bytes is not None:
+            config['message.max.bytes'] = self.message_max_bytes
+        if self.error_callback is not None:
+            config["error_cb"] = self.error_callback
+        if self.auth is not None:
+            config.update(self.auth())
+        return config
+
+
+class Serializable(abc.ABC):
+    def serialize(self) -> bytes:
+        raise NotImplementedError()

adc_streaming-0.0.0.dist-info/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2019, Mario Juric
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

adc_streaming-0.0.0.dist-info/METADATA

@@ -0,0 +1,85 @@
+Metadata-Version: 2.1
+Name: adc-streaming
+Version: 0.0.0
+Summary: Astronomy Data Commons streaming client libraries
+Home-page: https://github.com/astronomy-commons/adc-streaming
+Author: Astronomy Data Commons Team
+Author-email: swnelson@uw.edu
+License: BSD
+Platform: UNKNOWN
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Development Status :: 3 - Alpha
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS :: MacOS X
+Description-Content-Type: text/markdown
+Requires-Dist: confluent-kafka (>=2.11.0)
+Requires-Dist: tqdm
+Requires-Dist: certifi (>=2020.04.05.1)
+Requires-Dist: dataclasses ; python_version < "3.7"
+Requires-Dist: importlib-metadata ; python_version < "3.8"
+Requires-Dist: typing-extensions ; python_version < "3.8"
+Provides-Extra: dev
+Requires-Dist: autopep8 ; extra == 'dev'
+Requires-Dist: docker ; extra == 'dev'
+Requires-Dist: flake8 ; extra == 'dev'
+Requires-Dist: isort ; extra == 'dev'
+Requires-Dist: pytest ; extra == 'dev'
+Requires-Dist: pytest-timeout ; extra == 'dev'
+Requires-Dist: pytest-integration ; extra == 'dev'
+Requires-Dist: sphinx ; extra == 'dev'
+Requires-Dist: sphinx-rtd-theme ; extra == 'dev'
+Requires-Dist: twine ; extra == 'dev'
+
+# Astronomy Data Commons Streaming Client Libraries
+
+Libraries making it easy to access astronomy data commons resources.
+
+## Developer notes
+
+### Setup
+
+To prepare for development, run `pip install --editable ".[dev]"` from within
+the repo directory. This will install all dependencies, including those using
+during development workflows.
+
+This project expects you to use a `pip`-centric workflow for development on the
+project itself. If you're using conda, then use the conda environment's `pip` to
+install development dependencies, as described above.
+
+Integration tests require Docker to run a Kafka broker. The broker might have
+network problems on OSX if you use Docker Desktop; run the tests in a Linux
+virtual machine (like with VirtualBox) to get around this.
+
+### Code Workflow
+
+Write code, making changes.
+
+Use `make format` to reformat your code to comply with PEP8.
+
+Use `make lint` to catch common mistakes.
+
+Use `make test-quick` to run fast unit tests.
+
+Use `make test` to run the full slow test suite, including integration tests.
+
+Once satisfied with all four of those, push your changes and open a PR.
+
+### Tag, build, and upload to PyPI and Conda
+
+Tag a new version:
+```
+git tag -s -a v0.x.x
+```
+
+Build and release:
+
+```
+make pypi-dist
+make pypi-dist-check
+make pypi-upload
+make conda-build
+make conda-upload
+```
+
+

adc_streaming-0.0.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+adc/__init__.py,sha256=IGK6EiDe-qcBjT56a39nxC8SAaGbnKNUvoMbyvfrsbQ,332
+adc/auth.py,sha256=k1WX81lh56V7rFJU-BPtZAdnvftU_qVjqM-8lPnuVTs,2857
+adc/consumer.py,sha256=sfwH3KDmRDeUdryEzX2XCzKBvt3w11vrl3UqRstahog,15723
+adc/errors.py,sha256=3ummsAbA1fa23jLif0kdoYt-1GbtDo8gUqnVY5opfig,3731
+adc/io.py,sha256=RdfQuhmPLL9wVvUWaggRbeGKUVq5gMFp-QMO4Du-z9o,2541
+adc/kafka.py,sha256=jsFypqYbfJk_CMpznjgUK4lSL7QRuXzYKGObx5OQWtg,933
+adc/oidc.py,sha256=WRb3Z8-psffSwfYtOEV6HsmtEfEMmmm90JaFwOUdPkA,1072
+adc/producer.py,sha256=KwNlZBccFhHfTUsgEZhJ4HG8-ExelmwF0K_HzyUCwGo,6875
+adc_streaming-0.0.0.dist-info/LICENSE,sha256=BzSloDstU84TGMF0A6S-daRS2sVHoDgo3e7yh666Hgg,1519
+adc_streaming-0.0.0.dist-info/METADATA,sha256=gbYADPcYuwxHRjOCTYVYJyxoLksDPc0gMX12Embodc4,2602
+adc_streaming-0.0.0.dist-info/WHEEL,sha256=pqI-DBMA-Z6OTNov1nVxs7mwm6Yj2kHZGNp_6krVn1E,92
+adc_streaming-0.0.0.dist-info/top_level.txt,sha256=8quD7Fop9Wwsv5j59wJl1NsuB_dGLESpqC4Tk7V2Qew,4
+adc_streaming-0.0.0.dist-info/RECORD,,

adc_streaming-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.33.1)
+Root-Is-Purelib: true
+Tag: py2-none-any
+

adc_streaming-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+adc