sql_cli/sql/generators/
random_generators.rs

1use super::TableGenerator;
2use crate::data::datatable::{DataColumn, DataRow, DataTable, DataValue};
3use anyhow::Result;
4use rand::rngs::StdRng;
5use rand::{Rng, SeedableRng};
6use std::sync::Arc;
7
8/// Generate random integers in a range
9pub struct RandomIntegers;
10
11impl TableGenerator for RandomIntegers {
12    fn name(&self) -> &str {
13        "RANDOM_INT"
14    }
15
16    fn columns(&self) -> Vec<DataColumn> {
17        vec![DataColumn::new("id"), DataColumn::new("value")]
18    }
19
20    fn generate(&self, args: Vec<DataValue>) -> Result<Arc<DataTable>> {
21        if args.len() < 3 || args.len() > 4 {
22            return Err(anyhow::anyhow!(
23                "RANDOM_INT expects 3-4 arguments (count, min, max, [seed])"
24            ));
25        }
26
27        let count = match &args[0] {
28            DataValue::Integer(n) => *n as usize,
29            DataValue::Float(f) => *f as usize,
30            _ => return Err(anyhow::anyhow!("RANDOM_INT count must be a number")),
31        };
32
33        let min = match &args[1] {
34            DataValue::Integer(n) => *n,
35            DataValue::Float(f) => *f as i64,
36            _ => return Err(anyhow::anyhow!("RANDOM_INT min must be a number")),
37        };
38
39        let max = match &args[2] {
40            DataValue::Integer(n) => *n,
41            DataValue::Float(f) => *f as i64,
42            _ => return Err(anyhow::anyhow!("RANDOM_INT max must be a number")),
43        };
44
45        if min > max {
46            return Err(anyhow::anyhow!("RANDOM_INT min must be <= max"));
47        }
48
49        // Optional seed for reproducibility
50        let mut rng: StdRng = if args.len() == 4 {
51            let seed = match &args[3] {
52                DataValue::Integer(n) => *n as u64,
53                DataValue::Float(f) => *f as u64,
54                _ => return Err(anyhow::anyhow!("RANDOM_INT seed must be a number")),
55            };
56            StdRng::seed_from_u64(seed)
57        } else {
58            StdRng::from_entropy()
59        };
60
61        let mut table = DataTable::new("random_int");
62        table.add_column(DataColumn::new("id"));
63        table.add_column(DataColumn::new("value"));
64
65        for i in 0..count {
66            let value = rng.gen_range(min..=max);
67            table
68                .add_row(DataRow::new(vec![
69                    DataValue::Integer(i as i64),
70                    DataValue::Integer(value),
71                ]))
72                .unwrap();
73        }
74
75        Ok(Arc::new(table))
76    }
77
78    fn description(&self) -> &str {
79        "Generate random integers in range [min, max]"
80    }
81
82    fn arg_count(&self) -> usize {
83        3 // or 4 with seed
84    }
85}
86
87/// Generate random floats in a range
88pub struct RandomFloats;
89
90impl TableGenerator for RandomFloats {
91    fn name(&self) -> &str {
92        "RANDOM_FLOAT"
93    }
94
95    fn columns(&self) -> Vec<DataColumn> {
96        vec![DataColumn::new("id"), DataColumn::new("value")]
97    }
98
99    fn generate(&self, args: Vec<DataValue>) -> Result<Arc<DataTable>> {
100        if args.len() < 3 || args.len() > 4 {
101            return Err(anyhow::anyhow!(
102                "RANDOM_FLOAT expects 3-4 arguments (count, min, max, [seed])"
103            ));
104        }
105
106        let count = match &args[0] {
107            DataValue::Integer(n) => *n as usize,
108            DataValue::Float(f) => *f as usize,
109            _ => return Err(anyhow::anyhow!("RANDOM_FLOAT count must be a number")),
110        };
111
112        let min = match &args[1] {
113            DataValue::Integer(n) => *n as f64,
114            DataValue::Float(f) => *f,
115            _ => return Err(anyhow::anyhow!("RANDOM_FLOAT min must be a number")),
116        };
117
118        let max = match &args[2] {
119            DataValue::Integer(n) => *n as f64,
120            DataValue::Float(f) => *f,
121            _ => return Err(anyhow::anyhow!("RANDOM_FLOAT max must be a number")),
122        };
123
124        if min > max {
125            return Err(anyhow::anyhow!("RANDOM_FLOAT min must be <= max"));
126        }
127
128        let mut rng: StdRng = if args.len() == 4 {
129            let seed = match &args[3] {
130                DataValue::Integer(n) => *n as u64,
131                DataValue::Float(f) => *f as u64,
132                _ => return Err(anyhow::anyhow!("RANDOM_FLOAT seed must be a number")),
133            };
134            StdRng::seed_from_u64(seed)
135        } else {
136            StdRng::from_entropy()
137        };
138
139        let mut table = DataTable::new("random_float");
140        table.add_column(DataColumn::new("id"));
141        table.add_column(DataColumn::new("value"));
142
143        for i in 0..count {
144            let value = rng.gen_range(min..=max);
145            table
146                .add_row(DataRow::new(vec![
147                    DataValue::Integer(i as i64),
148                    DataValue::Float(value),
149                ]))
150                .unwrap();
151        }
152
153        Ok(Arc::new(table))
154    }
155
156    fn description(&self) -> &str {
157        "Generate random floating-point numbers in range [min, max]"
158    }
159
160    fn arg_count(&self) -> usize {
161        3 // or 4 with seed
162    }
163}
164
165/// Generate UUIDs
166pub struct GenerateUUIDs;
167
168impl TableGenerator for GenerateUUIDs {
169    fn name(&self) -> &str {
170        "GENERATE_UUID"
171    }
172
173    fn columns(&self) -> Vec<DataColumn> {
174        vec![DataColumn::new("id"), DataColumn::new("uuid")]
175    }
176
177    fn generate(&self, args: Vec<DataValue>) -> Result<Arc<DataTable>> {
178        if args.len() != 1 {
179            return Err(anyhow::anyhow!("GENERATE_UUID expects 1 argument (count)"));
180        }
181
182        let count = match &args[0] {
183            DataValue::Integer(n) => *n as usize,
184            DataValue::Float(f) => *f as usize,
185            _ => return Err(anyhow::anyhow!("GENERATE_UUID count must be a number")),
186        };
187
188        let mut table = DataTable::new("uuids");
189        table.add_column(DataColumn::new("id"));
190        table.add_column(DataColumn::new("uuid"));
191
192        for i in 0..count {
193            let uuid = uuid::Uuid::new_v4().to_string();
194            table
195                .add_row(DataRow::new(vec![
196                    DataValue::Integer(i as i64),
197                    DataValue::String(uuid),
198                ]))
199                .unwrap();
200        }
201
202        Ok(Arc::new(table))
203    }
204
205    fn description(&self) -> &str {
206        "Generate UUID v4 strings"
207    }
208
209    fn arg_count(&self) -> usize {
210        1
211    }
212}