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