Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions temporalio/bridge/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,15 @@ class ClientHttpConnectProxyConfig:
basic_auth: tuple[str, str] | None


@dataclass
class ClientDnsLoadBalancingConfig:
"""Python representation of the Rust struct for configuring DNS load
balancing.
"""

resolution_interval_millis: int


@dataclass
class ClientConfig:
"""Python representation of the Rust struct for configuring the client."""
Expand All @@ -71,6 +80,7 @@ class ClientConfig:
client_name: str
client_version: str
http_connect_proxy_config: ClientHttpConnectProxyConfig | None
dns_load_balancing_config: ClientDnsLoadBalancingConfig | None


@dataclass
Expand Down
28 changes: 23 additions & 5 deletions temporalio/bridge/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ pub struct ClientConfig {
retry_config: Option<ClientRetryConfig>,
keep_alive_config: Option<ClientKeepAliveConfig>,
http_connect_proxy_config: Option<ClientHttpConnectProxyConfig>,
dns_load_balancing_config: Option<ClientDnsLoadBalancingConfig>,
}

#[derive(FromPyObject)]
Expand Down Expand Up @@ -67,6 +68,11 @@ struct ClientHttpConnectProxyConfig {
pub basic_auth: Option<(String, String)>,
}

#[derive(FromPyObject)]
struct ClientDnsLoadBalancingConfig {
pub resolution_interval_millis: u64,
}

#[derive(FromPyObject)]
pub(crate) struct RpcCall {
pub(crate) rpc: String,
Expand Down Expand Up @@ -236,6 +242,14 @@ impl ClientConfig {
) -> PyResult<ConnectionOptions> {
let (ascii_headers, binary_headers) = partition_headers(self.metadata);
let has_proxy = self.http_connect_proxy_config.is_some();
// Core rejects DNS load balancing alongside an HTTP CONNECT proxy, so
// suppress DNS LB whenever a proxy is configured to keep the
// pre-existing behavior even if a caller leaves the default.
let dns_load_balancing = if has_proxy {
None
} else {
self.dns_load_balancing_config.map(Into::into)
};
let conn_opts = ConnectionOptions::new(
Url::parse(&self.target_url)
.map_err(|err| PyValueError::new_err(format!("invalid target URL: {err}")))?,
Expand All @@ -249,11 +263,7 @@ impl ClientConfig {
)
.keep_alive(self.keep_alive_config.map(Into::into))
.maybe_http_connect_proxy(self.http_connect_proxy_config.map(Into::into))
.dns_load_balancing(if has_proxy {
None
} else {
Some(DnsLoadBalancingOptions::default())
})
.dns_load_balancing(dns_load_balancing)
.headers(ascii_headers)
.binary_headers(binary_headers)
.maybe_api_key(self.api_key)
Expand Down Expand Up @@ -322,3 +332,11 @@ impl From<ClientHttpConnectProxyConfig> for HttpConnectProxyOptions {
}
}
}

impl From<ClientDnsLoadBalancingConfig> for DnsLoadBalancingOptions {
fn from(conf: ClientDnsLoadBalancingConfig) -> Self {
let mut opts = DnsLoadBalancingOptions::default();
opts.resolution_interval = Duration::from_millis(conf.resolution_interval_millis);
opts
}
}
18 changes: 18 additions & 0 deletions temporalio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@
)
from temporalio.service import (
ConnectConfig,
DnsLoadBalancingConfig,
HttpConnectProxyConfig,
KeepAliveConfig,
RetryConfig,
Expand Down Expand Up @@ -139,6 +140,8 @@ async def connect(
lazy: bool = False,
runtime: temporalio.runtime.Runtime | None = None,
http_connect_proxy_config: HttpConnectProxyConfig | None = None,
dns_load_balancing_config: DnsLoadBalancingConfig
| None = DnsLoadBalancingConfig.default,
header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC,
) -> Self:
"""Connect to a Temporal server.
Expand Down Expand Up @@ -194,6 +197,11 @@ async def connect(
used for workers.
runtime: The runtime for this client, or the default if unset.
http_connect_proxy_config: Configuration for HTTP CONNECT proxy.
dns_load_balancing_config: DNS load balancing configuration for the
client connection. Default is to re-resolve DNS every 30s. Can
be set to ``None`` to disable. Silently disabled when
``http_connect_proxy_config`` is set, since the two are mutually
exclusive.
header_codec_behavior: Encoding behavior for headers sent by the client.
"""
connect_config = temporalio.service.ConnectConfig(
Expand All @@ -207,6 +215,7 @@ async def connect(
lazy=lazy,
runtime=runtime,
http_connect_proxy_config=http_connect_proxy_config,
dns_load_balancing_config=dns_load_balancing_config,
)

def make_lambda(
Expand Down Expand Up @@ -2873,6 +2882,7 @@ class ClientConnectConfig(TypedDict, total=False):
lazy: bool
runtime: temporalio.runtime.Runtime | None
http_connect_proxy_config: HttpConnectProxyConfig | None
dns_load_balancing_config: DnsLoadBalancingConfig | None
header_codec_behavior: HeaderCodecBehavior


Expand Down Expand Up @@ -9774,6 +9784,8 @@ async def connect(
lazy: bool = False,
runtime: temporalio.runtime.Runtime | None = None,
http_connect_proxy_config: HttpConnectProxyConfig | None = None,
dns_load_balancing_config: DnsLoadBalancingConfig
| None = DnsLoadBalancingConfig.default,
) -> CloudOperationsClient:
"""Connect to a Temporal Cloud Operations API.

Expand Down Expand Up @@ -9810,6 +9822,11 @@ async def connect(
used for workers.
runtime: The runtime for this client, or the default if unset.
http_connect_proxy_config: Configuration for HTTP CONNECT proxy.
dns_load_balancing_config: DNS load balancing configuration for the
client connection. Default is to re-resolve DNS every 30s. Can
be set to ``None`` to disable. Silently disabled when
``http_connect_proxy_config`` is set, since the two are mutually
exclusive.
"""
# Add version if given
if version:
Expand All @@ -9826,6 +9843,7 @@ async def connect(
lazy=lazy,
runtime=runtime,
http_connect_proxy_config=http_connect_proxy_config,
dns_load_balancing_config=dns_load_balancing_config,
)
return CloudOperationsClient(
await temporalio.service.ServiceClient.connect(connect_config)
Expand Down
34 changes: 34 additions & 0 deletions temporalio/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,32 @@ def _to_bridge_config(
)


@dataclass(frozen=True)
class DnsLoadBalancingConfig:
"""DNS load balancing configuration for client connections.

When enabled, Core periodically re-resolves the target host's DNS records
and round-robins requests across the resolved addresses. Cannot be used
together with :py:class:`HttpConnectProxyConfig` -- DNS load balancing is
silently disabled when an HTTP CONNECT proxy is configured.
"""

resolution_interval_millis: int = 30000
"""How often to re-resolve DNS, in milliseconds."""
default: ClassVar[DnsLoadBalancingConfig]
"""Default DNS load balancing config."""

def _to_bridge_config(
self,
) -> temporalio.bridge.client.ClientDnsLoadBalancingConfig:
return temporalio.bridge.client.ClientDnsLoadBalancingConfig(
resolution_interval_millis=self.resolution_interval_millis,
)


DnsLoadBalancingConfig.default = DnsLoadBalancingConfig()


@dataclass
class ConnectConfig:
"""Config for connecting to the server."""
Expand All @@ -146,6 +172,9 @@ class ConnectConfig:
lazy: bool = False
runtime: temporalio.runtime.Runtime | None = None
http_connect_proxy_config: HttpConnectProxyConfig | None = None
dns_load_balancing_config: DnsLoadBalancingConfig | None = (
DnsLoadBalancingConfig.default
)

def __post_init__(self) -> None:
"""Set extra defaults on unset properties."""
Expand Down Expand Up @@ -203,6 +232,11 @@ def _to_bridge_config(self) -> temporalio.bridge.client.ClientConfig:
if self.http_connect_proxy_config
else None
),
dns_load_balancing_config=(
self.dns_load_balancing_config._to_bridge_config()
if self.dns_load_balancing_config
else None
),
)


Expand Down
23 changes: 23 additions & 0 deletions tests/test_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,29 @@ def test_connect_config_tls_explicit_config_preserved():
assert config.tls == tls_config


def test_connect_config_dns_load_balancing_custom():
"""Custom DnsLoadBalancingConfig is forwarded to the bridge unchanged."""
config = temporalio.service.ConnectConfig(
target_host="localhost:7233",
dns_load_balancing_config=temporalio.service.DnsLoadBalancingConfig(
resolution_interval_millis=5000,
),
)
bridge_config = config._to_bridge_config()
assert bridge_config.dns_load_balancing_config is not None
assert bridge_config.dns_load_balancing_config.resolution_interval_millis == 5000


def test_connect_config_dns_load_balancing_disabled():
"""Setting dns_load_balancing_config=None forwards None to the bridge."""
config = temporalio.service.ConnectConfig(
target_host="localhost:7233",
dns_load_balancing_config=None,
)
bridge_config = config._to_bridge_config()
assert bridge_config.dns_load_balancing_config is None


async def test_rpc_execution_not_unknown(client: Client):
"""
Execute each rpc method and expect a failure, but ensure the failure is not that the rpc method is unknown
Expand Down
Loading