Skip to main content

mz_deploy/client/
auto_scaling.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Conversions between the `AUTO SCALING STRATEGY` cluster option and the
11//! structured [`AutoScalingStrategy`] policy.
12//!
13//! Three representations meet here:
14//!
15//! - The SQL option in a `CREATE CLUSTER` / `ALTER CLUSTER` statement
16//!   (`ClusterAutoScalingStrategyOptionValue` AST).
17//! - The structured policy (`mz_sql::plan::AutoScalingStrategy`), used for
18//!   drift comparison.
19//! - The `strategy` jsonb column of
20//!   `mz_internal.mz_cluster_auto_scaling_strategies`, whose serde shape
21//!   matches the plan type.
22
23use mz_repr::adt::interval::Interval;
24use mz_sql::plan::{AutoScalingStrategy, OnHydration, TryFromValue};
25use mz_sql_parser::ast::{
26    ClusterAutoScalingStrategyOptionValue, ClusterOption, ClusterOptionName,
27    OnHydrationOptionValue, Raw, Value, WithOptionValue,
28};
29use std::time::Duration;
30
31/// Convert the AST option value into a structured policy. An empty block
32/// `AUTO SCALING STRATEGY = ()` maps to `None` (autoscaling disabled), the
33/// same normalization the server planner applies.
34pub(crate) fn strategy_from_option_value(
35    value: &ClusterAutoScalingStrategyOptionValue,
36) -> Result<Option<AutoScalingStrategy>, String> {
37    let Some(on_hydration) = &value.on_hydration else {
38        return Ok(None);
39    };
40
41    let hydration_size = String::try_from_value(on_hydration.hydration_size.clone())
42        .map_err(|e| format!("invalid HYDRATION SIZE: {}", e))?;
43
44    let linger_duration = on_hydration
45        .linger_duration
46        .clone()
47        .map(Duration::try_from_value)
48        .transpose()
49        .map_err(|e| format!("invalid LINGER DURATION: {}", e))?;
50
51    Ok(Some(AutoScalingStrategy {
52        on_hydration: Some(OnHydration {
53            hydration_size,
54            linger_duration,
55        }),
56    }))
57}
58
59/// Render a structured policy back into a `CREATE CLUSTER` / `ALTER CLUSTER
60/// ... SET` option, e.g.
61/// `AUTO SCALING STRATEGY = (ON HYDRATION (HYDRATION SIZE = '3200cc', LINGER DURATION = '00:10:00'))`.
62pub(crate) fn strategy_to_cluster_option(strategy: &AutoScalingStrategy) -> ClusterOption<Raw> {
63    let on_hydration = strategy
64        .on_hydration
65        .as_ref()
66        .map(|on_hydration| OnHydrationOptionValue {
67            hydration_size: Value::String(on_hydration.hydration_size.clone()),
68            linger_duration: on_hydration.linger_duration.map(|d| {
69                // Every linger duration we see was planned from an interval
70                // (either by the server, for values read back from the
71                // catalog, or by `strategy_from_option_value`, which enforces
72                // interval range), so the conversion back cannot fail.
73                let interval = Interval::from_duration(&d)
74                    .expect("linger duration is convertible back to an interval");
75                Value::String(interval.to_string())
76            }),
77        });
78    ClusterOption {
79        name: ClusterOptionName::AutoScalingStrategy,
80        value: Some(WithOptionValue::ClusterAutoScalingStrategyOptionValue(
81            ClusterAutoScalingStrategyOptionValue { on_hydration },
82        )),
83    }
84}
85
86/// Parse the `strategy` jsonb column of
87/// `mz_internal.mz_cluster_auto_scaling_strategies` (fetched as text). The
88/// column is JSON `null` for a cluster whose policy was just removed while a
89/// burst still lingers, which maps to `None` like a missing row.
90pub(crate) fn strategy_from_catalog_json(
91    json: &str,
92) -> Result<Option<AutoScalingStrategy>, String> {
93    serde_json::from_str::<Option<AutoScalingStrategy>>(json)
94        .map_err(|e| format!("cannot parse autoscaling strategy {:?}: {}", json, e))
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use mz_ore::assert_none;
101    use mz_sql_parser::ast::display::AstDisplay;
102    use mz_sql_parser::parser::parse_statements;
103
104    fn strategy(hydration_size: &str, linger: Option<Duration>) -> AutoScalingStrategy {
105        AutoScalingStrategy {
106            on_hydration: Some(OnHydration {
107                hydration_size: hydration_size.to_string(),
108                linger_duration: linger,
109            }),
110        }
111    }
112
113    /// Extract the AUTO SCALING STRATEGY option value from a CREATE CLUSTER.
114    fn option_value_of(sql: &str) -> Option<ClusterAutoScalingStrategyOptionValue> {
115        let stmts = parse_statements(sql).unwrap();
116        let create = match &stmts[0].ast {
117            mz_sql_parser::ast::Statement::CreateCluster(c) => c.clone(),
118            other => panic!("expected CREATE CLUSTER, got {:?}", other),
119        };
120        create
121            .options
122            .into_iter()
123            .find_map(|opt| match (opt.name, opt.value) {
124                (
125                    ClusterOptionName::AutoScalingStrategy,
126                    Some(WithOptionValue::ClusterAutoScalingStrategyOptionValue(v)),
127                ) => Some(v),
128                _ => None,
129            })
130    }
131
132    #[mz_ore::test]
133    fn test_strategy_from_option_value() {
134        let value = option_value_of(
135            "CREATE CLUSTER c (SIZE = '25cc', AUTO SCALING STRATEGY = \
136             (ON HYDRATION (HYDRATION SIZE = '100cc', LINGER DURATION = '600s')))",
137        )
138        .unwrap();
139        assert_eq!(
140            strategy_from_option_value(&value).unwrap(),
141            Some(strategy("100cc", Some(Duration::from_secs(600))))
142        );
143
144        let value = option_value_of(
145            "CREATE CLUSTER c (SIZE = '25cc', AUTO SCALING STRATEGY = \
146             (ON HYDRATION (HYDRATION SIZE = '100cc')))",
147        )
148        .unwrap();
149        assert_eq!(
150            strategy_from_option_value(&value).unwrap(),
151            Some(strategy("100cc", None))
152        );
153
154        // An empty block disables autoscaling.
155        let value = option_value_of("CREATE CLUSTER c (SIZE = '25cc', AUTO SCALING STRATEGY = ())")
156            .unwrap();
157        assert_none!(strategy_from_option_value(&value).unwrap());
158
159        // A non-interval linger duration is a conversion error.
160        let value = option_value_of(
161            "CREATE CLUSTER c (SIZE = '25cc', AUTO SCALING STRATEGY = \
162             (ON HYDRATION (HYDRATION SIZE = '100cc', LINGER DURATION = 'bogus')))",
163        )
164        .unwrap();
165        assert!(strategy_from_option_value(&value).is_err());
166    }
167
168    #[mz_ore::test]
169    fn test_strategy_option_round_trips_through_parser() {
170        for strategy in [
171            strategy("100cc", Some(Duration::from_secs(600))),
172            strategy("100cc", None),
173        ] {
174            let rendered = strategy_to_cluster_option(&strategy).to_ast_string_simple();
175            let sql = format!("CREATE CLUSTER c (SIZE = '25cc', {})", rendered);
176            let value = option_value_of(&sql)
177                .unwrap_or_else(|| panic!("rendered option did not parse: {}", sql));
178            assert_eq!(
179                strategy_from_option_value(&value).unwrap(),
180                Some(strategy),
181                "round trip through {}",
182                rendered
183            );
184        }
185    }
186
187    #[mz_ore::test]
188    fn test_strategy_from_catalog_json() {
189        // The serde shape of the durable AutoScalingStrategy, as surfaced by
190        // mz_internal.mz_cluster_auto_scaling_strategies.
191        let json = r#"{"on_hydration": {"hydration_size": "100cc", "linger_duration": {"secs": 600, "nanos": 0}}}"#;
192        assert_eq!(
193            strategy_from_catalog_json(json).unwrap(),
194            Some(strategy("100cc", Some(Duration::from_secs(600))))
195        );
196
197        let json = r#"{"on_hydration": {"hydration_size": "100cc", "linger_duration": null}}"#;
198        assert_eq!(
199            strategy_from_catalog_json(json).unwrap(),
200            Some(strategy("100cc", None))
201        );
202
203        // JSON null: a burst-only row for a just-removed policy.
204        assert_none!(strategy_from_catalog_json("null").unwrap());
205
206        assert!(strategy_from_catalog_json("not json").is_err());
207    }
208}