1use std::fmt::Debug;
13use std::num::NonZeroI64;
14use std::sync::Arc;
15use std::time::{Duration, Instant};
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, ReconfigurationTarget, ReplicaConfig, ReplicaLocation,
38 ReplicaSystemConfiguration, Role, RoleAuth, Schema, SourceReference, SourceReferences,
39 StorageCollectionMetadata, SystemConfiguration, SystemObjectDescription, SystemObjectMapping,
40 UnfinalizedShard,
41};
42pub use crate::durable::persist::shard_id;
43use crate::durable::persist::{Timestamp, UnopenedPersistCatalogState};
44pub use crate::durable::transaction::Transaction;
45use crate::durable::transaction::TransactionBatch;
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 current_upper(&mut self) -> Timestamp;
279}
280
281#[async_trait]
283#[allow(mismatched_lifetime_syntaxes)]
284pub trait DurableCatalogState: ReadOnlyDurableCatalogState {
285 fn is_read_only(&self) -> bool;
287
288 fn is_savepoint(&self) -> bool;
290
291 async fn mark_bootstrap_complete(&mut self);
293
294 async fn transaction(&mut self) -> Result<Transaction, CatalogError>;
296
297 fn transaction_from_snapshot(
303 &mut self,
304 snapshot: Snapshot,
305 ) -> Result<Transaction, CatalogError>;
306
307 async fn commit_transaction(
315 &mut self,
316 txn_batch: TransactionBatch,
317 commit_ts: Timestamp,
318 ) -> Result<Timestamp, CatalogError>;
319
320 async fn advance_upper(&mut self, new_upper: Timestamp) -> Result<(), CatalogError>;
325
326 #[mz_ore::instrument(level = "debug")]
330 async fn allocate_id(
331 &mut self,
332 id_type: &str,
333 amount: u64,
334 commit_ts: Timestamp,
335 ) -> Result<Vec<u64>, CatalogError> {
336 let start = Instant::now();
337 if amount == 0 {
338 return Ok(Vec::new());
339 }
340 let mut txn = self.transaction().await?;
341 let ids = txn.get_and_increment_id_by(id_type.to_string(), amount)?;
342 txn.commit_internal(commit_ts).await?;
343 self.metrics()
344 .allocate_id_seconds
345 .observe(start.elapsed().as_secs_f64());
346 Ok(ids)
347 }
348
349 async fn allocate_user_ids(
353 &mut self,
354 amount: u64,
355 commit_ts: Timestamp,
356 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
357 let ids = self
358 .allocate_id(USER_ITEM_ALLOC_KEY, amount, commit_ts)
359 .await?;
360 let ids = ids
361 .iter()
362 .map(|id| (CatalogItemId::User(*id), GlobalId::User(*id)))
363 .collect();
364 Ok(ids)
365 }
366
367 async fn allocate_user_id(
371 &mut self,
372 commit_ts: Timestamp,
373 ) -> Result<(CatalogItemId, GlobalId), CatalogError> {
374 let id = self.allocate_id(USER_ITEM_ALLOC_KEY, 1, commit_ts).await?;
375 let id = id.into_element();
376 Ok((CatalogItemId::User(id), GlobalId::User(id)))
377 }
378
379 async fn allocate_user_cluster_id(
383 &mut self,
384 commit_ts: Timestamp,
385 ) -> Result<ClusterId, CatalogError> {
386 let id = self
387 .allocate_id(USER_CLUSTER_ID_ALLOC_KEY, 1, commit_ts)
388 .await?
389 .into_element();
390 Ok(ClusterId::user(id).ok_or(SqlCatalogError::IdExhaustion)?)
391 }
392
393 async fn allocate_user_replica_ids(
397 &mut self,
398 amount: u64,
399 commit_ts: Timestamp,
400 ) -> Result<Vec<ReplicaId>, CatalogError> {
401 let ids = self
402 .allocate_id(USER_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
403 .await?;
404 let ids = ids.into_iter().map(ReplicaId::User).collect();
405 Ok(ids)
406 }
407
408 async fn allocate_system_replica_ids(
412 &mut self,
413 amount: u64,
414 commit_ts: Timestamp,
415 ) -> Result<Vec<ReplicaId>, CatalogError> {
416 let ids = self
417 .allocate_id(SYSTEM_REPLICA_ID_ALLOC_KEY, amount, commit_ts)
418 .await?;
419 let ids = ids.into_iter().map(ReplicaId::System).collect();
420 Ok(ids)
421 }
422
423 fn shard_id(&self) -> ShardId;
424}
425
426pub fn persist_desc() -> RelationDesc {
429 RelationDesc::builder()
430 .with_column("data", SqlScalarType::Jsonb.nullable(false))
431 .finish()
432}
433
434#[derive(Debug, Clone)]
436pub struct TestCatalogStateBuilder {
437 persist_client: PersistClient,
438 organization_id: Uuid,
439 version: semver::Version,
440 deploy_generation: Option<u64>,
441 metrics: Arc<Metrics>,
442}
443
444impl TestCatalogStateBuilder {
445 pub fn new(persist_client: PersistClient) -> Self {
446 Self {
447 persist_client,
448 organization_id: Uuid::new_v4(),
449 version: semver::Version::new(0, 0, 0),
450 deploy_generation: None,
451 metrics: Arc::new(Metrics::new(&MetricsRegistry::new())),
452 }
453 }
454
455 pub fn with_organization_id(mut self, organization_id: Uuid) -> Self {
456 self.organization_id = organization_id;
457 self
458 }
459
460 pub fn with_version(mut self, version: semver::Version) -> Self {
461 self.version = version;
462 self
463 }
464
465 pub fn with_deploy_generation(mut self, deploy_generation: u64) -> Self {
466 self.deploy_generation = Some(deploy_generation);
467 self
468 }
469
470 pub fn with_default_deploy_generation(self) -> Self {
471 self.with_deploy_generation(0)
472 }
473
474 pub fn with_metrics(mut self, metrics: Arc<Metrics>) -> Self {
475 self.metrics = metrics;
476 self
477 }
478
479 pub async fn build(self) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
480 persist_backed_catalog_state(
481 self.persist_client,
482 self.organization_id,
483 self.version,
484 self.deploy_generation,
485 self.metrics,
486 )
487 .await
488 }
489
490 pub async fn unwrap_build(self) -> Box<dyn OpenableDurableCatalogState> {
491 self.expect_build("failed to build").await
492 }
493
494 pub async fn expect_build(self, msg: &str) -> Box<dyn OpenableDurableCatalogState> {
495 self.build().await.expect(msg)
496 }
497}
498
499pub async fn persist_backed_catalog_state(
503 persist_client: PersistClient,
504 organization_id: Uuid,
505 version: semver::Version,
506 deploy_generation: Option<u64>,
507 metrics: Arc<Metrics>,
508) -> Result<Box<dyn OpenableDurableCatalogState>, DurableCatalogError> {
509 let state = UnopenedPersistCatalogState::new(
510 persist_client,
511 organization_id,
512 version,
513 deploy_generation,
514 metrics,
515 )
516 .await?;
517 Ok(Box::new(state))
518}
519
520pub fn test_bootstrap_args() -> BootstrapArgs {
521 BootstrapArgs {
522 default_cluster_replica_size: "scale=1,workers=1".into(),
523 default_cluster_replication_factor: 1,
524 bootstrap_role: None,
525 cluster_replica_size_map: ClusterReplicaSizeMap::for_tests(),
526 }
527}