fix: address reviewer feedback - pure should_emit, pytest fixture, read loop counter, Docker build

Co-Authored-By: bot_apk <apk@cognition.ai>

Author: devin-ai-integration[bot]

Dockerfile

@@ -17,8 +17,11 @@ COPY dist/*.whl ./dist/
 RUN poetry config virtualenvs.create false \
     && poetry install --only main --no-interaction --no-ansi || true
 
-# Build and install the package (with metrics extra for DogStatsD support)
-RUN pip install "dist/*.whl[metrics]"
+# Build and install the package with the metrics extra (for DogStatsD support).
+# Two-step install: first install the wheel, then reinstall with extras using
+# the resolved filename (shell glob + [extras] syntax don't mix in a single arg).
+RUN pip install dist/*.whl \
+    && pip install "$(ls dist/*.whl)[metrics]"
 
 # Recreate the original structure
 RUN mkdir -p source_declarative_manifest \

airbyte_cdk/entrypoint.py

@@ -283,11 +283,16 @@ def read(
 
         # The Airbyte protocol dictates that counts be expressed as float/double to better protect against integer overflows
         stream_message_counter: DefaultDict[HashableStreamDescriptor, float] = defaultdict(float)
+        _emit_check_counter = 0
         try:
             for message in self.source.read(self.logger, config, catalog, state):
                 yield self.handle_record_counts(message, stream_message_counter)
-                # Periodically emit memory metrics (every 30s by default)
-                metrics_client.maybe_emit_memory_metrics()
+                # Check for metric emission every 1000 records to avoid
+                # a time.monotonic() syscall on every single record.
+                _emit_check_counter += 1
+                if _emit_check_counter >= 1000:
+                    _emit_check_counter = 0
+                    metrics_client.maybe_emit_memory_metrics()
         finally:
             # Emit final memory metrics snapshot (runs even if the read loop raises)
             metrics_client.emit_memory_metrics()

airbyte_cdk/metrics/__init__.py

@@ -180,12 +180,9 @@ def should_emit(self, interval_seconds: float = DEFAULT_EMISSION_INTERVAL_SECOND
         Check if enough time has passed since the last emission to emit again.
 
         Returns True if at least interval_seconds have elapsed since the last emission.
+        This is a pure query — it does not update any internal state.
         """
-        now = time.monotonic()
-        if now - self._last_emission_time >= interval_seconds:
-            self._last_emission_time = now
-            return True
-        return False
+        return time.monotonic() - self._last_emission_time >= interval_seconds
 
     def maybe_emit_memory_metrics(
         self, interval_seconds: float = DEFAULT_EMISSION_INTERVAL_SECONDS

unit_tests/metrics/test_metrics_client.py

@@ -13,32 +13,36 @@
 
 from airbyte_cdk.metrics.memory import MemoryInfo
 
-# Ensure mock datadog module is available for tests regardless of whether
-# the optional `datadog` package is installed.
-_mock_dogstatsd_cls = MagicMock()
-
-if "datadog" not in sys.modules:
-    _mock_datadog = types.ModuleType("datadog")
-    _mock_dogstatsd_mod = types.ModuleType("datadog.dogstatsd")
-    _mock_dogstatsd_mod.DogStatsd = _mock_dogstatsd_cls  # type: ignore[attr-defined]
-    _mock_datadog.dogstatsd = _mock_dogstatsd_mod  # type: ignore[attr-defined]
-    sys.modules["datadog"] = _mock_datadog
-    sys.modules["datadog.dogstatsd"] = _mock_dogstatsd_mod
+
+@pytest.fixture(autouse=True)
+def _mock_datadog(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
+    """Provide a mock datadog module so tests work regardless of whether the
+    optional ``datadog`` package is installed.  ``monkeypatch`` automatically
+    restores ``sys.modules`` after each test, preventing cross-test pollution."""
+    mock_cls = MagicMock()
+    mock_mod = types.ModuleType("datadog.dogstatsd")
+    mock_mod.DogStatsd = mock_cls  # type: ignore[attr-defined]
+    mock_datadog = types.ModuleType("datadog")
+    mock_datadog.dogstatsd = mock_mod  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "datadog", mock_datadog)
+    monkeypatch.setitem(sys.modules, "datadog.dogstatsd", mock_mod)
+    return mock_cls
+
 
 import airbyte_cdk.metrics as metrics_module  # noqa: E402
 from airbyte_cdk.metrics import MetricsClient, get_metrics_client  # noqa: E402
 
 
-def _make_enabled_client() -> tuple[MetricsClient, MagicMock]:
+def _make_enabled_client(mock_dogstatsd_cls: MagicMock) -> tuple[MetricsClient, MagicMock]:
     """Helper to create an enabled MetricsClient with a mock DogStatsd instance."""
     mock_instance = MagicMock()
-    _mock_dogstatsd_cls.reset_mock()
-    _mock_dogstatsd_cls.return_value = mock_instance
+    mock_dogstatsd_cls.reset_mock()
+    mock_dogstatsd_cls.return_value = mock_instance
 
     client = MetricsClient()
     with (
         patch.dict("os.environ", {"DD_AGENT_HOST": "localhost"}, clear=True),
-        patch("datadog.dogstatsd.DogStatsd", _mock_dogstatsd_cls),
+        patch("datadog.dogstatsd.DogStatsd", mock_dogstatsd_cls),
     ):
         client.initialize()
     return client, mock_instance
@@ -51,8 +55,8 @@ def test_disabled_when_dd_agent_host_not_set(self) -> None:
             client.initialize()
         assert not client.enabled
 
-    def test_enabled_when_dd_agent_host_set(self) -> None:
-        client, _ = _make_enabled_client()
+    def test_enabled_when_dd_agent_host_set(self, _mock_datadog: MagicMock) -> None:
+        client, _ = _make_enabled_client(_mock_datadog)
         assert client.enabled
 
     def test_initialize_idempotent(self) -> None:
@@ -73,10 +77,10 @@ def test_disabled_when_datadog_import_fails(self) -> None:
 
 
 class TestMetricsClientTags:
-    def test_builds_tags_from_env(self) -> None:
+    def test_builds_tags_from_env(self, _mock_datadog: MagicMock) -> None:
         mock_instance = MagicMock()
-        _mock_dogstatsd_cls.reset_mock()
-        _mock_dogstatsd_cls.return_value = mock_instance
+        _mock_datadog.reset_mock()
+        _mock_datadog.return_value = mock_instance
 
         client = MetricsClient()
         env = {
@@ -88,7 +92,7 @@ def test_builds_tags_from_env(self) -> None:
         }
         with (
             patch.dict("os.environ", env, clear=True),
-            patch("datadog.dogstatsd.DogStatsd", _mock_dogstatsd_cls),
+            patch("datadog.dogstatsd.DogStatsd", _mock_datadog),
         ):
             client.initialize()
 
@@ -104,30 +108,30 @@ def test_gauge_noop_when_disabled(self) -> None:
         # Should not raise even when not initialized
         client.gauge("test.metric", 42.0)
 
-    def test_gauge_emits_when_enabled(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_gauge_emits_when_enabled(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
 
         client.gauge("test.metric", 42.0)
         mock_instance.gauge.assert_called_once_with("test.metric", 42.0, tags=client._tags)
 
-    def test_gauge_with_extra_tags(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_gauge_with_extra_tags(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
 
         client.gauge("test.metric", 42.0, extra_tags=["stream:users"])
         call_tags = mock_instance.gauge.call_args[1]["tags"]
         assert "stream:users" in call_tags
 
-    def test_gauge_swallows_exceptions(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_gauge_swallows_exceptions(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
         mock_instance.gauge.side_effect = Exception("network error")
 
         # Should not raise
         client.gauge("test.metric", 42.0)
 
 
 class TestEmitMemoryMetrics:
-    def test_emits_all_metrics_when_enabled(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_emits_all_metrics_when_enabled(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
 
         mock_info = MemoryInfo(usage_bytes=100_000_000, limit_bytes=200_000_000)
         with patch("airbyte_cdk.metrics.get_memory_info", return_value=mock_info):
@@ -138,8 +142,8 @@ def test_emits_all_metrics_when_enabled(self) -> None:
         assert gauge_calls["cdk.memory.limit_bytes"] == 200_000_000.0
         assert gauge_calls["cdk.memory.usage_percent"] == pytest.approx(0.5)
 
-    def test_skips_limit_when_unknown(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_skips_limit_when_unknown(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
 
         mock_info = MemoryInfo(usage_bytes=100_000_000, limit_bytes=None)
         with patch("airbyte_cdk.metrics.get_memory_info", return_value=mock_info):
@@ -164,18 +168,22 @@ def test_emits_on_first_call(self) -> None:
     def test_does_not_emit_before_interval(self) -> None:
         client = MetricsClient()
         assert client.should_emit(interval_seconds=30.0)
+        # Manually advance the timestamp (emit_memory_metrics is a no-op
+        # when the client is disabled, so set it directly).
+        client._last_emission_time = time.monotonic()
         assert not client.should_emit(interval_seconds=30.0)
 
     def test_emits_after_interval(self) -> None:
         client = MetricsClient()
         assert client.should_emit(interval_seconds=0.01)
+        client._last_emission_time = time.monotonic()
         time.sleep(0.02)
         assert client.should_emit(interval_seconds=0.01)
 
 
 class TestMaybeEmitMemoryMetrics:
-    def test_emits_on_interval(self) -> None:
-        client, mock_instance = _make_enabled_client()
+    def test_emits_on_interval(self, _mock_datadog: MagicMock) -> None:
+        client, mock_instance = _make_enabled_client(_mock_datadog)
 
         mock_info = MemoryInfo(usage_bytes=100, limit_bytes=200)
         with patch("airbyte_cdk.metrics.get_memory_info", return_value=mock_info):