1use 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
33pub(crate) struct ClusterDefinition {
35 pub name: String,
37 pub create_stmt: CreateClusterStatement<Raw>,
39 pub grants: Vec<GrantPrivilegesStatement<Raw>>,
41 pub comments: Vec<CommentStatement<Raw>>,
43}
44
45pub(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 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 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 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, };
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 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
133fn 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 match &s.target {
156 GrantTargetSpecification::Object {
157 object_type: ObjectType::Cluster,
158 object_spec_inner: GrantTargetSpecificationInner::Objects { names },
159 } => {
160 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 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 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 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 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
284fn apply_cluster_suffix(def: &mut ClusterDefinition, suffix: &str) {
289 def.name = format!("{}{}", def.name, suffix);
291 def.create_stmt.name = Ident::new(&def.name).expect("valid cluster identifier");
292
293 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 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
318fn suffixed_ident(ident: &Ident, suffix: &str) -> Ident {
320 Ident::new(&format!("{}{}", ident, suffix)).expect("valid cluster identifier")
321}
322
323pub(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
335pub(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
347pub(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 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 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 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 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 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}