diff --git a/changelog/added/extension-extraordinaire.md b/changelog/added/extension-extraordinaire.md new file mode 100644 index 000000000..f4dd976cd --- /dev/null +++ b/changelog/added/extension-extraordinaire.md @@ -0,0 +1 @@ +Add experimental envconfig extensions: `ClientConfig.load_client_connect_config` accepts `extensions` that read their own config from `profile..ext.` TOML tables (exposed as `ClientConfigProfile.ext` and never rejected by strict mode) and env vars, and contribute client plugins to the returned config's `plugins`. diff --git a/temporalio/bridge/src/envconfig.rs b/temporalio/bridge/src/envconfig.rs index 0b21d1987..a21bb935a 100644 --- a/temporalio/bridge/src/envconfig.rs +++ b/temporalio/bridge/src/envconfig.rs @@ -1,12 +1,12 @@ use pyo3::prelude::*; use pyo3::{ exceptions::PyRuntimeError, - types::{PyBytes, PyDict}, + types::{PyBytes, PyDict, PyList}, }; use std::collections::HashMap; use temporalio_common::envconfig::{ load_client_config as core_load_client_config, - load_client_config_profile as core_load_client_config_profile, + load_client_config_profile as core_load_client_config_profile, toml, ClientConfig as CoreClientConfig, ClientConfigCodec, ClientConfigProfile as CoreClientConfigProfile, ClientConfigTLS as CoreClientConfigTLS, DataSource, LoadClientConfigOptions, LoadClientConfigProfileOptions, @@ -78,9 +78,81 @@ fn profile_to_dict(py: Python, profile: &CoreClientConfigProfile) -> PyResult PyResult> { + let dict = PyDict::new(py); + for (k, v) in table { + dict.set_item(k, toml_value_to_py(py, v)?)?; + } Ok(dict.into()) } +// Values are converted the same way as by Python's tomllib. +fn toml_value_to_py(py: Python, value: &toml::Value) -> PyResult> { + Ok(match value { + toml::Value::String(s) => s.into_pyobject(py)?.into_any().unbind(), + toml::Value::Integer(i) => i.into_pyobject(py)?.into_any().unbind(), + toml::Value::Float(f) => f.into_pyobject(py)?.into_any().unbind(), + toml::Value::Boolean(b) => b.into_pyobject(py)?.to_owned().into_any().unbind(), + toml::Value::Datetime(dt) => toml_datetime_to_py(py, dt)?, + toml::Value::Array(a) => { + let list = PyList::empty(py); + for v in a { + list.append(toml_value_to_py(py, v)?)?; + } + list.into_any().unbind() + } + toml::Value::Table(t) => toml_table_to_py(py, t)?, + }) +} + +// Offset datetimes become aware datetimes, local datetimes become naive datetimes, local dates +// become dates, and local times become times. Fractional seconds are truncated to microseconds. +// Built via the datetime module since the datetime C API isn't part of the limited API. +fn toml_datetime_to_py(py: Python, dt: &toml::value::Datetime) -> PyResult> { + let datetime = py.import("datetime")?; + let time = dt.time.map(|t| { + ( + t.hour, + t.minute, + t.second.unwrap_or(0), + t.nanosecond.unwrap_or(0) / 1000, + ) + }); + let obj = match (dt.date, time) { + (Some(d), Some((hour, minute, second, micros))) => { + let tz = match dt.offset { + None => py.None().into_bound(py), + Some(toml::value::Offset::Z) => datetime.getattr("timezone")?.getattr("utc")?, + Some(toml::value::Offset::Custom { minutes }) => { + let delta = datetime + .getattr("timedelta")? + .call1((0, i32::from(minutes) * 60))?; + datetime.getattr("timezone")?.call1((delta,))? + } + }; + datetime + .getattr("datetime")? + .call1((d.year, d.month, d.day, hour, minute, second, micros, tz))? + } + (Some(d), None) => datetime.getattr("date")?.call1((d.year, d.month, d.day))?, + (None, Some((hour, minute, second, micros))) => datetime + .getattr("time")? + .call1((hour, minute, second, micros))?, + (None, None) => return Err(PyRuntimeError::new_err("empty TOML datetime")), + }; + Ok(obj.unbind()) +} + fn core_config_to_dict(py: Python, core_config: &CoreClientConfig) -> PyResult> { let profiles_dict = PyDict::new(py); for (name, profile) in &core_config.profiles { diff --git a/temporalio/envconfig.py b/temporalio/envconfig.py index 8eb0d5314..2b1e8d85a 100644 --- a/temporalio/envconfig.py +++ b/temporalio/envconfig.py @@ -6,13 +6,16 @@ from __future__ import annotations -from collections.abc import Mapping +import os +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal, TypeAlias, cast from typing_extensions import Self, TypedDict +import temporalio.client import temporalio.service from temporalio.bridge.temporal_sdk_bridge import envconfig as _bridge_envconfig @@ -40,6 +43,7 @@ class ClientConfigProfileDict(TypedDict, total=False): api_key: str tls: ClientConfigTLSDict grpc_meta: Mapping[str, str] + ext: Mapping[str, Mapping[str, Any]] def _from_dict_to_source(d: Mapping[str, Any] | None) -> DataSource | None: @@ -175,6 +179,7 @@ class ClientConnectConfig(TypedDict, total=False): api_key: str tls: bool | temporalio.service.TLSConfig rpc_metadata: Mapping[str, str] + plugins: Sequence[temporalio.client.Plugin] @dataclass(frozen=True) @@ -196,6 +201,16 @@ class ClientConfigProfile: """TLS configuration.""" grpc_meta: Mapping[str, str] = field(default_factory=dict) """gRPC metadata.""" + ext: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + """Raw extension config from ``profile..ext``, keyed by extension + name. Each extension decodes its own table, which should be treated as + read-only. Values are converted as by :py:mod:`tomllib`, so TOML datetimes + are :py:class:`datetime.datetime`, :py:class:`datetime.date`, or + :py:class:`datetime.time` objects. Not modified by environment variables. + + .. warning:: + This API is experimental and unstable. + """ @classmethod def from_dict(cls, d: ClientConfigProfileDict) -> Self: @@ -206,6 +221,7 @@ def from_dict(cls, d: ClientConfigProfileDict) -> Self: api_key=d.get("api_key"), tls=ClientConfigTLS.from_dict(d.get("tls")), grpc_meta=d.get("grpc_meta") or {}, + ext=d.get("ext") or {}, ) def to_dict(self) -> ClientConfigProfileDict: @@ -221,6 +237,8 @@ def to_dict(self) -> ClientConfigProfileDict: d["tls"] = self.tls.to_dict() if self.grpc_meta: d["grpc_meta"] = self.grpc_meta + if self.ext: + d["ext"] = self.ext return d def to_client_connect_config(self) -> ClientConnectConfig: @@ -294,6 +312,61 @@ def load( return ClientConfigProfile.from_dict(raw_profile) +@dataclass(frozen=True) +class ClientConfigExtensionRequest: + """Input to :py:meth:`ClientConfigExtension.client_plugin`. + + .. warning:: + This API is experimental and unstable. + """ + + profile: ClientConfigProfile + """The loaded profile, with environment variable overrides applied.""" + env: Mapping[str, str] | None + """Env to read extension env vars from, or None if env loading is disabled.""" + + +class ClientConfigExtension(ABC): + """Extension that adds config to client environment configuration. + + An extension owns the TOML table under ``profile..ext.``, + where ```` is :py:meth:`name`, and turns its config into a client + plugin. Extensions are recommended to use a ``TEMPORAL_EXT__`` prefix + for any environment variables they read. Pass extensions to + :py:meth:`ClientConfig.load_client_connect_config`. + + .. warning:: + This API is experimental and unstable. + """ + + @abstractmethod + def name(self) -> str: + """The extension's table key under ``profile..ext``. Must be + unique among loaded extensions. + """ + + @abstractmethod + def client_plugin( + self, req: ClientConfigExtensionRequest + ) -> temporalio.client.Plugin | None: + """Read this extension's config from the profile and env, and return a + plugin that applies it, or None if there's nothing to apply. + """ + + +def _validate_extensions(extensions: Sequence[ClientConfigExtension]) -> None: + names: set[str] = set() + for i, extension in enumerate(extensions): + if extension is None: # type: ignore[reportUnnecessaryComparison] + raise ValueError(f"client config extension at index {i} is None") + name = extension.name() + if not name: + raise ValueError(f"client config extension at index {i} has empty name") + if name in names: + raise ValueError(f'duplicate client config extension name "{name}"') + names.add(name) + + @dataclass class ClientConfig: """Client configuration loaded from TOML and environment variables. @@ -371,12 +444,14 @@ def load_client_connect_config( disable_env: bool = False, config_file_strict: bool = False, override_env_vars: Mapping[str, str] | None = None, + extensions: Sequence[ClientConfigExtension] = [], ) -> ClientConnectConfig: """Load a single client profile and convert to connect config. This is a convenience function that combines loading a profile and - converting it to a connect config dictionary. This will use the current - process's environment for overrides unless disabled. + converting it to a connect config dictionary, plus any extensions. This + will use the current process's environment for overrides unless + disabled. Args: profile: The profile to load from the config. Defaults to "default". @@ -393,11 +468,22 @@ def load_client_connect_config( environment is used. To use a specific set of environment variables, provide them here. To disable environment variable loading, set ``disable_env`` to true. + extensions: Extensions to apply, in order. Each extension is given + the loaded profile and env, and any plugin it returns is + appended to the ``plugins`` of the returned config. Extension + names must be unique. To also pass other plugins to + :py:meth:`temporalio.client.Client.connect`, combine them with + these rather than passing ``plugins`` separately. + + .. warning:: + This parameter is experimental and unstable. Returns: TypedDict of keyword arguments for :py:meth:`temporalio.client.Client.connect`. """ + _validate_extensions(extensions) + config_source: DataSource | None = None if config_file and not disable_file: config_source = Path(config_file) @@ -410,4 +496,25 @@ def load_client_connect_config( config_file_strict=config_file_strict, override_env_vars=override_env_vars, ) - return prof.to_client_connect_config() + config = prof.to_client_connect_config() + + if extensions: + req = ClientConfigExtensionRequest( + profile=prof, + env=None + if disable_env + else (os.environ if override_env_vars is None else override_env_vars), + ) + plugins: list[temporalio.client.Plugin] = [] + for extension in extensions: + try: + plugin = extension.client_plugin(req) + except Exception as err: + raise RuntimeError( + f'client config extension "{extension.name()}": {err}' + ) from err + if plugin is not None: + plugins.append(plugin) + if plugins: + config["plugins"] = plugins + return config diff --git a/tests/test_envconfig_extensions.py b/tests/test_envconfig_extensions.py new file mode 100644 index 000000000..5d376433e --- /dev/null +++ b/tests/test_envconfig_extensions.py @@ -0,0 +1,417 @@ +import dataclasses +import math +import os +import textwrap +from collections.abc import Awaitable, Callable +from datetime import date, datetime, time, timedelta, timezone +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +import temporalio.client +from temporalio.envconfig import ( + ClientConfig, + ClientConfigExtension, + ClientConfigExtensionRequest, + ClientConfigProfile, +) +from temporalio.service import ConnectConfig, ServiceClient + + +class _TestPlugin(temporalio.client.Plugin): + def __init__(self, name: str) -> None: + self._name = name + + def name(self) -> str: + return self._name + + def configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + return config + + async def connect_service_client( + self, + config: ConnectConfig, + next: Callable[[ConnectConfig], Awaitable[ServiceClient]], + ) -> ServiceClient: + return await next(config) + + +class _TestExtension(ClientConfigExtension): + def __init__( + self, + name: str, + client_plugin: Callable[ + [ClientConfigExtensionRequest], temporalio.client.Plugin | None + ], + ) -> None: + self._name = name + self._client_plugin = client_plugin + + def name(self) -> str: + return self._name + + def client_plugin( + self, req: ClientConfigExtensionRequest + ) -> temporalio.client.Plugin | None: + return self._client_plugin(req) + + +def _write_config(tmp_path: Path, content: str) -> str: + config_file = tmp_path / "config.toml" + config_file.write_text(textwrap.dedent(content)) + return str(config_file) + + +TOML_CONFIG_EXT = textwrap.dedent( + """ + [profile.default] + address = "default-address" + [profile.default.ext.first] + value = "first-value" + nested = { list = [1, 2], flag = true } + [profile.default.ext.second] + [profile.default.ext.second.deeper] + value = 5 + + [profile.other] + address = "other-address" + """ +) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_load_profile_with_ext(strict: bool): + profile = ClientConfigProfile.load( + config_source=TOML_CONFIG_EXT, + config_file_strict=strict, + override_env_vars={}, + ) + assert profile.address == "default-address" + assert profile.ext == { + "first": {"value": "first-value", "nested": {"list": [1, 2], "flag": True}}, + "second": {"deeper": {"value": 5}}, + } + + config = ClientConfig.load(config_source=TOML_CONFIG_EXT, config_file_strict=strict) + assert config.profiles["default"].ext == profile.ext + assert config.profiles["other"].ext == {} + + +def test_load_profile_ext_value_conversion(): + profile = ClientConfigProfile.load( + config_source=textwrap.dedent( + """ + [profile.default.ext.my_ext] + str = "value" + int = 5 + float = 1.5 + nan = nan + inf = -inf + bool = true + datetime = 1979-05-27T07:32:00Z + offset_datetime = 1979-05-27T00:32:00.999999999-07:00 + local_datetime = 1979-05-27T07:32:00 + local_date = 1979-05-27 + local_time = 07:32:00.5 + array = [1, "two", 1979-05-27] + nested = { key = "nested_value" } + """ + ), + disable_env=True, + ) + ext = dict(profile.ext["my_ext"]) + assert math.isnan(ext.pop("nan")) + assert ext == { + "str": "value", + "int": 5, + "float": 1.5, + "inf": -math.inf, + "bool": True, + "datetime": datetime(1979, 5, 27, 7, 32, tzinfo=timezone.utc), + "offset_datetime": datetime( + 1979, 5, 27, 0, 32, 0, 999999, tzinfo=timezone(timedelta(hours=-7)) + ), + "local_datetime": datetime(1979, 5, 27, 7, 32), + "local_date": date(1979, 5, 27), + "local_time": time(7, 32, 0, 500000), + "array": [1, "two", date(1979, 5, 27)], + "nested": {"key": "nested_value"}, + } + + +def test_strict_mode_still_rejects_unknown_non_ext_keys(): + with pytest.raises(RuntimeError, match="unrecognized_field"): + ClientConfigProfile.load( + config_source=textwrap.dedent( + """ + [profile.default] + unrecognized_field = "should-fail" + [profile.default.ext.first] + value = "first-value" + """ + ), + config_file_strict=True, + ) + + +def test_ext_not_modified_by_env(): + profile = ClientConfigProfile.load( + config_source='[profile.default.ext.foo]\nsome_field = "value"', + override_env_vars={"TEMPORAL_EXT_FOO_SOME_FIELD": "env-value"}, + ) + assert profile.ext == {"foo": {"some_field": "value"}} + + +def test_profile_ext_to_from_dict(): + profile = ClientConfigProfile( + address="some-address", ext={"foo": {"some_field": "value"}} + ) + profile_dict = profile.to_dict() + assert profile_dict == { + "address": "some-address", + "ext": {"foo": {"some_field": "value"}}, + } + assert ClientConfigProfile.from_dict(profile_dict) == profile + + +def test_load_client_connect_config_calls_extensions_in_order(tmp_path: Path): + env = {"TEMPORAL_ADDRESS": "env-address", "TEMPORAL_EXT_FIRST_VALUE": "env-value"} + calls: list[str] = [] + + def first_plugin(req: ClientConfigExtensionRequest) -> temporalio.client.Plugin: + calls.append("first") + # Env vars have already been applied to the profile + assert req.profile.address == "env-address" + assert req.profile.ext.get("first") == {"value": "file-value"} + assert req.env is env + assert req.env is not None + return _TestPlugin(f"first-{req.env['TEMPORAL_EXT_FIRST_VALUE']}") + + def noop_plugin(req: ClientConfigExtensionRequest) -> None: + calls.append("noop") + assert "noop" not in req.profile.ext + return None + + def second_plugin(_: ClientConfigExtensionRequest) -> temporalio.client.Plugin: + calls.append("second") + return _TestPlugin("second") + + config = ClientConfig.load_client_connect_config( + config_file=_write_config( + tmp_path, + """ + [profile.default] + address = "file-address" + [profile.default.ext.first] + value = "file-value" + """, + ), + config_file_strict=True, + override_env_vars=env, + extensions=[ + _TestExtension("second", second_plugin), + _TestExtension("first", first_plugin), + _TestExtension("noop", noop_plugin), + ], + ) + # Extension order is respected + assert calls == ["second", "first", "noop"] + assert [p.name() for p in config.get("plugins", [])] == [ + "second", + "first-env-value", + ] + assert config.get("target_host") == "env-address" + + +def test_load_client_connect_config_no_plugins_without_extensions(tmp_path: Path): + config_file = _write_config(tmp_path, TOML_CONFIG_EXT) + config = ClientConfig.load_client_connect_config( + config_file=config_file, override_env_vars={} + ) + assert "plugins" not in config + config = ClientConfig.load_client_connect_config( + config_file=config_file, + override_env_vars={}, + extensions=[_TestExtension("first", lambda _: None)], + ) + assert "plugins" not in config + + +def test_load_client_connect_config_passes_env_to_extensions(tmp_path: Path): + config_file = _write_config(tmp_path, "[profile.default]") + got_envs: list[object] = [] + + def record_env(req: ClientConfigExtensionRequest) -> None: + got_envs.append(req.env) + return None + + extensions = [_TestExtension("ext", record_env)] + env = {"TEMPORAL_EXT_EXT_VALUE": "value"} + + # Defaults to process env + ClientConfig.load_client_connect_config( + config_file=config_file, extensions=extensions + ) + # Uses override env when given + ClientConfig.load_client_connect_config( + config_file=config_file, override_env_vars=env, extensions=extensions + ) + # None when env is disabled + ClientConfig.load_client_connect_config( + config_file=config_file, + disable_env=True, + override_env_vars=env, + extensions=extensions, + ) + + assert len(got_envs) == 3 + assert got_envs[0] is os.environ + assert got_envs[1] is env + assert got_envs[2] is None + + +@pytest.mark.parametrize( + "names,expected_message", + [ + (["a", "b", "a"], 'duplicate client config extension name "a"'), + (["a", ""], "client config extension at index 1 has empty name"), + ([None], "client config extension at index 0 is None"), + ], +) +def test_load_client_connect_config_rejects_invalid_extensions( + tmp_path: Path, names: list[str | None], expected_message: str +): + def fail(_: ClientConfigExtensionRequest) -> None: + pytest.fail("extension should not be called") + + extensions = [ + None if name is None else _TestExtension(name, fail) for name in names + ] + with pytest.raises(ValueError) as err: + ClientConfig.load_client_connect_config( + config_file=_write_config(tmp_path, "[profile.default]"), + override_env_vars={}, + extensions=extensions, # type: ignore[arg-type] + ) + assert str(err.value) == expected_message + + +def test_load_client_connect_config_wraps_extension_errors(tmp_path: Path): + ext_err = ValueError("bad config") + + def raise_err(_: ClientConfigExtensionRequest) -> None: + raise ext_err + + with pytest.raises(RuntimeError) as err: + ClientConfig.load_client_connect_config( + config_file=_write_config(tmp_path, "[profile.default]"), + override_env_vars={}, + extensions=[_TestExtension("bad", raise_err)], + ) + assert str(err.value) == 'client config extension "bad": bad config' + assert err.value.__cause__ is ext_err + + +class ApiKeyFileExtension(ClientConfigExtension): + """Example extension that reads an API key from a file, configured like: + + .. code-block:: toml + + [profile.default.ext.api_key_file] + path = "/path/to/api-key" + + or with the ``TEMPORAL_EXT_API_KEY_FILE_PATH`` env var, which takes + precedence. + """ + + def name(self) -> str: + return "api_key_file" + + def client_plugin( + self, req: ClientConfigExtensionRequest + ) -> temporalio.client.Plugin | None: + key_path = req.profile.ext.get(self.name(), {}).get("path") + if key_path is not None and not isinstance(key_path, str): + raise TypeError(f"path must be a string, got {type(key_path).__name__}") + if req.env is not None: + key_path = req.env.get("TEMPORAL_EXT_API_KEY_FILE_PATH", key_path) + if key_path is None: + return None + return _ApiKeyPlugin(Path(key_path).read_text().strip()) + + +class _ApiKeyPlugin(temporalio.client.Plugin): + def __init__(self, api_key: str) -> None: + self._api_key = api_key + + def configure_client( + self, config: temporalio.client.ClientConfig + ) -> temporalio.client.ClientConfig: + return config + + async def connect_service_client( + self, + config: ConnectConfig, + next: Callable[[ConnectConfig], Awaitable[ServiceClient]], + ) -> ServiceClient: + return await next( + dataclasses.replace( + config, + api_key=self._api_key, + tls=True if config.tls is None else config.tls, + ) + ) + + +async def test_example_extension(tmp_path: Path): + key_file = tmp_path / "api-key" + key_file.write_text("my-api-key\n") + config = ClientConfig.load_client_connect_config( + config_file=_write_config( + tmp_path, + f""" + [profile.default] + address = "my-address" + [profile.default.ext.api_key_file] + path = {str(key_file)!r} + """, + ), + config_file_strict=True, + override_env_vars={}, + extensions=[ApiKeyFileExtension()], + ) + plugins = config.get("plugins", []) + assert len(plugins) == 1 + next = AsyncMock() + await plugins[0].connect_service_client( + ConnectConfig(target_host="my-address"), next + ) + connect_config = next.call_args.args[0] + assert connect_config.target_host == "my-address" + assert connect_config.api_key == "my-api-key" + assert connect_config.tls is True + + # Nothing configured + unconfigured = ClientConfig.load_client_connect_config( + config_file=_write_config(tmp_path, "[profile.default]"), + override_env_vars={}, + extensions=[ApiKeyFileExtension()], + ) + assert "plugins" not in unconfigured + + # Invalid config + with pytest.raises(RuntimeError) as err: + ClientConfig.load_client_connect_config( + config_file=_write_config( + tmp_path, "[profile.default.ext.api_key_file]\npath = 5" + ), + override_env_vars={}, + extensions=[ApiKeyFileExtension()], + ) + assert ( + str(err.value) + == 'client config extension "api_key_file": path must be a string, got int' + )