Skip to main content

mz_deploy/project/
clusters.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//! Cluster definition loading and validation.
11//!
12//! Loads cluster definitions from `<root>/clusters/` directory. Each `.sql` file
13//! defines a single cluster with a required `CREATE CLUSTER` statement and optional
14//! `GRANT` and `COMMENT` statements.
15
16use crate::client::auto_scaling::strategy_from_option_value;
17use crate::project::error::{
18    LoadError, ProjectError, ValidationError, ValidationErrorKind, ValidationErrors,
19};
20use crate::project::syntax::parser::{
21    LocatedStatement, parse_statements_with_context, statement_type_name,
22};
23use crate::project::syntax::profile_files::collect_all_sql_files;
24use mz_sql::plan::AutoScalingStrategy;
25use mz_sql_parser::ast::{
26    ClusterOptionName, CommentObjectType, CommentStatement, CreateClusterStatement,
27    GrantPrivilegesStatement, GrantTargetSpecification, GrantTargetSpecificationInner, Ident,
28    ObjectType, Raw, RawClusterName, Statement, UnresolvedObjectName, WithOptionValue,
29};
30use std::collections::BTreeMap;
31use std::path::Path;
32
33/// A parsed cluster definition from a `.sql` file in the `clusters/` directory.
34pub(crate) struct ClusterDefinition {
35    /// Cluster name (derived from filename and validated against CREATE statement).
36    pub name: String,
37    /// The CREATE CLUSTER statement.
38    pub create_stmt: CreateClusterStatement<Raw>,
39    /// Optional GRANT statements targeting this cluster.
40    pub grants: Vec<GrantPrivilegesStatement<Raw>>,
41    /// Optional COMMENT statements targeting this cluster.
42    pub comments: Vec<CommentStatement<Raw>>,
43}
44
45/// Load all cluster definitions from `<root>/clusters/`.
46///
47/// Returns an empty vec if `clusters/` doesn't exist (the directory is optional).
48/// If `profile_suffix` is provided, each cluster definition is rewritten with the
49/// suffix appended to all cluster name references (CREATE, GRANT, COMMENT).
50pub(crate) fn load_clusters(
51    root: &Path,
52    profile: &str,
53    profile_suffix: Option<&str>,
54    variables: &BTreeMap<String, String>,
55) -> Result<Vec<ClusterDefinition>, ProjectError> {
56    let clusters_dir = root.join("clusters");
57
58    if !clusters_dir.exists() {
59        return Ok(vec![]);
60    }
61
62    if !clusters_dir.is_dir() {
63        return Err(LoadError::RootNotDirectory { path: clusters_dir }.into());
64    }
65
66    let all_files = collect_all_sql_files(&clusters_dir)?;
67
68    let mut definitions = Vec::new();
69    let mut errors = Vec::new();
70
71    for object_files in all_files {
72        let expected_name = &object_files.name;
73
74        // Validate all variants independently
75        let mut all_variant_paths = Vec::new();
76        if let Some(ref default_path) = object_files.default {
77            all_variant_paths.push((default_path.clone(), None));
78        }
79        for (prof, override_path) in &object_files.overrides {
80            all_variant_paths.push((override_path.clone(), Some(prof.as_str())));
81        }
82
83        // Validate each variant independently
84        for (path, _) in &all_variant_paths {
85            let sql = std::fs::read_to_string(path).map_err(|e| LoadError::FileReadFailed {
86                path: path.clone(),
87                source: e,
88            })?;
89            let located = parse_statements_with_context(&sql, path.clone(), variables, true)?;
90
91            if let Err(mut errs) = classify_cluster_statements(expected_name, path, located) {
92                errors.append(&mut errs);
93            }
94        }
95
96        // Resolve the active variant: prefer profile match, fall back to default
97        let active_path = object_files
98            .overrides
99            .get(profile)
100            .or(object_files.default.as_ref());
101
102        let active_path = match active_path {
103            Some(p) => p.clone(),
104            None => continue, // no variant for this profile
105        };
106
107        let sql = std::fs::read_to_string(&active_path).map_err(|e| LoadError::FileReadFailed {
108            path: active_path.clone(),
109            source: e,
110        })?;
111        let located = parse_statements_with_context(&sql, active_path.clone(), variables, true)?;
112
113        match classify_cluster_statements(expected_name, &active_path, located) {
114            Ok(def) => definitions.push(def),
115            Err(mut errs) => errors.append(&mut errs),
116        }
117    }
118
119    if !errors.is_empty() {
120        return Err(ValidationErrors::new(errors).into());
121    }
122
123    // Apply cluster suffix after validation (filename-vs-declared-name check uses original names)
124    if let Some(suffix) = profile_suffix {
125        for def in &mut definitions {
126            apply_cluster_suffix(def, suffix);
127        }
128    }
129
130    Ok(definitions)
131}
132
133/// Classify parsed statements into a `ClusterDefinition`, returning validation errors.
134fn classify_cluster_statements(
135    expected_name: &str,
136    path: &Path,
137    located_statements: Vec<LocatedStatement>,
138) -> Result<ClusterDefinition, Vec<ValidationError>> {
139    let mut create_stmts: Vec<(CreateClusterStatement<Raw>, usize)> = Vec::new();
140    let mut grants: Vec<GrantPrivilegesStatement<Raw>> = Vec::new();
141    let mut comments: Vec<CommentStatement<Raw>> = Vec::new();
142    let mut errors = Vec::new();
143
144    for LocatedStatement {
145        ast: stmt,
146        byte_offset,
147    } in located_statements
148    {
149        match stmt {
150            Statement::CreateCluster(s) => {
151                create_stmts.push((s, byte_offset));
152            }
153            Statement::GrantPrivileges(s) => {
154                // Validate that the grant targets a cluster
155                match &s.target {
156                    GrantTargetSpecification::Object {
157                        object_type: ObjectType::Cluster,
158                        object_spec_inner: GrantTargetSpecificationInner::Objects { names },
159                    } => {
160                        // Validate cluster name matches
161                        for name in names {
162                            let target_name = name.to_string();
163                            if target_name.to_lowercase() != expected_name.to_lowercase() {
164                                errors.push(ValidationError::with_file_sql_and_offset(
165                                    ValidationErrorKind::ClusterGrantTargetMismatch {
166                                        target: target_name,
167                                        cluster_name: expected_name.to_string(),
168                                    },
169                                    path.to_path_buf(),
170                                    s.to_string(),
171                                    byte_offset,
172                                ));
173                            }
174                        }
175                        grants.push(s);
176                    }
177                    _ => {
178                        errors.push(ValidationError::with_file_sql_and_offset(
179                            ValidationErrorKind::InvalidClusterStatement {
180                                statement_type: "GRANT (not targeting a cluster)".to_string(),
181                                cluster_name: expected_name.to_string(),
182                            },
183                            path.to_path_buf(),
184                            s.to_string(),
185                            byte_offset,
186                        ));
187                    }
188                }
189            }
190            Statement::Comment(s) => {
191                // Validate that the comment targets a cluster
192                match &s.object {
193                    CommentObjectType::Cluster { name } => {
194                        let target_name = match name {
195                            RawClusterName::Unresolved(ident) => ident.to_string(),
196                            RawClusterName::Resolved(id) => id.clone(),
197                        };
198                        if target_name.to_lowercase() != expected_name.to_lowercase() {
199                            errors.push(ValidationError::with_file_sql_and_offset(
200                                ValidationErrorKind::ClusterCommentTargetMismatch {
201                                    target: target_name,
202                                    cluster_name: expected_name.to_string(),
203                                },
204                                path.to_path_buf(),
205                                s.to_string(),
206                                byte_offset,
207                            ));
208                        }
209                        comments.push(s);
210                    }
211                    _ => {
212                        errors.push(ValidationError::with_file_sql_and_offset(
213                            ValidationErrorKind::InvalidClusterStatement {
214                                statement_type: "COMMENT (not targeting a cluster)".to_string(),
215                                cluster_name: expected_name.to_string(),
216                            },
217                            path.to_path_buf(),
218                            s.to_string(),
219                            byte_offset,
220                        ));
221                    }
222                }
223            }
224            other => {
225                errors.push(ValidationError::with_file_sql_and_offset(
226                    ValidationErrorKind::InvalidClusterStatement {
227                        statement_type: statement_type_name(&other).to_string(),
228                        cluster_name: expected_name.to_string(),
229                    },
230                    path.to_path_buf(),
231                    other.to_string(),
232                    byte_offset,
233                ));
234            }
235        }
236    }
237
238    // Validate exactly one CREATE CLUSTER (file-level errors)
239    if create_stmts.is_empty() {
240        errors.push(ValidationError::with_file(
241            ValidationErrorKind::ClusterMissingCreateStatement {
242                cluster_name: expected_name.to_string(),
243            },
244            path.to_path_buf(),
245        ));
246    } else if create_stmts.len() > 1 {
247        // Point to the second CREATE CLUSTER
248        errors.push(ValidationError::with_file_and_offset(
249            ValidationErrorKind::ClusterMultipleCreateStatements {
250                cluster_name: expected_name.to_string(),
251            },
252            path.to_path_buf(),
253            create_stmts[1].1,
254        ));
255    }
256
257    if !errors.is_empty() {
258        return Err(errors);
259    }
260
261    let (create_stmt, create_offset) = create_stmts.into_iter().next().unwrap();
262
263    // Validate cluster name matches filename
264    let declared_name = create_stmt.name.to_string();
265    if declared_name.to_lowercase() != expected_name.to_lowercase() {
266        return Err(vec![ValidationError::with_file_and_offset(
267            ValidationErrorKind::ClusterNameMismatch {
268                declared: declared_name,
269                expected: expected_name.to_string(),
270            },
271            path.to_path_buf(),
272            create_offset,
273        )]);
274    }
275
276    Ok(ClusterDefinition {
277        name: expected_name.to_string(),
278        create_stmt,
279        grants,
280        comments,
281    })
282}
283
284/// Apply a suffix to all cluster name references within a `ClusterDefinition`.
285///
286/// Rewrites the definition name, the CREATE statement name, GRANT target names,
287/// and COMMENT target names.
288fn apply_cluster_suffix(def: &mut ClusterDefinition, suffix: &str) {
289    // Rewrite definition name first, then reference it for the CREATE statement
290    def.name = format!("{}{}", def.name, suffix);
291    def.create_stmt.name = Ident::new(&def.name).expect("valid cluster identifier");
292
293    // Rewrite GRANT target cluster names
294    for grant in &mut def.grants {
295        if let GrantTargetSpecification::Object {
296            object_type: ObjectType::Cluster,
297            object_spec_inner: GrantTargetSpecificationInner::Objects { names },
298        } = &mut grant.target
299        {
300            for name in names {
301                if let UnresolvedObjectName::Cluster(ident) = name {
302                    *ident = suffixed_ident(ident, suffix);
303                }
304            }
305        }
306    }
307
308    // Rewrite COMMENT target cluster names
309    for comment in &mut def.comments {
310        if let CommentObjectType::Cluster { name } = &mut comment.object {
311            if let RawClusterName::Unresolved(ident) = name {
312                *ident = suffixed_ident(ident, suffix);
313            }
314        }
315    }
316}
317
318/// Append a suffix to an `Ident`, returning a new `Ident`.
319fn suffixed_ident(ident: &Ident, suffix: &str) -> Ident {
320    Ident::new(&format!("{}{}", ident, suffix)).expect("valid cluster identifier")
321}
322
323/// Extract the desired SIZE from a CreateClusterStatement's options.
324pub(crate) fn extract_size(create_stmt: &CreateClusterStatement<Raw>) -> Option<String> {
325    for opt in &create_stmt.options {
326        if opt.name == ClusterOptionName::Size {
327            if let Some(WithOptionValue::Value(ref v)) = opt.value {
328                return Some(v.to_string().trim_matches('\'').to_string());
329            }
330        }
331    }
332    None
333}
334
335/// Extract the desired REPLICATION FACTOR from a CreateClusterStatement's options.
336pub(crate) fn extract_replication_factor(create_stmt: &CreateClusterStatement<Raw>) -> Option<u32> {
337    for opt in &create_stmt.options {
338        if opt.name == ClusterOptionName::ReplicationFactor {
339            if let Some(WithOptionValue::Value(ref v)) = opt.value {
340                return v.to_string().parse::<u32>().ok();
341            }
342        }
343    }
344    None
345}
346
347/// Extract the desired autoscaling policy from a CreateClusterStatement's
348/// options. An absent option and an empty `AUTO SCALING STRATEGY = ()` block
349/// both map to `None` (no policy), matching the server planner's
350/// normalization.
351pub(crate) fn extract_auto_scaling_strategy(
352    create_stmt: &CreateClusterStatement<Raw>,
353) -> Result<Option<AutoScalingStrategy>, String> {
354    for opt in &create_stmt.options {
355        if opt.name == ClusterOptionName::AutoScalingStrategy {
356            return match &opt.value {
357                Some(WithOptionValue::ClusterAutoScalingStrategyOptionValue(value)) => {
358                    strategy_from_option_value(value)
359                }
360                _ => Err("invalid AUTO SCALING STRATEGY value".to_string()),
361            };
362        }
363    }
364    Ok(None)
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370    use std::fs;
371    use tempfile::TempDir;
372
373    fn create_test_dir() -> TempDir {
374        TempDir::new().unwrap()
375    }
376
377    #[mz_ore::test]
378    fn test_load_clusters_no_directory() {
379        let dir = create_test_dir();
380        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
381        assert!(
382            result.is_empty(),
383            "should return empty vec when clusters/ doesn't exist"
384        );
385    }
386
387    #[mz_ore::test]
388    fn test_load_clusters_basic() {
389        let dir = create_test_dir();
390        let clusters_dir = dir.path().join("clusters");
391        fs::create_dir(&clusters_dir).unwrap();
392
393        fs::write(
394            clusters_dir.join("analytics.sql"),
395            "CREATE CLUSTER analytics (SIZE = '100cc', REPLICATION FACTOR = 1);\n\
396             GRANT USAGE ON CLUSTER analytics TO analyst_role;\n\
397             COMMENT ON CLUSTER analytics IS 'Analytics workloads';",
398        )
399        .unwrap();
400
401        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
402        assert_eq!(result.len(), 1);
403        assert_eq!(result[0].name, "analytics");
404        assert_eq!(result[0].grants.len(), 1);
405        assert_eq!(result[0].comments.len(), 1);
406
407        // Verify extracted options
408        assert_eq!(
409            extract_size(&result[0].create_stmt),
410            Some("100cc".to_string())
411        );
412        assert_eq!(extract_replication_factor(&result[0].create_stmt), Some(1));
413    }
414
415    #[mz_ore::test]
416    fn test_load_clusters_create_only() {
417        let dir = create_test_dir();
418        let clusters_dir = dir.path().join("clusters");
419        fs::create_dir(&clusters_dir).unwrap();
420
421        fs::write(
422            clusters_dir.join("quickstart.sql"),
423            "CREATE CLUSTER quickstart (SIZE = '25cc');",
424        )
425        .unwrap();
426
427        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
428        assert_eq!(result.len(), 1);
429        assert_eq!(result[0].name, "quickstart");
430        assert!(result[0].grants.is_empty());
431        assert!(result[0].comments.is_empty());
432    }
433
434    #[mz_ore::test]
435    fn test_load_clusters_name_mismatch() {
436        let dir = create_test_dir();
437        let clusters_dir = dir.path().join("clusters");
438        fs::create_dir(&clusters_dir).unwrap();
439
440        fs::write(
441            clusters_dir.join("analytics.sql"),
442            "CREATE CLUSTER wrong_name (SIZE = '100cc');",
443        )
444        .unwrap();
445
446        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
447        assert!(
448            result.is_err(),
449            "should error when cluster name doesn't match filename"
450        );
451    }
452
453    #[mz_ore::test]
454    fn test_load_clusters_missing_create() {
455        let dir = create_test_dir();
456        let clusters_dir = dir.path().join("clusters");
457        fs::create_dir(&clusters_dir).unwrap();
458
459        fs::write(
460            clusters_dir.join("analytics.sql"),
461            "GRANT USAGE ON CLUSTER analytics TO analyst_role;",
462        )
463        .unwrap();
464
465        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
466        assert!(
467            result.is_err(),
468            "should error when no CREATE CLUSTER statement"
469        );
470    }
471
472    #[mz_ore::test]
473    fn test_load_clusters_unsupported_statement() {
474        let dir = create_test_dir();
475        let clusters_dir = dir.path().join("clusters");
476        fs::create_dir(&clusters_dir).unwrap();
477
478        fs::write(
479            clusters_dir.join("analytics.sql"),
480            "CREATE CLUSTER analytics (SIZE = '100cc');\n\
481             CREATE TABLE foo (id INT);",
482        )
483        .unwrap();
484
485        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
486        assert!(
487            result.is_err(),
488            "should error on unsupported statement type"
489        );
490    }
491
492    #[mz_ore::test]
493    fn test_load_clusters_grant_target_mismatch() {
494        let dir = create_test_dir();
495        let clusters_dir = dir.path().join("clusters");
496        fs::create_dir(&clusters_dir).unwrap();
497
498        fs::write(
499            clusters_dir.join("analytics.sql"),
500            "CREATE CLUSTER analytics (SIZE = '100cc');\n\
501             GRANT USAGE ON CLUSTER other_cluster TO analyst_role;",
502        )
503        .unwrap();
504
505        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
506        assert!(
507            result.is_err(),
508            "should error when grant targets wrong cluster"
509        );
510    }
511
512    #[mz_ore::test]
513    fn test_load_clusters_comment_target_mismatch() {
514        let dir = create_test_dir();
515        let clusters_dir = dir.path().join("clusters");
516        fs::create_dir(&clusters_dir).unwrap();
517
518        fs::write(
519            clusters_dir.join("analytics.sql"),
520            "CREATE CLUSTER analytics (SIZE = '100cc');\n\
521             COMMENT ON CLUSTER other_cluster IS 'wrong target';",
522        )
523        .unwrap();
524
525        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
526        assert!(
527            result.is_err(),
528            "should error when comment targets wrong cluster"
529        );
530    }
531
532    #[mz_ore::test]
533    fn test_load_clusters_multiple_files() {
534        let dir = create_test_dir();
535        let clusters_dir = dir.path().join("clusters");
536        fs::create_dir(&clusters_dir).unwrap();
537
538        fs::write(
539            clusters_dir.join("analytics.sql"),
540            "CREATE CLUSTER analytics (SIZE = '100cc');",
541        )
542        .unwrap();
543
544        fs::write(
545            clusters_dir.join("quickstart.sql"),
546            "CREATE CLUSTER quickstart (SIZE = '25cc', REPLICATION FACTOR = 2);",
547        )
548        .unwrap();
549
550        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
551        assert_eq!(result.len(), 2);
552        // Sorted by filename
553        assert_eq!(result[0].name, "analytics");
554        assert_eq!(result[1].name, "quickstart");
555    }
556
557    #[mz_ore::test]
558    fn test_load_clusters_multi_variant_valid() {
559        let dir = create_test_dir();
560        let clusters_dir = dir.path().join("clusters");
561        fs::create_dir(&clusters_dir).unwrap();
562
563        fs::write(
564            clusters_dir.join("analytics.sql"),
565            "CREATE CLUSTER analytics (SIZE = '100cc');",
566        )
567        .unwrap();
568        fs::write(
569            clusters_dir.join("analytics#staging.sql"),
570            "CREATE CLUSTER analytics (SIZE = '25cc');",
571        )
572        .unwrap();
573
574        // With staging profile, should pick staging variant
575        let result = load_clusters(dir.path(), "staging", None, &BTreeMap::new()).unwrap();
576        assert_eq!(result.len(), 1);
577        assert_eq!(result[0].name, "analytics");
578        assert_eq!(
579            extract_size(&result[0].create_stmt),
580            Some("25cc".to_string())
581        );
582    }
583
584    #[mz_ore::test]
585    fn test_load_clusters_multi_variant_fallback_default() {
586        let dir = create_test_dir();
587        let clusters_dir = dir.path().join("clusters");
588        fs::create_dir(&clusters_dir).unwrap();
589
590        fs::write(
591            clusters_dir.join("analytics.sql"),
592            "CREATE CLUSTER analytics (SIZE = '100cc');",
593        )
594        .unwrap();
595        fs::write(
596            clusters_dir.join("analytics#staging.sql"),
597            "CREATE CLUSTER analytics (SIZE = '25cc');",
598        )
599        .unwrap();
600
601        // With prod profile (no match), should fall back to default
602        let result = load_clusters(dir.path(), "prod", None, &BTreeMap::new()).unwrap();
603        assert_eq!(result.len(), 1);
604        assert_eq!(result[0].name, "analytics");
605        assert_eq!(
606            extract_size(&result[0].create_stmt),
607            Some("100cc".to_string())
608        );
609    }
610
611    #[mz_ore::test]
612    fn test_load_clusters_auto_scaling_strategy() {
613        use std::time::Duration;
614
615        let dir = create_test_dir();
616        let clusters_dir = dir.path().join("clusters");
617        fs::create_dir(&clusters_dir).unwrap();
618
619        fs::write(
620            clusters_dir.join("analytics.sql"),
621            "CREATE CLUSTER analytics (SIZE = '100cc', AUTO SCALING STRATEGY = \
622             (ON HYDRATION (HYDRATION SIZE = '3200cc', LINGER DURATION = '600s')));",
623        )
624        .unwrap();
625
626        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
627        assert_eq!(result.len(), 1);
628        let strategy = extract_auto_scaling_strategy(&result[0].create_stmt)
629            .unwrap()
630            .unwrap();
631        let on_hydration = strategy.on_hydration.unwrap();
632        assert_eq!(on_hydration.hydration_size, "3200cc");
633        assert_eq!(on_hydration.linger_duration, Some(Duration::from_secs(600)));
634    }
635
636    #[mz_ore::test]
637    fn test_extract_auto_scaling_strategy_absent_and_empty() {
638        let dir = create_test_dir();
639        let clusters_dir = dir.path().join("clusters");
640        fs::create_dir(&clusters_dir).unwrap();
641
642        fs::write(
643            clusters_dir.join("plain.sql"),
644            "CREATE CLUSTER plain (SIZE = '100cc');",
645        )
646        .unwrap();
647        fs::write(
648            clusters_dir.join("disabled.sql"),
649            "CREATE CLUSTER disabled (SIZE = '100cc', AUTO SCALING STRATEGY = ());",
650        )
651        .unwrap();
652
653        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new()).unwrap();
654        assert_eq!(result.len(), 2);
655        for def in &result {
656            assert_eq!(
657                extract_auto_scaling_strategy(&def.create_stmt).unwrap(),
658                None,
659                "cluster '{}' should have no policy",
660                def.name
661            );
662        }
663    }
664
665    #[mz_ore::test]
666    fn test_load_clusters_invalid_variant_errors() {
667        let dir = create_test_dir();
668        let clusters_dir = dir.path().join("clusters");
669        fs::create_dir(&clusters_dir).unwrap();
670
671        fs::write(
672            clusters_dir.join("analytics.sql"),
673            "CREATE CLUSTER analytics (SIZE = '100cc');",
674        )
675        .unwrap();
676        // Invalid staging variant: name mismatch
677        fs::write(
678            clusters_dir.join("analytics#staging.sql"),
679            "CREATE CLUSTER wrong_name (SIZE = '25cc');",
680        )
681        .unwrap();
682
683        let result = load_clusters(dir.path(), "default", None, &BTreeMap::new());
684        assert!(
685            result.is_err(),
686            "invalid variant should error even when not active profile"
687        );
688    }
689}