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

sql_cli/sql/functions/
mod.rs
