sql_cli/sql/functions/
mod.rs

1use anyhow::{anyhow, Result};
2use std::collections::HashMap;
3use std::fmt;
4use std::sync::Arc;
5
6use crate::data::datatable::DataValue;
7
8pub mod astronomy;
9pub mod chemistry;
10pub mod comparison;
11pub mod constants;
12pub mod mathematics;
13pub mod string_methods;
14
15// Re-export MethodFunction trait
16pub use string_methods::MethodFunction;
17
18/// Category of SQL functions for organization and discovery
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum FunctionCategory {
21    Constant,     // Mathematical and physical constants
22    Mathematical, // Mathematical operations
23    Astronomical, // Astronomical constants and calculations
24    Chemical,     // Chemical elements and properties
25    Date,         // Date/time operations
26    String,       // String manipulation
27    Aggregate,    // Aggregation functions
28}
29
30impl fmt::Display for FunctionCategory {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        match self {
33            FunctionCategory::Constant => write!(f, "Constant"),
34            FunctionCategory::Mathematical => write!(f, "Mathematical"),
35            FunctionCategory::Astronomical => write!(f, "Astronomical"),
36            FunctionCategory::Chemical => write!(f, "Chemical"),
37            FunctionCategory::Date => write!(f, "Date"),
38            FunctionCategory::String => write!(f, "String"),
39            FunctionCategory::Aggregate => write!(f, "Aggregate"),
40        }
41    }
42}
43
44/// Describes the number of arguments a function accepts
45#[derive(Debug, Clone)]
46pub enum ArgCount {
47    /// Exactly n arguments
48    Fixed(usize),
49    /// Between min and max arguments (inclusive)
50    Range(usize, usize),
51    /// Any number of arguments
52    Variadic,
53}
54
55impl ArgCount {
56    pub fn is_valid(&self, count: usize) -> bool {
57        match self {
58            ArgCount::Fixed(n) => count == *n,
59            ArgCount::Range(min, max) => count >= *min && count <= *max,
60            ArgCount::Variadic => true,
61        }
62    }
63
64    pub fn description(&self) -> String {
65        match self {
66            ArgCount::Fixed(0) => "no arguments".to_string(),
67            ArgCount::Fixed(1) => "1 argument".to_string(),
68            ArgCount::Fixed(n) => format!("{} arguments", n),
69            ArgCount::Range(min, max) => format!("{} to {} arguments", min, max),
70            ArgCount::Variadic => "any number of arguments".to_string(),
71        }
72    }
73}
74
75/// Signature of a SQL function including metadata
76#[derive(Debug, Clone)]
77pub struct FunctionSignature {
78    pub name: &'static str,
79    pub category: FunctionCategory,
80    pub arg_count: ArgCount,
81    pub description: &'static str,
82    pub returns: &'static str,
83    pub examples: Vec<&'static str>,
84}
85
86/// Trait that all SQL functions must implement
87pub trait SqlFunction: Send + Sync {
88    /// Get the function's signature and metadata
89    fn signature(&self) -> FunctionSignature;
90
91    /// Evaluate the function with the given arguments
92    fn evaluate(&self, args: &[DataValue]) -> Result<DataValue>;
93
94    /// Validate arguments before evaluation (default implementation checks count)
95    fn validate_args(&self, args: &[DataValue]) -> Result<()> {
96        let sig = self.signature();
97        if !sig.arg_count.is_valid(args.len()) {
98            return Err(anyhow!(
99                "{}() expects {}, got {}",
100                sig.name,
101                sig.arg_count.description(),
102                args.len()
103            ));
104        }
105        Ok(())
106    }
107}
108
109/// Registry for all SQL functions
110pub struct FunctionRegistry {
111    functions: HashMap<String, Box<dyn SqlFunction>>,
112    by_category: HashMap<FunctionCategory, Vec<String>>,
113    methods: HashMap<String, Arc<dyn MethodFunction>>,
114}
115
116impl FunctionRegistry {
117    /// Create a new registry with all built-in functions
118    pub fn new() -> Self {
119        let mut registry = Self {
120            functions: HashMap::new(),
121            by_category: HashMap::new(),
122            methods: HashMap::new(),
123        };
124
125        // Register all built-in functions
126        registry.register_constants();
127        registry.register_astronomical_functions();
128        registry.register_chemical_functions();
129        registry.register_mathematical_functions();
130        registry.register_string_methods();
131        registry.register_comparison_functions();
132
133        registry
134    }
135
136    /// Register a function in the registry
137    pub fn register(&mut self, func: Box<dyn SqlFunction>) {
138        let sig = func.signature();
139        let name = sig.name.to_uppercase();
140        let category = sig.category;
141
142        // Add to main registry
143        self.functions.insert(name.clone(), func);
144
145        // Add to category index
146        self.by_category
147            .entry(category)
148            .or_insert_with(Vec::new)
149            .push(name);
150    }
151
152    /// Get a function by name (case-insensitive)
153    pub fn get(&self, name: &str) -> Option<&dyn SqlFunction> {
154        self.functions.get(&name.to_uppercase()).map(|b| b.as_ref())
155    }
156
157    /// Check if a function exists
158    pub fn contains(&self, name: &str) -> bool {
159        self.functions.contains_key(&name.to_uppercase())
160    }
161
162    /// Get all functions matching a prefix (for autocomplete)
163    pub fn autocomplete(&self, prefix: &str) -> Vec<FunctionSignature> {
164        let prefix_upper = prefix.to_uppercase();
165        self.functions
166            .iter()
167            .filter(|(name, _)| name.starts_with(&prefix_upper))
168            .map(|(_, func)| func.signature())
169            .collect()
170    }
171
172    /// Get all functions in a category
173    pub fn get_by_category(&self, category: FunctionCategory) -> Vec<FunctionSignature> {
174        self.by_category
175            .get(&category)
176            .map(|names| {
177                names
178                    .iter()
179                    .filter_map(|name| self.functions.get(name))
180                    .map(|func| func.signature())
181                    .collect()
182            })
183            .unwrap_or_default()
184    }
185
186    /// Get all available functions
187    pub fn all_functions(&self) -> Vec<FunctionSignature> {
188        self.functions
189            .values()
190            .map(|func| func.signature())
191            .collect()
192    }
193
194    /// Register a method function
195    pub fn register_method(&mut self, method: Arc<dyn MethodFunction>) {
196        let method_name = method.method_name().to_uppercase();
197        self.methods.insert(method_name, method);
198    }
199
200    /// Get a method function by name
201    pub fn get_method(&self, name: &str) -> Option<Arc<dyn MethodFunction>> {
202        // Try exact match first
203        if let Some(method) = self.methods.get(&name.to_uppercase()) {
204            return Some(Arc::clone(method));
205        }
206
207        // Try to find a method that handles this name
208        for method in self.methods.values() {
209            if method.handles_method(name) {
210                return Some(Arc::clone(method));
211            }
212        }
213
214        None
215    }
216
217    /// Check if a method exists
218    pub fn has_method(&self, name: &str) -> bool {
219        self.get_method(name).is_some()
220    }
221
222    /// Generate markdown documentation for all functions
223    pub fn generate_markdown_docs(&self) -> String {
224        use std::fmt::Write;
225        let mut doc = String::new();
226
227        writeln!(&mut doc, "# SQL CLI Function Reference\n").unwrap();
228        writeln!(
229            &mut doc,
230            "This document is auto-generated from the function registry.\n"
231        )
232        .unwrap();
233
234        // Get all categories in a deterministic order
235        let mut categories: Vec<FunctionCategory> = self.by_category.keys().copied().collect();
236        categories.sort_by_key(|c| format!("{:?}", c));
237
238        for category in categories {
239            let functions = self.get_by_category(category);
240            if functions.is_empty() {
241                continue;
242            }
243
244            writeln!(&mut doc, "## {} Functions\n", category).unwrap();
245
246            // Sort functions by name for consistent output
247            let mut functions = functions;
248            functions.sort_by_key(|f| f.name);
249
250            for func in functions {
251                writeln!(&mut doc, "### {}()\n", func.name).unwrap();
252                writeln!(&mut doc, "**Description:** {}\n", func.description).unwrap();
253                writeln!(
254                    &mut doc,
255                    "**Arguments:** {}\n",
256                    func.arg_count.description()
257                )
258                .unwrap();
259                writeln!(&mut doc, "**Returns:** {}\n", func.returns).unwrap();
260
261                if !func.examples.is_empty() {
262                    writeln!(&mut doc, "**Examples:**").unwrap();
263                    writeln!(&mut doc, "```sql").unwrap();
264                    for example in &func.examples {
265                        writeln!(&mut doc, "{}", example).unwrap();
266                    }
267                    writeln!(&mut doc, "```\n").unwrap();
268                }
269            }
270        }
271
272        doc
273    }
274
275    /// Generate help text for a specific function
276    pub fn generate_function_help(&self, name: &str) -> Option<String> {
277        self.get(name).map(|func| {
278            let sig = func.signature();
279            let mut help = String::new();
280            use std::fmt::Write;
281
282            writeln!(&mut help, "Function: {}()", sig.name).unwrap();
283            writeln!(&mut help, "Category: {}", sig.category).unwrap();
284            writeln!(&mut help, "Description: {}", sig.description).unwrap();
285            writeln!(&mut help, "Arguments: {}", sig.arg_count.description()).unwrap();
286            writeln!(&mut help, "Returns: {}", sig.returns).unwrap();
287
288            if !sig.examples.is_empty() {
289                writeln!(&mut help, "\nExamples:").unwrap();
290                for example in &sig.examples {
291                    writeln!(&mut help, "  {}", example).unwrap();
292                }
293            }
294
295            help
296        })
297    }
298
299    /// List all available functions with brief descriptions
300    pub fn list_functions(&self) -> String {
301        use std::fmt::Write;
302        let mut list = String::new();
303
304        writeln!(&mut list, "Available SQL Functions:\n").unwrap();
305
306        let mut categories: Vec<FunctionCategory> = self.by_category.keys().copied().collect();
307        categories.sort_by_key(|c| format!("{:?}", c));
308
309        for category in categories {
310            let functions = self.get_by_category(category);
311            if functions.is_empty() {
312                continue;
313            }
314
315            writeln!(&mut list, "{} Functions:", category).unwrap();
316
317            let mut functions = functions;
318            functions.sort_by_key(|f| f.name);
319
320            for func in functions {
321                writeln!(
322                    &mut list,
323                    "  {:20} - {}",
324                    format!("{}()", func.name),
325                    func.description
326                )
327                .unwrap();
328            }
329            writeln!(&mut list).unwrap();
330        }
331
332        list
333    }
334
335    /// Register constant functions
336    fn register_constants(&mut self) {
337        use constants::*;
338
339        self.register(Box::new(PiFunction));
340        self.register(Box::new(EFunction));
341        self.register(Box::new(MeFunction)); // Mass of electron
342        self.register(Box::new(MassElectronFunction)); // Alias for ME
343    }
344
345    /// Register astronomical functions
346    fn register_astronomical_functions(&mut self) {
347        use astronomy::*;
348
349        self.register(Box::new(MassEarthFunction));
350        self.register(Box::new(MassSunFunction));
351        self.register(Box::new(MassMoonFunction));
352        self.register(Box::new(AuFunction)); // Astronomical unit
353        self.register(Box::new(LightYearFunction));
354        self.register(Box::new(ParsecFunction));
355
356        // Planetary masses
357        self.register(Box::new(MassMercuryFunction));
358        self.register(Box::new(MassVenusFunction));
359        self.register(Box::new(MassMarsFunction));
360        self.register(Box::new(MassJupiterFunction));
361        self.register(Box::new(MassSaturnFunction));
362        self.register(Box::new(MassUranusFunction));
363        self.register(Box::new(MassNeptuneFunction));
364
365        // Solar body radius functions
366        self.register(Box::new(RadiusSunFunction));
367        self.register(Box::new(RadiusEarthFunction));
368        self.register(Box::new(RadiusMoonFunction));
369        self.register(Box::new(RadiusMercuryFunction));
370        self.register(Box::new(RadiusVenusFunction));
371        self.register(Box::new(RadiusMarsFunction));
372        self.register(Box::new(RadiusJupiterFunction));
373        self.register(Box::new(RadiusSaturnFunction));
374        self.register(Box::new(RadiusUranusFunction));
375        self.register(Box::new(RadiusNeptuneFunction));
376    }
377
378    /// Register chemical functions
379    fn register_chemical_functions(&mut self) {
380        use chemistry::*;
381
382        self.register(Box::new(AvogadroFunction));
383        self.register(Box::new(AtomicMassFunction));
384        self.register(Box::new(AtomicNumberFunction));
385    }
386
387    /// Register string method functions
388    fn register_string_methods(&mut self) {
389        string_methods::register_string_methods(self);
390    }
391
392    /// Register comparison functions
393    fn register_comparison_functions(&mut self) {
394        comparison::register_comparison_functions(self);
395    }
396
397    /// Register mathematical functions
398    fn register_mathematical_functions(&mut self) {
399        use mathematics::*;
400
401        self.register(Box::new(PrimeFunction));
402        self.register(Box::new(IsPrimeFunction));
403        self.register(Box::new(PrimeCountFunction));
404        self.register(Box::new(NextPrimeFunction));
405        self.register(Box::new(PrevPrimeFunction));
406    }
407}
408
409impl Default for FunctionRegistry {
410    fn default() -> Self {
411        Self::new()
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn test_registry_creation() {
421        let registry = FunctionRegistry::new();
422
423        // Check that some known functions exist
424        assert!(registry.contains("PI"));
425        assert!(registry.contains("MASS_EARTH"));
426        assert!(registry.contains("ME"));
427    }
428
429    #[test]
430    fn test_case_insensitive_lookup() {
431        let registry = FunctionRegistry::new();
432
433        assert!(registry.get("pi").is_some());
434        assert!(registry.get("PI").is_some());
435        assert!(registry.get("Pi").is_some());
436    }
437
438    #[test]
439    fn test_autocomplete() {
440        let registry = FunctionRegistry::new();
441
442        let mass_functions = registry.autocomplete("MASS");
443        assert!(!mass_functions.is_empty());
444
445        // Should include MASS_EARTH, MASS_SUN, etc.
446        let names: Vec<&str> = mass_functions.iter().map(|sig| sig.name).collect();
447        assert!(names.contains(&"MASS_EARTH"));
448        assert!(names.contains(&"MASS_SUN"));
449    }
450}
sql_cli/sql/functions/mod.rs

sql_cli/sql/functions/
mod.rs
