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

sql_cli/sql/functions/
mod.rs
