dynamic configuration lookup: wait until the configs are fully read from the topic (#1036)

Author: daniil-quix

quixstreams/dataframe/joins/lookups/quix_configuration_service/lookup.py

@@ -2,13 +2,14 @@
 import logging
 import threading
 import time
-from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Literal, Mapping, Optional, Union
 
 import httpx
 import orjson
+from confluent_kafka import TopicPartition
 
 from quixstreams.kafka import ConnectionConfig, Consumer
+from quixstreams.kafka.exceptions import KafkaConsumerException
 from quixstreams.models import HeadersMapping, Topic
 from quixstreams.platforms.quix.env import QUIX_ENVIRONMENT
 
@@ -119,19 +120,17 @@ def __init__(
             self._version_data, maxsize=cache_size
         )
 
-        self._consumer = Consumer(
+        self._config_consumer = Consumer(
             broker_address=broker_address,
             consumer_group=consumer_group,
             auto_offset_reset="earliest",
             auto_commit_enable=False,
             extra_config=consumer_extra_config,
         )
 
-        self._start()
+        self._start_consumer_thread()
 
-        self._fields_by_type: dict[int, dict[str, dict[str, BaseField]]] = defaultdict(
-            dict
-        )
+        self._fields_by_type: dict[int, dict[str, dict[str, BaseField]]] = {}
 
     def json_field(
         self,
@@ -203,7 +202,7 @@ def _fetch_version_content(self, version: ConfigurationVersion) -> Optional[byte
             version.failed()
             return None
 
-    def _start(self) -> None:
+    def _start_consumer_thread(self) -> None:
         """
         Start the enrichment process in a background thread and wait for initialization to complete.
         """
@@ -214,26 +213,31 @@ def _start(self) -> None:
 
     def _consumer_thread(self) -> None:
         """
-        Background thread for consuming configuration events from Kafka and updating internal state.
+        Background thread for consuming configuration events from Kafka
+        and updating internal state.
         """
-        assigned = False
-
-        def on_assign(consumer: Consumer, partitions: list[tuple[str, int]]) -> None:
-            """
-            Callback for partition assignment.
-            """
-            nonlocal assigned
-            assigned = True
-
         try:
-            self._consumer.subscribe(topics=[self._topic.name], on_assign=on_assign)
+            # Assign all available partitions of the config updates topic
+            # bypassing the consumer group protocol.
+            tps = [
+                TopicPartition(topic=self._topic.name, partition=i)
+                for i in range(self._topic.broker_config.num_partitions or 0)
+            ]
+            self._config_consumer.assign(tps)
 
             while True:
-                message = self._consumer.poll(timeout=self._consumer_poll_timeout)
+                # Check if the consumer processed all the available config messages
+                # and set the "started" event.
+                if not self._started.is_set() and self._configs_ready():
+                    self._started.set()
+
+                message = self._config_consumer.poll(
+                    timeout=self._consumer_poll_timeout
+                )
                 if message is None:
-                    if assigned and not self._started.is_set():
-                        self._started.set()
                     continue
+                elif message.error():
+                    raise KafkaConsumerException(error=message.error())
 
                 value = message.value()
                 if value is None:
@@ -267,7 +271,9 @@ def _process_config_event(self, event: Event) -> None:
 
         :raises: RuntimeError: If the event type is unknown.
         """
-        logger.info(f"Processing event: {event['event']} for ID: {event['id']}")
+        logger.info(
+            f'Processing update for configuration ID "{event["id"]}" ({event["event"]})'
+        )
         if event["event"] in {"created", "updated"}:
             if event["id"] not in self._configurations:
                 logger.debug(f"Creating new configuration for ID: {event['id']}")
@@ -329,28 +335,35 @@ def _find_version(
 
         :returns: The valid configuration version, or None if not found.
         """
-        logger.debug(
-            f"Fetching data for type: {type}, on: {on}, timestamp: {timestamp}"
-        )
+
         configuration = self._configurations.get(self._config_id(type, on))
-        if not configuration:
+        if configuration is None:
             logger.debug(
-                f"No configuration found for type: {type}, on: {on}. Trying wildcard."
+                "No configuration found for type: %s, on: %s. Trying wildcard.",
+                type,
+                on,
             )
             configuration = self._configurations.get(self._config_id(type, "*"))
-        if not configuration:
-            logger.debug(f"No configuration found for type: {type}, on: *")
+        if configuration is None:
+            logger.debug("No configuration found for type: %s, on: *", type)
             return None
 
         version = configuration.find_valid_version(timestamp)
         if version is None:
             logger.debug(
-                f"No valid version found for type: {type}, on: {on}, timestamp: {timestamp}"
+                "No valid configuration version found for type: %s, on: %s, timestamp: %s",
+                type,
+                on,
+                timestamp,
             )
             return None
 
         logger.debug(
-            f"Found valid version '{version.version}' for type: {type}, on: {on}, timestamp: {timestamp}"
+            "Found valid configuration version '%s' for type: %s, on: %s, timestamp: %s",
+            version.version,
+            type,
+            on,
+            timestamp,
         )
         return version
 
@@ -416,18 +429,42 @@ def join(
         start = time.time()
         logger.debug(f"Joining message with key: {on}, timestamp: {timestamp}")
 
-        fields_by_type = self._fields_by_type.get(id(fields))
+        fields_ids = id(fields)
+        fields_by_type = self._fields_by_type.get(fields_ids)
+
         if fields_by_type is None:
-            fields_by_type = defaultdict(dict)
+            fields_by_type = {}
             for key, field in fields.items():
-                fields_by_type[field.type][key] = field
-            self._fields_by_type[id(fields)] = fields_by_type
+                fields_by_type.setdefault(field.type, {})[key] = field
+            self._fields_by_type[fields_ids] = fields_by_type
+
+        for type_, fields in fields_by_type.items():
+            version = self._find_version(type_, on, timestamp)
 
-        for type, fields in fields_by_type.items():
-            version = self._find_version(type, on, timestamp)
             if version is not None and version.retry_at < start:
                 self._version_data_cached.remove(version, fields)
 
             value.update(self._version_data_cached(version, fields))
 
         logger.debug("Join took %.2f ms", (time.time() - start) * 1000)
+
+    def _configs_ready(self) -> bool:
+        """
+        Return True if the configs are loaded from the topic until the available HWM.
+        If there are outstanding messages in the config topic or the assignment is not
+        available yet, return False.
+        """
+
+        if not self._config_consumer.assignment():
+            return False
+
+        positions = self._config_consumer.position(self._config_consumer.assignment())
+        for position in positions:
+            # Check if the consumer reached the end of the configuration topic
+            # for each assigned partition.
+            _, hwm = self._config_consumer.get_watermark_offsets(
+                partition=position, cached=True
+            )
+            if hwm < 0 or (hwm > 0 and position.offset < hwm):
+                return False
+        return True

tests/test_quixstreams/test_dataframe/test_joins/test_lookup_quix_config.py

@@ -1,5 +1,6 @@
 import sys
 import time
+import uuid
 from datetime import datetime
 from typing import Optional
 from unittest.mock import Mock, patch
@@ -30,7 +31,7 @@ def lookup():
         patch(
             "quixstreams.dataframe.joins.lookups.quix_configuration_service.lookup.httpx.Client"
         ),
-        patch.object(Lookup, "_start"),
+        patch.object(Lookup, "_start_consumer_thread"),
     ):
         lookup = Lookup(
             topic=mock_topic,
@@ -59,36 +60,36 @@ def create_configuration_version(
     )
 
 
+def create_event(
+    id: str = "test-config",
+    version: int = 1,
+    valid_from: Optional[str] = None,
+    sha256sum: str = "test-hash",
+    content_url: str = "http://example.com/config",
+) -> Event:
+    """Helper method to create test events."""
+    return {
+        "id": id,
+        "event": "created",
+        "contentUrl": content_url,
+        "metadata": {
+            "type": "test-type",
+            "target_key": "test-key",
+            "valid_from": valid_from,
+            "category": "test-category",
+            "version": version,
+            "created_at": "2025-01-01T00:00:00Z",
+            "sha256sum": sha256sum,
+        },
+    }
+
+
 class TestConfiguration:
     """Test suite for the Configuration class."""
 
-    def create_event(
-        self,
-        id: str = "test-config",
-        version: int = 1,
-        valid_from: Optional[str] = None,
-        sha256sum: str = "test-hash",
-        content_url: str = "http://example.com/config",
-    ) -> Event:
-        """Helper method to create test events."""
-        return {
-            "id": id,
-            "event": "created",
-            "contentUrl": content_url,
-            "metadata": {
-                "type": "test-type",
-                "target_key": "test-key",
-                "valid_from": valid_from,
-                "category": "test-category",
-                "version": version,
-                "created_at": "2025-01-01T00:00:00Z",
-                "sha256sum": sha256sum,
-            },
-        }
-
     def test_from_event(self):
         """Test creating a Configuration from an Event."""
-        event = self.create_event(version=1, valid_from="2025-01-01T12:00:00Z")
+        event = create_event(version=1, valid_from="2025-01-01T12:00:00Z")
         config = Configuration.from_event(event)
 
         assert len(config.versions) == 1
@@ -469,10 +470,10 @@ def test_frozen_dataclass_immutability(self):
 
         # Should not be able to modify frozen fields
         with pytest.raises((AttributeError, TypeError)):
-            version.id = "modified"  # type: ignore
+            setattr(version, "id", "modified")
 
         with pytest.raises((AttributeError, TypeError)):
-            version.version = 2  # type: ignore
+            setattr(version, "version", 2)
 
         # But should be able to modify non-frozen fields via the methods
         version.failed()
@@ -1270,3 +1271,38 @@ def test_join_bytes_field_with_encoded_string(self, lookup):
 
             # Verify new data is added
             assert value["binary_data"].decode("utf-8") == test_content
+
+
+class TestLookupInit:
+    """
+    Test that Lookup loads all available configurations from the config topic
+    on init.
+    """
+
+    def test_lookup_init_config_topic_empty(self, app_factory):
+        app = app_factory()
+        config_topic = app.topic("configs")
+
+        lookup = Lookup(
+            topic=config_topic,
+            app_config=app.config,
+        )
+        assert not lookup._configurations
+
+    def test_lookup_init_config_topic_not_empty(self, app_factory):
+        app = app_factory()
+        config_topic = app.topic("configs")
+
+        num_configs = 50
+        with app.get_producer() as producer:
+            for _ in range(num_configs):
+                config_event = create_event(id=str(uuid.uuid4()))
+                producer.produce(
+                    topic=config_topic.name, value=orjson.dumps(config_event)
+                )
+
+        lookup = Lookup(
+            topic=config_topic,
+            app_config=app.config,
+        )
+        assert len(lookup._configurations) == num_configs