adc-streaming 2.3.2__tar.gz → 2.5.0__tar.gz

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/.github/workflows/build.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - name: Check out the code

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: adc-streaming
-Version: 2.3.2
+Version: 2.5.0
 Summary: Astronomy Data Commons streaming client libraries
 Home-page: https://github.com/astronomy-commons/adc-streaming
 Author: Astronomy Data Commons Team

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/adc/__init__.py

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

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/adc/auth.py

@@ -38,6 +38,8 @@ class SASLAuth(object):
     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.
@@ -63,6 +65,9 @@ class SASLAuth(object):
                 "security.protocol": "SASL_SSL",
                 "ssl.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"}
 

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/adc/consumer.py

@@ -1,10 +1,13 @@
 import dataclasses
 import enum
 import logging
-from datetime import timedelta
 import threading
-from typing import Dict, Iterable, Iterator, List, Optional, Set, Union
 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
@@ -14,6 +17,18 @@ from .errors import ErrorCallback, log_client_errors
 from .oidc import set_oauth_cb
 
 
+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
@@ -23,7 +38,8 @@ class Consumer:
         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.
+        # 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()
@@ -86,6 +102,27 @@ class Consumer:
         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.
@@ -95,7 +132,8 @@ class Consumer:
     def stream(self,
                autocommit: bool = True,
                batch_size: int = 100,
-               batch_timeout: timedelta = timedelta(seconds=1.0)
+               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.
@@ -112,6 +150,15 @@ class Consumer:
         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
@@ -119,6 +166,11 @@ class Consumer:
         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:
@@ -145,7 +197,7 @@ class Consumer:
                             self.mark_done(m, asynchronous=True)
                         yield m
                     else:
-                        raise(confluent_kafka.KafkaException(err))
+                        raise (confluent_kafka.KafkaException(err))
             finally:
                 if autocommit:
                     self._consumer.commit(asynchronous=True)
@@ -191,7 +243,7 @@ class Consumer:
                             # Done with all partitions for the topic, remove it
                             del active_partitions[m.topic()]
                     else:
-                        raise(confluent_kafka.KafkaException(err))
+                        raise (confluent_kafka.KafkaException(err))
             finally:
                 if autocommit:
                     self._consumer.commit(asynchronous=True)
@@ -202,7 +254,10 @@ class Consumer:
         self._consumer.close()
 
 
-class ConsumerStartPosition(enum.Enum):
+# 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
 
@@ -210,6 +265,11 @@ class ConsumerStartPosition(enum.Enum):
         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]
@@ -225,8 +285,13 @@ class ConsumerConfig:
     # was marked done with consumer.mark_done, regardless of this setting. This
     # is only used when the position in the stream is unknown.
     #
-    # This is specified as a logical offset via a ConsumerStartPosition value.
-    start_at: ConsumerStartPosition = ConsumerStartPosition.EARLIEST
+    # 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
@@ -269,13 +334,13 @@ class ConsumerConfig:
             "reconnect.backoff.max.ms": as_ms(self.reconnect_max_time),
             "reconnect.backoff.ms": as_ms(self.reconnect_backoff_time),
         }
-        if self.start_at is ConsumerStartPosition.EARLIEST:
+        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 ConsumerStartPosition.LATEST:
+        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 -

adc-streaming-2.5.0/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-streaming-2.3.2 → adc-streaming-2.5.0}/adc/oidc.py

@@ -18,7 +18,7 @@ def set_oauth_cb(config):
 
     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')

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/adc/producer.py

@@ -1,9 +1,9 @@
 import abc
-from ast import comprehension
 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:
@@ -27,22 +27,32 @@ class 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.
+        # 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) -> 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()
-        self.logger.debug("writing message to %s", self.conf.topic)
+        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)
         if delivery_callback is not None:
-            self._producer.produce(self.conf.topic, msg, headers=headers,
+            self._producer.produce(topic, msg, headers=headers, key=key,
                                    on_delivery=delivery_callback)
         else:
-            self._producer.produce(self.conf.topic, msg, headers=headers)
+            self._producer.produce(topic, msg, headers=headers, key=key,)
 
     def flush(self, timeout: timedelta = timedelta(seconds=10)) -> int:
         """Attempt to flush enqueued messages. Return the number of messages still
@@ -78,7 +88,7 @@ class Producer:
 @dataclasses.dataclass
 class ProducerConfig:
     broker_urls: List[str]
-    topic: str
+    topic: Optional[str]
     auth: Optional[SASLAuth] = None
     error_callback: Optional[ErrorCallback] = log_client_errors
 
@@ -104,7 +114,8 @@ class ProducerConfig:
     # 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
+    compression_type: Optional[Union[Literal['gzip'], Literal['snappy'],
+                                     Literal['lz4'], Literal['zstd']]] = None
 
     # maximum message size, before compression
     message_max_bytes: Optional[int] = None

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/adc_streaming.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: adc-streaming
-Version: 2.3.2
+Version: 2.5.0
 Summary: Astronomy Data Commons streaming client libraries
 Home-page: https://github.com/astronomy-commons/adc-streaming
 Author: Astronomy Data Commons Team

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/setup.cfg

@@ -13,6 +13,7 @@ exclude = setup.py,
 [tool:pytest]
 log_cli = True
 log_cli_level = INFO
+testpaths = tests
 
 [egg_info]
 tag_build = 

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/tests/test_auth.py

@@ -1,6 +1,6 @@
 import pytest
 
-from adc.auth import SASLAuth, SASLMethod
+from adc.auth import SASLAuth
 
 
 @pytest.mark.parametrize('auth,expected_config', [

{adc-streaming-2.3.2 → adc-streaming-2.5.0}/tests/test_kafka_integration.py

@@ -2,8 +2,8 @@ import logging
 import tempfile
 import time
 import unittest
-from datetime import timedelta
-from typing import List
+from datetime import datetime, timedelta
+from typing import List, Optional
 
 import docker
 import pytest
@@ -59,10 +59,34 @@ class KafkaIntegrationTestCase(unittest.TestCase):
         self.assertEqual(msg.topic(), topic)
         self.assertEqual(msg.value(), b"can you hear me?")
 
-    @unittest.skip("skipping due to bug in librdkafka")
-    def test_consume_from_end(self):
+    def test_message_with_key(self):
+        """Try writing a message into the Kafka broker, and try pulling the same
+        message back out.
+
+        """
+        topic = "test_message_with_key"
+        # Push one message in...
+        simple_write_msg(self.kafka, topic, "can you hear me?", key="test_msg")
+        # ... and pull it back out.
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+        ))
+        consumer.subscribe(topic)
+        stream = consumer.stream()
+
+        msg = next(stream)
+        if msg.error() is not None:
+            raise Exception(msg.error())
+
+        self.assertEqual(msg.topic(), topic)
+        self.assertEqual(msg.value(), b"can you hear me?")
+        self.assertEqual(msg.key(), b"test_msg")
+
+    def test_reset_to_end(self):
         # Write a few messages.
-        topic = "test_consume_from_end"
+        topic = "test_reset_to_end"
         simple_write_msgs(self.kafka, topic, [
             "message 1",
             "message 2",
@@ -79,15 +103,16 @@ class KafkaIntegrationTestCase(unittest.TestCase):
         stream = consumer.stream()
 
         # Now add messages after the "end"
+        time.sleep(0.5)
         simple_write_msg(self.kafka, topic, "message 4")
-
+        time.sleep(0.5)
         msg = next(stream)
         self.assertEqual(msg.topic(), topic)
         self.assertEqual(msg.value(), b"message 4")
 
-    def test_consume_from_beginning(self):
+    def test_reset_to_beginning(self):
         # Write a few messages.
-        topic = "test_consume_from_beginning"
+        topic = "test_reset_to_beginning"
         batch = [
             "message 1",
             "message 2",
@@ -164,24 +189,153 @@ class KafkaIntegrationTestCase(unittest.TestCase):
             self.assertEqual(actual.value().decode(), expected)
 
         # Start second consumer, also reading from earliest offset.
-        consumer_2 = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+        consumer_2a = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
             broker_urls=[self.kafka.address],
             group_id="test_consumer_2",
             auth=self.kafka.auth,
             read_forever=False,
             start_at=adc.consumer.ConsumerStartPosition.EARLIEST,
         ))
-        consumer_2.subscribe(topic)
-        stream_2 = consumer_2.stream()
-        msgs_2 = [msg for msg in stream_2]
-
-        # Now check that messages from both batches are processed.
-        assert consumer_2._stop_event.is_set()
-        self.assertEqual(len(batch_1 + batch_2), len(msgs_2))
-        for expected, actual in zip(batch_1 + batch_2, msgs_2):
+        consumer_2a.subscribe(topic)
+
+        # read the topic using consumer_2a
+        stream_2a = consumer_2a.stream()
+        msgs_2a = [pair for pair in zip(batch_1, stream_2a)]
+        # end iteration early after batch 1
+        stream_2a.close()
+
+        # check that messages from only batch 1 were consumed
+        self.assertEqual(len(batch_1), len(msgs_2a))
+        for expected, actual in msgs_2a:
+            self.assertEqual(actual.topic(), topic)
+            self.assertEqual(actual.value().decode(), expected)
+
+        # commit autocommited indices
+        consumer_2a.close()
+
+        # Start another consumer with the same groupid
+        consumer_2b = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer_2",
+            auth=self.kafka.auth,
+            read_forever=False,
+            start_at=adc.consumer.ConsumerStartPosition.EARLIEST,
+        ))
+        consumer_2b.subscribe(topic)
+        stream_2b = consumer_2b.stream(autocommit=False, start_at=adc.consumer.LogicalOffset.STORED)
+
+        # read the rest using consumer_2b
+        msgs_2b = [msg for msg in stream_2b]
+
+        # Now check that messages from only batch_2 were read.
+        assert consumer_2b._stop_event.is_set()
+        self.assertEqual(len(batch_2), len(msgs_2b))
+        for expected, actual in zip(batch_2, msgs_2b):
+            self.assertEqual(actual.topic(), topic)
+            self.assertEqual(actual.value().decode(), expected)
+
+    def test_consume_from_beginning(self):
+        # Write a few messages.
+        topic = "test_consume_from_beginning"
+        batch = [
+            "message 1",
+            "message 2",
+            "message 3",
+            "message 4",
+        ]
+        simple_write_msgs(self.kafka, topic, batch)
+
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+            read_forever=False,
+            # Make reading start at the end by default
+            start_at=adc.consumer.ConsumerStartPosition.LATEST,
+        ))
+        consumer.subscribe(topic)
+        # Request reading from the beginning
+        stream = consumer.stream(start_at=adc.consumer.LogicalOffset.BEGINNING)
+        msgs = [msg for msg in stream]
+
+        assert consumer._stop_event.is_set()
+        self.assertEqual(len(batch), len(msgs))
+        for expected, actual in zip(batch, msgs):
             self.assertEqual(actual.topic(), topic)
             self.assertEqual(actual.value().decode(), expected)
 
+        # Read again from the beginning
+        stream = consumer.stream(start_at=adc.consumer.LogicalOffset.BEGINNING)
+        msgs = [msg for msg in stream]
+
+        assert consumer._stop_event.is_set()
+        self.assertEqual(len(batch), len(msgs))
+        for expected, actual in zip(batch, msgs):
+            self.assertEqual(actual.topic(), topic)
+            self.assertEqual(actual.value().decode(), expected)
+
+    def test_consume_from_end(self):
+        # Write a few messages.
+        topic = "test_consume_from_end"
+        simple_write_msgs(self.kafka, topic, [
+            "message 1",
+            "message 2",
+            "message 3",
+        ])
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+            # Make reading start at the beginning by default
+            start_at=adc.consumer.ConsumerStartPosition.EARLIEST,
+        ))
+        consumer.subscribe(topic)
+        # Request reading from the end
+        stream = consumer.stream(start_at=adc.consumer.LogicalOffset.END)
+
+        # Now add messages after the "end"
+        time.sleep(0.5)
+        simple_write_msg(self.kafka, topic, "message 4")
+        time.sleep(0.5)
+        msg = next(stream)
+        self.assertEqual(msg.topic(), topic)
+        self.assertEqual(msg.value(), b"message 4")
+
+    def test_consume_from_datetime(self):
+        # Write a few messages.
+        topic = "test_consume_from_datetime"
+        simple_write_msgs(self.kafka, topic, [
+            "message 1",
+            "message 2",
+            "message 3",
+        ])
+        # Wait a while, write, and wait some more
+        time.sleep(2)
+        client_middle_time = datetime.now()
+        time.sleep(2)
+        simple_write_msg(self.kafka, topic, "message 4")
+        time.sleep(1)
+
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+            read_forever=False,
+            start_at=adc.consumer.ConsumerStartPosition.EARLIEST,
+        ))
+        consumer.subscribe(topic)
+        stream = consumer.stream()
+        timestamps = [datetime.fromtimestamp(msg.timestamp()[1] / 1000.0) for msg in stream]
+
+        middle_time = timestamps[2] + (timestamps[3] - timestamps[2]) / 2
+        diff = middle_time - client_middle_time
+        logger.info(f"Difference between client and received timestamps: {diff!s}")
+
+        stream = consumer.stream(start_at=middle_time)
+        msg = next(stream)
+        self.assertEqual(msg.topic(), topic)
+        self.assertEqual(msg.value(), b"message 4")
+
     def test_consume_not_forever(self):
         topic = "test_consume_not_forever"
         simple_write_msg(self.kafka, topic, "message 1")
@@ -248,6 +402,43 @@ class KafkaIntegrationTestCase(unittest.TestCase):
         self.assertEqual(messages[1].value(), b"message 2")
         self.assertEqual(messages[2].value(), b"message 3")
 
+    def test_multi_topic_handling(self):
+        """Use a single producer object to write messages to multiple topics,
+        and check that a consumer can receive them all.
+
+        """
+        topics = ["test_multi_1", "test_multi_2"]
+
+        # Push some messages in
+        producer = adc.producer.Producer(adc.producer.ProducerConfig(
+            broker_urls=[self.kafka.address],
+            topic=None,
+            auth=self.kafka.auth,
+        ))
+        for i in range(0, 8):
+            producer.write(str(i), topic=topics[i % 2])
+        producer.flush()
+        logger.info("messages sent")
+
+        # check that we receive the messages from the right topics
+        consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+            broker_urls=[self.kafka.address],
+            group_id="test_consumer",
+            auth=self.kafka.auth,
+        ))
+        consumer.subscribe(topics)
+        stream = consumer.stream()
+        total_messages = 0
+        for msg in stream:
+            if msg.error() is not None:
+                raise Exception(msg.error())
+            idx = int(msg.value())
+            self.assertEqual(msg.topic(), topics[idx % 2])
+            total_messages += 1
+            if total_messages == 8:
+                break
+        self.assertEqual(total_messages, 8)
+
 
 class KafkaDockerConnection:
     """Holds connection information for communicating with a Kafka broker running
@@ -285,6 +476,10 @@ class KafkaDockerConnection:
         self.auth = adc.auth.SASLAuth(
             user="test", password="test-pass",
             ssl_ca_location=self.certfile.name,
+            # disable endpoint verification because the docker service generates a certificate with
+            # a useless subject (the container ID) which can never match the hostname used to
+            # connect (typically 0.0.0.0)
+            ssl_endpoint_identification_algorithm="none",
         )
 
     def poll_for_kafka_broker_address(self, maxiter=20, sleep=timedelta(milliseconds=500)):
@@ -308,6 +503,8 @@ class KafkaDockerConnection:
         if not addrs:
             return None
         ip = addrs[0]['HostIp']
+        if len(ip) == 0:
+            ip = "localhost"
         port = addrs[0]['HostPort']
         return f"{ip}:{port}"
 
@@ -373,8 +570,16 @@ class KafkaDockerConnection:
             detach=True,
             auto_remove=True,
             network=self.net.name,
-            # Setting None below the OS pick an ephemeral port.
-            ports={"9092/tcp": None},
+            # Kafka insists on redirecting consumers to one of its advertised listeners,
+            # which it will get wrong if it is running in a private container network.
+            # To fix this, we need to tell it what to advertise, which means we must
+            # know what port will be visible from the host system, and we cannot use an
+            # ephemeral port, which would be known to us only after the container is
+            # started. Since we have to pick something, pick 9092, which means that
+            # these tests cannot run if there is already an instance of Kafka running on
+            # the same host.
+            ports={"9092/tcp": 9092},
+            command=["/root/runServer", "--advertisedListener", "SASL_SSL://localhost:9092"],
         )
 
     def get_or_create_docker_network(self):
@@ -389,13 +594,13 @@ class KafkaDockerConnection:
         return self.docker_client.networks.create(name="adc-integration-test")
 
 
-def simple_write_msg(conn: KafkaDockerConnection, topic: str, msg: str):
+def simple_write_msg(conn: KafkaDockerConnection, topic: str, msg: str, key: Optional[str] = None):
     producer = adc.producer.Producer(adc.producer.ProducerConfig(
         broker_urls=[conn.address],
         topic=topic,
         auth=conn.auth,
     ))
-    producer.write(msg)
+    producer.write(msg, key=key)
     producer.flush()
 
 

adc-streaming-2.3.2/adc/errors.py

@@ -1,54 +0,0 @@
-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)
-    elif msg.error() is not None:
-        raise KafkaException.from_kafka_error(msg.error())
-
-
-class KafkaException(Exception):
-    @classmethod
-    def from_kafka_error(cls, error):
-        return cls(error)
-
-    def __init__(self, error):
-        self.error = error
-        self.name = error.name()
-        self.reason = error.str()
-        self.retriable = error.retriable()
-        self.fatal = error.fatal()
-        msg = f"Error communicating with Kafka: code={self.name} {self.reason}"
-        super(KafkaException, self).__init__(msg)