1use llkv_result::{Error, Result};
8use rustc_hash::FxHashMap;
9use sqlparser::ast::{
10 Expr as SqlExpr, FunctionArg, FunctionArgExpr, GroupByExpr, ObjectName, ObjectNamePart, Select,
11 SelectItem, SelectItemQualifiedWildcardKind, TableAlias, TableFactor, UnaryOperator, Value,
12 ValueWithSpan,
13};
14
15use crate::PlanValue;
16
17pub fn plan_value_from_sql_expr(expr: &SqlExpr) -> Result<PlanValue> {
42 match expr {
43 SqlExpr::Value(value) => plan_value_from_sql_value(value),
44 SqlExpr::UnaryOp {
45 op: UnaryOperator::Minus,
46 expr,
47 } => match plan_value_from_sql_expr(expr)? {
48 PlanValue::Integer(v) => Ok(PlanValue::Integer(-v)),
49 PlanValue::Float(v) => Ok(PlanValue::Float(-v)),
50 PlanValue::Null | PlanValue::String(_) | PlanValue::Struct(_) => Err(
51 Error::InvalidArgumentError("cannot negate non-numeric literal".into()),
52 ),
53 },
54 SqlExpr::UnaryOp {
55 op: UnaryOperator::Plus,
56 expr,
57 } => plan_value_from_sql_expr(expr),
58 SqlExpr::Nested(inner) => plan_value_from_sql_expr(inner),
59 SqlExpr::Dictionary(fields) => {
60 let mut map = FxHashMap::with_capacity_and_hasher(fields.len(), Default::default());
61 for field in fields {
62 let key = field.key.value.clone();
63 let value = plan_value_from_sql_expr(&field.value)?;
64 map.insert(key, value);
65 }
66 Ok(PlanValue::Struct(map))
67 }
68 other => Err(Error::InvalidArgumentError(format!(
69 "unsupported literal expression: {other:?}"
70 ))),
71 }
72}
73
74pub fn plan_value_from_sql_value(value: &ValueWithSpan) -> Result<PlanValue> {
98 match &value.value {
99 Value::Null => Ok(PlanValue::Null),
100 Value::Number(text, _) => {
101 if text.contains(['.', 'e', 'E']) {
102 let parsed = text.parse::<f64>().map_err(|err| {
103 Error::InvalidArgumentError(format!("invalid float literal: {err}"))
104 })?;
105 Ok(PlanValue::Float(parsed))
106 } else {
107 let parsed = text.parse::<i64>().map_err(|err| {
108 Error::InvalidArgumentError(format!("invalid integer literal: {err}"))
109 })?;
110 Ok(PlanValue::Integer(parsed))
111 }
112 }
113 Value::Boolean(_) => Err(Error::InvalidArgumentError(
114 "BOOLEAN literals are not supported yet".into(),
115 )),
116 other => {
117 if let Some(text) = other.clone().into_string() {
118 Ok(PlanValue::String(text))
119 } else {
120 Err(Error::InvalidArgumentError(format!(
121 "unsupported literal: {other:?}"
122 )))
123 }
124 }
125 }
126}
127
128#[derive(Clone)]
134pub struct RangeSelectRows {
135 rows: Vec<Vec<PlanValue>>,
136}
137
138impl RangeSelectRows {
139 pub fn into_rows(self) -> Vec<Vec<PlanValue>> {
141 self.rows
142 }
143}
144
145#[derive(Clone)]
146enum RangeProjection {
147 Column,
148 Literal(PlanValue),
149}
150
151#[derive(Clone)]
152struct RangeSpec {
153 start: i64,
154 #[allow(dead_code)] end: i64,
156 row_count: usize,
157 column_name_lower: String,
158 table_alias_lower: Option<String>,
159}
160
161impl RangeSpec {
162 fn matches_identifier(&self, ident: &str) -> bool {
163 let lower = ident.to_ascii_lowercase();
164 lower == self.column_name_lower || lower == "range"
165 }
166
167 fn matches_table_alias(&self, ident: &str) -> bool {
168 let lower = ident.to_ascii_lowercase();
169 match &self.table_alias_lower {
170 Some(alias) => lower == *alias,
171 None => lower == "range",
172 }
173 }
174
175 fn matches_object_name(&self, name: &ObjectName) -> bool {
176 if name.0.len() != 1 {
177 return false;
178 }
179 match &name.0[0] {
180 ObjectNamePart::Identifier(ident) => self.matches_table_alias(&ident.value),
181 _ => false,
182 }
183 }
184}
185
186pub fn extract_rows_from_range(select: &Select) -> Result<Option<RangeSelectRows>> {
202 let spec = match parse_range_spec(select)? {
203 Some(spec) => spec,
204 None => return Ok(None),
205 };
206
207 if select.selection.is_some() {
208 return Err(Error::InvalidArgumentError(
209 "WHERE clauses are not supported for range() SELECT statements".into(),
210 ));
211 }
212 if select.having.is_some()
213 || !select.named_window.is_empty()
214 || select.qualify.is_some()
215 || select.distinct.is_some()
216 || select.top.is_some()
217 || select.into.is_some()
218 || select.prewhere.is_some()
219 || !select.lateral_views.is_empty()
220 || select.value_table_mode.is_some()
221 || !group_by_is_empty(&select.group_by)
222 {
223 return Err(Error::InvalidArgumentError(
224 "advanced SELECT clauses are not supported for range() SELECT statements".into(),
225 ));
226 }
227
228 let mut projections: Vec<RangeProjection> = Vec::with_capacity(select.projection.len());
229
230 if select.projection.is_empty() {
232 projections.push(RangeProjection::Column);
233 } else {
234 for item in &select.projection {
235 let projection = match item {
236 SelectItem::Wildcard(_) => RangeProjection::Column,
237 SelectItem::QualifiedWildcard(kind, _) => match kind {
238 SelectItemQualifiedWildcardKind::ObjectName(object_name) => {
239 if spec.matches_object_name(object_name) {
240 RangeProjection::Column
241 } else {
242 return Err(Error::InvalidArgumentError(
243 "qualified wildcard must reference the range() source".into(),
244 ));
245 }
246 }
247 SelectItemQualifiedWildcardKind::Expr(_) => {
248 return Err(Error::InvalidArgumentError(
249 "expression-qualified wildcards are not supported for range() SELECT statements".into(),
250 ));
251 }
252 },
253 SelectItem::UnnamedExpr(expr) => build_range_projection_expr(expr, &spec)?,
254 SelectItem::ExprWithAlias { expr, .. } => build_range_projection_expr(expr, &spec)?,
255 };
256 projections.push(projection);
257 }
258 }
259
260 let mut rows: Vec<Vec<PlanValue>> = Vec::with_capacity(spec.row_count);
261 for idx in 0..spec.row_count {
262 let mut row: Vec<PlanValue> = Vec::with_capacity(projections.len());
263 let value = spec.start + (idx as i64);
264 for projection in &projections {
265 match projection {
266 RangeProjection::Column => row.push(PlanValue::Integer(value)),
267 RangeProjection::Literal(value) => row.push(value.clone()),
268 }
269 }
270 rows.push(row);
271 }
272
273 Ok(Some(RangeSelectRows { rows }))
274}
275
276fn build_range_projection_expr(expr: &SqlExpr, spec: &RangeSpec) -> Result<RangeProjection> {
277 match expr {
278 SqlExpr::Identifier(ident) => {
279 if spec.matches_identifier(&ident.value) {
280 Ok(RangeProjection::Column)
281 } else {
282 Err(Error::InvalidArgumentError(format!(
283 "unknown column '{}' in range() SELECT",
284 ident.value
285 )))
286 }
287 }
288 SqlExpr::CompoundIdentifier(parts) => {
289 if parts.len() == 2
290 && spec.matches_table_alias(&parts[0].value)
291 && spec.matches_identifier(&parts[1].value)
292 {
293 Ok(RangeProjection::Column)
294 } else {
295 Err(Error::InvalidArgumentError(
296 "compound identifiers must reference the range() source".into(),
297 ))
298 }
299 }
300 SqlExpr::Wildcard(_) | SqlExpr::QualifiedWildcard(_, _) => unreachable!(),
301 other => Ok(RangeProjection::Literal(plan_value_from_sql_expr(other)?)),
302 }
303}
304
305fn parse_range_spec(select: &Select) -> Result<Option<RangeSpec>> {
306 if select.from.len() != 1 {
307 return Ok(None);
308 }
309 let item = &select.from[0];
310 if !item.joins.is_empty() {
311 return Err(Error::InvalidArgumentError(
312 "JOIN clauses are not supported for range() SELECT statements".into(),
313 ));
314 }
315
316 match &item.relation {
317 TableFactor::Function {
318 lateral,
319 name,
320 args,
321 alias,
322 } => {
323 if *lateral {
324 return Err(Error::InvalidArgumentError(
325 "LATERAL range() is not supported".into(),
326 ));
327 }
328 parse_range_spec_from_args(name, args, alias)
329 }
330 TableFactor::Table {
331 name,
332 alias,
333 args: Some(table_args),
334 with_ordinality,
335 ..
336 } => {
337 if *with_ordinality {
338 return Err(Error::InvalidArgumentError(
339 "WITH ORDINALITY is not supported for range()".into(),
340 ));
341 }
342 if table_args.settings.is_some() {
343 return Err(Error::InvalidArgumentError(
344 "range() SETTINGS clause is not supported".into(),
345 ));
346 }
347 parse_range_spec_from_args(name, &table_args.args, alias)
348 }
349 _ => Ok(None),
350 }
351}
352
353fn parse_range_spec_from_args(
354 name: &ObjectName,
355 args: &[FunctionArg],
356 alias: &Option<TableAlias>,
357) -> Result<Option<RangeSpec>> {
358 if name.0.len() != 1 {
359 return Ok(None);
360 }
361 let func_name = match &name.0[0] {
362 ObjectNamePart::Identifier(ident) => ident.value.to_ascii_lowercase(),
363 _ => return Ok(None),
364 };
365 if func_name != "range" {
366 return Ok(None);
367 }
368
369 if args.is_empty() || args.len() > 2 {
370 return Err(Error::InvalidArgumentError(
371 "range() requires one or two arguments".into(),
372 ));
373 }
374
375 let extract_int = |arg: &FunctionArg| -> Result<i64> {
377 let arg_expr = match arg {
378 FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => expr,
379 FunctionArg::Unnamed(FunctionArgExpr::QualifiedWildcard(_))
380 | FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => {
381 return Err(Error::InvalidArgumentError(
382 "range() argument must be an integer literal".into(),
383 ));
384 }
385 FunctionArg::Named { .. } | FunctionArg::ExprNamed { .. } => {
386 return Err(Error::InvalidArgumentError(
387 "named arguments are not supported for range()".into(),
388 ));
389 }
390 };
391
392 let value = plan_value_from_sql_expr(arg_expr)?;
393 match value {
394 PlanValue::Integer(v) => Ok(v),
395 _ => Err(Error::InvalidArgumentError(
396 "range() argument must be an integer literal".into(),
397 )),
398 }
399 };
400
401 let (start, end, row_count) = if args.len() == 1 {
402 let count = extract_int(&args[0])?;
404 if count < 0 {
405 return Err(Error::InvalidArgumentError(
406 "range() argument must be non-negative".into(),
407 ));
408 }
409 (0, count, count as usize)
410 } else {
411 let start = extract_int(&args[0])?;
413 let end = extract_int(&args[1])?;
414 if end < start {
415 return Err(Error::InvalidArgumentError(
416 "range() end must be >= start".into(),
417 ));
418 }
419 let row_count = (end - start) as usize;
420 (start, end, row_count)
421 };
422
423 let column_name_lower = alias
424 .as_ref()
425 .and_then(|a| {
426 a.columns
427 .first()
428 .map(|col| col.name.value.to_ascii_lowercase())
429 })
430 .unwrap_or_else(|| "range".to_string());
431 let table_alias_lower = alias.as_ref().map(|a| a.name.value.to_ascii_lowercase());
432
433 Ok(Some(RangeSpec {
434 start,
435 end,
436 row_count,
437 column_name_lower,
438 table_alias_lower,
439 }))
440}
441
442fn group_by_is_empty(expr: &GroupByExpr) -> bool {
443 matches!(
444 expr,
445 GroupByExpr::Expressions(exprs, modifiers)
446 if exprs.is_empty() && modifiers.is_empty()
447 )
448}
449
450#[cfg(test)]
451mod tests {
452 use super::*;
453 use sqlparser::ast::{Expr as SqlExpr, Value, ValueWithSpan};
454
455 fn value_with_span(v: Value) -> ValueWithSpan {
456 ValueWithSpan {
457 value: v,
458 span: sqlparser::tokenizer::Span::empty(),
459 }
460 }
461
462 #[test]
463 fn test_null_value() {
464 let value = value_with_span(Value::Null);
465 assert_eq!(plan_value_from_sql_value(&value).unwrap(), PlanValue::Null);
466 }
467
468 #[test]
469 fn test_integer_value() {
470 let value = value_with_span(Value::Number("42".to_string(), false));
471 assert_eq!(
472 plan_value_from_sql_value(&value).unwrap(),
473 PlanValue::Integer(42)
474 );
475 }
476
477 #[test]
478 fn test_negative_integer() {
479 let value = value_with_span(Value::Number("42".to_string(), false));
480 let expr = SqlExpr::UnaryOp {
481 op: UnaryOperator::Minus,
482 expr: Box::new(SqlExpr::Value(value)),
483 };
484 assert_eq!(
485 plan_value_from_sql_expr(&expr).unwrap(),
486 PlanValue::Integer(-42)
487 );
488 }
489
490 #[test]
491 #[allow(clippy::approx_constant)]
492 fn test_float_value() {
493 let value = value_with_span(Value::Number("3.14".to_string(), false));
494 assert_eq!(
495 plan_value_from_sql_value(&value).unwrap(),
496 PlanValue::Float(3.14)
497 );
498 }
499
500 #[test]
501 fn test_string_value() {
502 let value = value_with_span(Value::SingleQuotedString("hello".to_string()));
503 assert_eq!(
504 plan_value_from_sql_value(&value).unwrap(),
505 PlanValue::String("hello".to_string())
506 );
507 }
508
509 #[test]
510 fn test_nested_expression() {
511 let value = value_with_span(Value::Number("100".to_string(), false));
512 let expr = SqlExpr::Nested(Box::new(SqlExpr::Value(value)));
513 assert_eq!(
514 plan_value_from_sql_expr(&expr).unwrap(),
515 PlanValue::Integer(100)
516 );
517 }
518
519 #[test]
520 fn test_plus_operator() {
521 let value = value_with_span(Value::Number("50".to_string(), false));
522 let expr = SqlExpr::UnaryOp {
523 op: UnaryOperator::Plus,
524 expr: Box::new(SqlExpr::Value(value)),
525 };
526 assert_eq!(
527 plan_value_from_sql_expr(&expr).unwrap(),
528 PlanValue::Integer(50)
529 );
530 }
531
532 #[test]
533 fn test_cannot_negate_string() {
534 let value = value_with_span(Value::SingleQuotedString("test".to_string()));
535 let expr = SqlExpr::UnaryOp {
536 op: UnaryOperator::Minus,
537 expr: Box::new(SqlExpr::Value(value)),
538 };
539 assert!(plan_value_from_sql_expr(&expr).is_err());
540 }
541
542 #[test]
543 fn test_boolean_not_supported() {
544 let value = value_with_span(Value::Boolean(true));
545 assert!(plan_value_from_sql_value(&value).is_err());
546 }
547}