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

sql_cli/sql/functions/
mod.rs
