mz_expr_parser/
catalog.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;
11
12use mz_ore::cast::CastFrom;
13use mz_repr::explain::{DummyHumanizer, ExprHumanizer};
14use mz_repr::{GlobalId, RelationType, ScalarType};
15
16/// A catalog that holds types of objects previously created for the unit test.
17///
18/// This is for the purpose of allowing `MirRelationExpr`s to refer to them
19/// later.
20#[derive(Debug, Default)]
21pub struct TestCatalog {
22    objects: BTreeMap<String, (GlobalId, Vec<String>, RelationType)>,
23    names: BTreeMap<GlobalId, String>,
24}
25
26impl<'a> TestCatalog {
27    /// Registers an object in the catalog.
28    ///
29    /// Specifying `transient` as true allows the object to be deleted by
30    /// [Self::remove_transient_objects].
31    ///
32    /// Returns the GlobalId assigned by the catalog to the object.
33    ///
34    /// Errors if an object of the same name is already in the catalog.
35    pub fn insert(
36        &mut self,
37        name: &str,
38        cols: Vec<String>,
39        typ: RelationType,
40        transient: bool,
41    ) -> Result<GlobalId, String> {
42        if self.objects.contains_key(name) {
43            return Err(format!("Object {name} already exists in catalog"));
44        }
45        let id = if transient {
46            GlobalId::Transient(u64::cast_from(self.objects.len()))
47        } else {
48            GlobalId::User(u64::cast_from(self.objects.len()))
49        };
50        self.objects.insert(name.to_string(), (id, cols, typ));
51        self.names.insert(id, name.to_string());
52        Ok(id)
53    }
54
55    pub fn get(&'a self, name: &str) -> Option<&'a (GlobalId, Vec<String>, RelationType)> {
56        self.objects.get(name)
57    }
58
59    /// Looks up the name of the object referred to as `id`.
60    pub fn get_source_name(&'a self, id: &GlobalId) -> Option<&'a String> {
61        self.names.get(id)
62    }
63
64    /// Clears all transient objects from the catalog.
65    pub fn remove_transient_objects(&mut self) {
66        self.objects
67            .retain(|_, (id, _, _)| !matches!(id, GlobalId::Transient(_)));
68        self.names
69            .retain(|k, _| !matches!(k, GlobalId::Transient(_)));
70    }
71}
72
73impl ExprHumanizer for TestCatalog {
74    fn humanize_id(&self, id: GlobalId) -> Option<String> {
75        self.names.get(&id).map(|s| s.to_string())
76    }
77
78    fn humanize_id_unqualified(&self, id: GlobalId) -> Option<String> {
79        self.names.get(&id).map(|s| s.to_string())
80    }
81
82    fn humanize_id_parts(&self, id: GlobalId) -> Option<Vec<String>> {
83        self.humanize_id_unqualified(id).map(|name| vec![name])
84    }
85
86    fn humanize_scalar_type(&self, ty: &ScalarType, postgres_compat: bool) -> String {
87        DummyHumanizer.humanize_scalar_type(ty, postgres_compat)
88    }
89
90    fn column_names_for_id(&self, id: GlobalId) -> Option<Vec<String>> {
91        let src_name = self.get_source_name(&id)?;
92        self.objects.get(src_name).map(|(_, cols, _)| cols.clone())
93    }
94
95    fn humanize_column(&self, id: GlobalId, column: usize) -> Option<String> {
96        let src_name = self.get_source_name(&id)?;
97        self.objects
98            .get(src_name)
99            .map(|(_, cols, _)| cols[column].clone())
100    }
101
102    fn id_exists(&self, id: GlobalId) -> bool {
103        self.names.contains_key(&id)
104    }
105}
mz_expr_parser/catalog.rs

mz_expr_parser/
catalog.rs
