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, ClusterReplicaSize,
36 ClusterSystemConfiguration, ClusterVariant, ClusterVariantManaged, Comment, Database,
37 DefaultPrivilege, IndexColumn, IndexColumns, IntrospectionSourceIndex, Item, ItemColumn,
38 ItemColumns, ItemMetadata, NetworkPolicy, ReconfigurationState, ReconfigurationStatus,
39 ReconfigurationTarget, ReplicaConfig, ReplicaLocation, ReplicaSystemConfiguration, Role,
40 RoleAuth, Schema, SourceReference, SourceReferences, StorageCollectionMetadata,
41 SystemConfiguration, SystemObjectDescription, SystemObjectMapping, UnfinalizedShard,
42 managed_cluster_replica_name,
43};
44pub use crate::durable::persist::shard_id;
45use crate::durable::persist::{Timestamp, UnopenedPersistCatalogState};
46use crate::durable::transaction::TransactionBatch;
47pub use crate::durable::transaction::{DryRunTransaction, Transaction};
48pub use crate::durable::upgrade::CATALOG_VERSION;
49use crate::memory;
50
51pub mod debug;
52mod error;
53pub mod initialize;
54mod metrics;
55pub mod objects;
56mod persist;
57mod traits;
58mod transaction;
59mod upgrade;
60
61pub const DATABASE_ID_ALLOC_KEY: &str = "database";
62pub const SCHEMA_ID_ALLOC_KEY: &str = "schema";
63pub const USER_ITEM_ALLOC_KEY: &str = "user";
64pub const SYSTEM_ITEM_ALLOC_KEY: &str = "system";
65pub const USER_ROLE_ID_ALLOC_KEY: &str = "user_role";
66pub const USER_CLUSTER_ID_ALLOC_KEY: &str = "user_compute";
67pub const SYSTEM_CLUSTER_ID_ALLOC_KEY: &str = "system_compute";
68pub const USER_REPLICA_ID_ALLOC_KEY: &str = "replica";
69pub const SYSTEM_REPLICA_ID_ALLOC_KEY: &str = "system_replica";
70pub const AUDIT_LOG_ID_ALLOC_KEY: &str = "auditlog";
71pub const STORAGE_USAGE_ID_ALLOC_KEY: &str = "storage_usage";
72pub const USER_NETWORK_POLICY_ID_ALLOC_KEY: &str = "user_network_policy";
73pub const USER_CLUSTER_REPLICA_SIZE_ID_ALLOC_KEY: &str = "user_cluster_replica_size";
74pub const SYSTEM_CLUSTER_REPLICA_SIZE_ID_ALLOC_KEY: &str = "system_cluster_replica_size";
75pub const OID_ALLOC_KEY: &str = "oid";
76pub(crate) const CATALOG_CONTENT_VERSION_KEY: &str = "catalog_content_version";
77pub const BUILTIN_MIGRATION_SHARD_KEY: &str = "builtin_migration_shard";
78pub const EXPRESSION_CACHE_SHARD_KEY: &str = "expression_cache_shard";
79pub const MOCK_AUTHENTICATION_NONCE_KEY: &str = "mock_authentication_nonce";
80
81#[derive(Clone, Debug)]
82pub struct BootstrapArgs {
83 pub cluster_replica_size_map: ClusterReplicaSizeMap,
84 pub default_cluster_replica_size: String,
85 pub default_cluster_replication_factor: u32,
86 pub bootstrap_role: Option<String>,
87}
88
89pub type Epoch = NonZeroI64;
90
91#[async_trait]
95pub trait OpenableDurableCatalogState: Debug + Send {
96 async fn open_savepoint(
113 mut self: Box<Self>,
114 initial_ts: Timestamp,
115 bootstrap_args: &BootstrapArgs,
116 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
117
118 async fn open_read_only(
124 mut self: Box<Self>,
125 bootstrap_args: &BootstrapArgs,
126 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
127
128 async fn open(
134 mut self: Box<Self>,
135 initial_ts: Timestamp,
136 bootstrap_args: &BootstrapArgs,
137 ) -> Result<Box<dyn DurableCatalogState>, CatalogError>;
138
139 async fn open_debug(mut self: Box<Self>) -> Result<DebugCatalogState, CatalogError>;
142
143 async fn is_initialized(&mut self) -> Result<bool, CatalogError>;
145
146 async fn epoch(&mut self) -> Result<Epoch, CatalogError>;
156
157 async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError>;
160
161 async fn get_0dt_deployment_max_wait(&mut self) -> Result<Option<Duration>, CatalogError>;
167
168 async fn get_0dt_deployment_ddl_check_interval(
174 &mut self,
175 ) -> Result<Option<Duration>, CatalogError>;
176
177 async fn get_enable_0dt_deployment_panic_after_timeout(
184 &mut self,
185 ) -> Result<Option<bool>, CatalogError>;
186
187 async fn has_system_config_synced_once(&mut self) -> Result<bool, DurableCatalogError>;
189
190 async fn trace_unconsolidated(&mut self) -> Result<Trace, CatalogError>;
192
193 async fn trace_consolidated(&mut self) -> Result<Trace, CatalogError>;
195
196 async fn expire(self: Box<Self>);
198}
199
200#[async_trait]
202pub trait ReadOnlyDurableCatalogState: Debug + Send + Sync {
203 fn epoch(&self) -> Epoch;
213
214 fn metrics(&self) -> &Metrics;
216
217 async fn expire(self: Box<Self>);
219
220 fn is_bootstrap_complete(&self) -> bool;
222
223 async fn get_audit_logs(&mut self) -> Result<Vec<VersionedEvent>, CatalogError>;
229
230 async fn get_next_id(&mut self, id_type: &str) -> Result<u64, CatalogError>;
232
233 async fn get_next_user_item_id(&mut self) -> Result<u64, CatalogError> {
235 self.get_next_id(USER_ITEM_ALLOC_KEY).await
236 }
237
238 async fn get_next_system_item_id(&mut self) -> Result<u64, CatalogError> {
240 self.get_next_id(SYSTEM_ITEM_ALLOC_KEY).await
241 }
242
243 async fn get_next_system_replica_id(&mut self) -> Result<u64, CatalogError> {
245 self.get_next_id(SYSTEM_REPLICA_ID_ALLOC_KEY).await
246 }
247
248 async fn get_next_user_replica_id(&mut self) -> Result<u64, CatalogError> {
250 self.get_next_id(USER_REPLICA_ID_ALLOC_KEY).await
251 }
252
253 async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError>;
255
256 async fn snapshot(&mut self) -> Result<Snapshot, CatalogError>;
258
259 async fn sync_to_current_updates(
265 &mut self,
266 ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError>;
267
268 async fn sync_updates(
277 &mut self,
278 target_upper: Timestamp,
279 ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError>;
280
281 async fn ensure_not_out_of_sync(&mut self, target_upper: Timestamp)
285 -> Result<(), CatalogError>;
286
287 async fn current_upper(&mut self) -> Timestamp;
289}
290
291#[async_trait]
293#[allow(mismatched_lifetime_syntaxes)]
294pub trait DurableCatalogState: ReadOnlyDurableCatalogState {
295 fn is_read_only(&self) -> bool;
297
298 fn is_savepoint(&self) -> bool;
300
301 async fn mark_bootstrap_complete(&mut self);
303
304 async fn transaction(&mut self) -> Result<Transaction, CatalogError>;
308
309 fn transaction_from_snapshot(
313 &mut self,
314 snapshot: Snapshot,
315 ) -> Result<DryRunTransaction, CatalogError>;
316
317 async fn commit_transaction(
325 &mut self,
326 txn_batch: TransactionBatch,
327 commit_ts: Timestamp,
328 ) -> Result<Timestamp, CatalogError>;
329
330 async fn advance_upper(&mut self, new_upper: Timestamp) -> Result<(), CatalogError>;
336
337 async fn allocate_id(
342 &mut self,
343 id_type: &str,
344 amount: u64,
345 commit_ts: Timestamp,
346 ) -> Result<Vec<u64>, CatalogError>;
347
348 async fn allocate_user_ids(
352 &mut self,
353 amount: u64,
354 commit_ts: Timestamp,
355 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
356 let ids = self
357 .allocate_id(USER_ITEM_ALLOC_KEY, amount, commit_ts)
358 .await?;
359 let ids = ids
360 .iter()
361 .map(|id| (CatalogItemId::User(*id), GlobalId::User(*id)))
362 .collect();
363 Ok(ids)
364 }
365
366 async fn allocate_user_id(
370 &mut self,
371 commit_ts: Timestamp,
372 ) -> Result<(CatalogItemId, GlobalId), CatalogError> {
373 let id = self.allocate_id(USER_ITEM_ALLOC_KEY, 1, commit_ts).await?;
374 let id = id.into_element();
375 Ok((CatalogItemId::User(id), GlobalId::User(id)))
376 }
377
378 async fn allocate_user_cluster_id(
382 &mut self,
383 commit_ts: Timestamp,
384 ) -> Result<ClusterId, CatalogError> {
385 let id = self
386 .allocate_id(USER_CLUSTER_ID_ALLOC_KEY, 1, commit_ts)
387 .await?
388 .into_element();
389 Ok(ClusterId::user(id).ok_or(SqlCatalogError::IdExhaustion)?)
390 }
391
392 async fn allocate_user_replica_ids(
396 &mut self,
397 amount: u64,
398 commit_ts: Timestamp,
399 ) -> Result<Vec<ReplicaId>, CatalogError> {
400 let ids = self
401 .allocate_id(USER_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
402 .await?;
403 let ids = ids.into_iter().map(ReplicaId::User).collect();
404 Ok(ids)
405 }
406
407 async fn allocate_system_replica_ids(
411 &mut self,
412 amount: u64,
413 commit_ts: Timestamp,
414 ) -> Result<Vec<ReplicaId>, CatalogError> {
415 let ids = self
416 .allocate_id(SYSTEM_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
417 .await?;
418 let ids = ids.into_iter().map(ReplicaId::System).collect();
419 Ok(ids)
420 }
421
422 fn shard_id(&self) -> ShardId;
423}
424
425pub fn persist_desc() -> RelationDesc {
428 RelationDesc::builder()
429 .with_column("data", SqlScalarType::Jsonb.nullable(false))
430 .finish()
431}
432
433#[derive(Debug, Clone)]
435pub struct TestCatalogStateBuilder {
436 persist_client: PersistClient,
437 organization_id: Uuid,
438 version: semver::Version,
439 deploy_generation: Option<u64>,
440 metrics: Arc<Metrics>,
441}
442
443impl TestCatalogStateBuilder {
444 pub fn new(persist_client: PersistClient) -> Self {
445 Self {
446 persist_client,
447 organization_id: Uuid::new_v4(),
448 version: semver::Version::new(0, 0, 0),
449 deploy_generation: None,
450 metrics: Arc::new(Metrics::new(&MetricsRegistry::new())),
451 }
452 }
453
454 pub fn with_organization_id(mut self, organization_id: Uuid) -> Self {
455 self.organization_id = organization_id;
456 self
457 }
458
459 pub fn with_version(mut self, version: semver::Version) -> Self {
460 self.version = version;
461 self
462 }
463
464 pub fn with_deploy_generation(mut self, deploy_generation: u64) -> Self {
465 self.deploy_generation = Some(deploy_generation);
466 self
467 }
468
469 pub fn with_default_deploy_generation(self) -> Self {
470 self.with_deploy_generation(0)
471 }
472
473 pub fn with_metrics(mut self, metrics: Arc<Metrics>) -> Self {
474 self.metrics = metrics;
475 self
476 }
477
478 pub async fn build(self) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
479 persist_backed_catalog_state(
480 self.persist_client,
481 self.organization_id,
482 self.version,
483 self.deploy_generation,
484 self.metrics,
485 )
486 .await
487 }
488
489 pub async fn unwrap_build(self) -> Box<dyn OpenableDurableCatalogState> {
490 self.expect_build("failed to build").await
491 }
492
493 pub async fn expect_build(self, msg: &str) -> Box<dyn OpenableDurableCatalogState> {
494 self.build().await.expect(msg)
495 }
496}
497
498pub async fn persist_backed_catalog_state(
502 persist_client: PersistClient,
503 organization_id: Uuid,
504 version: semver::Version,
505 deploy_generation: Option<u64>,
506 metrics: Arc<Metrics>,
507) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
508 let state = UnopenedPersistCatalogState::new(
509 persist_client,
510 organization_id,
511 version,
512 deploy_generation,
513 metrics,
514 )
515 .await?;
516 Ok(Box::new(state))
517}
518
519pub fn test_bootstrap_args() -> BootstrapArgs {
520 BootstrapArgs {
521 default_cluster_replica_size: "scale=1,workers=1".into(),
522 default_cluster_replication_factor: 1,
523 bootstrap_role: None,
524 cluster_replica_size_map: ClusterReplicaSizeMap::for_tests(),
525 }
526}