risingwave_frontend/catalog/

root_catalog.rs

// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use itertools::Itertools;
use risingwave_common::catalog::{FunctionId, IndexId, StreamJobStatus, TableId};
use risingwave_common::id::ObjectId;
use risingwave_common::session_config::{SearchPath, USER_NAME_WILD_CARD};
use risingwave_common::types::DataType;
use risingwave_connector::sink::catalog::SinkCatalog;
use risingwave_pb::catalog::{
    PbConnection, PbDatabase, PbFunction, PbIndex, PbSchema, PbSecret, PbSink, PbSource,
    PbSubscription, PbTable, PbView,
};
use risingwave_pb::common::PbObjectType;
use risingwave_pb::hummock::HummockVersionStats;
use risingwave_pb::meta::ObjectDependency as PbObjectDependency;

use super::function_catalog::FunctionCatalog;
use super::source_catalog::SourceCatalog;
use super::subscription_catalog::{SubscriptionCatalog, SubscriptionState};
use super::view_catalog::ViewCatalog;
use super::{
    CatalogError, CatalogResult, ConnectionId, SecretId, SinkId, SourceId, SubscriptionId, ViewId,
};
use crate::catalog::connection_catalog::ConnectionCatalog;
use crate::catalog::database_catalog::DatabaseCatalog;
use crate::catalog::schema_catalog::SchemaCatalog;
use crate::catalog::secret_catalog::SecretCatalog;
use crate::catalog::system_catalog::{
    SystemTableCatalog, get_sys_tables_in_schema, get_sys_views_in_schema,
};
use crate::catalog::table_catalog::TableCatalog;
use crate::catalog::{DatabaseId, IndexCatalog, SchemaId};
use crate::expr::{Expr, ExprImpl};

#[derive(Copy, Clone)]
pub enum SchemaPath<'a> {
    Name(&'a str),
    /// (`search_path`, `user_name`).
    Path(&'a SearchPath, &'a str),
}

impl<'a> SchemaPath<'a> {
    pub fn new(
        schema_name: Option<&'a str>,
        search_path: &'a SearchPath,
        user_name: &'a str,
    ) -> Self {
        match schema_name {
            Some(schema_name) => SchemaPath::Name(schema_name),
            None => SchemaPath::Path(search_path, user_name),
        }
    }

    /// Call function `f` for each schema name. Return the first `Some` result.
    pub fn try_find<T, E>(
        &self,
        mut f: impl FnMut(&str) -> Result<Option<T>, E>,
    ) -> Result<Option<(T, &'a str)>, E> {
        match self {
            SchemaPath::Name(schema_name) => Ok(f(schema_name)?.map(|t| (t, *schema_name))),
            SchemaPath::Path(search_path, user_name) => {
                for schema_name in search_path.path() {
                    let mut schema_name: &str = schema_name;
                    if schema_name == USER_NAME_WILD_CARD {
                        schema_name = user_name;
                    }
                    if let Ok(Some(res)) = f(schema_name) {
                        return Ok(Some((res, schema_name)));
                    }
                }
                Ok(None)
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct ObjectDependency {
    pub object_id: ObjectId,
    pub referenced_object_id: ObjectId,
    pub referenced_object_type: PbObjectType,
}

impl From<PbObjectDependency> for ObjectDependency {
    fn from(dependency: PbObjectDependency) -> Self {
        let referenced_object_type = PbObjectType::try_from(dependency.referenced_object_type)
            .unwrap_or(PbObjectType::Unspecified);
        Self {
            object_id: dependency.object_id,
            referenced_object_id: dependency.referenced_object_id,
            referenced_object_type,
        }
    }
}

/// Root catalog of database catalog. It manages all database/schema/table in memory on frontend.
/// It is protected by a `RwLock`. Only [`crate::observer::FrontendObserverNode`]
/// will acquire the write lock and sync it with the meta catalog. In other situations, it is
/// read only.
///
/// - catalog (root catalog)
///   - database catalog
///     - schema catalog
///       - function catalog (i.e., user defined function)
///       - table/sink/source/index/view catalog
///        - column catalog
pub struct Catalog {
    database_by_name: HashMap<String, DatabaseCatalog>,
    db_name_by_id: HashMap<DatabaseId, String>,
    /// all table catalogs in the cluster identified by universal unique table id.
    table_by_id: HashMap<TableId, Arc<TableCatalog>>,
    table_stats: HummockVersionStats,
    object_dependencies: HashMap<ObjectId, Vec<ObjectDependency>>,
}

#[expect(clippy::derivable_impls)]
impl Default for Catalog {
    fn default() -> Self {
        Self {
            database_by_name: HashMap::new(),
            db_name_by_id: HashMap::new(),
            table_by_id: HashMap::new(),
            table_stats: HummockVersionStats::default(),
            object_dependencies: HashMap::new(),
        }
    }
}

impl Catalog {
    fn get_database_mut(&mut self, db_id: DatabaseId) -> Option<&mut DatabaseCatalog> {
        let name = self.db_name_by_id.get(&db_id)?;
        self.database_by_name.get_mut(name)
    }

    pub fn clear(&mut self) {
        self.database_by_name.clear();
        self.db_name_by_id.clear();
        self.table_by_id.clear();
        self.object_dependencies.clear();
    }

    pub fn create_database(&mut self, db: &PbDatabase) {
        let name = db.name.clone();
        let id = db.id;

        self.database_by_name
            .try_insert(name.clone(), db.into())
            .unwrap();
        self.db_name_by_id.try_insert(id, name).unwrap();
    }

    pub fn create_schema(&mut self, proto: &PbSchema) {
        let database_id = proto.database_id;
        let id = proto.id;
        self.get_database_mut(database_id)
            .unwrap()
            .create_schema(proto);

        for sys_table in get_sys_tables_in_schema(proto.name.as_str()) {
            self.get_database_mut(database_id)
                .unwrap()
                .get_schema_mut(id)
                .unwrap()
                .create_sys_table(sys_table);
        }
        for mut sys_view in get_sys_views_in_schema(proto.name.as_str()) {
            sys_view.database_id = database_id;
            sys_view.schema_id = id;
            self.get_database_mut(database_id)
                .unwrap()
                .get_schema_mut(id)
                .unwrap()
                .create_sys_view(Arc::new(sys_view));
        }
    }

    pub fn create_table(&mut self, proto: &PbTable) {
        let table = self
            .get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_table(proto);
        self.table_by_id.insert(proto.id, table);
    }

    pub fn create_index(&mut self, proto: &PbIndex) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_index(proto);
    }

    pub fn create_source(&mut self, proto: &PbSource) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_source(proto);
    }

    pub fn create_sink(&mut self, proto: &PbSink) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_sink(proto);
    }

    pub fn create_subscription(&mut self, proto: &PbSubscription) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_subscription(proto);
    }

    pub fn create_secret(&mut self, proto: &PbSecret) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_secret(proto);
    }

    pub fn create_view(&mut self, proto: &PbView) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_view(proto);
    }

    pub fn create_function(&mut self, proto: &PbFunction) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_function(proto);
    }

    pub fn create_connection(&mut self, proto: &PbConnection) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .create_connection(proto);
    }

    pub fn drop_connection(
        &mut self,
        db_id: DatabaseId,
        schema_id: SchemaId,
        connection_id: ConnectionId,
    ) {
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_connection(connection_id);
    }

    pub fn update_connection(&mut self, proto: &PbConnection) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_connection_by_id(proto.id).is_some() {
            schema.update_connection(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_connection(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id
                        && schema.get_connection_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_connection(proto.id);
        }
    }

    pub fn update_secret(&mut self, proto: &PbSecret) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        let secret_id = proto.id;
        if schema.get_secret_by_id(secret_id).is_some() {
            schema.update_secret(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_secret(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_secret_by_id(secret_id).is_some()
                })
                .unwrap()
                .drop_secret(secret_id);
        }
    }

    pub fn drop_database(&mut self, db_id: DatabaseId) {
        let name = self.db_name_by_id.remove(&db_id).unwrap();
        let database = self.database_by_name.remove(&name).unwrap();
        let object_ids: HashSet<ObjectId> = database.iter_object_ids().collect();
        self.remove_object_dependencies_for_objects(&object_ids);
        database.iter_all_table_ids().for_each(|table| {
            self.table_by_id.remove(&table);
        });
    }

    pub fn drop_schema(&mut self, db_id: DatabaseId, schema_id: SchemaId) {
        let object_ids = self
            .get_database_by_id(db_id)
            .ok()
            .and_then(|db| db.get_schema_by_id(schema_id))
            .map(|schema| schema.iter_object_ids().collect::<HashSet<_>>())
            .unwrap_or_default();
        self.remove_object_dependencies_for_objects(&object_ids);
        self.get_database_mut(db_id).unwrap().drop_schema(schema_id);
    }

    pub fn drop_table(&mut self, db_id: DatabaseId, schema_id: SchemaId, tb_id: TableId) {
        self.remove_object_dependencies_for_object(tb_id.as_object_id());
        self.table_by_id.remove(&tb_id);
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_table(tb_id);
    }

    pub fn update_table(&mut self, proto: &PbTable) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        let table = if schema.get_table_by_id(proto.id).is_some() {
            schema.update_table(proto)
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            let new_table = schema.create_table(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id
                        && schema.get_created_table_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_table(proto.id);
            new_table
        };

        self.table_by_id.insert(proto.id, table);
    }

    pub fn update_database(&mut self, proto: &PbDatabase) {
        let id = proto.id;
        let name = proto.name.clone();

        let old_database_name = self.db_name_by_id.get(&id).unwrap().to_owned();
        if old_database_name != name {
            let mut database = self.database_by_name.remove(&old_database_name).unwrap();
            database.name.clone_from(&name);
            database.owner = proto.owner;
            database.resource_group.clone_from(&proto.resource_group);
            database.barrier_interval_ms = proto.barrier_interval_ms;
            database.checkpoint_frequency = proto.checkpoint_frequency;
            self.database_by_name.insert(name.clone(), database);
            self.db_name_by_id.insert(id, name);
        } else {
            let database = self.get_database_mut(id).unwrap();
            database.name = name;
            database.owner = proto.owner;
            database.resource_group = proto.resource_group.clone();
            database.barrier_interval_ms = proto.barrier_interval_ms;
            database.checkpoint_frequency = proto.checkpoint_frequency;
        }
    }

    pub fn update_schema(&mut self, proto: &PbSchema) {
        self.get_database_mut(proto.database_id)
            .unwrap()
            .update_schema(proto);
    }

    pub fn update_index(&mut self, proto: &PbIndex) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_index_by_id(proto.id).is_some() {
            schema.update_index(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_index(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_index_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_index(proto.id);
        }
    }

    pub fn drop_source(&mut self, db_id: DatabaseId, schema_id: SchemaId, source_id: SourceId) {
        self.remove_object_dependencies_for_object(source_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_source(source_id);
    }

    pub fn update_source(&mut self, proto: &PbSource) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_source_by_id(proto.id).is_some() {
            schema.update_source(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_source(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_source_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_source(proto.id);
        }
    }

    pub fn drop_sink(&mut self, db_id: DatabaseId, schema_id: SchemaId, sink_id: SinkId) {
        self.remove_object_dependencies_for_object(sink_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_sink(sink_id);
    }

    pub fn drop_secret(&mut self, db_id: DatabaseId, schema_id: SchemaId, secret_id: SecretId) {
        self.remove_object_dependencies_for_object(secret_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_secret(secret_id);
    }

    pub fn update_sink(&mut self, proto: &PbSink) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_sink_by_id(proto.id).is_some() {
            schema.update_sink(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_sink(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_sink_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_sink(proto.id);
        }
    }

    pub fn drop_subscription(
        &mut self,
        db_id: DatabaseId,
        schema_id: SchemaId,
        subscription_id: SubscriptionId,
    ) {
        self.remove_object_dependencies_for_object(subscription_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_subscription(subscription_id);
    }

    pub fn update_subscription(&mut self, proto: &PbSubscription) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_subscription_by_id(proto.id).is_some() {
            schema.update_subscription(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_subscription(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id
                        && schema.get_subscription_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_subscription(proto.id);
        }
    }

    pub fn drop_index(&mut self, db_id: DatabaseId, schema_id: SchemaId, index_id: IndexId) {
        self.remove_object_dependencies_for_object(index_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_index(index_id);
    }

    pub fn drop_view(&mut self, db_id: DatabaseId, schema_id: SchemaId, view_id: ViewId) {
        self.remove_object_dependencies_for_object(view_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_view(view_id);
    }

    pub fn update_view(&mut self, proto: &PbView) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_view_by_id(proto.id).is_some() {
            schema.update_view(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_view(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_view_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_view(proto.id);
        }
    }

    pub fn drop_function(
        &mut self,
        db_id: DatabaseId,
        schema_id: SchemaId,
        function_id: FunctionId,
    ) {
        self.remove_object_dependencies_for_object(function_id.as_object_id());
        self.get_database_mut(db_id)
            .unwrap()
            .get_schema_mut(schema_id)
            .unwrap()
            .drop_function(function_id);
    }

    pub fn update_function(&mut self, proto: &PbFunction) {
        let database = self.get_database_mut(proto.database_id).unwrap();
        let schema = database.get_schema_mut(proto.schema_id).unwrap();
        if schema.get_function_by_id(proto.id).is_some() {
            schema.update_function(proto);
        } else {
            // Enter this branch when schema is changed by `ALTER ... SET SCHEMA ...` statement.
            schema.create_function(proto);
            database
                .iter_schemas_mut()
                .find(|schema| {
                    schema.id() != proto.schema_id && schema.get_function_by_id(proto.id).is_some()
                })
                .unwrap()
                .drop_function(proto.id);
        }

        self.get_database_mut(proto.database_id)
            .unwrap()
            .get_schema_mut(proto.schema_id)
            .unwrap()
            .update_function(proto);
    }

    pub fn get_database_by_name(&self, db_name: &str) -> CatalogResult<&DatabaseCatalog> {
        self.database_by_name
            .get(db_name)
            .ok_or_else(|| CatalogError::not_found("database", db_name))
    }

    pub fn get_database_by_id(&self, db_id: DatabaseId) -> CatalogResult<&DatabaseCatalog> {
        let db_name = self
            .db_name_by_id
            .get(&db_id)
            .ok_or_else(|| CatalogError::not_found("db_id", db_id.to_string()))?;
        self.database_by_name
            .get(db_name)
            .ok_or_else(|| CatalogError::not_found("database", db_name))
    }

    pub fn find_schema_secret_by_secret_id(
        &self,
        db_name: &str,
        secret_id: SecretId,
    ) -> CatalogResult<(String, String)> {
        let db = self.get_database_by_name(db_name)?;
        let schema_secret = db
            .iter_schemas()
            .find_map(|schema| {
                schema
                    .get_secret_by_id(secret_id)
                    .map(|secret| (schema.name(), secret.name.clone()))
            })
            .ok_or_else(|| CatalogError::not_found("secret", secret_id.to_string()))?;
        Ok(schema_secret)
    }

    pub fn get_all_schema_names(&self, db_name: &str) -> CatalogResult<Vec<String>> {
        Ok(self.get_database_by_name(db_name)?.get_all_schema_names())
    }

    pub fn iter_schemas(
        &self,
        db_name: &str,
    ) -> CatalogResult<impl Iterator<Item = &SchemaCatalog>> {
        Ok(self.get_database_by_name(db_name)?.iter_schemas())
    }

    pub fn get_all_database_names(&self) -> Vec<String> {
        self.database_by_name.keys().cloned().collect_vec()
    }

    pub fn iter_databases(&self) -> impl Iterator<Item = &DatabaseCatalog> {
        self.database_by_name.values()
    }

    pub fn get_schema_by_name(
        &self,
        db_name: &str,
        schema_name: &str,
    ) -> CatalogResult<&SchemaCatalog> {
        self.get_database_by_name(db_name)?
            .get_schema_by_name(schema_name)
            .ok_or_else(|| CatalogError::not_found("schema", schema_name))
    }

    pub fn get_table_name_by_id(&self, table_id: TableId) -> CatalogResult<String> {
        self.get_any_table_by_id(table_id)
            .map(|table| table.name.clone())
    }

    pub fn get_schema_by_id(
        &self,
        db_id: DatabaseId,
        schema_id: SchemaId,
    ) -> CatalogResult<&SchemaCatalog> {
        self.get_database_by_id(db_id)?
            .get_schema_by_id(schema_id)
            .ok_or_else(|| CatalogError::not_found("schema_id", schema_id.to_string()))
    }

    /// Refer to [`SearchPath`].
    pub fn first_valid_schema(
        &self,
        db_name: &str,
        search_path: &SearchPath,
        user_name: &str,
    ) -> CatalogResult<&SchemaCatalog> {
        for path in search_path.real_path() {
            let mut schema_name: &str = path;
            if schema_name == USER_NAME_WILD_CARD {
                schema_name = user_name;
            }

            if let schema_catalog @ Ok(_) = self.get_schema_by_name(db_name, schema_name) {
                return schema_catalog;
            }
        }
        Err(CatalogError::not_found(
            "first valid schema",
            "no schema has been selected to create in".to_owned(),
        ))
    }

    pub fn get_source_by_id<'a>(
        &self,
        db_name: &'a str,
        schema_path: SchemaPath<'a>,
        source_id: SourceId,
    ) -> CatalogResult<(&Arc<SourceCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_source_by_id(source_id))
            })?
            .ok_or_else(|| CatalogError::not_found("source", source_id.to_string()))
    }

    /// Finds a source by ID across all schemas in the given database.
    pub fn get_source_by_id_with_db(
        &self,
        db_name: &str,
        source_id: SourceId,
    ) -> CatalogResult<&Arc<SourceCatalog>> {
        for schema in self.get_database_by_name(db_name)?.iter_schemas() {
            if let Some(source) = schema.get_source_by_id(source_id) {
                return Ok(source);
            }
        }
        Err(CatalogError::not_found("source", source_id.to_string()))
    }

    pub fn get_table_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        table_name: &str,
        bind_creating: bool,
    ) -> CatalogResult<(&Arc<TableCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_table_by_name(table_name, bind_creating))
            })?
            .ok_or_else(|| CatalogError::not_found("table", table_name))
    }

    /// Used to get `TableCatalog` for Materialized Views, Tables and Indexes.
    /// Retrieves all tables, created or creating.
    pub fn get_any_table_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        table_name: &str,
    ) -> CatalogResult<(&Arc<TableCatalog>, &'a str)> {
        self.get_table_by_name(db_name, schema_path, table_name, true)
    }

    /// Used to get `TableCatalog` for Materialized Views, Tables and Indexes.
    /// Retrieves only created tables.
    pub fn get_created_table_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        table_name: &str,
    ) -> CatalogResult<(&Arc<TableCatalog>, &'a str)> {
        self.get_table_by_name(db_name, schema_path, table_name, false)
    }

    pub fn get_any_table_by_id(&self, table_id: TableId) -> CatalogResult<&Arc<TableCatalog>> {
        self.table_by_id
            .get(&table_id)
            .ok_or_else(|| CatalogError::not_found("table id", table_id.to_string()))
    }

    /// This function is similar to `get_table_by_id` expect that a table must be in a given database.
    pub fn get_created_table_by_id_with_db(
        &self,
        db_name: &str,
        table_id: u32,
    ) -> CatalogResult<&Arc<TableCatalog>> {
        let table_id = TableId::from(table_id);
        for schema in self.get_database_by_name(db_name)?.iter_schemas() {
            if let Some(table) = schema.get_created_table_by_id(table_id) {
                return Ok(table);
            }
        }
        Err(CatalogError::not_found("table id", table_id.to_string()))
    }

    pub fn iter_tables(&self) -> impl Iterator<Item = &Arc<TableCatalog>> {
        self.table_by_id.values()
    }

    pub fn iter_backfilling_internal_tables(&self) -> impl Iterator<Item = &Arc<TableCatalog>> {
        self.table_by_id
            .values()
            .filter(|t| t.is_internal_table() && !t.is_created())
    }

    // Used by test_utils only.
    pub fn alter_table_name_by_id(&mut self, table_id: TableId, table_name: &str) {
        let mut found = false;
        for database in self.database_by_name.values() {
            if !found {
                for schema in database.iter_schemas() {
                    if schema.iter_user_table().any(|t| t.id() == table_id) {
                        found = true;
                        break;
                    }
                }
            }
        }

        if found {
            let mut table = self.get_any_table_by_id(table_id).unwrap().to_prost();
            table.name = table_name.to_owned();
            self.update_table(&table);
        }
    }

    #[cfg(test)]
    pub fn insert_table_id_mapping(&mut self, table_id: TableId, fragment_id: super::FragmentId) {
        self.table_by_id.insert(
            table_id,
            Arc::new(TableCatalog {
                fragment_id,
                ..Default::default()
            }),
        );
    }

    pub fn get_sys_table_by_name(
        &self,
        db_name: &str,
        schema_name: &str,
        table_name: &str,
    ) -> CatalogResult<&Arc<SystemTableCatalog>> {
        self.get_schema_by_name(db_name, schema_name)?
            .get_system_table_by_name(table_name)
            .ok_or_else(|| CatalogError::not_found("table", table_name))
    }

    pub fn get_source_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        source_name: &str,
    ) -> CatalogResult<(&Arc<SourceCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_source_by_name(source_name))
            })?
            .ok_or_else(|| CatalogError::not_found("source", source_name))
    }

    pub fn get_sink_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        sink_name: &str,
        bind_creating: bool,
    ) -> CatalogResult<(&Arc<SinkCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_sink_by_name(sink_name, bind_creating))
            })?
            .ok_or_else(|| CatalogError::not_found("sink", sink_name))
    }

    pub fn get_any_sink_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        sink_name: &str,
    ) -> CatalogResult<(&Arc<SinkCatalog>, &'a str)> {
        self.get_sink_by_name(db_name, schema_path, sink_name, true)
    }

    pub fn get_created_sink_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        sink_name: &str,
    ) -> CatalogResult<(&Arc<SinkCatalog>, &'a str)> {
        self.get_sink_by_name(db_name, schema_path, sink_name, false)
    }

    pub fn get_subscription_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        subscription_name: &str,
    ) -> CatalogResult<(&Arc<SubscriptionCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_subscription_by_name(subscription_name))
            })?
            .ok_or_else(|| CatalogError::not_found("subscription", subscription_name))
    }

    pub fn get_index_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        index_name: &str,
    ) -> CatalogResult<(&Arc<IndexCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_created_index_by_name(index_name))
            })?
            .ok_or_else(|| CatalogError::not_found("index", index_name))
    }

    pub fn get_any_index_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        index_name: &str,
    ) -> CatalogResult<(&Arc<IndexCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_any_index_by_name(index_name))
            })?
            .ok_or_else(|| CatalogError::not_found("index", index_name))
    }

    pub fn get_index_by_id(
        &self,
        db_name: &str,
        index_id: u32,
    ) -> CatalogResult<&Arc<IndexCatalog>> {
        let index_id = IndexId::from(index_id);
        for schema in self.get_database_by_name(db_name)?.iter_schemas() {
            if let Some(index) = schema.get_index_by_id(index_id) {
                return Ok(index);
            }
        }
        Err(CatalogError::not_found("index", index_id.to_string()))
    }

    pub fn get_view_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        view_name: &str,
    ) -> CatalogResult<(&Arc<ViewCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_view_by_name(view_name))
            })?
            .ok_or_else(|| CatalogError::not_found("view", view_name))
    }

    pub fn get_view_by_id(&self, db_name: &str, view_id: u32) -> CatalogResult<Arc<ViewCatalog>> {
        for schema in self.get_database_by_name(db_name)?.iter_schemas() {
            if let Some(view) = schema.get_view_by_id(ViewId::from(view_id)) {
                return Ok(view.clone());
            }
        }
        Err(CatalogError::not_found("view", view_id.to_string()))
    }

    pub fn get_secret_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        secret_name: &str,
    ) -> CatalogResult<(&Arc<SecretCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_secret_by_name(secret_name))
            })?
            .ok_or_else(|| CatalogError::not_found("secret", secret_name))
    }

    pub fn get_connection_by_id(
        &self,
        db_name: &str,
        connection_id: ConnectionId,
    ) -> CatalogResult<&Arc<ConnectionCatalog>> {
        for schema in self.get_database_by_name(db_name)?.iter_schemas() {
            if let Some(conn) = schema.get_connection_by_id(connection_id) {
                return Ok(conn);
            }
        }
        Err(CatalogError::not_found(
            "connection",
            connection_id.to_string(),
        ))
    }

    pub fn get_connection_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        connection_name: &str,
    ) -> CatalogResult<(&Arc<ConnectionCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_connection_by_name(connection_name))
            })?
            .ok_or_else(|| CatalogError::not_found("connection", connection_name))
    }

    pub fn get_function_by_name_inputs<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        function_name: &str,
        inputs: &mut [ExprImpl],
    ) -> CatalogResult<(&Arc<FunctionCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_function_by_name_inputs(function_name, inputs))
            })?
            .ok_or_else(|| {
                CatalogError::not_found(
                    "function",
                    format!(
                        "{}({})",
                        function_name,
                        inputs
                            .iter()
                            .map(|a| a.return_type().to_string())
                            .join(", ")
                    ),
                )
            })
    }

    pub fn get_function_by_name_args<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        function_name: &str,
        args: &[DataType],
    ) -> CatalogResult<(&Arc<FunctionCatalog>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_function_by_name_args(function_name, args))
            })?
            .ok_or_else(|| {
                CatalogError::not_found(
                    "function",
                    format!(
                        "{}({})",
                        function_name,
                        args.iter().map(|a| a.to_string()).join(", ")
                    ),
                )
            })
    }

    /// Gets all functions with the given name.
    pub fn get_functions_by_name<'a>(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'a>,
        function_name: &str,
    ) -> CatalogResult<(Vec<&Arc<FunctionCatalog>>, &'a str)> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                Ok(self
                    .get_schema_by_name(db_name, schema_name)?
                    .get_functions_by_name(function_name))
            })?
            .ok_or_else(|| CatalogError::not_found("function", function_name))
    }

    /// Check if the name duplicates with existing table, materialized view or source.
    pub fn check_relation_name_duplicated(
        &self,
        db_name: &str,
        schema_name: &str,
        relation_name: &str,
    ) -> CatalogResult<()> {
        let schema = self.get_schema_by_name(db_name, schema_name)?;

        if let Some(table) = schema.get_any_table_by_name(relation_name) {
            let is_creating = table.stream_job_status == StreamJobStatus::Creating;
            if table.is_index() {
                Err(CatalogError::duplicated_under_creation(
                    "index",
                    relation_name.to_owned(),
                    is_creating,
                ))
            } else if table.is_mview() {
                Err(CatalogError::duplicated_under_creation(
                    "materialized view",
                    relation_name.to_owned(),
                    is_creating,
                ))
            } else {
                Err(CatalogError::duplicated_under_creation(
                    "table",
                    relation_name.to_owned(),
                    is_creating,
                ))
            }
        } else if schema.get_source_by_name(relation_name).is_some() {
            Err(CatalogError::duplicated("source", relation_name))
        } else if let Some(sink) = schema.get_any_sink_by_name(relation_name) {
            Err(CatalogError::duplicated_under_creation(
                "sink",
                relation_name.to_owned(),
                !sink.is_created(),
            ))
        } else if schema.get_view_by_name(relation_name).is_some() {
            Err(CatalogError::duplicated("view", relation_name))
        } else if let Some(subscription) = schema.get_subscription_by_name(relation_name) {
            let is_not_created = subscription.subscription_state != SubscriptionState::Created;
            Err(CatalogError::duplicated_under_creation(
                "subscription",
                relation_name.to_owned(),
                is_not_created,
            ))
        } else {
            Ok(())
        }
    }

    pub fn check_function_name_duplicated(
        &self,
        db_name: &str,
        schema_name: &str,
        function_name: &str,
        arg_types: &[DataType],
    ) -> CatalogResult<()> {
        let schema = self.get_schema_by_name(db_name, schema_name)?;

        if schema
            .get_function_by_name_args(function_name, arg_types)
            .is_some()
        {
            let name = format!(
                "{function_name}({})",
                arg_types.iter().map(|t| t.to_string()).join(",")
            );
            Err(CatalogError::duplicated("function", name))
        } else {
            Ok(())
        }
    }

    /// Check if the name duplicates with existing connection.
    pub fn check_connection_name_duplicated(
        &self,
        db_name: &str,
        schema_name: &str,
        connection_name: &str,
    ) -> CatalogResult<()> {
        let schema = self.get_schema_by_name(db_name, schema_name)?;

        if schema.get_connection_by_name(connection_name).is_some() {
            Err(CatalogError::duplicated("connection", connection_name))
        } else {
            Ok(())
        }
    }

    pub fn check_secret_name_duplicated(
        &self,
        db_name: &str,
        schema_name: &str,
        secret_name: &str,
    ) -> CatalogResult<()> {
        let schema = self.get_schema_by_name(db_name, schema_name)?;

        if schema.get_secret_by_name(secret_name).is_some() {
            Err(CatalogError::duplicated("secret", secret_name))
        } else {
            Ok(())
        }
    }

    pub fn table_stats(&self) -> &HummockVersionStats {
        &self.table_stats
    }

    pub fn set_table_stats(&mut self, table_stats: HummockVersionStats) {
        self.table_stats = table_stats;
    }

    pub fn set_object_dependencies(&mut self, dependencies: Vec<PbObjectDependency>) {
        self.object_dependencies.clear();
        self.insert_object_dependencies(dependencies);
    }

    pub fn insert_object_dependencies(&mut self, dependencies: Vec<PbObjectDependency>) {
        let mut grouped: HashMap<ObjectId, Vec<ObjectDependency>> = HashMap::new();
        for dependency in dependencies {
            let dependency = ObjectDependency::from(dependency);
            grouped
                .entry(dependency.object_id)
                .or_default()
                .push(dependency);
        }
        for (object_id, deps) in grouped {
            self.object_dependencies.insert(object_id, deps);
        }
    }

    pub fn iter_object_dependencies(&self) -> impl Iterator<Item = &ObjectDependency> {
        self.object_dependencies
            .values()
            .flat_map(|deps| deps.iter())
    }

    fn remove_object_dependencies_for_object(&mut self, object_id: ObjectId) {
        let mut object_ids = HashSet::new();
        object_ids.insert(object_id);
        self.remove_object_dependencies_for_objects(&object_ids);
    }

    fn remove_object_dependencies_for_objects(&mut self, object_ids: &HashSet<ObjectId>) {
        if object_ids.is_empty() {
            return;
        }
        self.object_dependencies.retain(|object_id, deps| {
            if object_ids.contains(object_id) {
                return false;
            }
            deps.retain(|dep| !object_ids.contains(&dep.referenced_object_id));
            !deps.is_empty()
        });
    }

    pub fn get_all_indexes_related_to_object(
        &self,
        db_id: DatabaseId,
        schema_id: SchemaId,
        mv_id: TableId,
    ) -> Vec<Arc<IndexCatalog>> {
        self.get_database_by_id(db_id)
            .unwrap()
            .get_schema_by_id(schema_id)
            .unwrap()
            .get_any_indexes_by_table_id(mv_id)
    }

    pub fn get_id_by_class_name(
        &self,
        db_name: &str,
        schema_path: SchemaPath<'_>,
        class_name: &str,
    ) -> CatalogResult<ObjectId> {
        schema_path
            .try_find(|schema_name| -> CatalogResult<_> {
                let schema = self.get_schema_by_name(db_name, schema_name)?;

                if let Some(item) = schema.get_system_table_by_name(class_name) {
                    Ok(Some(item.id().as_object_id()))
                } else if let Some(item) = schema.get_any_table_by_name(class_name) {
                    Ok(Some(item.id().as_object_id()))
                } else if let Some(item) = schema.get_any_index_by_name(class_name) {
                    Ok(Some(item.id.as_object_id()))
                } else if let Some(item) = schema.get_source_by_name(class_name) {
                    Ok(Some(item.id.as_object_id()))
                } else if let Some(item) = schema.get_view_by_name(class_name) {
                    Ok(Some(item.id.as_object_id()))
                } else if let Some(item) = schema.get_any_sink_by_name(class_name) {
                    Ok(Some(item.id.as_object_id()))
                } else {
                    Ok(None)
                }
            })?
            .map(|(id, _)| id)
            .ok_or_else(|| CatalogError::not_found("class", class_name))
    }
}

#[cfg(test)]
mod tests {
    use risingwave_pb::catalog::{
        PbCreateType, PbDatabase, PbSchema, PbSink, PbSinkType, PbStreamJobStatus,
    };

    use super::Catalog;

    fn test_sink(id: u32, name: &str, database_id: u32, schema_id: u32) -> PbSink {
        PbSink {
            id: id.into(),
            name: name.to_owned(),
            database_id: database_id.into(),
            schema_id: schema_id.into(),
            sink_type: PbSinkType::AppendOnly.into(),
            stream_job_status: PbStreamJobStatus::Creating.into(),
            create_type: PbCreateType::Background.into(),
            ..Default::default()
        }
    }

    #[test]
    fn test_update_sink_set_schema_drops_old_catalog() {
        let database_id = 1;
        let old_schema_id = 2;
        let new_schema_id = 3;
        let sink_id = 4;
        let mut catalog = Catalog::default();
        catalog.create_database(&PbDatabase {
            id: database_id.into(),
            name: "dev".to_owned(),
            ..Default::default()
        });
        catalog.create_schema(&PbSchema {
            id: old_schema_id.into(),
            database_id: database_id.into(),
            name: "public".to_owned(),
            ..Default::default()
        });
        catalog.create_schema(&PbSchema {
            id: new_schema_id.into(),
            database_id: database_id.into(),
            name: "archive".to_owned(),
            ..Default::default()
        });
        catalog.create_sink(&test_sink(
            sink_id,
            "replace_kafka_sink",
            database_id,
            old_schema_id,
        ));

        catalog.update_sink(&test_sink(
            sink_id,
            "replace_kafka_sink",
            database_id,
            new_schema_id,
        ));

        let old_schema = catalog
            .get_schema_by_id(database_id.into(), old_schema_id.into())
            .unwrap();
        assert!(old_schema.get_sink_by_id(sink_id.into()).is_none());
        let new_schema = catalog
            .get_schema_by_id(database_id.into(), new_schema_id.into())
            .unwrap();
        assert!(new_schema.get_sink_by_id(sink_id.into()).is_some());
    }
}