Skip to main content

sql_cli/sql/generators/
mod.rs

1use crate::data::datatable::{DataColumn, DataRow, DataTable, DataValue};
2use anyhow::Result;
3use std::sync::Arc;
4
5pub mod ascii_art;
6pub mod file_readers;
7pub mod literal_generators;
8pub mod math_generators;
9pub mod prime_generators;
10pub mod random_generators;
11pub mod sequence_generators;
12pub mod string_generators;
13
14/// Trait for table-generating functions that produce rows dynamically
15pub trait TableGenerator: Send + Sync {
16    /// Get the name of the generator function (e.g., "GENERATE_PRIMES")
17    fn name(&self) -> &str;
18
19    /// Get the column definitions for the generated table
20    fn columns(&self) -> Vec<DataColumn>;
21
22    /// Generate the table based on the provided arguments
23    /// Arguments are evaluated expressions from the SQL query
24    fn generate(&self, args: Vec<DataValue>) -> Result<Arc<DataTable>>;
25
26    /// Get a description of what this generator does
27    fn description(&self) -> &str;
28
29    /// Get the expected number of arguments
30    fn arg_count(&self) -> usize;
31}
32
33/// Registry for table generator functions
34pub struct GeneratorRegistry {
35    generators: std::collections::HashMap<String, Box<dyn TableGenerator>>,
36}
37
38impl GeneratorRegistry {
39    pub fn new() -> Self {
40        let mut registry = Self {
41            generators: std::collections::HashMap::new(),
42        };
43        registry.register_default_generators();
44        registry
45    }
46
47    fn register_default_generators(&mut self) {
48        use ascii_art::{AsciiArt, Banner, BigText};
49        use file_readers::{Grep, ReadCsv, ReadJson, ReadJsonl, ReadStdin, ReadText, ReadWords};
50        use literal_generators::{Array, Values};
51        use math_generators::{Collatz, Factorials, PascalTriangle, Squares, TriangularNumbers};
52        use prime_generators::{Fibonacci, GeneratePrimes, PrimeFactors};
53        use random_generators::{GenerateUUIDs, RandomFloats, RandomIntegers};
54        use sequence_generators::{Dates, Range, Series};
55        use string_generators::{Chars, Lines, Split, Tokenize};
56
57        // Literal value generators
58        self.register(Box::new(Values));
59        self.register(Box::new(Array));
60
61        // Sequence generators
62        self.register(Box::new(Range));
63        self.register(Box::new(Series));
64        self.register(Box::new(Dates));
65
66        // String generators
67        self.register(Box::new(Split));
68        self.register(Box::new(Tokenize));
69        self.register(Box::new(Chars));
70        self.register(Box::new(Lines));
71
72        // File readers
73        self.register(Box::new(ReadText));
74        self.register(Box::new(ReadWords));
75        self.register(Box::new(ReadJsonl));
76        self.register(Box::new(ReadJson));
77        self.register(Box::new(ReadCsv));
78        self.register(Box::new(ReadStdin));
79        self.register(Box::new(Grep));
80
81        // Prime and number theory generators
82        self.register(Box::new(GeneratePrimes));
83        self.register(Box::new(PrimeFactors));
84        self.register(Box::new(Fibonacci));
85
86        // Mathematical sequence generators
87        self.register(Box::new(Collatz));
88        self.register(Box::new(PascalTriangle));
89        self.register(Box::new(TriangularNumbers));
90        self.register(Box::new(Squares));
91        self.register(Box::new(Factorials));
92
93        // Random generators
94        self.register(Box::new(RandomIntegers));
95        self.register(Box::new(RandomFloats));
96        self.register(Box::new(GenerateUUIDs));
97
98        // ASCII Art generators
99        self.register(Box::new(AsciiArt));
100        self.register(Box::new(BigText));
101        self.register(Box::new(Banner));
102    }
103
104    pub fn register(&mut self, generator: Box<dyn TableGenerator>) {
105        self.generators
106            .insert(generator.name().to_uppercase(), generator);
107    }
108
109    pub fn get(&self, name: &str) -> Option<&Box<dyn TableGenerator>> {
110        self.generators.get(&name.to_uppercase())
111    }
112
113    pub fn list(&self) -> Vec<&str> {
114        let mut names: Vec<&str> = self.generators.keys().map(|s| s.as_str()).collect();
115        names.sort();
116        names
117    }
118
119    pub fn list_generators_formatted(&self) -> String {
120        let mut output = String::new();
121        output.push_str("=== Available Generator Functions ===\n\n");
122
123        // Group generators by category
124        let mut sequence_gens = Vec::new();
125        let mut string_gens = Vec::new();
126        let mut file_gens = Vec::new();
127        let mut math_gens = Vec::new();
128        let mut random_gens = Vec::new();
129        let mut utility_gens = Vec::new();
130        let mut ascii_gens = Vec::new();
131
132        for (name, gen) in &self.generators {
133            let entry = format!("  {} - {}", name, gen.description());
134
135            if name == "RANGE" || name == "SERIES" || name == "DATES" {
136                sequence_gens.push(entry);
137            } else if name == "SPLIT" || name == "TOKENIZE" || name == "CHARS" || name == "LINES" {
138                string_gens.push(entry);
139            } else if name == "READ_TEXT"
140                || name == "READ_WORDS"
141                || name == "READ_JSONL"
142                || name == "READ_JSON"
143                || name == "READ_CSV"
144                || name == "READ_STDIN"
145                || name == "GREP"
146            {
147                file_gens.push(entry);
148            } else if name == "ASCII_ART" || name == "BIG_TEXT" || name == "BANNER" {
149                ascii_gens.push(entry);
150            } else if name.starts_with("RANDOM_") {
151                random_gens.push(entry);
152            } else if name == "GENERATE_UUID" {
153                utility_gens.push(entry);
154            } else {
155                math_gens.push(entry);
156            }
157        }
158
159        if !sequence_gens.is_empty() {
160            sequence_gens.sort();
161            output.push_str("Sequence Generators:\n");
162            for entry in sequence_gens {
163                output.push_str(&format!("{}\n", entry));
164            }
165            output.push('\n');
166        }
167
168        if !string_gens.is_empty() {
169            string_gens.sort();
170            output.push_str("String Generators:\n");
171            for entry in string_gens {
172                output.push_str(&format!("{}\n", entry));
173            }
174            output.push('\n');
175        }
176
177        if !file_gens.is_empty() {
178            file_gens.sort();
179            output.push_str("File Readers:\n");
180            for entry in file_gens {
181                output.push_str(&format!("{}\n", entry));
182            }
183            output.push('\n');
184        }
185
186        if !math_gens.is_empty() {
187            math_gens.sort();
188            output.push_str("Mathematical Generators:\n");
189            for entry in math_gens {
190                output.push_str(&format!("{}\n", entry));
191            }
192            output.push('\n');
193        }
194
195        if !random_gens.is_empty() {
196            random_gens.sort();
197            output.push_str("Random Generators:\n");
198            for entry in random_gens {
199                output.push_str(&format!("{}\n", entry));
200            }
201            output.push('\n');
202        }
203
204        if !ascii_gens.is_empty() {
205            ascii_gens.sort();
206            output.push_str("ASCII Art Generators:\n");
207            for entry in ascii_gens {
208                output.push_str(&format!("{}\n", entry));
209            }
210            output.push('\n');
211        }
212
213        if !utility_gens.is_empty() {
214            utility_gens.sort();
215            output.push_str("Utility Generators:\n");
216            for entry in utility_gens {
217                output.push_str(&format!("{}\n", entry));
218            }
219            output.push('\n');
220        }
221
222        output.push_str("Use: SELECT * FROM <generator>(<args>)\n");
223        output.push_str("Example: SELECT * FROM GENERATE_PRIMES(100)\n");
224        output
225    }
226
227    pub fn get_generator_help(&self, name: &str) -> Option<String> {
228        self.generators.get(&name.to_uppercase()).map(|gen| {
229            let mut help = String::new();
230            help.push_str(&format!("=== {} ===\n\n", name.to_uppercase()));
231            help.push_str(&format!("Description: {}\n", gen.description()));
232            help.push_str(&format!(
233                "Arguments: {} argument(s) expected\n",
234                gen.arg_count()
235            ));
236            help.push_str("\nColumns:\n");
237            for col in gen.columns() {
238                help.push_str(&format!("  - {}\n", col.name));
239            }
240            help.push_str("\nExample:\n");
241            help.push_str(&format!("  SELECT * FROM {}(", name.to_uppercase()));
242
243            // Add example arguments based on the generator
244            match name.to_uppercase().as_str() {
245                "GENERATE_PRIMES" => help.push_str("100"),
246                "FIBONACCI" => help.push_str("20"),
247                "PRIME_FACTORS" => help.push_str("1260"),
248                "COLLATZ" => help.push_str("7"),
249                "PASCAL_TRIANGLE" => help.push_str("5"),
250                "TRIANGULAR" | "SQUARES" | "FACTORIALS" => help.push_str("10"),
251                "RANDOM_INT" => help.push_str("10, 1, 100, 42"),
252                "RANDOM_FLOAT" => help.push_str("10, 0, 1, 42"),
253                "GENERATE_UUID" => help.push_str("5"),
254                "READ_TEXT" => help.push_str("'data/file.txt'"),
255                "READ_WORDS" => help.push_str("'data/file.txt', 3, 'lower'"),
256                "READ_JSONL" => help.push_str("'data/events.jsonl'"),
257                "READ_CSV" => help.push_str("'data/trades.csv'"),
258                "READ_STDIN" => help.push_str("'^ERROR'"),
259                "GREP" => help.push_str("'data/file.txt', 'pattern'"),
260                "ASCII_ART" | "BIG_TEXT" => help.push_str("'SQL-CLI'"),
261                "BANNER" => help.push_str("'HELLO', '*'"),
262                _ => help.push_str("..."),
263            }
264            help.push_str(");\n");
265            help
266        })
267    }
268}
269
270/// Helper function to create a single-column table
271pub fn create_single_column_table(
272    name: &str,
273    column_name: &str,
274    values: Vec<DataValue>,
275) -> Arc<DataTable> {
276    let mut table = DataTable::new(name);
277    table.add_column(DataColumn::new(column_name));
278
279    for value in values {
280        table.add_row(DataRow::new(vec![value])).unwrap();
281    }
282
283    Arc::new(table)
284}
285
286/// Helper function to create a two-column table
287pub fn create_two_column_table(
288    name: &str,
289    col1_name: &str,
290    col2_name: &str,
291    rows: Vec<(DataValue, DataValue)>,
292) -> Arc<DataTable> {
293    let mut table = DataTable::new(name);
294    table.add_column(DataColumn::new(col1_name));
295    table.add_column(DataColumn::new(col2_name));
296
297    for (val1, val2) in rows {
298        table.add_row(DataRow::new(vec![val1, val2])).unwrap();
299    }
300
301    Arc::new(table)
302}