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(Derivative)]
80#[derivative(Debug)]
81pub struct Transaction<'a> {
82 #[derivative(Debug = "ignore")]
83 #[derivative(PartialEq = "ignore")]
84 durable_catalog: &'a mut dyn DurableCatalogState,
85 databases: TableTransaction<DatabaseKey, DatabaseValue>,
86 schemas: TableTransaction<SchemaKey, SchemaValue>,
87 items: TableTransaction<ItemKey, ItemValue>,
88 comments: TableTransaction<CommentKey, CommentValue>,
89 roles: TableTransaction<RoleKey, RoleValue>,
90 role_auth: TableTransaction<RoleAuthKey, RoleAuthValue>,
91 clusters: TableTransaction<ClusterKey, ClusterValue>,
92 cluster_replicas: TableTransaction<ClusterReplicaKey, ClusterReplicaValue>,
93 introspection_sources:
94 TableTransaction<ClusterIntrospectionSourceIndexKey, ClusterIntrospectionSourceIndexValue>,
95 id_allocator: TableTransaction<IdAllocKey, IdAllocValue>,
96 configs: TableTransaction<ConfigKey, ConfigValue>,
97 settings: TableTransaction<SettingKey, SettingValue>,
98 system_gid_mapping: TableTransaction<GidMappingKey, GidMappingValue>,
99 system_configurations: TableTransaction<ServerConfigurationKey, ServerConfigurationValue>,
100 cluster_system_configurations:
101 TableTransaction<ClusterSystemConfigurationKey, ClusterSystemConfigurationValue>,
102 replica_system_configurations:
103 TableTransaction<ReplicaSystemConfigurationKey, ReplicaSystemConfigurationValue>,
104 default_privileges: TableTransaction<DefaultPrivilegesKey, DefaultPrivilegesValue>,
105 source_references: TableTransaction<SourceReferencesKey, SourceReferencesValue>,
106 system_privileges: TableTransaction<SystemPrivilegesKey, SystemPrivilegesValue>,
107 network_policies: TableTransaction<NetworkPolicyKey, NetworkPolicyValue>,
108 storage_collection_metadata:
109 TableTransaction<StorageCollectionMetadataKey, StorageCollectionMetadataValue>,
110 unfinalized_shards: TableTransaction<UnfinalizedShardKey, ()>,
111 txn_wal_shard: TableTransaction<(), TxnWalShardValue>,
112 audit_log_updates: Vec<(AuditLogKey, Diff, Timestamp)>,
115 upper: mz_repr::Timestamp,
117 op_id: Timestamp,
119}
120
121impl<'a> Transaction<'a> {
122 pub fn new(
123 durable_catalog: &'a mut dyn DurableCatalogState,
124 Snapshot {
125 databases,
126 schemas,
127 roles,
128 role_auth,
129 items,
130 comments,
131 clusters,
132 network_policies,
133 cluster_replicas,
134 introspection_sources,
135 id_allocator,
136 configs,
137 settings,
138 source_references,
139 system_object_mappings,
140 system_configurations,
141 cluster_system_configurations,
142 replica_system_configurations,
143 default_privileges,
144 system_privileges,
145 storage_collection_metadata,
146 unfinalized_shards,
147 txn_wal_shard,
148 }: Snapshot,
149 upper: mz_repr::Timestamp,
150 ) -> Result<Transaction<'a>, CatalogError> {
151 let database_unique_fn: fn(&DatabaseValue, &DatabaseValue) -> bool =
154 |a, b| a.name == b.name;
155 let schema_unique_fn: fn(&SchemaValue, &SchemaValue) -> bool =
156 |a, b| a.database_id == b.database_id && a.name == b.name;
157 let role_key: fn(&RoleValue, &RoleValue) -> bool = |a, b| a.name == b.name;
158 let cluster_unique_fn: fn(&ClusterValue, &ClusterValue) -> bool = |a, b| a.name == b.name;
159 let network_policy_unique_fn: fn(&NetworkPolicyValue, &NetworkPolicyValue) -> bool =
160 |a, b| a.name == b.name;
161 let cluster_replica_unique_fn: fn(&ClusterReplicaValue, &ClusterReplicaValue) -> bool =
162 |a, b| a.cluster_id == b.cluster_id && a.name == b.name;
163
164 Ok(Transaction {
165 durable_catalog,
166 databases: TableTransaction::new_with_uniqueness_fn(
167 databases,
168 database_unique_fn,
169 database_unique_fn,
170 )?,
171 schemas: TableTransaction::new_with_uniqueness_fn(
172 schemas,
173 schema_unique_fn,
174 schema_unique_fn,
175 )?,
176 items: TableTransaction::new_with_uniqueness_fn(
177 items,
178 |a: &ItemValue, b| {
179 a.schema_id == b.schema_id && a.name == b.name && {
180 let a_type = a.item_type();
182 let b_type = b.item_type();
183 (a_type != CatalogItemType::Type && b_type != CatalogItemType::Type)
184 || (a_type == CatalogItemType::Type && b_type.conflicts_with_type())
185 || (b_type == CatalogItemType::Type && a_type.conflicts_with_type())
186 }
187 },
188 |prev: &ItemValue, next| {
189 prev.schema_id == next.schema_id
190 && prev.name == next.name
191 && prev.item_type() == next.item_type()
193 },
194 )?,
195 comments: TableTransaction::new(comments)?,
196 roles: TableTransaction::new_with_uniqueness_fn(roles, role_key, role_key)?,
197 role_auth: TableTransaction::new(role_auth)?,
198 clusters: TableTransaction::new_with_uniqueness_fn(
199 clusters,
200 cluster_unique_fn,
201 cluster_unique_fn,
202 )?,
203 network_policies: TableTransaction::new_with_uniqueness_fn(
204 network_policies,
205 network_policy_unique_fn,
206 network_policy_unique_fn,
207 )?,
208 cluster_replicas: TableTransaction::new_with_uniqueness_fn(
209 cluster_replicas,
210 cluster_replica_unique_fn,
211 cluster_replica_unique_fn,
212 )?,
213 introspection_sources: TableTransaction::new(introspection_sources)?,
214 id_allocator: TableTransaction::new(id_allocator)?,
215 configs: TableTransaction::new(configs)?,
216 settings: TableTransaction::new(settings)?,
217 source_references: TableTransaction::new(source_references)?,
218 system_gid_mapping: TableTransaction::new(system_object_mappings)?,
219 system_configurations: TableTransaction::new(system_configurations)?,
220 cluster_system_configurations: TableTransaction::new(cluster_system_configurations)?,
221 replica_system_configurations: TableTransaction::new(replica_system_configurations)?,
222 default_privileges: TableTransaction::new(default_privileges)?,
223 system_privileges: TableTransaction::new(system_privileges)?,
224 storage_collection_metadata: TableTransaction::new(storage_collection_metadata)?,
225 unfinalized_shards: TableTransaction::new(unfinalized_shards)?,
226 txn_wal_shard: TableTransaction::new(txn_wal_shard)?,
230 audit_log_updates: Vec::new(),
231 upper,
232 op_id: 0,
233 })
234 }
235
236 pub fn get_item(&self, id: &CatalogItemId) -> Option<Item> {
237 let key = ItemKey { id: *id };
238 self.items
239 .get(&key)
240 .map(|v| DurableType::from_key_value(key, v.clone()))
241 }
242
243 pub fn get_items(&self) -> impl Iterator<Item = Item> + use<> {
244 self.items
245 .items()
246 .into_iter()
247 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
248 .sorted_by_key(|Item { id, .. }| *id)
249 }
250
251 pub fn insert_audit_log_event(&mut self, event: VersionedEvent) {
252 self.insert_audit_log_events([event]);
253 }
254
255 pub fn insert_audit_log_events(&mut self, events: impl IntoIterator<Item = VersionedEvent>) {
256 let events = events
257 .into_iter()
258 .map(|event| (AuditLogKey { event }, Diff::ONE, self.op_id));
259 self.audit_log_updates.extend(events);
260 }
261
262 pub fn insert_user_database(
263 &mut self,
264 database_name: &str,
265 owner_id: RoleId,
266 privileges: Vec<MzAclItem>,
267 temporary_oids: &HashSet<u32>,
268 ) -> Result<(DatabaseId, u32), CatalogError> {
269 let id = self.get_and_increment_id(DATABASE_ID_ALLOC_KEY.to_string())?;
270 let id = DatabaseId::User(id);
271 let oid = self.allocate_oid(temporary_oids)?;
272 self.insert_database(id, database_name, owner_id, privileges, oid)?;
273 Ok((id, oid))
274 }
275
276 pub(crate) fn insert_database(
277 &mut self,
278 id: DatabaseId,
279 database_name: &str,
280 owner_id: RoleId,
281 privileges: Vec<MzAclItem>,
282 oid: u32,
283 ) -> Result<u32, CatalogError> {
284 match self.databases.insert(
285 DatabaseKey { id },
286 DatabaseValue {
287 name: database_name.to_string(),
288 owner_id,
289 privileges,
290 oid,
291 },
292 self.op_id,
293 ) {
294 Ok(_) => Ok(oid),
295 Err(_) => Err(SqlCatalogError::DatabaseAlreadyExists(database_name.to_owned()).into()),
296 }
297 }
298
299 pub fn insert_user_schema(
300 &mut self,
301 database_id: DatabaseId,
302 schema_name: &str,
303 owner_id: RoleId,
304 privileges: Vec<MzAclItem>,
305 temporary_oids: &HashSet<u32>,
306 ) -> Result<(SchemaId, u32), CatalogError> {
307 let id = self.get_and_increment_id(SCHEMA_ID_ALLOC_KEY.to_string())?;
308 let id = SchemaId::User(id);
309 let oid = self.allocate_oid(temporary_oids)?;
310 self.insert_schema(
311 id,
312 Some(database_id),
313 schema_name.to_string(),
314 owner_id,
315 privileges,
316 oid,
317 )?;
318 Ok((id, oid))
319 }
320
321 pub fn insert_system_schema(
322 &mut self,
323 schema_id: u64,
324 schema_name: &str,
325 owner_id: RoleId,
326 privileges: Vec<MzAclItem>,
327 oid: u32,
328 ) -> Result<(), CatalogError> {
329 let id = SchemaId::System(schema_id);
330 self.insert_schema(id, None, schema_name.to_string(), owner_id, privileges, oid)
331 }
332
333 pub(crate) fn insert_schema(
334 &mut self,
335 schema_id: SchemaId,
336 database_id: Option<DatabaseId>,
337 schema_name: String,
338 owner_id: RoleId,
339 privileges: Vec<MzAclItem>,
340 oid: u32,
341 ) -> Result<(), CatalogError> {
342 match self.schemas.insert(
343 SchemaKey { id: schema_id },
344 SchemaValue {
345 database_id,
346 name: schema_name.clone(),
347 owner_id,
348 privileges,
349 oid,
350 },
351 self.op_id,
352 ) {
353 Ok(_) => Ok(()),
354 Err(_) => Err(SqlCatalogError::SchemaAlreadyExists(schema_name).into()),
355 }
356 }
357
358 pub fn insert_builtin_role(
359 &mut self,
360 id: RoleId,
361 name: String,
362 attributes: RoleAttributesRaw,
363 membership: RoleMembership,
364 vars: RoleVars,
365 oid: u32,
366 ) -> Result<RoleId, CatalogError> {
367 soft_assert_or_log!(id.is_builtin(), "ID {id:?} is not builtin");
368 self.insert_role(id, name, attributes, membership, vars, oid)?;
369 Ok(id)
370 }
371
372 pub fn insert_user_role(
373 &mut self,
374 name: String,
375 attributes: RoleAttributesRaw,
376 membership: RoleMembership,
377 vars: RoleVars,
378 temporary_oids: &HashSet<u32>,
379 ) -> Result<(RoleId, u32), CatalogError> {
380 let id = self.get_and_increment_id(USER_ROLE_ID_ALLOC_KEY.to_string())?;
381 let id = RoleId::User(id);
382 let oid = self.allocate_oid(temporary_oids)?;
383 self.insert_role(id, name, attributes, membership, vars, oid)?;
384 Ok((id, oid))
385 }
386
387 fn insert_role(
388 &mut self,
389 id: RoleId,
390 name: String,
391 attributes: RoleAttributesRaw,
392 membership: RoleMembership,
393 vars: RoleVars,
394 oid: u32,
395 ) -> Result<(), CatalogError> {
396 if let Some(ref password) = attributes.password {
397 let hash = mz_auth::hash::scram256_hash(
398 password,
399 &attributes
400 .scram_iterations
401 .or_else(|| {
402 soft_panic_or_log!(
403 "Hash iterations must be set if a password is provided."
404 );
405 None
406 })
407 .unwrap_or_else(|| NonZeroU32::new(600_000).expect("known valid")),
410 )
411 .expect("password hash should be valid");
412 match self.role_auth.insert(
413 RoleAuthKey { role_id: id },
414 RoleAuthValue {
415 password_hash: Some(hash),
416 updated_at: SYSTEM_TIME(),
417 },
418 self.op_id,
419 ) {
420 Ok(_) => {}
421 Err(_) => {
422 return Err(SqlCatalogError::RoleAlreadyExists(name).into());
423 }
424 }
425 }
426
427 match self.roles.insert(
428 RoleKey { id },
429 RoleValue {
430 name: name.clone(),
431 attributes: attributes.into(),
432 membership,
433 vars,
434 oid,
435 },
436 self.op_id,
437 ) {
438 Ok(_) => Ok(()),
439 Err(_) => Err(SqlCatalogError::RoleAlreadyExists(name).into()),
440 }
441 }
442
443 pub fn insert_user_cluster(
445 &mut self,
446 cluster_id: ClusterId,
447 cluster_name: &str,
448 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
449 owner_id: RoleId,
450 privileges: Vec<MzAclItem>,
451 config: ClusterConfig,
452 temporary_oids: &HashSet<u32>,
453 ) -> Result<(), CatalogError> {
454 self.insert_cluster(
455 cluster_id,
456 cluster_name,
457 introspection_source_indexes,
458 owner_id,
459 privileges,
460 config,
461 temporary_oids,
462 )
463 }
464
465 pub fn insert_system_cluster(
467 &mut self,
468 cluster_name: &str,
469 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
470 privileges: Vec<MzAclItem>,
471 owner_id: RoleId,
472 config: ClusterConfig,
473 temporary_oids: &HashSet<u32>,
474 ) -> Result<ClusterId, CatalogError> {
475 let cluster_id = self.get_and_increment_id(SYSTEM_CLUSTER_ID_ALLOC_KEY.to_string())?;
476 let cluster_id = ClusterId::system(cluster_id).ok_or(SqlCatalogError::IdExhaustion)?;
477 self.insert_cluster(
478 cluster_id,
479 cluster_name,
480 introspection_source_indexes,
481 owner_id,
482 privileges,
483 config,
484 temporary_oids,
485 )?;
486 Ok(cluster_id)
487 }
488
489 fn insert_cluster(
490 &mut self,
491 cluster_id: ClusterId,
492 cluster_name: &str,
493 introspection_source_indexes: Vec<(&'static BuiltinLog, CatalogItemId, GlobalId)>,
494 owner_id: RoleId,
495 privileges: Vec<MzAclItem>,
496 config: ClusterConfig,
497 temporary_oids: &HashSet<u32>,
498 ) -> Result<(), CatalogError> {
499 if let Err(_) = self.clusters.insert(
500 ClusterKey { id: cluster_id },
501 ClusterValue {
502 name: cluster_name.to_string(),
503 owner_id,
504 privileges,
505 config,
506 },
507 self.op_id,
508 ) {
509 return Err(SqlCatalogError::ClusterAlreadyExists(cluster_name.to_owned()).into());
510 };
511
512 let amount = usize_to_u64(introspection_source_indexes.len());
513 let oids = self.allocate_oids(amount, temporary_oids)?;
514 let introspection_source_indexes: Vec<_> = introspection_source_indexes
515 .into_iter()
516 .zip_eq(oids)
517 .map(|((builtin, item_id, index_id), oid)| (builtin, item_id, index_id, oid))
518 .collect();
519 for (builtin, item_id, index_id, oid) in introspection_source_indexes {
520 let introspection_source_index = IntrospectionSourceIndex {
521 cluster_id,
522 name: builtin.name.to_string(),
523 item_id,
524 index_id,
525 oid,
526 };
527 let (key, value) = introspection_source_index.into_key_value();
528 self.introspection_sources
529 .insert(key, value, self.op_id)
530 .expect("no uniqueness violation");
531 }
532
533 Ok(())
534 }
535
536 pub fn rename_cluster(
537 &mut self,
538 cluster_id: ClusterId,
539 cluster_name: &str,
540 cluster_to_name: &str,
541 ) -> Result<(), CatalogError> {
542 let key = ClusterKey { id: cluster_id };
543
544 match self.clusters.update(
545 |k, v| {
546 if *k == key {
547 let mut value = v.clone();
548 value.name = cluster_to_name.to_string();
549 Some(value)
550 } else {
551 None
552 }
553 },
554 self.op_id,
555 )? {
556 Diff::ZERO => Err(SqlCatalogError::UnknownCluster(cluster_name.to_string()).into()),
557 Diff::ONE => Ok(()),
558 n => panic!(
559 "Expected to update single cluster {cluster_name} ({cluster_id}), updated {n}"
560 ),
561 }
562 }
563
564 pub fn rename_cluster_replica(
565 &mut self,
566 replica_id: ReplicaId,
567 replica_name: &QualifiedReplica,
568 replica_to_name: &str,
569 ) -> Result<(), CatalogError> {
570 let key = ClusterReplicaKey { id: replica_id };
571
572 match self.cluster_replicas.update(
573 |k, v| {
574 if *k == key {
575 let mut value = v.clone();
576 value.name = replica_to_name.to_string();
577 Some(value)
578 } else {
579 None
580 }
581 },
582 self.op_id,
583 )? {
584 Diff::ZERO => {
585 Err(SqlCatalogError::UnknownClusterReplica(replica_name.to_string()).into())
586 }
587 Diff::ONE => Ok(()),
588 n => panic!(
589 "Expected to update single cluster replica {replica_name} ({replica_id}), updated {n}"
590 ),
591 }
592 }
593
594 pub fn insert_cluster_replica_with_id(
595 &mut self,
596 cluster_id: ClusterId,
597 replica_id: ReplicaId,
598 replica_name: &str,
599 config: ReplicaConfig,
600 owner_id: RoleId,
601 ) -> Result<(), CatalogError> {
602 if let Err(_) = self.cluster_replicas.insert(
603 ClusterReplicaKey { id: replica_id },
604 ClusterReplicaValue {
605 cluster_id,
606 name: replica_name.into(),
607 config,
608 owner_id,
609 },
610 self.op_id,
611 ) {
612 let cluster = self
613 .clusters
614 .get(&ClusterKey { id: cluster_id })
615 .expect("cluster exists");
616 return Err(SqlCatalogError::DuplicateReplica(
617 replica_name.to_string(),
618 cluster.name.to_string(),
619 )
620 .into());
621 };
622 Ok(())
623 }
624
625 pub fn insert_user_network_policy(
626 &mut self,
627 name: String,
628 rules: Vec<NetworkPolicyRule>,
629 privileges: Vec<MzAclItem>,
630 owner_id: RoleId,
631 temporary_oids: &HashSet<u32>,
632 ) -> Result<NetworkPolicyId, CatalogError> {
633 let oid = self.allocate_oid(temporary_oids)?;
634 let id = self.get_and_increment_id(USER_NETWORK_POLICY_ID_ALLOC_KEY.to_string())?;
635 let id = NetworkPolicyId::User(id);
636 self.insert_network_policy(id, name, rules, privileges, owner_id, oid)
637 }
638
639 pub fn insert_network_policy(
640 &mut self,
641 id: NetworkPolicyId,
642 name: String,
643 rules: Vec<NetworkPolicyRule>,
644 privileges: Vec<MzAclItem>,
645 owner_id: RoleId,
646 oid: u32,
647 ) -> Result<NetworkPolicyId, CatalogError> {
648 match self.network_policies.insert(
649 NetworkPolicyKey { id },
650 NetworkPolicyValue {
651 name: name.clone(),
652 rules,
653 privileges,
654 owner_id,
655 oid,
656 },
657 self.op_id,
658 ) {
659 Ok(_) => Ok(id),
660 Err(_) => Err(SqlCatalogError::NetworkPolicyAlreadyExists(name).into()),
661 }
662 }
663
664 pub fn update_introspection_source_index_gids(
669 &mut self,
670 mappings: impl Iterator<
671 Item = (
672 ClusterId,
673 impl Iterator<Item = (String, CatalogItemId, GlobalId, u32)>,
674 ),
675 >,
676 ) -> Result<(), CatalogError> {
677 for (cluster_id, updates) in mappings {
678 for (name, item_id, index_id, oid) in updates {
679 let introspection_source_index = IntrospectionSourceIndex {
680 cluster_id,
681 name,
682 item_id,
683 index_id,
684 oid,
685 };
686 let (key, value) = introspection_source_index.into_key_value();
687
688 let prev = self
689 .introspection_sources
690 .set(key, Some(value), self.op_id)?;
691 if prev.is_none() {
692 return Err(SqlCatalogError::FailedBuiltinSchemaMigration(format!(
693 "{index_id}"
694 ))
695 .into());
696 }
697 }
698 }
699 Ok(())
700 }
701
702 pub fn insert_user_item(
703 &mut self,
704 id: CatalogItemId,
705 global_id: GlobalId,
706 schema_id: SchemaId,
707 item_name: &str,
708 create_sql: String,
709 owner_id: RoleId,
710 privileges: Vec<MzAclItem>,
711 temporary_oids: &HashSet<u32>,
712 versions: BTreeMap<RelationVersion, GlobalId>,
713 ) -> Result<u32, CatalogError> {
714 let oid = self.allocate_oid(temporary_oids)?;
715 self.insert_item(
716 id, oid, global_id, schema_id, item_name, create_sql, owner_id, privileges, versions,
717 )?;
718 Ok(oid)
719 }
720
721 pub fn insert_item(
722 &mut self,
723 id: CatalogItemId,
724 oid: u32,
725 global_id: GlobalId,
726 schema_id: SchemaId,
727 item_name: &str,
728 create_sql: String,
729 owner_id: RoleId,
730 privileges: Vec<MzAclItem>,
731 extra_versions: BTreeMap<RelationVersion, GlobalId>,
732 ) -> Result<(), CatalogError> {
733 match self.items.insert(
734 ItemKey { id },
735 ItemValue {
736 schema_id,
737 name: item_name.to_string(),
738 create_sql,
739 owner_id,
740 privileges,
741 oid,
742 global_id,
743 extra_versions,
744 },
745 self.op_id,
746 ) {
747 Ok(_) => Ok(()),
748 Err(_) => Err(SqlCatalogError::ItemAlreadyExists(id, item_name.to_owned()).into()),
749 }
750 }
751
752 pub fn get_and_increment_id(&mut self, key: String) -> Result<u64, CatalogError> {
753 Ok(self.get_and_increment_id_by(key, 1)?.into_element())
754 }
755
756 pub fn get_and_increment_id_by(
757 &mut self,
758 key: String,
759 amount: u64,
760 ) -> Result<Vec<u64>, CatalogError> {
761 assert!(
762 key != SYSTEM_ITEM_ALLOC_KEY || !self.durable_catalog.is_bootstrap_complete(),
763 "system item IDs cannot be allocated outside of bootstrap"
764 );
765
766 let current_id = self
767 .id_allocator
768 .items()
769 .get(&IdAllocKey { name: key.clone() })
770 .unwrap_or_else(|| panic!("{key} id allocator missing"))
771 .next_id;
772 let next_id = current_id
773 .checked_add(amount)
774 .ok_or(SqlCatalogError::IdExhaustion)?;
775 let prev = self.id_allocator.set(
776 IdAllocKey { name: key },
777 Some(IdAllocValue { next_id }),
778 self.op_id,
779 )?;
780 assert_eq!(
781 prev,
782 Some(IdAllocValue {
783 next_id: current_id
784 })
785 );
786 Ok((current_id..next_id).collect())
787 }
788
789 pub fn allocate_system_item_ids(
790 &mut self,
791 amount: u64,
792 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
793 assert!(
794 !self.durable_catalog.is_bootstrap_complete(),
795 "we can only allocate system item IDs during bootstrap"
796 );
797 Ok(self
798 .get_and_increment_id_by(SYSTEM_ITEM_ALLOC_KEY.to_string(), amount)?
799 .into_iter()
800 .map(|x| (CatalogItemId::System(x), GlobalId::System(x)))
802 .collect())
803 }
804
805 pub fn allocate_introspection_source_index_id(
837 cluster_id: &ClusterId,
838 log_variant: LogVariant,
839 ) -> (CatalogItemId, GlobalId) {
840 let cluster_variant: u8 = match cluster_id {
841 ClusterId::System(_) => 1,
842 ClusterId::User(_) => 2,
843 };
844 let cluster_id: u64 = cluster_id.inner_id();
845 const CLUSTER_ID_MASK: u64 = 0xFFFF << 48;
846 assert_eq!(
847 CLUSTER_ID_MASK & cluster_id,
848 0,
849 "invalid cluster ID: {cluster_id}"
850 );
851 let log_variant: u8 = match log_variant {
852 LogVariant::Timely(TimelyLog::Operates) => 1,
853 LogVariant::Timely(TimelyLog::Channels) => 2,
854 LogVariant::Timely(TimelyLog::Elapsed) => 3,
855 LogVariant::Timely(TimelyLog::Histogram) => 4,
856 LogVariant::Timely(TimelyLog::Addresses) => 5,
857 LogVariant::Timely(TimelyLog::Parks) => 6,
858 LogVariant::Timely(TimelyLog::MessagesSent) => 7,
859 LogVariant::Timely(TimelyLog::MessagesReceived) => 8,
860 LogVariant::Timely(TimelyLog::Reachability) => 9,
861 LogVariant::Timely(TimelyLog::BatchesSent) => 10,
862 LogVariant::Timely(TimelyLog::BatchesReceived) => 11,
863 LogVariant::Differential(DifferentialLog::ArrangementBatches) => 12,
864 LogVariant::Differential(DifferentialLog::ArrangementRecords) => 13,
865 LogVariant::Differential(DifferentialLog::Sharing) => 14,
866 LogVariant::Differential(DifferentialLog::BatcherRecords) => 15,
867 LogVariant::Differential(DifferentialLog::BatcherSize) => 16,
868 LogVariant::Differential(DifferentialLog::BatcherCapacity) => 17,
869 LogVariant::Differential(DifferentialLog::BatcherAllocations) => 18,
870 LogVariant::Compute(ComputeLog::DataflowCurrent) => 19,
871 LogVariant::Compute(ComputeLog::FrontierCurrent) => 20,
872 LogVariant::Compute(ComputeLog::PeekCurrent) => 21,
873 LogVariant::Compute(ComputeLog::PeekDuration) => 22,
874 LogVariant::Compute(ComputeLog::ImportFrontierCurrent) => 23,
875 LogVariant::Compute(ComputeLog::ArrangementHeapSize) => 24,
876 LogVariant::Compute(ComputeLog::ArrangementHeapCapacity) => 25,
877 LogVariant::Compute(ComputeLog::ArrangementHeapAllocations) => 26,
878 LogVariant::Compute(ComputeLog::ErrorCount) => 28,
879 LogVariant::Compute(ComputeLog::HydrationTime) => 29,
880 LogVariant::Compute(ComputeLog::LirMapping) => 30,
881 LogVariant::Compute(ComputeLog::DataflowGlobal) => 31,
882 LogVariant::Compute(ComputeLog::OperatorHydrationStatus) => 32,
883 LogVariant::Compute(ComputeLog::PrometheusMetrics) => 33,
884 };
885
886 let mut id: u64 = u64::from(cluster_variant) << 56;
887 id |= cluster_id << 8;
888 id |= u64::from(log_variant);
889
890 (
891 CatalogItemId::IntrospectionSourceIndex(id),
892 GlobalId::IntrospectionSourceIndex(id),
893 )
894 }
895
896 pub fn allocate_user_item_ids(
897 &mut self,
898 amount: u64,
899 ) -> Result<Vec<(CatalogItemId, GlobalId)>, CatalogError> {
900 Ok(self
901 .get_and_increment_id_by(USER_ITEM_ALLOC_KEY.to_string(), amount)?
902 .into_iter()
903 .map(|x| (CatalogItemId::User(x), GlobalId::User(x)))
905 .collect())
906 }
907
908 pub fn allocate_system_replica_id(&mut self) -> Result<ReplicaId, CatalogError> {
909 let id = self.get_and_increment_id(SYSTEM_REPLICA_ID_ALLOC_KEY.to_string())?;
910 Ok(ReplicaId::System(id))
911 }
912
913 pub fn allocate_audit_log_id(&mut self) -> Result<u64, CatalogError> {
914 self.get_and_increment_id(AUDIT_LOG_ID_ALLOC_KEY.to_string())
915 }
916
917 #[mz_ore::instrument]
920 fn allocate_oids(
921 &mut self,
922 amount: u64,
923 temporary_oids: &HashSet<u32>,
924 ) -> Result<Vec<u32>, CatalogError> {
925 struct UserOid(u32);
928
929 impl UserOid {
930 fn new(oid: u32) -> Result<UserOid, anyhow::Error> {
931 if oid < FIRST_USER_OID {
932 Err(anyhow!("invalid user OID {oid}"))
933 } else {
934 Ok(UserOid(oid))
935 }
936 }
937 }
938
939 impl std::ops::AddAssign<u32> for UserOid {
940 fn add_assign(&mut self, rhs: u32) {
941 let (res, overflow) = self.0.overflowing_add(rhs);
942 self.0 = if overflow { FIRST_USER_OID + res } else { res };
943 }
944 }
945
946 if amount > u32::MAX.into() {
947 return Err(CatalogError::Catalog(SqlCatalogError::OidExhaustion));
948 }
949
950 let mut allocated_oids = HashSet::with_capacity(
956 self.databases.len()
957 + self.schemas.len()
958 + self.roles.len()
959 + self.items.len()
960 + self.introspection_sources.len()
961 + temporary_oids.len(),
962 );
963 self.databases.for_values(|_, value| {
964 allocated_oids.insert(value.oid);
965 });
966 self.schemas.for_values(|_, value| {
967 allocated_oids.insert(value.oid);
968 });
969 self.roles.for_values(|_, value| {
970 allocated_oids.insert(value.oid);
971 });
972 self.items.for_values(|_, value| {
973 allocated_oids.insert(value.oid);
974 });
975 self.introspection_sources.for_values(|_, value| {
976 allocated_oids.insert(value.oid);
977 });
978
979 let is_allocated = |oid| allocated_oids.contains(&oid) || temporary_oids.contains(&oid);
980
981 let start_oid: u32 = self
982 .id_allocator
983 .items()
984 .get(&IdAllocKey {
985 name: OID_ALLOC_KEY.to_string(),
986 })
987 .unwrap_or_else(|| panic!("{OID_ALLOC_KEY} id allocator missing"))
988 .next_id
989 .try_into()
990 .expect("we should never persist an oid outside of the u32 range");
991 let mut current_oid = UserOid::new(start_oid)
992 .expect("we should never persist an oid outside of user OID range");
993 let mut oids = Vec::new();
994 while oids.len() < u64_to_usize(amount) {
995 if !is_allocated(current_oid.0) {
996 oids.push(current_oid.0);
997 }
998 current_oid += 1;
999
1000 if current_oid.0 == start_oid && oids.len() < u64_to_usize(amount) {
1001 return Err(CatalogError::Catalog(SqlCatalogError::OidExhaustion));
1003 }
1004 }
1005
1006 let next_id = current_oid.0;
1007 let prev = self.id_allocator.set(
1008 IdAllocKey {
1009 name: OID_ALLOC_KEY.to_string(),
1010 },
1011 Some(IdAllocValue {
1012 next_id: next_id.into(),
1013 }),
1014 self.op_id,
1015 )?;
1016 assert_eq!(
1017 prev,
1018 Some(IdAllocValue {
1019 next_id: start_oid.into(),
1020 })
1021 );
1022
1023 Ok(oids)
1024 }
1025
1026 pub fn allocate_oid(&mut self, temporary_oids: &HashSet<u32>) -> Result<u32, CatalogError> {
1029 self.allocate_oids(1, temporary_oids)
1030 .map(|oids| oids.into_element())
1031 }
1032
1033 pub fn current_snapshot(&self) -> Snapshot {
1041 Snapshot {
1042 databases: self.databases.current_items_proto(),
1043 schemas: self.schemas.current_items_proto(),
1044 roles: self.roles.current_items_proto(),
1045 role_auth: self.role_auth.current_items_proto(),
1046 items: self.items.current_items_proto(),
1047 comments: self.comments.current_items_proto(),
1048 clusters: self.clusters.current_items_proto(),
1049 network_policies: self.network_policies.current_items_proto(),
1050 cluster_replicas: self.cluster_replicas.current_items_proto(),
1051 introspection_sources: self.introspection_sources.current_items_proto(),
1052 id_allocator: self.id_allocator.current_items_proto(),
1053 configs: self.configs.current_items_proto(),
1054 settings: self.settings.current_items_proto(),
1055 system_object_mappings: self.system_gid_mapping.current_items_proto(),
1056 system_configurations: self.system_configurations.current_items_proto(),
1057 cluster_system_configurations: self.cluster_system_configurations.current_items_proto(),
1058 replica_system_configurations: self.replica_system_configurations.current_items_proto(),
1059 default_privileges: self.default_privileges.current_items_proto(),
1060 source_references: self.source_references.current_items_proto(),
1061 system_privileges: self.system_privileges.current_items_proto(),
1062 storage_collection_metadata: self.storage_collection_metadata.current_items_proto(),
1063 unfinalized_shards: self.unfinalized_shards.current_items_proto(),
1064 txn_wal_shard: self.txn_wal_shard.current_items_proto(),
1065 }
1066 }
1067
1068 pub(crate) fn insert_id_allocator(
1069 &mut self,
1070 name: String,
1071 next_id: u64,
1072 ) -> Result<(), CatalogError> {
1073 match self.id_allocator.insert(
1074 IdAllocKey { name: name.clone() },
1075 IdAllocValue { next_id },
1076 self.op_id,
1077 ) {
1078 Ok(_) => Ok(()),
1079 Err(_) => Err(SqlCatalogError::IdAllocatorAlreadyExists(name).into()),
1080 }
1081 }
1082
1083 pub fn remove_database(&mut self, id: &DatabaseId) -> Result<(), CatalogError> {
1090 let prev = self
1091 .databases
1092 .set(DatabaseKey { id: *id }, None, self.op_id)?;
1093 if prev.is_some() {
1094 Ok(())
1095 } else {
1096 Err(SqlCatalogError::UnknownDatabase(id.to_string()).into())
1097 }
1098 }
1099
1100 pub fn remove_databases(
1107 &mut self,
1108 databases: &BTreeSet<DatabaseId>,
1109 ) -> Result<(), CatalogError> {
1110 if databases.is_empty() {
1111 return Ok(());
1112 }
1113
1114 let to_remove = databases
1115 .iter()
1116 .map(|id| (DatabaseKey { id: *id }, None))
1117 .collect();
1118 let mut prev = self.databases.set_many(to_remove, self.op_id)?;
1119 prev.retain(|_k, val| val.is_none());
1120
1121 if !prev.is_empty() {
1122 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1123 return Err(SqlCatalogError::UnknownDatabase(err).into());
1124 }
1125
1126 Ok(())
1127 }
1128
1129 pub fn remove_schema(
1136 &mut self,
1137 database_id: &Option<DatabaseId>,
1138 schema_id: &SchemaId,
1139 ) -> Result<(), CatalogError> {
1140 let prev = self
1141 .schemas
1142 .set(SchemaKey { id: *schema_id }, None, self.op_id)?;
1143 if prev.is_some() {
1144 Ok(())
1145 } else {
1146 let database_name = match database_id {
1147 Some(id) => format!("{id}."),
1148 None => "".to_string(),
1149 };
1150 Err(SqlCatalogError::UnknownSchema(format!("{}.{}", database_name, schema_id)).into())
1151 }
1152 }
1153
1154 pub fn remove_schemas(
1161 &mut self,
1162 schemas: &BTreeMap<SchemaId, ResolvedDatabaseSpecifier>,
1163 ) -> Result<(), CatalogError> {
1164 if schemas.is_empty() {
1165 return Ok(());
1166 }
1167
1168 let to_remove = schemas
1169 .iter()
1170 .map(|(schema_id, _)| (SchemaKey { id: *schema_id }, None))
1171 .collect();
1172 let mut prev = self.schemas.set_many(to_remove, self.op_id)?;
1173 prev.retain(|_k, v| v.is_none());
1174
1175 if !prev.is_empty() {
1176 let err = prev
1177 .keys()
1178 .map(|k| {
1179 let db_spec = schemas.get(&k.id).expect("should_exist");
1180 let db_name = match db_spec {
1181 ResolvedDatabaseSpecifier::Id(id) => format!("{id}."),
1182 ResolvedDatabaseSpecifier::Ambient => "".to_string(),
1183 };
1184 format!("{}.{}", db_name, k.id)
1185 })
1186 .join(", ");
1187
1188 return Err(SqlCatalogError::UnknownSchema(err).into());
1189 }
1190
1191 Ok(())
1192 }
1193
1194 pub fn remove_source_references(
1195 &mut self,
1196 source_id: CatalogItemId,
1197 ) -> Result<(), CatalogError> {
1198 let deleted = self
1199 .source_references
1200 .delete_by_key(SourceReferencesKey { source_id }, self.op_id)
1201 .is_some();
1202 if deleted {
1203 Ok(())
1204 } else {
1205 Err(SqlCatalogError::UnknownItem(source_id.to_string()).into())
1206 }
1207 }
1208
1209 pub fn remove_user_roles(&mut self, roles: &BTreeSet<RoleId>) -> Result<(), CatalogError> {
1216 assert!(
1217 roles.iter().all(|id| id.is_user()),
1218 "cannot delete non-user roles"
1219 );
1220 self.remove_roles(roles)
1221 }
1222
1223 pub fn remove_roles(&mut self, roles: &BTreeSet<RoleId>) -> Result<(), CatalogError> {
1230 if roles.is_empty() {
1231 return Ok(());
1232 }
1233
1234 let to_remove_keys = roles
1235 .iter()
1236 .map(|role_id| RoleKey { id: *role_id })
1237 .collect::<Vec<_>>();
1238
1239 let to_remove_roles = to_remove_keys
1240 .iter()
1241 .map(|role_key| (role_key.clone(), None))
1242 .collect();
1243
1244 let mut prev = self.roles.set_many(to_remove_roles, self.op_id)?;
1245
1246 let to_remove_role_auth = to_remove_keys
1247 .iter()
1248 .map(|role_key| {
1249 (
1250 RoleAuthKey {
1251 role_id: role_key.id,
1252 },
1253 None,
1254 )
1255 })
1256 .collect();
1257
1258 let mut role_auth_prev = self.role_auth.set_many(to_remove_role_auth, self.op_id)?;
1259
1260 prev.retain(|_k, v| v.is_none());
1261 if !prev.is_empty() {
1262 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1263 return Err(SqlCatalogError::UnknownRole(err).into());
1264 }
1265
1266 role_auth_prev.retain(|_k, v| v.is_none());
1267 Ok(())
1271 }
1272
1273 pub fn remove_clusters(&mut self, clusters: &BTreeSet<ClusterId>) -> Result<(), CatalogError> {
1280 if clusters.is_empty() {
1281 return Ok(());
1282 }
1283
1284 let to_remove = clusters
1285 .iter()
1286 .map(|cluster_id| (ClusterKey { id: *cluster_id }, None))
1287 .collect();
1288 let mut prev = self.clusters.set_many(to_remove, self.op_id)?;
1289
1290 prev.retain(|_k, v| v.is_none());
1291 if !prev.is_empty() {
1292 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1293 return Err(SqlCatalogError::UnknownCluster(err).into());
1294 }
1295
1296 self.cluster_replicas
1302 .delete(|_k, v| clusters.contains(&v.cluster_id), self.op_id);
1303 self.introspection_sources
1304 .delete(|k, _v| clusters.contains(&k.cluster_id), self.op_id);
1305
1306 Ok(())
1307 }
1308
1309 pub fn remove_cluster_replica(&mut self, id: ReplicaId) -> Result<(), CatalogError> {
1316 let deleted = self
1317 .cluster_replicas
1318 .delete_by_key(ClusterReplicaKey { id }, self.op_id)
1319 .is_some();
1320 if deleted {
1321 Ok(())
1322 } else {
1323 Err(SqlCatalogError::UnknownClusterReplica(id.to_string()).into())
1324 }
1325 }
1326
1327 pub fn remove_cluster_replicas(
1334 &mut self,
1335 replicas: &BTreeSet<ReplicaId>,
1336 ) -> Result<(), CatalogError> {
1337 if replicas.is_empty() {
1338 return Ok(());
1339 }
1340
1341 let to_remove = replicas
1342 .iter()
1343 .map(|replica_id| (ClusterReplicaKey { id: *replica_id }, None))
1344 .collect();
1345 let mut prev = self.cluster_replicas.set_many(to_remove, self.op_id)?;
1346
1347 prev.retain(|_k, v| v.is_none());
1348 if !prev.is_empty() {
1349 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1350 return Err(SqlCatalogError::UnknownClusterReplica(err).into());
1351 }
1352
1353 Ok(())
1354 }
1355
1356 pub fn remove_item(&mut self, id: CatalogItemId) -> Result<(), CatalogError> {
1363 let prev = self.items.set(ItemKey { id }, None, self.op_id)?;
1364 if prev.is_some() {
1365 Ok(())
1366 } else {
1367 Err(SqlCatalogError::UnknownItem(id.to_string()).into())
1368 }
1369 }
1370
1371 pub fn remove_items(&mut self, ids: &BTreeSet<CatalogItemId>) -> Result<(), CatalogError> {
1378 if ids.is_empty() {
1379 return Ok(());
1380 }
1381
1382 let ks: Vec<_> = ids.clone().into_iter().map(|id| ItemKey { id }).collect();
1383 let n = self.items.delete_by_keys(ks, self.op_id).len();
1384 if n == ids.len() {
1385 Ok(())
1386 } else {
1387 let item_ids = self.items.items().keys().map(|k| k.id).collect();
1388 let mut unknown = ids.difference(&item_ids);
1389 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1390 }
1391 }
1392
1393 pub fn remove_system_object_mappings(
1400 &mut self,
1401 descriptions: BTreeSet<SystemObjectDescription>,
1402 ) -> Result<(), CatalogError> {
1403 if descriptions.is_empty() {
1404 return Ok(());
1405 }
1406
1407 let ks: Vec<_> = descriptions
1408 .clone()
1409 .into_iter()
1410 .map(|desc| GidMappingKey {
1411 schema_name: desc.schema_name,
1412 object_type: desc.object_type,
1413 object_name: desc.object_name,
1414 })
1415 .collect();
1416 let n = self.system_gid_mapping.delete_by_keys(ks, self.op_id).len();
1417
1418 if n == descriptions.len() {
1419 Ok(())
1420 } else {
1421 let item_descriptions = self
1422 .system_gid_mapping
1423 .items()
1424 .keys()
1425 .map(|k| SystemObjectDescription {
1426 schema_name: k.schema_name.clone(),
1427 object_type: k.object_type.clone(),
1428 object_name: k.object_name.clone(),
1429 })
1430 .collect();
1431 let mut unknown = descriptions.difference(&item_descriptions).map(|desc| {
1432 format!(
1433 "{} {}.{}",
1434 desc.object_type, desc.schema_name, desc.object_name
1435 )
1436 });
1437 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1438 }
1439 }
1440
1441 pub fn remove_introspection_source_indexes(
1448 &mut self,
1449 introspection_source_indexes: BTreeSet<(ClusterId, String)>,
1450 ) -> Result<(), CatalogError> {
1451 if introspection_source_indexes.is_empty() {
1452 return Ok(());
1453 }
1454
1455 let ks: Vec<_> = introspection_source_indexes
1456 .clone()
1457 .into_iter()
1458 .map(|(cluster_id, name)| ClusterIntrospectionSourceIndexKey { cluster_id, name })
1459 .collect();
1460 let n = self
1461 .introspection_sources
1462 .delete_by_keys(ks, self.op_id)
1463 .len();
1464 if n == introspection_source_indexes.len() {
1465 Ok(())
1466 } else {
1467 let txn_indexes = self
1468 .introspection_sources
1469 .items()
1470 .keys()
1471 .map(|k| (k.cluster_id, k.name.clone()))
1472 .collect();
1473 let mut unknown = introspection_source_indexes
1474 .difference(&txn_indexes)
1475 .map(|(cluster_id, name)| format!("{cluster_id} {name}"));
1476 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1477 }
1478 }
1479
1480 pub fn update_item(&mut self, id: CatalogItemId, item: Item) -> Result<(), CatalogError> {
1487 let updated =
1488 self.items
1489 .update_by_key(ItemKey { id }, item.into_key_value().1, self.op_id)?;
1490 if updated {
1491 Ok(())
1492 } else {
1493 Err(SqlCatalogError::UnknownItem(id.to_string()).into())
1494 }
1495 }
1496
1497 pub fn update_items(
1505 &mut self,
1506 items: BTreeMap<CatalogItemId, Item>,
1507 ) -> Result<(), CatalogError> {
1508 if items.is_empty() {
1509 return Ok(());
1510 }
1511
1512 let update_ids: BTreeSet<_> = items.keys().cloned().collect();
1513 let kvs: Vec<_> = items
1514 .clone()
1515 .into_iter()
1516 .map(|(id, item)| (ItemKey { id }, item.into_key_value().1))
1517 .collect();
1518 let n = self.items.update_by_keys(kvs, self.op_id)?;
1519 let n = usize::try_from(n.into_inner()).expect("Must be positive and fit in usize");
1520 if n == update_ids.len() {
1521 Ok(())
1522 } else {
1523 let item_ids: BTreeSet<_> = self.items.items().keys().map(|k| k.id).collect();
1524 let mut unknown = update_ids.difference(&item_ids);
1525 Err(SqlCatalogError::UnknownItem(unknown.join(", ")).into())
1526 }
1527 }
1528
1529 pub fn update_role(
1537 &mut self,
1538 id: RoleId,
1539 role: Role,
1540 password: PasswordAction,
1541 ) -> Result<(), CatalogError> {
1542 let key = RoleKey { id };
1543 if self.roles.get(&key).is_some() {
1544 let auth_key = RoleAuthKey { role_id: id };
1545
1546 match password {
1547 PasswordAction::Set(new_password) => {
1548 let hash = mz_auth::hash::scram256_hash(
1549 &new_password.password,
1550 &new_password.scram_iterations,
1551 )
1552 .expect("password hash should be valid");
1553 let value = RoleAuthValue {
1554 password_hash: Some(hash),
1555 updated_at: SYSTEM_TIME(),
1556 };
1557
1558 if self.role_auth.get(&auth_key).is_some() {
1559 self.role_auth
1560 .update_by_key(auth_key.clone(), value, self.op_id)?;
1561 } else {
1562 self.role_auth.insert(auth_key.clone(), value, self.op_id)?;
1563 }
1564 }
1565 PasswordAction::Clear => {
1566 let value = RoleAuthValue {
1567 password_hash: None,
1568 updated_at: SYSTEM_TIME(),
1569 };
1570 if self.role_auth.get(&auth_key).is_some() {
1571 self.role_auth
1572 .update_by_key(auth_key.clone(), value, self.op_id)?;
1573 }
1574 }
1575 PasswordAction::NoChange => {}
1576 }
1577
1578 self.roles
1579 .update_by_key(key, role.into_key_value().1, self.op_id)?;
1580
1581 Ok(())
1582 } else {
1583 Err(SqlCatalogError::UnknownRole(id.to_string()).into())
1584 }
1585 }
1586
1587 pub fn update_roles_without_auth(
1598 &mut self,
1599 roles: BTreeMap<RoleId, Role>,
1600 ) -> Result<(), CatalogError> {
1601 if roles.is_empty() {
1602 return Ok(());
1603 }
1604
1605 let update_role_ids: BTreeSet<_> = roles.keys().cloned().collect();
1606 let kvs: Vec<_> = roles
1607 .into_iter()
1608 .map(|(id, role)| (RoleKey { id }, role.into_key_value().1))
1609 .collect();
1610 let n = self.roles.update_by_keys(kvs, self.op_id)?;
1611 let n = usize::try_from(n.into_inner()).expect("Must be positive and fit in usize");
1612
1613 if n == update_role_ids.len() {
1614 Ok(())
1615 } else {
1616 let role_ids: BTreeSet<_> = self.roles.items().keys().map(|k| k.id).collect();
1617 let mut unknown = update_role_ids.difference(&role_ids);
1618 Err(SqlCatalogError::UnknownRole(unknown.join(", ")).into())
1619 }
1620 }
1621
1622 pub fn update_system_object_mappings(
1627 &mut self,
1628 mappings: BTreeMap<CatalogItemId, SystemObjectMapping>,
1629 ) -> Result<(), CatalogError> {
1630 if mappings.is_empty() {
1631 return Ok(());
1632 }
1633
1634 let n = self.system_gid_mapping.update(
1635 |_k, v| {
1636 if let Some(mapping) = mappings.get(&CatalogItemId::from(v.catalog_id)) {
1637 let (_, new_value) = mapping.clone().into_key_value();
1638 Some(new_value)
1639 } else {
1640 None
1641 }
1642 },
1643 self.op_id,
1644 )?;
1645
1646 if usize::try_from(n.into_inner()).expect("update diff should fit into usize")
1647 != mappings.len()
1648 {
1649 let id_str = mappings.keys().map(|id| id.to_string()).join(",");
1650 return Err(SqlCatalogError::FailedBuiltinSchemaMigration(id_str).into());
1651 }
1652
1653 Ok(())
1654 }
1655
1656 pub fn update_cluster(&mut self, id: ClusterId, cluster: Cluster) -> Result<(), CatalogError> {
1663 let updated = self.clusters.update_by_key(
1664 ClusterKey { id },
1665 cluster.into_key_value().1,
1666 self.op_id,
1667 )?;
1668 if updated {
1669 Ok(())
1670 } else {
1671 Err(SqlCatalogError::UnknownCluster(id.to_string()).into())
1672 }
1673 }
1674
1675 pub fn update_cluster_replica(
1682 &mut self,
1683 replica_id: ReplicaId,
1684 replica: ClusterReplica,
1685 ) -> Result<(), CatalogError> {
1686 let updated = self.cluster_replicas.update_by_key(
1687 ClusterReplicaKey { id: replica_id },
1688 replica.into_key_value().1,
1689 self.op_id,
1690 )?;
1691 if updated {
1692 Ok(())
1693 } else {
1694 Err(SqlCatalogError::UnknownClusterReplica(replica_id.to_string()).into())
1695 }
1696 }
1697
1698 pub fn update_database(
1705 &mut self,
1706 id: DatabaseId,
1707 database: Database,
1708 ) -> Result<(), CatalogError> {
1709 let updated = self.databases.update_by_key(
1710 DatabaseKey { id },
1711 database.into_key_value().1,
1712 self.op_id,
1713 )?;
1714 if updated {
1715 Ok(())
1716 } else {
1717 Err(SqlCatalogError::UnknownDatabase(id.to_string()).into())
1718 }
1719 }
1720
1721 pub fn update_schema(
1728 &mut self,
1729 schema_id: SchemaId,
1730 schema: Schema,
1731 ) -> Result<(), CatalogError> {
1732 let updated = self.schemas.update_by_key(
1733 SchemaKey { id: schema_id },
1734 schema.into_key_value().1,
1735 self.op_id,
1736 )?;
1737 if updated {
1738 Ok(())
1739 } else {
1740 Err(SqlCatalogError::UnknownSchema(schema_id.to_string()).into())
1741 }
1742 }
1743
1744 pub fn update_network_policy(
1751 &mut self,
1752 id: NetworkPolicyId,
1753 network_policy: NetworkPolicy,
1754 ) -> Result<(), CatalogError> {
1755 let updated = self.network_policies.update_by_key(
1756 NetworkPolicyKey { id },
1757 network_policy.into_key_value().1,
1758 self.op_id,
1759 )?;
1760 if updated {
1761 Ok(())
1762 } else {
1763 Err(SqlCatalogError::UnknownNetworkPolicy(id.to_string()).into())
1764 }
1765 }
1766 pub fn remove_network_policies(
1773 &mut self,
1774 network_policies: &BTreeSet<NetworkPolicyId>,
1775 ) -> Result<(), CatalogError> {
1776 if network_policies.is_empty() {
1777 return Ok(());
1778 }
1779
1780 let to_remove = network_policies
1781 .iter()
1782 .map(|policy_id| (NetworkPolicyKey { id: *policy_id }, None))
1783 .collect();
1784 let mut prev = self.network_policies.set_many(to_remove, self.op_id)?;
1785 assert!(
1786 prev.iter().all(|(k, _)| k.id.is_user()),
1787 "cannot delete non-user network policy"
1788 );
1789
1790 prev.retain(|_k, v| v.is_none());
1791 if !prev.is_empty() {
1792 let err = prev.keys().map(|k| k.id.to_string()).join(", ");
1793 return Err(SqlCatalogError::UnknownNetworkPolicy(err).into());
1794 }
1795
1796 Ok(())
1797 }
1798 pub fn set_default_privilege(
1802 &mut self,
1803 role_id: RoleId,
1804 database_id: Option<DatabaseId>,
1805 schema_id: Option<SchemaId>,
1806 object_type: ObjectType,
1807 grantee: RoleId,
1808 privileges: Option<AclMode>,
1809 ) -> Result<(), CatalogError> {
1810 self.default_privileges.set(
1811 DefaultPrivilegesKey {
1812 role_id,
1813 database_id,
1814 schema_id,
1815 object_type,
1816 grantee,
1817 },
1818 privileges.map(|privileges| DefaultPrivilegesValue { privileges }),
1819 self.op_id,
1820 )?;
1821 Ok(())
1822 }
1823
1824 pub fn set_default_privileges(
1826 &mut self,
1827 default_privileges: Vec<DefaultPrivilege>,
1828 ) -> Result<(), CatalogError> {
1829 if default_privileges.is_empty() {
1830 return Ok(());
1831 }
1832
1833 let default_privileges = default_privileges
1834 .into_iter()
1835 .map(DurableType::into_key_value)
1836 .map(|(k, v)| (k, Some(v)))
1837 .collect();
1838 self.default_privileges
1839 .set_many(default_privileges, self.op_id)?;
1840 Ok(())
1841 }
1842
1843 pub fn set_system_privilege(
1847 &mut self,
1848 grantee: RoleId,
1849 grantor: RoleId,
1850 acl_mode: Option<AclMode>,
1851 ) -> Result<(), CatalogError> {
1852 self.system_privileges.set(
1853 SystemPrivilegesKey { grantee, grantor },
1854 acl_mode.map(|acl_mode| SystemPrivilegesValue { acl_mode }),
1855 self.op_id,
1856 )?;
1857 Ok(())
1858 }
1859
1860 pub fn set_system_privileges(
1862 &mut self,
1863 system_privileges: Vec<MzAclItem>,
1864 ) -> Result<(), CatalogError> {
1865 if system_privileges.is_empty() {
1866 return Ok(());
1867 }
1868
1869 let system_privileges = system_privileges
1870 .into_iter()
1871 .map(DurableType::into_key_value)
1872 .map(|(k, v)| (k, Some(v)))
1873 .collect();
1874 self.system_privileges
1875 .set_many(system_privileges, self.op_id)?;
1876 Ok(())
1877 }
1878
1879 pub fn set_setting(&mut self, name: String, value: Option<String>) -> Result<(), CatalogError> {
1881 self.settings.set(
1882 SettingKey { name },
1883 value.map(|value| SettingValue { value }),
1884 self.op_id,
1885 )?;
1886 Ok(())
1887 }
1888
1889 pub fn set_catalog_content_version(&mut self, version: String) -> Result<(), CatalogError> {
1890 self.set_setting(CATALOG_CONTENT_VERSION_KEY.to_string(), Some(version))
1891 }
1892
1893 pub fn insert_introspection_source_indexes(
1895 &mut self,
1896 introspection_source_indexes: Vec<(ClusterId, String, CatalogItemId, GlobalId)>,
1897 temporary_oids: &HashSet<u32>,
1898 ) -> Result<(), CatalogError> {
1899 if introspection_source_indexes.is_empty() {
1900 return Ok(());
1901 }
1902
1903 let amount = usize_to_u64(introspection_source_indexes.len());
1904 let oids = self.allocate_oids(amount, temporary_oids)?;
1905 let introspection_source_indexes: Vec<_> = introspection_source_indexes
1906 .into_iter()
1907 .zip_eq(oids)
1908 .map(
1909 |((cluster_id, name, item_id, index_id), oid)| IntrospectionSourceIndex {
1910 cluster_id,
1911 name,
1912 item_id,
1913 index_id,
1914 oid,
1915 },
1916 )
1917 .collect();
1918
1919 for introspection_source_index in introspection_source_indexes {
1920 let (key, value) = introspection_source_index.into_key_value();
1921 self.introspection_sources.insert(key, value, self.op_id)?;
1922 }
1923
1924 Ok(())
1925 }
1926
1927 pub fn set_system_object_mappings(
1929 &mut self,
1930 mappings: Vec<SystemObjectMapping>,
1931 ) -> Result<(), CatalogError> {
1932 if mappings.is_empty() {
1933 return Ok(());
1934 }
1935
1936 let mappings = mappings
1937 .into_iter()
1938 .map(DurableType::into_key_value)
1939 .map(|(k, v)| (k, Some(v)))
1940 .collect();
1941 self.system_gid_mapping.set_many(mappings, self.op_id)?;
1942 Ok(())
1943 }
1944
1945 pub fn set_replicas(&mut self, replicas: Vec<ClusterReplica>) -> Result<(), CatalogError> {
1947 if replicas.is_empty() {
1948 return Ok(());
1949 }
1950
1951 let replicas = replicas
1952 .into_iter()
1953 .map(DurableType::into_key_value)
1954 .map(|(k, v)| (k, Some(v)))
1955 .collect();
1956 self.cluster_replicas.set_many(replicas, self.op_id)?;
1957 Ok(())
1958 }
1959
1960 pub fn set_config(&mut self, key: String, value: Option<u64>) -> Result<(), CatalogError> {
1962 match value {
1963 Some(value) => {
1964 let config = Config { key, value };
1965 let (key, value) = config.into_key_value();
1966 self.configs.set(key, Some(value), self.op_id)?;
1967 }
1968 None => {
1969 self.configs.set(ConfigKey { key }, None, self.op_id)?;
1970 }
1971 }
1972 Ok(())
1973 }
1974
1975 pub fn get_config(&self, key: String) -> Option<u64> {
1977 self.configs
1978 .get(&ConfigKey { key })
1979 .map(|entry| entry.value)
1980 }
1981
1982 pub fn get_setting(&self, name: String) -> Option<&str> {
1984 self.settings
1985 .get(&SettingKey { name })
1986 .map(|entry| &*entry.value)
1987 }
1988
1989 pub fn get_builtin_migration_shard(&self) -> Option<ShardId> {
1990 self.get_setting(BUILTIN_MIGRATION_SHARD_KEY.to_string())
1991 .map(|shard_id| shard_id.parse().expect("valid ShardId"))
1992 }
1993
1994 pub fn set_builtin_migration_shard(&mut self, shard_id: ShardId) -> Result<(), CatalogError> {
1995 self.set_setting(
1996 BUILTIN_MIGRATION_SHARD_KEY.to_string(),
1997 Some(shard_id.to_string()),
1998 )
1999 }
2000
2001 pub fn get_expression_cache_shard(&self) -> Option<ShardId> {
2002 self.get_setting(EXPRESSION_CACHE_SHARD_KEY.to_string())
2003 .map(|shard_id| shard_id.parse().expect("valid ShardId"))
2004 }
2005
2006 pub fn set_expression_cache_shard(&mut self, shard_id: ShardId) -> Result<(), CatalogError> {
2007 self.set_setting(
2008 EXPRESSION_CACHE_SHARD_KEY.to_string(),
2009 Some(shard_id.to_string()),
2010 )
2011 }
2012
2013 pub fn set_0dt_deployment_max_wait(&mut self, value: Duration) -> Result<(), CatalogError> {
2019 self.set_config(
2020 WITH_0DT_DEPLOYMENT_MAX_WAIT.into(),
2021 Some(
2022 value
2023 .as_millis()
2024 .try_into()
2025 .expect("max wait fits into u64"),
2026 ),
2027 )
2028 }
2029
2030 pub fn set_0dt_deployment_ddl_check_interval(
2037 &mut self,
2038 value: Duration,
2039 ) -> Result<(), CatalogError> {
2040 self.set_config(
2041 WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL.into(),
2042 Some(
2043 value
2044 .as_millis()
2045 .try_into()
2046 .expect("ddl check interval fits into u64"),
2047 ),
2048 )
2049 }
2050
2051 pub fn set_enable_0dt_deployment_panic_after_timeout(
2057 &mut self,
2058 value: bool,
2059 ) -> Result<(), CatalogError> {
2060 self.set_config(
2061 ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT.into(),
2062 Some(u64::from(value)),
2063 )
2064 }
2065
2066 pub fn reset_0dt_deployment_max_wait(&mut self) -> Result<(), CatalogError> {
2072 self.set_config(WITH_0DT_DEPLOYMENT_MAX_WAIT.into(), None)
2073 }
2074
2075 pub fn reset_0dt_deployment_ddl_check_interval(&mut self) -> Result<(), CatalogError> {
2082 self.set_config(WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL.into(), None)
2083 }
2084
2085 pub fn reset_enable_0dt_deployment_panic_after_timeout(&mut self) -> Result<(), CatalogError> {
2092 self.set_config(ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT.into(), None)
2093 }
2094
2095 pub fn set_system_config_synced_once(&mut self) -> Result<(), CatalogError> {
2097 self.set_config(SYSTEM_CONFIG_SYNCED_KEY.into(), Some(1))
2098 }
2099
2100 pub fn update_comment(
2101 &mut self,
2102 object_id: CommentObjectId,
2103 sub_component: Option<usize>,
2104 comment: Option<String>,
2105 ) -> Result<(), CatalogError> {
2106 let key = CommentKey {
2107 object_id,
2108 sub_component,
2109 };
2110 let value = comment.map(|c| CommentValue { comment: c });
2111 self.comments.set(key, value, self.op_id)?;
2112
2113 Ok(())
2114 }
2115
2116 pub fn drop_comments(
2117 &mut self,
2118 object_ids: &BTreeSet<CommentObjectId>,
2119 ) -> Result<(), CatalogError> {
2120 if object_ids.is_empty() {
2121 return Ok(());
2122 }
2123
2124 self.comments
2125 .delete(|k, _v| object_ids.contains(&k.object_id), self.op_id);
2126 Ok(())
2127 }
2128
2129 pub fn update_source_references(
2130 &mut self,
2131 source_id: CatalogItemId,
2132 references: Vec<SourceReference>,
2133 updated_at: u64,
2134 ) -> Result<(), CatalogError> {
2135 let key = SourceReferencesKey { source_id };
2136 let value = SourceReferencesValue {
2137 references,
2138 updated_at,
2139 };
2140 self.source_references.set(key, Some(value), self.op_id)?;
2141 Ok(())
2142 }
2143
2144 pub fn upsert_system_config(&mut self, name: &str, value: String) -> Result<(), CatalogError> {
2146 let key = ServerConfigurationKey {
2147 name: name.to_string(),
2148 };
2149 let value = ServerConfigurationValue { value };
2150 self.system_configurations
2151 .set(key, Some(value), self.op_id)?;
2152 Ok(())
2153 }
2154
2155 pub fn remove_system_config(&mut self, name: &str) {
2157 let key = ServerConfigurationKey {
2158 name: name.to_string(),
2159 };
2160 self.system_configurations
2161 .set(key, None, self.op_id)
2162 .expect("cannot have uniqueness violation");
2163 }
2164
2165 pub fn clear_system_configs(&mut self) {
2167 self.system_configurations.delete(|_k, _v| true, self.op_id);
2168 }
2169
2170 pub fn get_cluster_system_configurations(
2172 &self,
2173 ) -> impl Iterator<Item = ClusterSystemConfiguration> + use<'_> {
2174 self.cluster_system_configurations
2175 .items()
2176 .into_iter()
2177 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2178 }
2179
2180 pub fn upsert_cluster_system_config(
2183 &mut self,
2184 cluster_id: ClusterId,
2185 name: &str,
2186 value: String,
2187 ) -> Result<(), CatalogError> {
2188 let key = ClusterSystemConfigurationKey {
2189 cluster_id,
2190 name: name.to_string(),
2191 };
2192 let value = ClusterSystemConfigurationValue { value };
2193 self.cluster_system_configurations
2194 .set(key, Some(value), self.op_id)?;
2195 Ok(())
2196 }
2197
2198 pub fn remove_cluster_system_config(&mut self, cluster_id: ClusterId, name: &str) {
2201 let key = ClusterSystemConfigurationKey {
2202 cluster_id,
2203 name: name.to_string(),
2204 };
2205 self.cluster_system_configurations
2206 .set(key, None, self.op_id)
2207 .expect("cannot have uniqueness violation");
2208 }
2209
2210 pub fn get_replica_system_configurations(
2212 &self,
2213 ) -> impl Iterator<Item = ReplicaSystemConfiguration> + use<'_> {
2214 self.replica_system_configurations
2215 .items()
2216 .into_iter()
2217 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2218 }
2219
2220 pub fn upsert_replica_system_config(
2223 &mut self,
2224 replica_id: ReplicaId,
2225 name: &str,
2226 value: String,
2227 ) -> Result<(), CatalogError> {
2228 let key = ReplicaSystemConfigurationKey {
2229 replica_id,
2230 name: name.to_string(),
2231 };
2232 let value = ReplicaSystemConfigurationValue { value };
2233 self.replica_system_configurations
2234 .set(key, Some(value), self.op_id)?;
2235 Ok(())
2236 }
2237
2238 pub fn remove_replica_system_config(&mut self, replica_id: ReplicaId, name: &str) {
2241 let key = ReplicaSystemConfigurationKey {
2242 replica_id,
2243 name: name.to_string(),
2244 };
2245 self.replica_system_configurations
2246 .set(key, None, self.op_id)
2247 .expect("cannot have uniqueness violation");
2248 }
2249
2250 pub(crate) fn insert_config(&mut self, key: String, value: u64) -> Result<(), CatalogError> {
2251 match self.configs.insert(
2252 ConfigKey { key: key.clone() },
2253 ConfigValue { value },
2254 self.op_id,
2255 ) {
2256 Ok(_) => Ok(()),
2257 Err(_) => Err(SqlCatalogError::ConfigAlreadyExists(key).into()),
2258 }
2259 }
2260
2261 pub fn get_clusters(&self) -> impl Iterator<Item = Cluster> + use<'_> {
2262 self.clusters
2263 .items()
2264 .into_iter()
2265 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2266 }
2267
2268 pub fn get_cluster_replicas(&self) -> impl Iterator<Item = ClusterReplica> + use<'_> {
2269 self.cluster_replicas
2270 .items()
2271 .into_iter()
2272 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2273 }
2274
2275 pub fn get_databases(&self) -> impl Iterator<Item = Database> + use<'_> {
2276 self.databases
2277 .items()
2278 .into_iter()
2279 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2280 }
2281
2282 pub fn get_roles(&self) -> impl Iterator<Item = Role> + use<'_> {
2283 self.roles
2284 .items()
2285 .into_iter()
2286 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2287 }
2288
2289 pub fn get_network_policies(&self) -> impl Iterator<Item = NetworkPolicy> + use<'_> {
2290 self.network_policies
2291 .items()
2292 .into_iter()
2293 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2294 }
2295
2296 pub fn get_system_object_mappings(
2297 &self,
2298 ) -> impl Iterator<Item = SystemObjectMapping> + use<'_> {
2299 self.system_gid_mapping
2300 .items()
2301 .into_iter()
2302 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2303 }
2304
2305 pub fn get_schemas(&self) -> impl Iterator<Item = Schema> + use<'_> {
2306 self.schemas
2307 .items()
2308 .into_iter()
2309 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2310 }
2311
2312 pub fn get_system_configurations(&self) -> impl Iterator<Item = SystemConfiguration> + use<'_> {
2313 self.system_configurations
2314 .items()
2315 .into_iter()
2316 .map(|(k, v)| DurableType::from_key_value(k.clone(), v.clone()))
2317 }
2318
2319 pub fn get_schema(&self, id: &SchemaId) -> Option<Schema> {
2320 let key = SchemaKey { id: *id };
2321 self.schemas
2322 .get(&key)
2323 .map(|v| DurableType::from_key_value(key, v.clone()))
2324 }
2325
2326 pub fn get_introspection_source_indexes(
2327 &self,
2328 cluster_id: ClusterId,
2329 ) -> BTreeMap<&str, (GlobalId, u32)> {
2330 self.introspection_sources
2331 .items()
2332 .into_iter()
2333 .filter(|(k, _v)| k.cluster_id == cluster_id)
2334 .map(|(k, v)| (k.name.as_str(), (v.global_id.into(), v.oid)))
2335 .collect()
2336 }
2337
2338 pub fn get_catalog_content_version(&self) -> Option<&str> {
2339 self.settings
2340 .get(&SettingKey {
2341 name: CATALOG_CONTENT_VERSION_KEY.to_string(),
2342 })
2343 .map(|value| &*value.value)
2344 }
2345
2346 pub fn get_authentication_mock_nonce(&self) -> Option<String> {
2347 self.settings
2348 .get(&SettingKey {
2349 name: MOCK_AUTHENTICATION_NONCE_KEY.to_string(),
2350 })
2351 .map(|value| value.value.clone())
2352 }
2353
2354 #[must_use]
2360 pub fn get_and_commit_op_updates(&mut self) -> Vec<StateUpdate> {
2361 let updates = self.get_op_updates();
2362 self.commit_op();
2363 updates
2364 }
2365
2366 fn get_op_updates(&self) -> Vec<StateUpdate> {
2367 fn get_collection_op_updates<'a, T>(
2368 table_txn: &'a TableTransaction<T::Key, T::Value>,
2369 kind_fn: impl Fn(T) -> StateUpdateKind + 'a,
2370 op: Timestamp,
2371 ) -> impl Iterator<Item = (StateUpdateKind, StateDiff)> + 'a
2372 where
2373 T::Key: Ord + Eq + Clone + Debug,
2374 T::Value: Ord + Clone + Debug,
2375 T: DurableType,
2376 {
2377 table_txn
2378 .pending
2379 .iter()
2380 .flat_map(|(k, vs)| vs.into_iter().map(move |v| (k, v)))
2381 .filter_map(move |(k, v)| {
2382 if v.ts == op {
2383 let key = k.clone();
2384 let value = v.value.clone();
2385 let diff = v.diff.clone().try_into().expect("invalid diff");
2386 let update = DurableType::from_key_value(key, value);
2387 let kind = kind_fn(update);
2388 Some((kind, diff))
2389 } else {
2390 None
2391 }
2392 })
2393 }
2394
2395 fn get_large_collection_op_updates<'a, T>(
2396 collection: &'a Vec<(T::Key, Diff, Timestamp)>,
2397 kind_fn: impl Fn(T) -> StateUpdateKind + 'a,
2398 op: Timestamp,
2399 ) -> impl Iterator<Item = (StateUpdateKind, StateDiff)> + 'a
2400 where
2401 T::Key: Ord + Eq + Clone + Debug,
2402 T: DurableType<Value = ()>,
2403 {
2404 collection.iter().filter_map(move |(k, diff, ts)| {
2405 if *ts == op {
2406 let key = k.clone();
2407 let diff = diff.clone().try_into().expect("invalid diff");
2408 let update = DurableType::from_key_value(key, ());
2409 let kind = kind_fn(update);
2410 Some((kind, diff))
2411 } else {
2412 None
2413 }
2414 })
2415 }
2416
2417 let Transaction {
2418 durable_catalog: _,
2419 databases,
2420 schemas,
2421 items,
2422 comments,
2423 roles,
2424 role_auth,
2425 clusters,
2426 network_policies,
2427 cluster_replicas,
2428 introspection_sources,
2429 system_gid_mapping,
2430 system_configurations,
2431 cluster_system_configurations,
2432 replica_system_configurations,
2433 default_privileges,
2434 source_references,
2435 system_privileges,
2436 audit_log_updates,
2437 storage_collection_metadata,
2438 unfinalized_shards,
2439 id_allocator: _,
2441 configs: _,
2442 settings: _,
2443 txn_wal_shard: _,
2444 upper,
2445 op_id: _,
2446 } = &self;
2447
2448 let updates = std::iter::empty()
2449 .chain(get_collection_op_updates(
2450 roles,
2451 StateUpdateKind::Role,
2452 self.op_id,
2453 ))
2454 .chain(get_collection_op_updates(
2455 role_auth,
2456 StateUpdateKind::RoleAuth,
2457 self.op_id,
2458 ))
2459 .chain(get_collection_op_updates(
2460 databases,
2461 StateUpdateKind::Database,
2462 self.op_id,
2463 ))
2464 .chain(get_collection_op_updates(
2465 schemas,
2466 StateUpdateKind::Schema,
2467 self.op_id,
2468 ))
2469 .chain(get_collection_op_updates(
2470 default_privileges,
2471 StateUpdateKind::DefaultPrivilege,
2472 self.op_id,
2473 ))
2474 .chain(get_collection_op_updates(
2475 system_privileges,
2476 StateUpdateKind::SystemPrivilege,
2477 self.op_id,
2478 ))
2479 .chain(get_collection_op_updates(
2480 system_configurations,
2481 StateUpdateKind::SystemConfiguration,
2482 self.op_id,
2483 ))
2484 .chain(get_collection_op_updates(
2485 cluster_system_configurations,
2486 StateUpdateKind::ClusterSystemConfiguration,
2487 self.op_id,
2488 ))
2489 .chain(get_collection_op_updates(
2490 replica_system_configurations,
2491 StateUpdateKind::ReplicaSystemConfiguration,
2492 self.op_id,
2493 ))
2494 .chain(get_collection_op_updates(
2495 clusters,
2496 StateUpdateKind::Cluster,
2497 self.op_id,
2498 ))
2499 .chain(get_collection_op_updates(
2500 network_policies,
2501 StateUpdateKind::NetworkPolicy,
2502 self.op_id,
2503 ))
2504 .chain(get_collection_op_updates(
2505 introspection_sources,
2506 StateUpdateKind::IntrospectionSourceIndex,
2507 self.op_id,
2508 ))
2509 .chain(get_collection_op_updates(
2510 cluster_replicas,
2511 StateUpdateKind::ClusterReplica,
2512 self.op_id,
2513 ))
2514 .chain(get_collection_op_updates(
2515 system_gid_mapping,
2516 StateUpdateKind::SystemObjectMapping,
2517 self.op_id,
2518 ))
2519 .chain(get_collection_op_updates(
2520 items,
2521 StateUpdateKind::Item,
2522 self.op_id,
2523 ))
2524 .chain(get_collection_op_updates(
2525 comments,
2526 StateUpdateKind::Comment,
2527 self.op_id,
2528 ))
2529 .chain(get_collection_op_updates(
2530 source_references,
2531 StateUpdateKind::SourceReferences,
2532 self.op_id,
2533 ))
2534 .chain(get_collection_op_updates(
2535 storage_collection_metadata,
2536 StateUpdateKind::StorageCollectionMetadata,
2537 self.op_id,
2538 ))
2539 .chain(get_collection_op_updates(
2540 unfinalized_shards,
2541 StateUpdateKind::UnfinalizedShard,
2542 self.op_id,
2543 ))
2544 .chain(get_large_collection_op_updates(
2545 audit_log_updates,
2546 StateUpdateKind::AuditLog,
2547 self.op_id,
2548 ))
2549 .map(|(kind, diff)| StateUpdate {
2550 kind,
2551 ts: upper.clone(),
2552 diff,
2553 })
2554 .collect();
2555
2556 updates
2557 }
2558
2559 pub fn is_savepoint(&self) -> bool {
2560 self.durable_catalog.is_savepoint()
2561 }
2562
2563 fn commit_op(&mut self) {
2564 self.op_id += 1;
2565 }
2566
2567 pub fn op_id(&self) -> Timestamp {
2568 self.op_id
2569 }
2570
2571 pub fn upper(&self) -> mz_repr::Timestamp {
2572 self.upper
2573 }
2574
2575 pub(crate) fn into_parts(self) -> (TransactionBatch, &'a mut dyn DurableCatalogState) {
2576 let audit_log_updates = self
2577 .audit_log_updates
2578 .into_iter()
2579 .map(|(k, diff, _op)| (k.into_proto(), (), diff))
2580 .collect();
2581
2582 let txn_batch = TransactionBatch {
2583 databases: self.databases.pending(),
2584 schemas: self.schemas.pending(),
2585 items: self.items.pending(),
2586 comments: self.comments.pending(),
2587 roles: self.roles.pending(),
2588 role_auth: self.role_auth.pending(),
2589 clusters: self.clusters.pending(),
2590 cluster_replicas: self.cluster_replicas.pending(),
2591 network_policies: self.network_policies.pending(),
2592 introspection_sources: self.introspection_sources.pending(),
2593 id_allocator: self.id_allocator.pending(),
2594 configs: self.configs.pending(),
2595 source_references: self.source_references.pending(),
2596 settings: self.settings.pending(),
2597 system_gid_mapping: self.system_gid_mapping.pending(),
2598 system_configurations: self.system_configurations.pending(),
2599 cluster_system_configurations: self.cluster_system_configurations.pending(),
2600 replica_system_configurations: self.replica_system_configurations.pending(),
2601 default_privileges: self.default_privileges.pending(),
2602 system_privileges: self.system_privileges.pending(),
2603 storage_collection_metadata: self.storage_collection_metadata.pending(),
2604 unfinalized_shards: self.unfinalized_shards.pending(),
2605 txn_wal_shard: self.txn_wal_shard.pending(),
2606 audit_log_updates,
2607 upper: self.upper,
2608 };
2609 (txn_batch, self.durable_catalog)
2610 }
2611
2612 #[mz_ore::instrument(level = "debug")]
2624 pub(crate) async fn commit_internal(
2625 self,
2626 commit_ts: mz_repr::Timestamp,
2627 ) -> Result<(&'a mut dyn DurableCatalogState, mz_repr::Timestamp), CatalogError> {
2628 let (mut txn_batch, durable_catalog) = self.into_parts();
2629 let TransactionBatch {
2630 databases,
2631 schemas,
2632 items,
2633 comments,
2634 roles,
2635 role_auth,
2636 clusters,
2637 cluster_replicas,
2638 network_policies,
2639 introspection_sources,
2640 id_allocator,
2641 configs,
2642 source_references,
2643 settings,
2644 system_gid_mapping,
2645 system_configurations,
2646 cluster_system_configurations,
2647 replica_system_configurations,
2648 default_privileges,
2649 system_privileges,
2650 storage_collection_metadata,
2651 unfinalized_shards,
2652 txn_wal_shard,
2653 audit_log_updates,
2654 upper: _,
2655 } = &mut txn_batch;
2656 differential_dataflow::consolidation::consolidate_updates(databases);
2659 differential_dataflow::consolidation::consolidate_updates(schemas);
2660 differential_dataflow::consolidation::consolidate_updates(items);
2661 differential_dataflow::consolidation::consolidate_updates(comments);
2662 differential_dataflow::consolidation::consolidate_updates(roles);
2663 differential_dataflow::consolidation::consolidate_updates(role_auth);
2664 differential_dataflow::consolidation::consolidate_updates(clusters);
2665 differential_dataflow::consolidation::consolidate_updates(cluster_replicas);
2666 differential_dataflow::consolidation::consolidate_updates(network_policies);
2667 differential_dataflow::consolidation::consolidate_updates(introspection_sources);
2668 differential_dataflow::consolidation::consolidate_updates(id_allocator);
2669 differential_dataflow::consolidation::consolidate_updates(configs);
2670 differential_dataflow::consolidation::consolidate_updates(settings);
2671 differential_dataflow::consolidation::consolidate_updates(source_references);
2672 differential_dataflow::consolidation::consolidate_updates(system_gid_mapping);
2673 differential_dataflow::consolidation::consolidate_updates(system_configurations);
2674 differential_dataflow::consolidation::consolidate_updates(cluster_system_configurations);
2675 differential_dataflow::consolidation::consolidate_updates(replica_system_configurations);
2676 differential_dataflow::consolidation::consolidate_updates(default_privileges);
2677 differential_dataflow::consolidation::consolidate_updates(system_privileges);
2678 differential_dataflow::consolidation::consolidate_updates(storage_collection_metadata);
2679 differential_dataflow::consolidation::consolidate_updates(unfinalized_shards);
2680 differential_dataflow::consolidation::consolidate_updates(txn_wal_shard);
2681 differential_dataflow::consolidation::consolidate_updates(audit_log_updates);
2682
2683 let upper = durable_catalog
2684 .commit_transaction(txn_batch, commit_ts)
2685 .await?;
2686 Ok((durable_catalog, upper))
2687 }
2688
2689 #[mz_ore::instrument(level = "debug")]
2706 pub async fn commit(self, commit_ts: mz_repr::Timestamp) -> Result<(), CatalogError> {
2707 let op_updates = self.get_op_updates();
2708 assert!(
2709 op_updates.is_empty(),
2710 "unconsumed transaction updates: {op_updates:?}"
2711 );
2712
2713 let (durable_storage, upper) = self.commit_internal(commit_ts).await?;
2714 let updates = durable_storage.sync_updates(upper).await?;
2716 soft_assert_no_log!(
2721 durable_storage.is_read_only() || updates.iter().all(|update| update.ts == commit_ts),
2722 "unconsumed updates existed before transaction commit: commit_ts={commit_ts:?}, updates:{updates:?}"
2723 );
2724 Ok(())
2725 }
2726}
2727
2728use crate::durable::async_trait;
2729
2730use super::objects::{RoleAuthKey, RoleAuthValue};
2731
2732#[async_trait]
2733impl StorageTxn for Transaction<'_> {
2734 fn get_collection_metadata(&self) -> BTreeMap<GlobalId, ShardId> {
2735 self.storage_collection_metadata
2736 .items()
2737 .into_iter()
2738 .map(
2739 |(
2740 StorageCollectionMetadataKey { id },
2741 StorageCollectionMetadataValue { shard },
2742 )| { (*id, shard.clone()) },
2743 )
2744 .collect()
2745 }
2746
2747 fn insert_collection_metadata(
2748 &mut self,
2749 metadata: BTreeMap<GlobalId, ShardId>,
2750 ) -> Result<(), StorageError> {
2751 for (id, shard) in metadata {
2752 self.storage_collection_metadata
2753 .insert(
2754 StorageCollectionMetadataKey { id },
2755 StorageCollectionMetadataValue {
2756 shard: shard.clone(),
2757 },
2758 self.op_id,
2759 )
2760 .map_err(|err| match err {
2761 DurableCatalogError::DuplicateKey => {
2762 StorageError::CollectionMetadataAlreadyExists(id)
2763 }
2764 DurableCatalogError::UniquenessViolation => {
2765 StorageError::PersistShardAlreadyInUse(shard)
2766 }
2767 err => StorageError::Generic(anyhow::anyhow!(err)),
2768 })?;
2769 }
2770 Ok(())
2771 }
2772
2773 fn delete_collection_metadata(&mut self, ids: BTreeSet<GlobalId>) -> Vec<(GlobalId, ShardId)> {
2774 let ks: Vec<_> = ids
2775 .into_iter()
2776 .map(|id| StorageCollectionMetadataKey { id })
2777 .collect();
2778 self.storage_collection_metadata
2779 .delete_by_keys(ks, self.op_id)
2780 .into_iter()
2781 .map(
2782 |(
2783 StorageCollectionMetadataKey { id },
2784 StorageCollectionMetadataValue { shard },
2785 )| (id, shard),
2786 )
2787 .collect()
2788 }
2789
2790 fn get_unfinalized_shards(&self) -> BTreeSet<ShardId> {
2791 self.unfinalized_shards
2792 .items()
2793 .into_iter()
2794 .map(|(UnfinalizedShardKey { shard }, ())| *shard)
2795 .collect()
2796 }
2797
2798 fn insert_unfinalized_shards(&mut self, s: BTreeSet<ShardId>) -> Result<(), StorageError> {
2799 for shard in s {
2800 match self
2801 .unfinalized_shards
2802 .insert(UnfinalizedShardKey { shard }, (), self.op_id)
2803 {
2804 Ok(()) | Err(DurableCatalogError::DuplicateKey) => {}
2806 Err(e) => Err(StorageError::Generic(anyhow::anyhow!(e)))?,
2807 };
2808 }
2809 Ok(())
2810 }
2811
2812 fn mark_shards_as_finalized(&mut self, shards: BTreeSet<ShardId>) {
2813 let ks: Vec<_> = shards
2814 .into_iter()
2815 .map(|shard| UnfinalizedShardKey { shard })
2816 .collect();
2817 let _ = self.unfinalized_shards.delete_by_keys(ks, self.op_id);
2818 }
2819
2820 fn get_txn_wal_shard(&self) -> Option<ShardId> {
2821 self.txn_wal_shard
2822 .values()
2823 .iter()
2824 .next()
2825 .map(|TxnWalShardValue { shard }| *shard)
2826 }
2827
2828 fn write_txn_wal_shard(&mut self, shard: ShardId) -> Result<(), StorageError> {
2829 self.txn_wal_shard
2830 .insert((), TxnWalShardValue { shard }, self.op_id)
2831 .map_err(|err| match err {
2832 DurableCatalogError::DuplicateKey => StorageError::TxnWalShardAlreadyExists,
2833 err => StorageError::Generic(anyhow::anyhow!(err)),
2834 })
2835 }
2836}
2837
2838#[derive(Debug, Clone, Default, PartialEq)]
2840pub struct TransactionBatch {
2841 pub(crate) databases: Vec<(proto::DatabaseKey, proto::DatabaseValue, Diff)>,
2842 pub(crate) schemas: Vec<(proto::SchemaKey, proto::SchemaValue, Diff)>,
2843 pub(crate) items: Vec<(proto::ItemKey, proto::ItemValue, Diff)>,
2844 pub(crate) comments: Vec<(proto::CommentKey, proto::CommentValue, Diff)>,
2845 pub(crate) roles: Vec<(proto::RoleKey, proto::RoleValue, Diff)>,
2846 pub(crate) role_auth: Vec<(proto::RoleAuthKey, proto::RoleAuthValue, Diff)>,
2847 pub(crate) clusters: Vec<(proto::ClusterKey, proto::ClusterValue, Diff)>,
2848 pub(crate) cluster_replicas: Vec<(proto::ClusterReplicaKey, proto::ClusterReplicaValue, Diff)>,
2849 pub(crate) network_policies: Vec<(proto::NetworkPolicyKey, proto::NetworkPolicyValue, Diff)>,
2850 pub(crate) introspection_sources: Vec<(
2851 proto::ClusterIntrospectionSourceIndexKey,
2852 proto::ClusterIntrospectionSourceIndexValue,
2853 Diff,
2854 )>,
2855 pub(crate) id_allocator: Vec<(proto::IdAllocKey, proto::IdAllocValue, Diff)>,
2856 pub(crate) configs: Vec<(proto::ConfigKey, proto::ConfigValue, Diff)>,
2857 pub(crate) settings: Vec<(proto::SettingKey, proto::SettingValue, Diff)>,
2858 pub(crate) system_gid_mapping: Vec<(proto::GidMappingKey, proto::GidMappingValue, Diff)>,
2859 pub(crate) system_configurations: Vec<(
2860 proto::ServerConfigurationKey,
2861 proto::ServerConfigurationValue,
2862 Diff,
2863 )>,
2864 pub(crate) cluster_system_configurations: Vec<(
2865 proto::ClusterSystemConfigurationKey,
2866 proto::ClusterSystemConfigurationValue,
2867 Diff,
2868 )>,
2869 pub(crate) replica_system_configurations: Vec<(
2870 proto::ReplicaSystemConfigurationKey,
2871 proto::ReplicaSystemConfigurationValue,
2872 Diff,
2873 )>,
2874 pub(crate) default_privileges: Vec<(
2875 proto::DefaultPrivilegesKey,
2876 proto::DefaultPrivilegesValue,
2877 Diff,
2878 )>,
2879 pub(crate) source_references: Vec<(
2880 proto::SourceReferencesKey,
2881 proto::SourceReferencesValue,
2882 Diff,
2883 )>,
2884 pub(crate) system_privileges: Vec<(
2885 proto::SystemPrivilegesKey,
2886 proto::SystemPrivilegesValue,
2887 Diff,
2888 )>,
2889 pub(crate) storage_collection_metadata: Vec<(
2890 proto::StorageCollectionMetadataKey,
2891 proto::StorageCollectionMetadataValue,
2892 Diff,
2893 )>,
2894 pub(crate) unfinalized_shards: Vec<(proto::UnfinalizedShardKey, (), Diff)>,
2895 pub(crate) txn_wal_shard: Vec<((), proto::TxnWalShardValue, Diff)>,
2896 pub(crate) audit_log_updates: Vec<(proto::AuditLogKey, (), Diff)>,
2897 pub(crate) upper: mz_repr::Timestamp,
2899}
2900
2901impl TransactionBatch {
2902 pub fn is_empty(&self) -> bool {
2903 let TransactionBatch {
2904 databases,
2905 schemas,
2906 items,
2907 comments,
2908 roles,
2909 role_auth,
2910 clusters,
2911 cluster_replicas,
2912 network_policies,
2913 introspection_sources,
2914 id_allocator,
2915 configs,
2916 settings,
2917 source_references,
2918 system_gid_mapping,
2919 system_configurations,
2920 cluster_system_configurations,
2921 replica_system_configurations,
2922 default_privileges,
2923 system_privileges,
2924 storage_collection_metadata,
2925 unfinalized_shards,
2926 txn_wal_shard,
2927 audit_log_updates,
2928 upper: _,
2929 } = self;
2930 databases.is_empty()
2931 && schemas.is_empty()
2932 && items.is_empty()
2933 && comments.is_empty()
2934 && roles.is_empty()
2935 && role_auth.is_empty()
2936 && clusters.is_empty()
2937 && cluster_replicas.is_empty()
2938 && network_policies.is_empty()
2939 && introspection_sources.is_empty()
2940 && id_allocator.is_empty()
2941 && configs.is_empty()
2942 && settings.is_empty()
2943 && source_references.is_empty()
2944 && system_gid_mapping.is_empty()
2945 && system_configurations.is_empty()
2946 && cluster_system_configurations.is_empty()
2947 && replica_system_configurations.is_empty()
2948 && default_privileges.is_empty()
2949 && system_privileges.is_empty()
2950 && storage_collection_metadata.is_empty()
2951 && unfinalized_shards.is_empty()
2952 && txn_wal_shard.is_empty()
2953 && audit_log_updates.is_empty()
2954 }
2955}
2956
2957#[derive(Debug, Clone, PartialEq, Eq)]
2958struct TransactionUpdate<V> {
2959 value: V,
2960 ts: Timestamp,
2961 diff: Diff,
2962}
2963
2964trait UniqueName {
2966 const HAS_UNIQUE_NAME: bool;
2969 fn unique_name(&self) -> &str;
2971}
2972
2973mod unique_name {
2974 use crate::durable::objects::*;
2975
2976 macro_rules! impl_unique_name {
2977 ($($t:ty),* $(,)?) => {
2978 $(
2979 impl crate::durable::transaction::UniqueName for $t {
2980 const HAS_UNIQUE_NAME: bool = true;
2981 fn unique_name(&self) -> &str {
2982 &self.name
2983 }
2984 }
2985 )*
2986 };
2987 }
2988
2989 macro_rules! impl_no_unique_name {
2990 ($($t:ty),* $(,)?) => {
2991 $(
2992 impl crate::durable::transaction::UniqueName for $t {
2993 const HAS_UNIQUE_NAME: bool = false;
2994 fn unique_name(&self) -> &str {
2995 ""
2996 }
2997 }
2998 )*
2999 };
3000 }
3001
3002 impl_unique_name! {
3003 ClusterReplicaValue,
3004 ClusterValue,
3005 DatabaseValue,
3006 ItemValue,
3007 NetworkPolicyValue,
3008 RoleValue,
3009 SchemaValue,
3010 }
3011
3012 impl_no_unique_name!(
3013 (),
3014 ClusterIntrospectionSourceIndexValue,
3015 ClusterSystemConfigurationValue,
3016 CommentValue,
3017 ConfigValue,
3018 DefaultPrivilegesValue,
3019 GidMappingValue,
3020 IdAllocValue,
3021 ReplicaSystemConfigurationValue,
3022 ServerConfigurationValue,
3023 SettingValue,
3024 SourceReferencesValue,
3025 StorageCollectionMetadataValue,
3026 SystemPrivilegesValue,
3027 TxnWalShardValue,
3028 RoleAuthValue,
3029 );
3030
3031 #[cfg(test)]
3032 mod test {
3033 impl_no_unique_name!(String,);
3034 }
3035}
3036
3037#[derive(Debug)]
3045struct UniquenessCheck<V> {
3046 violation: fn(a: &V, b: &V) -> bool,
3047 is_unique_key_unchanged_after_update: fn(prev: &V, next: &V) -> bool,
3048}
3049
3050#[derive(Debug)]
3060struct TableTransaction<K, V> {
3061 initial: BTreeMap<K, V>,
3062 pending: BTreeMap<K, Vec<TransactionUpdate<V>>>,
3065 uniqueness_check: Option<UniquenessCheck<V>>,
3067}
3068
3069impl<K, V> TableTransaction<K, V>
3070where
3071 K: Ord + Eq + Clone + Debug,
3072 V: Ord + Clone + Debug + UniqueName,
3073{
3074 fn new<KP, VP>(initial: BTreeMap<KP, VP>) -> Result<Self, TryFromProtoError>
3081 where
3082 K: RustType<KP>,
3083 V: RustType<VP>,
3084 {
3085 let initial = initial
3086 .into_iter()
3087 .map(RustType::from_proto)
3088 .collect::<Result<_, _>>()?;
3089
3090 Ok(Self {
3091 initial,
3092 pending: BTreeMap::new(),
3093 uniqueness_check: None,
3094 })
3095 }
3096
3097 fn new_with_uniqueness_fn<KP, VP>(
3100 initial: BTreeMap<KP, VP>,
3101 uniqueness_violation: fn(a: &V, b: &V) -> bool,
3102 is_unique_key_unchanged_after_update: fn(prev: &V, next: &V) -> bool,
3103 ) -> Result<Self, TryFromProtoError>
3104 where
3105 K: RustType<KP>,
3106 V: RustType<VP>,
3107 {
3108 let initial = initial
3109 .into_iter()
3110 .map(RustType::from_proto)
3111 .collect::<Result<_, _>>()?;
3112
3113 Ok(Self {
3114 initial,
3115 pending: BTreeMap::new(),
3116 uniqueness_check: Some(UniquenessCheck {
3117 violation: uniqueness_violation,
3118 is_unique_key_unchanged_after_update,
3119 }),
3120 })
3121 }
3122
3123 fn pending<KP, VP>(self) -> Vec<(KP, VP, Diff)>
3126 where
3127 K: RustType<KP>,
3128 V: RustType<VP>,
3129 {
3130 soft_assert_no_log!(self.verify().is_ok());
3131 self.pending
3134 .into_iter()
3135 .flat_map(|(k, v)| {
3136 let mut v: Vec<_> = v
3137 .into_iter()
3138 .map(|TransactionUpdate { value, ts: _, diff }| (value, diff))
3139 .collect();
3140 differential_dataflow::consolidation::consolidate(&mut v);
3141 v.into_iter().map(move |(v, diff)| (k.clone(), v, diff))
3142 })
3143 .map(|(key, val, diff)| (key.into_proto(), val.into_proto(), diff))
3144 .collect()
3145 }
3146
3147 fn verify(&self) -> Result<(), DurableCatalogError> {
3152 if let Some(check) = &self.uniqueness_check {
3153 let items = self.values();
3155 if V::HAS_UNIQUE_NAME {
3156 let by_name: BTreeMap<_, _> = items
3157 .iter()
3158 .enumerate()
3159 .map(|(v, vi)| (vi.unique_name(), (v, vi)))
3160 .collect();
3161 for (i, vi) in items.iter().enumerate() {
3162 if let Some((j, vj)) = by_name.get(vi.unique_name()) {
3163 if i != *j && (check.violation)(vi, *vj) {
3164 return Err(DurableCatalogError::UniquenessViolation);
3165 }
3166 }
3167 }
3168 } else {
3169 for (i, vi) in items.iter().enumerate() {
3170 for (j, vj) in items.iter().enumerate() {
3171 if i != j && (check.violation)(vi, vj) {
3172 return Err(DurableCatalogError::UniquenessViolation);
3173 }
3174 }
3175 }
3176 }
3177 }
3178 soft_assert_no_log!(
3179 self.pending
3180 .values()
3181 .all(|pending| { pending.is_sorted_by(|a, b| a.ts <= b.ts) }),
3182 "pending should be sorted by timestamp: {:?}",
3183 self.pending
3184 );
3185 Ok(())
3186 }
3187
3188 fn verify_keys<'a>(
3193 &self,
3194 keys: impl IntoIterator<Item = &'a K>,
3195 ) -> Result<(), DurableCatalogError>
3196 where
3197 K: 'a,
3198 {
3199 if let Some(check) = &self.uniqueness_check {
3200 let entries: Vec<_> = keys
3201 .into_iter()
3202 .filter_map(|key| self.get(key).map(|value| (key, value)))
3203 .collect();
3204 for (ki, vi) in self.items() {
3206 for (kj, vj) in &entries {
3207 if ki != *kj && (check.violation)(vi, vj) {
3208 return Err(DurableCatalogError::UniquenessViolation);
3209 }
3210 }
3211 }
3212 }
3213 soft_assert_no_log!(self.verify().is_ok());
3214 Ok(())
3215 }
3216
3217 fn for_values<'a, F: FnMut(&'a K, &'a V)>(&'a self, mut f: F) {
3220 let mut seen = BTreeSet::new();
3221 for k in self.pending.keys() {
3222 seen.insert(k);
3223 let v = self.get(k);
3224 if let Some(v) = v {
3227 f(k, v);
3228 }
3229 }
3230 for (k, v) in self.initial.iter() {
3231 if !seen.contains(k) {
3233 f(k, v);
3234 }
3235 }
3236 }
3237
3238 fn get(&self, k: &K) -> Option<&V> {
3240 let pending = self.pending.get(k).map(Vec::as_slice).unwrap_or_default();
3241 let mut updates = Vec::with_capacity(pending.len() + 1);
3242 if let Some(initial) = self.initial.get(k) {
3243 updates.push((initial, Diff::ONE));
3244 }
3245 updates.extend(
3246 pending
3247 .into_iter()
3248 .map(|TransactionUpdate { value, ts: _, diff }| (value, *diff)),
3249 );
3250
3251 differential_dataflow::consolidation::consolidate(&mut updates);
3252 assert!(updates.len() <= 1);
3253 updates.into_iter().next().map(|(v, _)| v)
3254 }
3255
3256 #[cfg(test)]
3261 fn items_cloned(&self) -> BTreeMap<K, V> {
3262 let mut items = BTreeMap::new();
3263 self.for_values(|k, v| {
3264 items.insert(k.clone(), v.clone());
3265 });
3266 items
3267 }
3268
3269 fn current_items_proto<KP, VP>(&self) -> BTreeMap<KP, VP>
3273 where
3274 K: RustType<KP>,
3275 V: RustType<VP>,
3276 KP: Ord,
3277 {
3278 let mut items = BTreeMap::new();
3279 self.for_values(|k, v| {
3280 items.insert(k.into_proto(), v.into_proto());
3281 });
3282 items
3283 }
3284
3285 fn items(&self) -> BTreeMap<&K, &V> {
3288 let mut items = BTreeMap::new();
3289 self.for_values(|k, v| {
3290 items.insert(k, v);
3291 });
3292 items
3293 }
3294
3295 fn values(&self) -> BTreeSet<&V> {
3297 let mut items = BTreeSet::new();
3298 self.for_values(|_, v| {
3299 items.insert(v);
3300 });
3301 items
3302 }
3303
3304 fn len(&self) -> usize {
3306 let mut count = 0;
3307 self.for_values(|_, _| {
3308 count += 1;
3309 });
3310 count
3311 }
3312
3313 fn for_values_mut<F: FnMut(&mut BTreeMap<K, Vec<TransactionUpdate<V>>>, &K, &V)>(
3317 &mut self,
3318 mut f: F,
3319 ) {
3320 let mut pending = BTreeMap::new();
3321 self.for_values(|k, v| f(&mut pending, k, v));
3322 for (k, updates) in pending {
3323 self.pending.entry(k).or_default().extend(updates);
3324 }
3325 }
3326
3327 fn insert(&mut self, k: K, v: V, ts: Timestamp) -> Result<(), DurableCatalogError> {
3331 let mut violation = None;
3332 let uniqueness_violation = self.uniqueness_check.as_ref().map(|check| check.violation);
3333 self.for_values(|for_k, for_v| {
3334 if &k == for_k {
3335 violation = Some(DurableCatalogError::DuplicateKey);
3336 }
3337 if let Some(uniqueness_violation) = uniqueness_violation {
3338 if uniqueness_violation(for_v, &v) {
3339 violation = Some(DurableCatalogError::UniquenessViolation);
3340 }
3341 }
3342 });
3343 if let Some(violation) = violation {
3344 return Err(violation);
3345 }
3346 self.pending.entry(k).or_default().push(TransactionUpdate {
3347 value: v,
3348 ts,
3349 diff: Diff::ONE,
3350 });
3351 soft_assert_no_log!(self.verify().is_ok());
3352 Ok(())
3353 }
3354
3355 fn update<F: Fn(&K, &V) -> Option<V>>(
3364 &mut self,
3365 f: F,
3366 ts: Timestamp,
3367 ) -> Result<Diff, DurableCatalogError> {
3368 let mut changed = Diff::ZERO;
3369 let mut keys = BTreeSet::new();
3370 let pending = self.pending.clone();
3372 self.for_values_mut(|p, k, v| {
3373 if let Some(next) = f(k, v) {
3374 changed += Diff::ONE;
3375 keys.insert(k.clone());
3376 let updates = p.entry(k.clone()).or_default();
3377 updates.push(TransactionUpdate {
3378 value: v.clone(),
3379 ts,
3380 diff: Diff::MINUS_ONE,
3381 });
3382 updates.push(TransactionUpdate {
3383 value: next,
3384 ts,
3385 diff: Diff::ONE,
3386 });
3387 }
3388 });
3389 if let Err(err) = self.verify_keys(&keys) {
3391 self.pending = pending;
3392 Err(err)
3393 } else {
3394 Ok(changed)
3395 }
3396 }
3397
3398 fn update_by_key(&mut self, k: K, v: V, ts: Timestamp) -> Result<bool, DurableCatalogError> {
3403 if let Some(cur_v) = self.get(&k) {
3404 if v != *cur_v {
3405 self.set(k, Some(v), ts)?;
3406 }
3407 Ok(true)
3408 } else {
3409 Ok(false)
3410 }
3411 }
3412
3413 fn update_by_keys(
3418 &mut self,
3419 kvs: impl IntoIterator<Item = (K, V)>,
3420 ts: Timestamp,
3421 ) -> Result<Diff, DurableCatalogError> {
3422 let kvs: Vec<_> = kvs
3423 .into_iter()
3424 .filter_map(|(k, v)| match self.get(&k) {
3425 Some(cur_v) => Some((*cur_v == v, k, v)),
3427 None => None,
3428 })
3429 .collect();
3430 let changed = kvs.len();
3431 let changed =
3432 Diff::try_from(changed).map_err(|e| DurableCatalogError::Internal(e.to_string()))?;
3433 let kvs = kvs
3434 .into_iter()
3435 .filter(|(no_op, _, _)| !no_op)
3437 .map(|(_, k, v)| (k, Some(v)))
3438 .collect();
3439 self.set_many(kvs, ts)?;
3440 Ok(changed)
3441 }
3442
3443 fn update_needs_uniqueness_check(&self, prev: Option<&V>, next: Option<&V>) -> bool {
3445 match (&self.uniqueness_check, prev, next) {
3446 (None, _, _) | (_, _, None) => false,
3448 (Some(check), Some(prev), Some(next)) => {
3450 !(check.is_unique_key_unchanged_after_update)(prev, next)
3451 }
3452 (Some(_), None, Some(_)) => true,
3454 }
3455 }
3456
3457 fn set(&mut self, k: K, v: Option<V>, ts: Timestamp) -> Result<Option<V>, DurableCatalogError> {
3464 let prev = self.get(&k).cloned();
3465 let needs_uniqueness_check = self.update_needs_uniqueness_check(prev.as_ref(), v.as_ref());
3466 let entry = self.pending.entry(k.clone()).or_default();
3467 let restore_len = entry.len();
3468
3469 match (v, prev.clone()) {
3470 (Some(v), Some(prev)) => {
3471 entry.push(TransactionUpdate {
3472 value: prev,
3473 ts,
3474 diff: Diff::MINUS_ONE,
3475 });
3476 entry.push(TransactionUpdate {
3477 value: v,
3478 ts,
3479 diff: Diff::ONE,
3480 });
3481 }
3482 (Some(v), None) => {
3483 entry.push(TransactionUpdate {
3484 value: v,
3485 ts,
3486 diff: Diff::ONE,
3487 });
3488 }
3489 (None, Some(prev)) => {
3490 entry.push(TransactionUpdate {
3491 value: prev,
3492 ts,
3493 diff: Diff::MINUS_ONE,
3494 });
3495 }
3496 (None, None) => {}
3497 }
3498
3499 if needs_uniqueness_check {
3501 if let Err(err) = self.verify_keys([&k]) {
3502 let pending = self.pending.get_mut(&k).expect("inserted above");
3505 pending.truncate(restore_len);
3506 return Err(err);
3507 }
3508 }
3509 Ok(prev)
3510 }
3511
3512 fn set_many(
3517 &mut self,
3518 kvs: BTreeMap<K, Option<V>>,
3519 ts: Timestamp,
3520 ) -> Result<BTreeMap<K, Option<V>>, DurableCatalogError> {
3521 if kvs.is_empty() {
3522 return Ok(BTreeMap::new());
3523 }
3524
3525 let mut prevs = BTreeMap::new();
3526 let mut restores = BTreeMap::new();
3527 let mut keys_to_verify_uniqueness = Vec::new();
3529
3530 for (k, v) in kvs {
3531 let prev = self.get(&k).cloned();
3532 if self.update_needs_uniqueness_check(prev.as_ref(), v.as_ref()) {
3533 keys_to_verify_uniqueness.push(k.clone());
3534 }
3535 let entry = self.pending.entry(k.clone()).or_default();
3536 restores.insert(k.clone(), entry.len());
3537
3538 match (v, prev.clone()) {
3539 (Some(v), Some(prev)) => {
3540 entry.push(TransactionUpdate {
3541 value: prev,
3542 ts,
3543 diff: Diff::MINUS_ONE,
3544 });
3545 entry.push(TransactionUpdate {
3546 value: v,
3547 ts,
3548 diff: Diff::ONE,
3549 });
3550 }
3551 (Some(v), None) => {
3552 entry.push(TransactionUpdate {
3553 value: v,
3554 ts,
3555 diff: Diff::ONE,
3556 });
3557 }
3558 (None, Some(prev)) => {
3559 entry.push(TransactionUpdate {
3560 value: prev,
3561 ts,
3562 diff: Diff::MINUS_ONE,
3563 });
3564 }
3565 (None, None) => {}
3566 }
3567
3568 prevs.insert(k, prev);
3569 }
3570
3571 if let Err(err) = self.verify_keys(keys_to_verify_uniqueness.iter()) {
3573 for (k, restore_len) in restores {
3574 let pending = self.pending.get_mut(&k).expect("inserted above");
3577 pending.truncate(restore_len);
3578 }
3579 Err(err)
3580 } else {
3581 Ok(prevs)
3582 }
3583 }
3584
3585 fn delete<F: Fn(&K, &V) -> bool>(&mut self, f: F, ts: Timestamp) -> Vec<(K, V)> {
3591 let mut deleted = Vec::new();
3592 self.for_values_mut(|p, k, v| {
3593 if f(k, v) {
3594 deleted.push((k.clone(), v.clone()));
3595 p.entry(k.clone()).or_default().push(TransactionUpdate {
3596 value: v.clone(),
3597 ts,
3598 diff: Diff::MINUS_ONE,
3599 });
3600 }
3601 });
3602 soft_assert_no_log!(self.verify().is_ok());
3603 deleted
3604 }
3605
3606 fn delete_by_key(&mut self, k: K, ts: Timestamp) -> Option<V> {
3610 self.set(k, None, ts)
3611 .expect("deleting an entry cannot violate uniqueness")
3612 }
3613
3614 fn delete_by_keys(&mut self, ks: impl IntoIterator<Item = K>, ts: Timestamp) -> Vec<(K, V)> {
3618 let kvs = ks.into_iter().map(|k| (k, None)).collect();
3619 let prevs = self
3620 .set_many(kvs, ts)
3621 .expect("deleting entries cannot violate uniqueness");
3622 prevs
3623 .into_iter()
3624 .filter_map(|(k, v)| v.map(|v| (k, v)))
3625 .collect()
3626 }
3627}
3628
3629#[cfg(test)]
3630#[allow(clippy::unwrap_used)]
3631mod tests {
3632 use super::*;
3633
3634 use mz_controller::clusters::ReplicaLogging;
3635 use mz_ore::now::SYSTEM_TIME;
3636 use mz_ore::{assert_none, assert_ok};
3637 use mz_persist_client::cache::PersistClientCache;
3638 use mz_persist_types::PersistLocation;
3639 use semver::Version;
3640
3641 use crate::durable::{
3642 ReplicaConfig, ReplicaLocation, TestCatalogStateBuilder, test_bootstrap_args,
3643 };
3644 use crate::memory;
3645
3646 #[mz_ore::test]
3647 fn test_table_transaction_simple() {
3648 fn uniqueness_violation(a: &String, b: &String) -> bool {
3649 a == b
3650 }
3651 let mut table = TableTransaction::new_with_uniqueness_fn(
3652 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "a".to_string())]),
3653 uniqueness_violation,
3654 uniqueness_violation,
3655 )
3656 .unwrap();
3657
3658 assert_ok!(table.insert(2i64.to_le_bytes().to_vec(), "b".to_string(), 0));
3661 assert_ok!(table.insert(3i64.to_le_bytes().to_vec(), "c".to_string(), 0));
3662 assert!(
3663 table
3664 .insert(1i64.to_le_bytes().to_vec(), "c".to_string(), 0)
3665 .is_err()
3666 );
3667 assert!(
3668 table
3669 .insert(4i64.to_le_bytes().to_vec(), "c".to_string(), 0)
3670 .is_err()
3671 );
3672 }
3673
3674 #[mz_ore::test]
3675 fn test_skip_scan_when_unique_key_unchanged() {
3676 fn first_char_same(prev: &String, next: &String) -> bool {
3677 prev.chars().next() == next.chars().next()
3678 }
3679
3680 fn panic_uniqueness_violation(_: &String, _: &String) -> bool {
3682 panic!("uniqueness scan ran for an update that kept the same unique key");
3683 }
3684 let mut table = TableTransaction::new_with_uniqueness_fn(
3685 BTreeMap::from([
3686 (1i64.to_le_bytes().to_vec(), "a1".to_string()),
3687 (2i64.to_le_bytes().to_vec(), "b1".to_string()),
3688 ]),
3689 panic_uniqueness_violation,
3690 first_char_same,
3692 )
3693 .unwrap();
3694 assert!(
3697 table
3698 .update_by_key(1i64.to_le_bytes().to_vec(), "a2".to_string(), 0)
3699 .unwrap()
3700 );
3701
3702 fn real_uniqueness_violation(a: &String, b: &String) -> bool {
3704 a.chars().next() == b.chars().next()
3705 }
3706 let mut table = TableTransaction::new_with_uniqueness_fn(
3707 BTreeMap::from([
3708 (1i64.to_le_bytes().to_vec(), "a1".to_string()),
3709 (2i64.to_le_bytes().to_vec(), "b1".to_string()),
3710 ]),
3711 real_uniqueness_violation,
3712 first_char_same,
3713 )
3714 .unwrap();
3715 assert!(
3718 table
3719 .update_by_key(1i64.to_le_bytes().to_vec(), "b2".to_string(), 0)
3720 .is_err()
3721 );
3722 }
3723
3724 #[mz_ore::test]
3725 fn test_table_transaction() {
3726 fn uniqueness_violation(a: &String, b: &String) -> bool {
3727 a == b
3728 }
3729 let mut table: BTreeMap<Vec<u8>, String> = BTreeMap::new();
3730
3731 fn commit(
3732 table: &mut BTreeMap<Vec<u8>, String>,
3733 mut pending: Vec<(Vec<u8>, String, Diff)>,
3734 ) {
3735 pending.sort_by(|a, b| a.2.cmp(&b.2));
3737 for (k, v, diff) in pending {
3738 if diff == Diff::MINUS_ONE {
3739 let prev = table.remove(&k);
3740 assert_eq!(prev, Some(v));
3741 } else if diff == Diff::ONE {
3742 let prev = table.insert(k, v);
3743 assert_eq!(prev, None);
3744 } else {
3745 panic!("unexpected diff: {diff}");
3746 }
3747 }
3748 }
3749
3750 table.insert(1i64.to_le_bytes().to_vec(), "v1".to_string());
3751 table.insert(2i64.to_le_bytes().to_vec(), "v2".to_string());
3752 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3753 table.clone(),
3754 uniqueness_violation,
3755 uniqueness_violation,
3756 )
3757 .unwrap();
3758 assert_eq!(table_txn.items_cloned(), table);
3759 assert_eq!(table_txn.delete(|_k, _v| false, 0).len(), 0);
3760 assert_eq!(table_txn.delete(|_k, v| v == "v2", 1).len(), 1);
3761 assert_eq!(
3762 table_txn.items_cloned(),
3763 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v1".to_string())])
3764 );
3765 assert_eq!(
3766 table_txn
3767 .update(|_k, _v| Some("v3".to_string()), 2)
3768 .unwrap(),
3769 Diff::ONE
3770 );
3771
3772 table_txn
3774 .insert(3i64.to_le_bytes().to_vec(), "v3".to_string(), 3)
3775 .unwrap_err();
3776
3777 table_txn
3778 .insert(3i64.to_le_bytes().to_vec(), "v4".to_string(), 4)
3779 .unwrap();
3780 assert_eq!(
3781 table_txn.items_cloned(),
3782 BTreeMap::from([
3783 (1i64.to_le_bytes().to_vec(), "v3".to_string()),
3784 (3i64.to_le_bytes().to_vec(), "v4".to_string()),
3785 ])
3786 );
3787 let err = table_txn
3788 .update(|_k, _v| Some("v1".to_string()), 5)
3789 .unwrap_err();
3790 assert!(
3791 matches!(err, DurableCatalogError::UniquenessViolation),
3792 "unexpected err: {err:?}"
3793 );
3794 let pending = table_txn.pending();
3795 assert_eq!(
3796 pending,
3797 vec![
3798 (
3799 1i64.to_le_bytes().to_vec(),
3800 "v1".to_string(),
3801 Diff::MINUS_ONE
3802 ),
3803 (1i64.to_le_bytes().to_vec(), "v3".to_string(), Diff::ONE),
3804 (
3805 2i64.to_le_bytes().to_vec(),
3806 "v2".to_string(),
3807 Diff::MINUS_ONE
3808 ),
3809 (3i64.to_le_bytes().to_vec(), "v4".to_string(), Diff::ONE),
3810 ]
3811 );
3812 commit(&mut table, pending);
3813 assert_eq!(
3814 table,
3815 BTreeMap::from([
3816 (1i64.to_le_bytes().to_vec(), "v3".to_string()),
3817 (3i64.to_le_bytes().to_vec(), "v4".to_string())
3818 ])
3819 );
3820
3821 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3822 table.clone(),
3823 uniqueness_violation,
3824 uniqueness_violation,
3825 )
3826 .unwrap();
3827 assert_eq!(
3829 table_txn.delete(|k, _v| k == &1i64.to_le_bytes(), 0).len(),
3830 1
3831 );
3832 table_txn
3833 .insert(1i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3834 .unwrap();
3835 table_txn
3837 .insert(5i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3838 .unwrap_err();
3839 table_txn
3841 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 0)
3842 .unwrap_err();
3843 assert_eq!(
3844 table_txn.delete(|k, _v| k == &1i64.to_le_bytes(), 0).len(),
3845 1
3846 );
3847 table_txn
3849 .insert(5i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3850 .unwrap();
3851 table_txn
3852 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 0)
3853 .unwrap();
3854 let pending = table_txn.pending();
3855 assert_eq!(
3856 pending,
3857 vec![
3858 (
3859 1i64.to_le_bytes().to_vec(),
3860 "v3".to_string(),
3861 Diff::MINUS_ONE
3862 ),
3863 (1i64.to_le_bytes().to_vec(), "v5".to_string(), Diff::ONE),
3864 (5i64.to_le_bytes().to_vec(), "v3".to_string(), Diff::ONE),
3865 ]
3866 );
3867 commit(&mut table, pending);
3868 assert_eq!(
3869 table,
3870 BTreeMap::from([
3871 (1i64.to_le_bytes().to_vec(), "v5".to_string()),
3872 (3i64.to_le_bytes().to_vec(), "v4".to_string()),
3873 (5i64.to_le_bytes().to_vec(), "v3".to_string()),
3874 ])
3875 );
3876
3877 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3878 table.clone(),
3879 uniqueness_violation,
3880 uniqueness_violation,
3881 )
3882 .unwrap();
3883 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 3);
3884 table_txn
3885 .insert(1i64.to_le_bytes().to_vec(), "v1".to_string(), 0)
3886 .unwrap();
3887
3888 commit(&mut table, table_txn.pending());
3889 assert_eq!(
3890 table,
3891 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v1".to_string()),])
3892 );
3893
3894 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3895 table.clone(),
3896 uniqueness_violation,
3897 uniqueness_violation,
3898 )
3899 .unwrap();
3900 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 1);
3901 table_txn
3902 .insert(1i64.to_le_bytes().to_vec(), "v2".to_string(), 0)
3903 .unwrap();
3904 commit(&mut table, table_txn.pending());
3905 assert_eq!(
3906 table,
3907 BTreeMap::from([(1i64.to_le_bytes().to_vec(), "v2".to_string()),])
3908 );
3909
3910 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3912 table.clone(),
3913 uniqueness_violation,
3914 uniqueness_violation,
3915 )
3916 .unwrap();
3917 assert_eq!(table_txn.delete(|_k, _v| true, 0).len(), 1);
3918 table_txn
3919 .insert(1i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3920 .unwrap();
3921 table_txn
3922 .insert(1i64.to_le_bytes().to_vec(), "v4".to_string(), 1)
3923 .unwrap_err();
3924 assert_eq!(table_txn.delete(|_k, _v| true, 1).len(), 1);
3925 table_txn
3926 .insert(1i64.to_le_bytes().to_vec(), "v5".to_string(), 1)
3927 .unwrap();
3928 commit(&mut table, table_txn.pending());
3929 assert_eq!(
3930 table.clone().into_iter().collect::<Vec<_>>(),
3931 vec![(1i64.to_le_bytes().to_vec(), "v5".to_string())]
3932 );
3933
3934 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3936 table.clone(),
3937 uniqueness_violation,
3938 uniqueness_violation,
3939 )
3940 .unwrap();
3941 table_txn
3943 .set(2i64.to_le_bytes().to_vec(), Some("v5".to_string()), 0)
3944 .unwrap_err();
3945 table_txn
3946 .set(3i64.to_le_bytes().to_vec(), Some("v6".to_string()), 1)
3947 .unwrap();
3948 table_txn.set(2i64.to_le_bytes().to_vec(), None, 2).unwrap();
3949 table_txn.set(1i64.to_le_bytes().to_vec(), None, 2).unwrap();
3950 let pending = table_txn.pending();
3951 assert_eq!(
3952 pending,
3953 vec![
3954 (
3955 1i64.to_le_bytes().to_vec(),
3956 "v5".to_string(),
3957 Diff::MINUS_ONE
3958 ),
3959 (3i64.to_le_bytes().to_vec(), "v6".to_string(), Diff::ONE),
3960 ]
3961 );
3962 commit(&mut table, pending);
3963 assert_eq!(
3964 table,
3965 BTreeMap::from([(3i64.to_le_bytes().to_vec(), "v6".to_string())])
3966 );
3967
3968 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3970 table.clone(),
3971 uniqueness_violation,
3972 uniqueness_violation,
3973 )
3974 .unwrap();
3975 table_txn
3976 .set(3i64.to_le_bytes().to_vec(), Some("v6".to_string()), 0)
3977 .unwrap();
3978 let pending = table_txn.pending::<Vec<u8>, String>();
3979 assert!(pending.is_empty());
3980
3981 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3983 table.clone(),
3984 uniqueness_violation,
3985 uniqueness_violation,
3986 )
3987 .unwrap();
3988 table_txn
3990 .set_many(
3991 BTreeMap::from([
3992 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
3993 (42i64.to_le_bytes().to_vec(), Some("v1".to_string())),
3994 ]),
3995 0,
3996 )
3997 .unwrap_err();
3998 table_txn
3999 .set_many(
4000 BTreeMap::from([
4001 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
4002 (3i64.to_le_bytes().to_vec(), Some("v1".to_string())),
4003 ]),
4004 1,
4005 )
4006 .unwrap();
4007 table_txn
4008 .set_many(
4009 BTreeMap::from([
4010 (42i64.to_le_bytes().to_vec(), Some("v7".to_string())),
4011 (3i64.to_le_bytes().to_vec(), None),
4012 ]),
4013 2,
4014 )
4015 .unwrap();
4016 let pending = table_txn.pending();
4017 assert_eq!(
4018 pending,
4019 vec![
4020 (1i64.to_le_bytes().to_vec(), "v6".to_string(), Diff::ONE),
4021 (
4022 3i64.to_le_bytes().to_vec(),
4023 "v6".to_string(),
4024 Diff::MINUS_ONE
4025 ),
4026 (42i64.to_le_bytes().to_vec(), "v7".to_string(), Diff::ONE),
4027 ]
4028 );
4029 commit(&mut table, pending);
4030 assert_eq!(
4031 table,
4032 BTreeMap::from([
4033 (1i64.to_le_bytes().to_vec(), "v6".to_string()),
4034 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4035 ])
4036 );
4037
4038 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4040 table.clone(),
4041 uniqueness_violation,
4042 uniqueness_violation,
4043 )
4044 .unwrap();
4045 table_txn
4046 .set_many(
4047 BTreeMap::from([
4048 (1i64.to_le_bytes().to_vec(), Some("v6".to_string())),
4049 (42i64.to_le_bytes().to_vec(), Some("v7".to_string())),
4050 ]),
4051 0,
4052 )
4053 .unwrap();
4054 let pending = table_txn.pending::<Vec<u8>, String>();
4055 assert!(pending.is_empty());
4056 commit(&mut table, pending);
4057 assert_eq!(
4058 table,
4059 BTreeMap::from([
4060 (1i64.to_le_bytes().to_vec(), "v6".to_string()),
4061 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4062 ])
4063 );
4064
4065 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4067 table.clone(),
4068 uniqueness_violation,
4069 uniqueness_violation,
4070 )
4071 .unwrap();
4072 table_txn
4074 .update_by_key(1i64.to_le_bytes().to_vec(), "v7".to_string(), 0)
4075 .unwrap_err();
4076 assert!(
4077 table_txn
4078 .update_by_key(1i64.to_le_bytes().to_vec(), "v8".to_string(), 1)
4079 .unwrap()
4080 );
4081 assert!(
4082 !table_txn
4083 .update_by_key(5i64.to_le_bytes().to_vec(), "v8".to_string(), 2)
4084 .unwrap()
4085 );
4086 let pending = table_txn.pending();
4087 assert_eq!(
4088 pending,
4089 vec![
4090 (
4091 1i64.to_le_bytes().to_vec(),
4092 "v6".to_string(),
4093 Diff::MINUS_ONE
4094 ),
4095 (1i64.to_le_bytes().to_vec(), "v8".to_string(), Diff::ONE),
4096 ]
4097 );
4098 commit(&mut table, pending);
4099 assert_eq!(
4100 table,
4101 BTreeMap::from([
4102 (1i64.to_le_bytes().to_vec(), "v8".to_string()),
4103 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4104 ])
4105 );
4106
4107 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4109 table.clone(),
4110 uniqueness_violation,
4111 uniqueness_violation,
4112 )
4113 .unwrap();
4114 assert!(
4115 table_txn
4116 .update_by_key(1i64.to_le_bytes().to_vec(), "v8".to_string(), 0)
4117 .unwrap()
4118 );
4119 let pending = table_txn.pending::<Vec<u8>, String>();
4120 assert!(pending.is_empty());
4121 commit(&mut table, pending);
4122 assert_eq!(
4123 table,
4124 BTreeMap::from([
4125 (1i64.to_le_bytes().to_vec(), "v8".to_string()),
4126 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4127 ])
4128 );
4129
4130 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4132 table.clone(),
4133 uniqueness_violation,
4134 uniqueness_violation,
4135 )
4136 .unwrap();
4137 table_txn
4139 .update_by_keys(
4140 [
4141 (1i64.to_le_bytes().to_vec(), "v7".to_string()),
4142 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4143 ],
4144 0,
4145 )
4146 .unwrap_err();
4147 let n = table_txn
4148 .update_by_keys(
4149 [
4150 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4151 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4152 ],
4153 1,
4154 )
4155 .unwrap();
4156 assert_eq!(n, Diff::ONE);
4157 let n = table_txn
4158 .update_by_keys(
4159 [
4160 (15i64.to_le_bytes().to_vec(), "v9".to_string()),
4161 (5i64.to_le_bytes().to_vec(), "v7".to_string()),
4162 ],
4163 2,
4164 )
4165 .unwrap();
4166 assert_eq!(n, Diff::ZERO);
4167 let pending = table_txn.pending();
4168 assert_eq!(
4169 pending,
4170 vec![
4171 (
4172 1i64.to_le_bytes().to_vec(),
4173 "v8".to_string(),
4174 Diff::MINUS_ONE
4175 ),
4176 (1i64.to_le_bytes().to_vec(), "v9".to_string(), Diff::ONE),
4177 ]
4178 );
4179 commit(&mut table, pending);
4180 assert_eq!(
4181 table,
4182 BTreeMap::from([
4183 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4184 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4185 ])
4186 );
4187
4188 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4190 table.clone(),
4191 uniqueness_violation,
4192 uniqueness_violation,
4193 )
4194 .unwrap();
4195 let n = table_txn
4196 .update_by_keys(
4197 [
4198 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4199 (42i64.to_le_bytes().to_vec(), "v7".to_string()),
4200 ],
4201 0,
4202 )
4203 .unwrap();
4204 assert_eq!(n, Diff::from(2));
4205 let pending = table_txn.pending::<Vec<u8>, String>();
4206 assert!(pending.is_empty());
4207 commit(&mut table, pending);
4208 assert_eq!(
4209 table,
4210 BTreeMap::from([
4211 (1i64.to_le_bytes().to_vec(), "v9".to_string()),
4212 (42i64.to_le_bytes().to_vec(), "v7".to_string())
4213 ])
4214 );
4215
4216 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4218 table.clone(),
4219 uniqueness_violation,
4220 uniqueness_violation,
4221 )
4222 .unwrap();
4223 let prev = table_txn.delete_by_key(1i64.to_le_bytes().to_vec(), 0);
4224 assert_eq!(prev, Some("v9".to_string()));
4225 let prev = table_txn.delete_by_key(5i64.to_le_bytes().to_vec(), 1);
4226 assert_none!(prev);
4227 let prev = table_txn.delete_by_key(1i64.to_le_bytes().to_vec(), 2);
4228 assert_none!(prev);
4229 let pending = table_txn.pending();
4230 assert_eq!(
4231 pending,
4232 vec![(
4233 1i64.to_le_bytes().to_vec(),
4234 "v9".to_string(),
4235 Diff::MINUS_ONE
4236 ),]
4237 );
4238 commit(&mut table, pending);
4239 assert_eq!(
4240 table,
4241 BTreeMap::from([(42i64.to_le_bytes().to_vec(), "v7".to_string())])
4242 );
4243
4244 let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4246 table.clone(),
4247 uniqueness_violation,
4248 uniqueness_violation,
4249 )
4250 .unwrap();
4251 let prevs = table_txn.delete_by_keys(
4252 [42i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4253 0,
4254 );
4255 assert_eq!(
4256 prevs,
4257 vec![(42i64.to_le_bytes().to_vec(), "v7".to_string())]
4258 );
4259 let prevs = table_txn.delete_by_keys(
4260 [42i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4261 1,
4262 );
4263 assert_eq!(prevs, vec![]);
4264 let prevs = table_txn.delete_by_keys(
4265 [10i64.to_le_bytes().to_vec(), 55i64.to_le_bytes().to_vec()],
4266 2,
4267 );
4268 assert_eq!(prevs, vec![]);
4269 let pending = table_txn.pending();
4270 assert_eq!(
4271 pending,
4272 vec![(
4273 42i64.to_le_bytes().to_vec(),
4274 "v7".to_string(),
4275 Diff::MINUS_ONE
4276 ),]
4277 );
4278 commit(&mut table, pending);
4279 assert_eq!(table, BTreeMap::new());
4280 }
4281
4282 #[mz_ore::test(tokio::test)]
4283 #[cfg_attr(miri, ignore)] async fn test_savepoint() {
4285 const VERSION: Version = Version::new(26, 0, 0);
4286 let mut persist_cache = PersistClientCache::new_no_metrics();
4287 persist_cache.cfg.build_version = VERSION;
4288 let persist_client = persist_cache
4289 .open(PersistLocation::new_in_mem())
4290 .await
4291 .unwrap();
4292 let state_builder = TestCatalogStateBuilder::new(persist_client)
4293 .with_default_deploy_generation()
4294 .with_version(VERSION);
4295
4296 let _ = state_builder
4298 .clone()
4299 .unwrap_build()
4300 .await
4301 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4302 .await
4303 .unwrap();
4304 let mut savepoint_state = state_builder
4305 .unwrap_build()
4306 .await
4307 .open_savepoint(SYSTEM_TIME().into(), &test_bootstrap_args())
4308 .await
4309 .unwrap();
4310
4311 let initial_snapshot = savepoint_state.sync_to_current_updates().await.unwrap();
4312 assert!(!initial_snapshot.is_empty());
4313
4314 let db_name = "db";
4315 let db_owner = RoleId::User(42);
4316 let db_privileges = Vec::new();
4317 let mut txn = savepoint_state.transaction().await.unwrap();
4318 let (db_id, db_oid) = txn
4319 .insert_user_database(db_name, db_owner, db_privileges.clone(), &HashSet::new())
4320 .unwrap();
4321 let commit_ts = txn.upper();
4322 txn.commit_internal(commit_ts).await.unwrap();
4323 let updates = savepoint_state.sync_to_current_updates().await.unwrap();
4324 let update = updates.into_element();
4325
4326 assert_eq!(update.diff, StateDiff::Addition);
4327
4328 let db = match update.kind {
4329 memory::objects::StateUpdateKind::Database(db) => db,
4330 update => panic!("unexpected update: {update:?}"),
4331 };
4332
4333 assert_eq!(db_id, db.id);
4334 assert_eq!(db_oid, db.oid);
4335 assert_eq!(db_name, db.name);
4336 assert_eq!(db_owner, db.owner_id);
4337 assert_eq!(db_privileges, db.privileges);
4338 }
4339
4340 #[mz_ore::test(tokio::test)]
4344 #[cfg_attr(miri, ignore)] async fn test_insert_replica_with_id_does_not_consume_allocator() {
4346 const VERSION: Version = Version::new(26, 0, 0);
4347 let mut persist_cache = PersistClientCache::new_no_metrics();
4348 persist_cache.cfg.build_version = VERSION;
4349 let persist_client = persist_cache
4350 .open(PersistLocation::new_in_mem())
4351 .await
4352 .unwrap();
4353 let state_builder = TestCatalogStateBuilder::new(persist_client)
4354 .with_default_deploy_generation()
4355 .with_version(VERSION);
4356 let mut state = state_builder
4357 .unwrap_build()
4358 .await
4359 .open(SYSTEM_TIME().into(), &test_bootstrap_args())
4360 .await
4361 .unwrap();
4362
4363 let cluster_id = ClusterId::User(1);
4366 let owner_id = RoleId::User(1);
4367 let config = ReplicaConfig {
4368 location: ReplicaLocation::Managed {
4369 size: "1".to_string(),
4370 availability_zones: Vec::new(),
4371 internal: false,
4372 billed_as: None,
4373 pending: false,
4374 },
4375 logging: ReplicaLogging {
4376 log_logging: false,
4377 interval: Some(Duration::from_secs(1)),
4378 },
4379 };
4380
4381 let commit_ts = state.current_upper().await;
4383 let a = state
4384 .allocate_user_replica_ids(1, commit_ts)
4385 .await
4386 .unwrap()
4387 .into_element();
4388 assert!(a.is_user());
4389
4390 let mut txn = state.transaction().await.unwrap();
4392 txn.insert_cluster_replica_with_id(cluster_id, a, "explicit", config, owner_id)
4393 .unwrap();
4394 let commit_ts = txn.upper();
4395 txn.commit_internal(commit_ts).await.unwrap();
4396
4397 let commit_ts = state.current_upper().await;
4399 let b = state
4400 .allocate_user_replica_ids(1, commit_ts)
4401 .await
4402 .unwrap()
4403 .into_element();
4404
4405 assert_eq!(b.inner_id(), a.inner_id() + 1);
4408
4409 let txn = state.transaction().await.unwrap();
4411 let found = txn
4412 .get_cluster_replicas()
4413 .any(|replica| replica.replica_id == a);
4414 assert!(found, "explicitly inserted replica {a} not found");
4415 }
4416
4417 #[mz_ore::test]
4418 fn test_allocate_introspection_source_index_id() {
4419 let cluster_variant: u8 = 0b0000_0001;
4420 let cluster_id_inner: u64 =
4421 0b0000_0000_1100_0101_1100_0011_1010_1101_0000_1011_1111_1001_0110_1010;
4422 let timely_messages_received_log_variant: u8 = 0b0000_1000;
4423
4424 let cluster_id = ClusterId::System(cluster_id_inner);
4425 let log_variant = LogVariant::Timely(TimelyLog::MessagesReceived);
4426
4427 let introspection_source_index_id: u64 =
4428 0b0000_0001_1100_0101_1100_0011_1010_1101_0000_1011_1111_1001_0110_1010_0000_1000;
4429
4430 {
4432 let mut cluster_variant_mask = 0xFF << 56;
4433 cluster_variant_mask &= introspection_source_index_id;
4434 cluster_variant_mask >>= 56;
4435 assert_eq!(cluster_variant_mask, u64::from(cluster_variant));
4436 }
4437
4438 {
4440 let mut cluster_id_inner_mask = 0xFFFF_FFFF_FFFF << 8;
4441 cluster_id_inner_mask &= introspection_source_index_id;
4442 cluster_id_inner_mask >>= 8;
4443 assert_eq!(cluster_id_inner_mask, cluster_id_inner);
4444 }
4445
4446 {
4448 let mut log_variant_mask = 0xFF;
4449 log_variant_mask &= introspection_source_index_id;
4450 assert_eq!(
4451 log_variant_mask,
4452 u64::from(timely_messages_received_log_variant)
4453 );
4454 }
4455
4456 let (catalog_item_id, global_id) =
4457 Transaction::allocate_introspection_source_index_id(&cluster_id, log_variant);
4458
4459 assert_eq!(
4460 catalog_item_id,
4461 CatalogItemId::IntrospectionSourceIndex(introspection_source_index_id)
4462 );
4463 assert_eq!(
4464 global_id,
4465 GlobalId::IntrospectionSourceIndex(introspection_source_index_id)
4466 );
4467 }
4468}