1use std::sync::Arc;
10use uqa_core::{
11 memory::{
12 Budgeted, BudgetedVec, MemoryBudget, MemoryError, MemoryReservation, Produced,
13 ProductionControl,
14 },
15 CancellationToken, QueryCancelled, Value, ValueRetentionError,
16};
17
18use crate::ast::{ColumnDef, ColumnType, Expr, FunctionBinding};
19
20mod columns;
21mod expressions;
22
23#[derive(Debug, thiserror::Error)]
24pub enum CatalogRetentionError {
25 #[error(transparent)]
26 Memory(#[from] MemoryError),
27 #[error(transparent)]
28 Cancelled(#[from] QueryCancelled),
29 #[error("validated column expression contains a subquery")]
30 UnexpectedSubquery,
31}
32
33impl From<ValueRetentionError> for CatalogRetentionError {
34 fn from(error: ValueRetentionError) -> Self {
35 match error {
36 ValueRetentionError::Memory(error) => Self::Memory(error),
37 ValueRetentionError::Cancelled(error) => Self::Cancelled(error),
38 }
39 }
40}
41
42impl From<CatalogRetentionError> for crate::SQLError {
43 fn from(error: CatalogRetentionError) -> Self {
44 match error {
45 CatalogRetentionError::Memory(error) => Self::Routine {
46 sqlstate: "53200".into(),
47 message: error.to_string(),
48 },
49 CatalogRetentionError::Cancelled(error) => Self::Cancelled(error),
50 CatalogRetentionError::UnexpectedSubquery => Self::Internal(error.to_string()),
51 }
52 }
53}
54
55type Result<T> = std::result::Result<T, CatalogRetentionError>;
56
57#[derive(Debug, Clone)]
59pub struct RetainedColumns(Arc<Budgeted<Arc<Vec<ColumnDef>>>>);
60
61impl RetainedColumns {
62 pub fn capture(
64 columns: &Arc<Vec<ColumnDef>>,
65 budget: &MemoryBudget,
66 cancellation: &CancellationToken,
67 ) -> Result<Self> {
68 let mut walker = Walker::new(budget, cancellation);
69 walker.charge(size_of::<Vec<ColumnDef>>())?;
70 walker.children(columns, Node::Column)?;
71 let memory = walker.finish()?;
72 cancellation.check()?;
73 Ok(Self(
74 Budgeted::new(Arc::clone(columns), memory).into_shared()?,
75 ))
76 }
77
78 pub fn as_slice(&self) -> &[ColumnDef] {
79 &self.0
80 }
81
82 pub fn reserved_bytes(&self) -> usize {
83 self.0.reserved_bytes()
84 }
85}
86
87impl std::ops::Deref for RetainedColumns {
88 type Target = [ColumnDef];
89
90 fn deref(&self) -> &Self::Target {
91 self.as_slice()
92 }
93}
94
95impl ColumnDef {
96 pub fn reserve_retained_payload(
98 &self,
99 budget: &MemoryBudget,
100 cancellation: &CancellationToken,
101 ) -> Result<MemoryReservation> {
102 Walker::new(budget, cancellation).root(Node::Column(self))
103 }
104}
105
106impl ColumnType {
107 pub(crate) fn retain_external_with_control(
109 self,
110 control: &ProductionControl<'_>,
111 ) -> Result<Produced<Self>> {
112 struct Handoff {
113 value: ColumnType,
114 memory: Option<MemoryReservation>,
115 }
116 let mut output = Handoff {
117 value: self,
118 memory: control.empty_reservation(),
119 };
120 control.check()?;
121 if let Some(budget) = control.budget() {
122 let mut walker = Walker::with_control(budget, *control);
123 let result = walker
124 .visit(Node::Type(&output.value))
125 .and_then(|()| walker.drain());
126 let Walker {
127 memory, pending, ..
128 } = walker;
129 drop(pending);
130 output
131 .memory
132 .as_mut()
133 .expect("controlled catalog handoff")
134 .absorb(memory);
135 result?;
136 }
137 Ok(control.finish(output.value, output.memory)?)
138 }
139
140 pub fn reserve_retained_payload(
142 &self,
143 budget: &MemoryBudget,
144 cancellation: &CancellationToken,
145 ) -> Result<MemoryReservation> {
146 Walker::new(budget, cancellation).root(Node::Type(self))
147 }
148}
149
150impl Expr {
151 pub fn reserve_column_payload(
153 &self,
154 budget: &MemoryBudget,
155 cancellation: &CancellationToken,
156 ) -> Result<MemoryReservation> {
157 Walker::new(budget, cancellation).root(Node::Expr(self))
158 }
159}
160
161enum Node<'a> {
162 Column(&'a ColumnDef),
163 Type(&'a ColumnType),
164 Expr(&'a Expr),
165 Binding(&'a FunctionBinding),
166}
167
168struct Walker<'a> {
169 memory: MemoryReservation,
170 pending: BudgetedVec<Node<'a>>,
171 control: ProductionControl<'a>,
172}
173
174impl<'a> Walker<'a> {
175 fn new(budget: &'a MemoryBudget, cancellation: &'a CancellationToken) -> Self {
176 Self::with_control(
177 budget,
178 ProductionControl::new(budget, cancellation, cancellation),
179 )
180 }
181
182 fn with_control(budget: &'a MemoryBudget, control: ProductionControl<'a>) -> Self {
183 Self {
184 memory: budget.empty_reservation(),
185 pending: BudgetedVec::new(budget),
186 control,
187 }
188 }
189
190 fn root(mut self, root: Node<'a>) -> Result<MemoryReservation> {
191 self.visit(root)?;
192 self.finish()
193 }
194
195 fn finish(mut self) -> Result<MemoryReservation> {
196 self.drain()?;
197 Ok(self.memory)
198 }
199
200 fn drain(&mut self) -> Result<()> {
201 self.control.check_cancellation()?;
202 while let Some(node) = self.pending.pop() {
203 self.visit(node)?;
204 }
205 Ok(())
206 }
207
208 fn visit(&mut self, node: Node<'a>) -> Result<()> {
209 self.control.check_cancellation()?;
210 match node {
211 Node::Column(column) => self.column(column),
212 Node::Type(ty) => self.ty(ty),
213 Node::Expr(expr) => self.expr(expr),
214 Node::Binding(binding) => self.binding(binding),
215 }
216 }
217
218 fn node(&mut self, node: Node<'a>) -> Result<()> {
219 self.control.check_cancellation()?;
220 self.pending.push(node)?;
221 Ok(())
222 }
223
224 fn charge(&mut self, bytes: usize) -> Result<()> {
225 self.control.check_cancellation()?;
226 self.memory.grow(bytes)?;
227 Ok(())
228 }
229
230 fn buffer<T>(&mut self, capacity: usize) -> Result<()> {
231 self.charge(
232 capacity
233 .checked_mul(size_of::<T>())
234 .ok_or(MemoryError::SizeOverflow)?,
235 )
236 }
237
238 fn children<T>(&mut self, items: &'a Vec<T>, node: fn(&'a T) -> Node<'a>) -> Result<()> {
239 self.buffer::<T>(items.capacity())?;
240 for item in items {
241 self.node(node(item))?;
242 }
243 Ok(())
244 }
245
246 fn boxed<T>(&mut self, item: &'a T, node: fn(&'a T) -> Node<'a>) -> Result<()> {
247 self.charge(size_of::<T>())?;
248 self.node(node(item))
249 }
250
251 fn text(&mut self, text: &String) -> Result<()> {
252 self.charge(text.capacity())
253 }
254
255 fn optional_text(&mut self, text: Option<&String>) -> Result<()> {
256 if let Some(text) = text {
257 self.text(text)?;
258 }
259 Ok(())
260 }
261
262 fn texts(&mut self, texts: &Vec<String>) -> Result<()> {
263 self.buffer::<String>(texts.capacity())?;
264 for text in texts {
265 self.text(text)?;
266 }
267 Ok(())
268 }
269
270 fn value(&mut self, value: &Value) -> Result<()> {
271 let memory = value.reserve_retained_payload_with_check(self.memory.budget(), || {
272 self.control.check_cancellation()
273 })?;
274 self.memory.absorb(memory);
275 Ok(())
276 }
277
278 fn optional_expr(&mut self, expr: Option<&'a Expr>) -> Result<()> {
279 if let Some(expr) = expr {
280 self.node(Node::Expr(expr))?;
281 }
282 Ok(())
283 }
284
285 fn optional_boxed_expr(&mut self, expr: Option<&'a Expr>) -> Result<()> {
286 if let Some(expr) = expr {
287 self.boxed(expr, Node::Expr)?;
288 }
289 Ok(())
290 }
291}
292
293#[cfg(test)]
294mod tests;