risingwave_frontend/catalog/
database_catalog.rs1use std::collections::HashMap;
16
17use itertools::Itertools;
18use risingwave_pb::catalog::{PbDatabase, PbSchema};
19use risingwave_pb::user::grant_privilege::Object;
20
21use super::OwnedByUserCatalog;
22use crate::catalog::schema_catalog::SchemaCatalog;
23use crate::catalog::{DatabaseId, SchemaId, TableId};
24use crate::user::UserId;
25
26#[derive(Clone, Debug)]
27pub struct DatabaseCatalog {
28 id: DatabaseId,
29 pub name: String,
30 schema_by_name: HashMap<String, SchemaCatalog>,
31 schema_name_by_id: HashMap<SchemaId, String>,
32 pub owner: u32,
33 pub resource_group: String,
34}
35
36impl DatabaseCatalog {
37 pub fn create_schema(&mut self, proto: &PbSchema) {
38 let name = proto.name.clone();
39 let id = proto.id;
40 let schema = proto.into();
41 self.schema_by_name
42 .try_insert(name.clone(), schema)
43 .unwrap();
44 self.schema_name_by_id.try_insert(id, name).unwrap();
45 }
46
47 pub fn drop_schema(&mut self, schema_id: SchemaId) {
48 let name = self.schema_name_by_id.remove(&schema_id).unwrap();
49 self.schema_by_name.remove(&name).unwrap();
50 }
51
52 pub fn get_all_schema_names(&self) -> Vec<String> {
53 self.schema_by_name.keys().cloned().collect_vec()
54 }
55
56 pub fn iter_all_table_ids(&self) -> impl Iterator<Item = TableId> + '_ {
57 self.schema_by_name
58 .values()
59 .flat_map(|schema| schema.iter_all().map(|t| t.id()))
60 }
61
62 pub fn iter_schemas(&self) -> impl Iterator<Item = &SchemaCatalog> {
63 self.schema_by_name.values()
64 }
65
66 pub fn iter_schemas_mut(&mut self) -> impl Iterator<Item = &mut SchemaCatalog> {
67 self.schema_by_name.values_mut()
68 }
69
70 pub fn get_schema_by_name(&self, name: &str) -> Option<&SchemaCatalog> {
71 self.schema_by_name.get(name)
72 }
73
74 pub fn get_schema_by_id(&self, schema_id: &SchemaId) -> Option<&SchemaCatalog> {
75 self.schema_by_name
76 .get(self.schema_name_by_id.get(schema_id)?)
77 }
78
79 pub fn get_schema_mut(&mut self, schema_id: SchemaId) -> Option<&mut SchemaCatalog> {
80 let name = self.schema_name_by_id.get(&schema_id).unwrap();
81 self.schema_by_name.get_mut(name)
82 }
83
84 pub fn get_grant_object_by_oid(&self, oid: u32) -> Option<Object> {
85 for schema in self.schema_by_name.values() {
86 let object = schema.get_grant_object_by_oid(oid);
87 if object.is_some() {
88 return object;
89 }
90 }
91 None
92 }
93
94 pub fn update_schema(&mut self, prost: &PbSchema) {
95 let id = prost.id;
96 let name = prost.name.clone();
97
98 let old_schema_name = self.schema_name_by_id.get(&id).unwrap().to_owned();
99 if old_schema_name != name {
100 let mut schema = self.schema_by_name.remove(&old_schema_name).unwrap();
101 schema.name.clone_from(&name);
102 schema.database_id = prost.database_id;
103 schema.owner = prost.owner;
104 self.schema_by_name.insert(name.clone(), schema);
105 self.schema_name_by_id.insert(id, name);
106 } else {
107 let schema = self.get_schema_mut(id).unwrap();
108 schema.name.clone_from(&name);
109 schema.database_id = prost.database_id;
110 schema.owner = prost.owner;
111 };
112 }
113
114 pub fn id(&self) -> DatabaseId {
115 self.id
116 }
117
118 pub fn name(&self) -> &str {
119 &self.name
120 }
121}
122
123impl OwnedByUserCatalog for DatabaseCatalog {
124 fn owner(&self) -> UserId {
125 self.owner
126 }
127}
128
129impl From<&PbDatabase> for DatabaseCatalog {
130 fn from(db: &PbDatabase) -> Self {
131 Self {
132 id: db.id,
133 name: db.name.clone(),
134 schema_by_name: HashMap::new(),
135 schema_name_by_id: HashMap::new(),
136 owner: db.owner,
137 resource_group: db.resource_group.clone(),
138 }
139 }
140}