adapter: colocate temporary schemas with ephemeral owner registration (SQL-150)

Author: SangJunBak

src/adapter/src/catalog.rs

@@ -56,11 +56,9 @@ use mz_license_keys::ValidatedLicenseKey;
 use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::{EpochMillis, NowFn, SYSTEM_TIME};
 use mz_ore::result::ResultExt as _;
-use mz_ore::soft_assert_or_log;
 use mz_persist_client::PersistClient;
 use mz_repr::adt::mz_acl_item::{AclMode, PrivilegeMap};
 use mz_repr::explain::ExprHumanizer;
-use mz_repr::namespaces::MZ_TEMP_SCHEMA;
 use mz_repr::network_policy_id::NetworkPolicyId;
 use mz_repr::optimize::OptimizerFeatures;
 use mz_repr::role_id::RoleId;
@@ -1164,75 +1162,42 @@ impl Catalog {
         self.state.try_get_role_auth_by_id(id)
     }
 
-    /// Creates a new schema in the `Catalog` for temporary items
-    /// indicated by the TEMPORARY or TEMP keywords.
-    pub fn create_temporary_schema(
-        &mut self,
-        conn_id: &ConnectionId,
-        owner_id: RoleId,
-    ) -> Result<(), Error> {
-        self.state.create_temporary_schema(conn_id, owner_id)
+    /// Registers the connection's temporary namespace: the `uuid` <->
+    /// `conn_id` mapping used to stamp and apply durable temporary items
+    /// owned by the session. The `mz_temp` schema itself materializes when
+    /// the first temporary item is applied.
+    ///
+    /// The coordinator calls this at a session's first temporary-item
+    /// creation, strictly before the transaction that persists the item, and
+    /// guards on [`CatalogState::has_temporary_namespace`], so registering
+    /// an already-registered namespace is a bug.
+    pub fn register_temporary_namespace(&mut self, conn_id: &ConnectionId, uuid: Uuid) {
+        self.state
+            .temporary_namespaces
+            .register(conn_id.clone(), uuid);
     }
 
     fn item_exists_in_temp_schemas(&self, conn_id: &ConnectionId, item_name: &str) -> bool {
-        // Temporary schemas are created lazily, so it's valid for one to not exist yet.
+        // A temporary namespace is registered at the connection's first
+        // temporary-item creation, so it's valid for one to not exist yet.
         self.state
-            .temporary_schemas
-            .get(conn_id)
+            .temporary_namespaces
+            .schema(conn_id)
             .map(|schema| schema.items.contains_key(item_name))
             .unwrap_or(false)
     }
 
-    /// Drops schema for connection if it exists. Returns an error if it exists and has items.
-    /// Returns Ok if conn_id's temp schema does not exist.
-    pub fn drop_temporary_schema(&mut self, conn_id: &ConnectionId) -> Result<(), Error> {
-        let Some(schema) = self.state.temporary_schemas.remove(conn_id) else {
-            return Ok(());
-        };
-        if !schema.items.is_empty() {
-            return Err(Error::new(ErrorKind::SchemaNotEmpty(MZ_TEMP_SCHEMA.into())));
-        }
-        Ok(())
-    }
-
-    /// Registers the session as an ephemeral owner: the `uuid` <-> `conn_id`
-    /// mapping used to stamp and apply durable temporary items owned by the
-    /// session.
-    ///
-    /// The coordinator calls this at a session's first temporary-item
-    /// creation, strictly before the transaction that persists the item, and
-    /// guards on [`CatalogState::is_ephemeral_owner`], so registering an
-    /// already-registered connection is a bug.
-    pub fn register_ephemeral_owner(&mut self, uuid: Uuid, conn_id: ConnectionId) {
-        let prev = self
-            .state
-            .ephemeral_owner_conns_by_uuid
-            .insert(uuid, conn_id.clone());
-        soft_assert_or_log!(
-            prev.is_none(),
-            "duplicate ephemeral owner registration for {uuid}"
-        );
-        self.state
-            .ephemeral_owner_uuids_by_conn
-            .insert(conn_id, uuid);
-    }
-
-    /// Removes the ephemeral-owner registration for `conn_id`.
+    /// Removes the connection's temporary namespace, if it has one. Returns
+    /// Ok if none exists.
     ///
-    /// Callers guard on [`CatalogState::is_ephemeral_owner`], so
-    /// unregistering an unregistered connection is a bug. They must also
-    /// only do this after the transaction dropping the session's temporary
-    /// items has been applied, since applying an ephemeral item update
-    /// resolves the owning connection through this mapping.
-    pub fn unregister_ephemeral_owner(&mut self, conn_id: &ConnectionId) {
-        let uuid = self.state.ephemeral_owner_uuids_by_conn.remove(conn_id);
-        soft_assert_or_log!(
-            uuid.is_some(),
-            "no ephemeral owner registration for {conn_id}"
-        );
-        if let Some(uuid) = uuid {
-            self.state.ephemeral_owner_conns_by_uuid.remove(&uuid);
-        }
+    /// Callers must only do this after the transaction dropping the
+    /// session's temporary items has been applied, since applying an
+    /// ephemeral item update resolves the owning connection through the
+    /// namespace. If the schema still contains items (the drop transaction
+    /// failed, e.g. because this process is being fenced out), this returns
+    /// an error and leaves the namespace fully in place.
+    pub fn drop_temporary_namespace(&mut self, conn_id: &ConnectionId) -> Result<(), Error> {
+        self.state.temporary_namespaces.unregister(conn_id)
     }
 
     pub(crate) fn object_dependents(
@@ -2077,7 +2042,7 @@ impl SessionCatalog for ConnCatalog<'_> {
                 self.state
                     .ambient_schemas_by_id
                     .values()
-                    .chain(self.state.temporary_schemas.values())
+                    .chain(self.state.temporary_namespaces.schemas())
                     .map(|schema| schema as &dyn CatalogSchema),
             )
             .collect()
@@ -2811,6 +2776,10 @@ mod tests {
                 ]
             );
 
+            // A session's temporary namespace is registered at its first
+            // temporary-item creation, so until then `mz_temp` drops out of
+            // the resolved search path and temporary resolution comes only
+            // from the implicit leading temporary schema.
             let mut session = Session::dummy();
             session
                 .vars_mut()
@@ -2822,25 +2791,14 @@ mod tests {
                 )
                 .expect("failed to set search_path");
             let conn_catalog = catalog.for_session(&session);
-            assert_ne!(
-                conn_catalog.effective_search_path(false),
-                conn_catalog.search_path
-            );
-            assert_ne!(
-                conn_catalog.effective_search_path(true),
-                conn_catalog.search_path
-            );
+            assert_eq!(conn_catalog.search_path, vec![]);
             assert_eq!(
                 conn_catalog.effective_search_path(false),
-                vec![
-                    mz_catalog_schema.clone(),
-                    pg_catalog_schema.clone(),
-                    mz_temp_schema.clone()
-                ]
+                vec![mz_catalog_schema.clone(), pg_catalog_schema.clone()]
             );
             assert_eq!(
                 conn_catalog.effective_search_path(true),
-                vec![mz_catalog_schema, pg_catalog_schema, mz_temp_schema]
+                vec![mz_temp_schema, mz_catalog_schema, pg_catalog_schema]
             );
             catalog.expire().await;
         })

src/adapter/src/catalog/apply.rs

@@ -1166,7 +1166,7 @@ impl CatalogState {
         };
         match update.diff {
             // An addition is local when the owning session is connected here.
-            StateDiff::Addition => !self.ephemeral_owner_conns_by_uuid.contains_key(&owner),
+            StateDiff::Addition => !self.temporary_namespaces.contains_uuid(&owner),
             // A retraction must be applied iff the matching addition was
             // applied here, and entry presence records exactly that. A
             // read-only catalog never reaches this arm: it only applies its
@@ -1207,32 +1207,25 @@ impl CatalogState {
                 // session's connection, not in the schema named by
                 // `schema_id` (which is the temporary schema sentinel).
                 // Updates for sessions not connected to this process are
-                // filtered out in `apply_updates_inner`, so the mapping must
-                // exist here.
+                // filtered out in `apply_updates_inner`, so the namespace
+                // (registered before the transaction that created the item)
+                // must exist here. Its schema materializes in `insert_entry`.
                 let conn_id = ephemeral_owner_session.map(|owner| {
-                    self.ephemeral_owner_conns_by_uuid
-                        .get(&owner)
+                    self.temporary_namespaces
+                        .conn_for_uuid(&owner)
                         .cloned()
                         .unwrap_or_else(|| {
                             panic!("no session record applied for temporary item owner {owner}")
                         })
                 });
                 let name = match &conn_id {
-                    Some(conn_id) => {
-                        // Lazily create the temporary schema if it doesn't
-                        // exist yet.
-                        if !self.temporary_schemas.contains_key(conn_id) {
-                            self.create_temporary_schema(conn_id, owner_id)
-                                .expect("failed to create temporary schema");
-                        }
-                        QualifiedItemName {
-                            qualifiers: ItemQualifiers {
-                                database_spec: ResolvedDatabaseSpecifier::Ambient,
-                                schema_spec: SchemaSpecifier::Temporary,
-                            },
-                            item: name.clone(),
-                        }
-                    }
+                    Some(_) => QualifiedItemName {
+                        qualifiers: ItemQualifiers {
+                            database_spec: ResolvedDatabaseSpecifier::Ambient,
+                            schema_spec: SchemaSpecifier::Temporary,
+                        },
+                        item: name.clone(),
+                    },
                     None => {
                         let schema = self.find_non_temp_schema(&schema_id);
                         QualifiedItemName {
@@ -1698,8 +1691,8 @@ impl CatalogState {
         // Keep in sync with `get_schemas`
         match (database_spec, schema_spec) {
             (ResolvedDatabaseSpecifier::Ambient, SchemaSpecifier::Temporary) => self
-                .temporary_schemas
-                .get_mut(conn_id)
+                .temporary_namespaces
+                .schema_mut(conn_id)
                 .expect("catalog out of sync"),
             (ResolvedDatabaseSpecifier::Ambient, SchemaSpecifier::Id(id)) => self
                 .ambient_schemas_by_id
@@ -2004,13 +1997,12 @@ impl CatalogState {
             self.entry_by_global_id.insert(gid, entry.id());
         }
         let conn_id = entry.item().conn_id().unwrap_or(&SYSTEM_CONN_ID);
-        // Lazily create the temporary schema if this is a temporary item and the schema
-        // doesn't exist yet.
-        if entry.name().qualifiers.schema_spec == SchemaSpecifier::Temporary
-            && !self.temporary_schemas.contains_key(conn_id)
-        {
-            self.create_temporary_schema(conn_id, entry.owner_id)
-                .expect("failed to create temporary schema");
+        // Materialize the temporary schema at the first applied temporary
+        // item. The owning session's namespace was registered before the
+        // transaction that created the item.
+        if entry.name().qualifiers.schema_spec == SchemaSpecifier::Temporary {
+            self.temporary_namespaces
+                .ensure_schema(conn_id, entry.owner_id);
         }
         let schema = self.get_schema_mut(
             &entry.name().qualifiers.database_spec,

src/adapter/src/catalog/open.rs

@@ -25,7 +25,6 @@ use mz_audit_log::{
     CreateOrDropClusterReplicaReasonV1, EventDetails, EventType, ObjectType, VersionedEvent,
 };
 use mz_auth::hash::scram256_hash;
-use mz_catalog::SYSTEM_CONN_ID;
 use mz_catalog::builtin::{
     BUILTIN_CLUSTERS, BUILTIN_PREFIXES, BUILTIN_ROLES, BUILTINS, Builtin, Fingerprint,
     MZ_CATALOG_RAW, RUNTIME_ALTERABLE_FINGERPRINT_SENTINEL,
@@ -155,9 +154,7 @@ impl Catalog {
             comments: Arc::new(CommentsMap::default()),
             source_references: imbl::OrdMap::new(),
             storage_metadata: Arc::new(StorageMetadata::default()),
-            temporary_schemas: imbl::OrdMap::new(),
-            ephemeral_owner_conns_by_uuid: imbl::OrdMap::new(),
-            ephemeral_owner_uuids_by_conn: imbl::OrdMap::new(),
+            temporary_namespaces: Default::default(),
             mock_authentication_nonce: Default::default(),
             config: mz_sql::catalog::CatalogConfig {
                 start_time: to_datetime((config.now)()),
@@ -251,7 +248,6 @@ impl Catalog {
                     Err(e) => return Err(e.into()),
                 };
             }
-            state.create_temporary_schema(&SYSTEM_CONN_ID, MZ_SYSTEM_ROLE_ID)?;
         }
 
         // Make life easier by consolidating all updates, so that we end up with only positive