alopex_embedded/
owned_sql.rs

1//! Owned, bounded SQL stream planning for embedded-local consumers.
2//!
3//! This module deliberately does not reuse `execute_sql_streaming`: that legacy helper creates
4//! a borrowed iterator and commits its source transaction before returning it. The public Python
5//! stream surface instead obtains one [`OwnedSqlStreamPlan`] here, opens its cursor through an
6//! owned core session, and advances the cursor until a terminal lease transition.
7
8use alopex_core::kv::{OwnedKVScan, OwnedKVTransaction};
9use alopex_core::{Key, Result as CoreResult, Value};
10use alopex_sql::catalog::{
11    Catalog, CatalogOverlay, ColumnMetadata, StorageType, TableMetadata, TxnCatalogView,
12};
13use alopex_sql::executor::evaluator::{evaluate, EvalContext};
14use alopex_sql::executor::ColumnInfo;
15use alopex_sql::planner::typed_expr::{Projection, TypedExpr};
16use alopex_sql::{AlopexDialect, LogicalPlan, Parser, Planner, RowCodec, SqlError, SqlValue};
17
18use crate::{Database, Error, Result};
19
20const UNSUPPORTED_STREAMING_SQL: &str = "unsupported_streaming_sql";
21
22/// A preflight-validated local SQL stream plan.
23///
24/// Only a single table scan (or literal-only `SELECT`) followed by a row-local filter and a
25/// slice is representable. The plan owns every catalog-derived value it needs after opening, so
26/// no catalog lock or borrowed SQL transaction is retained by a live stream.
27#[derive(Clone, Debug)]
28pub struct OwnedSqlStreamPlan {
29    source: OwnedSqlSource,
30    filter: Option<TypedExpr>,
31    projection: Projection,
32    columns: Vec<ColumnInfo>,
33    offset_remaining: u64,
34    limit_remaining: Option<u64>,
35}
36
37#[derive(Clone, Debug)]
38enum OwnedSqlSource {
39    Table(Box<TableMetadata>),
40    Literal,
41}
42
43/// One row-processing outcome after consuming a physical cursor entry.
44#[derive(Debug, Clone, PartialEq)]
45pub enum OwnedSqlRowOutcome {
46    /// The physical entry was filtered or skipped by OFFSET.
47    Skip,
48    /// One projected SQL row is ready for the caller.
49    Row(Vec<SqlValue>),
50    /// A `LIMIT` was reached and the cursor must not be advanced again.
51    Exhausted,
52}
53
54impl OwnedSqlStreamPlan {
55    /// Parse and validate the documented v0.8 local SELECT subset without opening a session.
56    pub fn preflight(database: &Database, sql: &str) -> Result<Self> {
57        let dialect = AlopexDialect;
58        let statements = Parser::parse_sql(&dialect, sql).map_err(SqlError::from)?;
59        if statements.len() != 1 {
60            return Err(unsupported("exactly one SELECT statement is required"));
61        }
62
63        let plan = {
64            let catalog = database
65                .sql_catalog
66                .read()
67                .map_err(|_| Error::CatalogLockPoisoned)?;
68            Planner::new(&*catalog)
69                .plan(&statements[0])
70                .map_err(SqlError::from)?
71        };
72        let catalog = database
73            .sql_catalog
74            .read()
75            .map_err(|_| Error::CatalogLockPoisoned)?;
76        Self::from_plan(&*catalog, plan)
77    }
78
79    /// Parse and validate a local SELECT against an uncommitted catalog overlay.
80    ///
81    /// The returned plan copies the table metadata it needs, so neither the catalog guard nor the
82    /// overlay borrow survives into the cursor.  This preserves transaction-local DDL visibility
83    /// without coupling a live Python stream to a borrowed SQL transaction.
84    pub fn preflight_in_transaction(
85        database: &Database,
86        overlay: &CatalogOverlay,
87        sql: &str,
88    ) -> Result<Self> {
89        let dialect = AlopexDialect;
90        let statements = Parser::parse_sql(&dialect, sql).map_err(SqlError::from)?;
91        if statements.len() != 1 {
92            return Err(unsupported("exactly one SELECT statement is required"));
93        }
94
95        let catalog = database
96            .sql_catalog
97            .read()
98            .map_err(|_| Error::CatalogLockPoisoned)?;
99        let view = TxnCatalogView::new(&*catalog, overlay);
100        let plan = Planner::new(&view)
101            .plan(&statements[0])
102            .map_err(SqlError::from)?;
103        Self::from_plan(&view, plan)
104    }
105
106    /// Return the output schema before a cursor is opened.
107    pub fn columns(&self) -> &[ColumnInfo] {
108        &self.columns
109    }
110
111    /// Open the physical cursor through the owned core transaction supplied by the session lease.
112    pub fn open_cursor(
113        &self,
114        transaction: &mut dyn OwnedKVTransaction,
115    ) -> CoreResult<Box<dyn OwnedKVScan>> {
116        match &self.source {
117            OwnedSqlSource::Table(table) => {
118                transaction.scan_prefix(&alopex_sql::KeyEncoder::table_prefix(table.table_id))
119            }
120            OwnedSqlSource::Literal => Ok(Box::new(OneRowCursor { emitted: false })),
121        }
122    }
123
124    /// Whether the logical slice is already exhausted without another storage read.
125    pub fn is_exhausted(&self) -> bool {
126        self.limit_remaining == Some(0)
127    }
128
129    /// Decode, filter, slice, and project a single owned cursor entry.
130    pub fn process_entry(&mut self, key: Key, value: Value) -> Result<OwnedSqlRowOutcome> {
131        if self.is_exhausted() {
132            return Ok(OwnedSqlRowOutcome::Exhausted);
133        }
134
135        let row = match &self.source {
136            OwnedSqlSource::Table(table) => {
137                let (table_id, _) = alopex_sql::KeyEncoder::decode_row_key(&key)
138                    .map_err(|error| Error::Sql(SqlError::from(error)))?;
139                if table_id != table.table_id {
140                    return Err(Error::Sql(SqlError::Execution {
141                        message: "owned table cursor yielded a row from another table".to_string(),
142                        code: "ALOPEX-E020",
143                    }));
144                }
145                RowCodec::decode(&value).map_err(|error| Error::Sql(SqlError::from(error)))?
146            }
147            OwnedSqlSource::Literal => Vec::new(),
148        };
149
150        if let Some(predicate) = &self.filter {
151            let context = EvalContext::new(&row);
152            if !matches!(
153                evaluate(predicate, &context).map_err(sql_execution_error)?,
154                SqlValue::Boolean(true)
155            ) {
156                return Ok(OwnedSqlRowOutcome::Skip);
157            }
158        }
159
160        if self.offset_remaining > 0 {
161            self.offset_remaining -= 1;
162            return Ok(OwnedSqlRowOutcome::Skip);
163        }
164
165        let projected = project_row(&self.projection, &row)?;
166        if let Some(remaining) = &mut self.limit_remaining {
167            *remaining = remaining.saturating_sub(1);
168        }
169        Ok(OwnedSqlRowOutcome::Row(projected))
170    }
171
172    fn from_plan(catalog: &impl Catalog, plan: LogicalPlan) -> Result<Self> {
173        let mut filter = None;
174        let mut limit = None;
175        let mut offset = 0;
176        let scan = unwrap_stream_nodes(plan, &mut filter, &mut limit, &mut offset)?;
177        let LogicalPlan::Scan { table, projection } = scan else {
178            return Err(unsupported(
179                "only a table scan or literal SELECT is streamable",
180            ));
181        };
182
183        let source = if table == alopex_sql::ast::dml::LITERAL_TABLE {
184            OwnedSqlSource::Literal
185        } else {
186            let table_meta = catalog
187                .get_table(&table)
188                .cloned()
189                .ok_or_else(|| Error::TableNotFound(table.clone()))?;
190            if table_meta.storage_options.storage_type != StorageType::Row {
191                return Err(unsupported(
192                    "columnar tables require the LocalScan.columnar_segment streaming path",
193                ));
194            }
195            OwnedSqlSource::Table(Box::new(table_meta))
196        };
197        let schema = match &source {
198            OwnedSqlSource::Table(table) => table.columns.clone(),
199            OwnedSqlSource::Literal => Vec::new(),
200        };
201        let columns = columns_for(&projection, &schema)?;
202
203        Ok(Self {
204            source,
205            filter,
206            projection,
207            columns,
208            offset_remaining: offset,
209            limit_remaining: limit,
210        })
211    }
212}
213
214/// An owned one-entry cursor used only for literal-only SELECT queries. The plan recognizes the
215/// empty key/value pair as its synthetic row; no caller can observe this physical representation.
216struct OneRowCursor {
217    emitted: bool,
218}
219
220impl OwnedKVScan for OneRowCursor {
221    fn next_entry(&mut self) -> CoreResult<Option<(Key, Value)>> {
222        if self.emitted {
223            Ok(None)
224        } else {
225            self.emitted = true;
226            Ok(Some((Vec::new(), Vec::new())))
227        }
228    }
229}
230
231fn unwrap_stream_nodes(
232    plan: LogicalPlan,
233    filter: &mut Option<TypedExpr>,
234    limit: &mut Option<u64>,
235    offset: &mut u64,
236) -> Result<LogicalPlan> {
237    match plan {
238        LogicalPlan::Limit {
239            input,
240            limit: next_limit,
241            offset: next_offset,
242            ties,
243        } => {
244            if ties.is_some() {
245                return Err(unsupported("FETCH ... WITH TIES is not streamable"));
246            }
247            if limit.is_some() || *offset != 0 {
248                return Err(unsupported("multiple slice nodes are not streamable"));
249            }
250            *limit = next_limit;
251            *offset = next_offset.unwrap_or(0);
252            unwrap_stream_nodes(*input, filter, limit, offset)
253        }
254        LogicalPlan::Filter { input, predicate } => {
255            if filter.replace(predicate).is_some() {
256                return Err(unsupported("multiple filter nodes are not streamable"));
257            }
258            unwrap_stream_nodes(*input, filter, limit, offset)
259        }
260        LogicalPlan::Scan { .. } => Ok(plan),
261        _ => Err(unsupported(
262            "streaming supports only SELECT with one table, row-local WHERE, projection, LIMIT, and OFFSET",
263        )),
264    }
265}
266
267fn columns_for(projection: &Projection, schema: &[ColumnMetadata]) -> Result<Vec<ColumnInfo>> {
268    match projection {
269        Projection::All(names) => names
270            .iter()
271            .map(|name| {
272                let column = schema
273                    .iter()
274                    .find(|column| column.name == *name)
275                    .ok_or_else(|| {
276                        Error::Sql(SqlError::Execution {
277                            message: format!("column not found: {name}"),
278                            code: "ALOPEX-E020",
279                        })
280                    })?;
281                Ok(ColumnInfo::new(&column.name, column.data_type.clone()))
282            })
283            .collect(),
284        Projection::Columns(columns) => Ok(columns
285            .iter()
286            .map(|column| {
287                let name = column
288                    .alias
289                    .clone()
290                    .unwrap_or_else(|| match &column.expr.kind {
291                        alopex_sql::TypedExprKind::ColumnRef { column, .. } => column.clone(),
292                        _ => "?column?".to_string(),
293                    });
294                ColumnInfo::new(name, column.expr.resolved_type.clone())
295            })
296            .collect()),
297    }
298}
299
300fn project_row(projection: &Projection, row: &[SqlValue]) -> Result<Vec<SqlValue>> {
301    match projection {
302        Projection::All(names) => {
303            if names.len() != row.len() {
304                return Err(Error::Sql(SqlError::Execution {
305                    message: "owned stream projection no longer matches the table schema"
306                        .to_string(),
307                    code: "ALOPEX-E020",
308                }));
309            }
310            Ok(row.to_vec())
311        }
312        Projection::Columns(columns) => {
313            let context = EvalContext::new(row);
314            columns
315                .iter()
316                .map(|column| evaluate(&column.expr, &context).map_err(sql_execution_error))
317                .collect()
318        }
319    }
320}
321
322fn sql_execution_error(error: alopex_sql::ExecutorError) -> Error {
323    Error::Sql(SqlError::from(error))
324}
325
326fn unsupported(message: impl Into<String>) -> Error {
327    Error::Sql(SqlError::Execution {
328        message: message.into(),
329        code: UNSUPPORTED_STREAMING_SQL,
330    })
331}
332
333#[cfg(test)]
334mod tests {
335    use alopex_core::txn::OwnedLeaseOutcome;
336    use alopex_core::TxnMode;
337
338    use super::{OwnedSqlRowOutcome, OwnedSqlStreamPlan};
339    use crate::Database;
340
341    #[test]
342    fn owned_sql_plan_streams_filter_projection_and_slice_without_borrowed_transaction() {
343        let database = Database::new();
344        database
345            .execute_sql("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, enabled BOOLEAN)")
346            .unwrap();
347        database
348            .execute_sql(
349                "INSERT INTO users (id, name, enabled) VALUES (1, 'one', true), (2, 'two', false), (3, 'three', true)",
350            )
351            .unwrap();
352
353        let mut plan = OwnedSqlStreamPlan::preflight(
354            &database,
355            "SELECT name, id + 10 AS next_id FROM users WHERE enabled = true LIMIT 1 OFFSET 1",
356        )
357        .unwrap();
358        assert_eq!(
359            plan.columns()
360                .iter()
361                .map(|column| column.name.as_str())
362                .collect::<Vec<_>>(),
363            vec!["name", "next_id"]
364        );
365
366        let session = database.begin_owned_read(Default::default()).unwrap();
367        let lease = session.acquire_lease().unwrap();
368        let mut cursor = lease
369            .with_transaction(|transaction| plan.open_cursor(transaction))
370            .unwrap();
371        let mut rows = Vec::new();
372        while let Some((key, value)) = cursor.next_entry().unwrap() {
373            if let OwnedSqlRowOutcome::Row(row) = plan.process_entry(key, value).unwrap() {
374                rows.push(row);
375            }
376        }
377        cursor.close().unwrap();
378        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
379        assert_eq!(
380            rows,
381            vec![vec![
382                alopex_sql::SqlValue::Text("three".to_string()),
383                alopex_sql::SqlValue::Integer(13),
384            ]]
385        );
386    }
387
388    #[test]
389    fn unsupported_streaming_sql_is_rejected_before_an_owned_session_is_opened() {
390        let database = Database::new();
391        database
392            .execute_sql("CREATE TABLE users (id INTEGER PRIMARY KEY)")
393            .unwrap();
394        let error = OwnedSqlStreamPlan::preflight(&database, "SELECT id FROM users ORDER BY id")
395            .unwrap_err();
396        assert_eq!(error.sql_error_code(), Some("unsupported_streaming_sql"));
397
398        let session = database
399            .begin_owned_transaction(TxnMode::ReadWrite)
400            .unwrap();
401        let lease = session.acquire_lease().unwrap();
402        lease
403            .with_transaction(|transaction| {
404                transaction.put(b"still-open".to_vec(), b"yes".to_vec())
405            })
406            .unwrap();
407        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
408        session.rollback().unwrap();
409    }
410
411    #[test]
412    fn literal_only_select_uses_the_owned_values_cursor() {
413        let database = Database::new();
414        let mut plan = OwnedSqlStreamPlan::preflight(&database, "SELECT 1 + 2 AS value").unwrap();
415        let session = database.begin_owned_read(Default::default()).unwrap();
416        let lease = session.acquire_lease().unwrap();
417        let mut cursor = lease
418            .with_transaction(|transaction| plan.open_cursor(transaction))
419            .unwrap();
420        let (key, value) = cursor.next_entry().unwrap().unwrap();
421        assert_eq!(
422            plan.process_entry(key, value).unwrap(),
423            OwnedSqlRowOutcome::Row(vec![alopex_sql::SqlValue::Integer(3)])
424        );
425        assert!(cursor.next_entry().unwrap().is_none());
426        cursor.close().unwrap();
427        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
428    }
429}
alopex_embedded/owned_sql.rs

alopex_embedded/
owned_sql.rs
