Skip to main content

mz_catalog/durable/
transaction.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use 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/// A [`Transaction`] batches multiple catalog operations together and commits them atomically.
78/// An operation also logically groups multiple catalog updates together.
79#[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    // Don't make this a table transaction so that it's not read into the
113    // in-memory cache.
114    audit_log_updates: Vec<(AuditLogKey, Diff, Timestamp)>,
115    /// The upper of `durable_catalog` at the start of the transaction.
116    upper: mz_repr::Timestamp,
117    /// The ID of the current operation of this transaction.
118    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        // For these collections uniqueness is plain equality of the name fields, so the same
152        // predicate answers both "do these two conflict?" and "did this update keep the same key?".
153        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                        // `item_type` is slow, only compute if needed.
181                        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                        // `item_type` is slow, only compute it once name and schema match.
192                        && 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            // Uniqueness violations for this value occur at the key rather than
227            // the value (the key is the unit struct `()` so this is a singleton
228            // value).
229            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                    // This should never happen, but rather than panicking we'll
408                    // set a known secure value as a fallback.
409                    .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    /// Panics if any introspection source id is not a system id
444    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    /// Panics if any introspection source id is not a system id
466    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    /// Updates persisted information about persisted introspection source
665    /// indexes.
666    ///
667    /// Panics if provided id is not a system id.
668    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            // TODO(alter_table): Use separate ID allocators.
801            .map(|x| (CatalogItemId::System(x), GlobalId::System(x)))
802            .collect())
803    }
804
805    /// Allocates an ID for an introspection source index. These IDs are deterministically derived
806    /// from the `cluster_id` and `log_variant`.
807    ///
808    /// Introspection source indexes are a special edge case of items. They are considered system
809    /// items, but they are the only system item that can be created by the user at any time. All
810    /// other system items can only be created by the system during the startup of an upgrade.
811    ///
812    /// Furthermore, all other system item IDs are allocated deterministically in the same order
813    /// during startup. Therefore, all read-only `environmentd` processes during an upgrade will
814    /// allocate the same system IDs to the same items, and due to the way catalog fencing works,
815    /// only one of them can successfully write the IDs down to the catalog. This removes the need
816    /// for `environmentd` processes to coordinate system IDs allocated during read-only mode.
817    ///
818    /// Since introspection IDs can be allocated at any time, read-only instances would either need
819    /// to coordinate across processes when allocating a new ID or allocate them deterministically.
820    /// We opted to allocate the IDs deterministically to avoid the overhead of coordination.
821    ///
822    /// Introspection source index IDs are 64 bit integers, with the following format (not to
823    /// scale):
824    ///
825    /// -------------------------------------------------------------
826    /// | Cluster ID Variant | Cluster ID Inner Value | Log Variant |
827    /// |--------------------|------------------------|-------------|
828    /// |       8-bits       |         48-bits        |   8-bits    |
829    /// -------------------------------------------------------------
830    ///
831    /// Cluster ID Variant:      A unique number indicating the variant of cluster the index belongs
832    ///                          to.
833    /// Cluster ID Inner Value:  A per variant unique number indicating the cluster the index
834    ///                          belongs to.
835    /// Log Variant:             A unique number indicating the log variant this index is on.
836    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            // TODO(alter_table): Use separate ID allocators.
904            .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    /// Allocates `amount` OIDs. OIDs can be recycled if they aren't currently assigned to any
918    /// object.
919    #[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 representing an OID for a user object. Allocated OIDs can be recycled, so when we've
926        /// allocated [`u32::MAX`] we'll wrap back around to [`FIRST_USER_OID`].
927        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        // This is potentially slow to do everytime we allocate an OID. A faster approach might be
951        // to have an ID allocator that is updated everytime an OID is allocated or de-allocated.
952        // However, benchmarking shows that this doesn't make a noticeable difference and the other
953        // approach requires making sure that allocator always stays in-sync which can be
954        // error-prone. If DDL starts slowing down, this is a good place to try and optimize.
955        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                // We've exhausted all possible OIDs and still don't have `amount`.
1002                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    /// Allocates a single OID. OIDs can be recycled if they aren't currently assigned to any
1027    /// object.
1028    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    /// Exports the current state of this transaction as a [`Snapshot`].
1034    ///
1035    /// This merges each `TableTransaction`'s initial data with its pending
1036    /// changes to produce the current view, then converts back to proto types.
1037    /// Used to persist transaction state between incremental DDL dry runs so
1038    /// the next dry run's fresh `Transaction` starts in sync with the
1039    /// accumulated `CatalogState`.
1040    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    /// Removes the database `id` from the transaction.
1084    ///
1085    /// Returns an error if `id` is not found.
1086    ///
1087    /// Runtime is linear with respect to the total number of databases in the catalog.
1088    /// DO NOT call this function in a loop, use [`Self::remove_databases`] instead.
1089    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    /// Removes all databases in `databases` from the transaction.
1101    ///
1102    /// Returns an error if any id in `databases` is not found.
1103    ///
1104    /// NOTE: On error, there still may be some databases removed from the transaction. It
1105    /// is up to the caller to either abort the transaction or commit.
1106    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    /// Removes the schema identified by `database_id` and `schema_id` from the transaction.
1130    ///
1131    /// Returns an error if `(database_id, schema_id)` is not found.
1132    ///
1133    /// Runtime is linear with respect to the total number of schemas in the catalog.
1134    /// DO NOT call this function in a loop, use [`Self::remove_schemas`] instead.
1135    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    /// Removes all schemas in `schemas` from the transaction.
1155    ///
1156    /// Returns an error if any id in `schemas` is not found.
1157    ///
1158    /// NOTE: On error, there still may be some schemas removed from the transaction. It
1159    /// is up to the caller to either abort the transaction or commit.
1160    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    /// Removes all user roles in `roles` from the transaction.
1210    ///
1211    /// Returns an error if any id in `roles` is not found.
1212    ///
1213    /// NOTE: On error, there still may be some roles removed from the transaction. It
1214    /// is up to the caller to either abort the transaction or commit.
1215    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    /// Removes all roles in `roles` from the transaction.
1224    ///
1225    /// Returns an error if any id in `roles` is not found.
1226    ///
1227    /// NOTE: On error, there still may be some roles removed from the transaction. It
1228    /// is up to the caller to either abort the transaction or commit.
1229    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        // The reason we don't to the same check as above is that the role auth table
1268        // is not required to have all roles in the role table.
1269
1270        Ok(())
1271    }
1272
1273    /// Removes all cluster in `clusters` from the transaction.
1274    ///
1275    /// Returns an error if any id in `clusters` is not found.
1276    ///
1277    /// NOTE: On error, there still may be some clusters removed from the transaction. It is up to
1278    /// the caller to either abort the transaction or commit.
1279    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        // Cascade delete introspection sources and cluster replicas.
1297        //
1298        // TODO(benesch): this doesn't seem right. Cascade deletions should
1299        // be entirely the domain of the higher catalog layer, not the
1300        // storage layer.
1301        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    /// Removes the cluster replica `id` from the transaction.
1310    ///
1311    /// Returns an error if `id` is not found.
1312    ///
1313    /// Runtime is linear with respect to the total number of cluster replicas in the catalog.
1314    /// DO NOT call this function in a loop, use [`Self::remove_cluster_replicas`] instead.
1315    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    /// Removes all cluster replicas in `replicas` from the transaction.
1328    ///
1329    /// Returns an error if any id in `replicas` is not found.
1330    ///
1331    /// NOTE: On error, there still may be some cluster replicas removed from the transaction. It
1332    /// is up to the caller to either abort the transaction or commit.
1333    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    /// Removes item `id` from the transaction.
1357    ///
1358    /// Returns an error if `id` is not found.
1359    ///
1360    /// Runtime is linear with respect to the total number of items in the catalog.
1361    /// DO NOT call this function in a loop, use [`Self::remove_items`] instead.
1362    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    /// Removes all items in `ids` from the transaction.
1372    ///
1373    /// Returns an error if any id in `ids` is not found.
1374    ///
1375    /// NOTE: On error, there still may be some items removed from the transaction. It is
1376    /// up to the caller to either abort the transaction or commit.
1377    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    /// Removes all system object mappings in `descriptions` from the transaction.
1394    ///
1395    /// Returns an error if any description in `descriptions` is not found.
1396    ///
1397    /// NOTE: On error, there still may be some items removed from the transaction. It is
1398    /// up to the caller to either abort the transaction or commit.
1399    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    /// Removes all introspection source indexes in `indexes` from the transaction.
1442    ///
1443    /// Returns an error if any index in `indexes` is not found.
1444    ///
1445    /// NOTE: On error, there still may be some indexes removed from the transaction. It is
1446    /// up to the caller to either abort the transaction or commit.
1447    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    /// Updates item `id` in the transaction to `item_name` and `item`.
1481    ///
1482    /// Returns an error if `id` is not found.
1483    ///
1484    /// Runtime is linear with respect to the total number of items in the catalog.
1485    /// DO NOT call this function in a loop, use [`Self::update_items`] instead.
1486    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    /// Updates all items with ids matching the keys of `items` in the transaction, to the
1498    /// corresponding value in `items`.
1499    ///
1500    /// Returns an error if any id in `items` is not found.
1501    ///
1502    /// NOTE: On error, there still may be some items updated in the transaction. It is
1503    /// up to the caller to either abort the transaction or commit.
1504    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    /// Updates role `id` in the transaction to `role`.
1530    ///
1531    /// Returns an error if `id` is not found.
1532    ///
1533    /// Runtime is linear with respect to the total number of items in the catalog.
1534    /// DO NOT call this function in a loop, implement and use some `Self::update_roles` instead.
1535    /// You should model it after [`Self::update_items`].
1536    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    /// Updates all [`Role`]s with ids matching the keys of `roles` in the transaction, to the
1588    /// corresponding value in `roles`.
1589    ///
1590    /// This function does *not* write role_authentication information to the catalog.
1591    /// It is purely for updating the role itself.
1592    ///
1593    /// Returns an error if any id in `roles` is not found.
1594    ///
1595    /// NOTE: On error, there still may be some roles updated in the transaction. It is
1596    /// up to the caller to either abort the transaction or commit.
1597    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    /// Updates persisted mapping from system objects to global IDs and fingerprints. Each element
1623    /// of `mappings` should be (old-global-id, new-system-object-mapping).
1624    ///
1625    /// Panics if provided id is not a system id.
1626    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    /// Updates cluster `id` in the transaction to `cluster`.
1657    ///
1658    /// Returns an error if `id` is not found.
1659    ///
1660    /// Runtime is linear with respect to the total number of clusters in the catalog.
1661    /// DO NOT call this function in a loop.
1662    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    /// Updates cluster replica `replica_id` in the transaction to `replica`.
1676    ///
1677    /// Returns an error if `replica_id` is not found.
1678    ///
1679    /// Runtime is linear with respect to the total number of cluster replicas in the catalog.
1680    /// DO NOT call this function in a loop.
1681    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    /// Updates database `id` in the transaction to `database`.
1699    ///
1700    /// Returns an error if `id` is not found.
1701    ///
1702    /// Runtime is linear with respect to the total number of databases in the catalog.
1703    /// DO NOT call this function in a loop.
1704    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    /// Updates schema `schema_id` in the transaction to `schema`.
1722    ///
1723    /// Returns an error if `schema_id` is not found.
1724    ///
1725    /// Runtime is linear with respect to the total number of schemas in the catalog.
1726    /// DO NOT call this function in a loop.
1727    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    /// Updates `network_policy_id` in the transaction to `network policy`.
1745    ///
1746    /// Returns an error if `id` is not found.
1747    ///
1748    /// Runtime is linear with respect to the total number of databases in the catalog.
1749    /// DO NOT call this function in a loop.
1750    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    /// Removes all network policies in `network policies` from the transaction.
1767    ///
1768    /// Returns an error if any id in `network policy` is not found.
1769    ///
1770    /// NOTE: On error, there still may be some roles removed from the transaction. It
1771    /// is up to the caller to either abort the transaction or commit.
1772    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    /// Set persisted default privilege.
1799    ///
1800    /// DO NOT call this function in a loop, use [`Self::set_default_privileges`] instead.
1801    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    /// Set persisted default privileges.
1825    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    /// Set persisted system privilege.
1844    ///
1845    /// DO NOT call this function in a loop, use [`Self::set_system_privileges`] instead.
1846    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    /// Set persisted system privileges.
1861    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    /// Set persisted setting.
1880    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    /// Insert persisted introspection source index.
1894    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    /// Set persisted system object mappings.
1928    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    /// Set persisted replica.
1946    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    /// Set persisted configuration.
1961    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    /// Get the value of a persisted config.
1976    pub fn get_config(&self, key: String) -> Option<u64> {
1977        self.configs
1978            .get(&ConfigKey { key })
1979            .map(|entry| entry.value)
1980    }
1981
1982    /// Get the value of a persisted setting.
1983    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    /// Updates the catalog `with_0dt_deployment_max_wait` "config" value to
2014    /// match the `with_0dt_deployment_max_wait` "system var" value.
2015    ///
2016    /// These are mirrored so that we can toggle the flag with Launch Darkly,
2017    /// but use it in boot before Launch Darkly is available.
2018    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    /// Updates the catalog `with_0dt_deployment_ddl_check_interval` "config"
2031    /// value to match the `with_0dt_deployment_ddl_check_interval` "system var"
2032    /// value.
2033    ///
2034    /// These are mirrored so that we can toggle the flag with Launch Darkly,
2035    /// but use it in boot before Launch Darkly is available.
2036    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    /// Updates the catalog `0dt_deployment_panic_after_timeout` "config" value to
2052    /// match the `0dt_deployment_panic_after_timeout` "system var" value.
2053    ///
2054    /// These are mirrored so that we can toggle the flag with Launch Darkly,
2055    /// but use it in boot before Launch Darkly is available.
2056    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    /// Removes the catalog `with_0dt_deployment_max_wait` "config" value to
2067    /// match the `with_0dt_deployment_max_wait` "system var" value.
2068    ///
2069    /// These are mirrored so that we can toggle the flag with LaunchDarkly,
2070    /// but use it in boot before LaunchDarkly is available.
2071    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    /// Removes the catalog `with_0dt_deployment_ddl_check_interval` "config"
2076    /// value to match the `with_0dt_deployment_ddl_check_interval` "system var"
2077    /// value.
2078    ///
2079    /// These are mirrored so that we can toggle the flag with LaunchDarkly, but
2080    /// use it in boot before LaunchDarkly is available.
2081    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    /// Removes the catalog `enable_0dt_deployment_panic_after_timeout` "config"
2086    /// value to match the `enable_0dt_deployment_panic_after_timeout` "system
2087    /// var" value.
2088    ///
2089    /// These are mirrored so that we can toggle the flag with LaunchDarkly, but
2090    /// use it in boot before LaunchDarkly is available.
2091    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    /// Updates the catalog `system_config_synced` "config" value to true.
2096    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    /// Upserts persisted system configuration `name` to `value`.
2145    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    /// Removes persisted system configuration `name`.
2156    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    /// Removes all persisted system configurations.
2166    pub fn clear_system_configs(&mut self) {
2167        self.system_configurations.delete(|_k, _v| true, self.op_id);
2168    }
2169
2170    /// Returns the persisted cluster-coherent scoped system configurations.
2171    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    /// Upserts the persisted cluster-coherent scoped configuration `name` = `value`
2181    /// for `cluster_id`.
2182    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    /// Removes the persisted cluster-coherent scoped configuration `name` for
2199    /// `cluster_id`.
2200    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    /// Returns the persisted replica-local scoped system configurations.
2211    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    /// Upserts the persisted replica-local scoped configuration `name` = `value`
2221    /// for `replica_id`.
2222    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    /// Removes the persisted replica-local scoped configuration `name` for
2239    /// `replica_id`.
2240    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    /// Commit the current operation within the transaction. This does not cause anything to be
2355    /// written durably, but signals to the current transaction that we are moving on to the next
2356    /// operation.
2357    ///
2358    /// Returns the updates of the committed operation.
2359    #[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            // Not representable as a `StateUpdate`.
2440            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    /// Commits the storage transaction to durable storage. Any error returned outside read-only
2613    /// mode indicates the catalog may be in an indeterminate state and needs to be fully re-read
2614    /// before proceeding. In general, this must be fatal to the calling process. We do not
2615    /// panic/halt inside this function itself so that errors can bubble up during initialization.
2616    ///
2617    /// The transaction is committed at `commit_ts`.
2618    ///
2619    /// Returns what the upper was directly after the transaction committed.
2620    ///
2621    /// In read-only mode, this will return an error for non-empty transactions indicating that the
2622    /// catalog is not writeable.
2623    #[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        // Consolidate in memory because it will likely be faster than consolidating after the
2657        // transaction has been made durable.
2658        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    /// Commits the storage transaction to durable storage. Any error returned outside read-only
2690    /// mode indicates the catalog may be in an indeterminate state and needs to be fully re-read
2691    /// before proceeding. In general, this must be fatal to the calling process. We do not
2692    /// panic/halt inside this function itself so that errors can bubble up during initialization.
2693    ///
2694    /// In read-only mode, this will return an error for non-empty transactions indicating that the
2695    /// catalog is not writeable.
2696    ///
2697    /// IMPORTANT: It is assumed that the committer of this transaction has already applied all
2698    /// updates from this transaction. Therefore, updates from this transaction will not be returned
2699    /// when calling [`crate::durable::ReadOnlyDurableCatalogState::sync_to_current_updates`] or
2700    /// [`crate::durable::ReadOnlyDurableCatalogState::sync_updates`].
2701    ///
2702    /// An alternative implementation would be for the caller to explicitly consume their updates
2703    /// after committing and only then apply the updates in-memory. While this removes assumptions
2704    /// about the caller in this method, in practice it results in duplicate work on every commit.
2705    #[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        // Drain all the updates from the commit since it is assumed that they were already applied.
2715        let updates = durable_storage.sync_updates(upper).await?;
2716        // Writable and savepoint catalogs should have consumed all updates before committing a
2717        // transaction, otherwise the commit was performed with an out of date state.
2718        // Read-only catalogs can only commit empty transactions, so they don't need to consume all
2719        // updates before committing.
2720        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                // Inserting duplicate keys has no effect.
2805                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/// Describes a set of changes to apply as the result of a catalog transaction.
2839#[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    /// The upper of the catalog when the transaction started.
2898    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
2964/// Utility trait to check for plan validity.
2965trait UniqueName {
2966    /// Does the item have a unique name? If yes, we can check for name equality in validity
2967    /// checking.
2968    const HAS_UNIQUE_NAME: bool;
2969    /// The unique name, only returns a meaningful name if [`Self::HAS_UNIQUE_NAME`] is `true`.
2970    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/// A collection's uniqueness constraint.
3038///
3039/// `violation` reports whether two values conflict, i.e. cannot both exist.
3040///
3041/// `is_unique_key_unchanged_after_update` reports whether an update left
3042/// unchanged every field `violation` reads. It is used as an optimization
3043/// since `violation` is an expensive operation to run.
3044#[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/// TableTransaction emulates some features of a typical SQL transaction over
3051/// table for a Collection.
3052///
3053/// It supports:
3054/// - uniqueness constraints
3055/// - transactional reads and writes (including read-your-writes before commit)
3056///
3057/// `K` is the primary key type. Multiple entries with the same key are disallowed.
3058/// `V` is the an arbitrary value type.
3059#[derive(Debug)]
3060struct TableTransaction<K, V> {
3061    initial: BTreeMap<K, V>,
3062    // The desired updates to keys after commit.
3063    // Invariant: Value is sorted by `ts`.
3064    pending: BTreeMap<K, Vec<TransactionUpdate<V>>>,
3065    // `None` for collections with no uniqueness constraint.
3066    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    /// Create a new TableTransaction with initial data.
3075    ///
3076    /// Internally the catalog serializes data as protobuf. All fields in a proto message are
3077    /// optional, which makes using them in Rust cumbersome. Generic parameters `KP` and `VP` are
3078    /// protobuf types which deserialize to `K` and `V` that a [`TableTransaction`] is generic
3079    /// over.
3080    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    /// Like [`Self::new`], but with the collection's uniqueness constraint.
3098    /// See [`UniquenessCheck`] for more details on the uniqueness constraint.
3099    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    /// Consumes and returns the pending changes and their diffs. `Diff` is
3124    /// guaranteed to be 1 or -1.
3125    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        // Pending describes the desired final state for some keys. K,V pairs should be
3132        // retracted if they already exist and were deleted or are being updated.
3133        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    /// Verifies that no items in `self` violate `self.uniqueness_check`.
3148    ///
3149    /// Runtime is O(n^2), where n is the number of items in `self`, if
3150    /// [`UniqueName::HAS_UNIQUE_NAME`] is false for `V`. Prefer using [`Self::verify_keys`].
3151    fn verify(&self) -> Result<(), DurableCatalogError> {
3152        if let Some(check) = &self.uniqueness_check {
3153            // Compare each value to each other value and ensure they are unique.
3154            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    /// Verifies that no items in `self` violate `self.uniqueness_check` with `keys`.
3189    ///
3190    /// Runtime is O(n * k), where n is the number of items in `self` and k is the number of
3191    /// items in `keys`.
3192    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            // Compare each value in `entries` to each value in `self` and ensure they are unique.
3205            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    /// Iterates over the items viewable in the current transaction in arbitrary
3218    /// order and applies `f` on all key, value pairs.
3219    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            // Deleted items don't exist so shouldn't be visited, but still suppress
3225            // visiting the key later.
3226            if let Some(v) = v {
3227                f(k, v);
3228            }
3229        }
3230        for (k, v) in self.initial.iter() {
3231            // Add on initial items that don't have updates.
3232            if !seen.contains(k) {
3233                f(k, v);
3234            }
3235        }
3236    }
3237
3238    /// Returns the current value of `k`.
3239    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    /// Returns the items viewable in the current transaction. The items are
3257    /// cloned, so this is an expensive operation. Prefer using [`Self::items`], or
3258    /// [`Self::for_values`].
3259    // Used by tests.
3260    #[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    /// Returns the current items as proto-typed key-value pairs, suitable for
3270    /// constructing a [`Snapshot`]. This merges `initial` and `pending` to
3271    /// produce the current view and converts back to proto types.
3272    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    /// Returns the items viewable in the current transaction as references. Returns a map
3286    /// of references.
3287    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    /// Returns the values viewable in the current transaction as references.
3296    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    /// Returns the number of items viewable in the current transaction.
3305    fn len(&self) -> usize {
3306        let mut count = 0;
3307        self.for_values(|_, _| {
3308            count += 1;
3309        });
3310        count
3311    }
3312
3313    /// Iterates over the items viewable in the current transaction, and provides a
3314    /// map where additional pending items can be inserted, which will be appended
3315    /// to current pending items. Does not verify uniqueness.
3316    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    /// Inserts a new k,v pair.
3328    ///
3329    /// Returns an error if the uniqueness check failed or the key already exists.
3330    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    /// Updates k, v pairs. `f` is a function that can return `Some(V)` if the
3356    /// value should be updated, otherwise `None`. Returns the number of changed
3357    /// entries.
3358    ///
3359    /// Returns an error if the uniqueness check failed.
3360    ///
3361    /// Prefer using [`Self::update_by_key`] or [`Self::update_by_keys`], which generally have
3362    /// better performance.
3363    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        // Keep a copy of pending in case of uniqueness violation.
3371        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        // Check for uniqueness violation.
3390        if let Err(err) = self.verify_keys(&keys) {
3391            self.pending = pending;
3392            Err(err)
3393        } else {
3394            Ok(changed)
3395        }
3396    }
3397
3398    /// Updates `k`, `v` pair if `k` already exists in `self`.
3399    ///
3400    /// Returns `true` if `k` was updated, `false` otherwise.
3401    /// Returns an error if the uniqueness check failed.
3402    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    /// Updates k, v pairs. Keys that don't already exist in `self` are ignored.
3414    ///
3415    /// Returns the number of changed entries.
3416    /// Returns an error if the uniqueness check failed.
3417    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                // Record if updating this entry would be a no-op.
3426                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 out no-ops to save some work.
3436            .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    /// Whether changing a key's value from `prev` to `next` can introduce a uniqueness violation
3444    fn update_needs_uniqueness_check(&self, prev: Option<&V>, next: Option<&V>) -> bool {
3445        match (&self.uniqueness_check, prev, next) {
3446            // No uniqueness constraint, or a delete: nothing to check.
3447            (None, _, _) | (_, _, None) => false,
3448            // An update: a scan is only needed if it changed the unique key.
3449            (Some(check), Some(prev), Some(next)) => {
3450                !(check.is_unique_key_unchanged_after_update)(prev, next)
3451            }
3452            // An insert: must check.
3453            (Some(_), None, Some(_)) => true,
3454        }
3455    }
3456
3457    /// Set the value for a key. Returns the previous entry if the key existed,
3458    /// otherwise None.
3459    ///
3460    /// Returns an error if the uniqueness check failed.
3461    ///
3462    /// DO NOT call this function in a loop, use [`Self::set_many`] instead.
3463    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        // Check for uniqueness violation
3500        if needs_uniqueness_check {
3501            if let Err(err) = self.verify_keys([&k]) {
3502                // Revert self.pending to the state it was in before calling this
3503                // function.
3504                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    /// Set the values for many keys. Returns the previous entry for each key if the key existed,
3513    /// otherwise None.
3514    ///
3515    /// Returns an error if any uniqueness check failed.
3516    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        // Only the keys whose update can introduce a uniqueness violation need scanning.
3528        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        // Check for uniqueness violation.
3572        if let Err(err) = self.verify_keys(keys_to_verify_uniqueness.iter()) {
3573            for (k, restore_len) in restores {
3574                // Revert self.pending to the state it was in before calling this
3575                // function.
3576                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    /// Deletes items for which `f` returns true. Returns the keys and values of
3586    /// the deleted entries.
3587    ///
3588    /// Prefer using [`Self::delete_by_key`] or [`Self::delete_by_keys`], which generally have
3589    /// better performance.
3590    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    /// Deletes item with key `k`.
3607    ///
3608    /// Returns the value of the deleted entry, if it existed.
3609    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    /// Deletes items with key in `ks`.
3615    ///
3616    /// Returns the keys and values of the deleted entries.
3617    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        // Ideally, we compare for errors here, but it's hard/impossible to implement PartialEq
3659        // for DurableCatalogError.
3660        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        // Panics if the uniqueness scan ever runs
3681        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            // We treat the first character as the unique key.
3691            first_char_same,
3692        )
3693        .unwrap();
3694        // Same first char, so the unique key didn't change and the update must not scan
3695        // for uniqueness.
3696        assert!(
3697            table
3698                .update_by_key(1i64.to_le_bytes().to_vec(), "a2".to_string(), 0)
3699                .unwrap()
3700        );
3701
3702        // An update that changes the unique key must still run the scan and detect a collision.
3703        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        // Changing key 1's first char from 'a' to 'b' changes its unique key and collides
3716        // with key 2.
3717        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            // Sort by diff so that we process retractions first.
3736            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        // Uniqueness violation.
3773        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        // Deleting then creating an item that has a uniqueness violation should work.
3828        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        // Uniqueness violation in value.
3836        table_txn
3837            .insert(5i64.to_le_bytes().to_vec(), "v3".to_string(), 0)
3838            .unwrap_err();
3839        // Key already exists, expect error.
3840        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        // Both the inserts work now because the key and uniqueness violation are gone.
3848        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        // Verify we don't try to delete v3 or v4 during commit.
3911        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        // Test `set`.
3935        let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3936            table.clone(),
3937            uniqueness_violation,
3938            uniqueness_violation,
3939        )
3940        .unwrap();
3941        // Uniqueness violation.
3942        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        // Duplicate `set`.
3969        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        // Test `set_many`.
3982        let mut table_txn = TableTransaction::new_with_uniqueness_fn(
3983            table.clone(),
3984            uniqueness_violation,
3985            uniqueness_violation,
3986        )
3987        .unwrap();
3988        // Uniqueness violation.
3989        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        // Duplicate `set_many`.
4039        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        // Test `update_by_key`
4066        let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4067            table.clone(),
4068            uniqueness_violation,
4069            uniqueness_violation,
4070        )
4071        .unwrap();
4072        // Uniqueness violation.
4073        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        // Duplicate `update_by_key`.
4108        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        // Test `update_by_keys`
4131        let mut table_txn = TableTransaction::new_with_uniqueness_fn(
4132            table.clone(),
4133            uniqueness_violation,
4134            uniqueness_violation,
4135        )
4136        .unwrap();
4137        // Uniqueness violation.
4138        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        // Duplicate `update_by_keys`.
4189        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        // Test `delete_by_key`
4217        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        // Test `delete_by_keys`
4245        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)] //  unsupported operation: can't call foreign function `TLS_client_method` on OS `linux`
4284    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        // Initialize catalog.
4297        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    /// Regression test for DB-147: inserting a replica with an explicit id must not consume the
4341    /// `IdAlloc` counter, and the durable allocator must advance independently so a later
4342    /// allocation never collides with an explicitly inserted id.
4343    #[mz_ore::test(tokio::test)]
4344    #[cfg_attr(miri, ignore)] //  unsupported operation: can't call foreign function `TLS_client_method` on OS `linux`
4345    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        // The cluster does not need to exist: `insert_cluster_replica_with_id` only writes a
4364        // `cluster_replicas` row and does not validate the referenced cluster.
4365        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        // Step 1: allocate one user replica id out-of-band via the durable allocator.
4382        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        // Step 2: insert a replica with that explicit id and commit.
4391        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        // Step 3: allocate one more user replica id.
4398        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        // Step 4: the explicit insert consumed zero ids, so the allocator's next value is the
4406        // direct successor of the first allocation.
4407        assert_eq!(b.inner_id(), a.inner_id() + 1);
4408
4409        // Step 5: the explicitly inserted replica is present in the committed state.
4410        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        // Sanity check that `introspection_source_index_id` contains `cluster_variant`.
4431        {
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        // Sanity check that `introspection_source_index_id` contains `cluster_id_inner`.
4439        {
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        // Sanity check that `introspection_source_index_id` contains `timely_messages_received_log_variant`.
4447        {
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}