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

sql_cli/sql/functions/
mod.rs
