Skip to main content

gluesql_core/plan/statement/query/
limit.rs

1use {
2    super::{
3        DistinctPlan, OffsetPlan, ProjectPlan, SelectOrderByPlan, ValuesOrderByPlan, ValuesPlan,
4    },
5    crate::plan::ExprPlan,
6    serde::{Deserialize, Serialize},
7};
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub struct LimitPlan {
11    pub input: LimitInputPlan,
12    pub count: ExprPlan,
13}
14
15impl LimitPlan {
16    pub(super) fn project(&self) -> Option<&ProjectPlan> {
17        match &self.input {
18            LimitInputPlan::Project(project) => Some(project),
19            LimitInputPlan::Values(_) | LimitInputPlan::ValuesOrderBy(_) => None,
20            LimitInputPlan::SelectOrderBy(order_by) => Some(&order_by.input),
21            LimitInputPlan::Distinct(distinct) => Some(distinct.project()),
22            LimitInputPlan::Offset(offset) => offset.project(),
23        }
24    }
25}
26
27#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
28pub enum LimitInputPlan {
29    Project(ProjectPlan),
30    Values(ValuesPlan),
31    SelectOrderBy(SelectOrderByPlan),
32    ValuesOrderBy(ValuesOrderByPlan),
33    Distinct(DistinctPlan),
34    Offset(OffsetPlan),
35}
36
37#[cfg(test)]
38mod tests {
39    use {
40        super::{LimitInputPlan, LimitPlan},
41        crate::{
42            ast::Literal,
43            plan::{ExprPlan, OffsetInputPlan, OffsetPlan, ValuesPlan},
44        },
45    };
46
47    #[test]
48    fn limit_accepts_values_input() {
49        let plan = LimitPlan {
50            input: LimitInputPlan::Values(ValuesPlan(Vec::new())),
51            count: count(3),
52        };
53
54        assert!(matches!(plan.input, LimitInputPlan::Values(_)));
55    }
56
57    fn count(value: i64) -> ExprPlan {
58        ExprPlan::Literal(Literal::Number(value.into()))
59    }
60
61    #[test]
62    fn limit_accepts_offset_input() {
63        let plan = LimitPlan {
64            input: LimitInputPlan::Offset(OffsetPlan {
65                input: OffsetInputPlan::Values(ValuesPlan(Vec::new())),
66                count: count(2),
67            }),
68            count: count(3),
69        };
70
71        assert!(matches!(
72            plan,
73            LimitPlan {
74                input: LimitInputPlan::Offset(_),
75                count: actual,
76            } if actual == count(3)
77        ));
78    }
79}