1use 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#[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#[derive(Debug, Clone, PartialEq)]
45pub enum OwnedSqlRowOutcome {
46 Skip,
48 Row(Vec<SqlValue>),
50 Exhausted,
52}
53
54impl OwnedSqlStreamPlan {
55 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 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 pub fn columns(&self) -> &[ColumnInfo] {
108 &self.columns
109 }
110
111 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 pub fn is_exhausted(&self) -> bool {
126 self.limit_remaining == Some(0)
127 }
128
129 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
214struct 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}