Skip to main content

uqa_sql/schema/
retention.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Cancellable admission of immutable catalog column generations and their owned SQL payloads.
8
9use 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/// An admitted immutable column generation. Clones share both the original definitions and one reservation; they do not copy ASTs or reacquire an allowance. Strings, vector capacities, boxed payloads and nested Core values are charged. Allocator/reference-count bookkeeping and Core value map node slack are outside the payload allowance.
58#[derive(Debug, Clone)]
59pub struct RetainedColumns(Arc<Budgeted<Arc<Vec<ColumnDef>>>>);
60
61impl RetainedColumns {
62    /// Admit an existing, SQL-validated generation before acquiring shared ownership. The source remains unchanged on cancellation or quota failure. Subqueries violate the stored column-expression invariant enforced by the default, CHECK and generated-column validators and return a typed error.
63    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    /// Admit owned payloads below this inline column definition without copying them. The returned lease must remain with any retained owner.
97    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    /// Retain a type returned by an external catalog resolver. Its constructor belongs to that resolver; SQL retains the exposed payload before deriving names or traversing the type.
108    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    /// Admit nested type names and boxed base types, excluding this type's inline layout.
141    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    /// Admit a validated column expression's owned payload, including function bindings and Core values, excluding its inline layout. Traversal is iterative and uses the same allowance. Subqueries violate the stored column-expression invariant and return a typed error.
152    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;