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

sql_cli/sql/functions/
mod.rs
