feat(envconfig): add extension support

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.<profile>.ext.<name>` 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`.

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<Py
     if !profile.grpc_meta.is_empty() {
         dict.set_item("grpc_meta", &profile.grpc_meta)?;
     }
+    if !profile.ext.is_empty() {
+        let ext = PyDict::new(py);
+        for (name, table) in &profile.ext {
+            ext.set_item(name, toml_table_to_py(py, table)?)?;
+        }
+        dict.set_item("ext", ext)?;
+    }
+    Ok(dict.into())
+}
+
+fn toml_table_to_py(py: Python, table: &toml::Table) -> PyResult<Py<PyAny>> {
+    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<Py<PyAny>> {
+    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<Py<PyAny>> {
+    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<Py<PyAny>> {
     let profiles_dict = PyDict::new(py);
     for (name, profile) in &core_config.profiles {

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.<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.<profile>.ext.<name>``,
+    where ``<name>`` is :py:meth:`name`, and turns its config into a client
+    plugin. Extensions are recommended to use a ``TEMPORAL_EXT_<NAME>_`` 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.<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