sql_cli/sql/functions/
mod.rs

1use anyhow::{anyhow, Result};
2use std::collections::HashMap;
3use std::fmt;
4
5use crate::data::datatable::DataValue;
6
7pub mod astronomy;
8pub mod chemistry;
9pub mod constants;
10
11/// Category of SQL functions for organization and discovery
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub enum FunctionCategory {
14    Constant,     // Mathematical and physical constants
15    Mathematical, // Mathematical operations
16    Astronomical, // Astronomical constants and calculations
17    Chemical,     // Chemical elements and properties
18    Date,         // Date/time operations
19    String,       // String manipulation
20    Aggregate,    // Aggregation functions
21}
22
23impl fmt::Display for FunctionCategory {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            FunctionCategory::Constant => write!(f, "Constant"),
27            FunctionCategory::Mathematical => write!(f, "Mathematical"),
28            FunctionCategory::Astronomical => write!(f, "Astronomical"),
29            FunctionCategory::Chemical => write!(f, "Chemical"),
30            FunctionCategory::Date => write!(f, "Date"),
31            FunctionCategory::String => write!(f, "String"),
32            FunctionCategory::Aggregate => write!(f, "Aggregate"),
33        }
34    }
35}
36
37/// Describes the number of arguments a function accepts
38#[derive(Debug, Clone)]
39pub enum ArgCount {
40    /// Exactly n arguments
41    Fixed(usize),
42    /// Between min and max arguments (inclusive)
43    Range(usize, usize),
44    /// Any number of arguments
45    Variadic,
46}
47
48impl ArgCount {
49    pub fn is_valid(&self, count: usize) -> bool {
50        match self {
51            ArgCount::Fixed(n) => count == *n,
52            ArgCount::Range(min, max) => count >= *min && count <= *max,
53            ArgCount::Variadic => true,
54        }
55    }
56
57    pub fn description(&self) -> String {
58        match self {
59            ArgCount::Fixed(0) => "no arguments".to_string(),
60            ArgCount::Fixed(1) => "1 argument".to_string(),
61            ArgCount::Fixed(n) => format!("{} arguments", n),
62            ArgCount::Range(min, max) => format!("{} to {} arguments", min, max),
63            ArgCount::Variadic => "any number of arguments".to_string(),
64        }
65    }
66}
67
68/// Signature of a SQL function including metadata
69#[derive(Debug, Clone)]
70pub struct FunctionSignature {
71    pub name: &'static str,
72    pub category: FunctionCategory,
73    pub arg_count: ArgCount,
74    pub description: &'static str,
75    pub returns: &'static str,
76    pub examples: Vec<&'static str>,
77}
78
79/// Trait that all SQL functions must implement
80pub trait SqlFunction: Send + Sync {
81    /// Get the function's signature and metadata
82    fn signature(&self) -> FunctionSignature;
83
84    /// Evaluate the function with the given arguments
85    fn evaluate(&self, args: &[DataValue]) -> Result<DataValue>;
86
87    /// Validate arguments before evaluation (default implementation checks count)
88    fn validate_args(&self, args: &[DataValue]) -> Result<()> {
89        let sig = self.signature();
90        if !sig.arg_count.is_valid(args.len()) {
91            return Err(anyhow!(
92                "{}() expects {}, got {}",
93                sig.name,
94                sig.arg_count.description(),
95                args.len()
96            ));
97        }
98        Ok(())
99    }
100}
101
102/// Registry for all SQL functions
103pub struct FunctionRegistry {
104    functions: HashMap<String, Box<dyn SqlFunction>>,
105    by_category: HashMap<FunctionCategory, Vec<String>>,
106}
107
108impl FunctionRegistry {
109    /// Create a new registry with all built-in functions
110    pub fn new() -> Self {
111        let mut registry = Self {
112            functions: HashMap::new(),
113            by_category: HashMap::new(),
114        };
115
116        // Register all built-in functions
117        registry.register_constants();
118        registry.register_astronomical_functions();
119        registry.register_chemical_functions();
120
121        registry
122    }
123
124    /// Register a function in the registry
125    pub fn register(&mut self, func: Box<dyn SqlFunction>) {
126        let sig = func.signature();
127        let name = sig.name.to_uppercase();
128        let category = sig.category;
129
130        // Add to main registry
131        self.functions.insert(name.clone(), func);
132
133        // Add to category index
134        self.by_category
135            .entry(category)
136            .or_insert_with(Vec::new)
137            .push(name);
138    }
139
140    /// Get a function by name (case-insensitive)
141    pub fn get(&self, name: &str) -> Option<&dyn SqlFunction> {
142        self.functions.get(&name.to_uppercase()).map(|b| b.as_ref())
143    }
144
145    /// Check if a function exists
146    pub fn contains(&self, name: &str) -> bool {
147        self.functions.contains_key(&name.to_uppercase())
148    }
149
150    /// Get all functions matching a prefix (for autocomplete)
151    pub fn autocomplete(&self, prefix: &str) -> Vec<FunctionSignature> {
152        let prefix_upper = prefix.to_uppercase();
153        self.functions
154            .iter()
155            .filter(|(name, _)| name.starts_with(&prefix_upper))
156            .map(|(_, func)| func.signature())
157            .collect()
158    }
159
160    /// Get all functions in a category
161    pub fn get_by_category(&self, category: FunctionCategory) -> Vec<FunctionSignature> {
162        self.by_category
163            .get(&category)
164            .map(|names| {
165                names
166                    .iter()
167                    .filter_map(|name| self.functions.get(name))
168                    .map(|func| func.signature())
169                    .collect()
170            })
171            .unwrap_or_default()
172    }
173
174    /// Get all available functions
175    pub fn all_functions(&self) -> Vec<FunctionSignature> {
176        self.functions
177            .values()
178            .map(|func| func.signature())
179            .collect()
180    }
181
182    /// Register constant functions
183    fn register_constants(&mut self) {
184        use constants::*;
185
186        self.register(Box::new(PiFunction));
187        self.register(Box::new(EFunction));
188        self.register(Box::new(MeFunction)); // Mass of electron
189    }
190
191    /// Register astronomical functions
192    fn register_astronomical_functions(&mut self) {
193        use astronomy::*;
194
195        self.register(Box::new(MassEarthFunction));
196        self.register(Box::new(MassSunFunction));
197        self.register(Box::new(MassMoonFunction));
198        self.register(Box::new(AuFunction)); // Astronomical unit
199        self.register(Box::new(LightYearFunction));
200        self.register(Box::new(ParsecFunction));
201
202        // Planetary masses
203        self.register(Box::new(MassMercuryFunction));
204        self.register(Box::new(MassVenusFunction));
205        self.register(Box::new(MassMarsFunction));
206        self.register(Box::new(MassJupiterFunction));
207        self.register(Box::new(MassSaturnFunction));
208        self.register(Box::new(MassUranusFunction));
209        self.register(Box::new(MassNeptuneFunction));
210
211        // Solar body radius functions
212        self.register(Box::new(RadiusSunFunction));
213        self.register(Box::new(RadiusEarthFunction));
214        self.register(Box::new(RadiusMoonFunction));
215        self.register(Box::new(RadiusMercuryFunction));
216        self.register(Box::new(RadiusVenusFunction));
217        self.register(Box::new(RadiusMarsFunction));
218        self.register(Box::new(RadiusJupiterFunction));
219        self.register(Box::new(RadiusSaturnFunction));
220        self.register(Box::new(RadiusUranusFunction));
221        self.register(Box::new(RadiusNeptuneFunction));
222    }
223
224    /// Register chemical functions
225    fn register_chemical_functions(&mut self) {
226        use chemistry::*;
227
228        self.register(Box::new(AvogadroFunction));
229        self.register(Box::new(AtomicMassFunction));
230        self.register(Box::new(AtomicNumberFunction));
231    }
232}
233
234impl Default for FunctionRegistry {
235    fn default() -> Self {
236        Self::new()
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn test_registry_creation() {
246        let registry = FunctionRegistry::new();
247
248        // Check that some known functions exist
249        assert!(registry.contains("PI"));
250        assert!(registry.contains("MASS_EARTH"));
251        assert!(registry.contains("ME"));
252    }
253
254    #[test]
255    fn test_case_insensitive_lookup() {
256        let registry = FunctionRegistry::new();
257
258        assert!(registry.get("pi").is_some());
259        assert!(registry.get("PI").is_some());
260        assert!(registry.get("Pi").is_some());
261    }
262
263    #[test]
264    fn test_autocomplete() {
265        let registry = FunctionRegistry::new();
266
267        let mass_functions = registry.autocomplete("MASS");
268        assert!(!mass_functions.is_empty());
269
270        // Should include MASS_EARTH, MASS_SUN, etc.
271        let names: Vec<&str> = mass_functions.iter().map(|sig| sig.name).collect();
272        assert!(names.contains(&"MASS_EARTH"));
273        assert!(names.contains(&"MASS_SUN"));
274    }
275}
sql_cli/sql/functions/mod.rs

sql_cli/sql/functions/
mod.rs
