add redis publisher/consumer settings (#18160)

Author: zzstoatzz

src/integrations/prefect-redis/prefect_redis/messaging.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import json
 import socket
@@ -8,6 +10,7 @@
 from functools import partial
 from types import TracebackType
 from typing import (
+    Annotated,
     Any,
     AsyncGenerator,
     Awaitable,
@@ -20,6 +23,7 @@
 )
 
 import orjson
+from pydantic import BeforeValidator, Field
 from redis.asyncio import Redis
 from redis.exceptions import ResponseError
 from typing_extensions import Self
@@ -33,12 +37,84 @@
     StopConsumer,
 )
 from prefect.server.utilities.messaging import Publisher as _Publisher
+from prefect.settings.base import PrefectBaseSettings, build_settings_config
 from prefect_redis.client import get_async_redis_client
 
 logger = get_logger(__name__)
 
 M = TypeVar("M", bound=Message)
 
+
+def _interpret_string_as_timedelta_seconds(value: timedelta | str) -> timedelta:
+    """Interpret a string as a timedelta in seconds."""
+    if isinstance(value, str):
+        return timedelta(seconds=int(value))
+    return value
+
+
+TimeDelta = Annotated[
+    Union[str, timedelta],
+    BeforeValidator(_interpret_string_as_timedelta_seconds),
+]
+
+
+class RedisMessagingPublisherSettings(PrefectBaseSettings):
+    """Settings for the Redis messaging publisher.
+
+    No settings are required to be set by the user but any of the settings can be
+    overridden by the user using environment variables.
+
+    Example:
+        ```
+        PREFECT_REDIS_MESSAGING_PUBLISHER_BATCH_SIZE=10
+        PREFECT_REDIS_MESSAGING_PUBLISHER_PUBLISH_EVERY=10
+        PREFECT_REDIS_MESSAGING_PUBLISHER_DEDUPLICATE_BY=message_id
+        ```
+    """
+
+    model_config = build_settings_config(
+        (
+            "redis",
+            "messaging",
+            "publisher",
+        ),
+    )
+    batch_size: int = Field(default=5)
+    publish_every: TimeDelta = Field(default=timedelta(seconds=10))
+    deduplicate_by: Optional[str] = Field(default=None)
+
+
+class RedisMessagingConsumerSettings(PrefectBaseSettings):
+    """Settings for the Redis messaging consumer.
+
+    No settings are required to be set by the user but any of the settings can be
+    overridden by the user using environment variables.
+
+    Example:
+        ```
+        PREFECT_REDIS_MESSAGING_CONSUMER_BLOCK=10
+        PREFECT_REDIS_MESSAGING_CONSUMER_MIN_IDLE_TIME=10
+        PREFECT_REDIS_MESSAGING_CONSUMER_MAX_RETRIES=3
+        PREFECT_REDIS_MESSAGING_CONSUMER_TRIM_EVERY=60
+        ```
+    """
+
+    model_config = build_settings_config(
+        (
+            "redis",
+            "messaging",
+            "consumer",
+        ),
+    )
+    block: TimeDelta = Field(default=timedelta(seconds=1))
+    min_idle_time: TimeDelta = Field(default=timedelta(seconds=0))
+    max_retries: int = Field(default=3)
+    trim_every: TimeDelta = Field(default=timedelta(seconds=60))
+    should_process_pending_messages: bool = Field(default=True)
+    starting_message_id: str = Field(default="0")
+    automatically_acknowledge: bool = Field(default=True)
+
+
 MESSAGE_DEDUPLICATION_LOOKBACK = timedelta(minutes=5)
 
 
@@ -130,14 +206,20 @@ def __init__(
         topic: str,
         cache: _Cache,
         deduplicate_by: Optional[str] = None,
-        batch_size: int = 5,
+        batch_size: Optional[int] = None,
         publish_every: Optional[timedelta] = None,
     ):
+        settings = RedisMessagingPublisherSettings()
+
         self.stream = topic  # Use topic as stream name
         self.cache = cache
-        self.deduplicate_by = deduplicate_by
-        self.batch_size = batch_size
-        self.publish_every = publish_every
+        self.deduplicate_by = (
+            deduplicate_by if deduplicate_by is not None else settings.deduplicate_by
+        )
+        self.batch_size = batch_size if batch_size is not None else settings.batch_size
+        self.publish_every = (
+            publish_every if publish_every is not None else settings.publish_every
+        )
         self._periodic_task: Optional[asyncio.Task[None]] = None
 
     async def __aenter__(self) -> Self:
@@ -220,27 +302,45 @@ def __init__(
         topic: str,
         name: Optional[str] = None,
         group: Optional[str] = None,
-        block: timedelta = timedelta(seconds=1),
-        min_idle_time: timedelta = timedelta(seconds=0),
-        should_process_pending_messages: bool = True,
-        starting_message_id: str = "0",
-        automatically_acknowledge: bool = True,
-        max_retries: int = 3,
-        trim_every: timedelta = timedelta(seconds=60),
+        block: Optional[timedelta] = None,
+        min_idle_time: Optional[timedelta] = None,
+        should_process_pending_messages: Optional[bool] = None,
+        starting_message_id: Optional[str] = None,
+        automatically_acknowledge: Optional[bool] = None,
+        max_retries: Optional[int] = None,
+        trim_every: Optional[timedelta] = None,
     ):
+        settings = RedisMessagingConsumerSettings()
+
         self.name = name or topic
         self.stream = topic  # Use topic as stream name
         self.group = group or topic  # Use topic as default group name
-        self.block = block
-        self.min_idle_time = min_idle_time
-        self.should_process_pending_messages = should_process_pending_messages
-        self.starting_message_id = starting_message_id
-        self.automatically_acknowledge = automatically_acknowledge
+        self.block = block if block is not None else settings.block
+        self.min_idle_time = (
+            min_idle_time if min_idle_time is not None else settings.min_idle_time
+        )
+        self.should_process_pending_messages = (
+            should_process_pending_messages
+            if should_process_pending_messages is not None
+            else settings.should_process_pending_messages
+        )
+        self.starting_message_id = (
+            starting_message_id
+            if starting_message_id is not None
+            else settings.starting_message_id
+        )
+        self.automatically_acknowledge = (
+            automatically_acknowledge
+            if automatically_acknowledge is not None
+            else settings.automatically_acknowledge
+        )
+        self.trim_every = trim_every if trim_every is not None else settings.trim_every
 
-        self.subscription = Subscription(max_retries=max_retries)
+        self.subscription = Subscription(
+            max_retries=max_retries if max_retries is not None else settings.max_retries
+        )
         self._retry_counts: dict[str, int] = {}
 
-        self.trim_every = trim_every
         self._last_trimmed: Optional[float] = None
 
     async def _ensure_stream_and_group(self, redis_client: Redis) -> None:

src/integrations/prefect-redis/tests/test_messaging.py

@@ -12,6 +12,8 @@
     Consumer,
     Message,
     Publisher,
+    RedisMessagingConsumerSettings,
+    RedisMessagingPublisherSettings,
     StopConsumer,
 )
 from redis.asyncio import Redis
@@ -442,3 +444,39 @@ async def handler(message: Message):
         await consumer_task
 
     assert captured_events == [emitted_event]
+
+
+class TestRedisMessagingSettings:
+    """Test the Redis messaging settings."""
+
+    def test_publisher_settings(self):
+        """Test Redis publisher settings."""
+        settings = RedisMessagingPublisherSettings()
+        assert settings.batch_size == 5
+        assert settings.publish_every == timedelta(seconds=10)
+        assert settings.deduplicate_by is None
+
+    def test_consumer_settings(self):
+        """Test Redis consumer settings."""
+        settings = RedisMessagingConsumerSettings()
+        assert settings.block == timedelta(seconds=1)
+        assert settings.min_idle_time == timedelta(seconds=0)
+        assert settings.max_retries == 3
+        assert settings.trim_every == timedelta(seconds=60)
+        assert settings.should_process_pending_messages is True
+        assert settings.starting_message_id == "0"
+        assert settings.automatically_acknowledge is True
+
+    def test_publisher_settings_can_be_overridden(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """Test that Redis publisher settings can be overridden."""
+        monkeypatch.setenv("PREFECT_REDIS_MESSAGING_PUBLISHER_BATCH_SIZE", "10")
+        settings = RedisMessagingPublisherSettings()
+        assert settings.batch_size == 10
+
+    def test_consumer_settings_can_be_overridden(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that Redis consumer settings can be overridden."""
+        monkeypatch.setenv("PREFECT_REDIS_MESSAGING_CONSUMER_BLOCK", "10")
+        settings = RedisMessagingConsumerSettings()
+        assert settings.block == timedelta(seconds=10)