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

sql_cli/sql/functions/
mod.rs
