From 2085eb9d928a25f4e1f4e161433739b4e3dece40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Fri, 22 May 2026 12:44:20 +0200 Subject: [PATCH] feature(datashare-python): support exporting prometheus metrics --- datashare-python/datashare_python/config.py | 10 ++++++++++ datashare-python/tests/test_config.py | 19 ++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/datashare-python/datashare_python/config.py b/datashare-python/datashare_python/config.py index 34b60032..7adf301d 100644 --- a/datashare-python/datashare_python/config.py +++ b/datashare-python/datashare_python/config.py @@ -13,6 +13,7 @@ DefaultPayloadConverter, JSONPlainPayloadConverter, ) +from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig import datashare_python @@ -64,13 +65,22 @@ def to_task_client(self) -> DatashareTaskClient: class TemporalClientConfig(BaseModel): host: str = "temporal:7233" namespace: str = "datashare-default" + prometheus_host: str | None = None + _client: TemporalClient | None = PrivateAttr(default=None) async def to_client(self) -> TemporalClient: if self._client is None: + runtime = None + if self.prometheus_host is not None: + telemetry_config = TelemetryConfig( + metrics=PrometheusConfig(bind_address="0.0.0.0:9000") + ) + runtime = Runtime(telemetry=telemetry_config) self._client = await TemporalClient.connect( target_host=self.host, namespace=self.namespace, + runtime=runtime, data_converter=PYDANTIC_DATA_CONVERTER, ) return self._client diff --git a/datashare-python/tests/test_config.py b/datashare-python/tests/test_config.py index ca2ac8c8..d658337f 100644 --- a/datashare-python/tests/test_config.py +++ b/datashare-python/tests/test_config.py @@ -1,6 +1,9 @@ import os +from unittest.mock import AsyncMock, patch -from datashare_python.config import WorkerConfig +from datashare_python.config import ( + WorkerConfig, +) def test_worker_config_loggers_from_env(reset_env) -> None: # noqa: ANN001, ARG001 @@ -11,3 +14,17 @@ def test_worker_config_loggers_from_env(reset_env) -> None: # noqa: ANN001, ARG config = WorkerConfig() # Then assert config.logging.loggers["datashare_python"] == "WARNING" + + +async def test_worker_config_should_export_prometheus_metrics(reset_env) -> None: # noqa: ANN001, ARG001 + # Given + prometheus_host = "0.0.0.0:9000" + os.environ["DS_WORKER_TEMPORAL__PROMETHEUS_HOST"] = prometheus_host + config = WorkerConfig() + assert config.temporal.prometheus_host == prometheus_host + # When + mock_connect = AsyncMock() + with patch("datashare_python.config.TemporalClient.connect", mock_connect): + await config.temporal.to_client() + # Then + assert mock_connect.await_args_list[0].kwargs["runtime"] is not None