Skip to main content

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        } => {
243            if limit.is_some() || *offset != 0 {
244                return Err(unsupported("multiple slice nodes are not streamable"));
245            }
246            *limit = next_limit;
247            *offset = next_offset.unwrap_or(0);
248            unwrap_stream_nodes(*input, filter, limit, offset)
249        }
250        LogicalPlan::Filter { input, predicate } => {
251            if filter.replace(predicate).is_some() {
252                return Err(unsupported("multiple filter nodes are not streamable"));
253            }
254            unwrap_stream_nodes(*input, filter, limit, offset)
255        }
256        LogicalPlan::Scan { .. } => Ok(plan),
257        _ => Err(unsupported(
258            "streaming supports only SELECT with one table, row-local WHERE, projection, LIMIT, and OFFSET",
259        )),
260    }
261}
262
263fn columns_for(projection: &Projection, schema: &[ColumnMetadata]) -> Result<Vec<ColumnInfo>> {
264    match projection {
265        Projection::All(names) => names
266            .iter()
267            .map(|name| {
268                let column = schema
269                    .iter()
270                    .find(|column| column.name == *name)
271                    .ok_or_else(|| {
272                        Error::Sql(SqlError::Execution {
273                            message: format!("column not found: {name}"),
274                            code: "ALOPEX-E020",
275                        })
276                    })?;
277                Ok(ColumnInfo::new(&column.name, column.data_type.clone()))
278            })
279            .collect(),
280        Projection::Columns(columns) => Ok(columns
281            .iter()
282            .map(|column| {
283                let name = column
284                    .alias
285                    .clone()
286                    .unwrap_or_else(|| match &column.expr.kind {
287                        alopex_sql::TypedExprKind::ColumnRef { column, .. } => column.clone(),
288                        _ => "?column?".to_string(),
289                    });
290                ColumnInfo::new(name, column.expr.resolved_type.clone())
291            })
292            .collect()),
293    }
294}
295
296fn project_row(projection: &Projection, row: &[SqlValue]) -> Result<Vec<SqlValue>> {
297    match projection {
298        Projection::All(names) => {
299            if names.len() != row.len() {
300                return Err(Error::Sql(SqlError::Execution {
301                    message: "owned stream projection no longer matches the table schema"
302                        .to_string(),
303                    code: "ALOPEX-E020",
304                }));
305            }
306            Ok(row.to_vec())
307        }
308        Projection::Columns(columns) => {
309            let context = EvalContext::new(row);
310            columns
311                .iter()
312                .map(|column| evaluate(&column.expr, &context).map_err(sql_execution_error))
313                .collect()
314        }
315    }
316}
317
318fn sql_execution_error(error: alopex_sql::ExecutorError) -> Error {
319    Error::Sql(SqlError::from(error))
320}
321
322fn unsupported(message: impl Into<String>) -> Error {
323    Error::Sql(SqlError::Execution {
324        message: message.into(),
325        code: UNSUPPORTED_STREAMING_SQL,
326    })
327}
328
329#[cfg(test)]
330mod tests {
331    use alopex_core::txn::OwnedLeaseOutcome;
332    use alopex_core::TxnMode;
333
334    use super::{OwnedSqlRowOutcome, OwnedSqlStreamPlan};
335    use crate::Database;
336
337    #[test]
338    fn owned_sql_plan_streams_filter_projection_and_slice_without_borrowed_transaction() {
339        let database = Database::new();
340        database
341            .execute_sql("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, enabled BOOLEAN)")
342            .unwrap();
343        database
344            .execute_sql(
345                "INSERT INTO users (id, name, enabled) VALUES (1, 'one', true), (2, 'two', false), (3, 'three', true)",
346            )
347            .unwrap();
348
349        let mut plan = OwnedSqlStreamPlan::preflight(
350            &database,
351            "SELECT name, id + 10 AS next_id FROM users WHERE enabled = true LIMIT 1 OFFSET 1",
352        )
353        .unwrap();
354        assert_eq!(
355            plan.columns()
356                .iter()
357                .map(|column| column.name.as_str())
358                .collect::<Vec<_>>(),
359            vec!["name", "next_id"]
360        );
361
362        let session = database.begin_owned_read(Default::default()).unwrap();
363        let lease = session.acquire_lease().unwrap();
364        let mut cursor = lease
365            .with_transaction(|transaction| plan.open_cursor(transaction))
366            .unwrap();
367        let mut rows = Vec::new();
368        while let Some((key, value)) = cursor.next_entry().unwrap() {
369            if let OwnedSqlRowOutcome::Row(row) = plan.process_entry(key, value).unwrap() {
370                rows.push(row);
371            }
372        }
373        cursor.close().unwrap();
374        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
375        assert_eq!(
376            rows,
377            vec![vec![
378                alopex_sql::SqlValue::Text("three".to_string()),
379                alopex_sql::SqlValue::Integer(13),
380            ]]
381        );
382    }
383
384    #[test]
385    fn unsupported_streaming_sql_is_rejected_before_an_owned_session_is_opened() {
386        let database = Database::new();
387        database
388            .execute_sql("CREATE TABLE users (id INTEGER PRIMARY KEY)")
389            .unwrap();
390        let error = OwnedSqlStreamPlan::preflight(&database, "SELECT id FROM users ORDER BY id")
391            .unwrap_err();
392        assert_eq!(error.sql_error_code(), Some("unsupported_streaming_sql"));
393
394        let session = database
395            .begin_owned_transaction(TxnMode::ReadWrite)
396            .unwrap();
397        let lease = session.acquire_lease().unwrap();
398        lease
399            .with_transaction(|transaction| {
400                transaction.put(b"still-open".to_vec(), b"yes".to_vec())
401            })
402            .unwrap();
403        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
404        session.rollback().unwrap();
405    }
406
407    #[test]
408    fn literal_only_select_uses_the_owned_values_cursor() {
409        let database = Database::new();
410        let mut plan = OwnedSqlStreamPlan::preflight(&database, "SELECT 1 + 2 AS value").unwrap();
411        let session = database.begin_owned_read(Default::default()).unwrap();
412        let lease = session.acquire_lease().unwrap();
413        let mut cursor = lease
414            .with_transaction(|transaction| plan.open_cursor(transaction))
415            .unwrap();
416        let (key, value) = cursor.next_entry().unwrap().unwrap();
417        assert_eq!(
418            plan.process_entry(key, value).unwrap(),
419            OwnedSqlRowOutcome::Row(vec![alopex_sql::SqlValue::Integer(3)])
420        );
421        assert!(cursor.next_entry().unwrap().is_none());
422        cursor.close().unwrap();
423        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
424    }
425}