1use sqlparser::ast::{
2 DuplicateTreatment, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, JoinConstraint,
3 JoinOperator, LimitClause, OrderByKind, Query, Select, SelectItem, SetExpr, Statement,
4 TableFactor, TableWithJoins,
5};
6
7use crate::error::{Result, SQLRiteError};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum AggregateFn {
12 Count,
13 Sum,
14 Avg,
15 Min,
16 Max,
17}
18
19impl AggregateFn {
20 pub fn as_str(self) -> &'static str {
21 match self {
22 AggregateFn::Count => "COUNT",
23 AggregateFn::Sum => "SUM",
24 AggregateFn::Avg => "AVG",
25 AggregateFn::Min => "MIN",
26 AggregateFn::Max => "MAX",
27 }
28 }
29
30 fn from_name(name: &str) -> Option<Self> {
31 match name.to_ascii_lowercase().as_str() {
32 "count" => Some(AggregateFn::Count),
33 "sum" => Some(AggregateFn::Sum),
34 "avg" => Some(AggregateFn::Avg),
35 "min" => Some(AggregateFn::Min),
36 "max" => Some(AggregateFn::Max),
37 _ => None,
38 }
39 }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
44pub enum AggregateArg {
45 Star,
46 Column(String),
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct AggregateCall {
52 pub func: AggregateFn,
53 pub arg: AggregateArg,
54 pub distinct: bool,
56}
57
58impl AggregateCall {
59 pub fn display_name(&self) -> String {
63 let inner = match &self.arg {
64 AggregateArg::Star => "*".to_string(),
65 AggregateArg::Column(c) => {
66 if self.distinct {
67 format!("DISTINCT {c}")
68 } else {
69 c.clone()
70 }
71 }
72 };
73 format!("{}({inner})", self.func.as_str())
74 }
75}
76
77#[derive(Debug, Clone)]
79pub struct ProjectionItem {
80 pub kind: ProjectionKind,
81 pub alias: Option<String>,
83}
84
85impl ProjectionItem {
86 pub fn output_name(&self) -> String {
92 if let Some(a) = &self.alias {
93 return a.clone();
94 }
95 match &self.kind {
96 ProjectionKind::Column { name, .. } => name.clone(),
97 ProjectionKind::Aggregate(a) => a.display_name(),
98 }
99 }
100}
101
102#[derive(Debug, Clone)]
104pub enum ProjectionKind {
105 Column {
111 qualifier: Option<String>,
112 name: String,
113 },
114 Aggregate(AggregateCall),
116}
117
118#[derive(Debug, Clone)]
120pub enum Projection {
121 All,
123 Items(Vec<ProjectionItem>),
126}
127
128#[derive(Debug, Clone)]
136pub struct OrderByClause {
137 pub expr: Expr,
138 pub ascending: bool,
139}
140
141#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub enum JoinType {
148 Inner,
149 LeftOuter,
150 RightOuter,
151 FullOuter,
152}
153
154impl JoinType {
155 pub fn as_str(self) -> &'static str {
156 match self {
157 JoinType::Inner => "INNER",
158 JoinType::LeftOuter => "LEFT OUTER",
159 JoinType::RightOuter => "RIGHT OUTER",
160 JoinType::FullOuter => "FULL OUTER",
161 }
162 }
163}
164
165#[derive(Debug, Clone)]
170pub struct JoinClause {
171 pub join_type: JoinType,
172 pub right_table: String,
173 pub right_alias: Option<String>,
177 pub on: Expr,
180}
181
182#[derive(Debug, Clone)]
184pub struct SelectQuery {
185 pub table_name: String,
186 pub table_alias: Option<String>,
190 pub joins: Vec<JoinClause>,
193 pub projection: Projection,
194 pub selection: Option<Expr>,
196 pub order_by: Option<OrderByClause>,
197 pub limit: Option<usize>,
198 pub distinct: bool,
200 pub group_by: Vec<String>,
202}
203
204impl SelectQuery {
205 pub fn new(statement: &Statement) -> Result<Self> {
206 let Statement::Query(query) = statement else {
207 return Err(SQLRiteError::Internal(
208 "Error parsing SELECT: expected a Query statement".to_string(),
209 ));
210 };
211
212 let Query {
213 body,
214 order_by,
215 limit_clause,
216 ..
217 } = query.as_ref();
218
219 let SetExpr::Select(select) = body.as_ref() else {
220 return Err(SQLRiteError::NotImplemented(
221 "Only simple SELECT queries are supported (no UNION / VALUES / CTEs yet)"
222 .to_string(),
223 ));
224 };
225 let Select {
226 projection,
227 from,
228 selection,
229 distinct,
230 group_by,
231 having,
232 ..
233 } = select.as_ref();
234
235 let distinct_flag = match distinct {
239 None => false,
240 Some(sqlparser::ast::Distinct::Distinct) => true,
241 Some(sqlparser::ast::Distinct::All) => false,
242 Some(sqlparser::ast::Distinct::On(_)) => {
243 return Err(SQLRiteError::NotImplemented(
244 "SELECT DISTINCT ON (...) is not supported".to_string(),
245 ));
246 }
247 };
248 if having.is_some() {
249 return Err(SQLRiteError::NotImplemented(
250 "HAVING is not supported yet".to_string(),
251 ));
252 }
253 let group_by_cols: Vec<String> = match group_by {
258 sqlparser::ast::GroupByExpr::Expressions(exprs, _) => {
259 let mut out = Vec::with_capacity(exprs.len());
260 for e in exprs {
261 let col = match e {
262 Expr::Identifier(ident) => ident.value.clone(),
263 Expr::CompoundIdentifier(parts) => {
264 parts.last().map(|p| p.value.clone()).ok_or_else(|| {
265 SQLRiteError::Internal("empty compound identifier".to_string())
266 })?
267 }
268 other => {
269 return Err(SQLRiteError::NotImplemented(format!(
270 "GROUP BY only supports bare column references for now, got {other:?}"
271 )));
272 }
273 };
274 out.push(col);
275 }
276 out
277 }
278 _ => {
279 return Err(SQLRiteError::NotImplemented(
280 "GROUP BY ALL is not supported".to_string(),
281 ));
282 }
283 };
284
285 let (table_name, table_alias, joins) = extract_from_clause(from)?;
286 let projection = parse_projection(projection)?;
287 let order_by = parse_order_by(order_by.as_ref())?;
288 let limit = parse_limit(limit_clause.as_ref())?;
289
290 if !group_by_cols.is_empty()
294 && let Projection::Items(items) = &projection
295 {
296 for item in items {
297 if let ProjectionKind::Column { name: c, .. } = &item.kind
298 && !group_by_cols.contains(c)
299 {
300 return Err(SQLRiteError::Internal(format!(
301 "column '{c}' must appear in GROUP BY or be used in an aggregate function"
302 )));
303 }
304 }
305 }
306
307 if !joins.is_empty() {
312 let has_agg = matches!(
313 &projection,
314 Projection::Items(items)
315 if items.iter().any(|i| matches!(i.kind, ProjectionKind::Aggregate(_)))
316 );
317 if has_agg || !group_by_cols.is_empty() {
318 return Err(SQLRiteError::NotImplemented(
319 "GROUP BY / aggregate functions over JOIN results are not supported yet"
320 .to_string(),
321 ));
322 }
323 if distinct_flag {
324 return Err(SQLRiteError::NotImplemented(
325 "SELECT DISTINCT over JOIN results is not supported yet".to_string(),
326 ));
327 }
328 }
329
330 Ok(SelectQuery {
331 table_name,
332 table_alias,
333 joins,
334 projection,
335 selection: selection.clone(),
336 order_by,
337 limit,
338 distinct: distinct_flag,
339 group_by: group_by_cols,
340 })
341 }
342}
343
344fn extract_from_clause(
351 from: &[TableWithJoins],
352) -> Result<(String, Option<String>, Vec<JoinClause>)> {
353 if from.is_empty() {
354 return Err(SQLRiteError::Internal(
355 "SELECT requires a FROM clause".to_string(),
356 ));
357 }
358 if from.len() != 1 {
359 return Err(SQLRiteError::NotImplemented(
360 "comma-separated FROM lists are not supported — use explicit JOIN syntax".to_string(),
361 ));
362 }
363 let twj = &from[0];
364 let (table_name, table_alias) = extract_table_factor(&twj.relation)?;
365
366 let mut joins = Vec::with_capacity(twj.joins.len());
367 for j in &twj.joins {
368 let (right_table, right_alias) = extract_table_factor(&j.relation)?;
369 let (join_type, on_expr) = match &j.join_operator {
370 JoinOperator::Join(c) | JoinOperator::Inner(c) => (JoinType::Inner, parse_on(c)?),
372 JoinOperator::Left(c) | JoinOperator::LeftOuter(c) => {
373 (JoinType::LeftOuter, parse_on(c)?)
374 }
375 JoinOperator::Right(c) | JoinOperator::RightOuter(c) => {
376 (JoinType::RightOuter, parse_on(c)?)
377 }
378 JoinOperator::FullOuter(c) => (JoinType::FullOuter, parse_on(c)?),
379 other => {
380 return Err(SQLRiteError::NotImplemented(format!(
381 "join flavor {other:?} is not supported \
382 (only INNER / LEFT OUTER / RIGHT OUTER / FULL OUTER with ON)"
383 )));
384 }
385 };
386 joins.push(JoinClause {
387 join_type,
388 right_table,
389 right_alias,
390 on: on_expr,
391 });
392 }
393
394 Ok((table_name, table_alias, joins))
395}
396
397fn extract_table_factor(tf: &TableFactor) -> Result<(String, Option<String>)> {
398 match tf {
399 TableFactor::Table { name, alias, .. } => {
400 let table_name = name.to_string();
401 let alias_name = alias.as_ref().map(|a| a.name.value.clone());
402 if let Some(a) = alias.as_ref()
406 && !a.columns.is_empty()
407 {
408 return Err(SQLRiteError::NotImplemented(
409 "table alias column lists are not supported".to_string(),
410 ));
411 }
412 Ok((table_name, alias_name))
413 }
414 _ => Err(SQLRiteError::NotImplemented(
415 "only plain table references are supported in FROM / JOIN".to_string(),
416 )),
417 }
418}
419
420fn parse_on(constraint: &JoinConstraint) -> Result<Expr> {
421 match constraint {
422 JoinConstraint::On(expr) => Ok(expr.clone()),
423 JoinConstraint::Using(_) => Err(SQLRiteError::NotImplemented(
424 "JOIN ... USING (...) is not supported yet — use JOIN ... ON instead".to_string(),
425 )),
426 JoinConstraint::Natural => Err(SQLRiteError::NotImplemented(
427 "NATURAL JOIN is not supported".to_string(),
428 )),
429 JoinConstraint::None => Err(SQLRiteError::NotImplemented(
430 "JOIN without an ON condition is not supported (use INNER JOIN ... ON ...)".to_string(),
431 )),
432 }
433}
434
435fn parse_projection(items: &[SelectItem]) -> Result<Projection> {
436 if items.len() == 1
438 && let SelectItem::Wildcard(_) = &items[0]
439 {
440 return Ok(Projection::All);
441 }
442 let mut out = Vec::with_capacity(items.len());
443 for item in items {
444 out.push(parse_select_item(item)?);
445 }
446 Ok(Projection::Items(out))
447}
448
449fn parse_select_item(item: &SelectItem) -> Result<ProjectionItem> {
450 match item {
451 SelectItem::UnnamedExpr(expr) => parse_projection_expr(expr, None),
452 SelectItem::ExprWithAlias { expr, alias } => {
453 parse_projection_expr(expr, Some(alias.value.clone()))
454 }
455 SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => {
456 Err(SQLRiteError::NotImplemented(
457 "Wildcard mixed with other columns is not supported".to_string(),
458 ))
459 }
460 }
461}
462
463fn parse_projection_expr(expr: &Expr, alias: Option<String>) -> Result<ProjectionItem> {
464 match expr {
465 Expr::Identifier(ident) => Ok(ProjectionItem {
466 kind: ProjectionKind::Column {
467 qualifier: None,
468 name: ident.value.clone(),
469 },
470 alias,
471 }),
472 Expr::CompoundIdentifier(parts) => match parts.as_slice() {
473 [only] => Ok(ProjectionItem {
474 kind: ProjectionKind::Column {
475 qualifier: None,
476 name: only.value.clone(),
477 },
478 alias,
479 }),
480 [q, c] => Ok(ProjectionItem {
481 kind: ProjectionKind::Column {
482 qualifier: Some(q.value.clone()),
483 name: c.value.clone(),
484 },
485 alias,
486 }),
487 _ => Err(SQLRiteError::NotImplemented(format!(
488 "compound identifier with {} parts is not supported in projection",
489 parts.len()
490 ))),
491 },
492 Expr::Function(func) => {
493 let call = parse_aggregate_call(func)?;
494 Ok(ProjectionItem {
495 kind: ProjectionKind::Aggregate(call),
496 alias,
497 })
498 }
499 other => Err(SQLRiteError::NotImplemented(format!(
500 "Only bare column references and aggregate functions are supported in the projection list (got {other:?})"
501 ))),
502 }
503}
504
505fn parse_aggregate_call(func: &sqlparser::ast::Function) -> Result<AggregateCall> {
506 let name = match func.name.0.as_slice() {
509 [sqlparser::ast::ObjectNamePart::Identifier(ident)] => ident.value.clone(),
510 _ => {
511 return Err(SQLRiteError::NotImplemented(format!(
512 "qualified function names not supported: {:?}",
513 func.name
514 )));
515 }
516 };
517 let agg_fn = AggregateFn::from_name(&name).ok_or_else(|| {
518 SQLRiteError::NotImplemented(format!(
519 "function '{name}' is not supported in the projection list (only aggregate functions are: COUNT, SUM, AVG, MIN, MAX)"
520 ))
521 })?;
522
523 let arg_list = match &func.args {
526 FunctionArguments::List(l) => l,
527 _ => {
528 return Err(SQLRiteError::NotImplemented(format!(
529 "{name}(...) — unsupported argument shape"
530 )));
531 }
532 };
533
534 let distinct = matches!(
535 arg_list.duplicate_treatment,
536 Some(DuplicateTreatment::Distinct)
537 );
538
539 if !arg_list.clauses.is_empty() {
540 return Err(SQLRiteError::NotImplemented(format!(
541 "{name}(...) — extra argument clauses (ORDER BY / LIMIT inside the call) are not supported"
542 )));
543 }
544 if func.over.is_some() {
545 return Err(SQLRiteError::NotImplemented(
546 "window functions (OVER (...)) are not supported".to_string(),
547 ));
548 }
549 if func.filter.is_some() {
550 return Err(SQLRiteError::NotImplemented(
551 "FILTER (WHERE ...) on aggregates is not supported".to_string(),
552 ));
553 }
554 if !func.within_group.is_empty() {
555 return Err(SQLRiteError::NotImplemented(
556 "WITHIN GROUP on aggregates is not supported".to_string(),
557 ));
558 }
559
560 if arg_list.args.len() != 1 {
561 return Err(SQLRiteError::NotImplemented(format!(
562 "{name}(...) expects exactly one argument, got {}",
563 arg_list.args.len()
564 )));
565 }
566
567 let arg = match &arg_list.args[0] {
568 FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => AggregateArg::Star,
569 FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(ident))) => {
570 AggregateArg::Column(ident.value.clone())
571 }
572 FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::CompoundIdentifier(parts))) => {
573 let c = parts
574 .last()
575 .map(|p| p.value.clone())
576 .ok_or_else(|| SQLRiteError::Internal("empty compound identifier".to_string()))?;
577 AggregateArg::Column(c)
578 }
579 other => {
580 return Err(SQLRiteError::NotImplemented(format!(
581 "{name}(...) — argument must be `*` or a bare column reference (got {other:?})"
582 )));
583 }
584 };
585
586 if distinct && agg_fn != AggregateFn::Count {
590 return Err(SQLRiteError::NotImplemented(format!(
591 "DISTINCT is only supported on COUNT(...) for now, not {}",
592 agg_fn.as_str()
593 )));
594 }
595 if matches!(arg, AggregateArg::Star) && agg_fn != AggregateFn::Count {
596 return Err(SQLRiteError::NotImplemented(format!(
597 "{}(*) is not supported; use {}(<column>)",
598 agg_fn.as_str(),
599 agg_fn.as_str()
600 )));
601 }
602
603 Ok(AggregateCall {
604 func: agg_fn,
605 arg,
606 distinct,
607 })
608}
609
610fn parse_order_by(order_by: Option<&sqlparser::ast::OrderBy>) -> Result<Option<OrderByClause>> {
611 let Some(ob) = order_by else {
612 return Ok(None);
613 };
614 let exprs = match &ob.kind {
615 OrderByKind::Expressions(v) => v,
616 OrderByKind::All(_) => {
617 return Err(SQLRiteError::NotImplemented(
618 "ORDER BY ALL is not supported".to_string(),
619 ));
620 }
621 };
622 if exprs.len() != 1 {
623 return Err(SQLRiteError::NotImplemented(
624 "ORDER BY must have exactly one column for now".to_string(),
625 ));
626 }
627 let obe = &exprs[0];
628 let expr = obe.expr.clone();
634 let ascending = obe.options.asc.unwrap_or(true);
636 Ok(Some(OrderByClause { expr, ascending }))
637}
638
639fn parse_limit(limit: Option<&LimitClause>) -> Result<Option<usize>> {
640 let Some(lc) = limit else {
641 return Ok(None);
642 };
643 let limit_expr = match lc {
644 LimitClause::LimitOffset { limit, offset, .. } => {
645 if offset.is_some() {
646 return Err(SQLRiteError::NotImplemented(
647 "OFFSET is not supported yet".to_string(),
648 ));
649 }
650 limit.as_ref()
651 }
652 LimitClause::OffsetCommaLimit { .. } => {
653 return Err(SQLRiteError::NotImplemented(
654 "`LIMIT <offset>, <limit>` syntax is not supported yet".to_string(),
655 ));
656 }
657 };
658 let Some(expr) = limit_expr else {
659 return Ok(None);
660 };
661 let n = eval_const_usize(expr)?;
662 Ok(Some(n))
663}
664
665fn eval_const_usize(expr: &Expr) -> Result<usize> {
666 match expr {
667 Expr::Value(v) => match &v.value {
668 sqlparser::ast::Value::Number(n, _) => n.parse::<usize>().map_err(|e| {
669 SQLRiteError::Internal(format!("LIMIT must be a non-negative integer: {e}"))
670 }),
671 _ => Err(SQLRiteError::Internal(
672 "LIMIT must be an integer literal".to_string(),
673 )),
674 },
675 _ => Err(SQLRiteError::NotImplemented(
676 "LIMIT expression must be a literal number".to_string(),
677 )),
678 }
679}