Skip to main content

toolkit/
context.rs

1use serde::de::DeserializeOwned;
2use std::sync::Arc;
3use tokio_util::sync::CancellationToken;
4use uuid::Uuid;
5
6// Import configuration types from the config gear
7use crate::{
8    config::{ConfigError, ConfigProvider, gear_config_or_default},
9    gear_config_required,
10};
11
12#[cfg(feature = "db")]
13pub(crate) type DbManager = toolkit_db::DbManager;
14#[cfg(feature = "db")]
15pub(crate) type DbProvider = toolkit_db::DBProvider<toolkit_db::DbError>;
16
17#[derive(Clone)]
18#[must_use]
19pub struct GearCtx {
20    gear_name: Arc<str>,
21    instance_id: Uuid,
22    config_provider: Arc<dyn ConfigProvider>,
23    client_hub: Arc<crate::client_hub::ClientHub>,
24    cancellation_token: CancellationToken,
25    #[cfg(feature = "db")]
26    db: Option<DbProvider>,
27}
28
29/// Builder for creating gear-scoped contexts with resolved database handles.
30///
31/// Use [`GearContextBuilder::with_db_manager`] (feature `db`) to attach a
32/// `DbManager` so `for_gear` can resolve a per-gear `DbHandle`.
33#[must_use]
34pub struct GearContextBuilder {
35    instance_id: Uuid,
36    config_provider: Arc<dyn ConfigProvider>,
37    client_hub: Arc<crate::client_hub::ClientHub>,
38    root_token: CancellationToken,
39    #[cfg(feature = "db")]
40    db_manager: Option<Arc<DbManager>>, // internal only, never exposed to gears
41}
42
43impl GearContextBuilder {
44    pub fn new(
45        instance_id: Uuid,
46        config_provider: Arc<dyn ConfigProvider>,
47        client_hub: Arc<crate::client_hub::ClientHub>,
48        root_token: CancellationToken,
49    ) -> Self {
50        Self {
51            instance_id,
52            config_provider,
53            client_hub,
54            root_token,
55            #[cfg(feature = "db")]
56            db_manager: None,
57        }
58    }
59
60    /// Attach a `DbManager` used by [`for_gear`](Self::for_gear) to resolve
61    /// per-gear database handles.
62    #[cfg(feature = "db")]
63    pub fn with_db_manager(mut self, db_manager: Arc<DbManager>) -> Self {
64        self.db_manager = Some(db_manager);
65        self
66    }
67
68    /// Returns the process-level instance ID.
69    #[must_use]
70    pub fn instance_id(&self) -> Uuid {
71        self.instance_id
72    }
73
74    /// Build a gear-scoped context, resolving the `DbHandle` for the given
75    /// gear when the `db` feature is enabled.
76    ///
77    /// Kept `async` in both configurations so callers don't need cfg branches
78    /// around `.await`; under `not(feature = "db")` the future is ready on
79    /// first poll.
80    ///
81    /// # Errors
82    /// Returns an error if database resolution fails.
83    #[cfg_attr(not(feature = "db"), allow(clippy::unused_async))]
84    pub async fn for_gear(&self, gear_name: &str) -> anyhow::Result<GearCtx> {
85        let ctx = GearCtx::new(
86            Arc::<str>::from(gear_name),
87            self.instance_id,
88            self.config_provider.clone(),
89            self.client_hub.clone(),
90            self.root_token.child_token(),
91        );
92        #[cfg(feature = "db")]
93        let ctx = if let Some(mgr) = &self.db_manager
94            && let Some(handle) = mgr.get(gear_name).await?
95        {
96            ctx.with_db(toolkit_db::DBProvider::new(handle))
97        } else {
98            ctx
99        };
100        Ok(ctx)
101    }
102}
103
104impl GearCtx {
105    /// Create a new gear-scoped context with all required fields.
106    ///
107    /// Attach a database entrypoint with [`with_db`](Self::with_db) (feature `db`).
108    pub fn new(
109        gear_name: impl Into<Arc<str>>,
110        instance_id: Uuid,
111        config_provider: Arc<dyn ConfigProvider>,
112        client_hub: Arc<crate::client_hub::ClientHub>,
113        cancellation_token: CancellationToken,
114    ) -> Self {
115        Self {
116            gear_name: gear_name.into(),
117            instance_id,
118            config_provider,
119            client_hub,
120            cancellation_token,
121            #[cfg(feature = "db")]
122            db: None,
123        }
124    }
125
126    /// Attach the per-gear database entrypoint.
127    #[cfg(feature = "db")]
128    pub fn with_db(mut self, db: DbProvider) -> Self {
129        self.db = Some(db);
130        self
131    }
132
133    // ---- public read-only API for gears ----
134
135    #[inline]
136    #[must_use]
137    pub fn gear_name(&self) -> &str {
138        &self.gear_name
139    }
140
141    /// Returns the process-level instance ID.
142    ///
143    /// This is a unique identifier for this process instance, shared by all gears
144    /// in the same process. It is generated once at bootstrap.
145    #[inline]
146    #[must_use]
147    pub fn instance_id(&self) -> Uuid {
148        self.instance_id
149    }
150
151    #[inline]
152    #[must_use]
153    pub fn config_provider(&self) -> &dyn ConfigProvider {
154        &*self.config_provider
155    }
156
157    /// Get the `ClientHub` for dependency resolution.
158    #[inline]
159    #[must_use]
160    pub fn client_hub(&self) -> Arc<crate::client_hub::ClientHub> {
161        self.client_hub.clone()
162    }
163
164    #[inline]
165    #[must_use]
166    pub fn cancellation_token(&self) -> &CancellationToken {
167        &self.cancellation_token
168    }
169
170    /// Get a gear-scoped DB entrypoint for secure database operations.
171    ///
172    /// Returns `None` if no database is configured for this gear.
173    ///
174    /// # Security
175    ///
176    /// The returned `DBProvider<toolkit_db::DbError>`:
177    /// - Is cheap to clone (shares an internal `Db`)
178    /// - Provides `conn()` for non-transactional access (fails inside tx via guard)
179    /// - Provides `transaction(..)` for transactional operations
180    ///
181    /// # Example
182    ///
183    /// ```ignore
184    /// let db = ctx.db().ok_or_else(|| anyhow!("no db"))?;
185    /// let conn = db.conn()?;
186    /// let user = svc.get_user(&conn, &scope, id).await?;
187    /// ```
188    #[must_use]
189    #[cfg(feature = "db")]
190    pub fn db(&self) -> Option<toolkit_db::DBProvider<toolkit_db::DbError>> {
191        self.db.clone()
192    }
193
194    /// Get a database handle, returning an error if not configured.
195    ///
196    /// This is a convenience method that combines `db()` with an error for
197    /// gears that require database access.
198    ///
199    /// # Errors
200    ///
201    /// Returns an error if the database is not configured for this gear.
202    ///
203    /// # Example
204    ///
205    /// ```ignore
206    /// let db = ctx.db_required()?;
207    /// let conn = db.conn()?;
208    /// let user = svc.get_user(&conn, &scope, id).await?;
209    /// ```
210    #[cfg(feature = "db")]
211    pub fn db_required(&self) -> anyhow::Result<toolkit_db::DBProvider<toolkit_db::DbError>> {
212        self.db().ok_or_else(|| {
213            anyhow::anyhow!("Database is not configured for gear '{}'", self.gear_name)
214        })
215    }
216
217    /// Deserialize the gear's config section into `T`.
218    ///
219    /// This reads the `gears.<name>.config` object for the current gear and
220    /// deserializes it into the requested type.
221    ///
222    /// # Errors
223    /// Returns `ConfigError` if the gear config is missing or deserialization fails.
224    pub fn config<T: DeserializeOwned>(&self) -> Result<T, ConfigError> {
225        gear_config_required(self.config_provider.as_ref(), &self.gear_name)
226    }
227
228    /// Deserialize the gear's config section into T, or use defaults if missing.
229    ///
230    /// This method uses lenient configuration loading: if the gear is not present in config,
231    /// has no config section, or the gear entry is not an object, it returns `T::default()`.
232    /// This allows gears to exist without configuration sections in the main config file.
233    ///
234    /// It extracts the 'config' field from: `gears.<name> = { database: ..., config: ... }`
235    ///
236    /// # Example
237    ///
238    /// ```rust,ignore
239    /// #[derive(serde::Deserialize, Default)]
240    /// struct MyConfig {
241    ///     api_key: String,
242    ///     timeout_ms: u64,
243    /// }
244    ///
245    /// let config: MyConfig = ctx.config_or_default()?;
246    /// ```
247    ///
248    /// # Errors
249    /// Returns `ConfigError` if deserialization fails.
250    pub fn config_or_default<T: DeserializeOwned + Default>(&self) -> Result<T, ConfigError> {
251        gear_config_or_default(self.config_provider.as_ref(), &self.gear_name)
252    }
253
254    /// Like [`config()`](Self::config), but additionally expands `${VAR}` placeholders
255    /// in fields marked with `#[expand_vars]`.
256    ///
257    /// # Errors
258    /// Returns `ConfigError` if the gear config is missing, deserialization fails,
259    /// or environment variable expansion fails.
260    pub fn config_expanded<T>(&self) -> Result<T, ConfigError>
261    where
262        T: DeserializeOwned + crate::var_expand::ExpandVars,
263    {
264        let mut cfg: T = self.config()?;
265        cfg.expand_vars().map_err(|e| ConfigError::VarExpand {
266            gear: self.gear_name.to_string(),
267            source: e,
268        })?;
269        Ok(cfg)
270    }
271
272    /// Like [`config_or_default()`](Self::config_or_default), but additionally expands `${VAR}`
273    /// placeholders
274    /// in fields marked with `#[expand_vars]` (requires `#[derive(ExpandVars)]` on the config
275    /// struct).
276    ///
277    /// Gears that do not need environment variable expansion should use
278    /// [`config_or_default()`](Self::config_or_default).
279    ///
280    /// # Example
281    ///
282    /// ```rust,ignore
283    /// #[derive(serde::Deserialize, Default, ExpandVars)]
284    /// struct MyConfig {
285    ///     #[expand_vars]
286    ///     api_key: String,
287    ///     timeout_ms: u64,
288    /// }
289    ///
290    /// let config: MyConfig = ctx.config_expanded_or_default()?;
291    /// ```
292    ///
293    /// # Errors
294    /// Returns `ConfigError` if deserialization fails or if environment variable expansion fails.
295    pub fn config_expanded_or_default<T>(&self) -> Result<T, ConfigError>
296    where
297        T: DeserializeOwned + Default + crate::var_expand::ExpandVars,
298    {
299        let mut cfg: T = self.config_or_default()?;
300        cfg.expand_vars().map_err(|e| ConfigError::VarExpand {
301            gear: self.gear_name.to_string(),
302            source: e,
303        })?;
304        Ok(cfg)
305    }
306
307    /// Get the raw JSON value of the gear's config section.
308    /// Returns the 'config' field from: gears.<name> = { database: ..., config: ... }
309    #[must_use]
310    pub fn raw_config(&self) -> &serde_json::Value {
311        use std::sync::LazyLock;
312
313        static EMPTY: LazyLock<serde_json::Value> =
314            LazyLock::new(|| serde_json::Value::Object(serde_json::Map::new()));
315
316        if let Some(gear_raw) = self.config_provider.get_gear_config(&self.gear_name) {
317            // Try new structure first: gears.<name> = { database: ..., config: ... }
318            if let Some(obj) = gear_raw.as_object()
319                && let Some(config_section) = obj.get("config")
320            {
321                return config_section;
322            }
323        }
324        &EMPTY
325    }
326
327    /// Create a derivative context with the same references but no DB handle.
328    /// Useful for gears that don't require database access.
329    pub fn without_db(&self) -> GearCtx {
330        GearCtx {
331            gear_name: self.gear_name.clone(),
332            instance_id: self.instance_id,
333            config_provider: self.config_provider.clone(),
334            client_hub: self.client_hub.clone(),
335            cancellation_token: self.cancellation_token.clone(),
336            #[cfg(feature = "db")]
337            db: None,
338        }
339    }
340}
341
342#[cfg(test)]
343#[cfg_attr(coverage_nightly, coverage(off))]
344mod tests {
345    use super::*;
346    use serde::Deserialize;
347    use serde_json::json;
348    use std::collections::HashMap;
349
350    #[derive(Debug, PartialEq, Deserialize, Default)]
351    struct TestConfig {
352        #[serde(default)]
353        api_key: String,
354        #[serde(default)]
355        timeout_ms: u64,
356        #[serde(default)]
357        enabled: bool,
358    }
359
360    struct MockConfigProvider {
361        gears: HashMap<String, serde_json::Value>,
362    }
363
364    impl MockConfigProvider {
365        fn new() -> Self {
366            let mut gears = HashMap::new();
367
368            // Valid gear config
369            gears.insert(
370                "test_gear".to_owned(),
371                json!({
372                    "database": {
373                        "url": "postgres://localhost/test"
374                    },
375                    "config": {
376                        "api_key": "secret123",
377                        "timeout_ms": 5000,
378                        "enabled": true
379                    }
380                }),
381            );
382
383            Self { gears }
384        }
385    }
386
387    impl ConfigProvider for MockConfigProvider {
388        fn get_gear_config(&self, gear_name: &str) -> Option<&serde_json::Value> {
389            self.gears.get(gear_name)
390        }
391    }
392
393    #[test]
394    fn test_gear_ctx_config_with_valid_config() {
395        let provider = Arc::new(MockConfigProvider::new());
396        let ctx = GearCtx::new(
397            "test_gear",
398            Uuid::new_v4(),
399            provider,
400            Arc::new(crate::client_hub::ClientHub::default()),
401            CancellationToken::new(),
402        );
403
404        let result: Result<TestConfig, ConfigError> = ctx.config();
405        assert!(result.is_ok());
406
407        let config = result.unwrap();
408        assert_eq!(config.api_key, "secret123");
409        assert_eq!(config.timeout_ms, 5000);
410        assert!(config.enabled);
411    }
412
413    #[test]
414    fn test_gear_ctx_config_returns_error_for_missing_gear() {
415        let provider = Arc::new(MockConfigProvider::new());
416        let ctx = GearCtx::new(
417            "nonexistent_gear",
418            Uuid::new_v4(),
419            provider,
420            Arc::new(crate::client_hub::ClientHub::default()),
421            CancellationToken::new(),
422        );
423
424        let result: Result<TestConfig, ConfigError> = ctx.config();
425        assert!(matches!(
426            result,
427            Err(ConfigError::GearNotFound { ref gear }) if gear == "nonexistent_gear"
428        ));
429    }
430
431    #[test]
432    fn test_gear_ctx_config_or_default_returns_default_for_missing_gear() {
433        let provider = Arc::new(MockConfigProvider::new());
434        let ctx = GearCtx::new(
435            "nonexistent_gear",
436            Uuid::new_v4(),
437            provider,
438            Arc::new(crate::client_hub::ClientHub::default()),
439            CancellationToken::new(),
440        );
441
442        let result: Result<TestConfig, ConfigError> = ctx.config_or_default();
443        assert!(result.is_ok());
444
445        let config = result.unwrap();
446        assert_eq!(config, TestConfig::default());
447    }
448
449    #[test]
450    fn test_gear_ctx_instance_id() {
451        let provider = Arc::new(MockConfigProvider::new());
452        let instance_id = Uuid::new_v4();
453        let ctx = GearCtx::new(
454            "test_gear",
455            instance_id,
456            provider,
457            Arc::new(crate::client_hub::ClientHub::default()),
458            CancellationToken::new(),
459        );
460
461        assert_eq!(ctx.instance_id(), instance_id);
462    }
463
464    // --- config_expanded tests ---
465
466    #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
467    struct ExpandableConfig {
468        #[expand_vars]
469        #[serde(default)]
470        api_key: String,
471        #[expand_vars]
472        #[serde(default)]
473        endpoint: Option<String>,
474        #[serde(default)]
475        retries: u32,
476    }
477
478    fn make_ctx(gear_name: &str, config_json: serde_json::Value) -> GearCtx {
479        let mut gears = HashMap::new();
480        gears.insert(gear_name.to_owned(), config_json);
481        let provider = Arc::new(MockConfigProvider { gears });
482        GearCtx::new(
483            gear_name,
484            Uuid::new_v4(),
485            provider,
486            Arc::new(crate::client_hub::ClientHub::default()),
487            CancellationToken::new(),
488        )
489    }
490
491    #[test]
492    fn config_expanded_resolves_env_vars() {
493        let ctx = make_ctx(
494            "expand_mod",
495            json!({
496                "config": {
497                    "api_key": "${TOOLKIT_TEST_KEY}",
498                    "endpoint": "https://${TOOLKIT_TEST_HOST}/api",
499                    "retries": 3
500                }
501            }),
502        );
503
504        temp_env::with_vars(
505            [
506                ("TOOLKIT_TEST_KEY", Some("secret-42")),
507                ("TOOLKIT_TEST_HOST", Some("example.com")),
508            ],
509            || {
510                let cfg: ExpandableConfig = ctx.config_expanded().unwrap();
511                assert_eq!(cfg.api_key, "secret-42");
512                assert_eq!(cfg.endpoint.as_deref(), Some("https://example.com/api"));
513                assert_eq!(cfg.retries, 3);
514            },
515        );
516    }
517
518    #[test]
519    fn config_expanded_returns_error_on_missing_var() {
520        let ctx = make_ctx(
521            "expand_mod",
522            json!({
523                "config": {
524                    "api_key": "${TOOLKIT_TEST_MISSING_VAR_XYZ}"
525                }
526            }),
527        );
528
529        temp_env::with_vars([("TOOLKIT_TEST_MISSING_VAR_XYZ", None::<&str>)], || {
530            let err = ctx.config_expanded::<ExpandableConfig>().unwrap_err();
531            assert!(
532                matches!(err, ConfigError::VarExpand { ref gear, .. } if gear == "expand_mod"),
533                "expected EnvExpand error, got: {err:?}"
534            );
535        });
536    }
537
538    #[test]
539    fn config_expanded_skips_none_option_fields() {
540        let ctx = make_ctx(
541            "expand_mod",
542            json!({
543                "config": {
544                    "api_key": "literal-key",
545                    "retries": 5
546                }
547            }),
548        );
549
550        let cfg: ExpandableConfig = ctx.config_expanded().unwrap();
551        assert_eq!(cfg.api_key, "literal-key");
552        assert_eq!(cfg.endpoint, None);
553        assert_eq!(cfg.retries, 5);
554    }
555
556    #[test]
557    fn config_expanded_returns_error_when_missing() {
558        let ctx = make_ctx("missing_mod", json!({}));
559        let err = ctx.config_expanded::<ExpandableConfig>().unwrap_err();
560        assert!(matches!(
561            err,
562            ConfigError::MissingConfigSection { ref gear } if gear == "missing_mod"
563        ));
564    }
565
566    #[test]
567    fn config_expanded_or_default_falls_back_to_default_when_missing() {
568        let ctx = make_ctx("missing_mod", json!({}));
569        let cfg: ExpandableConfig = ctx.config_expanded_or_default().unwrap();
570        assert_eq!(cfg, ExpandableConfig::default());
571    }
572
573    // --- nested struct expansion ---
574
575    #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
576    struct NestedProvider {
577        #[expand_vars]
578        #[serde(default)]
579        host: String,
580        #[expand_vars]
581        #[serde(default)]
582        token: Option<String>,
583        #[expand_vars]
584        #[serde(default)]
585        auth_config: Option<HashMap<String, String>>,
586        #[serde(default)]
587        port: u16,
588    }
589
590    #[derive(Debug, PartialEq, Deserialize, Default, toolkit_macros::ExpandVars)]
591    struct NestedConfig {
592        #[expand_vars]
593        #[serde(default)]
594        name: String,
595        #[expand_vars]
596        #[serde(default)]
597        providers: HashMap<String, NestedProvider>,
598        #[expand_vars]
599        #[serde(default)]
600        tags: Vec<String>,
601    }
602
603    #[test]
604    fn config_expanded_resolves_nested_structs() {
605        let ctx = make_ctx(
606            "nested_mod",
607            json!({
608                "config": {
609                    "name": "${TOOLKIT_NESTED_NAME}",
610                    "providers": {
611                        "primary": {
612                            "host": "${TOOLKIT_NESTED_HOST}",
613                            "token": "${TOOLKIT_NESTED_TOKEN}",
614                            "auth_config": {
615                                "header": "X-Api-Key",
616                                "secret_ref": "${TOOLKIT_NESTED_SECRET}"
617                            },
618                            "port": 443
619                        }
620                    },
621                    "tags": ["${TOOLKIT_NESTED_TAG}", "literal"]
622                }
623            }),
624        );
625
626        temp_env::with_vars(
627            [
628                ("TOOLKIT_NESTED_NAME", Some("my-service")),
629                ("TOOLKIT_NESTED_HOST", Some("api.example.com")),
630                ("TOOLKIT_NESTED_TOKEN", Some("sk-secret")),
631                ("TOOLKIT_NESTED_SECRET", Some("key-12345")),
632                ("TOOLKIT_NESTED_TAG", Some("production")),
633            ],
634            || {
635                let cfg: NestedConfig = ctx.config_expanded().unwrap();
636                assert_eq!(cfg.name, "my-service");
637                assert_eq!(cfg.tags, vec!["production", "literal"]);
638
639                let primary = cfg.providers.get("primary").expect("primary provider");
640                assert_eq!(primary.host, "api.example.com");
641                assert_eq!(primary.token.as_deref(), Some("sk-secret"));
642                assert_eq!(primary.port, 443);
643
644                let auth = primary.auth_config.as_ref().expect("auth_config present");
645                assert_eq!(auth.get("header").map(String::as_str), Some("X-Api-Key"));
646                assert_eq!(
647                    auth.get("secret_ref").map(String::as_str),
648                    Some("key-12345")
649                );
650            },
651        );
652    }
653
654    #[test]
655    fn config_expanded_nested_missing_var_returns_error() {
656        let ctx = make_ctx(
657            "nested_mod",
658            json!({
659                "config": {
660                    "name": "ok",
661                    "providers": {
662                        "bad": { "host": "${TOOLKIT_NESTED_GONE}", "port": 80 }
663                    }
664                }
665            }),
666        );
667
668        temp_env::with_vars([("TOOLKIT_NESTED_GONE", None::<&str>)], || {
669            let err = ctx.config_expanded::<NestedConfig>().unwrap_err();
670            assert!(
671                matches!(err, ConfigError::VarExpand { ref gear, .. } if gear == "nested_mod"),
672                "expected EnvExpand, got: {err:?}"
673            );
674        });
675    }
676}