1use std::fmt::Debug;
13use std::num::NonZeroI64;
14use std::sync::Arc;
15use std::time::Duration;
16
17use async_trait::async_trait;
18use mz_audit_log::VersionedEvent;
19use mz_controller_types::{ClusterId, ReplicaId};
20use mz_ore::collections::CollectionExt;
21use mz_ore::metrics::MetricsRegistry;
22use mz_persist_client::PersistClient;
23use mz_persist_types::ShardId;
24use mz_repr::{CatalogItemId, GlobalId, RelationDesc, SqlScalarType};
25use mz_sql::catalog::CatalogError as SqlCatalogError;
26use uuid::Uuid;
27
28use crate::config::ClusterReplicaSizeMap;
29use crate::durable::debug::{DebugCatalogState, Trace};
30pub use crate::durable::error::{CatalogError, DurableCatalogError, FenceError};
31pub use crate::durable::metrics::Metrics;
32pub use crate::durable::objects::Snapshot;
33pub use crate::durable::objects::state_update::StateUpdate;
34pub use crate::durable::objects::{
35 BurstState, Cluster, ClusterConfig, ClusterReplica, ClusterSystemConfiguration, ClusterVariant,
36 ClusterVariantManaged, Comment, Database, DefaultPrivilege, IntrospectionSourceIndex, Item,
37 NetworkPolicy, ReconfigurationState, ReconfigurationStatus, ReconfigurationTarget,
38 ReplicaConfig, ReplicaLocation, ReplicaSystemConfiguration, Role, RoleAuth, Schema,
39 SourceReference, SourceReferences, StorageCollectionMetadata, SystemConfiguration,
40 SystemObjectDescription, SystemObjectMapping, UnfinalizedShard,
41};
42pub use crate::durable::persist::shard_id;
43use crate::durable::persist::{Timestamp, UnopenedPersistCatalogState};
44use crate::durable::transaction::TransactionBatch;
45pub use crate::durable::transaction::{DryRunTransaction, Transaction};
46pub use crate::durable::upgrade::CATALOG_VERSION;
47use crate::memory;
48
49pub mod debug;
50mod error;
51pub mod initialize;
52mod metrics;
53pub mod objects;
54mod persist;
55mod traits;
56mod transaction;
57mod upgrade;
58
59pub const DATABASE_ID_ALLOC_KEY: &str = "database";
60pub const SCHEMA_ID_ALLOC_KEY: &str = "schema";
61pub const USER_ITEM_ALLOC_KEY: &str = "user";
62pub const SYSTEM_ITEM_ALLOC_KEY: &str = "system";
63pub const USER_ROLE_ID_ALLOC_KEY: &str = "user_role";
64pub const USER_CLUSTER_ID_ALLOC_KEY: &str = "user_compute";
65pub const SYSTEM_CLUSTER_ID_ALLOC_KEY: &str = "system_compute";
66pub const USER_REPLICA_ID_ALLOC_KEY: &str = "replica";
67pub const SYSTEM_REPLICA_ID_ALLOC_KEY: &str = "system_replica";
68pub const AUDIT_LOG_ID_ALLOC_KEY: &str = "auditlog";
69pub const STORAGE_USAGE_ID_ALLOC_KEY: &str = "storage_usage";
70pub const USER_NETWORK_POLICY_ID_ALLOC_KEY: &str = "user_network_policy";
71pub const OID_ALLOC_KEY: &str = "oid";
72pub(crate) const CATALOG_CONTENT_VERSION_KEY: &str = "catalog_content_version";
73pub const BUILTIN_MIGRATION_SHARD_KEY: &str = "builtin_migration_shard";
74pub const EXPRESSION_CACHE_SHARD_KEY: &str = "expression_cache_shard";
75pub const MOCK_AUTHENTICATION_NONCE_KEY: &str = "mock_authentication_nonce";
76
77#[derive(Clone, Debug)]
78pub struct BootstrapArgs {
79 pub cluster_replica_size_map: ClusterReplicaSizeMap,
80 pub default_cluster_replica_size: String,
81 pub default_cluster_replication_factor: u32,
82 pub bootstrap_role: Option<String>,
83}
84
85pub type Epoch = NonZeroI64;
86
87#[async_trait]
91pub trait OpenableDurableCatalogState: Debug + Send {
92 async fn open_savepoint(
109 mut self: Box<Self>,
110 initial_ts: Timestamp,
111 bootstrap_args: &BootstrapArgs,
112 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
113
114 async fn open_read_only(
120 mut self: Box<Self>,
121 bootstrap_args: &BootstrapArgs,
122 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
123
124 async fn open(
130 mut self: Box<Self>,
131 initial_ts: Timestamp,
132 bootstrap_args: &BootstrapArgs,
133 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
134
135 async fn open_debug(mut self: Box<Self>) -> Result<DebugCatalogState, CatalogError>;
138
139 async fn is_initialized(&mut self) -> Result<bool, CatalogError>;
141
142 async fn epoch(&mut self) -> Result<Epoch, CatalogError>;
152
153 async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError>;
156
157 async fn get_0dt_deployment_max_wait(&mut self) -> Result<Option<Duration>, CatalogError>;
163
164 async fn get_0dt_deployment_ddl_check_interval(
170 &mut self,
171 ) -> Result<Option<Duration>, CatalogError>;
172
173 async fn get_enable_0dt_deployment_panic_after_timeout(
180 &mut self,
181 ) -> Result<Option<bool>, CatalogError>;
182
183 async fn has_system_config_synced_once(&mut self) -> Result<bool, DurableCatalogError>;
185
186 async fn trace_unconsolidated(&mut self) -> Result<Trace, CatalogError>;
188
189 async fn trace_consolidated(&mut self) -> Result<Trace, CatalogError>;
191
192 async fn expire(self: Box<Self>);
194}
195
196#[async_trait]
198pub trait ReadOnlyDurableCatalogState: Debug + Send + Sync {
199 fn epoch(&self) -> Epoch;
209
210 fn metrics(&self) -> &Metrics;
212
213 async fn expire(self: Box<Self>);
215
216 fn is_bootstrap_complete(&self) -> bool;
218
219 async fn get_audit_logs(&mut self) -> Result<Vec<VersionedEvent>, CatalogError>;
225
226 async fn get_next_id(&mut self, id_type: &str) -> Result<u64, CatalogError>;
228
229 async fn get_next_user_item_id(&mut self) -> Result<u64, CatalogError> {
231 self.get_next_id(USER_ITEM_ALLOC_KEY).await
232 }
233
234 async fn get_next_system_item_id(&mut self) -> Result<u64, CatalogError> {
236 self.get_next_id(SYSTEM_ITEM_ALLOC_KEY).await
237 }
238
239 async fn get_next_system_replica_id(&mut self) -> Result<u64, CatalogError> {
241 self.get_next_id(SYSTEM_REPLICA_ID_ALLOC_KEY).await
242 }
243
244 async fn get_next_user_replica_id(&mut self) -> Result<u64, CatalogError> {
246 self.get_next_id(USER_REPLICA_ID_ALLOC_KEY).await
247 }
248
249 async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError>;
251
252 async fn snapshot(&mut self) -> Result<Snapshot, CatalogError>;
254
255 async fn sync_to_current_updates(
261 &mut self,
262 ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError>;
263
264 async fn sync_updates(
273 &mut self,
274 target_upper: Timestamp,
275 ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError>;
276
277 async fn ensure_not_out_of_sync(&mut self, target_upper: Timestamp)
281 -> Result<(), CatalogError>;
282
283 async fn current_upper(&mut self) -> Timestamp;
285}
286
287#[async_trait]
289#[allow(mismatched_lifetime_syntaxes)]
290pub trait DurableCatalogState: ReadOnlyDurableCatalogState {
291 fn is_read_only(&self) -> bool;
293
294 fn is_savepoint(&self) -> bool;
296
297 async fn mark_bootstrap_complete(&mut self);
299
300 async fn transaction(&mut self) -> Result<Transaction, CatalogError>;
304
305 fn transaction_from_snapshot(
309 &mut self,
310 snapshot: Snapshot,
311 ) -> Result<DryRunTransaction, CatalogError>;
312
313 async fn commit_transaction(
321 &mut self,
322 txn_batch: TransactionBatch,
323 commit_ts: Timestamp,
324 ) -> Result<Timestamp, CatalogError>;
325
326 async fn advance_upper(&mut self, new_upper: Timestamp) -> Result<(), CatalogError>;
332
333 async fn allocate_id(
338 &mut self,
339 id_type: &str,
340 amount: u64,
341 commit_ts: Timestamp,
342 ) -> Result<Vec<u64>, CatalogError>;
343
344 async fn allocate_user_ids(
348 &mut self,
349 amount: u64,
350 commit_ts: Timestamp,
351 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
352 let ids = self
353 .allocate_id(USER_ITEM_ALLOC_KEY, amount, commit_ts)
354 .await?;
355 let ids = ids
356 .iter()
357 .map(|id| (CatalogItemId::User(*id), GlobalId::User(*id)))
358 .collect();
359 Ok(ids)
360 }
361
362 async fn allocate_user_id(
366 &mut self,
367 commit_ts: Timestamp,
368 ) -> Result<(CatalogItemId, GlobalId), CatalogError> {
369 let id = self.allocate_id(USER_ITEM_ALLOC_KEY, 1, commit_ts).await?;
370 let id = id.into_element();
371 Ok((CatalogItemId::User(id), GlobalId::User(id)))
372 }
373
374 async fn allocate_user_cluster_id(
378 &mut self,
379 commit_ts: Timestamp,
380 ) -> Result<ClusterId, CatalogError> {
381 let id = self
382 .allocate_id(USER_CLUSTER_ID_ALLOC_KEY, 1, commit_ts)
383 .await?
384 .into_element();
385 Ok(ClusterId::user(id).ok_or(SqlCatalogError::IdExhaustion)?)
386 }
387
388 async fn allocate_user_replica_ids(
392 &mut self,
393 amount: u64,
394 commit_ts: Timestamp,
395 ) -> Result<Vec<ReplicaId>, CatalogError> {
396 let ids = self
397 .allocate_id(USER_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
398 .await?;
399 let ids = ids.into_iter().map(ReplicaId::User).collect();
400 Ok(ids)
401 }
402
403 async fn allocate_system_replica_ids(
407 &mut self,
408 amount: u64,
409 commit_ts: Timestamp,
410 ) -> Result<Vec<ReplicaId>, CatalogError> {
411 let ids = self
412 .allocate_id(SYSTEM_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
413 .await?;
414 let ids = ids.into_iter().map(ReplicaId::System).collect();
415 Ok(ids)
416 }
417
418 fn shard_id(&self) -> ShardId;
419}
420
421pub fn persist_desc() -> RelationDesc {
424 RelationDesc::builder()
425 .with_column("data", SqlScalarType::Jsonb.nullable(false))
426 .finish()
427}
428
429#[derive(Debug, Clone)]
431pub struct TestCatalogStateBuilder {
432 persist_client: PersistClient,
433 organization_id: Uuid,
434 version: semver::Version,
435 deploy_generation: Option<u64>,
436 metrics: Arc<Metrics>,
437}
438
439impl TestCatalogStateBuilder {
440 pub fn new(persist_client: PersistClient) -> Self {
441 Self {
442 persist_client,
443 organization_id: Uuid::new_v4(),
444 version: semver::Version::new(0, 0, 0),
445 deploy_generation: None,
446 metrics: Arc::new(Metrics::new(&MetricsRegistry::new())),
447 }
448 }
449
450 pub fn with_organization_id(mut self, organization_id: Uuid) -> Self {
451 self.organization_id = organization_id;
452 self
453 }
454
455 pub fn with_version(mut self, version: semver::Version) -> Self {
456 self.version = version;
457 self
458 }
459
460 pub fn with_deploy_generation(mut self, deploy_generation: u64) -> Self {
461 self.deploy_generation = Some(deploy_generation);
462 self
463 }
464
465 pub fn with_default_deploy_generation(self) -> Self {
466 self.with_deploy_generation(0)
467 }
468
469 pub fn with_metrics(mut self, metrics: Arc<Metrics>) -> Self {
470 self.metrics = metrics;
471 self
472 }
473
474 pub async fn build(self) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
475 persist_backed_catalog_state(
476 self.persist_client,
477 self.organization_id,
478 self.version,
479 self.deploy_generation,
480 self.metrics,
481 )
482 .await
483 }
484
485 pub async fn unwrap_build(self) -> Box<dyn OpenableDurableCatalogState> {
486 self.expect_build("failed to build").await
487 }
488
489 pub async fn expect_build(self, msg: &str) -> Box<dyn OpenableDurableCatalogState> {
490 self.build().await.expect(msg)
491 }
492}
493
494pub async fn persist_backed_catalog_state(
498 persist_client: PersistClient,
499 organization_id: Uuid,
500 version: semver::Version,
501 deploy_generation: Option<u64>,
502 metrics: Arc<Metrics>,
503) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
504 let state = UnopenedPersistCatalogState::new(
505 persist_client,
506 organization_id,
507 version,
508 deploy_generation,
509 metrics,
510 )
511 .await?;
512 Ok(Box::new(state))
513}
514
515pub fn test_bootstrap_args() -> BootstrapArgs {
516 BootstrapArgs {
517 default_cluster_replica_size: "scale=1,workers=1".into(),
518 default_cluster_replication_factor: 1,
519 bootstrap_role: None,
520 cluster_replica_size_map: ClusterReplicaSizeMap::for_tests(),
521 }
522}