Skip to main content

vantage_table/table/impls/
columns.rs

1use indexmap::IndexMap;
2use vantage_types::Entity;
3
4use crate::{
5    column::core::ColumnType, prelude::ColumnLike, table::Table, traits::table_source::TableSource,
6};
7
8impl<T: TableSource, E: Entity<T::Value>> Table<T, E> {
9    /// Add a column to the table (accepts any typed column, converts to `Column<AnyType>`)
10    pub fn add_column<NewColumnType>(&mut self, column: T::Column<NewColumnType>)
11    where
12        NewColumnType: ColumnType,
13    {
14        let name = column.name().to_string();
15
16        if self.columns.contains_key(&name) {
17            panic!("Duplicate column: {}", name);
18        }
19
20        // Convert typed column to Column<AnyType> for storage
21        let any_column = self.data_source.to_any_column(column);
22        self.columns.insert(name, any_column);
23    }
24
25    /// Add a column using builder pattern
26    pub fn with_column<NewColumnType>(mut self, column: T::Column<NewColumnType>) -> Self
27    where
28        NewColumnType: ColumnType,
29    {
30        self.add_column(column);
31        self
32    }
33
34    /// Add a typed column to the table (mutable)
35    pub fn add_column_of<NewColumnType>(&mut self, name: impl Into<String>)
36    where
37        NewColumnType: ColumnType,
38    {
39        let column = self
40            .data_source
41            .create_column::<NewColumnType>(&name.into());
42        self.add_column(column);
43    }
44
45    /// Add an ID column — sets both the column and the id_field flag.
46    pub fn with_id_column(mut self, name: impl Into<String>) -> Self
47    where
48        T::Id: ColumnType,
49    {
50        let name = name.into();
51        self.id_field = Some(name.clone());
52        let column = self.data_source.create_column::<T::Id>(&name);
53        self.add_column(column);
54        self
55    }
56
57    /// Add a typed column to the table (builder pattern)
58    pub fn with_column_of<NewColumnType>(self, name: impl Into<String>) -> Self
59    where
60        NewColumnType: ColumnType,
61    {
62        let column = self
63            .data_source
64            .create_column::<NewColumnType>(&name.into());
65        self.with_column(column)
66    }
67
68    /// Get all columns as type-erased columns (`Column<AnyType>`)
69    pub fn columns(&self) -> &IndexMap<String, T::Column<T::AnyType>> {
70        &self.columns
71    }
72
73    /// Get a typed column by converting from stored `Column<AnyType>`
74    pub fn get_column<Type>(&self, name: &str) -> Option<T::Column<Type>>
75    where
76        Type: ColumnType,
77    {
78        let any_column = self.columns.get(name)?;
79        self.data_source
80            .convert_any_column::<Type>(any_column.clone())
81    }
82
83    /// Get an expression for a column or computed expression by name.
84    ///
85    /// If `name` matches a registered expression (from `with_expression`), evaluates
86    /// and returns it. Otherwise returns the column as an expression. Returns `None`
87    /// if the name doesn't match either.
88    pub fn get_column_expr(&self, name: &str) -> Option<vantage_expressions::Expression<T::Value>>
89    where
90        T::Column<T::AnyType>: vantage_expressions::Expressive<T::Value>,
91    {
92        if let Some(expr_fn) = self.expressions.get(name) {
93            Some(expr_fn(self))
94        } else {
95            use vantage_expressions::Expressive;
96            self.columns.get(name).map(|c| c.expr())
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::mocks::mock_column::MockColumn;
105    use crate::prelude::MockTableSource;
106    use serde_json::Value;
107    use vantage_types::EmptyEntity;
108
109    #[test]
110    fn test_add_column() {
111        let ds = MockTableSource::new();
112        let mut table = Table::<MockTableSource, EmptyEntity>::new("test", ds);
113
114        table.add_column(MockColumn::<String>::new("name"));
115
116        assert!(table.columns().contains_key("name"));
117        assert_eq!(table.columns().len(), 1);
118    }
119
120    #[test]
121    fn test_with_column() {
122        let ds = MockTableSource::new();
123        let table = Table::<MockTableSource, EmptyEntity>::new("test", ds)
124            .with_column(MockColumn::<Value>::new("name"))
125            .with_column(MockColumn::<i32>::new("email"));
126
127        assert!(table.columns().contains_key("name"));
128        assert!(table.columns().contains_key("email"));
129        assert_eq!(table.columns().len(), 2);
130    }
131
132    #[test]
133    #[should_panic(expected = "Duplicate column")]
134    fn test_duplicate_column_panics() {
135        let ds = MockTableSource::new();
136        let mut table = Table::<MockTableSource, EmptyEntity>::new("test", ds);
137
138        table.add_column(MockColumn::<String>::new("name"));
139        table.add_column(MockColumn::<String>::new("name")); // Should panic
140    }
141
142    #[test]
143    fn test_with_column_of() {
144        let ds = MockTableSource::new();
145        let table = Table::<MockTableSource, EmptyEntity>::new("test", ds)
146            .with_column_of::<String>("name")
147            .with_column_of::<i64>("age")
148            .with_column_of::<bool>("active");
149
150        assert!(table.columns().contains_key("name"));
151        assert!(table.columns().contains_key("age"));
152        assert!(table.columns().contains_key("active"));
153        assert_eq!(table.columns().len(), 3);
154    }
155
156    #[test]
157    fn test_add_column_of() {
158        let ds = MockTableSource::new();
159        let mut table = Table::<MockTableSource, EmptyEntity>::new("test", ds);
160
161        table.add_column_of::<String>("email");
162        table.add_column_of::<i64>("balance");
163
164        assert!(table.columns().contains_key("email"));
165        assert!(table.columns().contains_key("balance"));
166        assert_eq!(table.columns().len(), 2);
167    }
168
169    #[test]
170    fn test_columns_access() {
171        let ds = MockTableSource::new();
172        let table = Table::<MockTableSource, EmptyEntity>::new("test", ds)
173            .with_column_of::<String>("name")
174            .with_column_of::<i64>("age");
175
176        let columns = table.columns();
177        assert!(columns.contains_key("name"));
178        assert!(columns.contains_key("age"));
179        assert_eq!(columns.len(), 2);
180
181        let name_column = table.columns().get("name");
182        assert!(name_column.is_some());
183        assert_eq!(name_column.unwrap().name(), "name");
184
185        let missing_column = table.columns().get("missing");
186        assert!(missing_column.is_none());
187    }
188}