proof_of_sql/sql/parse/
query_expr.rs

1use super::{EnrichedExpr, FilterExecBuilder, QueryContextBuilder};
2use crate::{
3    base::database::SchemaAccessor,
4    sql::{
5        parse::ConversionResult,
6        postprocessing::{
7            GroupByPostprocessing, OrderByPostprocessing, OwnedTablePostprocessing,
8            SelectPostprocessing, SlicePostprocessing,
9        },
10        proof_plans::{DynProofPlan, GroupByExec},
11    },
12};
13use alloc::{fmt, vec, vec::Vec};
14use proof_of_sql_parser::{intermediate_ast::SetExpression, SelectStatement};
15use serde::{Deserialize, Serialize};
16use sqlparser::ast::Ident;
17
18#[derive(PartialEq, Serialize, Deserialize)]
19/// A `QueryExpr` represents a Proof of SQL query that can be executed against a database.
20/// It consists of a `DynProofPlan` for provable components and a vector of `OwnedTablePostprocessing` for the rest.
21pub struct QueryExpr {
22    proof_expr: DynProofPlan,
23    postprocessing: Vec<OwnedTablePostprocessing>,
24}
25
26// Implements fmt::Debug to aid in debugging QueryExpr.
27// Prints filter and postprocessing fields in a readable format.
28impl fmt::Debug for QueryExpr {
29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30        write!(
31            f,
32            "QueryExpr \n[{:#?},\n{:#?}\n]",
33            self.proof_expr, self.postprocessing
34        )
35    }
36}
37
38impl QueryExpr {
39    /// Creates a new `QueryExpr` with the given `DynProofPlan` and `OwnedTablePostprocessing`.
40    #[must_use]
41    pub fn new(proof_expr: DynProofPlan, postprocessing: Vec<OwnedTablePostprocessing>) -> Self {
42        Self {
43            proof_expr,
44            postprocessing,
45        }
46    }
47
48    /// Parse an intermediate AST `SelectStatement` into a `QueryExpr`.
49    pub fn try_new(
50        ast: SelectStatement,
51        default_schema: Ident,
52        schema_accessor: &dyn SchemaAccessor,
53    ) -> ConversionResult<Self> {
54        let context = match *ast.expr {
55            SetExpression::Query {
56                result_exprs,
57                from,
58                where_expr,
59                group_by,
60            } => QueryContextBuilder::new(schema_accessor)
61                .visit_table_expr(&from, default_schema)
62                .visit_group_by_exprs(group_by.into_iter().map(Ident::from).collect())?
63                .visit_result_exprs(result_exprs)?
64                .visit_where_expr(where_expr)?
65                .visit_order_by_exprs(ast.order_by.into_iter().map(Into::into).collect())?
66                .visit_slice_expr(ast.slice)
67                .build()?,
68        };
69        let result_aliased_exprs = context.get_aliased_result_exprs()?.to_vec();
70        let group_by = context.get_group_by_exprs();
71        // Figure out the basic postprocessing steps.
72        let mut postprocessing = vec![];
73        let order_bys = context.get_order_by_exprs();
74        if !order_bys.is_empty() {
75            postprocessing.push(OwnedTablePostprocessing::new_order_by(
76                OrderByPostprocessing::new(order_bys.to_vec()),
77            ));
78        }
79        if let Some(slice) = context.get_slice_expr() {
80            postprocessing.push(OwnedTablePostprocessing::new_slice(
81                SlicePostprocessing::new(Some(slice.number_rows), Some(slice.offset_value)),
82            ));
83        }
84        if context.has_agg() {
85            if let Some(group_by_expr) = Option::<GroupByExec>::try_from(&context)? {
86                Ok(Self {
87                    proof_expr: DynProofPlan::GroupBy(group_by_expr),
88                    postprocessing,
89                })
90            } else {
91                let raw_enriched_exprs = result_aliased_exprs
92                    .iter()
93                    .map(|aliased_expr| EnrichedExpr {
94                        residue_expression: aliased_expr.clone(),
95                        dyn_proof_expr: None,
96                    })
97                    .collect::<Vec<_>>();
98                let filter = FilterExecBuilder::new(context.get_column_mapping())
99                    .add_table_expr(context.get_table_ref().clone())
100                    .add_where_expr(context.get_where_expr().clone())?
101                    .add_result_columns(&raw_enriched_exprs)
102                    .build();
103
104                let group_by_postprocessing =
105                    GroupByPostprocessing::try_new(group_by.to_vec(), result_aliased_exprs)?;
106                postprocessing.insert(
107                    0,
108                    OwnedTablePostprocessing::new_group_by(group_by_postprocessing.clone()),
109                );
110                let remainder_exprs = group_by_postprocessing.remainder_exprs();
111                // Check whether we need to do select postprocessing.
112                // That is, if any of them is not simply a column reference.
113                if remainder_exprs
114                    .iter()
115                    .any(|expr| expr.try_as_identifier().is_none())
116                {
117                    postprocessing.insert(
118                        1,
119                        OwnedTablePostprocessing::new_select(SelectPostprocessing::new(
120                            remainder_exprs.to_vec(),
121                        )),
122                    );
123                }
124                Ok(Self {
125                    proof_expr: DynProofPlan::Filter(filter),
126                    postprocessing,
127                })
128            }
129        } else {
130            // No group by, so we need to do a filter.
131            let column_mapping = context.get_column_mapping();
132            let enriched_exprs = result_aliased_exprs
133                .iter()
134                .map(|aliased_expr| EnrichedExpr::new(aliased_expr.clone(), &column_mapping))
135                .collect::<Vec<_>>();
136            let select_exprs = enriched_exprs
137                .iter()
138                .map(|enriched_expr| enriched_expr.residue_expression.clone())
139                .collect::<Vec<_>>();
140            let filter = FilterExecBuilder::new(context.get_column_mapping())
141                .add_table_expr(context.get_table_ref().clone())
142                .add_where_expr(context.get_where_expr().clone())?
143                .add_result_columns(&enriched_exprs)
144                .build();
145            // Check whether we need to do select postprocessing.
146            if select_exprs
147                .iter()
148                .any(|expr| expr.try_as_identifier().is_none())
149            {
150                postprocessing.insert(
151                    0,
152                    OwnedTablePostprocessing::new_select(SelectPostprocessing::new(select_exprs)),
153                );
154            }
155            Ok(Self {
156                proof_expr: DynProofPlan::Filter(filter),
157                postprocessing,
158            })
159        }
160    }
161
162    /// Immutable access to this query's provable filter expression.
163    #[must_use]
164    pub fn proof_expr(&self) -> &DynProofPlan {
165        &self.proof_expr
166    }
167
168    /// Immutable access to this query's post-proof result transform expression.
169    #[must_use]
170    pub fn postprocessing(&self) -> &[OwnedTablePostprocessing] {
171        &self.postprocessing
172    }
173}
proof_of_sql/sql/parse/query_expr.rs

proof_of_sql/sql/parse/
query_expr.rs
