Skip to main content

sqlrite/sql/parser/
select.rs

1use sqlparser::ast::{
2    Expr, LimitClause, OrderByKind, Query, Select, SelectItem, SetExpr, Statement, TableFactor,
3    TableWithJoins,
4};
5
6use crate::error::{Result, SQLRiteError};
7
8/// What columns to project from a SELECT.
9#[derive(Debug, Clone, PartialEq)]
10pub enum Projection {
11    /// `SELECT *` — every column in the table, in declaration order.
12    All,
13    /// `SELECT a, b, c` — explicit list.
14    Columns(Vec<String>),
15}
16
17/// A parsed `ORDER BY` clause: a single sort key (expression), ascending
18/// by default. Phase 7b widened this from "bare column name" to
19/// "arbitrary expression" so KNN queries of the form
20/// `ORDER BY vec_distance_l2(col, [...]) LIMIT k` work end-to-end. The
21/// expression is evaluated per-row at execution time via `eval_expr`;
22/// the simple `ORDER BY col` form still works because that's just an
23/// `Expr::Identifier` taking the same path.
24#[derive(Debug, Clone)]
25pub struct OrderByClause {
26    pub expr: Expr,
27    pub ascending: bool,
28}
29
30/// A parsed, simplified SELECT query.
31#[derive(Debug, Clone)]
32pub struct SelectQuery {
33    pub table_name: String,
34    pub projection: Projection,
35    /// Raw sqlparser WHERE expression, evaluated by the executor at run time.
36    pub selection: Option<Expr>,
37    pub order_by: Option<OrderByClause>,
38    pub limit: Option<usize>,
39}
40
41impl SelectQuery {
42    pub fn new(statement: &Statement) -> Result<Self> {
43        let Statement::Query(query) = statement else {
44            return Err(SQLRiteError::Internal(
45                "Error parsing SELECT: expected a Query statement".to_string(),
46            ));
47        };
48
49        let Query {
50            body,
51            order_by,
52            limit_clause,
53            ..
54        } = query.as_ref();
55
56        let SetExpr::Select(select) = body.as_ref() else {
57            return Err(SQLRiteError::NotImplemented(
58                "Only simple SELECT queries are supported (no UNION / VALUES / CTEs yet)"
59                    .to_string(),
60            ));
61        };
62        let Select {
63            projection,
64            from,
65            selection,
66            distinct,
67            group_by,
68            having,
69            ..
70        } = select.as_ref();
71
72        if distinct.is_some() {
73            return Err(SQLRiteError::NotImplemented(
74                "SELECT DISTINCT is not supported yet".to_string(),
75            ));
76        }
77        if having.is_some() {
78            return Err(SQLRiteError::NotImplemented(
79                "HAVING is not supported yet".to_string(),
80            ));
81        }
82        // GroupByExpr::Expressions(v, _) with an empty v is the "no GROUP BY" shape.
83        if let sqlparser::ast::GroupByExpr::Expressions(exprs, _) = group_by {
84            if !exprs.is_empty() {
85                return Err(SQLRiteError::NotImplemented(
86                    "GROUP BY is not supported yet".to_string(),
87                ));
88            }
89        } else {
90            return Err(SQLRiteError::NotImplemented(
91                "GROUP BY ALL is not supported".to_string(),
92            ));
93        }
94
95        let table_name = extract_single_table_name(from)?;
96        let projection = parse_projection(projection)?;
97        let order_by = parse_order_by(order_by.as_ref())?;
98        let limit = parse_limit(limit_clause.as_ref())?;
99
100        Ok(SelectQuery {
101            table_name,
102            projection,
103            selection: selection.clone(),
104            order_by,
105            limit,
106        })
107    }
108}
109
110fn extract_single_table_name(from: &[TableWithJoins]) -> Result<String> {
111    if from.len() != 1 {
112        return Err(SQLRiteError::NotImplemented(
113            "SELECT from multiple tables (joins / comma-joins) is not supported yet".to_string(),
114        ));
115    }
116    let twj = &from[0];
117    if !twj.joins.is_empty() {
118        return Err(SQLRiteError::NotImplemented(
119            "JOIN is not supported yet".to_string(),
120        ));
121    }
122    match &twj.relation {
123        TableFactor::Table { name, .. } => Ok(name.to_string()),
124        _ => Err(SQLRiteError::NotImplemented(
125            "Only SELECT from a plain table is supported".to_string(),
126        )),
127    }
128}
129
130fn parse_projection(items: &[SelectItem]) -> Result<Projection> {
131    // Special-case `SELECT *`.
132    if items.len() == 1 {
133        if let SelectItem::Wildcard(_) = &items[0] {
134            return Ok(Projection::All);
135        }
136    }
137    let mut cols = Vec::with_capacity(items.len());
138    for item in items {
139        match item {
140            SelectItem::UnnamedExpr(Expr::Identifier(ident)) => cols.push(ident.value.clone()),
141            SelectItem::UnnamedExpr(Expr::CompoundIdentifier(parts)) => {
142                if let Some(last) = parts.last() {
143                    cols.push(last.value.clone());
144                } else {
145                    return Err(SQLRiteError::Internal(
146                        "empty qualified column reference".to_string(),
147                    ));
148                }
149            }
150            SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => {
151                return Err(SQLRiteError::NotImplemented(
152                    "Wildcard mixed with other columns is not supported".to_string(),
153                ));
154            }
155            SelectItem::ExprWithAlias { .. } | SelectItem::UnnamedExpr(_) => {
156                return Err(SQLRiteError::NotImplemented(
157                    "Only bare column references are supported in the projection list".to_string(),
158                ));
159            }
160        }
161    }
162    Ok(Projection::Columns(cols))
163}
164
165fn parse_order_by(order_by: Option<&sqlparser::ast::OrderBy>) -> Result<Option<OrderByClause>> {
166    let Some(ob) = order_by else {
167        return Ok(None);
168    };
169    let exprs = match &ob.kind {
170        OrderByKind::Expressions(v) => v,
171        OrderByKind::All(_) => {
172            return Err(SQLRiteError::NotImplemented(
173                "ORDER BY ALL is not supported".to_string(),
174            ));
175        }
176    };
177    if exprs.len() != 1 {
178        return Err(SQLRiteError::NotImplemented(
179            "ORDER BY must have exactly one column for now".to_string(),
180        ));
181    }
182    let obe = &exprs[0];
183    // Phase 7b: accept arbitrary expressions, not just bare column refs.
184    // The executor's `sort_rowids` evaluates this expression per row via
185    // `eval_expr`, which handles Identifier (column lookup), Function
186    // (vec_distance_*), arithmetic, etc. uniformly. The previous
187    // column-name-only restriction has been lifted.
188    let expr = obe.expr.clone();
189    // `asc == None` is the dialect default (ASC).
190    let ascending = obe.options.asc.unwrap_or(true);
191    Ok(Some(OrderByClause { expr, ascending }))
192}
193
194fn parse_limit(limit: Option<&LimitClause>) -> Result<Option<usize>> {
195    let Some(lc) = limit else {
196        return Ok(None);
197    };
198    let limit_expr = match lc {
199        LimitClause::LimitOffset { limit, offset, .. } => {
200            if offset.is_some() {
201                return Err(SQLRiteError::NotImplemented(
202                    "OFFSET is not supported yet".to_string(),
203                ));
204            }
205            limit.as_ref()
206        }
207        LimitClause::OffsetCommaLimit { .. } => {
208            return Err(SQLRiteError::NotImplemented(
209                "`LIMIT <offset>, <limit>` syntax is not supported yet".to_string(),
210            ));
211        }
212    };
213    let Some(expr) = limit_expr else {
214        return Ok(None);
215    };
216    let n = eval_const_usize(expr)?;
217    Ok(Some(n))
218}
219
220fn eval_const_usize(expr: &Expr) -> Result<usize> {
221    match expr {
222        Expr::Value(v) => match &v.value {
223            sqlparser::ast::Value::Number(n, _) => n.parse::<usize>().map_err(|e| {
224                SQLRiteError::Internal(format!("LIMIT must be a non-negative integer: {e}"))
225            }),
226            _ => Err(SQLRiteError::Internal(
227                "LIMIT must be an integer literal".to_string(),
228            )),
229        },
230        _ => Err(SQLRiteError::NotImplemented(
231            "LIMIT expression must be a literal number".to_string(),
232        )),
233    }
234}