Skip to content

Commit 1ac910c

Browse files
committed
feature(datashare-python): support exporting prometheus metrics
1 parent f7a6cc9 commit 1ac910c

2 files changed

Lines changed: 31 additions & 1 deletion

File tree

datashare-python/datashare_python/config.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
DefaultPayloadConverter,
1414
JSONPlainPayloadConverter,
1515
)
16+
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
1617

1718
import datashare_python
1819

@@ -64,13 +65,22 @@ def to_task_client(self) -> DatashareTaskClient:
6465
class TemporalClientConfig(BaseModel):
6566
host: str = "temporal:7233"
6667
namespace: str = "datashare-default"
68+
prometheus_host: str | None = None
69+
6770
_client: TemporalClient | None = PrivateAttr(default=None)
6871

6972
async def to_client(self) -> TemporalClient:
7073
if self._client is None:
74+
runtime = None
75+
if self.prometheus_host is None:
76+
telemetry_config = TelemetryConfig(
77+
metrics=PrometheusConfig(bind_address="0.0.0.0:9000")
78+
)
79+
runtime = Runtime(telemetry=telemetry_config)
7180
self._client = await TemporalClient.connect(
7281
target_host=self.host,
7382
namespace=self.namespace,
83+
runtime=runtime,
7484
data_converter=PYDANTIC_DATA_CONVERTER,
7585
)
7686
return self._client

datashare-python/tests/test_config.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
import os
2+
from unittest.mock import AsyncMock, patch
23

3-
from datashare_python.config import WorkerConfig
4+
from datashare_python.config import (
5+
DatashareClientConfig,
6+
TemporalClientConfig,
7+
WorkerConfig,
8+
)
49

510

611
def test_worker_config_loggers_from_env(reset_env) -> None: # noqa: ANN001, ARG001
@@ -11,3 +16,18 @@ def test_worker_config_loggers_from_env(reset_env) -> None: # noqa: ANN001, ARG
1116
config = WorkerConfig()
1217
# Then
1318
assert config.logging.loggers["datashare_python"] == "WARNING"
19+
20+
21+
async def test_worker_config_should_export_prometheus_metrics(reset_env) -> None: # noqa: ANN001, ARG001
22+
# Given
23+
prometheus_host = "0.0.0.0:9000"
24+
os.environ["DS_WORKER_TEMPORAL__PROMETHEUS_HOST"] = prometheus_host
25+
config = WorkerConfig(
26+
datashare=DatashareClientConfig(url="http://localhost:8080"),
27+
temporal=TemporalClientConfig(host="localhost:7233"),
28+
)
29+
# When
30+
mock_connect = AsyncMock()
31+
with patch("datashare_python.config.TemporalClient.connect", mock_connect):
32+
await config.temporal.to_client()
33+
assert mock_connect.await_args_list[0].kwargs["runtime"] is not None

0 commit comments

Comments
 (0)