python(feat): Support additional configurations for sift-stream-bindings

python/changelog.d/+stream-metrics-rate-limit-options.feat.md

@@ -0,0 +1,3 @@
+Added `metrics_streaming_interval` and `rate_limit_policy` to `create_ingestion_config_streaming_client` and `create_auto_register_streaming_client`.
+
+`metrics_streaming_interval` takes a `timedelta` and defaults to 500 ms. Pass `None` to disable metrics streaming. `rate_limit_policy` takes a `RateLimitPolicyPy` from `sift_stream_bindings` and controls how rate-limited control-plane calls, such as stream setup, `attach_run`, and flow registration, are handled. By default, the stream waits the delay Sift requests before retrying. When a rate limit is not retried, or the retries run out, these calls raise `RateLimitedErrorPy`, and `retry_delay.to_timedelta()` returns the delay Sift requested. This release requires `sift-stream-bindings` 0.6.1.

python/lib/sift_client/_internal/low_level_wrappers/ingestion.py

@@ -11,6 +11,7 @@
 import hashlib
 import logging
 from collections import namedtuple
+from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Iterable, cast
 
 from sift.ingestion_configs.v2.ingestion_configs_pb2 import (
@@ -35,6 +36,9 @@
 DEFAULT_INGESTION_CONFIG_PAGE_SIZE = 100
 """Default page size for ingestion config and flow list calls (flow configs can be large)."""
 
+DEFAULT_METRICS_STREAMING_INTERVAL = timedelta(milliseconds=500)
+"""Default interval at which SiftStream pushes stream metrics to Sift. Matches the SiftStream default."""
+
 if TYPE_CHECKING:
     from datetime import datetime
 
@@ -47,6 +51,7 @@
         IngestWithConfigDataStreamRequestPy,
         IngestWithConfigDataStreamRequestWrapperPy,
         MetadataPy,
+        RateLimitPolicyPy,
         RetryPolicyPy,
         RunFormPy,
         RunSelectorPy,
@@ -248,6 +253,8 @@ async def _build_sift_stream_instance(
     retry_policy: RetryPolicyPy | None = None,
     disk_backup_policy: DiskBackupPolicyPy | None = None,
     checkpoint_interval_seconds: int | None = None,
+    metrics_streaming_interval: timedelta | None = DEFAULT_METRICS_STREAMING_INTERVAL,
+    rate_limit_policy: RateLimitPolicyPy | None = None,
     enable_tls: bool = True,
     tracing_config: TracingConfig | None = None,
 ) -> SiftStreamPy:
@@ -273,6 +280,17 @@ async def _build_sift_stream_instance(
     if streaming_mode is ...:  # type: ignore[comparison-overlap]
         streaming_mode = StreamingMode.LIVE_WITH_BACKUPS
 
+    # None disables metrics streaming, so it passes through to the builder as-is. Zero is
+    # rejected here because SiftStream cannot schedule metrics on a zero interval.
+    metrics_interval: DurationPy | None = None
+    if metrics_streaming_interval is not None:
+        if metrics_streaming_interval <= timedelta(0):
+            raise ValueError(
+                "metrics_streaming_interval must be greater than zero, "
+                f"got {metrics_streaming_interval}"
+            )
+        metrics_interval = DurationPy.from_timedelta(metrics_streaming_interval)
+
     if not is_tracing_initialized():
         if tracing_config is None:
             tracing_config = TracingConfig.with_file()
@@ -295,18 +313,24 @@ async def _build_sift_stream_instance(
     config_builder.run_id = run_id
     config_builder.asset_tags = asset_tags
     config_builder.metadata = asset_metadata
+    if rate_limit_policy is not None:
+        config_builder.rate_limit_policy = rate_limit_policy
 
     if streaming_mode == StreamingMode.LIVE_ONLY:
-        return await config_builder.live_only().build()
+        lo_builder = config_builder.live_only()
+        lo_builder.metrics_streaming_interval = metrics_interval
+        return await lo_builder.build()
 
     elif streaming_mode == StreamingMode.FILE_BACKUP:
         fb_builder = config_builder.file_backup()
+        fb_builder.metrics_streaming_interval = metrics_interval
         if disk_backup_policy is not None:
             fb_builder.disk_backup_policy = disk_backup_policy
         return await fb_builder.build()
 
     else:  # LIVE_WITH_BACKUPS (default)
         lwb_builder = config_builder.live_with_backups()
+        lwb_builder.metrics_streaming_interval = metrics_interval
         if retry_policy is not None:
             lwb_builder.retry_policy = retry_policy
         if disk_backup_policy is not None:
@@ -339,6 +363,8 @@ async def create_sift_stream_instance(
         retry_policy: RetryPolicyPy | None = None,
         disk_backup_policy: DiskBackupPolicyPy | None = None,
         checkpoint_interval_seconds: int | None = None,
+        metrics_streaming_interval: timedelta | None = DEFAULT_METRICS_STREAMING_INTERVAL,
+        rate_limit_policy: RateLimitPolicyPy | None = None,
         enable_tls: bool = True,
         tracing_config: TracingConfig | None = None,
     ) -> IngestionConfigStreamingLowLevelClient:
@@ -354,6 +380,8 @@ async def create_sift_stream_instance(
             retry_policy=retry_policy,
             disk_backup_policy=disk_backup_policy,
             checkpoint_interval_seconds=checkpoint_interval_seconds,
+            metrics_streaming_interval=metrics_streaming_interval,
+            rate_limit_policy=rate_limit_policy,
             enable_tls=enable_tls,
             tracing_config=tracing_config,
         )

python/lib/sift_client/_tests/_internal/low_level_wrappers/test_ingestion_stream_builder.py

@@ -0,0 +1,99 @@
+"""Tests for how `_build_sift_stream_instance` configures the SiftStream builders."""
+
+from datetime import timedelta
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from sift_client._internal.low_level_wrappers.ingestion import _build_sift_stream_instance
+from sift_client.resources.ingestion import StreamingMode
+
+sift_stream_bindings = pytest.importorskip("sift_stream_bindings")
+
+MODE_BUILDER_METHODS = {
+    StreamingMode.LIVE_ONLY: "live_only",
+    StreamingMode.LIVE_WITH_BACKUPS: "live_with_backups",
+    StreamingMode.FILE_BACKUP: "file_backup",
+}
+
+
+@pytest.fixture
+def config_builder(monkeypatch):
+    """Replace `SiftStreamBuilderPy` with a mock and return the config builder it produces."""
+    sift_builder = MagicMock()
+    config_builder = sift_builder.ingestion_config.return_value
+    for method in MODE_BUILDER_METHODS.values():
+        getattr(config_builder, method).return_value.build = AsyncMock()
+
+    monkeypatch.setattr(
+        sift_stream_bindings, "SiftStreamBuilderPy", MagicMock(return_value=sift_builder)
+    )
+    monkeypatch.setattr(sift_stream_bindings, "is_tracing_initialized", lambda: True)
+    return config_builder
+
+
+async def _build(**kwargs):
+    await _build_sift_stream_instance(
+        api_key="api-key",
+        grpc_uri="localhost:50051",
+        ingestion_config_form=MagicMock(),
+        **kwargs,
+    )
+
+
+@pytest.mark.parametrize("streaming_mode", list(MODE_BUILDER_METHODS))
+class TestMetricsStreamingInterval:
+    @pytest.mark.asyncio
+    async def test_defaults_to_500_ms(self, config_builder, streaming_mode):
+        await _build(streaming_mode=streaming_mode)
+
+        mode_builder = getattr(config_builder, MODE_BUILDER_METHODS[streaming_mode]).return_value
+        assert mode_builder.metrics_streaming_interval.to_timedelta() == timedelta(milliseconds=500)
+        mode_builder.build.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_custom_interval(self, config_builder, streaming_mode):
+        await _build(
+            streaming_mode=streaming_mode,
+            metrics_streaming_interval=timedelta(seconds=2, milliseconds=250),
+        )
+
+        mode_builder = getattr(config_builder, MODE_BUILDER_METHODS[streaming_mode]).return_value
+        assert mode_builder.metrics_streaming_interval.to_timedelta() == timedelta(
+            seconds=2, milliseconds=250
+        )
+
+    @pytest.mark.asyncio
+    async def test_none_disables_metrics_streaming(self, config_builder, streaming_mode):
+        await _build(streaming_mode=streaming_mode, metrics_streaming_interval=None)
+
+        mode_builder = getattr(config_builder, MODE_BUILDER_METHODS[streaming_mode]).return_value
+        assert mode_builder.metrics_streaming_interval is None
+        mode_builder.build.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("interval", [timedelta(0), timedelta(seconds=-1)])
+    async def test_rejects_non_positive_interval(self, config_builder, streaming_mode, interval):
+        with pytest.raises(ValueError, match="greater than zero"):
+            await _build(streaming_mode=streaming_mode, metrics_streaming_interval=interval)
+
+        mode_builder = getattr(config_builder, MODE_BUILDER_METHODS[streaming_mode]).return_value
+        mode_builder.build.assert_not_awaited()
+
+
+class TestRateLimitPolicy:
+    @pytest.mark.asyncio
+    async def test_sets_policy_on_config_builder(self, config_builder):
+        policy = sift_stream_bindings.RateLimitPolicyPy.return_error()
+
+        await _build(rate_limit_policy=policy)
+
+        assert config_builder.rate_limit_policy is policy
+
+    @pytest.mark.asyncio
+    async def test_leaves_default_when_none(self, config_builder):
+        await _build()
+
+        assert not isinstance(
+            config_builder.rate_limit_policy, sift_stream_bindings.RateLimitPolicyPy
+        )

python/lib/sift_client/_tests/resources/test_ingestion.py

@@ -0,0 +1,44 @@
+"""Tests that rate-limit errors from SiftStream reach callers of the streaming clients unchanged."""
+
+from datetime import timedelta
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from sift_client.resources.ingestion import (
+    AutoRegisterStreamingClient,
+    IngestionConfigStreamingClient,
+)
+
+sift_stream_bindings = pytest.importorskip("sift_stream_bindings")
+
+
+def _rate_limited_error():
+    return sift_stream_bindings.RateLimitedErrorPy(
+        "[RateLimitedError] rate limited",
+        sift_stream_bindings.DurationPy.from_timedelta(timedelta(seconds=2, milliseconds=500)),
+    )
+
+
+class TestRateLimitedErrorPropagation:
+    @pytest.mark.asyncio
+    async def test_ingestion_config_client_attach_run(self):
+        low_level_client = MagicMock()
+        low_level_client.attach_run = AsyncMock(side_effect=_rate_limited_error())
+        client = IngestionConfigStreamingClient(MagicMock(), low_level_client)
+
+        with pytest.raises(sift_stream_bindings.RateLimitedErrorPy) as exc_info:
+            await client.attach_run("run-id")
+
+        assert exc_info.value.retry_delay.to_timedelta() == timedelta(seconds=2, milliseconds=500)
+
+    @pytest.mark.asyncio
+    async def test_auto_register_client_send(self):
+        stream = MagicMock()
+        stream.send = AsyncMock(side_effect=_rate_limited_error())
+        client = AutoRegisterStreamingClient(MagicMock(), stream)
+
+        with pytest.raises(sift_stream_bindings.RateLimitedErrorPy) as exc_info:
+            await client.send(MagicMock())
+
+        assert exc_info.value.retry_delay.to_timedelta() == timedelta(seconds=2, milliseconds=500)