1use std::collections::{BTreeMap, BTreeSet};
11use std::fmt::Debug;
12use std::num::NonZeroU32;
13use std::time::Duration;
14
15use anyhow::anyhow;
16use derivative::Derivative;
17use itertools::Itertools;
18use mz_audit_log::VersionedEvent;
19use mz_compute_client::logging::{ComputeLog, DifferentialLog, LogVariant, TimelyLog};
20use mz_controller_types::{ClusterId, ReplicaId};
21use mz_ore::cast::{u64_to_usize, usize_to_u64};
22use mz_ore::collections::{CollectionExt, HashSet};
23use mz_ore::now::SYSTEM_TIME;
24use mz_ore::vec::VecExt;
25use mz_ore::{soft_assert_no_log, soft_assert_or_log, soft_panic_or_log};
26use mz_persist_types::ShardId;
27use mz_pgrepr::oid::FIRST_USER_OID;
28use mz_proto::{RustType, TryFromProtoError};
29use mz_repr::adt::mz_acl_item::{AclMode, MzAclItem};
30use mz_repr::network_policy_id::NetworkPolicyId;
31use mz_repr::role_id::RoleId;
32use mz_repr::{CatalogItemId, Diff, GlobalId, RelationVersion};
33use mz_sql::catalog::{
34 CatalogError as SqlCatalogError, CatalogItemType, ObjectType, PasswordAction,
35 RoleAttributesRaw, RoleMembership, RoleVars,
36};
37use mz_sql::names::{CommentObjectId, DatabaseId, ResolvedDatabaseSpecifier, SchemaId};
38use mz_sql::plan::NetworkPolicyRule;
39use mz_sql_parser::ast::QualifiedReplica;
40use mz_storage_client::controller::StorageTxn;
41use mz_storage_types::controller::StorageError;
42use tracing::warn;
43
44use crate::builtin::BuiltinLog;
45use crate::durable::initialize::{
46 ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT, SYSTEM_CONFIG_SYNCED_KEY,
47 WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL, WITH_0DT_DEPLOYMENT_MAX_WAIT,
48};
49use crate::durable::objects::serialization::proto;
50use crate::durable::objects::{
51 AuditLogKey, Cluster, ClusterConfig, ClusterIntrospectionSourceIndexKey,
52 ClusterIntrospectionSourceIndexValue, ClusterKey, ClusterReplica, ClusterReplicaKey,
53 ClusterReplicaValue, ClusterSystemConfiguration, ClusterSystemConfigurationKey,
54 ClusterSystemConfigurationValue, ClusterValue, CommentKey, CommentValue, Config, ConfigKey,
55 ConfigValue, Database, DatabaseKey, DatabaseValue, DefaultPrivilegesKey,
56 DefaultPrivilegesValue, DurableType, GidMappingKey, GidMappingValue, IdAllocKey, IdAllocValue,
57 IntrospectionSourceIndex, Item, ItemKey, ItemValue, NetworkPolicyKey, NetworkPolicyValue,
58 ReplicaConfig, ReplicaSystemConfiguration, ReplicaSystemConfigurationKey,
59 ReplicaSystemConfigurationValue, Role, RoleKey, RoleValue, Schema, SchemaKey, SchemaValue,
60 ServerConfigurationKey, ServerConfigurationValue, SettingKey, SettingValue, SourceReference,
61 SourceReferencesKey, SourceReferencesValue, StorageCollectionMetadataKey,
62 StorageCollectionMetadataValue, SystemObjectDescription, SystemObjectMapping,
63 SystemPrivilegesKey, SystemPrivilegesValue, TxnWalShardValue, UnfinalizedShardKey,
64};
65use crate::durable::{
66 AUDIT_LOG_ID_ALLOC_KEY, BUILTIN_MIGRATION_SHARD_KEY, CATALOG_CONTENT_VERSION_KEY, CatalogError,
67 DATABASE_ID_ALLOC_KEY, DefaultPrivilege, DurableCatalogError, DurableCatalogState,
68 EXPRESSION_CACHE_SHARD_KEY, MOCK_AUTHENTICATION_NONCE_KEY, NetworkPolicy, OID_ALLOC_KEY,
69 SCHEMA_ID_ALLOC_KEY, SYSTEM_CLUSTER_ID_ALLOC_KEY, SYSTEM_ITEM_ALLOC_KEY,
70 SYSTEM_REPLICA_ID_ALLOC_KEY, Snapshot, SystemConfiguration, USER_ITEM_ALLOC_KEY,
71 USER_NETWORK_POLICY_ID_ALLOC_KEY, USER_ROLE_ID_ALLOC_KEY,
72};
73use crate::memory::objects::{StateDiff, StateUpdate, StateUpdateKind};
74
75type Timestamp = u64;
76
77#[derive(Debug, PartialEq)]
78struct CommitCapability;
79
80#[derive(Derivative)]
83#[derivative(Debug)]
84pub struct Transaction<'a> {
85 #[derivative(Debug = "ignore")]
86 #[derivative(PartialEq = "ignore")]
87 durable_catalog: &'a mut dyn DurableCatalogState,
88 databases: TableTransaction<DatabaseKey, DatabaseValue>,
89 schemas: TableTransaction<SchemaKey, SchemaValue>,
90 items: TableTransaction<ItemKey, ItemValue>,
91 comments: TableTransaction<CommentKey, CommentValue>,
92 roles: TableTransaction<RoleKey, RoleValue>,
93 role_auth: TableTransaction<RoleAuthKey, RoleAuthValue>,
94 clusters: TableTransaction<ClusterKey, ClusterValue>,
95 cluster_replicas: TableTransaction<ClusterReplicaKey, ClusterReplicaValue>,
96 introspection_sources:
97 TableTransaction<ClusterIntrospectionSourceIndexKey, ClusterIntrospectionSourceIndexValue>,
98 id_allocator: TableTransaction<IdAllocKey, IdAllocValue>,
99 configs: TableTransaction<ConfigKey, ConfigValue>,
100 settings: TableTransaction<SettingKey, SettingValue>,
101 system_gid_mapping: TableTransaction<GidMappingKey, GidMappingValue>,
102 system_configurations: TableTransaction<ServerConfigurationKey, ServerConfigurationValue>,
103 cluster_system_configurations:
104 TableTransaction<ClusterSystemConfigurationKey, ClusterSystemConfigurationValue>,
105 replica_system_configurations:
106 TableTransaction<ReplicaSystemConfigurationKey, ReplicaSystemConfigurationValue>,
107 default_privileges: TableTransaction<DefaultPrivilegesKey, DefaultPrivilegesValue>,
108 source_references: TableTransaction<SourceReferencesKey, SourceReferencesValue>,
109 system_privileges: TableTransaction<SystemPrivilegesKey, SystemPrivilegesValue>,
110 network_policies: TableTransaction<NetworkPolicyKey, NetworkPolicyValue>,
111 storage_collection_metadata:
112 TableTransaction<StorageCollectionMetadataKey, StorageCollectionMetadataValue>,
113 unfinalized_shards: TableTransaction<UnfinalizedShardKey, ()>,
114 txn_wal_shard: TableTransaction<(), TxnWalShardValue>,
115 audit_log_updates: Vec<(AuditLogKey, Diff, Timestamp)>,
118 upper: mz_repr::Timestamp,
120 op_id: Timestamp,
122 commit_capability: Option<CommitCapability>,
126}
127
128#[derive(Debug)]
130pub struct DryRunTransaction<'a> {
131 transaction: Transaction<'a>,
132}
133
134impl<'a> DryRunTransaction<'a> {
135 pub fn new(mut transaction: Transaction<'a>) -> Self {
137 transaction.commit_capability = None;
138 Self { transaction }
139 }
140
141 pub fn transaction_mut(&mut self) -> &mut Transaction<'a> {
143 &mut self.transaction
144 }
145
146 pub fn current_snapshot(&self) -> Snapshot {
148 self.transaction.current_snapshot()
149 }
150}
151
152impl<'a> Transaction<'a> {
153 pub(super) fn new(
154 durable_catalog: &'a mut dyn DurableCatalogState,
155 Snapshot {
156 databases,
157 schemas,
158 roles,
159 role_auth,
160 items,
161 comments,
162 clusters,
163 network_policies,
164 cluster_replicas,
165 introspection_sources,
166 id_allocator,
167 configs,
168 settings,
169 source_references,
170 system_object_mappings,
171 system_configurations,
172 cluster_system_configurations,
173 replica_system_configurations,
174 default_privileges,
175 system_privileges,
176 storage_collection_metadata,
177 unfinalized_shards,
178 txn_wal_shard,
179 }: Snapshot,
180 upper: mz_repr::Timestamp,
181 ) -> Result<Transaction<'a>, CatalogError> {
182 let database_unique_fn: fn(&DatabaseValue, &DatabaseValue) -> bool =
185 |a, b| a.name == b.name;
186 let schema_unique_fn: fn(&SchemaValue, &SchemaValue) -> bool =
187 |a, b| a.database_id == b.database_id && a.name == b.name;
188 let role_key: fn(&RoleValue, &RoleValue) -> bool = |a, b| a.name == b.name;
189 let cluster_unique_fn: fn(&ClusterValue, &ClusterValue) -> bool = |a, b| a.name == b.name;
190 let network_policy_unique_fn: fn(&NetworkPolicyValue, &NetworkPolicyValue) -> bool =
191 |a, b| a.name == b.name;
192 let cluster_replica_unique_fn: fn(&ClusterReplicaValue, &ClusterReplicaValue) -> bool =
193 |a, b| a.cluster_id == b.cluster_id && a.name == b.name;
194
195 Ok(Transaction {
196 durable_catalog,
197 databases: TableTransaction::new_with_uniqueness_fn(
198 databases,
199 database_unique_fn,
200 database_unique_fn,
201 )?,
202 schemas: TableTransaction::new_with_uniqueness_fn(
203 schemas,
204 schema_unique_fn,
205 schema_unique_fn,
206 )?,
207 items: TableTransaction::new_with_uniqueness_fn(
208 items,
209 |a: &ItemValue, b| {
210 a.schema_id == b.schema_id && a.name == b.name && {
211 let a_type = a.item_type();
213 let b_type = b.item_type();
214 (a_type != CatalogItemType::Type && b_type != CatalogItemType::Type)
215 || (a_type == CatalogItemType::Type && b_type.conflicts_with_type())
216 || (b_type == CatalogItemType::Type && a_type.conflicts_with_type())
217 }
218 },
219 |prev: &ItemValue, next| {
220 prev.schema_id == next.schema_id
221 && prev.name == next.name
222 && prev.item_type() == next.item_type()
224 },
225 )?,
226 comments: TableTransaction::new(comments)?,
227 roles: TableTransaction::new_with_uniqueness_fn(roles, role_key, role_key)?,
228 role_auth: TableTransaction::new(role_auth)?,
229 clusters: TableTransaction::new_with_uniqueness_fn(
230 clusters,
231 cluster_unique_fn,
232 cluster_unique_fn,
233 )?,
234 network_policies: TableTransaction::new_with_uniqueness_fn(
235 network_policies,
236 network_policy_unique_fn,
237 network_policy_unique_fn,
238 )?,
239 cluster_replicas: TableTransaction::new_with_uniqueness_fn(
240 cluster_replicas,
241 cluster_replica_unique_fn,
242 cluster_replica_unique_fn,
243 )?,
244 introspection_sources: TableTransaction::new(introspection_sources)?,
245 id_allocator: TableTransaction::new(id_allocator)?,
246 configs: TableTransaction::new(configs)?,
247 settings: TableTransaction::new(settings)?,
248 source_references: TableTransaction::new(source_references)?,
249 system_gid_mapping: TableTransaction::new(system_object_mappings)?,
250 system_configurations: TableTransaction::new(system_configurations)?,
251 cluster_system_configurations: TableTransaction::new(cluster_system_configurations)?,
252 replica_system_configurations: TableTransaction::new(replica_system_configurations)?,
253 default_privileges: TableTransaction::new(default_privileges)?,
254 system_privileges: TableTransaction::new(system_privileges)?,
255 storage_collection_metadata: TableTransaction::new(storage_collection_metadata)?,
256 unfinalized_shards: TableTransaction::new(unfinalized_shards)?,
257 txn_wal_shard: TableTransaction::new(txn_wal_shard)?,
261 audit_log_updates: Vec::new(),
262 upper,
263 op_id: 0,
264 commit_capability: Some(CommitCapability),
265 })
266 }
267
268 pub fn get_item(&self, id: &CatalogItemId) -> Option<Item> {
269 let key = ItemKey { id: *id };
270 self.items
271 .get(&key)
272 .map(|v| DurableType::from_key_value(key, v.clone()))
273 }
274
275 pub fn get_items(&self) -> impl Iterator<Item = Item> + use<> {
276 self.items
277 .items()
278 .into_iter()
279 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
280 .sorted_by_key(|Item { id, .. }| *id)
281 }
282
283 pub fn insert_audit_log_event(&mut self, event: VersionedEvent) {
284 self.insert_audit_log_events([event]);
285 }
286
287 pub fn insert_audit_log_events(&mut self, events: impl IntoIterator<Item = VersionedEvent>) {
288 let events = events
289 .into_iter()
290 .map(|event| (AuditLogKey { event }, Diff::ONE, self.op_id));
291 self.audit_log_updates.extend(events);
292 }
293
294 pub fn insert_user_database(
295 &mut self,
296 database_name: &str,
297 owner_id: RoleId,
298 privileges: Vec<MzAclItem>,
299 temporary_oids: &HashSet<u32>,
300 ) -> Result<(DatabaseId, u32), CatalogError> {
301 let id = self.get_and_increment_id(DATABASE_ID_ALLOC_KEY.to_string())?;
302 let id = DatabaseId::User(id);
303 let oid = self.allocate_oid(temporary_oids)?;
304 self.insert_database(id, database_name, owner_id, privileges, oid)?;
305 Ok((id, oid))
306 }
307
308 pub(crate) fn insert_database(
309 &mut self,
310 id: DatabaseId,
311 database_name: &str,
312 owner_id: RoleId,
313 privileges: Vec<MzAclItem>,
314 oid: u32,
315 ) -> Result<u32, CatalogError> {
316 match self.databases.insert(
317 DatabaseKey { id },
318 DatabaseValue {
319 name: database_name.to_string(),
320 owner_id,
321 privileges,
322 oid,
323 },
324 self.op_id,
325 ) {
326 Ok(_) => Ok(oid),
327 Err(_) => Err(SqlCatalogError::DatabaseAlreadyExists(database_name.to_owned()).into()),
328 }
329 }
330
331 pub fn insert_user_schema(
332 &mut self,
333 database_id: DatabaseId,
334 schema_name: &str,
335 owner_id: RoleId,
336 privileges: Vec<MzAclItem>,
337 temporary_oids: &HashSet<u32>,
338 ) -> Result<(SchemaId, u32), CatalogError> {
339 let id = self.get_and_increment_id(SCHEMA_ID_ALLOC_KEY.to_string())?;
340 let id = SchemaId::User(id);
341 let oid = self.allocate_oid(temporary_oids)?;
342 self.insert_schema(
343 id,
344 Some(database_id),
345 schema_name.to_string(),
346 owner_id,
347 privileges,
348 oid,
349 )?;
350 Ok((id, oid))
351 }
352
353 pub fn insert_system_schema(
354 &mut self,
355 schema_id: u64,
356 schema_name: &str,
357 owner_id: RoleId,
358 privileges: Vec<MzAclItem>,
359 oid: u32,
360 ) -> Result<(), CatalogError> {
361 let id = SchemaId::System(schema_id);
362 self.insert_schema(id, None, schema_name.to_string(), owner_id, privileges, oid)
363 }
364
365 pub(crate) fn insert_schema(
366 &mut self,
367 schema_id: SchemaId,
368 database_id: Option<DatabaseId>,
369 schema_name: String,
370 owner_id: RoleId,
371 privileges: Vec<MzAclItem>,
372 oid: u32,
373 ) -> Result<(), CatalogError> {
374 match self.schemas.insert(
375 SchemaKey { id: schema_id },
376 SchemaValue {
377 database_id,
378 name: schema_name.clone(),
379 owner_id,
380 privileges,
381 oid,
382 },
383 self.op_id,
384 ) {
385 Ok(_) => Ok(()),
386 Err(_) => Err(SqlCatalogError::SchemaAlreadyExists(schema_name).into()),
387 }
388 }
389
390 pub fn insert_builtin_role(
391 &mut self,
392 id: RoleId,
393 name: String,
394 attributes: RoleAttributesRaw,
395 membership: RoleMembership,
396 vars: RoleVars,
397 oid: u32,
398 ) -> Result<RoleId, CatalogError> {
399 soft_assert_or_log!(id.is_builtin(), "ID {id:?} is not builtin");
400 self.insert_role(id, name, attributes, membership, vars, oid)?;
401 Ok(id)
402 }
403
404 pub fn insert_user_role(
405 &mut self,
406 name: String,
407 attributes: RoleAttributesRaw,
408 membership: RoleMembership,
409 vars: RoleVars,
410 temporary_oids: &HashSet<u32>,
411 ) -> Result<(RoleId, u32), CatalogError> {
412 let id = self.get_and_increment_id(USER_ROLE_ID_ALLOC_KEY.to_string())?;
413 let id = RoleId::User(id);
414 let oid = self.allocate_oid(temporary_oids)?;
415 self.insert_role(id, name, attributes, membership, vars, oid)?;
416 Ok((id, oid))
417 }
418
419 fn insert_role(
420 &mut self,
421 id: RoleId,
422 name: String,
423 attributes: RoleAttributesRaw,
424 membership: RoleMembership,
425 vars: RoleVars,
426 oid: u32,
427 ) -> Result<(), CatalogError> {
428 if let Some(ref password) = attributes.password {
429 let hash = mz_auth::hash::scram256_hash(
430 password,
431 &attributes
432 .scram_iterations
433 .or_else(|| {
434 soft_panic_or_log!(
435 "Hash iterations must be set if a password is provided."
436 );
437 None
438 })
439 .unwrap_or_else(|| NonZeroU32::new(600_000).expect("known valid")),
442 )
443 .expect("password hash should be valid");
444 match self.role_auth.insert(
445 RoleAuthKey { role_id: id },
446 RoleAuthValue {
447 password_hash: Some(hash),
448 updated_at: SYSTEM_TIME(),
449 },
450 self.op_id,
451 ) {
452 Ok(_) => {}
453 Err(_) => {
454 return Err(SqlCatalogError::RoleAlreadyExists(name).into());
455 }
456 }
457 }
458
459 match self.roles.insert(
460 RoleKey { id },
461 RoleValue {
462 name: name.clone(),
463 attributes: attributes.into(),
464 membership,
465 vars,
466 oid,
467 },
468 self.op_id,
469 ) {
470 Ok(_) => Ok(()),
471 Err(_) => Err(SqlCatalogError::RoleAlreadyExists(name).into()),
472 }
473 }
474
475 pub fn insert_user_cluster(
477 &mut self,
478 cluster_id: ClusterId,
479 cluster_name: &str,
480 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
481 owner_id: RoleId,
482 privileges: Vec<MzAclItem>,
483 config: ClusterConfig,
484 temporary_oids: &HashSet<u32>,
485 ) -> Result<(), CatalogError> {
486 self.insert_cluster(
487 cluster_id,
488 cluster_name,
489 introspection_source_indexes,
490 owner_id,
491 privileges,
492 config,
493 temporary_oids,
494 )
495 }
496
497 pub fn insert_system_cluster(
499 &mut self,
500 cluster_name: &str,
501 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
502 privileges: Vec<MzAclItem>,
503 owner_id: RoleId,
504 config: ClusterConfig,
505 temporary_oids: &HashSet<u32>,
506 ) -> Result<ClusterId, CatalogError> {
507 let cluster_id = self.get_and_increment_id(SYSTEM_CLUSTER_ID_ALLOC_KEY.to_string())?;
508 let cluster_id = ClusterId::system(cluster_id).ok_or(SqlCatalogError::IdExhaustion)?;
509 self.insert_cluster(
510 cluster_id,
511 cluster_name,
512 introspection_source_indexes,
513 owner_id,
514 privileges,
515 config,
516 temporary_oids,
517 )?;
518 Ok(cluster_id)
519 }
520
521 fn insert_cluster(
522 &mut self,
523 cluster_id: ClusterId,
524 cluster_name: &str,
525 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
526 owner_id: RoleId,
527 privileges: Vec<MzAclItem>,
528 config: ClusterConfig,
529 temporary_oids: &HashSet<u32>,
530 ) -> Result<(), CatalogError> {
531 if let Err(_) = self.clusters.insert(
532 ClusterKey { id: cluster_id },
533 ClusterValue {
534 name: cluster_name.to_string(),
535 owner_id,
536 privileges,
537 config,
538 },
539 self.op_id,
540 ) {
541 return Err(SqlCatalogError::ClusterAlreadyExists(cluster_name.to_owned()).into());
542 };
543
544 let amount = usize_to_u64(introspection_source_indexes.len());
545 let oids = self.allocate_oids(amount, temporary_oids)?;
546 let introspection_source_indexes: Vec<_> = introspection_source_indexes
547 .into_iter()
548 .zip_eq(oids)
549 .map(|((builtin, item_id, index_id), oid)| (builtin, item_id, index_id, oid))
550 .collect();
551 for (builtin, item_id, index_id, oid) in introspection_source_indexes {
552 let introspection_source_index = IntrospectionSourceIndex {
553 cluster_id,
554 name: builtin.name.to_string(),
555 item_id,
556 index_id,
557 oid,
558 };
559 let (key, value) = introspection_source_index.into_key_value();
560 self.introspection_sources
561 .insert(key, value, self.op_id)
562 .expect("no uniqueness violation");
563 }
564
565 Ok(())
566 }
567
568 pub fn rename_cluster(
569 &mut self,
570 cluster_id: ClusterId,
571 cluster_name: &str,
572 cluster_to_name: &str,
573 ) -> Result<(), CatalogError> {
574 let key = ClusterKey { id: cluster_id };
575
576 match self.clusters.update(
577 |k, v| {
578 if *k == key {
579 let mut value = v.clone();
580 value.name = cluster_to_name.to_string();
581 Some(value)
582 } else {
583 None
584 }
585 },
586 self.op_id,
587 )? {
588 Diff::ZERO => Err(SqlCatalogError::UnknownCluster(cluster_name.to_string()).into()),
589 Diff::ONE => Ok(()),
590 n => panic!(
591 "Expected to update single cluster {cluster_name} ({cluster_id}), updated {n}"
592 ),
593 }
594 }
595
596 pub fn rename_cluster_replica(
597 &mut self,
598 replica_id: ReplicaId,
599 replica_name: &QualifiedReplica,
600 replica_to_name: &str,
601 ) -> Result<(), CatalogError> {
602 let key = ClusterReplicaKey { id: replica_id };
603
604 match self.cluster_replicas.update(
605 |k, v| {
606 if *k == key {
607 let mut value = v.clone();
608 value.name = replica_to_name.to_string();
609 Some(value)
610 } else {
611 None
612 }
613 },
614 self.op_id,
615 )? {
616 Diff::ZERO => {
617 Err(SqlCatalogError::UnknownClusterReplica(replica_name.to_string()).into())
618 }
619 Diff::ONE => Ok(()),
620 n => panic!(
621 "Expected to update single cluster replica {replica_name} ({replica_id}), updated {n}"
622 ),
623 }
624 }
625
626 pub fn insert_cluster_replica_with_id(
627 &mut self,
628 cluster_id: ClusterId,
629 replica_id: ReplicaId,
630 replica_name: &str,
631 config: ReplicaConfig,
632 owner_id: RoleId,
633 ) -> Result<(), CatalogError> {
634 if let Err(_) = self.cluster_replicas.insert(
635 ClusterReplicaKey { id: replica_id },
636 ClusterReplicaValue {
637 cluster_id,
638 name: replica_name.into(),
639 config,
640 owner_id,
641 },
642 self.op_id,
643 ) {
644 let cluster = self
645 .clusters
646 .get(&ClusterKey { id: cluster_id })
647 .expect("cluster exists");
648 return Err(SqlCatalogError::DuplicateReplica(
649 replica_name.to_string(),
650 cluster.name.to_string(),
651 )
652 .into());
653 };
654 Ok(())
655 }
656
657 pub fn insert_user_network_policy(
658 &mut self,
659 name: String,
660 rules: Vec<NetworkPolicyRule>,
661 privileges: Vec<MzAclItem>,
662 owner_id: RoleId,
663 temporary_oids: &HashSet<u32>,
664 ) -> Result<NetworkPolicyId, CatalogError> {
665 let oid = self.allocate_oid(temporary_oids)?;
666 let id = self.get_and_increment_id(USER_NETWORK_POLICY_ID_ALLOC_KEY.to_string())?;
667 let id = NetworkPolicyId::User(id);
668 self.insert_network_policy(id, name, rules, privileges, owner_id, oid)
669 }
670
671 pub fn insert_network_policy(
672 &mut self,
673 id: NetworkPolicyId,
674 name: String,
675 rules: Vec<NetworkPolicyRule>,
676 privileges: Vec<MzAclItem>,
677 owner_id: RoleId,
678 oid: u32,
679 ) -> Result<NetworkPolicyId, CatalogError> {
680 match self.network_policies.insert(
681 NetworkPolicyKey { id },
682 NetworkPolicyValue {
683 name: name.clone(),
684 rules,
685 privileges,
686 owner_id,
687 oid,
688 },
689 self.op_id,
690 ) {
691 Ok(_) => Ok(id),
692 Err(_) => Err(SqlCatalogError::NetworkPolicyAlreadyExists(name).into()),
693 }
694 }
695
696 pub fn update_introspection_source_index_gids(
701 &mut self,
702 mappings: impl Iterator<
703 Item = (
704 ClusterId,
705 impl Iterator<Item = (String, CatalogItemId, GlobalId, u32)>,
706 ),
707 >,
708 ) -> Result<(), CatalogError> {
709 for (cluster_id, updates) in mappings {
710 for (name, item_id, index_id, oid) in updates {
711 let introspection_source_index = IntrospectionSourceIndex {
712 cluster_id,
713 name,
714 item_id,
715 index_id,
716 oid,
717 };
718 let (key, value) = introspection_source_index.into_key_value();
719
720 let prev = self
721 .introspection_sources
722 .set(key, Some(value), self.op_id)?;
723 if prev.is_none() {
724 return Err(SqlCatalogError::FailedBuiltinSchemaMigration(format!(
725 "{index_id}"
726 ))
727 .into());
728 }
729 }
730 }
731 Ok(())
732 }
733
734 pub fn insert_user_item(
735 &mut self,
736 id: CatalogItemId,
737 global_id: GlobalId,
738 schema_id: SchemaId,
739 item_name: &str,
740 create_sql: String,
741 owner_id: RoleId,
742 privileges: Vec<MzAclItem>,
743 temporary_oids: &HashSet<u32>,
744 versions: BTreeMap<RelationVersion, GlobalId>,
745 ) -> Result<u32, CatalogError> {
746 let oid = self.allocate_oid(temporary_oids)?;
747 self.insert_item(
748 id, oid, global_id, schema_id, item_name, create_sql, owner_id, privileges, versions,
749 )?;
750 Ok(oid)
751 }
752
753 pub fn insert_item(
754 &mut self,
755 id: CatalogItemId,
756 oid: u32,
757 global_id: GlobalId,
758 schema_id: SchemaId,
759 item_name: &str,
760 create_sql: String,
761 owner_id: RoleId,
762 privileges: Vec<MzAclItem>,
763 extra_versions: BTreeMap<RelationVersion, GlobalId>,
764 ) -> Result<(), CatalogError> {
765 match self.items.insert(
766 ItemKey { id },
767 ItemValue {
768 schema_id,
769 name: item_name.to_string(),
770 create_sql,
771 owner_id,
772 privileges,
773 oid,
774 global_id,
775 extra_versions,
776 },
777 self.op_id,
778 ) {
779 Ok(_) => Ok(()),
780 Err(_) => Err(SqlCatalogError::ItemAlreadyExists(id, item_name.to_owned()).into()),
781 }
782 }
783
784 pub fn get_and_increment_id(&mut self, key: String) -> Result<u64, CatalogError> {
785 Ok(self.get_and_increment_id_by(key, 1)?.into_element())
786 }
787
788 pub fn get_and_increment_id_by(
789 &mut self,
790 key: String,
791 amount: u64,
792 ) -> Result<Vec<u64>, CatalogError> {
793 assert!(
794 key != SYSTEM_ITEM_ALLOC_KEY || !self.durable_catalog.is_bootstrap_complete(),
795 "system item IDs cannot be allocated outside of bootstrap"
796 );
797
798 let current_id = self
799 .id_allocator
800 .items()
801 .get(&IdAllocKey { name: key.clone() })
802 .unwrap_or_else(|| panic!("{key} id allocator missing"))
803 .next_id;
804 let next_id = current_id
805 .checked_add(amount)
806 .ok_or(SqlCatalogError::IdExhaustion)?;
807 let prev = self.id_allocator.set(
808 IdAllocKey { name: key },
809 Some(IdAllocValue { next_id }),
810 self.op_id,
811 )?;
812 assert_eq!(
813 prev,
814 Some(IdAllocValue {
815 next_id: current_id
816 })
817 );
818 Ok((current_id..next_id).collect())
819 }
820
821 pub fn allocate_system_item_ids(
822 &mut self,
823 amount: u64,
824 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
825 assert!(
826 !self.durable_catalog.is_bootstrap_complete(),
827 "we can only allocate system item IDs during bootstrap"
828 );
829 Ok(self
830 .get_and_increment_id_by(SYSTEM_ITEM_ALLOC_KEY.to_string(), amount)?
831 .into_iter()
832 .map(|x| (CatalogItemId::System(x), GlobalId::System(x)))
834 .collect())
835 }
836
837 pub fn allocate_introspection_source_index_id(
869 cluster_id: &ClusterId,
870 log_variant: LogVariant,
871 ) -> (CatalogItemId, GlobalId) {
872 let cluster_variant: u8 = match cluster_id {
873 ClusterId::System(_) => 1,
874 ClusterId::User(_) => 2,
875 };
876 let cluster_id: u64 = cluster_id.inner_id();
877 const CLUSTER_ID_MASK: u64 = 0xFFFF << 48;
878 assert_eq!(
879 CLUSTER_ID_MASK & cluster_id,
880 0,
881 "invalid cluster ID: {cluster_id}"
882 );
883 let log_variant: u8 = match log_variant {
884 LogVariant::Timely(TimelyLog::Operates) => 1,
885 LogVariant::Timely(TimelyLog::Channels) => 2,
886 LogVariant::Timely(TimelyLog::Elapsed) => 3,
887 LogVariant::Timely(TimelyLog::Histogram) => 4,
888 LogVariant::Timely(TimelyLog::Addresses) => 5,
889 LogVariant::Timely(TimelyLog::Parks) => 6,
890 LogVariant::Timely(TimelyLog::MessagesSent) => 7,
891 LogVariant::Timely(TimelyLog::MessagesReceived) => 8,
892 LogVariant::Timely(TimelyLog::Reachability) => 9,
893 LogVariant::Timely(TimelyLog::BatchesSent) => 10,
894 LogVariant::Timely(TimelyLog::BatchesReceived) => 11,
895 LogVariant::Differential(DifferentialLog::ArrangementBatches) => 12,
896 LogVariant::Differential(DifferentialLog::ArrangementRecords) => 13,
897 LogVariant::Differential(DifferentialLog::Sharing) => 14,
898 LogVariant::Differential(DifferentialLog::BatcherRecords) => 15,
899 LogVariant::Differential(DifferentialLog::BatcherSize) => 16,
900 LogVariant::Differential(DifferentialLog::BatcherCapacity) => 17,
901 LogVariant::Differential(DifferentialLog::BatcherAllocations) => 18,
902 LogVariant::Compute(ComputeLog::DataflowCurrent) => 19,
903 LogVariant::Compute(ComputeLog::FrontierCurrent) => 20,
904 LogVariant::Compute(ComputeLog::PeekCurrent) => 21,
905 LogVariant::Compute(ComputeLog::PeekDuration) => 22,
906 LogVariant::Compute(ComputeLog::ImportFrontierCurrent) => 23,
907 LogVariant::Compute(ComputeLog::ArrangementHeapSize) => 24,
908 LogVariant::Compute(ComputeLog::ArrangementHeapCapacity) => 25,
909 LogVariant::Compute(ComputeLog::ArrangementHeapAllocations) => 26,
910 LogVariant::Compute(ComputeLog::ErrorCount) => 28,
911 LogVariant::Compute(ComputeLog::HydrationTime) => 29,
912 LogVariant::Compute(ComputeLog::LirMapping) => 30,
913 LogVariant::Compute(ComputeLog::DataflowGlobal) => 31,
914 LogVariant::Compute(ComputeLog::OperatorHydrationStatus) => 32,
915 LogVariant::Compute(ComputeLog::PrometheusMetrics) => 33,
916 };
917
918 let mut id: u64 = u64::from(cluster_variant) << 56;
919 id |= cluster_id << 8;
920 id |= u64::from(log_variant);
921
922 (
923 CatalogItemId::IntrospectionSourceIndex(id),
924 GlobalId::IntrospectionSourceIndex(id),
925 )
926 }
927
928 pub fn allocate_user_item_ids(
929 &mut self,
930 amount: u64,
931 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
932 Ok(self
933 .get_and_increment_id_by(USER_ITEM_ALLOC_KEY.to_string(), amount)?
934 .into_iter()
935 .map(|x| (CatalogItemId::User(x), GlobalId::User(x)))
937 .collect())
938 }
939
940 pub fn allocate_system_replica_id(&mut self) -> Result<ReplicaId, CatalogError> {
941 let id = self.get_and_increment_id(SYSTEM_REPLICA_ID_ALLOC_KEY.to_string())?;
942 Ok(ReplicaId::System(id))
943 }
944
945 pub fn allocate_audit_log_id(&mut self) -> Result<u64, CatalogError> {
946 self.get_and_increment_id(AUDIT_LOG_ID_ALLOC_KEY.to_string())
947 }
948
949 #[mz_ore::instrument]
952 fn allocate_oids(
953 &mut self,
954 amount: u64,
955 temporary_oids: &HashSet<u32>,
956 ) -> Result<Vec<u32>, CatalogError> {
957 struct UserOid(u32);
960
961 impl UserOid {
962 fn new(oid: u32) -> Result<UserOid, anyhow::Error> {
963 if oid < FIRST_USER_OID {
964 Err(anyhow!("invalid user OID {oid}"))
965 } else {
966 Ok(UserOid(oid))
967 }
968 }
969 }
970
971 impl std::ops::AddAssign<u32> for UserOid {
972 fn add_assign(&mut self, rhs: u32) {
973 let (res, overflow) = self.0.overflowing_add(rhs);
974 self.0 = if overflow { FIRST_USER_OID + res } else { res };
975 }
976 }
977
978 if amount > u32::MAX.into() {
979 return Err(CatalogError::Catalog(SqlCatalogError::OidExhaustion));
980 }
981
982 let mut allocated_oids = HashSet::with_capacity(
988 self.databases.len()
989 + self.schemas.len()
990 + self.roles.len()
991 + self.items.len()
992 + self.introspection_sources.len()
993 + temporary_oids.len(),
994 );
995 self.databases.for_values(|_, value| {
996 allocated_oids.insert(value.oid);
997 });
998 self.schemas.for_values(|_, value| {
999 allocated_oids.insert(value.oid);
1000 });
1001 self.roles.for_values(|_, value| {
1002 allocated_oids.insert(value.oid);
1003 });
1004 self.items.for_values(|_, value| {
1005 allocated_oids.insert(value.oid);
1006 });
1007 self.introspection_sources.for_values(|_, value| {
1008 allocated_oids.insert(value.oid);
1009 });
1010
1011 let is_allocated = |oid| allocated_oids.contains(&oid) || temporary_oids.contains(&oid);
1012
1013 let start_oid: u32 = self
1014 .id_allocator
1015 .items()
1016 .get(&IdAllocKey {
1017 name: OID_ALLOC_KEY.to_string(),
1018 })
1019 .unwrap_or_else(|| panic!("{OID_ALLOC_KEY} id allocator missing"))
1020 .next_id
1021 .try_into()
1022 .expect("we should never persist an oid outside of the u32 range");
1023 let mut current_oid = UserOid::new(start_oid)
1024 .expect("we should never persist an oid outside of user OID range");
1025 let mut oids = Vec::new();
1026 while oids.len() < u64_to_usize(amount) {
1027 if !is_allocated(current_oid.0) {
1028 oids.push(current_oid.0);
1029 }
1030 current_oid += 1;
1031
1032 if current_oid.0 == start_oid && oids.len() < u64_to_usize(amount) {
1033 return Err(CatalogError::Catalog(SqlCatalogError::OidExhaustion));
1035 }
1036 }
1037
1038 let next_id = current_oid.0;
1039 let prev = self.id_allocator.set(
1040 IdAllocKey {
1041 name: OID_ALLOC_KEY.to_string(),
1042 },
1043 Some(IdAllocValue {
1044 next_id: next_id.into(),
1045 }),
1046 self.op_id,
1047 )?;
1048 assert_eq!(
1049 prev,
1050 Some(IdAllocValue {
1051 next_id: start_oid.into(),
1052 })
1053 );
1054
1055 Ok(oids)
1056 }
1057
1058 pub fn allocate_oid(&mut self, temporary_oids: &HashSet<u32>) -> Result<u32, CatalogError> {
1061 self.allocate_oids(1, temporary_oids)
1062 .map(|oids| oids.into_element())
1063 }
1064
1065 pub fn current_snapshot(&self) -> Snapshot {
1073 Snapshot {
1074 databases: self.databases.current_items_proto(),
1075 schemas: self.schemas.current_items_proto(),
1076 roles: self.roles.current_items_proto(),
1077 role_auth: self.role_auth.current_items_proto(),
1078 items: self.items.current_items_proto(),
1079 comments: self.comments.current_items_proto(),
1080 clusters: self.clusters.current_items_proto(),
1081 network_policies: self.network_policies.current_items_proto(),
1082 cluster_replicas: self.cluster_replicas.current_items_proto(),
1083 introspection_sources: self.introspection_sources.current_items_proto(),
1084 id_allocator: self.id_allocator.current_items_proto(),
1085 configs: self.configs.current_items_proto(),
1086 settings: self.settings.current_items_proto(),
1087 system_object_mappings: self.system_gid_mapping.current_items_proto(),
1088 system_configurations: self.system_configurations.current_items_proto(),
1089 cluster_system_configurations: self.cluster_system_configurations.current_items_proto(),
1090 replica_system_configurations: self.replica_system_configurations.current_items_proto(),
1091 default_privileges: self.default_privileges.current_items_proto(),
1092 source_references: self.source_references.current_items_proto(),
1093 system_privileges: self.system_privileges.current_items_proto(),
1094 storage_collection_metadata: self.storage_collection_metadata.current_items_proto(),
1095 unfinalized_shards: self.unfinalized_shards.current_items_proto(),
1096 txn_wal_shard: self.txn_wal_shard.current_items_proto(),
1097 }
1098 }
1099
1100 pub(crate) fn insert_id_allocator(
1101 &mut self,
1102 name: String,
1103 next_id: u64,
1104 ) -> Result<(), CatalogError> {
1105 match self.id_allocator.insert(
1106 IdAllocKey { name: name.clone() },
1107 IdAllocValue { next_id },
1108 self.op_id,
1109 ) {
1110 Ok(_) => Ok(()),
1111 Err(_) => Err(SqlCatalogError::IdAllocatorAlreadyExists(name).into()),
1112 }
1113 }
1114
1115 pub fn remove_database(&mut self, id: &DatabaseId) -> Result<(), CatalogError> {
1122 let prev = self
1123 .databases
1124 .set(DatabaseKey { id: *id }, None, self.op_id)?;
1125 if prev.is_some() {
1126 Ok(())
1127 } else {
1128 Err(SqlCatalogError::UnknownDatabase(id.to_string()).into())
1129 }
1130 }
1131
1132 pub fn remove_databases(
1139 &mut self,
1140 databases: &BTreeSet<DatabaseId>,
1141 ) -> Result<(), CatalogError> {
1142 if databases.is_empty() {
1143 return Ok(());
1144 }
1145
1146 let to_remove = databases
1147 .iter()
1148 .map(|id| (DatabaseKey { id: *id }, None))
1149 .collect();
1150 let mut prev = self.databases.set_many(to_remove, self.op_id)?;
1151 prev.retain(|_k, val| val.is_none());
1152
1153 if !prev.is_empty() {
1154 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1155 return Err(SqlCatalogError::UnknownDatabase(err).into());
1156 }
1157
1158 Ok(())
1159 }
1160
1161 pub fn remove_schema(
1168 &mut self,
1169 database_id: &Option<DatabaseId>,
1170 schema_id: &SchemaId,
1171 ) -> Result<(), CatalogError> {
1172 let prev = self
1173 .schemas
1174 .set(SchemaKey { id: *schema_id }, None, self.op_id)?;
1175 if prev.is_some() {
1176 Ok(())
1177 } else {
1178 let database_name = match database_id {
1179 Some(id) => format!("{id}."),
1180 None => "".to_string(),
1181 };
1182 Err(SqlCatalogError::UnknownSchema(format!("{}.{}", database_name, schema_id)).into())
1183 }
1184 }
1185
1186 pub fn remove_schemas(
1193 &mut self,
1194 schemas: &BTreeMap<SchemaId, ResolvedDatabaseSpecifier>,
1195 ) -> Result<(), CatalogError> {
1196 if schemas.is_empty() {
1197 return Ok(());
1198 }
1199
1200 let to_remove = schemas
1201 .iter()
1202 .map(|(schema_id, _)| (SchemaKey { id: *schema_id }, None))
1203 .collect();
1204 let mut prev = self.schemas.set_many(to_remove, self.op_id)?;
1205 prev.retain(|_k, v| v.is_none());
1206
1207 if !prev.is_empty() {
1208 let err = prev
1209 .keys()
1210 .map(|k| {
1211 let db_spec = schemas.get(&k.id).expect("should_exist");
1212 let db_name = match db_spec {
1213 ResolvedDatabaseSpecifier::Id(id) => format!("{id}."),
1214 ResolvedDatabaseSpecifier::Ambient => "".to_string(),
1215 };
1216 format!("{}.{}", db_name, k.id)
1217 })
1218 .join(", ");
1219
1220 return Err(SqlCatalogError::UnknownSchema(err).into());
1221 }
1222
1223 Ok(())
1224 }
1225
1226 pub fn remove_source_references(
1227 &mut self,
1228 source_id: CatalogItemId,
1229 ) -> Result<(), CatalogError> {
1230 let deleted = self
1231 .source_references
1232 .delete_by_key(SourceReferencesKey { source_id }, self.op_id)
1233 .is_some();
1234 if deleted {
1235 Ok(())
1236 } else {
1237 Err(SqlCatalogError::UnknownItem(source_id.to_string()).into())
1238 }
1239 }
1240
1241 pub fn remove_user_roles(&mut self, roles: &BTreeSet<RoleId>) -> Result<(), CatalogError> {
1248 assert!(
1249 roles.iter().all(|id| id.is_user()),
1250 "cannot delete non-user roles"
1251 );
1252 self.remove_roles(roles)
1253 }
1254
1255 pub fn remove_roles(&mut self, roles: &BTreeSet<RoleId>) -> Result<(), CatalogError> {
1262 if roles.is_empty() {
1263 return Ok(());
1264 }
1265
1266 let to_remove_keys = roles
1267 .iter()
1268 .map(|role_id| RoleKey { id: *role_id })
1269 .collect::<Vec<_>>();
1270
1271 let to_remove_roles = to_remove_keys
1272 .iter()
1273 .map(|role_key| (role_key.clone(), None))
1274 .collect();
1275
1276 let mut prev = self.roles.set_many(to_remove_roles, self.op_id)?;
1277
1278 let to_remove_role_auth = to_remove_keys
1279 .iter()
1280 .map(|role_key| {
1281 (
1282 RoleAuthKey {
1283 role_id: role_key.id,
1284 },
1285 None,
1286 )
1287 })
1288 .collect();
1289
1290 let mut role_auth_prev = self.role_auth.set_many(to_remove_role_auth, self.op_id)?;
1291
1292 prev.retain(|_k, v| v.is_none());
1293 if !prev.is_empty() {
1294 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1295 return Err(SqlCatalogError::UnknownRole(err).into());
1296 }
1297
1298 role_auth_prev.retain(|_k, v| v.is_none());
1299 Ok(())
1303 }
1304
1305 pub fn remove_clusters(&mut self, clusters: &BTreeSet<ClusterId>) -> Result<(), CatalogError> {
1312 if clusters.is_empty() {
1313 return Ok(());
1314 }
1315
1316 let to_remove = clusters
1317 .iter()
1318 .map(|cluster_id| (ClusterKey { id: *cluster_id }, None))
1319 .collect();
1320 let mut prev = self.clusters.set_many(to_remove, self.op_id)?;
1321
1322 prev.retain(|_k, v| v.is_none());
1323 if !prev.is_empty() {
1324 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1325 return Err(SqlCatalogError::UnknownCluster(err).into());
1326 }
1327
1328 self.cluster_replicas
1334 .delete(|_k, v| clusters.contains(&v.cluster_id), self.op_id);
1335 self.introspection_sources
1336 .delete(|k, _v| clusters.contains(&k.cluster_id), self.op_id);
1337
1338 Ok(())
1339 }
1340
1341 pub fn remove_cluster_replica(&mut self, id: ReplicaId) -> Result<(), CatalogError> {
1348 let deleted = self
1349 .cluster_replicas
1350 .delete_by_key(ClusterReplicaKey { id }, self.op_id)
1351 .is_some();
1352 if deleted {
1353 Ok(())
1354 } else {
1355 Err(SqlCatalogError::UnknownClusterReplica(id.to_string()).into())
1356 }
1357 }
1358
1359 pub fn remove_cluster_replicas(
1366 &mut self,
1367 replicas: &BTreeSet<ReplicaId>,
1368 ) -> Result<(), CatalogError> {
1369 if replicas.is_empty() {
1370 return Ok(());
1371 }
1372
1373 let to_remove = replicas
1374 .iter()
1375 .map(|replica_id| (ClusterReplicaKey { id: *replica_id }, None))
1376 .collect();
1377 let mut prev = self.cluster_replicas.set_many(to_remove, self.op_id)?;
1378
1379 prev.retain(|_k, v| v.is_none());
1380 if !prev.is_empty() {
1381 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1382 return Err(SqlCatalogError::UnknownClusterReplica(err).into());
1383 }
1384
1385 Ok(())
1386 }
1387
1388 pub fn remove_item(&mut self, id: CatalogItemId) -> Result<(), CatalogError> {
1395 let prev = self.items.set(ItemKey { id }, None, self.op_id)?;
1396 if prev.is_some() {
1397 Ok(())
1398 } else {
1399 Err(SqlCatalogError::UnknownItem(id.to_string()).into())
1400 }
1401 }
1402
1403 pub fn remove_items(&mut self, ids: &BTreeSet<CatalogItemId>) -> Result<(), CatalogError> {
1410 if ids.is_empty() {
1411 return Ok(());
1412 }
1413
1414 let ks: Vec<_> = ids.clone().into_iter().map(|id| ItemKey { id }).collect();
1415 let n = self.items.delete_by_keys(ks, self.op_id).len();
1416 if n == ids.len() {
1417 Ok(())
1418 } else {
1419 let item_ids = self.items.items().keys().map(|k| k.id).collect();
1420 let mut unknown = ids.difference(&item_ids);
1421 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1422 }
1423 }
1424
1425 pub fn remove_system_object_mappings(
1432 &mut self,
1433 descriptions: BTreeSet<SystemObjectDescription>,
1434 ) -> Result<(), CatalogError> {
1435 if descriptions.is_empty() {
1436 return Ok(());
1437 }
1438
1439 let ks: Vec<_> = descriptions
1440 .clone()
1441 .into_iter()
1442 .map(|desc| GidMappingKey {
1443 schema_name: desc.schema_name,
1444 object_type: desc.object_type,
1445 object_name: desc.object_name,
1446 })
1447 .collect();
1448 let n = self.system_gid_mapping.delete_by_keys(ks, self.op_id).len();
1449
1450 if n == descriptions.len() {
1451 Ok(())
1452 } else {
1453 let item_descriptions = self
1454 .system_gid_mapping
1455 .items()
1456 .keys()
1457 .map(|k| SystemObjectDescription {
1458 schema_name: k.schema_name.clone(),
1459 object_type: k.object_type.clone(),
1460 object_name: k.object_name.clone(),
1461 })
1462 .collect();
1463 let mut unknown = descriptions.difference(&item_descriptions).map(|desc| {
1464 format!(
1465 "{} {}.{}",
1466 desc.object_type, desc.schema_name, desc.object_name
1467 )
1468 });
1469 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1470 }
1471 }
1472
1473 pub fn remove_introspection_source_indexes(
1480 &mut self,
1481 introspection_source_indexes: BTreeSet<(ClusterId, String)>,
1482 ) -> Result<(), CatalogError> {
1483 if introspection_source_indexes.is_empty() {
1484 return Ok(());
1485 }
1486
1487 let ks: Vec<_> = introspection_source_indexes
1488 .clone()
1489 .into_iter()
1490 .map(|(cluster_id, name)| ClusterIntrospectionSourceIndexKey { cluster_id, name })
1491 .collect();
1492 let n = self
1493 .introspection_sources
1494 .delete_by_keys(ks, self.op_id)
1495 .len();
1496 if n == introspection_source_indexes.len() {
1497 Ok(())
1498 } else {
1499 let txn_indexes = self
1500 .introspection_sources
1501 .items()
1502 .keys()
1503 .map(|k| (k.cluster_id, k.name.clone()))
1504 .collect();
1505 let mut unknown = introspection_source_indexes
1506 .difference(&txn_indexes)
1507 .map(|(cluster_id, name)| format!("{cluster_id} {name}"));
1508 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1509 }
1510 }
1511
1512 pub fn update_item(&mut self, id: CatalogItemId, item: Item) -> Result<(), CatalogError> {
1519 let updated =
1520 self.items
1521 .update_by_key(ItemKey { id }, item.into_key_value().1, self.op_id)?;
1522 if updated {
1523 Ok(())
1524 } else {
1525 Err(SqlCatalogError::UnknownItem(id.to_string()).into())
1526 }
1527 }
1528
1529 pub fn update_items(
1537 &mut self,
1538 items: BTreeMap<CatalogItemId, Item>,
1539 ) -> Result<(), CatalogError> {
1540 if items.is_empty() {
1541 return Ok(());
1542 }
1543
1544 let update_ids: BTreeSet<_> = items.keys().cloned().collect();
1545 let kvs: Vec<_> = items
1546 .clone()
1547 .into_iter()
1548 .map(|(id, item)| (ItemKey { id }, item.into_key_value().1))
1549 .collect();
1550 let n = self.items.update_by_keys(kvs, self.op_id)?;
1551 let n = usize::try_from(n.into_inner()).expect("Must be positive and fit in usize");
1552 if n == update_ids.len() {
1553 Ok(())
1554 } else {
1555 let item_ids: BTreeSet<_> = self.items.items().keys().map(|k| k.id).collect();
1556 let mut unknown = update_ids.difference(&item_ids);
1557 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1558 }
1559 }
1560
1561 pub fn update_role(
1569 &mut self,
1570 id: RoleId,
1571 role: Role,
1572 password: PasswordAction,
1573 ) -> Result<(), CatalogError> {
1574 let key = RoleKey { id };
1575 if self.roles.get(&key).is_some() {
1576 let auth_key = RoleAuthKey { role_id: id };
1577
1578 match password {
1579 PasswordAction::Set(new_password) => {
1580 let hash = mz_auth::hash::scram256_hash(
1581 &new_password.password,
1582 &new_password.scram_iterations,
1583 )
1584 .expect("password hash should be valid");
1585 let value = RoleAuthValue {
1586 password_hash: Some(hash),
1587 updated_at: SYSTEM_TIME(),
1588 };
1589
1590 if self.role_auth.get(&auth_key).is_some() {
1591 self.role_auth
1592 .update_by_key(auth_key.clone(), value, self.op_id)?;
1593 } else {
1594 self.role_auth.insert(auth_key.clone(), value, self.op_id)?;
1595 }
1596 }
1597 PasswordAction::Clear => {
1598 let value = RoleAuthValue {
1599 password_hash: None,
1600 updated_at: SYSTEM_TIME(),
1601 };
1602 if self.role_auth.get(&auth_key).is_some() {
1603 self.role_auth
1604 .update_by_key(auth_key.clone(), value, self.op_id)?;
1605 }
1606 }
1607 PasswordAction::NoChange => {}
1608 }
1609
1610 self.roles
1611 .update_by_key(key, role.into_key_value().1, self.op_id)?;
1612
1613 Ok(())
1614 } else {
1615 Err(SqlCatalogError::UnknownRole(id.to_string()).into())
1616 }
1617 }
1618
1619 pub fn update_roles_without_auth(
1630 &mut self,
1631 roles: BTreeMap<RoleId, Role>,
1632 ) -> Result<(), CatalogError> {
1633 if roles.is_empty() {
1634 return Ok(());
1635 }
1636
1637 let update_role_ids: BTreeSet<_> = roles.keys().cloned().collect();
1638 let kvs: Vec<_> = roles
1639 .into_iter()
1640 .map(|(id, role)| (RoleKey { id }, role.into_key_value().1))
1641 .collect();
1642 let n = self.roles.update_by_keys(kvs, self.op_id)?;
1643 let n = usize::try_from(n.into_inner()).expect("Must be positive and fit in usize");
1644
1645 if n == update_role_ids.len() {
1646 Ok(())
1647 } else {
1648 let role_ids: BTreeSet<_> = self.roles.items().keys().map(|k| k.id).collect();
1649 let mut unknown = update_role_ids.difference(&role_ids);
1650 Err(SqlCatalogError::UnknownRole(unknown.join(", ")).into())
1651 }
1652 }
1653
1654 pub fn update_system_object_mappings(
1659 &mut self,
1660 mappings: BTreeMap<CatalogItemId, SystemObjectMapping>,
1661 ) -> Result<(), CatalogError> {
1662 if mappings.is_empty() {
1663 return Ok(());
1664 }
1665
1666 let n = self.system_gid_mapping.update(
1667 |_k, v| {
1668 if let Some(mapping) = mappings.get(&CatalogItemId::from(v.catalog_id)) {
1669 let (_, new_value) = mapping.clone().into_key_value();
1670 Some(new_value)
1671 } else {
1672 None
1673 }
1674 },
1675 self.op_id,
1676 )?;
1677
1678 if usize::try_from(n.into_inner()).expect("update diff should fit into usize")
1679 != mappings.len()
1680 {
1681 let id_str = mappings.keys().map(|id| id.to_string()).join(",");
1682 return Err(SqlCatalogError::FailedBuiltinSchemaMigration(id_str).into());
1683 }
1684
1685 Ok(())
1686 }
1687
1688 pub fn update_cluster(&mut self, id: ClusterId, cluster: Cluster) -> Result<(), CatalogError> {
1695 let updated = self.clusters.update_by_key(
1696 ClusterKey { id },
1697 cluster.into_key_value().1,
1698 self.op_id,
1699 )?;
1700 if updated {
1701 Ok(())
1702 } else {
1703 Err(SqlCatalogError::UnknownCluster(id.to_string()).into())
1704 }
1705 }
1706
1707 pub fn update_cluster_replica(
1714 &mut self,
1715 replica_id: ReplicaId,
1716 replica: ClusterReplica,
1717 ) -> Result<(), CatalogError> {
1718 let updated = self.cluster_replicas.update_by_key(
1719 ClusterReplicaKey { id: replica_id },
1720 replica.into_key_value().1,
1721 self.op_id,
1722 )?;
1723 if updated {
1724 Ok(())
1725 } else {
1726 Err(SqlCatalogError::UnknownClusterReplica(replica_id.to_string()).into())
1727 }
1728 }
1729
1730 pub fn update_database(
1737 &mut self,
1738 id: DatabaseId,
1739 database: Database,
1740 ) -> Result<(), CatalogError> {
1741 let updated = self.databases.update_by_key(
1742 DatabaseKey { id },
1743 database.into_key_value().1,
1744 self.op_id,
1745 )?;
1746 if updated {
1747 Ok(())
1748 } else {
1749 Err(SqlCatalogError::UnknownDatabase(id.to_string()).into())
1750 }
1751 }
1752
1753 pub fn update_schema(
1760 &mut self,
1761 schema_id: SchemaId,
1762 schema: Schema,
1763 ) -> Result<(), CatalogError> {
1764 let updated = self.schemas.update_by_key(
1765 SchemaKey { id: schema_id },
1766 schema.into_key_value().1,
1767 self.op_id,
1768 )?;
1769 if updated {
1770 Ok(())
1771 } else {
1772 Err(SqlCatalogError::UnknownSchema(schema_id.to_string()).into())
1773 }
1774 }
1775
1776 pub fn update_network_policy(
1783 &mut self,
1784 id: NetworkPolicyId,
1785 network_policy: NetworkPolicy,
1786 ) -> Result<(), CatalogError> {
1787 let updated = self.network_policies.update_by_key(
1788 NetworkPolicyKey { id },
1789 network_policy.into_key_value().1,
1790 self.op_id,
1791 )?;
1792 if updated {
1793 Ok(())
1794 } else {
1795 Err(SqlCatalogError::UnknownNetworkPolicy(id.to_string()).into())
1796 }
1797 }
1798 pub fn remove_network_policies(
1805 &mut self,
1806 network_policies: &BTreeSet<NetworkPolicyId>,
1807 ) -> Result<(), CatalogError> {
1808 if network_policies.is_empty() {
1809 return Ok(());
1810 }
1811
1812 let to_remove = network_policies
1813 .iter()
1814 .map(|policy_id| (NetworkPolicyKey { id: *policy_id }, None))
1815 .collect();
1816 let mut prev = self.network_policies.set_many(to_remove, self.op_id)?;
1817 assert!(
1818 prev.iter().all(|(k, _)| k.id.is_user()),
1819 "cannot delete non-user network policy"
1820 );
1821
1822 prev.retain(|_k, v| v.is_none());
1823 if !prev.is_empty() {
1824 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1825 return Err(SqlCatalogError::UnknownNetworkPolicy(err).into());
1826 }
1827
1828 Ok(())
1829 }
1830 pub fn set_default_privilege(
1834 &mut self,
1835 role_id: RoleId,
1836 database_id: Option<DatabaseId>,
1837 schema_id: Option<SchemaId>,
1838 object_type: ObjectType,
1839 grantee: RoleId,
1840 privileges: Option<AclMode>,
1841 ) -> Result<(), CatalogError> {
1842 self.default_privileges.set(
1843 DefaultPrivilegesKey {
1844 role_id,
1845 database_id,
1846 schema_id,
1847 object_type,
1848 grantee,
1849 },
1850 privileges.map(|privileges| DefaultPrivilegesValue { privileges }),
1851 self.op_id,
1852 )?;
1853 Ok(())
1854 }
1855
1856 pub fn set_default_privileges(
1858 &mut self,
1859 default_privileges: Vec<DefaultPrivilege>,
1860 ) -> Result<(), CatalogError> {
1861 if default_privileges.is_empty() {
1862 return Ok(());
1863 }
1864
1865 let default_privileges = default_privileges
1866 .into_iter()
1867 .map(DurableType::into_key_value)
1868 .map(|(k, v)| (k, Some(v)))
1869 .collect();
1870 self.default_privileges
1871 .set_many(default_privileges, self.op_id)?;
1872 Ok(())
1873 }
1874
1875 pub fn set_system_privilege(
1879 &mut self,
1880 grantee: RoleId,
1881 grantor: RoleId,
1882 acl_mode: Option<AclMode>,
1883 ) -> Result<(), CatalogError> {
1884 self.system_privileges.set(
1885 SystemPrivilegesKey { grantee, grantor },
1886 acl_mode.map(|acl_mode| SystemPrivilegesValue { acl_mode }),
1887 self.op_id,
1888 )?;
1889 Ok(())
1890 }
1891
1892 pub fn set_system_privileges(
1894 &mut self,
1895 system_privileges: Vec<MzAclItem>,
1896 ) -> Result<(), CatalogError> {
1897 if system_privileges.is_empty() {
1898 return Ok(());
1899 }
1900
1901 let system_privileges = system_privileges
1902 .into_iter()
1903 .map(DurableType::into_key_value)
1904 .map(|(k, v)| (k, Some(v)))
1905 .collect();
1906 self.system_privileges
1907 .set_many(system_privileges, self.op_id)?;
1908 Ok(())
1909 }
1910
1911 pub fn set_setting(&mut self, name: String, value: Option<String>) -> Result<(), CatalogError> {
1913 self.settings.set(
1914 SettingKey { name },
1915 value.map(|value| SettingValue { value }),
1916 self.op_id,
1917 )?;
1918 Ok(())
1919 }
1920
1921 pub fn set_catalog_content_version(&mut self, version: String) -> Result<(), CatalogError> {
1922 self.set_setting(CATALOG_CONTENT_VERSION_KEY.to_string(), Some(version))
1923 }
1924
1925 pub fn insert_introspection_source_indexes(
1927 &mut self,
1928 introspection_source_indexes: Vec<(ClusterId, String, CatalogItemId, GlobalId)>,
1929 temporary_oids: &HashSet<u32>,
1930 ) -> Result<(), CatalogError> {
1931 if introspection_source_indexes.is_empty() {
1932 return Ok(());
1933 }
1934
1935 let amount = usize_to_u64(introspection_source_indexes.len());
1936 let oids = self.allocate_oids(amount, temporary_oids)?;
1937 let introspection_source_indexes: Vec<_> = introspection_source_indexes
1938 .into_iter()
1939 .zip_eq(oids)
1940 .map(
1941 |((cluster_id, name, item_id, index_id), oid)| IntrospectionSourceIndex {
1942 cluster_id,
1943 name,
1944 item_id,
1945 index_id,
1946 oid,
1947 },
1948 )
1949 .collect();
1950
1951 for introspection_source_index in introspection_source_indexes {
1952 let (key, value) = introspection_source_index.into_key_value();
1953 self.introspection_sources.insert(key, value, self.op_id)?;
1954 }
1955
1956 Ok(())
1957 }
1958
1959 pub fn set_system_object_mappings(
1961 &mut self,
1962 mappings: Vec<SystemObjectMapping>,
1963 ) -> Result<(), CatalogError> {
1964 if mappings.is_empty() {
1965 return Ok(());
1966 }
1967
1968 let mappings = mappings
1969 .into_iter()
1970 .map(DurableType::into_key_value)
1971 .map(|(k, v)| (k, Some(v)))
1972 .collect();
1973 self.system_gid_mapping.set_many(mappings, self.op_id)?;
1974 Ok(())
1975 }
1976
1977 pub fn set_replicas(&mut self, replicas: Vec<ClusterReplica>) -> Result<(), CatalogError> {
1979 if replicas.is_empty() {
1980 return Ok(());
1981 }
1982
1983 let replicas = replicas
1984 .into_iter()
1985 .map(DurableType::into_key_value)
1986 .map(|(k, v)| (k, Some(v)))
1987 .collect();
1988 self.cluster_replicas.set_many(replicas, self.op_id)?;
1989 Ok(())
1990 }
1991
1992 pub fn set_config(&mut self, key: String, value: Option<u64>) -> Result<(), CatalogError> {
1994 match value {
1995 Some(value) => {
1996 let config = Config { key, value };
1997 let (key, value) = config.into_key_value();
1998 self.configs.set(key, Some(value), self.op_id)?;
1999 }
2000 None => {
2001 self.configs.set(ConfigKey { key }, None, self.op_id)?;
2002 }
2003 }
2004 Ok(())
2005 }
2006
2007 pub fn get_config(&self, key: String) -> Option<u64> {
2009 self.configs
2010 .get(&ConfigKey { key })
2011 .map(|entry| entry.value)
2012 }
2013
2014 pub fn get_setting(&self, name: String) -> Option<&str> {
2016 self.settings
2017 .get(&SettingKey { name })
2018 .map(|entry| &*entry.value)
2019 }
2020
2021 pub fn get_builtin_migration_shard(&self) -> Option<ShardId> {
2022 self.get_setting(BUILTIN_MIGRATION_SHARD_KEY.to_string())
2023 .map(|shard_id| shard_id.parse().expect("valid ShardId"))
2024 }
2025
2026 pub fn set_builtin_migration_shard(&mut self, shard_id: ShardId) -> Result<(), CatalogError> {
2027 self.set_setting(
2028 BUILTIN_MIGRATION_SHARD_KEY.to_string(),
2029 Some(shard_id.to_string()),
2030 )
2031 }
2032
2033 pub fn get_expression_cache_shard(&self) -> Option<ShardId> {
2034 self.get_setting(EXPRESSION_CACHE_SHARD_KEY.to_string())
2035 .map(|shard_id| shard_id.parse().expect("valid ShardId"))
2036 }
2037
2038 pub fn set_expression_cache_shard(&mut self, shard_id: ShardId) -> Result<(), CatalogError> {
2039 self.set_setting(
2040 EXPRESSION_CACHE_SHARD_KEY.to_string(),
2041 Some(shard_id.to_string()),
2042 )
2043 }
2044
2045 pub fn set_0dt_deployment_max_wait(&mut self, value: Duration) -> Result<(), CatalogError> {
2051 self.set_config(
2052 WITH_0DT_DEPLOYMENT_MAX_WAIT.into(),
2053 Some(
2054 value
2055 .as_millis()
2056 .try_into()
2057 .expect("max wait fits into u64"),
2058 ),
2059 )
2060 }
2061
2062 pub fn set_0dt_deployment_ddl_check_interval(
2069 &mut self,
2070 value: Duration,
2071 ) -> Result<(), CatalogError> {
2072 self.set_config(
2073 WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL.into(),
2074 Some(
2075 value
2076 .as_millis()
2077 .try_into()
2078 .expect("ddl check interval fits into u64"),
2079 ),
2080 )
2081 }
2082
2083 pub fn set_enable_0dt_deployment_panic_after_timeout(
2089 &mut self,
2090 value: bool,
2091 ) -> Result<(), CatalogError> {
2092 self.set_config(
2093 ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT.into(),
2094 Some(u64::from(value)),
2095 )
2096 }
2097
2098 pub fn reset_0dt_deployment_max_wait(&mut self) -> Result<(), CatalogError> {
2104 self.set_config(WITH_0DT_DEPLOYMENT_MAX_WAIT.into(), None)
2105 }
2106
2107 pub fn reset_0dt_deployment_ddl_check_interval(&mut self) -> Result<(), CatalogError> {
2114 self.set_config(WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL.into(), None)
2115 }
2116
2117 pub fn reset_enable_0dt_deployment_panic_after_timeout(&mut self) -> Result<(), CatalogError> {
2124 self.set_config(ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT.into(), None)
2125 }
2126
2127 pub fn set_system_config_synced_once(&mut self) -> Result<(), CatalogError> {
2129 self.set_config(SYSTEM_CONFIG_SYNCED_KEY.into(), Some(1))
2130 }
2131
2132 pub fn update_comment(
2133 &mut self,
2134 object_id: CommentObjectId,
2135 sub_component: Option<usize>,
2136 comment: Option<String>,
2137 ) -> Result<(), CatalogError> {
2138 let key = CommentKey {
2139 object_id,
2140 sub_component,
2141 };
2142 let value = comment.map(|c| CommentValue { comment: c });
2143 self.comments.set(key, value, self.op_id)?;
2144
2145 Ok(())
2146 }
2147
2148 pub fn drop_comments(
2149 &mut self,
2150 object_ids: &BTreeSet<CommentObjectId>,
2151 ) -> Result<(), CatalogError> {
2152 if object_ids.is_empty() {
2153 return Ok(());
2154 }
2155
2156 self.comments
2157 .delete(|k, _v| object_ids.contains(&k.object_id), self.op_id);
2158 Ok(())
2159 }
2160
2161 pub fn update_source_references(
2162 &mut self,
2163 source_id: CatalogItemId,
2164 references: Vec<SourceReference>,
2165 updated_at: u64,
2166 ) -> Result<(), CatalogError> {
2167 let key = SourceReferencesKey { source_id };
2168 let value = SourceReferencesValue {
2169 references,
2170 updated_at,
2171 };
2172 self.source_references.set(key, Some(value), self.op_id)?;
2173 Ok(())
2174 }
2175
2176 pub fn upsert_system_config(&mut self, name: &str, value: String) -> Result<(), CatalogError> {
2178 let key = ServerConfigurationKey {
2179 name: name.to_string(),
2180 };
2181 let value = ServerConfigurationValue { value };
2182 self.system_configurations
2183 .set(key, Some(value), self.op_id)?;
2184 Ok(())
2185 }
2186
2187 pub fn remove_system_config(&mut self, name: &str) {
2189 let key = ServerConfigurationKey {
2190 name: name.to_string(),
2191 };
2192 self.system_configurations
2193 .set(key, None, self.op_id)
2194 .expect("cannot have uniqueness violation");
2195 }
2196
2197 pub fn clear_system_configs(&mut self) {
2199 self.system_configurations.delete(|_k, _v| true, self.op_id);
2200 }
2201
2202 pub fn get_cluster_system_configurations(
2204 &self,
2205 ) -> impl Iterator<Item = ClusterSystemConfiguration> + use<'_> {
2206 self.cluster_system_configurations
2207 .items()
2208 .into_iter()
2209 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2210 }
2211
2212 pub fn upsert_cluster_system_config(
2215 &mut self,
2216 cluster_id: ClusterId,
2217 name: &str,
2218 value: String,
2219 ) -> Result<(), CatalogError> {
2220 let key = ClusterSystemConfigurationKey {
2221 cluster_id,
2222 name: name.to_string(),
2223 };
2224 let value = ClusterSystemConfigurationValue { value };
2225 self.cluster_system_configurations
2226 .set(key, Some(value), self.op_id)?;
2227 Ok(())
2228 }
2229
2230 pub fn remove_cluster_system_config(&mut self, cluster_id: ClusterId, name: &str) {
2233 let key = ClusterSystemConfigurationKey {
2234 cluster_id,
2235 name: name.to_string(),
2236 };
2237 self.cluster_system_configurations
2238 .set(key, None, self.op_id)
2239 .expect("cannot have uniqueness violation");
2240 }
2241
2242 pub fn get_replica_system_configurations(
2244 &self,
2245 ) -> impl Iterator<Item = ReplicaSystemConfiguration> + use<'_> {
2246 self.replica_system_configurations
2247 .items()
2248 .into_iter()
2249 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2250 }
2251
2252 pub fn upsert_replica_system_config(
2255 &mut self,
2256 replica_id: ReplicaId,
2257 name: &str,
2258 value: String,
2259 ) -> Result<(), CatalogError> {
2260 let key = ReplicaSystemConfigurationKey {
2261 replica_id,
2262 name: name.to_string(),
2263 };
2264 let value = ReplicaSystemConfigurationValue { value };
2265 self.replica_system_configurations
2266 .set(key, Some(value), self.op_id)?;
2267 Ok(())
2268 }
2269
2270 pub fn remove_replica_system_config(&mut self, replica_id: ReplicaId, name: &str) {
2273 let key = ReplicaSystemConfigurationKey {
2274 replica_id,
2275 name: name.to_string(),
2276 };
2277 self.replica_system_configurations
2278 .set(key, None, self.op_id)
2279 .expect("cannot have uniqueness violation");
2280 }
2281
2282 pub(crate) fn insert_config(&mut self, key: String, value: u64) -> Result<(), CatalogError> {
2283 match self.configs.insert(
2284 ConfigKey { key: key.clone() },
2285 ConfigValue { value },
2286 self.op_id,
2287 ) {
2288 Ok(_) => Ok(()),
2289 Err(_) => Err(SqlCatalogError::ConfigAlreadyExists(key).into()),
2290 }
2291 }
2292
2293 pub fn get_clusters(&self) -> impl Iterator<Item = Cluster> + use<'_> {
2294 self.clusters
2295 .items()
2296 .into_iter()
2297 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2298 }
2299
2300 pub fn get_cluster_replicas(&self) -> impl Iterator<Item = ClusterReplica> + use<'_> {
2301 self.cluster_replicas
2302 .items()
2303 .into_iter()
2304 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2305 }
2306
2307 pub fn get_databases(&self) -> impl Iterator<Item = Database> + use<'_> {
2308 self.databases
2309 .items()
2310 .into_iter()
2311 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2312 }
2313
2314 pub fn get_roles(&self) -> impl Iterator<Item = Role> + use<'_> {
2315 self.roles
2316 .items()
2317 .into_iter()
2318 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2319 }
2320
2321 pub fn get_network_policies(&self) -> impl Iterator<Item = NetworkPolicy> + use<'_> {
2322 self.network_policies
2323 .items()
2324 .into_iter()
2325 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2326 }
2327
2328 pub fn get_system_object_mappings(
2329 &self,
2330 ) -> impl Iterator<Item = SystemObjectMapping> + use<'_> {
2331 self.system_gid_mapping
2332 .items()
2333 .into_iter()
2334 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2335 }
2336
2337 pub fn get_schemas(&self) -> impl Iterator<Item = Schema> + use<'_> {
2338 self.schemas
2339 .items()
2340 .into_iter()
2341 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2342 }
2343
2344 pub fn get_system_configurations(&self) -> impl Iterator<Item = SystemConfiguration> + use<'_> {
2345 self.system_configurations
2346 .items()
2347 .into_iter()
2348 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2349 }
2350
2351 pub fn get_schema(&self, id: &SchemaId) -> Option<Schema> {
2352 let key = SchemaKey { id: *id };
2353 self.schemas
2354 .get(&key)
2355 .map(|v| DurableType::from_key_value(key, v.clone()))
2356 }
2357
2358 pub fn get_introspection_source_indexes(
2359 &self,
2360 cluster_id: ClusterId,
2361 ) -> BTreeMap<&str, (GlobalId, u32)> {
2362 self.introspection_sources
2363 .items()
2364 .into_iter()
2365 .filter(|(k, _v)| k.cluster_id == cluster_id)
2366 .map(|(k, v)| (k.name.as_str(), (v.global_id.into(), v.oid)))
2367 .collect()
2368 }
2369
2370 pub fn get_catalog_content_version(&self) -> Option<&str> {
2371 self.settings
2372 .get(&SettingKey {
2373 name: CATALOG_CONTENT_VERSION_KEY.to_string(),
2374 })
2375 .map(|value| &*value.value)
2376 }
2377
2378 pub fn get_authentication_mock_nonce(&self) -> Option<String> {
2379 self.settings
2380 .get(&SettingKey {
2381 name: MOCK_AUTHENTICATION_NONCE_KEY.to_string(),
2382 })
2383 .map(|value| value.value.clone())
2384 }
2385
2386 #[must_use]
2392 pub fn get_and_commit_op_updates(&mut self) -> Vec<StateUpdate> {
2393 let updates = self.get_op_updates();
2394 self.commit_op();
2395 updates
2396 }
2397
2398 fn get_op_updates(&self) -> Vec<StateUpdate> {
2399 fn get_collection_op_updates<'a, T>(
2400 table_txn: &'a TableTransaction<T::Key, T::Value>,
2401 kind_fn: impl Fn(T) -> StateUpdateKind + 'a,
2402 op: Timestamp,
2403 ) -> impl Iterator<Item = (StateUpdateKind, StateDiff)> + 'a
2404 where
2405 T::Key: Ord + Eq + Clone + Debug,
2406 T::Value: Ord + Clone + Debug,
2407 T: DurableType,
2408 {
2409 table_txn
2410 .pending
2411 .iter()
2412 .flat_map(|(k, vs)| vs.into_iter().map(move |v| (k, v)))
2413 .filter_map(move |(k, v)| {
2414 if v.ts == op {
2415 let key = k.clone();
2416 let value = v.value.clone();
2417 let diff = v.diff.clone().try_into().expect("invalid diff");
2418 let update = DurableType::from_key_value(key, value);
2419 let kind = kind_fn(update);
2420 Some((kind, diff))
2421 } else {
2422 None
2423 }
2424 })
2425 }
2426
2427 fn get_large_collection_op_updates<'a, T>(
2428 collection: &'a Vec<(T::Key, Diff, Timestamp)>,
2429 kind_fn: impl Fn(T) -> StateUpdateKind + 'a,
2430 op: Timestamp,
2431 ) -> impl Iterator<Item = (StateUpdateKind, StateDiff)> + 'a
2432 where
2433 T::Key: Ord + Eq + Clone + Debug,
2434 T: DurableType<Value = ()>,
2435 {
2436 collection.iter().filter_map(move |(k, diff, ts)| {
2437 if *ts == op {
2438 let key = k.clone();
2439 let diff = diff.clone().try_into().expect("invalid diff");
2440 let update = DurableType::from_key_value(key, ());
2441 let kind = kind_fn(update);
2442 Some((kind, diff))
2443 } else {
2444 None
2445 }
2446 })
2447 }
2448
2449 let Transaction {
2450 durable_catalog: _,
2451 databases,
2452 schemas,
2453 items,
2454 comments,
2455 roles,
2456 role_auth,
2457 clusters,
2458 network_policies,
2459 cluster_replicas,
2460 introspection_sources,
2461 system_gid_mapping,
2462 system_configurations,
2463 cluster_system_configurations,
2464 replica_system_configurations,
2465 default_privileges,
2466 source_references,
2467 system_privileges,
2468 audit_log_updates,
2469 storage_collection_metadata,
2470 unfinalized_shards,
2471 id_allocator: _,
2473 configs: _,
2474 settings: _,
2475 txn_wal_shard: _,
2476 upper,
2477 op_id: _,
2478 commit_capability: _,
2479 } = &self;
2480
2481 let updates = std::iter::empty()
2482 .chain(get_collection_op_updates(
2483 roles,
2484 StateUpdateKind::Role,
2485 self.op_id,
2486 ))
2487 .chain(get_collection_op_updates(
2488 role_auth,
2489 StateUpdateKind::RoleAuth,
2490 self.op_id,
2491 ))
2492 .chain(get_collection_op_updates(
2493 databases,
2494 StateUpdateKind::Database,
2495 self.op_id,
2496 ))
2497 .chain(get_collection_op_updates(
2498 schemas,
2499 StateUpdateKind::Schema,
2500 self.op_id,
2501 ))
2502 .chain(get_collection_op_updates(
2503 default_privileges,
2504 StateUpdateKind::DefaultPrivilege,
2505 self.op_id,
2506 ))
2507 .chain(get_collection_op_updates(
2508 system_privileges,
2509 StateUpdateKind::SystemPrivilege,
2510 self.op_id,
2511 ))
2512 .chain(get_collection_op_updates(
2513 system_configurations,
2514 StateUpdateKind::SystemConfiguration,
2515 self.op_id,
2516 ))
2517 .chain(get_collection_op_updates(
2518 cluster_system_configurations,
2519 StateUpdateKind::ClusterSystemConfiguration,
2520 self.op_id,
2521 ))
2522 .chain(get_collection_op_updates(
2523 replica_system_configurations,
2524 StateUpdateKind::ReplicaSystemConfiguration,
2525 self.op_id,
2526 ))
2527 .chain(get_collection_op_updates(
2528 clusters,
2529 StateUpdateKind::Cluster,
2530 self.op_id,
2531 ))
2532 .chain(get_collection_op_updates(
2533 network_policies,
2534 StateUpdateKind::NetworkPolicy,
2535 self.op_id,
2536 ))
2537 .chain(get_collection_op_updates(
2538 introspection_sources,
2539 StateUpdateKind::IntrospectionSourceIndex,
2540 self.op_id,
2541 ))
2542 .chain(get_collection_op_updates(
2543 cluster_replicas,
2544 StateUpdateKind::ClusterReplica,
2545 self.op_id,
2546 ))
2547 .chain(get_collection_op_updates(
2548 system_gid_mapping,
2549 StateUpdateKind::SystemObjectMapping,
2550 self.op_id,
2551 ))
2552 .chain(get_collection_op_updates(
2553 items,
2554 StateUpdateKind::Item,
2555 self.op_id,
2556 ))
2557 .chain(get_collection_op_updates(
2558 comments,
2559 StateUpdateKind::Comment,
2560 self.op_id,
2561 ))
2562 .chain(get_collection_op_updates(
2563 source_references,
2564 StateUpdateKind::SourceReferences,
2565 self.op_id,
2566 ))
2567 .chain(get_collection_op_updates(
2568 storage_collection_metadata,
2569 StateUpdateKind::StorageCollectionMetadata,
2570 self.op_id,
2571 ))
2572 .chain(get_collection_op_updates(
2573 unfinalized_shards,
2574 StateUpdateKind::UnfinalizedShard,
2575 self.op_id,
2576 ))
2577 .chain(get_large_collection_op_updates(
2578 audit_log_updates,
2579 StateUpdateKind::AuditLog,
2580 self.op_id,
2581 ))
2582 .map(|(kind, diff)| StateUpdate {
2583 kind,
2584 ts: upper.clone(),
2585 diff,
2586 })
2587 .collect();
2588
2589 updates
2590 }
2591
2592 pub fn is_savepoint(&self) -> bool {
2593 self.durable_catalog.is_savepoint()
2594 }
2595
2596 fn commit_op(&mut self) {
2597 self.op_id += 1;
2598 }
2599
2600 pub fn op_id(&self) -> Timestamp {
2601 self.op_id
2602 }
2603
2604 pub fn upper(&self) -> mz_repr::Timestamp {
2605 self.upper
2606 }
2607
2608 fn ensure_committable(&self) -> Result<(), CatalogError> {
2609 match self.commit_capability {
2610 Some(_) => Ok(()),
2611 None => Err(DurableCatalogError::DryRunTransaction.into()),
2612 }
2613 }
2614
2615 pub(super) async fn ensure_not_out_of_sync(&mut self) -> Result<(), CatalogError> {
2617 self.durable_catalog
2618 .ensure_not_out_of_sync(self.upper)
2619 .await
2620 }
2621
2622 pub(crate) fn into_parts(
2623 self,
2624 ) -> Result<(TransactionBatch, &'a mut dyn DurableCatalogState), CatalogError> {
2625 let commit_capability = self
2626 .commit_capability
2627 .ok_or(DurableCatalogError::DryRunTransaction)?;
2628 let audit_log_updates = self
2629 .audit_log_updates
2630 .into_iter()
2631 .map(|(k, diff, _op)| (k.into_proto(), (), diff))
2632 .collect();
2633
2634 let txn_batch = TransactionBatch {
2635 databases: self.databases.pending(),
2636 schemas: self.schemas.pending(),
2637 items: self.items.pending(),
2638 comments: self.comments.pending(),
2639 roles: self.roles.pending(),
2640 role_auth: self.role_auth.pending(),
2641 clusters: self.clusters.pending(),
2642 cluster_replicas: self.cluster_replicas.pending(),
2643 network_policies: self.network_policies.pending(),
2644 introspection_sources: self.introspection_sources.pending(),
2645 id_allocator: self.id_allocator.pending(),
2646 configs: self.configs.pending(),
2647 source_references: self.source_references.pending(),
2648 settings: self.settings.pending(),
2649 system_gid_mapping: self.system_gid_mapping.pending(),
2650 system_configurations: self.system_configurations.pending(),
2651 cluster_system_configurations: self.cluster_system_configurations.pending(),
2652 replica_system_configurations: self.replica_system_configurations.pending(),
2653 default_privileges: self.default_privileges.pending(),
2654 system_privileges: self.system_privileges.pending(),
2655 storage_collection_metadata: self.storage_collection_metadata.pending(),
2656 unfinalized_shards: self.unfinalized_shards.pending(),
2657 txn_wal_shard: self.txn_wal_shard.pending(),
2658 audit_log_updates,
2659 upper: self.upper,
2660 _commit_capability: commit_capability,
2661 };
2662 Ok((txn_batch, self.durable_catalog))
2663 }
2664
2665 #[mz_ore::instrument(level = "debug")]
2681 pub(crate) async fn commit_internal(
2682 self,
2683 commit_ts: mz_repr::Timestamp,
2684 ) -> Result<(&'a mut dyn DurableCatalogState, mz_repr::Timestamp), CatalogError> {
2685 self.ensure_committable()?;
2686 let (mut txn_batch, durable_catalog) = self.into_parts()?;
2687 let TransactionBatch {
2688 databases,
2689 schemas,
2690 items,
2691 comments,
2692 roles,
2693 role_auth,
2694 clusters,
2695 cluster_replicas,
2696 network_policies,
2697 introspection_sources,
2698 id_allocator,
2699 configs,
2700 source_references,
2701 settings,
2702 system_gid_mapping,
2703 system_configurations,
2704 cluster_system_configurations,
2705 replica_system_configurations,
2706 default_privileges,
2707 system_privileges,
2708 storage_collection_metadata,
2709 unfinalized_shards,
2710 txn_wal_shard,
2711 audit_log_updates,
2712 upper: _,
2713 _commit_capability: _,
2714 } = &mut txn_batch;
2715 differential_dataflow::consolidation::consolidate_updates(databases);
2718 differential_dataflow::consolidation::consolidate_updates(schemas);
2719 differential_dataflow::consolidation::consolidate_updates(items);
2720 differential_dataflow::consolidation::consolidate_updates(comments);
2721 differential_dataflow::consolidation::consolidate_updates(roles);
2722 differential_dataflow::consolidation::consolidate_updates(role_auth);
2723 differential_dataflow::consolidation::consolidate_updates(clusters);
2724 differential_dataflow::consolidation::consolidate_updates(cluster_replicas);
2725 differential_dataflow::consolidation::consolidate_updates(network_policies);
2726 differential_dataflow::consolidation::consolidate_updates(introspection_sources);
2727 differential_dataflow::consolidation::consolidate_updates(id_allocator);
2728 differential_dataflow::consolidation::consolidate_updates(configs);
2729 differential_dataflow::consolidation::consolidate_updates(settings);
2730 differential_dataflow::consolidation::consolidate_updates(source_references);
2731 differential_dataflow::consolidation::consolidate_updates(system_gid_mapping);
2732 differential_dataflow::consolidation::consolidate_updates(system_configurations);
2733 differential_dataflow::consolidation::consolidate_updates(cluster_system_configurations);
2734 differential_dataflow::consolidation::consolidate_updates(replica_system_configurations);
2735 differential_dataflow::consolidation::consolidate_updates(default_privileges);
2736 differential_dataflow::consolidation::consolidate_updates(system_privileges);
2737 differential_dataflow::consolidation::consolidate_updates(storage_collection_metadata);
2738 differential_dataflow::consolidation::consolidate_updates(unfinalized_shards);
2739 differential_dataflow::consolidation::consolidate_updates(txn_wal_shard);
2740 differential_dataflow::consolidation::consolidate_updates(audit_log_updates);
2741
2742 let upper = durable_catalog
2743 .commit_transaction(txn_batch, commit_ts)
2744 .await?;
2745 Ok((durable_catalog, upper))
2746 }
2747
2748 #[mz_ore::instrument(level = "debug")]
2769 pub async fn commit(self, commit_ts: mz_repr::Timestamp) -> Result<(), CatalogError> {
2770 self.ensure_committable()?;
2771 let op_updates = self.get_op_updates();
2772 assert!(
2773 op_updates.is_empty(),
2774 "unconsumed transaction updates: {op_updates:?}"
2775 );
2776
2777 let (durable_storage, upper) = self.commit_internal(commit_ts).await?;
2778 let updates = durable_storage.sync_updates(upper).await?;
2780 soft_assert_no_log!(
2789 durable_storage.is_read_only()
2790 || updates
2791 .iter()
2792 .all(|update| update.ts >= commit_ts && update.ts < upper),
2793 "unconsumed updates existed before transaction commit: commit_ts={commit_ts:?}, upper={upper:?}, updates:{updates:?}"
2794 );
2795 Ok(())
2796 }
2797}
2798
2799use crate::durable::async_trait;
2800
2801use super::objects::{RoleAuthKey, RoleAuthValue};
2802
2803#[async_trait]
2804impl StorageTxn for Transaction<'_> {
2805 fn get_collection_metadata(&self) -> BTreeMap<GlobalId, ShardId> {
2806 self.storage_collection_metadata
2807 .items()
2808 .into_iter()
2809 .map(
2810 |(
2811 StorageCollectionMetadataKey { id },
2812 StorageCollectionMetadataValue { shard },
2813 )| { (*id, shard.clone()) },
2814 )
2815 .collect()
2816 }
2817
2818 fn insert_collection_metadata(
2819 &mut self,
2820 metadata: BTreeMap<GlobalId, ShardId>,
2821 ) -> Result<(), StorageError> {
2822 for (id, shard) in metadata {
2823 self.storage_collection_metadata
2824 .insert(
2825 StorageCollectionMetadataKey { id },
2826 StorageCollectionMetadataValue {
2827 shard: shard.clone(),
2828 },
2829 self.op_id,
2830 )
2831 .map_err(|err| match err {
2832 DurableCatalogError::DuplicateKey => {
2833 StorageError::CollectionMetadataAlreadyExists(id)
2834 }
2835 DurableCatalogError::UniquenessViolation => {
2836 StorageError::PersistShardAlreadyInUse(shard)
2837 }
2838 err => StorageError::Generic(anyhow::anyhow!(err)),
2839 })?;
2840 }
2841 Ok(())
2842 }
2843
2844 fn delete_collection_metadata(&mut self, ids: BTreeSet<GlobalId>) -> Vec<(GlobalId, ShardId)> {
2845 let ks: Vec<_> = ids
2846 .into_iter()
2847 .map(|id| StorageCollectionMetadataKey { id })
2848 .collect();
2849 self.storage_collection_metadata
2850 .delete_by_keys(ks, self.op_id)
2851 .into_iter()
2852 .map(
2853 |(
2854 StorageCollectionMetadataKey { id },
2855 StorageCollectionMetadataValue { shard },
2856 )| (id, shard),
2857 )
2858 .collect()
2859 }
2860
2861 fn get_unfinalized_shards(&self) -> BTreeSet<ShardId> {
2862 self.unfinalized_shards
2863 .items()
2864 .into_iter()
2865 .map(|(UnfinalizedShardKey { shard }, ())| *shard)
2866 .collect()
2867 }
2868
2869 fn insert_unfinalized_shards(&mut self, s: BTreeSet<ShardId>) -> Result<(), StorageError> {
2870 for shard in s {
2871 match self
2872 .unfinalized_shards
2873 .insert(UnfinalizedShardKey { shard }, (), self.op_id)
2874 {
2875 Ok(()) | Err(DurableCatalogError::DuplicateKey) => {}
2877 Err(e) => Err(StorageError::Generic(anyhow::anyhow!(e)))?,
2878 };
2879 }
2880 Ok(())
2881 }
2882
2883 fn remove_unfinalized_shards(&mut self, shards: BTreeSet<ShardId>) {
2884 let ks: Vec<_> = shards
2885 .into_iter()
2886 .map(|shard| UnfinalizedShardKey { shard })
2887 .collect();
2888 let _ = self.unfinalized_shards.delete_by_keys(ks, self.op_id);
2889 }
2890
2891 fn get_txn_wal_shard(&self) -> Option<ShardId> {
2892 self.txn_wal_shard
2893 .values()
2894 .iter()
2895 .next()
2896 .map(|TxnWalShardValue { shard }| *shard)
2897 }
2898
2899 fn write_txn_wal_shard(&mut self, shard: ShardId) -> Result<(), StorageError> {
2900 self.txn_wal_shard
2901 .insert((), TxnWalShardValue { shard }, self.op_id)
2902 .map_err(|err| match err {
2903 DurableCatalogError::DuplicateKey => StorageError::TxnWalShardAlreadyExists,
2904 err => StorageError::Generic(anyhow::anyhow!(err)),
2905 })
2906 }
2907}
2908
2909#[derive(Debug, PartialEq)]
2911pub struct TransactionBatch {
2912 pub(crate) databases: Vec<(proto::DatabaseKey, proto::DatabaseValue, Diff)>,
2913 pub(crate) schemas: Vec<(proto::SchemaKey, proto::SchemaValue, Diff)>,
2914 pub(crate) items: Vec<(proto::ItemKey, proto::ItemValue, Diff)>,
2915 pub(crate) comments: Vec<(proto::CommentKey, proto::CommentValue, Diff)>,
2916 pub(crate) roles: Vec<(proto::RoleKey, proto::RoleValue, Diff)>,
2917 pub(crate) role_auth: Vec<(proto::RoleAuthKey, proto::RoleAuthValue, Diff)>,
2918 pub(crate) clusters: Vec<(proto::ClusterKey, proto::ClusterValue, Diff)>,
2919 pub(crate) cluster_replicas: Vec<(proto::ClusterReplicaKey, proto::ClusterReplicaValue, Diff)>,
2920 pub(crate) network_policies: Vec<(proto::NetworkPolicyKey, proto::NetworkPolicyValue, Diff)>,
2921 pub(crate) introspection_sources: Vec<(
2922 proto::ClusterIntrospectionSourceIndexKey,
2923 proto::ClusterIntrospectionSourceIndexValue,
2924 Diff,
2925 )>,
2926 pub(crate) id_allocator: Vec<(proto::IdAllocKey, proto::IdAllocValue, Diff)>,
2927 pub(crate) configs: Vec<(proto::ConfigKey, proto::ConfigValue, Diff)>,
2928 pub(crate) settings: Vec<(proto::SettingKey, proto::SettingValue, Diff)>,
2929 pub(crate) system_gid_mapping: Vec<(proto::GidMappingKey, proto::GidMappingValue, Diff)>,
2930 pub(crate) system_configurations: Vec<(
2931 proto::ServerConfigurationKey,
2932 proto::ServerConfigurationValue,
2933 Diff,
2934 )>,
2935 pub(crate) cluster_system_configurations: Vec<(
2936 proto::ClusterSystemConfigurationKey,
2937 proto::ClusterSystemConfigurationValue,
2938 Diff,
2939 )>,
2940 pub(crate) replica_system_configurations: Vec<(
2941 proto::ReplicaSystemConfigurationKey,
2942 proto::ReplicaSystemConfigurationValue,
2943 Diff,
2944 )>,
2945 pub(crate) default_privileges: Vec<(
2946 proto::DefaultPrivilegesKey,
2947 proto::DefaultPrivilegesValue,
2948 Diff,
2949 )>,
2950 pub(crate) source_references: Vec<(
2951 proto::SourceReferencesKey,
2952 proto::SourceReferencesValue,
2953 Diff,
2954 )>,
2955 pub(crate) system_privileges: Vec<(
2956 proto::SystemPrivilegesKey,
2957 proto::SystemPrivilegesValue,
2958 Diff,
2959 )>,
2960 pub(crate) storage_collection_metadata: Vec<(
2961 proto::StorageCollectionMetadataKey,
2962 proto::StorageCollectionMetadataValue,
2963 Diff,
2964 )>,
2965 pub(crate) unfinalized_shards: Vec<(proto::UnfinalizedShardKey, (), Diff)>,
2966 pub(crate) txn_wal_shard: Vec<((), proto::TxnWalShardValue, Diff)>,
2967 pub(crate) audit_log_updates: Vec<(proto::AuditLogKey, (), Diff)>,
2968 pub(crate) upper: mz_repr::Timestamp,
2970 _commit_capability: CommitCapability,
2973}
2974
2975impl TransactionBatch {
2976 pub fn is_empty(&self) -> bool {
2977 let TransactionBatch {
2978 databases,
2979 schemas,
2980 items,
2981 comments,
2982 roles,
2983 role_auth,
2984 clusters,
2985 cluster_replicas,
2986 network_policies,
2987 introspection_sources,
2988 id_allocator,
2989 configs,
2990 settings,
2991 source_references,
2992 system_gid_mapping,
2993 system_configurations,
2994 cluster_system_configurations,
2995 replica_system_configurations,
2996 default_privileges,
2997 system_privileges,
2998 storage_collection_metadata,
2999 unfinalized_shards,
3000 txn_wal_shard,
3001 audit_log_updates,
3002 upper: _,
3003 _commit_capability: _,
3004 } = self;
3005 databases.is_empty()
3006 && schemas.is_empty()
3007 && items.is_empty()
3008 && comments.is_empty()
3009 && roles.is_empty()
3010 && role_auth.is_empty()
3011 && clusters.is_empty()
3012 && cluster_replicas.is_empty()
3013 && network_policies.is_empty()
3014 && introspection_sources.is_empty()
3015 && id_allocator.is_empty()
3016 && configs.is_empty()
3017 && settings.is_empty()
3018 && source_references.is_empty()
3019 && system_gid_mapping.is_empty()
3020 && system_configurations.is_empty()
3021 && cluster_system_configurations.is_empty()
3022 && replica_system_configurations.is_empty()
3023 && default_privileges.is_empty()
3024 && system_privileges.is_empty()
3025 && storage_collection_metadata.is_empty()
3026 && unfinalized_shards.is_empty()
3027 && txn_wal_shard.is_empty()
3028 && audit_log_updates.is_empty()
3029 }
3030}
3031
3032#[derive(Debug, Clone, PartialEq, Eq)]
3033struct TransactionUpdate<V> {
3034 value: V,
3035 ts: Timestamp,
3036 diff: Diff,
3037}
3038
3039trait UniqueName {
3041 const HAS_UNIQUE_NAME: bool;
3044 fn unique_name(&self) -> &str;
3046}
3047
3048mod unique_name {
3049 use crate::durable::objects::*;
3050
3051 macro_rules! impl_unique_name {
3052 ($($t:ty),* $(,)?) => {
3053 $(
3054 impl crate::durable::transaction::UniqueName for $t {
3055 const HAS_UNIQUE_NAME: bool = true;
3056 fn unique_name(&self) -> &str {
3057 &self.name
3058 }
3059 }
3060 )*
3061 };
3062 }
3063
3064 macro_rules! impl_no_unique_name {
3065 ($($t:ty),* $(,)?) => {
3066 $(
3067 impl crate::durable::transaction::UniqueName for $t {
3068 const HAS_UNIQUE_NAME: bool = false;
3069 fn unique_name(&self) -> &str {
3070 ""
3071 }
3072 }
3073 )*
3074 };
3075 }
3076
3077 impl_unique_name! {
3078 ClusterReplicaValue,
3079 ClusterValue,
3080 DatabaseValue,
3081 ItemValue,
3082 NetworkPolicyValue,
3083 RoleValue,
3084 SchemaValue,
3085 }
3086
3087 impl_no_unique_name!(
3088 (),
3089 ClusterIntrospectionSourceIndexValue,
3090 ClusterSystemConfigurationValue,
3091 CommentValue,
3092 ConfigValue,
3093 DefaultPrivilegesValue,
3094 GidMappingValue,
3095 IdAllocValue,
3096 ReplicaSystemConfigurationValue,
3097 ServerConfigurationValue,
3098 SettingValue,
3099 SourceReferencesValue,
3100 StorageCollectionMetadataValue,
3101 SystemPrivilegesValue,
3102 TxnWalShardValue,
3103 RoleAuthValue,
3104 );
3105
3106 #[cfg(test)]
3107 mod test {
3108 impl_no_unique_name!(String,);
3109 }
3110}
3111
3112#[derive(Debug)]
3120struct UniquenessCheck<V> {
3121 violation: fn(a: &V, b: &V) -> bool,
3122 is_unique_key_unchanged_after_update: fn(prev: &V, next: &V) -> bool,
3123}
3124
3125#[derive(Debug)]
3135struct TableTransaction<K, V> {
3136 initial: BTreeMap<K, V>,
3137 pending: BTreeMap<K, Vec<TransactionUpdate<V>>>,
3140 uniqueness_check: Option<UniquenessCheck<V>>,
3142}
3143
3144impl<K, V> TableTransaction<K, V>
3145where
3146 K: Ord + Eq + Clone + Debug,
3147 V: Ord + Clone + Debug + UniqueName,
3148{
3149 fn new<KP, VP>(initial: BTreeMap<KP, VP>) -> Result<Self, TryFromProtoError>
3156 where
3157 K: RustType<KP>,
3158 V: RustType<VP>,
3159 {
3160 let initial = initial
3161 .into_iter()
3162 .map(RustType::from_proto)
3163 .collect::<Result<_, _>>()?;
3164
3165 Ok(Self {
3166 initial,
3167 pending: BTreeMap::new(),
3168 uniqueness_check: None,
3169 })
3170 }
3171
3172 fn new_with_uniqueness_fn<KP, VP>(
3175 initial: BTreeMap<KP, VP>,
3176 uniqueness_violation: fn(a: &V, b: &V) -> bool,
3177 is_unique_key_unchanged_after_update: fn(prev: &V, next: &V) -> bool,
3178 ) -> Result<Self, TryFromProtoError>
3179 where
3180 K: RustType<KP>,
3181 V: RustType<VP>,
3182 {
3183 let initial = initial
3184 .into_iter()
3185 .map(RustType::from_proto)
3186 .collect::<Result<_, _>>()?;
3187
3188 Ok(Self {
3189 initial,
3190 pending: BTreeMap::new(),
3191 uniqueness_check: Some(UniquenessCheck {
3192 violation: uniqueness_violation,
3193 is_unique_key_unchanged_after_update,
3194 }),
3195 })
3196 }
3197
3198 fn pending<KP, VP>(self) -> Vec<(KP, VP, Diff)>
3201 where
3202 K: RustType<KP>,
3203 V: RustType<VP>,
3204 {
3205 soft_assert_no_log!(self.verify().is_ok());
3206 self.pending
3209 .into_iter()
3210 .flat_map(|(k, v)| {
3211 let mut v: Vec<_> = v
3212 .into_iter()
3213 .map(|TransactionUpdate { value, ts: _, diff }| (value, diff))
3214 .collect();
3215 differential_dataflow::consolidation::consolidate(&mut v);
3216 v.into_iter().map(move |(v, diff)| (k.clone(), v, diff))
3217 })
3218 .map(|(key, val, diff)| (key.into_proto(), val.into_proto(), diff))
3219 .collect()
3220 }
3221
3222 fn verify(&self) -> Result<(), DurableCatalogError> {
3227 if let Some(check) = &self.uniqueness_check {
3228 let items = self.values();
3230 if V::HAS_UNIQUE_NAME {
3231 let by_name: BTreeMap<_, _> = items
3232 .iter()
3233 .enumerate()
3234 .map(|(v, vi)| (vi.unique_name(), (v, vi)))
3235 .collect();
3236 for (i, vi) in items.iter().enumerate() {
3237 if let Some((j, vj)) = by_name.get(vi.unique_name()) {
3238 if i != *j && (check.violation)(vi, *vj) {
3239 return Err(DurableCatalogError::UniquenessViolation);
3240 }
3241 }
3242 }
3243 } else {
3244 for (i, vi) in items.iter().enumerate() {
3245 for (j, vj) in items.iter().enumerate() {
3246 if i != j && (check.violation)(vi, vj) {
3247 return Err(DurableCatalogError::UniquenessViolation);
3248 }
3249 }
3250 }
3251 }
3252 }
3253 soft_assert_no_log!(
3254 self.pending
3255 .values()
3256 .all(|pending| { pending.is_sorted_by(|a, b| a.ts <= b.ts) }),
3257 "pending should be sorted by timestamp: {:?}",
3258 self.pending
3259 );
3260 Ok(())
3261 }
3262
3263 fn verify_keys<'a>(
3268 &self,
3269 keys: impl IntoIterator<Item = &'a K>,
3270 ) -> Result<(), DurableCatalogError>
3271 where
3272 K: 'a,
3273 {
3274 if let Some(check) = &self.uniqueness_check {
3275 let entries: Vec<_> = keys
3276 .into_iter()
3277 .filter_map(|key| self.get(key).map(|value| (key, value)))
3278 .collect();
3279 for (ki, vi) in self.items() {
3281 for (kj, vj) in &entries {
3282 if ki != *kj && (check.violation)(vi, vj) {
3283 return Err(DurableCatalogError::UniquenessViolation);
3284 }
3285 }
3286 }
3287 }
3288 soft_assert_no_log!(self.verify().is_ok());
3289 Ok(())
3290 }
3291
3292 fn for_values<'a, F: FnMut(&'a K, &'a V)>(&'a self, mut f: F) {
3295 let mut seen = BTreeSet::new();
3296 for k in self.pending.keys() {
3297 seen.insert(k);
3298 let v = self.get(k);
3299 if let Some(v) = v {
3302 f(k, v);
3303 }
3304 }
3305 for (k, v) in self.initial.iter() {
3306 if !seen.contains(k) {
3308 f(k, v);
3309 }
3310 }
3311 }
3312
3313 fn get(&self, k: &K) -> Option<&V> {
3315 let pending = self.pending.get(k).map(Vec::as_slice).unwrap_or_default();
3316 let mut updates = Vec::with_capacity(pending.len() + 1);
3317 if let Some(initial) = self.initial.get(k) {
3318 updates.push((initial, Diff::ONE));
3319 }
3320 updates.extend(
3321 pending
3322 .into_iter()
3323 .map(|TransactionUpdate { value, ts: _, diff }| (value, *diff)),
3324 );
3325
3326 differential_dataflow::consolidation::consolidate(&mut updates);
3327 assert!(updates.len() <= 1);
3328 updates.into_iter().next().map(|(v, _)| v)
3329 }
3330
3331 #[cfg(test)]
3336 fn items_cloned(&self) -> BTreeMap<K, V> {
3337 let mut items = BTreeMap::new();
3338 self.for_values(|k, v| {
3339 items.insert(k.clone(), v.clone());
3340 });
3341 items
3342 }
3343
3344 fn current_items_proto<KP, VP>(&self) -> BTreeMap<KP, VP>
3348 where
3349 K: RustType<KP>,
3350 V: RustType<VP>,
3351 KP: Ord,
3352 {
3353 let mut items = BTreeMap::new();
3354 self.for_values(|k, v| {
3355 items.insert(k.into_proto(), v.into_proto());
3356 });
3357 items
3358 }
3359
3360 fn items(&self) -> BTreeMap<&K, &V> {
3363 let mut items = BTreeMap::new();
3364 self.for_values(|k, v| {
3365 items.insert(k, v);
3366 });
3367 items
3368 }
3369
3370 fn values(&self) -> BTreeSet<&V> {
3372 let mut items = BTreeSet::new();
3373 self.for_values(|_, v| {
3374 items.insert(v);
3375 });
3376 items
3377 }
3378
3379 fn len(&self) -> usize {
3381 let mut count = 0;
3382 self.for_values(|_, _| {
3383 count += 1;
3384 });
3385 count
3386 }
3387
3388 fn for_values_mut<F: FnMut(&mut BTreeMap<K, Vec<TransactionUpdate<V>>>, &K, &V)>(
3392 &mut self,
3393 mut f: F,
3394 ) {
3395 let mut pending = BTreeMap::new();
3396 self.for_values(|k, v| f(&mut pending, k, v));
3397 for (k, updates) in pending {
3398 self.pending.entry(k).or_default().extend(updates);
3399 }
3400 }
3401
3402 fn insert(&mut self, k: K, v: V, ts: Timestamp) -> Result<(), DurableCatalogError> {
3406 let mut violation = None;
3407 let uniqueness_violation = self.uniqueness_check.as_ref().map(|check| check.violation);
3408 self.for_values(|for_k, for_v| {
3409 if &k == for_k {
3410 violation = Some(DurableCatalogError::DuplicateKey);
3411 }
3412 if let Some(uniqueness_violation) = uniqueness_violation {
3413 if uniqueness_violation(for_v, &v) {
3414 violation = Some(DurableCatalogError::UniquenessViolation);
3415 }
3416 }
3417 });
3418 if let Some(violation) = violation {
3419 return Err(violation);
3420 }
3421 self.pending.entry(k).or_default().push(TransactionUpdate {
3422 value: v,
3423 ts,
3424 diff: Diff::ONE,
3425 });
3426 soft_assert_no_log!(self.verify().is_ok());
3427 Ok(())
3428 }
3429
3430 fn update<F: Fn(&K, &V) -> Option<V>>(
3439 &mut self,
3440 f: F,
3441 ts: Timestamp,
3442 ) -> Result<Diff, DurableCatalogError> {
3443 let mut changed = Diff::ZERO;
3444 let mut keys = BTreeSet::new();
3445 let pending = self.pending.clone();
3447 self.for_values_mut(|p, k, v| {
3448 if let Some(next) = f(k, v) {
3449 changed += Diff::ONE;
3450 keys.insert(k.clone());
3451 let updates = p.entry(k.clone()).or_default();
3452 updates.push(TransactionUpdate {
3453 value: v.clone(),
3454 ts,
3455 diff: Diff::MINUS_ONE,
3456 });
3457 updates.push(TransactionUpdate {
3458 value: next,
3459 ts,
3460 diff: Diff::ONE,
3461 });
3462 }
3463 });
3464 if let Err(err) = self.verify_keys(&keys) {
3466 self.pending = pending;
3467 Err(err)
3468 } else {
3469 Ok(changed)
3470 }
3471 }
3472
3473 fn update_by_key(&mut self, k: K, v: V, ts: Timestamp) -> Result<bool, DurableCatalogError> {
3478 if let Some(cur_v) = self.get(&k) {
3479 if v != *cur_v {
3480 self.set(k, Some(v), ts)?;
3481 }
3482 Ok(true)
3483 } else {
3484 Ok(false)
3485 }
3486 }
3487
3488 fn update_by_keys(
3493 &mut self,
3494 kvs: impl IntoIterator<Item = (K, V)>,
3495 ts: Timestamp,
3496 ) -> Result<Diff, DurableCatalogError> {
3497 let kvs: Vec<_> = kvs
3498 .into_iter()
3499 .filter_map(|(k, v)| match self.get(&k) {
3500 Some(cur_v) => Some((*cur_v == v, k, v)),
3502 None => None,
3503 })
3504 .collect();
3505 let changed = kvs.len();
3506 let changed =
3507 Diff::try_from(changed).map_err(|e| DurableCatalogError::Internal(e.to_string()))?;
3508 let kvs = kvs
3509 .into_iter()
3510 .filter(|(no_op, _, _)| !no_op)
3512 .map(|(_, k, v)| (k, Some(v)))
3513 .collect();
3514 self.set_many(kvs, ts)?;
3515 Ok(changed)
3516 }
3517
3518 fn update_needs_uniqueness_check(&self, prev: Option<&V>, next: Option<&V>) -> bool {
3520 match (&self.uniqueness_check, prev, next) {
3521 (None, _, _) | (_, _, None) => false,
3523 (Some(check), Some(prev), Some(next)) => {
3525 !(check.is_unique_key_unchanged_after_update)(prev, next)
3526 }
3527 (Some(_), None, Some(_)) => true,
3529 }
3530 }
3531
3532 fn set(&mut self, k: K, v: Option<V>, ts: Timestamp) -> Result<Option<V>, DurableCatalogError> {
3539 let prev = self.get(&k).cloned();
3540 let needs_uniqueness_check = self.update_needs_uniqueness_check(prev.as_ref(), v.as_ref());
3541 let entry = self.pending.entry(k.clone()).or_default();
3542 let restore_len = entry.len();
3543
3544 match (v, prev.clone()) {
3545 (Some(v), Some(prev)) => {
3546 entry.push(TransactionUpdate {
3547 value: prev,
3548 ts,
3549 diff: Diff::MINUS_ONE,
3550 });
3551 entry.push(TransactionUpdate {
3552 value: v,
3553 ts,
3554 diff: Diff::ONE,
3555 });
3556 }
3557 (Some(v), None) => {
3558 entry.push(TransactionUpdate {
3559 value: v,
3560 ts,
3561 diff: Diff::ONE,
3562 });
3563 }
3564 (None, Some(prev)) => {
3565 entry.push(TransactionUpdate {
3566 value: prev,
3567 ts,
3568 diff: Diff::MINUS_ONE,
3569 });
3570 }
3571 (None, None) => {}
3572 }
3573
3574 if needs_uniqueness_check {
3576 if let Err(err) = self.verify_keys([&k]) {
3577 let pending = self.pending.get_mut(&k).expect("inserted above");
3580 pending.truncate(restore_len);
3581 return Err(err);
3582 }
3583 }
3584 Ok(prev)
3585 }
3586
3587 fn set_many(
3592 &mut self,
3593 kvs: BTreeMap<K, Option<V>>,
3594 ts: Timestamp,
3595 ) -> Result<BTreeMap<K, Option<V>>, DurableCatalogError> {
3596 if kvs.is_empty() {
3597 return Ok(BTreeMap::new());
3598 }
3599
3600 let mut prevs = BTreeMap::new();
3601 let mut restores = BTreeMap::new();
3602 let mut keys_to_verify_uniqueness = Vec::new();
3604
3605 for (k, v) in kvs {
3606 let prev = self.get(&k).cloned();
3607 if self.update_needs_uniqueness_check(prev.as_ref(), v.as_ref()) {
3608 keys_to_verify_uniqueness.push(k.clone());
3609 }
3610 let entry = self.pending.entry(k.clone()).or_default();
3611 restores.insert(k.clone(), entry.len());
3612
3613 match (v, prev.clone()) {
3614 (Some(v), Some(prev)) => {
3615 entry.push(TransactionUpdate {
3616 value: prev,
3617 ts,
3618 diff: Diff::MINUS_ONE,
3619 });
3620 entry.push(TransactionUpdate {
3621 value: v,
3622 ts,
3623 diff: Diff::ONE,
3624 });
3625 }
3626 (Some(v), None) => {
3627 entry.push(TransactionUpdate {
3628 value: v,
3629 ts,
3630 diff: Diff::ONE,
3631 });
3632 }
3633 (None, Some(prev)) => {
3634 entry.push(TransactionUpdate {
3635 value: prev,
3636 ts,
3637 diff: Diff::MINUS_ONE,
3638 });
3639 }
3640 (None, None) => {}
3641 }
3642
3643 prevs.insert(k, prev);
3644 }
3645
3646 if let Err(err) = self.verify_keys(keys_to_verify_uniqueness.iter()) {
3648 for (k, restore_len) in restores {
3649 let pending = self.pending.get_mut(&k).expect("inserted above");
3652 pending.truncate(restore_len);
3653 }
3654 Err(err)
3655 } else {
3656 Ok(prevs)
3657 }
3658 }
3659
3660 fn delete<F: Fn(&K, &V) -> bool>(&mut self, f: F, ts: Timestamp) -> Vec<(K, V)> {
3666 let mut deleted = Vec::new();
3667 self.for_values_mut(|p, k, v| {
3668 if f(k, v) {
3669 deleted.push((k.clone(), v.clone()));
3670 p.entry(k.clone()).or_default().push(TransactionUpdate {
3671 value: v.clone(),
3672 ts,
3673 diff: Diff::MINUS_ONE,
3674 });
3675 }
3676 });
3677 soft_assert_no_log!(self.verify().is_ok());
3678 deleted
3679 }
3680
3681 fn delete_by_key(&mut self, k: K, ts: Timestamp) -> Option<V> {
3685 self.set(k, None, ts)
3686 .expect("deleting an entry cannot violate uniqueness")
3687 }
3688
3689 fn delete_by_keys(&mut self, ks: impl IntoIterator<Item = K>, ts: Timestamp) -> Vec<(K, V)> {
3693 let kvs = ks.into_iter().map(|k| (k, None)).collect();
3694 let prevs = self
3695 .set_many(kvs, ts)
3696 .expect("deleting entries cannot violate uniqueness");
3697 prevs
3698 .into_iter()
3699 .filter_map(|(k, v)| v.map(|v| (k, v)))
3700 .collect()
3701 }
3702}
3703
3704#[cfg(test)]
3705#[allow(clippy::unwrap_used)]
3706mod tests {
3707 use super::*;
3708
3709 use mz_controller::clusters::ReplicaLogging;
3710 use mz_ore::now::SYSTEM_TIME;
3711 use mz_ore::{assert_none, assert_ok};
3712 use mz_persist_client::cache::PersistClientCache;
3713 use mz_persist_types::PersistLocation;
3714 use semver::Version;
3715
3716 use crate::durable::{
3717 ReplicaConfig, ReplicaLocation, TestCatalogStateBuilder, test_bootstrap_args,
3718 };
3719 use crate::memory;
3720
3721 #[mz_ore::test]
3722 fn test_table_transaction_simple() {
3723 fn uniqueness_violation(a: &String, b: &String) -> bool {
3724 a == b
3725 }
3726 let mut table = TableTransaction::new_with_uniqueness_fn(
3727 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "a".to_string())]),
3728 uniqueness_violation,
3729 uniqueness_violation,
3730 )
3731 .unwrap();
3732
3733 assert_ok!(table.insert(2i64.to_le_bytes().to_vec(), "b".to_string(), 0));
3736 assert_ok!(table.insert(3i64.to_le_bytes().to_vec(), "c".to_string(), 0));
3737 assert!(
3738 table
3739 .insert(1i64.to_le_bytes().to_vec(), "c".to_string(), 0)
3740 .is_err()
3741 );
3742 assert!(
3743 table
3744 .insert(4i64.to_le_bytes().to_vec(), "c".to_string(), 0)
3745 .is_err()
3746 );
3747 }
3748
3749 #[mz_ore::test]
3750 fn test_skip_scan_when_unique_key_unchanged() {
3751 fn first_char_same(prev: &String, next: &String) -> bool {
3752 prev.chars().next() == next.chars().next()
3753 }
3754
3755 fn panic_uniqueness_violation(_: &String, _: &String) -> bool {
3757 panic!("uniqueness scan ran for an update that kept the same unique key");
3758 }
3759 let mut table = TableTransaction::new_with_uniqueness_fn(
3760 BTreeMap::from([
3761 (1i64.to_le_bytes().to_vec(), "a1".to_string()),
3762 (2i64.to_le_bytes().to_vec(), "b1".to_string()),
3763 ]),
3764 panic_uniqueness_violation,
3765 first_char_same,
3767 )
3768 .unwrap();
3769 assert!(
3772 table
3773 .update_by_key(1i64.to_le_bytes().to_vec(), "a2".to_string(), 0)
3774 .unwrap()
3775 );
3776
3777 fn real_uniqueness_violation(a: &String, b: &String) -> bool {
3779 a.chars().next() == b.chars().next()
3780 }
3781 let mut table = TableTransaction::new_with_uniqueness_fn(
3782 BTreeMap::from([
3783 (1i64.to_le_bytes().to_vec(), "a1".to_string()),
3784 (2i64.to_le_bytes().to_vec(), "b1".to_string()),
3785 ]),
3786 real_uniqueness_violation,
3787 first_char_same,
3788 )
3789 .unwrap();
3790 assert!(
3793 table
3794 .update_by_key(1i64.to_le_bytes().to_vec(), "b2".to_string(), 0)
3795 .is_err()
3796 );
3797 }
3798
3799 #[mz_ore::test]
3800 fn test_table_transaction() {
3801 fn uniqueness_violation(a: &String, b: &String) -> bool {
3802 a == b
3803 }
3804 let mut table: BTreeMap<Vec<u8>, String> = BTreeMap::new();
3805
3806 fn commit(
3807 table: &mut BTreeMap<Vec<u8>, String>,
3808 mut pending: Vec<(Vec<u8>, String, Diff)>,
3809 ) {
3810 pending.sort_by(|a, b| a.2.cmp(&b.2));
3812 for (k, v, diff) in pending {
3813 if diff == Diff::MINUS_ONE {
3814 let prev = table.remove(&k);
3815 assert_eq!(prev, Some(v));
3816 } else if diff == Diff::ONE {
3817 let prev = table.insert(k, v);
3818 assert_eq!(prev, None);
3819 } else {
3820 panic!("unexpected diff: {diff}");
3821 }
3822 }
3823 }
3824
3825 table.insert(1i64.to_le_bytes().to_vec(), "v1".to_string());
3826 table.insert(2i64.to_le_bytes().to_vec(), "v2".to_string());
3827 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3828 table.clone(),
3829 uniqueness_violation,
3830 uniqueness_violation,
3831 )
3832 .unwrap();
3833 assert_eq!(table_txn.items_cloned(), table);
3834 assert_eq!(table_txn.delete(|_k, _v| false, 0).len(), 0);
3835 assert_eq!(table_txn.delete(|_k, v| v == "v2", 1).len(), 1);
3836 assert_eq!(
3837 table_txn.items_cloned(),
3838 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v1".to_string())])
3839 );
3840 assert_eq!(
3841 table_txn
3842 .update(|_k, _v| Some("v3".to_string()), 2)
3843 .unwrap(),
3844 Diff::ONE
3845 );
3846
3847 table_txn
3849 .insert(3i64.to_le_bytes().to_vec(), "v3".to_string(), 3)
3850 .unwrap_err();
3851
3852 table_txn
3853 .insert(3i64.to_le_bytes().to_vec(), "v4".to_string(), 4)
3854 .unwrap();
3855 assert_eq!(
3856 table_txn.items_cloned(),
3857 BTreeMap::from([
3858 (1i64.to_le_bytes().to_vec(), "v3".to_string()),
3859 (3i64.to_le_bytes().to_vec(), "v4".to_string()),
3860 ])
3861 );
3862 let err = table_txn
3863 .update(|_k, _v| Some("v1".to_string()), 5)
3864 .unwrap_err();
3865 assert!(
3866 matches!(err, DurableCatalogError::UniquenessViolation),
3867 "unexpected err: {err:?}"
3868 );
3869 let pending = table_txn.pending();
3870 assert_eq!(
3871 pending,
3872 vec![
3873 (
3874 1i64.to_le_bytes().to_vec(),
3875 "v1".to_string(),
3876 Diff::MINUS_ONE
3877 ),
3878 (1i64.to_le_bytes().to_vec(), "v3".to_string(), Diff::ONE),
3879 (
3880 2i64.to_le_bytes().to_vec(),
3881 "v2".to_string(),
3882 Diff::MINUS_ONE
3883 ),
3884 (3i64.to_le_bytes().to_vec(), "v4".to_string(), Diff::ONE),
3885 ]
3886 );
3887 commit(&mut table, pending);
3888 assert_eq!(
3889 table,
3890 BTreeMap::from([
3891 (1i64.to_le_bytes().to_vec(), "v3".to_string()),
3892 (3i64.to_le_bytes().to_vec(), "v4".to_string())
3893 ])
3894 );
3895
3896 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3897 table.clone(),
3898 uniqueness_violation,
3899 uniqueness_violation,
3900 )
3901 .unwrap();
3902 assert_eq!(
3904 table_txn.delete(|k, _v| k == &1i64.to_le_bytes(), 0).len(),
3905 1
3906 );
3907 table_txn
3908 .insert(1i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3909 .unwrap();
3910 table_txn
3912 .insert(5i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3913 .unwrap_err();
3914 table_txn
3916 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 0)
3917 .unwrap_err();
3918 assert_eq!(
3919 table_txn.delete(|k, _v| k == &1i64.to_le_bytes(), 0).len(),
3920 1
3921 );
3922 table_txn
3924 .insert(5i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3925 .unwrap();
3926 table_txn
3927 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 0)
3928 .unwrap();
3929 let pending = table_txn.pending();
3930 assert_eq!(
3931 pending,
3932 vec![
3933 (
3934 1i64.to_le_bytes().to_vec(),
3935 "v3".to_string(),
3936 Diff::MINUS_ONE
3937 ),
3938 (1i64.to_le_bytes().to_vec(), "v5".to_string(), Diff::ONE),
3939 (5i64.to_le_bytes().to_vec(), "v3".to_string(), Diff::ONE),
3940 ]
3941 );
3942 commit(&mut table, pending);
3943 assert_eq!(
3944 table,
3945 BTreeMap::from([
3946 (1i64.to_le_bytes().to_vec(), "v5".to_string()),
3947 (3i64.to_le_bytes().to_vec(), "v4".to_string()),
3948 (5i64.to_le_bytes().to_vec(), "v3".to_string()),
3949 ])
3950 );
3951
3952 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3953 table.clone(),
3954 uniqueness_violation,
3955 uniqueness_violation,
3956 )
3957 .unwrap();
3958 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 3);
3959 table_txn
3960 .insert(1i64.to_le_bytes().to_vec(), "v1".to_string(), 0)
3961 .unwrap();
3962
3963 commit(&mut table, table_txn.pending());
3964 assert_eq!(
3965 table,
3966 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v1".to_string()),])
3967 );
3968
3969 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3970 table.clone(),
3971 uniqueness_violation,
3972 uniqueness_violation,
3973 )
3974 .unwrap();
3975 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 1);
3976 table_txn
3977 .insert(1i64.to_le_bytes().to_vec(), "v2".to_string(), 0)
3978 .unwrap();
3979 commit(&mut table, table_txn.pending());
3980 assert_eq!(
3981 table,
3982 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v2".to_string()),])
3983 );
3984
3985 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3987 table.clone(),
3988 uniqueness_violation,
3989 uniqueness_violation,
3990 )
3991 .unwrap();
3992 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 1);
3993 table_txn
3994 .insert(1i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3995 .unwrap();
3996 table_txn
3997 .insert(1i64.to_le_bytes().to_vec(), "v4".to_string(), 1)
3998 .unwrap_err();
3999 assert_eq!(table_txn.delete(|_k, _v| true, 1).len(), 1);
4000 table_txn
4001 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 1)
4002 .unwrap();
4003 commit(&mut table, table_txn.pending());
4004 assert_eq!(
4005 table.clone().into_iter().collect::<Vec<_>>(),
4006 vec![(1i64.to_le_bytes().to_vec(), "v5".to_string())]
4007 );
4008
4009 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4011 table.clone(),
4012 uniqueness_violation,
4013 uniqueness_violation,
4014 )
4015 .unwrap();
4016 table_txn
4018 .set(2i64.to_le_bytes().to_vec(), Some("v5".to_string()), 0)
4019 .unwrap_err();
4020 table_txn
4021 .set(3i64.to_le_bytes().to_vec(), Some("v6".to_string()), 1)
4022 .unwrap();
4023 table_txn.set(2i64.to_le_bytes().to_vec(), None, 2).unwrap();
4024 table_txn.set(1i64.to_le_bytes().to_vec(), None, 2).unwrap();
4025 let pending = table_txn.pending();
4026 assert_eq!(
4027 pending,
4028 vec![
4029 (
4030 1i64.to_le_bytes().to_vec(),
4031 "v5".to_string(),
4032 Diff::MINUS_ONE
4033 ),
4034 (3i64.to_le_bytes().to_vec(), "v6".to_string(), Diff::ONE),
4035 ]
4036 );
4037 commit(&mut table, pending);
4038 assert_eq!(
4039 table,
4040 BTreeMap::from([(3i64.to_le_bytes().to_vec(), "v6".to_string())])
4041 );
4042
4043 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4045 table.clone(),
4046 uniqueness_violation,
4047 uniqueness_violation,
4048 )
4049 .unwrap();
4050 table_txn
4051 .set(3i64.to_le_bytes().to_vec(), Some("v6".to_string()), 0)
4052 .unwrap();
4053 let pending = table_txn.pending::<Vec<u8>, String>();
4054 assert!(pending.is_empty());
4055
4056 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4058 table.clone(),
4059 uniqueness_violation,
4060 uniqueness_violation,
4061 )
4062 .unwrap();
4063 table_txn
4065 .set_many(
4066 BTreeMap::from([
4067 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
4068 (42i64.to_le_bytes().to_vec(), Some("v1".to_string())),
4069 ]),
4070 0,
4071 )
4072 .unwrap_err();
4073 table_txn
4074 .set_many(
4075 BTreeMap::from([
4076 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
4077 (3i64.to_le_bytes().to_vec(), Some("v1".to_string())),
4078 ]),
4079 1,
4080 )
4081 .unwrap();
4082 table_txn
4083 .set_many(
4084 BTreeMap::from([
4085 (42i64.to_le_bytes().to_vec(), Some("v7".to_string())),
4086 (3i64.to_le_bytes().to_vec(), None),
4087 ]),
4088 2,
4089 )
4090 .unwrap();
4091 let pending = table_txn.pending();
4092 assert_eq!(
4093 pending,
4094 vec![
4095 (1i64.to_le_bytes().to_vec(), "v6".to_string(), Diff::ONE),
4096 (
4097 3i64.to_le_bytes().to_vec(),
4098 "v6".to_string(),
4099 Diff::MINUS_ONE
4100 ),
4101 (42i64.to_le_bytes().to_vec(), "v7".to_string(), Diff::ONE),
4102 ]
4103 );
4104 commit(&mut table, pending);
4105 assert_eq!(
4106 table,
4107 BTreeMap::from([
4108 (1i64.to_le_bytes().to_vec(), "v6".to_string()),
4109 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4110 ])
4111 );
4112
4113 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4115 table.clone(),
4116 uniqueness_violation,
4117 uniqueness_violation,
4118 )
4119 .unwrap();
4120 table_txn
4121 .set_many(
4122 BTreeMap::from([
4123 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
4124 (42i64.to_le_bytes().to_vec(), Some("v7".to_string())),
4125 ]),
4126 0,
4127 )
4128 .unwrap();
4129 let pending = table_txn.pending::<Vec<u8>, String>();
4130 assert!(pending.is_empty());
4131 commit(&mut table, pending);
4132 assert_eq!(
4133 table,
4134 BTreeMap::from([
4135 (1i64.to_le_bytes().to_vec(), "v6".to_string()),
4136 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4137 ])
4138 );
4139
4140 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4142 table.clone(),
4143 uniqueness_violation,
4144 uniqueness_violation,
4145 )
4146 .unwrap();
4147 table_txn
4149 .update_by_key(1i64.to_le_bytes().to_vec(), "v7".to_string(), 0)
4150 .unwrap_err();
4151 assert!(
4152 table_txn
4153 .update_by_key(1i64.to_le_bytes().to_vec(), "v8".to_string(), 1)
4154 .unwrap()
4155 );
4156 assert!(
4157 !table_txn
4158 .update_by_key(5i64.to_le_bytes().to_vec(), "v8".to_string(), 2)
4159 .unwrap()
4160 );
4161 let pending = table_txn.pending();
4162 assert_eq!(
4163 pending,
4164 vec![
4165 (
4166 1i64.to_le_bytes().to_vec(),
4167 "v6".to_string(),
4168 Diff::MINUS_ONE
4169 ),
4170 (1i64.to_le_bytes().to_vec(), "v8".to_string(), Diff::ONE),
4171 ]
4172 );
4173 commit(&mut table, pending);
4174 assert_eq!(
4175 table,
4176 BTreeMap::from([
4177 (1i64.to_le_bytes().to_vec(), "v8".to_string()),
4178 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4179 ])
4180 );
4181
4182 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4184 table.clone(),
4185 uniqueness_violation,
4186 uniqueness_violation,
4187 )
4188 .unwrap();
4189 assert!(
4190 table_txn
4191 .update_by_key(1i64.to_le_bytes().to_vec(), "v8".to_string(), 0)
4192 .unwrap()
4193 );
4194 let pending = table_txn.pending::<Vec<u8>, String>();
4195 assert!(pending.is_empty());
4196 commit(&mut table, pending);
4197 assert_eq!(
4198 table,
4199 BTreeMap::from([
4200 (1i64.to_le_bytes().to_vec(), "v8".to_string()),
4201 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4202 ])
4203 );
4204
4205 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4207 table.clone(),
4208 uniqueness_violation,
4209 uniqueness_violation,
4210 )
4211 .unwrap();
4212 table_txn
4214 .update_by_keys(
4215 [
4216 (1i64.to_le_bytes().to_vec(), "v7".to_string()),
4217 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4218 ],
4219 0,
4220 )
4221 .unwrap_err();
4222 let n = table_txn
4223 .update_by_keys(
4224 [
4225 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4226 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4227 ],
4228 1,
4229 )
4230 .unwrap();
4231 assert_eq!(n, Diff::ONE);
4232 let n = table_txn
4233 .update_by_keys(
4234 [
4235 (15i64.to_le_bytes().to_vec(), "v9".to_string()),
4236 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4237 ],
4238 2,
4239 )
4240 .unwrap();
4241 assert_eq!(n, Diff::ZERO);
4242 let pending = table_txn.pending();
4243 assert_eq!(
4244 pending,
4245 vec![
4246 (
4247 1i64.to_le_bytes().to_vec(),
4248 "v8".to_string(),
4249 Diff::MINUS_ONE
4250 ),
4251 (1i64.to_le_bytes().to_vec(), "v9".to_string(), Diff::ONE),
4252 ]
4253 );
4254 commit(&mut table, pending);
4255 assert_eq!(
4256 table,
4257 BTreeMap::from([
4258 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4259 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4260 ])
4261 );
4262
4263 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4265 table.clone(),
4266 uniqueness_violation,
4267 uniqueness_violation,
4268 )
4269 .unwrap();
4270 let n = table_txn
4271 .update_by_keys(
4272 [
4273 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4274 (42i64.to_le_bytes().to_vec(), "v7".to_string()),
4275 ],
4276 0,
4277 )
4278 .unwrap();
4279 assert_eq!(n, Diff::from(2));
4280 let pending = table_txn.pending::<Vec<u8>, String>();
4281 assert!(pending.is_empty());
4282 commit(&mut table, pending);
4283 assert_eq!(
4284 table,
4285 BTreeMap::from([
4286 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4287 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4288 ])
4289 );
4290
4291 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4293 table.clone(),
4294 uniqueness_violation,
4295 uniqueness_violation,
4296 )
4297 .unwrap();
4298 let prev = table_txn.delete_by_key(1i64.to_le_bytes().to_vec(), 0);
4299 assert_eq!(prev, Some("v9".to_string()));
4300 let prev = table_txn.delete_by_key(5i64.to_le_bytes().to_vec(), 1);
4301 assert_none!(prev);
4302 let prev = table_txn.delete_by_key(1i64.to_le_bytes().to_vec(), 2);
4303 assert_none!(prev);
4304 let pending = table_txn.pending();
4305 assert_eq!(
4306 pending,
4307 vec![(
4308 1i64.to_le_bytes().to_vec(),
4309 "v9".to_string(),
4310 Diff::MINUS_ONE
4311 ),]
4312 );
4313 commit(&mut table, pending);
4314 assert_eq!(
4315 table,
4316 BTreeMap::from([(42i64.to_le_bytes().to_vec(), "v7".to_string())])
4317 );
4318
4319 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4321 table.clone(),
4322 uniqueness_violation,
4323 uniqueness_violation,
4324 )
4325 .unwrap();
4326 let prevs = table_txn.delete_by_keys(
4327 [42i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4328 0,
4329 );
4330 assert_eq!(
4331 prevs,
4332 vec![(42i64.to_le_bytes().to_vec(), "v7".to_string())]
4333 );
4334 let prevs = table_txn.delete_by_keys(
4335 [42i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4336 1,
4337 );
4338 assert_eq!(prevs, vec![]);
4339 let prevs = table_txn.delete_by_keys(
4340 [10i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4341 2,
4342 );
4343 assert_eq!(prevs, vec![]);
4344 let pending = table_txn.pending();
4345 assert_eq!(
4346 pending,
4347 vec![(
4348 42i64.to_le_bytes().to_vec(),
4349 "v7".to_string(),
4350 Diff::MINUS_ONE
4351 ),]
4352 );
4353 commit(&mut table, pending);
4354 assert_eq!(table, BTreeMap::new());
4355 }
4356
4357 #[mz_ore::test(tokio::test)]
4358 #[cfg_attr(miri, ignore)] async fn test_savepoint() {
4360 const VERSION: Version = Version::new(26, 0, 0);
4361 let mut persist_cache = PersistClientCache::new_no_metrics();
4362 persist_cache.cfg.build_version = VERSION;
4363 let persist_client = persist_cache
4364 .open(PersistLocation::new_in_mem())
4365 .await
4366 .unwrap();
4367 let state_builder = TestCatalogStateBuilder::new(persist_client)
4368 .with_default_deploy_generation()
4369 .with_version(VERSION);
4370
4371 let _ = state_builder
4373 .clone()
4374 .unwrap_build()
4375 .await
4376 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4377 .await
4378 .unwrap();
4379 let mut savepoint_state = state_builder
4380 .unwrap_build()
4381 .await
4382 .open_savepoint(SYSTEM_TIME().into(), &test_bootstrap_args())
4383 .await
4384 .unwrap();
4385
4386 let initial_snapshot = savepoint_state.sync_to_current_updates().await.unwrap();
4387 assert!(!initial_snapshot.is_empty());
4388
4389 let db_name = "db";
4390 let db_owner = RoleId::User(42);
4391 let db_privileges = Vec::new();
4392 let mut txn = savepoint_state.transaction().await.unwrap();
4393 let (db_id, db_oid) = txn
4394 .insert_user_database(db_name, db_owner, db_privileges.clone(), &HashSet::new())
4395 .unwrap();
4396 let commit_ts = txn.upper();
4397 txn.commit_internal(commit_ts).await.unwrap();
4398 let updates = savepoint_state.sync_to_current_updates().await.unwrap();
4399 let update = updates.into_element();
4400
4401 assert_eq!(update.diff, StateDiff::Addition);
4402
4403 let db = match update.kind {
4404 memory::objects::StateUpdateKind::Database(db) => db,
4405 update => panic!("unexpected update: {update:?}"),
4406 };
4407
4408 assert_eq!(db_id, db.id);
4409 assert_eq!(db_oid, db.oid);
4410 assert_eq!(db_name, db.name);
4411 assert_eq!(db_owner, db.owner_id);
4412 assert_eq!(db_privileges, db.privileges);
4413 }
4414
4415 #[mz_ore::test(tokio::test)]
4416 #[cfg_attr(miri, ignore)] async fn test_dry_run_transaction_rejects_internal_commit() {
4418 const VERSION: Version = Version::new(26, 0, 0);
4419 let mut persist_cache = PersistClientCache::new_no_metrics();
4420 persist_cache.cfg.build_version = VERSION;
4421 let persist_client = persist_cache
4422 .open(PersistLocation::new_in_mem())
4423 .await
4424 .unwrap();
4425 let mut state = TestCatalogStateBuilder::new(persist_client)
4426 .with_default_deploy_generation()
4427 .with_version(VERSION)
4428 .unwrap_build()
4429 .await
4430 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4431 .await
4432 .unwrap();
4433 let _ = state.sync_to_current_updates().await.unwrap();
4434
4435 let initial_id = state.get_next_id(USER_ITEM_ALLOC_KEY).await.unwrap();
4436 let initial_upper = state.current_upper().await;
4437 let snapshot = state.snapshot().await.unwrap();
4438 let mut dry_run = state.transaction_from_snapshot(snapshot).unwrap();
4439 let ids = dry_run
4440 .transaction_mut()
4441 .get_and_increment_id_by(USER_ITEM_ALLOC_KEY.to_string(), 1)
4442 .unwrap();
4443 assert_eq!(ids, vec![initial_id]);
4444
4445 let transaction = dry_run.transaction;
4446 let err = transaction
4447 .commit_internal(initial_upper)
4448 .await
4449 .unwrap_err();
4450 assert!(matches!(
4451 err,
4452 CatalogError::Durable(DurableCatalogError::DryRunTransaction)
4453 ));
4454 assert_eq!(state.current_upper().await, initial_upper);
4455 assert_eq!(
4456 state.get_next_id(USER_ITEM_ALLOC_KEY).await.unwrap(),
4457 initial_id
4458 );
4459 }
4460
4461 #[mz_ore::test(tokio::test)]
4462 #[cfg_attr(miri, ignore)] async fn test_dry_run_transaction_rejects_into_parts_escape() {
4464 const VERSION: Version = Version::new(26, 0, 0);
4465 let mut persist_cache = PersistClientCache::new_no_metrics();
4466 persist_cache.cfg.build_version = VERSION;
4467 let persist_client = persist_cache
4468 .open(PersistLocation::new_in_mem())
4469 .await
4470 .unwrap();
4471 let mut dry_run_state = TestCatalogStateBuilder::new(persist_client.clone())
4472 .with_default_deploy_generation()
4473 .with_version(VERSION)
4474 .unwrap_build()
4475 .await
4476 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4477 .await
4478 .unwrap();
4479 let mut replacement_state = TestCatalogStateBuilder::new(persist_client)
4480 .with_default_deploy_generation()
4481 .with_version(VERSION)
4482 .unwrap_build()
4483 .await
4484 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4485 .await
4486 .unwrap();
4487 let _ = dry_run_state.sync_to_current_updates().await.unwrap();
4488 let _ = replacement_state.sync_to_current_updates().await.unwrap();
4489
4490 let initial_id = dry_run_state
4491 .get_next_id(USER_ITEM_ALLOC_KEY)
4492 .await
4493 .unwrap();
4494 let initial_upper = dry_run_state.current_upper().await;
4495 let snapshot = dry_run_state.snapshot().await.unwrap();
4496 let mut dry_run = dry_run_state.transaction_from_snapshot(snapshot).unwrap();
4497 let ids = dry_run
4498 .transaction_mut()
4499 .get_and_increment_id_by(USER_ITEM_ALLOC_KEY.to_string(), 1)
4500 .unwrap();
4501 assert_eq!(ids, vec![initial_id]);
4502
4503 let replacement = replacement_state.transaction().await.unwrap();
4504 let escaped = std::mem::replace(dry_run.transaction_mut(), replacement);
4505 drop(dry_run);
4506
4507 let err = match escaped.into_parts() {
4508 Ok(_) => panic!("dry-run transaction decomposed into committable parts"),
4509 Err(err) => err,
4510 };
4511 assert!(matches!(
4512 err,
4513 CatalogError::Durable(DurableCatalogError::DryRunTransaction)
4514 ));
4515 assert_eq!(dry_run_state.current_upper().await, initial_upper);
4516 assert_eq!(
4517 dry_run_state
4518 .get_next_id(USER_ITEM_ALLOC_KEY)
4519 .await
4520 .unwrap(),
4521 initial_id
4522 );
4523 }
4524
4525 #[mz_ore::test(tokio::test)]
4529 #[cfg_attr(miri, ignore)] async fn test_insert_replica_with_id_does_not_consume_allocator() {
4531 const VERSION: Version = Version::new(26, 0, 0);
4532 let mut persist_cache = PersistClientCache::new_no_metrics();
4533 persist_cache.cfg.build_version = VERSION;
4534 let persist_client = persist_cache
4535 .open(PersistLocation::new_in_mem())
4536 .await
4537 .unwrap();
4538 let state_builder = TestCatalogStateBuilder::new(persist_client)
4539 .with_default_deploy_generation()
4540 .with_version(VERSION);
4541 let mut state = state_builder
4542 .unwrap_build()
4543 .await
4544 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4545 .await
4546 .unwrap();
4547
4548 let cluster_id = ClusterId::User(1);
4551 let owner_id = RoleId::User(1);
4552 let config = ReplicaConfig {
4553 location: ReplicaLocation::Managed {
4554 size: "1".to_string(),
4555 availability_zones: Vec::new(),
4556 internal: false,
4557 billed_as: None,
4558 pending: false,
4559 },
4560 logging: ReplicaLogging {
4561 log_logging: false,
4562 interval: Some(Duration::from_secs(1)),
4563 },
4564 arrangement_compression: false,
4565 };
4566
4567 let commit_ts = state.current_upper().await;
4569 let a = state
4570 .allocate_user_replica_ids(1, commit_ts)
4571 .await
4572 .unwrap()
4573 .into_element();
4574 assert!(a.is_user());
4575
4576 let initial_updates = state.sync_to_current_updates().await.unwrap();
4577 assert!(!initial_updates.is_empty());
4578
4579 let mut txn = state.transaction().await.unwrap();
4581 txn.insert_cluster_replica_with_id(cluster_id, a, "explicit", config, owner_id)
4582 .unwrap();
4583 let commit_ts = txn.upper();
4584 txn.commit_internal(commit_ts).await.unwrap();
4585 let _ = state.sync_to_current_updates().await.unwrap();
4586
4587 let commit_ts = state.current_upper().await;
4589 let b = state
4590 .allocate_user_replica_ids(1, commit_ts)
4591 .await
4592 .unwrap()
4593 .into_element();
4594
4595 assert_eq!(b.inner_id(), a.inner_id() + 1);
4598
4599 let txn = state.transaction().await.unwrap();
4601 let found = txn
4602 .get_cluster_replicas()
4603 .any(|replica| replica.replica_id == a);
4604 assert!(found, "explicitly inserted replica {a} not found");
4605 }
4606
4607 #[mz_ore::test]
4608 fn test_allocate_introspection_source_index_id() {
4609 let cluster_variant: u8 = 0b0000_0001;
4610 let cluster_id_inner: u64 =
4611 0b0000_0000_1100_0101_1100_0011_1010_1101_0000_1011_1111_1001_0110_1010;
4612 let timely_messages_received_log_variant: u8 = 0b0000_1000;
4613
4614 let cluster_id = ClusterId::System(cluster_id_inner);
4615 let log_variant = LogVariant::Timely(TimelyLog::MessagesReceived);
4616
4617 let introspection_source_index_id: u64 =
4618 0b0000_0001_1100_0101_1100_0011_1010_1101_0000_1011_1111_1001_0110_1010_0000_1000;
4619
4620 {
4622 let mut cluster_variant_mask = 0xFF << 56;
4623 cluster_variant_mask &= introspection_source_index_id;
4624 cluster_variant_mask >>= 56;
4625 assert_eq!(cluster_variant_mask, u64::from(cluster_variant));
4626 }
4627
4628 {
4630 let mut cluster_id_inner_mask = 0xFFFF_FFFF_FFFF << 8;
4631 cluster_id_inner_mask &= introspection_source_index_id;
4632 cluster_id_inner_mask >>= 8;
4633 assert_eq!(cluster_id_inner_mask, cluster_id_inner);
4634 }
4635
4636 {
4638 let mut log_variant_mask = 0xFF;
4639 log_variant_mask &= introspection_source_index_id;
4640 assert_eq!(
4641 log_variant_mask,
4642 u64::from(timely_messages_received_log_variant)
4643 );
4644 }
4645
4646 let (catalog_item_id, global_id) =
4647 Transaction::allocate_introspection_source_index_id(&cluster_id, log_variant);
4648
4649 assert_eq!(
4650 catalog_item_id,
4651 CatalogItemId::IntrospectionSourceIndex(introspection_source_index_id)
4652 );
4653 assert_eq!(
4654 global_id,
4655 GlobalId::IntrospectionSourceIndex(introspection_source_index_id)
4656 );
4657 }
4658}