Skip to main content

fraiseql_core/runtime/
jsonb_strategy.rs

1//! JSONB field handling strategy selection
2//!
3//! Determines whether to project fields at database level or stream full JSONB
4
5use serde::Deserialize;
6
7/// Strategy for JSONB field handling
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
9pub enum JsonbStrategy {
10    /// Extract only requested fields using jsonb_build_object/JSON_OBJECT
11    #[default]
12    Project,
13    /// Stream full JSONB column, filter in application
14    Stream,
15}
16
17impl std::str::FromStr for JsonbStrategy {
18    type Err = String;
19
20    fn from_str(s: &str) -> Result<Self, Self::Err> {
21        match s.to_lowercase().as_str() {
22            "project" => Ok(JsonbStrategy::Project),
23            "stream" => Ok(JsonbStrategy::Stream),
24            other => {
25                Err(format!("Invalid JSONB strategy '{}', must be 'project' or 'stream'", other))
26            },
27        }
28    }
29}
30
31impl<'de> Deserialize<'de> for JsonbStrategy {
32    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
33    where
34        D: serde::Deserializer<'de>,
35    {
36        let s = String::deserialize(deserializer)?;
37        std::str::FromStr::from_str(&s).map_err(serde::de::Error::custom)
38    }
39}
40
41/// Configuration for JSONB optimization strategy
42#[derive(Debug, Clone)]
43pub struct JsonbOptimizationOptions {
44    /// Default strategy to use
45    pub default_strategy: JsonbStrategy,
46
47    /// Auto-switch threshold: if requesting >= this % of fields, use stream
48    pub auto_threshold_percent: u32,
49}
50
51impl Default for JsonbOptimizationOptions {
52    fn default() -> Self {
53        Self {
54            default_strategy:       JsonbStrategy::Project,
55            auto_threshold_percent: 80,
56        }
57    }
58}
59
60impl JsonbOptimizationOptions {
61    /// Choose strategy based on field count and configuration
62    pub fn choose_strategy(&self, requested_fields: usize, total_fields: usize) -> JsonbStrategy {
63        if total_fields == 0 {
64            return self.default_strategy;
65        }
66
67        let percent = (requested_fields as f64 / total_fields as f64) * 100.0;
68
69        if percent >= f64::from(self.auto_threshold_percent) {
70            JsonbStrategy::Stream
71        } else {
72            self.default_strategy
73        }
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    // ========================================================================
82    // Strategy Parsing Tests
83    // ========================================================================
84
85    #[test]
86    fn test_jsonb_strategy_from_str_project() {
87        let strategy: JsonbStrategy = "project".parse().unwrap();
88        assert_eq!(strategy, JsonbStrategy::Project);
89    }
90
91    #[test]
92    fn test_jsonb_strategy_from_str_stream() {
93        let strategy: JsonbStrategy = "stream".parse().unwrap();
94        assert_eq!(strategy, JsonbStrategy::Stream);
95    }
96
97    #[test]
98    fn test_jsonb_strategy_from_str_case_insensitive() {
99        assert_eq!("PROJECT".parse::<JsonbStrategy>().unwrap(), JsonbStrategy::Project);
100        assert_eq!("Stream".parse::<JsonbStrategy>().unwrap(), JsonbStrategy::Stream);
101        assert_eq!("pRoJeCt".parse::<JsonbStrategy>().unwrap(), JsonbStrategy::Project);
102    }
103
104    #[test]
105    fn test_jsonb_strategy_from_str_invalid() {
106        let result = "invalid".parse::<JsonbStrategy>();
107        assert!(result.is_err());
108        assert!(result.unwrap_err().contains("Invalid JSONB strategy"));
109    }
110
111    #[test]
112    fn test_jsonb_strategy_deserialize() {
113        let json = r#""project""#;
114        let strategy: JsonbStrategy = serde_json::from_str(json).unwrap();
115        assert_eq!(strategy, JsonbStrategy::Project);
116    }
117
118    #[test]
119    fn test_jsonb_strategy_deserialize_stream() {
120        let json = r#""stream""#;
121        let strategy: JsonbStrategy = serde_json::from_str(json).unwrap();
122        assert_eq!(strategy, JsonbStrategy::Stream);
123    }
124
125    // ========================================================================
126    // Strategy Selection Tests
127    // ========================================================================
128
129    #[test]
130    fn test_choose_strategy_below_threshold() {
131        let opts = JsonbOptimizationOptions {
132            default_strategy:       JsonbStrategy::Project,
133            auto_threshold_percent: 80,
134        };
135
136        let strategy = opts.choose_strategy(5, 10);
137        assert_eq!(strategy, JsonbStrategy::Project);
138    }
139
140    #[test]
141    fn test_choose_strategy_at_threshold() {
142        let opts = JsonbOptimizationOptions {
143            default_strategy:       JsonbStrategy::Project,
144            auto_threshold_percent: 80,
145        };
146
147        let strategy = opts.choose_strategy(8, 10);
148        assert_eq!(strategy, JsonbStrategy::Stream);
149    }
150
151    #[test]
152    fn test_choose_strategy_above_threshold() {
153        let opts = JsonbOptimizationOptions {
154            default_strategy:       JsonbStrategy::Project,
155            auto_threshold_percent: 80,
156        };
157
158        let strategy = opts.choose_strategy(9, 10);
159        assert_eq!(strategy, JsonbStrategy::Stream);
160    }
161
162    #[test]
163    fn test_choose_strategy_respects_default() {
164        let opts = JsonbOptimizationOptions {
165            default_strategy:       JsonbStrategy::Stream,
166            auto_threshold_percent: 80,
167        };
168
169        let strategy = opts.choose_strategy(2, 10);
170        assert_eq!(strategy, JsonbStrategy::Stream);
171    }
172
173    #[test]
174    fn test_choose_strategy_zero_total() {
175        let opts = JsonbOptimizationOptions::default();
176        let strategy = opts.choose_strategy(0, 0);
177        assert_eq!(strategy, JsonbStrategy::Project);
178    }
179}