catalog/memory/
manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::collections::hash_map::Entry;
17use std::collections::{HashMap, HashSet};
18use std::sync::{Arc, RwLock, Weak};
19
20use async_stream::{stream, try_stream};
21use common_catalog::build_db_string;
22use common_catalog::consts::{
23    DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME,
24    INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME,
25};
26use common_meta::key::flow::FlowMetadataManager;
27use common_meta::kv_backend::memory::MemoryKvBackend;
28use futures_util::stream::BoxStream;
29use session::context::QueryContext;
30use snafu::OptionExt;
31use table::TableRef;
32use table::metadata::{TableId, TableInfoRef};
33
34use crate::error::{CatalogNotFoundSnafu, Result, SchemaNotFoundSnafu, TableExistsSnafu};
35use crate::information_schema::InformationSchemaProvider;
36#[cfg(any(test, feature = "testing"))]
37use crate::process_manager::ProcessManagerRef;
38use crate::system_schema::SystemSchemaProvider;
39use crate::{CatalogManager, DeregisterTableRequest, RegisterSchemaRequest, RegisterTableRequest};
40
41type SchemaEntries = HashMap<String, HashMap<String, TableRef>>;
42
43/// Simple in-memory list of catalogs used for tests.
44#[derive(Clone)]
45pub struct MemoryCatalogManager {
46    /// Collection of catalogs containing schemas and ultimately Tables
47    catalogs: Arc<RwLock<HashMap<String, SchemaEntries>>>,
48}
49
50#[async_trait::async_trait]
51impl CatalogManager for MemoryCatalogManager {
52    fn as_any(&self) -> &dyn Any {
53        self
54    }
55
56    async fn catalog_names(&self) -> Result<Vec<String>> {
57        Ok(self.catalogs.read().unwrap().keys().cloned().collect())
58    }
59
60    async fn schema_names(
61        &self,
62        catalog: &str,
63        _query_ctx: Option<&QueryContext>,
64    ) -> Result<Vec<String>> {
65        Ok(self
66            .catalogs
67            .read()
68            .unwrap()
69            .get(catalog)
70            .with_context(|| CatalogNotFoundSnafu {
71                catalog_name: catalog,
72            })?
73            .keys()
74            .cloned()
75            .collect())
76    }
77
78    async fn table_names(
79        &self,
80        catalog: &str,
81        schema: &str,
82        _query_ctx: Option<&QueryContext>,
83    ) -> Result<Vec<String>> {
84        Ok(self
85            .catalogs
86            .read()
87            .unwrap()
88            .get(catalog)
89            .with_context(|| CatalogNotFoundSnafu {
90                catalog_name: catalog,
91            })?
92            .get(schema)
93            .with_context(|| SchemaNotFoundSnafu { catalog, schema })?
94            .keys()
95            .cloned()
96            .collect())
97    }
98
99    async fn catalog_exists(&self, catalog: &str) -> Result<bool> {
100        self.catalog_exist_sync(catalog)
101    }
102
103    async fn schema_exists(
104        &self,
105        catalog: &str,
106        schema: &str,
107        _query_ctx: Option<&QueryContext>,
108    ) -> Result<bool> {
109        self.schema_exist_sync(catalog, schema)
110    }
111
112    async fn table_exists(
113        &self,
114        catalog: &str,
115        schema: &str,
116        table: &str,
117        _query_ctx: Option<&QueryContext>,
118    ) -> Result<bool> {
119        let catalogs = self.catalogs.read().unwrap();
120        Ok(catalogs
121            .get(catalog)
122            .with_context(|| CatalogNotFoundSnafu {
123                catalog_name: catalog,
124            })?
125            .get(schema)
126            .with_context(|| SchemaNotFoundSnafu { catalog, schema })?
127            .contains_key(table))
128    }
129
130    async fn table(
131        &self,
132        catalog: &str,
133        schema: &str,
134        table_name: &str,
135        _query_ctx: Option<&QueryContext>,
136    ) -> Result<Option<TableRef>> {
137        let catalogs = self.catalogs.read().unwrap();
138        let result = catalogs
139            .get(catalog)
140            .and_then(|c| c.get(schema))
141            .and_then(|s| s.get(table_name))
142            .cloned();
143
144        Ok(result)
145    }
146
147    async fn table_info_by_id(&self, table_id: TableId) -> Result<Option<TableInfoRef>> {
148        Ok(self
149            .catalogs
150            .read()
151            .unwrap()
152            .values()
153            .flat_map(|schema_entries| schema_entries.values())
154            .flat_map(|tables| tables.values())
155            .find(|t| t.table_info().ident.table_id == table_id)
156            .map(|t| t.table_info()))
157    }
158
159    async fn tables_by_ids(
160        &self,
161        catalog: &str,
162        schema: &str,
163        table_ids: &[TableId],
164    ) -> Result<Vec<TableRef>> {
165        let catalogs = self.catalogs.read().unwrap();
166
167        let schemas = catalogs.get(catalog).context(CatalogNotFoundSnafu {
168            catalog_name: catalog,
169        })?;
170
171        let tables = schemas
172            .get(schema)
173            .context(SchemaNotFoundSnafu { catalog, schema })?;
174
175        let filter_ids: HashSet<_> = table_ids.iter().collect();
176        // It is very inefficient, but we do not need to optimize it since it will not be called in `MemoryCatalogManager`.
177        let tables = tables
178            .values()
179            .filter(|t| filter_ids.contains(&t.table_info().table_id()))
180            .cloned()
181            .collect::<Vec<_>>();
182
183        Ok(tables)
184    }
185
186    fn tables<'a>(
187        &'a self,
188        catalog: &'a str,
189        schema: &'a str,
190        _query_ctx: Option<&QueryContext>,
191    ) -> BoxStream<'a, Result<TableRef>> {
192        let catalogs = self.catalogs.read().unwrap();
193
194        let Some(schemas) = catalogs.get(catalog) else {
195            return Box::pin(stream!({
196                yield CatalogNotFoundSnafu {
197                    catalog_name: catalog,
198                }
199                .fail();
200            }));
201        };
202
203        let Some(tables) = schemas.get(schema) else {
204            return Box::pin(stream!({
205                yield SchemaNotFoundSnafu { catalog, schema }.fail();
206            }));
207        };
208
209        let tables = tables.values().cloned().collect::<Vec<_>>();
210
211        Box::pin(try_stream!({
212            for table in tables {
213                yield table;
214            }
215        }))
216    }
217}
218
219impl MemoryCatalogManager {
220    pub fn new() -> Arc<Self> {
221        Arc::new(Self {
222            catalogs: Default::default(),
223        })
224    }
225
226    /// Creates a manager with some default setups
227    /// (e.g. default catalog/schema and information schema)
228    pub fn with_default_setup() -> Arc<Self> {
229        let manager = Arc::new(Self {
230            catalogs: Default::default(),
231        });
232
233        // Safety: default catalog/schema is registered in order so no CatalogNotFound error will occur
234        manager.register_catalog_sync(DEFAULT_CATALOG_NAME).unwrap();
235        manager
236            .register_schema_sync(RegisterSchemaRequest {
237                catalog: DEFAULT_CATALOG_NAME.to_string(),
238                schema: DEFAULT_SCHEMA_NAME.to_string(),
239            })
240            .unwrap();
241        manager
242            .register_schema_sync(RegisterSchemaRequest {
243                catalog: DEFAULT_CATALOG_NAME.to_string(),
244                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
245            })
246            .unwrap();
247        manager
248            .register_schema_sync(RegisterSchemaRequest {
249                catalog: DEFAULT_CATALOG_NAME.to_string(),
250                schema: PG_CATALOG_NAME.to_string(),
251            })
252            .unwrap();
253        manager
254            .register_schema_sync(RegisterSchemaRequest {
255                catalog: DEFAULT_CATALOG_NAME.to_string(),
256                schema: INFORMATION_SCHEMA_NAME.to_string(),
257            })
258            .unwrap();
259
260        manager
261    }
262
263    #[cfg(any(test, feature = "testing"))]
264    pub fn register_process_list_table(&self, process_manager: ProcessManagerRef) {
265        let backend = Arc::new(MemoryKvBackend::new());
266        let manager = Arc::new(self.clone());
267        let information_schema_provider = InformationSchemaProvider::new(
268            DEFAULT_CATALOG_NAME.to_string(),
269            Arc::downgrade(&manager) as Weak<dyn CatalogManager>,
270            Arc::new(FlowMetadataManager::new(backend.clone())),
271            Some(process_manager),
272            backend,
273        );
274        let process_list = information_schema_provider
275            .table(crate::information_schema::PROCESS_LIST)
276            .unwrap();
277        self.catalogs
278            .write()
279            .unwrap()
280            .get_mut(DEFAULT_CATALOG_NAME)
281            .unwrap()
282            .get_mut(INFORMATION_SCHEMA_NAME)
283            .unwrap()
284            .insert(
285                crate::information_schema::PROCESS_LIST.to_string(),
286                process_list,
287            );
288    }
289
290    fn schema_exist_sync(&self, catalog: &str, schema: &str) -> Result<bool> {
291        Ok(self
292            .catalogs
293            .read()
294            .unwrap()
295            .get(catalog)
296            .with_context(|| CatalogNotFoundSnafu {
297                catalog_name: catalog,
298            })?
299            .contains_key(schema))
300    }
301
302    fn catalog_exist_sync(&self, catalog: &str) -> Result<bool> {
303        Ok(self.catalogs.read().unwrap().contains_key(catalog))
304    }
305
306    /// Registers a catalog if it does not exist and returns false if the schema exists.
307    pub fn register_catalog_sync(&self, name: &str) -> Result<bool> {
308        let name = name.to_string();
309
310        let mut catalogs = self.catalogs.write().unwrap();
311
312        match catalogs.entry(name.clone()) {
313            Entry::Vacant(e) => {
314                let arc_self = Arc::new(self.clone());
315                let catalog = arc_self.create_catalog_entry(name);
316                e.insert(catalog);
317                crate::metrics::METRIC_CATALOG_MANAGER_CATALOG_COUNT.inc();
318                Ok(true)
319            }
320            Entry::Occupied(_) => Ok(false),
321        }
322    }
323
324    pub fn deregister_table_sync(&self, request: DeregisterTableRequest) -> Result<()> {
325        let mut catalogs = self.catalogs.write().unwrap();
326        let schema = catalogs
327            .get_mut(&request.catalog)
328            .with_context(|| CatalogNotFoundSnafu {
329                catalog_name: &request.catalog,
330            })?
331            .get_mut(&request.schema)
332            .with_context(|| SchemaNotFoundSnafu {
333                catalog: &request.catalog,
334                schema: &request.schema,
335            })?;
336        let result = schema.remove(&request.table_name);
337        if result.is_some() {
338            crate::metrics::METRIC_CATALOG_MANAGER_TABLE_COUNT
339                .with_label_values(&[build_db_string(&request.catalog, &request.schema).as_str()])
340                .dec();
341        }
342        Ok(())
343    }
344
345    /// Registers a schema if it does not exist.
346    /// It returns an error if the catalog does not exist,
347    /// and returns false if the schema exists.
348    pub fn register_schema_sync(&self, request: RegisterSchemaRequest) -> Result<bool> {
349        let mut catalogs = self.catalogs.write().unwrap();
350        let catalog = catalogs
351            .get_mut(&request.catalog)
352            .with_context(|| CatalogNotFoundSnafu {
353                catalog_name: &request.catalog,
354            })?;
355
356        match catalog.entry(request.schema) {
357            Entry::Vacant(e) => {
358                e.insert(HashMap::new());
359                crate::metrics::METRIC_CATALOG_MANAGER_SCHEMA_COUNT.inc();
360                Ok(true)
361            }
362            Entry::Occupied(_) => Ok(false),
363        }
364    }
365
366    /// Registers a schema and returns an error if the catalog or schema does not exist.
367    pub fn register_table_sync(&self, request: RegisterTableRequest) -> Result<bool> {
368        let mut catalogs = self.catalogs.write().unwrap();
369        let schema = catalogs
370            .get_mut(&request.catalog)
371            .with_context(|| CatalogNotFoundSnafu {
372                catalog_name: &request.catalog,
373            })?
374            .get_mut(&request.schema)
375            .with_context(|| SchemaNotFoundSnafu {
376                catalog: &request.catalog,
377                schema: &request.schema,
378            })?;
379
380        if schema.contains_key(&request.table_name) {
381            return TableExistsSnafu {
382                table: &request.table_name,
383            }
384            .fail();
385        }
386        schema.insert(request.table_name, request.table);
387        crate::metrics::METRIC_CATALOG_MANAGER_TABLE_COUNT
388            .with_label_values(&[build_db_string(&request.catalog, &request.schema).as_str()])
389            .inc();
390        Ok(true)
391    }
392
393    fn create_catalog_entry(self: &Arc<Self>, catalog: String) -> SchemaEntries {
394        let backend = Arc::new(MemoryKvBackend::new());
395        let information_schema_provider = InformationSchemaProvider::new(
396            catalog,
397            Arc::downgrade(self) as Weak<dyn CatalogManager>,
398            Arc::new(FlowMetadataManager::new(backend.clone())),
399            None, // we don't need ProcessManager on regions server.
400            backend,
401        );
402        let information_schema = information_schema_provider.tables().clone();
403
404        let mut catalog = HashMap::new();
405        catalog.insert(INFORMATION_SCHEMA_NAME.to_string(), information_schema);
406        catalog
407    }
408
409    #[cfg(any(test, feature = "testing"))]
410    pub fn new_with_table(table: TableRef) -> Arc<Self> {
411        let manager = Self::with_default_setup();
412        let catalog = &table.table_info().catalog_name;
413        let schema = &table.table_info().schema_name;
414
415        if !manager.catalog_exist_sync(catalog).unwrap() {
416            manager.register_catalog_sync(catalog).unwrap();
417        }
418
419        if !manager.schema_exist_sync(catalog, schema).unwrap() {
420            manager
421                .register_schema_sync(RegisterSchemaRequest {
422                    catalog: catalog.clone(),
423                    schema: schema.clone(),
424                })
425                .unwrap();
426        }
427
428        let request = RegisterTableRequest {
429            catalog: catalog.clone(),
430            schema: schema.clone(),
431            table_name: table.table_info().name.clone(),
432            table_id: table.table_info().ident.table_id,
433            table,
434        };
435        let _ = manager.register_table_sync(request).unwrap();
436        manager
437    }
438}
439
440/// Create a memory catalog list contains a numbers table for test
441pub fn new_memory_catalog_manager() -> Result<Arc<MemoryCatalogManager>> {
442    Ok(MemoryCatalogManager::with_default_setup())
443}
444
445#[cfg(test)]
446mod tests {
447    use common_catalog::consts::*;
448    use futures_util::TryStreamExt;
449    use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable};
450
451    use super::*;
452
453    #[tokio::test]
454    async fn test_new_memory_catalog_list() {
455        let catalog_list = new_memory_catalog_manager().unwrap();
456
457        let register_request = RegisterTableRequest {
458            catalog: DEFAULT_CATALOG_NAME.to_string(),
459            schema: DEFAULT_SCHEMA_NAME.to_string(),
460            table_name: NUMBERS_TABLE_NAME.to_string(),
461            table_id: NUMBERS_TABLE_ID,
462            table: NumbersTable::table(NUMBERS_TABLE_ID),
463        };
464
465        catalog_list.register_table_sync(register_request).unwrap();
466        let table = catalog_list
467            .table(
468                DEFAULT_CATALOG_NAME,
469                DEFAULT_SCHEMA_NAME,
470                NUMBERS_TABLE_NAME,
471                None,
472            )
473            .await
474            .unwrap()
475            .unwrap();
476        let stream = catalog_list.tables(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, None);
477        let tables = stream.try_collect::<Vec<_>>().await.unwrap();
478        assert_eq!(tables.len(), 1);
479        assert_eq!(
480            table.table_info().table_id(),
481            tables[0].table_info().table_id()
482        );
483
484        assert!(
485            catalog_list
486                .table(
487                    DEFAULT_CATALOG_NAME,
488                    DEFAULT_SCHEMA_NAME,
489                    "not_exists",
490                    None
491                )
492                .await
493                .unwrap()
494                .is_none()
495        );
496    }
497
498    #[test]
499    pub fn test_register_catalog_sync() {
500        let list = MemoryCatalogManager::with_default_setup();
501        assert!(list.register_catalog_sync("test_catalog").unwrap());
502        assert!(!list.register_catalog_sync("test_catalog").unwrap());
503    }
504
505    #[tokio::test]
506    pub async fn test_catalog_deregister_table() {
507        let catalog = MemoryCatalogManager::with_default_setup();
508        let table_name = "foo_table";
509
510        let register_table_req = RegisterTableRequest {
511            catalog: DEFAULT_CATALOG_NAME.to_string(),
512            schema: DEFAULT_SCHEMA_NAME.to_string(),
513            table_name: table_name.to_string(),
514            table_id: 2333,
515            table: NumbersTable::table(2333),
516        };
517        catalog.register_table_sync(register_table_req).unwrap();
518        assert!(
519            catalog
520                .table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, table_name, None)
521                .await
522                .unwrap()
523                .is_some()
524        );
525
526        let deregister_table_req = DeregisterTableRequest {
527            catalog: DEFAULT_CATALOG_NAME.to_string(),
528            schema: DEFAULT_SCHEMA_NAME.to_string(),
529            table_name: table_name.to_string(),
530        };
531        catalog.deregister_table_sync(deregister_table_req).unwrap();
532        assert!(
533            catalog
534                .table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, table_name, None)
535                .await
536                .unwrap()
537                .is_none()
538        );
539    }
540}
catalog/memory/manager.rs

catalog/memory/
manager.rs
