1use std::collections::{BTreeMap, BTreeSet};
11use std::num::NonZero;
12
13use anyhow::bail;
14use bytesize::ByteSize;
15use ipnet::IpNet;
16use mz_adapter_types::bootstrap_builtin_cluster_config::BootstrapBuiltinClusterConfig;
17use mz_auth::password::Password;
18use mz_build_info::BuildInfo;
19use mz_cloud_resources::AwsExternalIdPrefix;
20use mz_controller::clusters::ReplicaAllocation;
21use mz_license_keys::ValidatedLicenseKey;
22use mz_orchestrator::MemoryLimit;
23use mz_ore::cast::CastFrom;
24use mz_ore::metrics::MetricsRegistry;
25use mz_persist_client::PersistClient;
26use mz_repr::CatalogItemId;
27use mz_repr::adt::numeric::Numeric;
28use mz_sql::catalog::CatalogError as SqlCatalogError;
29use mz_sql::catalog::EnvironmentId;
30use serde::Serialize;
31
32use crate::durable::{CatalogError, DurableCatalogState};
33
34const GIB: u64 = 1024 * 1024 * 1024;
35
36#[derive(Debug)]
38pub struct Config<'a> {
39 pub storage: Box<dyn DurableCatalogState>,
41 pub metrics_registry: &'a MetricsRegistry,
43 pub state: StateConfig,
44}
45
46#[derive(Debug)]
47pub struct StateConfig {
48 pub unsafe_mode: bool,
50 pub all_features: bool,
52 pub build_info: &'static BuildInfo,
54 pub environment_id: EnvironmentId,
56 pub read_only: bool,
58 pub now: mz_ore::now::NowFn,
60 pub boot_ts: mz_repr::Timestamp,
62 pub skip_migrations: bool,
64 pub cluster_replica_sizes: ClusterReplicaSizeMap,
66 pub builtin_system_cluster_config: BootstrapBuiltinClusterConfig,
68 pub builtin_catalog_server_cluster_config: BootstrapBuiltinClusterConfig,
70 pub builtin_probe_cluster_config: BootstrapBuiltinClusterConfig,
72 pub builtin_support_cluster_config: BootstrapBuiltinClusterConfig,
74 pub builtin_analytics_cluster_config: BootstrapBuiltinClusterConfig,
76 pub system_parameter_defaults: BTreeMap<String, String>,
78 pub remote_system_parameters: Option<BTreeMap<String, String>>,
81 pub availability_zones: Vec<String>,
83 pub egress_addresses: Vec<IpNet>,
85 pub aws_principal_context: Option<AwsPrincipalContext>,
87 pub aws_privatelink_availability_zones: Option<BTreeSet<String>>,
89 pub http_host_name: Option<String>,
91 pub connection_context: mz_storage_types::connections::ConnectionContext,
93 pub builtin_item_migration_config: BuiltinItemMigrationConfig,
94 pub persist_client: PersistClient,
95 pub enable_expression_cache_override: Option<bool>,
98 pub helm_chart_version: Option<String>,
100 pub external_login_password_mz_system: Option<Password>,
101 pub license_key: ValidatedLicenseKey,
102}
103
104#[derive(Debug)]
105pub struct BuiltinItemMigrationConfig {
106 pub persist_client: PersistClient,
107 pub read_only: bool,
108 pub force_migration: Option<String>,
109}
110
111#[derive(Debug, Clone, Serialize)]
112pub struct ClusterReplicaSizeMap(pub BTreeMap<String, ReplicaAllocation>);
113
114impl ClusterReplicaSizeMap {
115 pub fn parse_from_str(s: &str, credit_consumption_from_memory: bool) -> anyhow::Result<Self> {
116 let mut cluster_replica_sizes: BTreeMap<String, ReplicaAllocation> =
117 serde_json::from_str(s)?;
118 if credit_consumption_from_memory {
119 for (name, replica) in cluster_replica_sizes.iter_mut() {
120 let Some(memory_limit) = replica.memory_limit else {
121 bail!("No memory limit found in cluster definition for {name}");
122 };
123 let total_memory = memory_limit.0 * replica.scale.get();
124 replica.credits_per_hour = Numeric::from(total_memory.0) / Numeric::from(GIB);
125 }
126 }
127 Ok(Self(cluster_replica_sizes))
128 }
129
130 pub fn enabled_allocations(&self) -> impl Iterator<Item = (&String, &ReplicaAllocation)> {
132 self.0.iter().filter(|(_, a)| !a.disabled)
133 }
134
135 pub fn get_allocation_by_name(&self, name: &str) -> Result<&ReplicaAllocation, CatalogError> {
138 self.0.get(name).ok_or_else(|| {
139 CatalogError::Catalog(SqlCatalogError::UnknownClusterReplicaSize(name.into()))
140 })
141 }
142
143 pub fn for_tests() -> Self {
148 let mut inner = (0..=5)
170 .flat_map(|i| {
171 let workers = 1 << i;
172 [
173 (format!("scale=1,workers={workers}"), None),
174 (format!("scale=1,workers={workers},mem=4GiB"), Some(4)),
175 (format!("scale=1,workers={workers},mem=8GiB"), Some(8)),
176 (format!("scale=1,workers={workers},mem=16GiB"), Some(16)),
177 (format!("scale=1,workers={workers},mem=32GiB"), Some(32)),
178 ]
179 .map(|(name, memory_limit)| {
180 (
181 name,
182 ReplicaAllocation {
183 memory_limit: memory_limit.map(|gib| MemoryLimit(ByteSize::gib(gib))),
184 cpu_limit: None,
185 cpu_request: None,
186 disk_limit: None,
187 scale: NonZero::new(1).expect("not zero"),
188 workers: NonZero::new(workers).expect("not zero"),
189 credits_per_hour: 1.into(),
190 cpu_exclusive: false,
191 is_cc: false,
192 swap_enabled: false,
193 disabled: false,
194 selectors: BTreeMap::default(),
195 },
196 )
197 })
198 })
199 .collect::<BTreeMap<_, _>>();
200
201 for i in 1..=5 {
202 let scale = 1 << i;
203 inner.insert(
204 format!("scale={scale},workers=1"),
205 ReplicaAllocation {
206 memory_limit: None,
207 cpu_limit: None,
208 cpu_request: None,
209 disk_limit: None,
210 scale: NonZero::new(scale).expect("not zero"),
211 workers: NonZero::new(1).expect("not zero"),
212 credits_per_hour: scale.into(),
213 cpu_exclusive: false,
214 is_cc: false,
215 swap_enabled: false,
216 disabled: false,
217 selectors: BTreeMap::default(),
218 },
219 );
220
221 inner.insert(
222 format!("scale={scale},workers={scale}"),
223 ReplicaAllocation {
224 memory_limit: None,
225 cpu_limit: None,
226 cpu_request: None,
227 disk_limit: None,
228 scale: NonZero::new(scale).expect("not zero"),
229 workers: NonZero::new(scale.into()).expect("not zero"),
230 credits_per_hour: scale.into(),
231 cpu_exclusive: false,
232 is_cc: false,
233 swap_enabled: false,
234 disabled: false,
235 selectors: BTreeMap::default(),
236 },
237 );
238
239 inner.insert(
240 format!("scale=1,workers=8,mem={scale}GiB"),
241 ReplicaAllocation {
242 memory_limit: Some(MemoryLimit(ByteSize(u64::cast_from(scale) * (1 << 30)))),
243 cpu_limit: None,
244 cpu_request: None,
245 disk_limit: None,
246 scale: NonZero::new(1).expect("not zero"),
247 workers: NonZero::new(8).expect("not zero"),
248 credits_per_hour: 1.into(),
249 cpu_exclusive: false,
250 is_cc: false,
251 swap_enabled: false,
252 disabled: false,
253 selectors: BTreeMap::default(),
254 },
255 );
256 }
257
258 inner.insert(
259 "scale=2,workers=4".to_string(),
260 ReplicaAllocation {
261 memory_limit: None,
262 cpu_limit: None,
263 cpu_request: None,
264 disk_limit: None,
265 scale: NonZero::new(2).expect("not zero"),
266 workers: NonZero::new(4).expect("not zero"),
267 credits_per_hour: 2.into(),
268 cpu_exclusive: false,
269 is_cc: false,
270 swap_enabled: false,
271 disabled: false,
272 selectors: BTreeMap::default(),
273 },
274 );
275
276 inner.insert(
277 "free".to_string(),
278 ReplicaAllocation {
279 memory_limit: None,
280 cpu_limit: None,
281 cpu_request: None,
282 disk_limit: None,
283 scale: NonZero::new(1).expect("not zero"),
284 workers: NonZero::new(1).expect("not zero"),
285 credits_per_hour: 0.into(),
286 cpu_exclusive: false,
287 is_cc: true,
288 swap_enabled: false,
289 disabled: true,
290 selectors: BTreeMap::default(),
291 },
292 );
293
294 Self(inner)
295 }
296}
297
298#[derive(Debug, Clone, Serialize)]
303pub struct AwsPrincipalContext {
304 pub aws_account_id: String,
305 pub aws_external_id_prefix: AwsExternalIdPrefix,
306}
307
308impl AwsPrincipalContext {
309 pub fn to_principal_string(&self, aws_external_id_suffix: CatalogItemId) -> String {
310 format!(
311 "arn:aws:iam::{}:role/mz_{}_{}",
312 self.aws_account_id, self.aws_external_id_prefix, aws_external_id_suffix
313 )
314 }
315}
316
317#[cfg(test)]
318#[allow(clippy::unwrap_used)]
319mod tests {
320 use super::*;
321
322 #[mz_ore::test]
323 #[cfg_attr(miri, ignore)] fn cluster_replica_size_credits_from_memory() {
325 let s = r#"{
326 "test": {
327 "memory_limit": "1000MiB",
328 "scale": 2,
329 "workers": 10,
330 "credits_per_hour": "0"
331 }
332 }"#;
333 let map = ClusterReplicaSizeMap::parse_from_str(s, true).unwrap();
334
335 let alloc = map.get_allocation_by_name("test").unwrap();
336 let expected = Numeric::from(2000) / Numeric::from(1024);
337 assert_eq!(alloc.credits_per_hour, expected);
338 }
339}