1#![no_std]
6
7extern crate alloc;
8
9pub mod aggregate;
10pub mod describe;
11pub mod eval;
12pub mod json;
13pub mod memoize;
14pub mod plan_cache;
15pub mod publications;
16pub mod query_stats;
17pub mod reorder;
18pub mod selectivity;
19pub mod statistics;
20pub mod subscriptions;
21pub mod users;
22
23pub use crate::users::{Role, ScramSecrets, UserError, UserStore};
24
25use alloc::borrow::Cow;
26use alloc::boxed::Box;
27use alloc::collections::BTreeMap;
28use alloc::string::{String, ToString};
29use alloc::vec::Vec;
30use core::fmt;
31
32use spg_sql::ast::{
33 BinOp, ColumnDef, ColumnName, ColumnTypeName, CreateIndexStatement,
34 CreatePublicationStatement, CreateSubscriptionStatement, CreateTableStatement,
35 CreateUserStatement, Expr, FrameBound, FrameKind, FromClause, IndexMethod, InsertStatement,
36 JoinKind, Literal, OrderBy, SelectItem, SelectStatement, Statement, UnOp, UnionKind,
37 VecEncoding as SqlVecEncoding, WindowFrame,
38};
39use spg_sql::parser::{self, ParseError};
40use spg_storage::{
41 Catalog, ColumnSchema, CompactReport, DataType, IndexKey, IndexKind, Row, StorageError, Table,
42 TableSchema, Value, VecEncoding,
43};
44
45use crate::eval::{EvalContext, EvalError};
46
47#[derive(Debug, Clone, PartialEq)]
49#[non_exhaustive]
50pub enum QueryResult {
51 CommandOk {
60 affected: usize,
61 modified_catalog: bool,
62 },
63 Rows {
65 columns: Vec<ColumnSchema>,
66 rows: Vec<Row>,
67 },
68}
69
70#[derive(Debug, Clone, PartialEq)]
76#[non_exhaustive]
77pub enum EngineError {
78 Parse(ParseError),
79 Storage(StorageError),
80 Eval(EvalError),
81 Unsupported(String),
83 TransactionAlreadyOpen,
85 NoActiveTransaction,
87 WriteRequired,
92 RowLimitExceeded(usize),
95 Cancelled,
101}
102
103impl fmt::Display for EngineError {
104 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105 match self {
106 Self::Parse(e) => write!(f, "parse: {e}"),
107 Self::Storage(e) => write!(f, "storage: {e}"),
108 Self::Eval(e) => write!(f, "eval: {e}"),
109 Self::Unsupported(s) => write!(f, "unsupported: {s}"),
110 Self::TransactionAlreadyOpen => f.write_str("a transaction is already open"),
111 Self::NoActiveTransaction => f.write_str("no active transaction"),
112 Self::WriteRequired => {
113 f.write_str("statement requires a write lock (use execute, not execute_readonly)")
114 }
115 Self::RowLimitExceeded(n) => {
116 write!(f, "query exceeded max_query_rows={n}")
117 }
118 Self::Cancelled => f.write_str("query cancelled (timeout or client request)"),
119 }
120 }
121}
122
123impl From<ParseError> for EngineError {
124 fn from(e: ParseError) -> Self {
125 Self::Parse(e)
126 }
127}
128impl From<StorageError> for EngineError {
129 fn from(e: StorageError) -> Self {
130 Self::Storage(e)
131 }
132}
133impl From<EvalError> for EngineError {
134 fn from(e: EvalError) -> Self {
135 Self::Eval(e)
136 }
137}
138
139pub type ClockFn = fn() -> i64;
148
149pub type SaltFn = fn() -> [u8; 16];
156
157#[derive(Debug, Clone, Copy)]
168pub struct CancelToken<'a> {
169 flag: Option<&'a core::sync::atomic::AtomicBool>,
170}
171
172impl<'a> CancelToken<'a> {
173 #[must_use]
174 pub const fn none() -> Self {
175 Self { flag: None }
176 }
177
178 #[must_use]
179 pub const fn from_flag(f: &'a core::sync::atomic::AtomicBool) -> Self {
180 Self { flag: Some(f) }
181 }
182
183 #[must_use]
184 pub fn is_cancelled(self) -> bool {
185 self.flag
186 .is_some_and(|f| f.load(core::sync::atomic::Ordering::Relaxed))
187 }
188
189 #[inline]
193 pub fn check(self) -> Result<(), EngineError> {
194 if self.is_cancelled() {
195 Err(EngineError::Cancelled)
196 } else {
197 Ok(())
198 }
199 }
200}
201
202const ENVELOPE_MAGIC: &[u8; 8] = b"SPGENV01";
260const ENVELOPE_VERSION_V1: u8 = 1;
261const ENVELOPE_VERSION_V2: u8 = 2;
262const ENVELOPE_VERSION_V3: u8 = 3;
263const ENVELOPE_VERSION_V4: u8 = 4;
264const ENVELOPE_VERSION_V5: u8 = 5;
265
266fn build_envelope(
267 catalog: &[u8],
268 users: &[u8],
269 pubs: &[u8],
270 subs: &[u8],
271 stats: &[u8],
272) -> Vec<u8> {
273 let mut out = Vec::with_capacity(
274 8 + 1
275 + 4
276 + catalog.len()
277 + 4
278 + users.len()
279 + 4
280 + pubs.len()
281 + 4
282 + subs.len()
283 + 4
284 + stats.len()
285 + 4,
286 );
287 out.extend_from_slice(ENVELOPE_MAGIC);
288 out.push(ENVELOPE_VERSION_V5);
289 out.extend_from_slice(
290 &u32::try_from(catalog.len())
291 .expect("≤ 4G catalog")
292 .to_le_bytes(),
293 );
294 out.extend_from_slice(catalog);
295 out.extend_from_slice(
296 &u32::try_from(users.len())
297 .expect("≤ 4G users")
298 .to_le_bytes(),
299 );
300 out.extend_from_slice(users);
301 out.extend_from_slice(
302 &u32::try_from(pubs.len())
303 .expect("≤ 4G publications")
304 .to_le_bytes(),
305 );
306 out.extend_from_slice(pubs);
307 out.extend_from_slice(
308 &u32::try_from(subs.len())
309 .expect("≤ 4G subscriptions")
310 .to_le_bytes(),
311 );
312 out.extend_from_slice(subs);
313 out.extend_from_slice(
314 &u32::try_from(stats.len())
315 .expect("≤ 4G statistics")
316 .to_le_bytes(),
317 );
318 out.extend_from_slice(stats);
319 let crc = spg_crypto::crc32::crc32(&out);
320 out.extend_from_slice(&crc.to_le_bytes());
321 out
322}
323
324enum EnvelopeParse<'a> {
331 Bare,
332 Pair {
333 catalog: &'a [u8],
334 users: &'a [u8],
335 publications: Option<&'a [u8]>,
336 subscriptions: Option<&'a [u8]>,
337 statistics: Option<&'a [u8]>,
338 },
339 CrcMismatch {
340 expected: u32,
341 computed: u32,
342 },
343}
344
345fn split_envelope(buf: &[u8]) -> EnvelopeParse<'_> {
350 if buf.len() < 8 + 1 + 4 || &buf[..8] != ENVELOPE_MAGIC {
351 return EnvelopeParse::Bare;
352 }
353 let version = buf[8];
354 if !matches!(
355 version,
356 ENVELOPE_VERSION_V1
357 | ENVELOPE_VERSION_V2
358 | ENVELOPE_VERSION_V3
359 | ENVELOPE_VERSION_V4
360 | ENVELOPE_VERSION_V5
361 ) {
362 return EnvelopeParse::Bare;
363 }
364 let mut p = 9usize;
365 let Some(cat_len_bytes) = buf.get(p..p + 4) else {
366 return EnvelopeParse::Bare;
367 };
368 let Ok(cat_len_arr) = cat_len_bytes.try_into() else {
369 return EnvelopeParse::Bare;
370 };
371 let cat_len = u32::from_le_bytes(cat_len_arr) as usize;
372 p += 4;
373 if p + cat_len + 4 > buf.len() {
374 return EnvelopeParse::Bare;
375 }
376 let catalog = &buf[p..p + cat_len];
377 p += cat_len;
378 let Some(user_len_bytes) = buf.get(p..p + 4) else {
379 return EnvelopeParse::Bare;
380 };
381 let Ok(user_len_arr) = user_len_bytes.try_into() else {
382 return EnvelopeParse::Bare;
383 };
384 let user_len = u32::from_le_bytes(user_len_arr) as usize;
385 p += 4;
386 if p + user_len > buf.len() {
387 return EnvelopeParse::Bare;
388 }
389 let users = &buf[p..p + user_len];
390 p += user_len;
391 let publications = if matches!(
392 version,
393 ENVELOPE_VERSION_V3 | ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5
394 ) {
395 let Some(pubs_len_bytes) = buf.get(p..p + 4) else {
397 return EnvelopeParse::Bare;
398 };
399 let Ok(pubs_len_arr) = pubs_len_bytes.try_into() else {
400 return EnvelopeParse::Bare;
401 };
402 let pubs_len = u32::from_le_bytes(pubs_len_arr) as usize;
403 p += 4;
404 if p + pubs_len > buf.len() {
405 return EnvelopeParse::Bare;
406 }
407 let pubs_slice = &buf[p..p + pubs_len];
408 p += pubs_len;
409 Some(pubs_slice)
410 } else {
411 None
412 };
413 let subscriptions = if matches!(version, ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5) {
414 let Some(subs_len_bytes) = buf.get(p..p + 4) else {
416 return EnvelopeParse::Bare;
417 };
418 let Ok(subs_len_arr) = subs_len_bytes.try_into() else {
419 return EnvelopeParse::Bare;
420 };
421 let subs_len = u32::from_le_bytes(subs_len_arr) as usize;
422 p += 4;
423 if p + subs_len > buf.len() {
424 return EnvelopeParse::Bare;
425 }
426 let subs_slice = &buf[p..p + subs_len];
427 p += subs_len;
428 Some(subs_slice)
429 } else {
430 None
431 };
432 let statistics = if version == ENVELOPE_VERSION_V5 {
433 let Some(stats_len_bytes) = buf.get(p..p + 4) else {
435 return EnvelopeParse::Bare;
436 };
437 let Ok(stats_len_arr) = stats_len_bytes.try_into() else {
438 return EnvelopeParse::Bare;
439 };
440 let stats_len = u32::from_le_bytes(stats_len_arr) as usize;
441 p += 4;
442 if p + stats_len > buf.len() {
443 return EnvelopeParse::Bare;
444 }
445 let stats_slice = &buf[p..p + stats_len];
446 p += stats_len;
447 Some(stats_slice)
448 } else {
449 None
450 };
451 if matches!(
452 version,
453 ENVELOPE_VERSION_V2 | ENVELOPE_VERSION_V3 | ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5
454 ) {
455 if p + 4 != buf.len() {
456 return EnvelopeParse::Bare;
457 }
458 let Ok(crc_arr) = buf[p..p + 4].try_into() else {
459 return EnvelopeParse::Bare;
460 };
461 let expected = u32::from_le_bytes(crc_arr);
462 let computed = spg_crypto::crc32::crc32(&buf[..p]);
463 if expected != computed {
464 return EnvelopeParse::CrcMismatch { expected, computed };
465 }
466 } else if p != buf.len() {
467 return EnvelopeParse::Bare;
469 }
470 EnvelopeParse::Pair {
471 catalog,
472 users,
473 publications,
474 subscriptions,
475 statistics,
476 }
477}
478
479#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
489pub struct TxId(pub u64);
490
491pub const IMPLICIT_TX: TxId = TxId(0);
494
495pub const COMPACTION_TARGET_DEFAULT_BYTES: u64 = 4 * 1024 * 1024;
501
502#[derive(Debug, Default, Clone)]
507struct TxState {
508 catalog: Catalog,
513 savepoints: Vec<(String, Catalog)>,
519}
520
521#[derive(Debug, Clone)]
535pub struct CatalogSnapshot {
536 catalog: Catalog,
537 statistics: statistics::Statistics,
538 clock: Option<ClockFn>,
539 max_query_rows: Option<usize>,
540}
541
542#[derive(Debug, Default)]
543pub struct Engine {
544 catalog: Catalog,
547 tx_catalogs: BTreeMap<TxId, TxState>,
552 current_tx: Option<TxId>,
557 next_tx_id: u64,
560 clock: Option<ClockFn>,
563 salt_fn: Option<SaltFn>,
567 max_query_rows: Option<usize>,
573 users: UserStore,
579 publications: publications::Publications,
583 subscriptions: subscriptions::Subscriptions,
587 statistics: statistics::Statistics,
591 plan_cache: plan_cache::PlanCache,
595 query_stats: query_stats::QueryStats,
599 activity_provider: Option<ActivityProvider>,
606 audit_chain_provider: Option<AuditChainProvider>,
611 audit_verifier: Option<AuditVerifier>,
612 slow_query_threshold_us: Option<u64>,
618 slow_query_logger: Option<SlowQueryLogger>,
619}
620
621pub type SlowQueryLogger = fn(&str, u64);
625
626fn render_create_table(name: &str, columns: &[ColumnSchema]) -> String {
631 let mut out = alloc::format!("CREATE TABLE {name} (");
632 for (i, col) in columns.iter().enumerate() {
633 if i > 0 {
634 out.push_str(", ");
635 }
636 out.push_str(&col.name);
637 out.push(' ');
638 out.push_str(&render_data_type(col.ty));
639 if !col.nullable {
640 out.push_str(" NOT NULL");
641 }
642 if col.auto_increment {
643 out.push_str(" AUTO_INCREMENT");
644 }
645 }
646 out.push(')');
647 out
648}
649
650fn render_data_type(ty: DataType) -> String {
651 match ty {
652 DataType::SmallInt => "SMALLINT".into(),
653 DataType::Int => "INT".into(),
654 DataType::BigInt => "BIGINT".into(),
655 DataType::Float => "FLOAT".into(),
656 DataType::Text => "TEXT".into(),
657 DataType::Varchar(n) => alloc::format!("VARCHAR({n})"),
658 DataType::Char(n) => alloc::format!("CHAR({n})"),
659 DataType::Bool => "BOOL".into(),
660 DataType::Vector { dim, encoding } => match encoding {
661 spg_storage::VecEncoding::F32 => alloc::format!("VECTOR({dim})"),
662 spg_storage::VecEncoding::Sq8 => alloc::format!("VECTOR({dim}) USING SQ8"),
663 spg_storage::VecEncoding::F16 => alloc::format!("VECTOR({dim}) USING HALF"),
664 },
665 DataType::Numeric { precision, scale } => {
666 alloc::format!("NUMERIC({precision},{scale})")
667 }
668 DataType::Date => "DATE".into(),
669 DataType::Timestamp => "TIMESTAMP".into(),
670 DataType::Interval => "INTERVAL".into(),
671 DataType::Json => "JSON".into(),
672 DataType::Jsonb => "JSONB".into(),
673 DataType::Timestamptz => "TIMESTAMPTZ".into(),
674 DataType::Bytes => "BYTEA".into(),
675 DataType::TextArray => "TEXT[]".into(),
676 }
677}
678
679#[derive(Debug, Clone)]
683pub struct ActivityRow {
684 pub pid: u32,
685 pub user: String,
686 pub started_at_us: i64,
687 pub current_sql: String,
688 pub wait_event: String,
689 pub elapsed_us: i64,
690 pub in_transaction: bool,
691}
692
693pub type ActivityProvider = fn() -> Vec<ActivityRow>;
696
697#[derive(Debug, Clone)]
700pub struct AuditRow {
701 pub seq: i64,
702 pub ts_ms: i64,
703 pub prev_hash_hex: String,
704 pub entry_hash_hex: String,
705 pub sql: String,
706}
707
708pub type AuditChainProvider = fn() -> Vec<AuditRow>;
713pub type AuditVerifier = fn() -> (i64, i64);
714
715impl Engine {
716 pub fn new() -> Self {
717 Self {
718 catalog: Catalog::new(),
719 tx_catalogs: BTreeMap::new(),
720 current_tx: None,
721 next_tx_id: 1,
722 clock: None,
723 salt_fn: None,
724 max_query_rows: None,
725 users: UserStore::new(),
726 publications: publications::Publications::new(),
727 subscriptions: subscriptions::Subscriptions::new(),
728 statistics: statistics::Statistics::new(),
729 plan_cache: plan_cache::PlanCache::new(),
730 query_stats: query_stats::QueryStats::new(),
731 activity_provider: None,
732 audit_chain_provider: None,
733 audit_verifier: None,
734 slow_query_threshold_us: None,
735 slow_query_logger: None,
736 }
737 }
738
739 #[must_use]
748 pub fn clone_snapshot(&self) -> CatalogSnapshot {
749 CatalogSnapshot {
750 catalog: self.active_catalog().clone(),
751 statistics: self.statistics.clone(),
752 clock: self.clock,
753 max_query_rows: self.max_query_rows,
754 }
755 }
756
757 pub fn execute_readonly_on_snapshot(
765 snapshot: &CatalogSnapshot,
766 sql: &str,
767 ) -> Result<QueryResult, EngineError> {
768 Self::execute_readonly_on_snapshot_with_cancel(snapshot, sql, CancelToken::none())
769 }
770
771 pub fn execute_readonly_on_snapshot_with_cancel(
778 snapshot: &CatalogSnapshot,
779 sql: &str,
780 cancel: CancelToken<'_>,
781 ) -> Result<QueryResult, EngineError> {
782 let transient = Engine {
783 catalog: snapshot.catalog.clone(),
784 statistics: snapshot.statistics.clone(),
785 clock: snapshot.clock,
786 max_query_rows: snapshot.max_query_rows,
787 ..Engine::default()
788 };
789 transient.execute_readonly_with_cancel(sql, cancel)
790 }
791
792 pub fn restore(catalog: Catalog) -> Self {
795 Self {
796 catalog,
797 tx_catalogs: BTreeMap::new(),
798 current_tx: None,
799 next_tx_id: 1,
800 clock: None,
801 salt_fn: None,
802 max_query_rows: None,
803 users: UserStore::new(),
804 publications: publications::Publications::new(),
805 subscriptions: subscriptions::Subscriptions::new(),
806 statistics: statistics::Statistics::new(),
807 plan_cache: plan_cache::PlanCache::new(),
808 query_stats: query_stats::QueryStats::new(),
809 activity_provider: None,
810 audit_chain_provider: None,
811 audit_verifier: None,
812 slow_query_threshold_us: None,
813 slow_query_logger: None,
814 }
815 }
816
817 pub fn restore_envelope(buf: &[u8]) -> Result<Self, EngineError> {
824 match split_envelope(buf) {
825 EnvelopeParse::Pair {
826 catalog: catalog_bytes,
827 users: user_bytes,
828 publications: pub_bytes,
829 subscriptions: sub_bytes,
830 statistics: stats_bytes,
831 } => {
832 let catalog = Catalog::deserialize(catalog_bytes).map_err(EngineError::Storage)?;
833 let users = users::deserialize_users(user_bytes)
834 .map_err(|e| EngineError::Unsupported(alloc::format!("users restore: {e}")))?;
835 let publications = match pub_bytes {
836 Some(b) => publications::Publications::deserialize(b).map_err(|e| {
837 EngineError::Unsupported(alloc::format!("publications restore: {e:?}"))
838 })?,
839 None => publications::Publications::new(),
840 };
841 let subscriptions = match sub_bytes {
842 Some(b) => subscriptions::Subscriptions::deserialize(b).map_err(|e| {
843 EngineError::Unsupported(alloc::format!("subscriptions restore: {e:?}"))
844 })?,
845 None => subscriptions::Subscriptions::new(),
846 };
847 let statistics = match stats_bytes {
848 Some(b) => statistics::Statistics::deserialize(b).map_err(|e| {
849 EngineError::Unsupported(alloc::format!("statistics restore: {e:?}"))
850 })?,
851 None => statistics::Statistics::new(),
852 };
853 Ok(Self {
854 catalog,
855 tx_catalogs: BTreeMap::new(),
856 current_tx: None,
857 next_tx_id: 1,
858 clock: None,
859 salt_fn: None,
860 max_query_rows: None,
861 users,
862 publications,
863 subscriptions,
864 statistics,
865 plan_cache: plan_cache::PlanCache::new(),
866 query_stats: query_stats::QueryStats::new(),
867 activity_provider: None,
868 audit_chain_provider: None,
869 audit_verifier: None,
870 slow_query_threshold_us: None,
871 slow_query_logger: None,
872 })
873 }
874 EnvelopeParse::CrcMismatch { expected, computed } => {
875 Err(EngineError::Storage(StorageError::Corrupt(alloc::format!(
876 "snapshot envelope CRC32 mismatch (expected={expected:#010x}, computed={computed:#010x})"
877 ))))
878 }
879 EnvelopeParse::Bare => {
880 let catalog = Catalog::deserialize(buf).map_err(EngineError::Storage)?;
881 Ok(Self::restore(catalog))
882 }
883 }
884 }
885
886 pub const fn users(&self) -> &UserStore {
887 &self.users
888 }
889
890 pub fn create_user(
894 &mut self,
895 name: &str,
896 password: &str,
897 role: Role,
898 salt: [u8; 16],
899 ) -> Result<(), UserError> {
900 self.users.create(name, password, role, salt)?;
901 let scram_salt = self.salt_fn.map_or_else(
907 || {
908 let mut s = [0u8; users::SCRAM_SALT_LEN];
909 let digest = spg_crypto::hash(name.as_bytes());
910 s.copy_from_slice(&digest[16..32]);
913 s
914 },
915 |f| f(),
916 );
917 self.users
918 .enable_scram(name, password, scram_salt, users::SCRAM_DEFAULT_ITERS)?;
919 Ok(())
920 }
921
922 pub fn drop_user(&mut self, name: &str) -> Result<(), UserError> {
923 self.users.drop(name)
924 }
925
926 pub fn verify_user(&self, name: &str, password: &str) -> Option<Role> {
927 self.users.verify(name, password)
928 }
929
930 #[must_use]
933 pub const fn with_clock(mut self, clock: ClockFn) -> Self {
934 self.clock = Some(clock);
935 self
936 }
937
938 #[must_use]
941 pub const fn with_salt_fn(mut self, f: SaltFn) -> Self {
942 self.salt_fn = Some(f);
943 self
944 }
945
946 #[must_use]
952 pub const fn with_max_query_rows(mut self, n: usize) -> Self {
953 self.max_query_rows = Some(n);
954 self
955 }
956
957 pub const fn catalog(&self) -> &Catalog {
961 &self.catalog
962 }
963
964 pub fn snapshot(&self) -> Vec<u8> {
972 if self.users.is_empty()
973 && self.publications.is_empty()
974 && self.subscriptions.is_empty()
975 && self.statistics.is_empty()
976 {
977 self.catalog.serialize()
978 } else {
979 build_envelope(
980 &self.catalog.serialize(),
981 &users::serialize_users(&self.users),
982 &self.publications.serialize(),
983 &self.subscriptions.serialize(),
984 &self.statistics.serialize(),
985 )
986 }
987 }
988
989 pub fn in_transaction(&self) -> bool {
994 !self.tx_catalogs.is_empty()
995 }
996
997 pub fn alloc_tx_id(&mut self) -> TxId {
1006 let id = TxId(self.next_tx_id);
1007 self.next_tx_id = self.next_tx_id.saturating_add(1);
1008 id
1009 }
1010
1011 pub fn replace_catalog(&mut self, catalog: Catalog) {
1031 self.catalog = catalog;
1032 }
1033
1034 pub fn freeze_oldest_to_cold(
1042 &mut self,
1043 table_name: &str,
1044 index_name: &str,
1045 max_rows: usize,
1046 ) -> Result<spg_storage::FreezeReport, EngineError> {
1047 let report = self
1048 .active_catalog_mut()
1049 .freeze_oldest_to_cold(table_name, index_name, max_rows)
1050 .map_err(EngineError::Storage)?;
1051 if let Some(t) = self.active_catalog_mut().get_mut(table_name) {
1052 t.mark_cold_row_count_stale();
1053 }
1054 Ok(report)
1055 }
1056
1057 pub fn receive_cold_segment(
1071 &mut self,
1072 segment_id: u32,
1073 bytes: Vec<u8>,
1074 ) -> Result<(), EngineError> {
1075 let mut new_cat = self.catalog.clone();
1076 match new_cat.load_segment_bytes_at(segment_id, bytes) {
1077 Ok(()) => {
1078 self.replace_catalog(new_cat);
1079 Ok(())
1080 }
1081 Err(StorageError::Corrupt(msg)) if msg.contains("already occupied") => Ok(()),
1082 Err(e) => Err(EngineError::Storage(e)),
1083 }
1084 }
1085
1086 pub fn compact_cold_segments_with_target(
1100 &mut self,
1101 target_segment_bytes: u64,
1102 ) -> Result<Vec<(String, String, CompactReport)>, EngineError> {
1103 let table_names = self.active_catalog().table_names();
1104 let mut reports: Vec<(String, String, CompactReport)> = Vec::new();
1105 for tname in table_names {
1106 if is_internal_table_name(&tname) {
1107 continue;
1108 }
1109 let idx_names: Vec<String> = {
1110 let Some(t) = self.active_catalog().get(&tname) else {
1111 continue;
1112 };
1113 t.indices()
1114 .iter()
1115 .filter(|i| matches!(i.kind, IndexKind::BTree(_)))
1116 .map(|i| i.name.clone())
1117 .collect()
1118 };
1119 for iname in idx_names {
1120 let report = self
1121 .active_catalog_mut()
1122 .compact_cold_segments(&tname, &iname, target_segment_bytes)
1123 .map_err(EngineError::Storage)?;
1124 if report.merged_segment_id.is_some() {
1125 if let Some(t) = self.active_catalog_mut().get_mut(&tname) {
1126 t.mark_cold_row_count_stale();
1127 }
1128 reports.push((tname.clone(), iname, report));
1129 }
1130 }
1131 }
1132 Ok(reports)
1133 }
1134
1135 fn active_catalog(&self) -> &Catalog {
1136 match self.current_tx {
1137 Some(t) => self
1138 .tx_catalogs
1139 .get(&t)
1140 .map_or(&self.catalog, |s| &s.catalog),
1141 None => &self.catalog,
1142 }
1143 }
1144
1145 fn active_catalog_mut(&mut self) -> &mut Catalog {
1146 let tx = self.current_tx;
1147 match tx {
1148 Some(t) => match self.tx_catalogs.get_mut(&t) {
1149 Some(s) => &mut s.catalog,
1150 None => &mut self.catalog,
1151 },
1152 None => &mut self.catalog,
1153 }
1154 }
1155
1156 pub fn execute_readonly(&self, sql: &str) -> Result<QueryResult, EngineError> {
1168 self.execute_readonly_with_cancel(sql, CancelToken::none())
1169 }
1170
1171 pub fn execute_readonly_with_cancel(
1177 &self,
1178 sql: &str,
1179 cancel: CancelToken<'_>,
1180 ) -> Result<QueryResult, EngineError> {
1181 cancel.check()?;
1182 let mut stmt = parser::parse_statement(sql)?;
1183 let now_micros = self.clock.map(|f| f());
1184 rewrite_clock_calls(&mut stmt, now_micros);
1185 if let Statement::Select(s) = &mut stmt {
1186 resolve_order_by_position(s);
1187 reorder::reorder_joins(s, &self.catalog, &self.statistics);
1189 }
1190 let result = match stmt {
1191 Statement::Select(s) => self.exec_select_cancel(&s, cancel),
1192 Statement::ShowTables => Ok(self.exec_show_tables()),
1193 Statement::ShowColumns(table) => self.exec_show_columns(&table),
1194 Statement::ShowUsers => Ok(self.exec_show_users()),
1195 Statement::ShowPublications => Ok(self.exec_show_publications()),
1196 Statement::ShowSubscriptions => Ok(self.exec_show_subscriptions()),
1197 Statement::WaitForWalPosition { .. } => Err(EngineError::Unsupported(
1198 "WAIT FOR WAL POSITION must be handled by the server layer".into(),
1199 )),
1200 Statement::Explain(e) => self.exec_explain(&e, cancel),
1201 _ => Err(EngineError::WriteRequired),
1202 };
1203 self.enforce_row_limit(result)
1204 }
1205
1206 fn enforce_row_limit(
1210 &self,
1211 result: Result<QueryResult, EngineError>,
1212 ) -> Result<QueryResult, EngineError> {
1213 if let (Ok(QueryResult::Rows { rows, .. }), Some(cap)) = (&result, self.max_query_rows)
1214 && rows.len() > cap
1215 {
1216 return Err(EngineError::RowLimitExceeded(cap));
1217 }
1218 result
1219 }
1220
1221 pub fn execute(&mut self, sql: &str) -> Result<QueryResult, EngineError> {
1222 self.execute_in_with_cancel(sql, IMPLICIT_TX, CancelToken::none())
1223 }
1224
1225 pub fn execute_with_cancel(
1230 &mut self,
1231 sql: &str,
1232 cancel: CancelToken<'_>,
1233 ) -> Result<QueryResult, EngineError> {
1234 self.execute_in_with_cancel(sql, IMPLICIT_TX, cancel)
1235 }
1236
1237 pub fn execute_in(&mut self, sql: &str, tx_id: TxId) -> Result<QueryResult, EngineError> {
1244 self.execute_in_with_cancel(sql, tx_id, CancelToken::none())
1245 }
1246
1247 pub fn execute_in_with_cancel(
1253 &mut self,
1254 sql: &str,
1255 tx_id: TxId,
1256 cancel: CancelToken<'_>,
1257 ) -> Result<QueryResult, EngineError> {
1258 let saved = self.current_tx;
1259 self.current_tx = Some(tx_id);
1260 let result = self.execute_inner_with_cancel(sql, cancel);
1261 self.current_tx = saved;
1262 result
1263 }
1264
1265 pub fn prepare(&self, sql: &str) -> Result<Statement, ParseError> {
1277 let mut stmt = parser::parse_statement(sql)?;
1278 let now_micros = self.clock.map(|f| f());
1279 rewrite_clock_calls(&mut stmt, now_micros);
1280 if let Statement::Select(s) = &mut stmt {
1281 expand_group_by_all(s);
1285 resolve_order_by_position(s);
1286 reorder::reorder_joins(s, &self.catalog, &self.statistics);
1289 }
1290 Ok(stmt)
1291 }
1292
1293 pub fn prepare_cached(&mut self, sql: &str) -> Result<Statement, ParseError> {
1305 let current_version = self.statistics.version();
1308 if let Some(plan) = self.plan_cache.get(sql) {
1309 if plan.statistics_version == current_version {
1310 return Ok(plan.stmt.clone());
1311 }
1312 }
1314 self.plan_cache.evict(sql);
1315 let stmt = self.prepare(sql)?;
1316 let source_tables = plan_cache::collect_source_tables(&stmt);
1317 let plan = plan_cache::PreparedPlan {
1318 stmt: stmt.clone(),
1319 statistics_version: current_version,
1320 source_tables,
1321 describe_columns: alloc::vec::Vec::new(),
1322 };
1323 self.plan_cache.insert(String::from(sql), plan);
1324 Ok(stmt)
1325 }
1326
1327 pub fn plan_cache(&self) -> &plan_cache::PlanCache {
1329 &self.plan_cache
1330 }
1331
1332 pub fn plan_cache_mut(&mut self) -> &mut plan_cache::PlanCache {
1334 &mut self.plan_cache
1335 }
1336
1337 pub fn describe_prepared(
1343 &self,
1344 stmt: &Statement,
1345 ) -> (Vec<u32>, Vec<ColumnSchema>) {
1346 describe::describe_prepared(stmt, self.active_catalog())
1347 }
1348
1349 pub fn execute_prepared(
1359 &mut self,
1360 mut stmt: Statement,
1361 params: &[Value],
1362 ) -> Result<QueryResult, EngineError> {
1363 substitute_placeholders(&mut stmt, params)?;
1364 self.execute_stmt_with_cancel(stmt, CancelToken::none())
1365 }
1366
1367 fn execute_inner_with_cancel(
1368 &mut self,
1369 sql: &str,
1370 cancel: CancelToken<'_>,
1371 ) -> Result<QueryResult, EngineError> {
1372 cancel.check()?;
1373 let stmt = self.prepare(sql)?;
1374 let start_us = self.clock.map(|f| f());
1378 let result = self.execute_stmt_with_cancel(stmt, cancel);
1379 if let (Some(t0), Ok(_)) = (start_us, &result) {
1380 let now = self.clock.map_or(t0, |f| f());
1381 let elapsed = now.saturating_sub(t0).max(0) as u64;
1382 self.query_stats.record(sql, elapsed, now as u64);
1383 if let (Some(threshold), Some(logger)) =
1386 (self.slow_query_threshold_us, self.slow_query_logger)
1387 && elapsed >= threshold
1388 {
1389 logger(sql, elapsed);
1390 }
1391 }
1392 result
1393 }
1394
1395 fn execute_stmt_with_cancel(
1396 &mut self,
1397 stmt: Statement,
1398 cancel: CancelToken<'_>,
1399 ) -> Result<QueryResult, EngineError> {
1400 cancel.check()?;
1401 let result = match stmt {
1402 Statement::CreateTable(s) => self.exec_create_table(s),
1403 Statement::CreateExtension(_) => Ok(QueryResult::CommandOk {
1407 affected: 0,
1408 modified_catalog: false,
1409 }),
1410 Statement::DoBlock => Ok(QueryResult::CommandOk {
1413 affected: 0,
1414 modified_catalog: false,
1415 }),
1416 Statement::CreateIndex(s) => self.exec_create_index(s),
1417 Statement::Insert(s) => self.exec_insert(s),
1418 Statement::Update(s) => self.exec_update_cancel(&s, cancel),
1419 Statement::Delete(s) => self.exec_delete_cancel(&s, cancel),
1420 Statement::Select(s) => self.exec_select_cancel(&s, cancel),
1421 Statement::Begin => self.exec_begin(),
1422 Statement::Commit => self.exec_commit(),
1423 Statement::Rollback => self.exec_rollback(),
1424 Statement::Savepoint(name) => self.exec_savepoint(name),
1425 Statement::RollbackToSavepoint(name) => self.exec_rollback_to_savepoint(&name),
1426 Statement::ReleaseSavepoint(name) => self.exec_release_savepoint(&name),
1427 Statement::ShowTables => Ok(self.exec_show_tables()),
1428 Statement::ShowColumns(table) => self.exec_show_columns(&table),
1429 Statement::ShowUsers => Ok(self.exec_show_users()),
1430 Statement::ShowPublications => Ok(self.exec_show_publications()),
1431 Statement::ShowSubscriptions => Ok(self.exec_show_subscriptions()),
1432 Statement::CreateUser(s) => self.exec_create_user(&s),
1433 Statement::DropUser(name) => self.exec_drop_user(&name),
1434 Statement::Explain(e) => self.exec_explain(&e, cancel),
1435 Statement::AlterIndex(s) => self.exec_alter_index(s),
1436 Statement::AlterTable(s) => self.exec_alter_table(s),
1437 Statement::CreatePublication(s) => self.exec_create_publication(s),
1438 Statement::DropPublication(name) => self.exec_drop_publication(&name),
1439 Statement::CreateSubscription(s) => self.exec_create_subscription(s),
1440 Statement::DropSubscription(name) => self.exec_drop_subscription(&name),
1441 Statement::WaitForWalPosition { .. } => Err(EngineError::Unsupported(
1448 "WAIT FOR WAL POSITION must be handled by the server layer".into(),
1449 )),
1450 Statement::Analyze(target) => self.exec_analyze(target.as_deref()),
1452 Statement::CompactColdSegments => self.exec_compact_cold_segments(),
1454 };
1455 self.enforce_row_limit(result)
1456 }
1457
1458 fn exec_create_publication(
1466 &mut self,
1467 s: CreatePublicationStatement,
1468 ) -> Result<QueryResult, EngineError> {
1469 self.publications
1475 .create(s.name, s.scope)
1476 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE PUBLICATION: {e:?}")))?;
1477 Ok(QueryResult::CommandOk {
1478 affected: 1,
1479 modified_catalog: true,
1480 })
1481 }
1482
1483 fn exec_drop_publication(&mut self, name: &str) -> Result<QueryResult, EngineError> {
1488 let removed = self.publications.drop(name);
1489 Ok(QueryResult::CommandOk {
1490 affected: usize::from(removed),
1491 modified_catalog: removed,
1492 })
1493 }
1494
1495 pub const fn publications(&self) -> &publications::Publications {
1500 &self.publications
1501 }
1502
1503 fn exec_create_subscription(
1508 &mut self,
1509 s: CreateSubscriptionStatement,
1510 ) -> Result<QueryResult, EngineError> {
1511 let sub = subscriptions::Subscription {
1515 conn_str: s.conn_str,
1516 publications: s.publications,
1517 enabled: true,
1518 last_received_pos: 0,
1519 };
1520 self.subscriptions
1521 .create(s.name, sub)
1522 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE SUBSCRIPTION: {e:?}")))?;
1523 Ok(QueryResult::CommandOk {
1524 affected: 1,
1525 modified_catalog: true,
1526 })
1527 }
1528
1529 fn exec_drop_subscription(&mut self, name: &str) -> Result<QueryResult, EngineError> {
1537 let removed = self.subscriptions.drop(name);
1538 Ok(QueryResult::CommandOk {
1539 affected: usize::from(removed),
1540 modified_catalog: removed,
1541 })
1542 }
1543
1544 pub const fn subscriptions(&self) -> &subscriptions::Subscriptions {
1549 &self.subscriptions
1550 }
1551
1552 pub fn subscription_advance(&mut self, name: &str, pos: u64) -> bool {
1558 self.subscriptions.update_last_received_pos(name, pos)
1559 }
1560
1561 fn exec_show_subscriptions(&self) -> QueryResult {
1567 let columns = alloc::vec![
1568 ColumnSchema::new("name", DataType::Text, false),
1569 ColumnSchema::new("conn_str", DataType::Text, false),
1570 ColumnSchema::new("publications", DataType::Text, false),
1571 ColumnSchema::new("enabled", DataType::Bool, false),
1572 ColumnSchema::new("last_received_pos", DataType::BigInt, false),
1573 ];
1574 let rows: Vec<Row> = self
1575 .subscriptions
1576 .iter()
1577 .map(|(name, sub)| {
1578 Row::new(alloc::vec![
1579 Value::Text(name.clone()),
1580 Value::Text(sub.conn_str.clone()),
1581 Value::Text(sub.publications.join(", ")),
1582 Value::Bool(sub.enabled),
1583 Value::BigInt(i64::try_from(sub.last_received_pos).unwrap_or(i64::MAX)),
1584 ])
1585 })
1586 .collect();
1587 QueryResult::Rows { columns, rows }
1588 }
1589
1590 fn exec_spg_statistic(&self) -> QueryResult {
1595 let columns = alloc::vec![
1596 ColumnSchema::new("table_name", DataType::Text, false),
1597 ColumnSchema::new("column_name", DataType::Text, false),
1598 ColumnSchema::new("null_frac", DataType::Float, false),
1599 ColumnSchema::new("n_distinct", DataType::BigInt, false),
1600 ColumnSchema::new("histogram_bounds", DataType::Text, false),
1601 ColumnSchema::new("cold_row_count", DataType::BigInt, false),
1606 ];
1607 let rows: Vec<Row> = self
1608 .statistics
1609 .iter()
1610 .map(|((t, c), s)| {
1611 let cold = self
1612 .catalog
1613 .get(t)
1614 .map_or(0, |table| table.cold_row_count());
1615 Row::new(alloc::vec![
1616 Value::Text(t.clone()),
1617 Value::Text(c.clone()),
1618 Value::Float(f64::from(s.null_frac)),
1619 Value::BigInt(i64::try_from(s.n_distinct).unwrap_or(i64::MAX)),
1620 Value::Text(render_histogram_bounds(&s.histogram_bounds)),
1621 Value::BigInt(i64::try_from(cold).unwrap_or(i64::MAX)),
1622 ])
1623 })
1624 .collect();
1625 QueryResult::Rows { columns, rows }
1626 }
1627
1628 fn exec_spg_stat_replication(&self) -> QueryResult {
1635 let columns = alloc::vec![
1636 ColumnSchema::new("name", DataType::Text, false),
1637 ColumnSchema::new("conn_str", DataType::Text, false),
1638 ColumnSchema::new("publications", DataType::Text, false),
1639 ColumnSchema::new("last_received_pos", DataType::BigInt, false),
1640 ColumnSchema::new("enabled", DataType::Bool, false),
1641 ];
1642 let rows: Vec<Row> = self
1643 .subscriptions
1644 .iter()
1645 .map(|(name, sub)| {
1646 Row::new(alloc::vec![
1647 Value::Text(name.clone()),
1648 Value::Text(sub.conn_str.clone()),
1649 Value::Text(sub.publications.join(",")),
1650 Value::BigInt(i64::try_from(sub.last_received_pos).unwrap_or(i64::MAX)),
1651 Value::Bool(sub.enabled),
1652 ])
1653 })
1654 .collect();
1655 QueryResult::Rows { columns, rows }
1656 }
1657
1658 fn exec_spg_stat_segment(&self) -> QueryResult {
1670 let columns = alloc::vec![
1671 ColumnSchema::new("segment_id", DataType::BigInt, false),
1672 ColumnSchema::new("table_name", DataType::Text, false),
1673 ColumnSchema::new("num_rows", DataType::BigInt, false),
1674 ColumnSchema::new("num_pages", DataType::BigInt, false),
1675 ColumnSchema::new("total_bytes", DataType::BigInt, false),
1676 ];
1677 let mut segment_owners: alloc::collections::BTreeMap<u32, String> = BTreeMap::new();
1683 for tname in self.catalog.table_names() {
1684 if is_internal_table_name(&tname) {
1685 continue;
1686 }
1687 let Some(t) = self.catalog.get(&tname) else {
1688 continue;
1689 };
1690 for idx in t.indices() {
1691 if let spg_storage::IndexKind::BTree(map) = &idx.kind {
1692 for (_, locs) in map.iter() {
1693 for loc in locs {
1694 if let spg_storage::RowLocator::Cold { segment_id, .. } = loc {
1695 segment_owners.entry(*segment_id).or_insert_with(|| tname.clone());
1696 }
1697 }
1698 }
1699 }
1700 }
1701 }
1702 let rows: Vec<Row> = self
1703 .catalog
1704 .cold_segment_ids_global()
1705 .iter()
1706 .filter_map(|&id| {
1707 let seg = self.catalog.cold_segment(id)?;
1708 let meta = seg.meta();
1709 let owner = segment_owners
1710 .get(&id)
1711 .cloned()
1712 .unwrap_or_default();
1713 Some(Row::new(alloc::vec![
1714 Value::BigInt(i64::from(id)),
1715 Value::Text(owner),
1716 Value::BigInt(i64::try_from(meta.num_rows).unwrap_or(i64::MAX)),
1717 Value::BigInt(i64::from(meta.num_pages)),
1718 Value::BigInt(i64::try_from(meta.total_bytes).unwrap_or(i64::MAX)),
1719 ]))
1720 })
1721 .collect();
1722 QueryResult::Rows { columns, rows }
1723 }
1724
1725 fn exec_spg_stat_query(&self) -> QueryResult {
1731 let columns = alloc::vec![
1732 ColumnSchema::new("sql", DataType::Text, false),
1733 ColumnSchema::new("exec_count", DataType::BigInt, false),
1734 ColumnSchema::new("total_us", DataType::BigInt, false),
1735 ColumnSchema::new("mean_us", DataType::BigInt, false),
1736 ColumnSchema::new("max_us", DataType::BigInt, false),
1737 ColumnSchema::new("last_seen_us", DataType::BigInt, false),
1738 ];
1739 let rows: Vec<Row> = self
1740 .query_stats
1741 .snapshot()
1742 .into_iter()
1743 .map(|(sql, s)| {
1744 let mean = if s.exec_count == 0 {
1745 0
1746 } else {
1747 s.total_us / s.exec_count
1748 };
1749 Row::new(alloc::vec![
1750 Value::Text(sql),
1751 Value::BigInt(i64::try_from(s.exec_count).unwrap_or(i64::MAX)),
1752 Value::BigInt(i64::try_from(s.total_us).unwrap_or(i64::MAX)),
1753 Value::BigInt(i64::try_from(mean).unwrap_or(i64::MAX)),
1754 Value::BigInt(i64::try_from(s.max_us).unwrap_or(i64::MAX)),
1755 Value::BigInt(i64::try_from(s.last_seen_us).unwrap_or(i64::MAX)),
1756 ])
1757 })
1758 .collect();
1759 QueryResult::Rows { columns, rows }
1760 }
1761
1762 #[must_use]
1767 pub const fn with_activity_provider(mut self, f: ActivityProvider) -> Self {
1768 self.activity_provider = Some(f);
1769 self
1770 }
1771
1772 #[must_use]
1774 pub const fn with_audit_providers(
1775 mut self,
1776 chain: AuditChainProvider,
1777 verify: AuditVerifier,
1778 ) -> Self {
1779 self.audit_chain_provider = Some(chain);
1780 self.audit_verifier = Some(verify);
1781 self
1782 }
1783
1784 #[must_use]
1789 pub const fn with_slow_query_log(
1790 mut self,
1791 threshold_us: u64,
1792 logger: SlowQueryLogger,
1793 ) -> Self {
1794 self.slow_query_threshold_us = Some(threshold_us);
1795 self.slow_query_logger = Some(logger);
1796 self
1797 }
1798
1799 pub fn set_plan_cache_max(&mut self, n: usize) {
1803 self.plan_cache.set_max_entries(n);
1804 }
1805
1806 fn exec_spg_stat_activity(&self) -> QueryResult {
1811 let columns = alloc::vec![
1812 ColumnSchema::new("pid", DataType::Int, false),
1813 ColumnSchema::new("user", DataType::Text, false),
1814 ColumnSchema::new("started_at_us", DataType::BigInt, false),
1815 ColumnSchema::new("current_sql", DataType::Text, false),
1816 ColumnSchema::new("wait_event", DataType::Text, false),
1817 ColumnSchema::new("elapsed_us", DataType::BigInt, false),
1818 ColumnSchema::new("in_transaction", DataType::Bool, false),
1819 ];
1820 let rows: Vec<Row> = self
1821 .activity_provider
1822 .map(|f| f())
1823 .unwrap_or_default()
1824 .into_iter()
1825 .map(|r| {
1826 Row::new(alloc::vec![
1827 Value::Int(i32::try_from(r.pid).unwrap_or(i32::MAX)),
1828 Value::Text(r.user),
1829 Value::BigInt(r.started_at_us),
1830 Value::Text(r.current_sql),
1831 Value::Text(r.wait_event),
1832 Value::BigInt(r.elapsed_us),
1833 Value::Bool(r.in_transaction),
1834 ])
1835 })
1836 .collect();
1837 QueryResult::Rows { columns, rows }
1838 }
1839
1840 fn exec_spg_table_ddl(&self) -> QueryResult {
1844 let columns = alloc::vec![
1845 ColumnSchema::new("table_name", DataType::Text, false),
1846 ColumnSchema::new("ddl", DataType::Text, false),
1847 ];
1848 let rows: Vec<Row> = self
1849 .catalog
1850 .table_names()
1851 .into_iter()
1852 .filter(|n| !is_internal_table_name(n))
1853 .filter_map(|name| {
1854 let table = self.catalog.get(&name)?;
1855 let ddl = render_create_table(&name, &table.schema().columns);
1856 Some(Row::new(alloc::vec![
1857 Value::Text(name),
1858 Value::Text(ddl),
1859 ]))
1860 })
1861 .collect();
1862 QueryResult::Rows { columns, rows }
1863 }
1864
1865 fn exec_spg_role_ddl(&self) -> QueryResult {
1869 let columns = alloc::vec![
1870 ColumnSchema::new("role_name", DataType::Text, false),
1871 ColumnSchema::new("ddl", DataType::Text, false),
1872 ];
1873 let rows: Vec<Row> = self
1874 .users
1875 .iter()
1876 .map(|(name, rec)| {
1877 let ddl = alloc::format!(
1878 "CREATE USER {name} WITH PASSWORD '<redacted>' ROLE '{}'",
1879 rec.role.as_str(),
1880 );
1881 Row::new(alloc::vec![Value::Text(String::from(name)), Value::Text(ddl)])
1882 })
1883 .collect();
1884 QueryResult::Rows { columns, rows }
1885 }
1886
1887 fn exec_spg_database_ddl(&self) -> QueryResult {
1893 let columns = alloc::vec![ColumnSchema::new("ddl", DataType::Text, false)];
1894 let mut out = String::new();
1895 for (name, rec) in self.users.iter() {
1896 out.push_str(&alloc::format!(
1897 "CREATE USER {name} WITH PASSWORD '<redacted>' ROLE '{}';\n",
1898 rec.role.as_str(),
1899 ));
1900 }
1901 for name in self.catalog.table_names() {
1902 if is_internal_table_name(&name) {
1903 continue;
1904 }
1905 if let Some(table) = self.catalog.get(&name) {
1906 out.push_str(&render_create_table(&name, &table.schema().columns));
1907 out.push_str(";\n");
1908 }
1909 }
1910 QueryResult::Rows {
1911 columns,
1912 rows: alloc::vec![Row::new(alloc::vec![Value::Text(out)])],
1913 }
1914 }
1915
1916 fn exec_spg_audit_chain(&self) -> QueryResult {
1920 let columns = alloc::vec![
1921 ColumnSchema::new("seq", DataType::BigInt, false),
1922 ColumnSchema::new("ts_ms", DataType::BigInt, false),
1923 ColumnSchema::new("prev_hash", DataType::Text, false),
1924 ColumnSchema::new("entry_hash", DataType::Text, false),
1925 ColumnSchema::new("sql", DataType::Text, false),
1926 ];
1927 let rows: Vec<Row> = self
1928 .audit_chain_provider
1929 .map(|f| f())
1930 .unwrap_or_default()
1931 .into_iter()
1932 .map(|r| {
1933 Row::new(alloc::vec![
1934 Value::BigInt(r.seq),
1935 Value::BigInt(r.ts_ms),
1936 Value::Text(r.prev_hash_hex),
1937 Value::Text(r.entry_hash_hex),
1938 Value::Text(r.sql),
1939 ])
1940 })
1941 .collect();
1942 QueryResult::Rows { columns, rows }
1943 }
1944
1945 fn exec_spg_audit_verify(&self) -> QueryResult {
1951 let columns = alloc::vec![
1952 ColumnSchema::new("verified_count", DataType::BigInt, false),
1953 ColumnSchema::new("broken_at_seq", DataType::BigInt, false),
1954 ];
1955 let (verified, broken) = self.audit_verifier.map(|f| f()).unwrap_or((0, -1));
1956 let row = Row::new(alloc::vec![
1957 Value::BigInt(verified),
1958 Value::BigInt(broken),
1959 ]);
1960 QueryResult::Rows {
1961 columns,
1962 rows: alloc::vec![row],
1963 }
1964 }
1965
1966 pub fn query_stats(&self) -> &query_stats::QueryStats {
1968 &self.query_stats
1969 }
1970
1971 pub fn query_stats_mut(&mut self) -> &mut query_stats::QueryStats {
1973 &mut self.query_stats
1974 }
1975
1976 pub const fn statistics(&self) -> &statistics::Statistics {
1980 &self.statistics
1981 }
1982
1983 pub fn tables_needing_analyze(&self) -> Vec<String> {
1996 const MIN_ROWS: u64 = 100;
1997 let mut out = Vec::new();
1998 for name in self.catalog.table_names() {
1999 if is_internal_table_name(&name) {
2000 continue;
2001 }
2002 let Some(table) = self.catalog.get(&name) else {
2003 continue;
2004 };
2005 let row_count = table.rows().len() as u64;
2006 let modified = self.statistics.modified_since_last_analyze(&name);
2007 let base = row_count.max(MIN_ROWS);
2012 let threshold = base.saturating_add(9) / 10;
2013 if modified >= threshold {
2014 out.push(name);
2015 }
2016 }
2017 out
2018 }
2019
2020 fn exec_analyze(&mut self, target: Option<&str>) -> Result<QueryResult, EngineError> {
2031 let names: Vec<String> = if let Some(name) = target {
2032 if self.catalog.get(name).is_none() {
2034 return Err(EngineError::Storage(StorageError::TableNotFound {
2035 name: name.to_string(),
2036 }));
2037 }
2038 alloc::vec![name.to_string()]
2039 } else {
2040 self.catalog
2041 .table_names()
2042 .into_iter()
2043 .filter(|n| !is_internal_table_name(n))
2044 .collect()
2045 };
2046 let mut analysed = 0usize;
2047 for table_name in &names {
2048 self.analyze_one_table(table_name)?;
2049 analysed += 1;
2050 }
2051 if analysed > 0 {
2057 self.statistics.bump_version();
2058 if target.is_some() {
2059 for t in &names {
2060 self.plan_cache.evict_referencing(t);
2061 }
2062 } else {
2063 self.plan_cache.clear();
2064 }
2065 }
2066 Ok(QueryResult::CommandOk {
2067 affected: analysed,
2068 modified_catalog: true,
2069 })
2070 }
2071
2072 fn exec_compact_cold_segments(&mut self) -> Result<QueryResult, EngineError> {
2083 let target = COMPACTION_TARGET_DEFAULT_BYTES;
2084 let reports = self.compact_cold_segments_with_target(target)?;
2085 let columns = alloc::vec![
2086 ColumnSchema::new("table_name", DataType::Text, false),
2087 ColumnSchema::new("index_name", DataType::Text, false),
2088 ColumnSchema::new("sources_merged", DataType::BigInt, false),
2089 ColumnSchema::new("merged_segment_id", DataType::BigInt, false),
2090 ColumnSchema::new("merged_rows", DataType::BigInt, false),
2091 ColumnSchema::new("deleted_rows_pruned", DataType::BigInt, false),
2092 ColumnSchema::new("bytes_reclaimed_estimate", DataType::BigInt, false),
2093 ];
2094 let rows: Vec<Row> = reports
2095 .into_iter()
2096 .map(|(tname, iname, report)| {
2097 Row::new(alloc::vec![
2098 Value::Text(tname),
2099 Value::Text(iname),
2100 Value::BigInt(i64::try_from(report.sources.len()).unwrap_or(i64::MAX)),
2101 Value::BigInt(i64::from(report.merged_segment_id.unwrap_or(0))),
2102 Value::BigInt(i64::try_from(report.merged_rows).unwrap_or(i64::MAX)),
2103 Value::BigInt(
2104 i64::try_from(report.deleted_rows_pruned).unwrap_or(i64::MAX),
2105 ),
2106 Value::BigInt(
2107 i64::try_from(report.bytes_reclaimed_estimate).unwrap_or(i64::MAX),
2108 ),
2109 ])
2110 })
2111 .collect();
2112 Ok(QueryResult::Rows { columns, rows })
2113 }
2114
2115 fn analyze_one_table(&mut self, table_name: &str) -> Result<(), EngineError> {
2120 let table = self.catalog.get(table_name).ok_or_else(|| {
2121 EngineError::Storage(StorageError::TableNotFound {
2122 name: table_name.to_string(),
2123 })
2124 })?;
2125 let schema = table.schema().clone();
2126 let row_count = table.rows().len();
2127 self.statistics.clear_table(table_name);
2132 for (col_pos, col_schema) in schema.columns.iter().enumerate() {
2133 if matches!(col_schema.ty, DataType::Vector { .. }) {
2136 continue;
2137 }
2138 let mut non_null_values: Vec<Value> = Vec::with_capacity(row_count);
2139 let mut nulls: u64 = 0;
2140 for row in table.rows() {
2141 match row.values.get(col_pos) {
2142 Some(Value::Null) | None => nulls += 1,
2143 Some(v) => non_null_values.push(v.clone()),
2144 }
2145 }
2146 non_null_values.sort_by(|a, b| sort_values_for_histogram(a, b));
2151 let non_null: Vec<String> = non_null_values
2152 .iter()
2153 .map(canonical_value_repr)
2154 .collect();
2155 let null_frac = if row_count == 0 {
2156 0.0
2157 } else {
2158 #[allow(clippy::cast_precision_loss)]
2159 let f = nulls as f32 / row_count as f32;
2160 f
2161 };
2162 let n_distinct = statistics::estimate_n_distinct(&non_null);
2163 let histogram_bounds = statistics::build_histogram(&non_null);
2164 self.statistics.set(
2165 table_name.to_string(),
2166 col_schema.name.clone(),
2167 statistics::ColumnStats {
2168 null_frac,
2169 n_distinct,
2170 histogram_bounds,
2171 },
2172 );
2173 }
2174 self.statistics.reset_modified(table_name);
2175 let cold_count = {
2181 let table = self
2182 .active_catalog()
2183 .get(table_name)
2184 .expect("table still present");
2185 table.count_cold_locators()
2186 };
2187 let table_mut = self
2188 .active_catalog_mut()
2189 .get_mut(table_name)
2190 .expect("table still present");
2191 table_mut.set_cold_row_count(cold_count);
2192 Ok(())
2193 }
2194
2195 fn exec_show_publications(&self) -> QueryResult {
2207 let columns = alloc::vec![
2208 ColumnSchema::new("name", DataType::Text, false),
2209 ColumnSchema::new("scope", DataType::Text, false),
2210 ColumnSchema::new("table_count", DataType::Int, true),
2211 ];
2212 let rows: Vec<Row> = self
2213 .publications
2214 .iter()
2215 .map(|(name, scope)| {
2216 let (scope_str, count_val) = match scope {
2217 spg_sql::ast::PublicationScope::AllTables => {
2218 ("FOR ALL TABLES".to_string(), Value::Null)
2219 }
2220 spg_sql::ast::PublicationScope::ForTables(ts) => (
2221 alloc::format!("FOR TABLE {}", ts.join(", ")),
2222 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
2223 ),
2224 spg_sql::ast::PublicationScope::AllTablesExcept(ts) => (
2225 alloc::format!("FOR ALL TABLES EXCEPT {}", ts.join(", ")),
2226 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
2227 ),
2228 };
2229 Row::new(alloc::vec![
2230 Value::Text(name.clone()),
2231 Value::Text(scope_str),
2232 count_val,
2233 ])
2234 })
2235 .collect();
2236 QueryResult::Rows { columns, rows }
2237 }
2238
2239 fn exec_show_users(&self) -> QueryResult {
2241 let columns = alloc::vec![
2242 ColumnSchema::new("name", DataType::Text, false),
2243 ColumnSchema::new("role", DataType::Text, false),
2244 ];
2245 let rows: Vec<Row> = self
2246 .users
2247 .iter()
2248 .map(|(name, rec)| {
2249 Row::new(alloc::vec![
2250 Value::Text(name.to_string()),
2251 Value::Text(rec.role.as_str().to_string()),
2252 ])
2253 })
2254 .collect();
2255 QueryResult::Rows { columns, rows }
2256 }
2257
2258 fn exec_create_user(&mut self, s: &CreateUserStatement) -> Result<QueryResult, EngineError> {
2259 if self.in_transaction() {
2260 return Err(EngineError::Unsupported(
2261 "CREATE USER is not allowed inside a transaction".into(),
2262 ));
2263 }
2264 let role = users::Role::parse(&s.role).ok_or_else(|| {
2265 EngineError::Unsupported(alloc::format!("invalid role: {:?}", s.role))
2266 })?;
2267 let salt = self.salt_fn.map_or_else(
2271 || {
2272 let mut s_bytes = [0u8; 16];
2273 let digest = spg_crypto::hash(s.name.as_bytes());
2274 s_bytes.copy_from_slice(&digest[..16]);
2275 s_bytes
2276 },
2277 |f| f(),
2278 );
2279 self.users
2280 .create(&s.name, &s.password, role, salt)
2281 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE USER: {e}")))?;
2282 Ok(QueryResult::CommandOk {
2283 affected: 1,
2284 modified_catalog: true,
2285 })
2286 }
2287
2288 fn exec_drop_user(&mut self, name: &str) -> Result<QueryResult, EngineError> {
2289 if self.in_transaction() {
2290 return Err(EngineError::Unsupported(
2291 "DROP USER is not allowed inside a transaction".into(),
2292 ));
2293 }
2294 self.users
2295 .drop(name)
2296 .map_err(|e| EngineError::Unsupported(alloc::format!("DROP USER: {e}")))?;
2297 Ok(QueryResult::CommandOk {
2298 affected: 1,
2299 modified_catalog: true,
2300 })
2301 }
2302
2303 fn exec_update_cancel(
2310 &mut self,
2311 stmt: &spg_sql::ast::UpdateStatement,
2312 cancel: CancelToken<'_>,
2313 ) -> Result<QueryResult, EngineError> {
2314 if let Some(w) = &stmt.where_ {
2322 let schema_cols = self
2323 .active_catalog()
2324 .get(&stmt.table)
2325 .ok_or_else(|| {
2326 EngineError::Storage(StorageError::TableNotFound {
2327 name: stmt.table.clone(),
2328 })
2329 })?
2330 .schema()
2331 .columns
2332 .clone();
2333 if let Some((col_pos, key)) = try_pk_predicate(w, &schema_cols, stmt.table.as_str())
2334 && let Some(idx_name) = self
2335 .active_catalog()
2336 .get(&stmt.table)
2337 .and_then(|t| t.index_on(col_pos).map(|i| i.name.clone()))
2338 {
2339 let _ = self
2343 .active_catalog_mut()
2344 .promote_cold_row(&stmt.table, &idx_name, &key);
2345 }
2346 }
2347
2348 let table = self
2349 .active_catalog_mut()
2350 .get_mut(&stmt.table)
2351 .ok_or_else(|| {
2352 EngineError::Storage(StorageError::TableNotFound {
2353 name: stmt.table.clone(),
2354 })
2355 })?;
2356 let schema_cols: Vec<ColumnSchema> = table.schema().columns.clone();
2357 let mut targets: Vec<(usize, &Expr)> = Vec::with_capacity(stmt.assignments.len());
2361 for (col, expr) in &stmt.assignments {
2362 let pos = schema_cols
2363 .iter()
2364 .position(|c| c.name == *col)
2365 .ok_or_else(|| {
2366 EngineError::Eval(EvalError::ColumnNotFound { name: col.clone() })
2367 })?;
2368 targets.push((pos, expr));
2369 }
2370 let ctx = EvalContext::new(&schema_cols, Some(stmt.table.as_str()));
2371 let mut planned: Vec<(usize, Vec<Value>)> = Vec::new();
2377 for (i, row) in table.rows().iter().enumerate() {
2378 if i.is_multiple_of(256) {
2382 cancel.check()?;
2383 }
2384 if let Some(w) = &stmt.where_ {
2385 let cond = eval::eval_expr(w, row, &ctx)?;
2386 if !matches!(cond, Value::Bool(true)) {
2387 continue;
2388 }
2389 }
2390 let mut new_vals = row.values.clone();
2391 for (pos, expr) in &targets {
2392 let v = eval::eval_expr(expr, row, &ctx)?;
2393 new_vals[*pos] =
2394 coerce_value(v, schema_cols[*pos].ty, &schema_cols[*pos].name, *pos)?;
2395 }
2396 planned.push((i, new_vals));
2397 }
2398 let plan_with_old: Vec<(usize, Vec<Value>, Vec<Value>)> = planned
2402 .iter()
2403 .map(|(pos, new_vals)| (*pos, table.rows()[*pos].values.clone(), new_vals.clone()))
2404 .collect();
2405 let self_fks = table.schema().foreign_keys.clone();
2406 let affected = planned.len();
2407 let _ = table;
2409 if !self_fks.is_empty() {
2413 let new_rows: Vec<Vec<Value>> = planned
2414 .iter()
2415 .map(|(_pos, new_vals)| new_vals.clone())
2416 .collect();
2417 enforce_fk_inserts(self.active_catalog(), &stmt.table, &self_fks, &new_rows)?;
2418 }
2419 let child_plan = plan_fk_parent_updates(self.active_catalog(), &stmt.table, &plan_with_old)?;
2423 for step in &child_plan {
2425 apply_fk_child_step(self.active_catalog_mut(), step)?;
2426 }
2427 let table = self
2429 .active_catalog_mut()
2430 .get_mut(&stmt.table)
2431 .ok_or_else(|| {
2432 EngineError::Storage(StorageError::TableNotFound {
2433 name: stmt.table.clone(),
2434 })
2435 })?;
2436 let updated_for_returning: Vec<Vec<Value>> =
2438 if stmt.returning.is_some() {
2439 planned.iter().map(|(_pos, vals)| vals.clone()).collect()
2440 } else {
2441 Vec::new()
2442 };
2443 for (pos, vals) in planned {
2444 table.update_row(pos, vals)?;
2445 }
2446 let _ = table;
2447 if !self.in_transaction() && affected > 0 {
2449 self.statistics
2450 .record_modifications(&stmt.table, affected as u64);
2451 }
2452 if let Some(items) = &stmt.returning {
2454 return self.build_returning_rows(
2455 &stmt.table,
2456 items,
2457 updated_for_returning,
2458 );
2459 }
2460 Ok(QueryResult::CommandOk {
2461 affected,
2462 modified_catalog: !self.in_transaction(),
2463 })
2464 }
2465
2466 fn exec_delete_cancel(
2470 &mut self,
2471 stmt: &spg_sql::ast::DeleteStatement,
2472 cancel: CancelToken<'_>,
2473 ) -> Result<QueryResult, EngineError> {
2474 let mut cold_shadow_count: usize = 0;
2482 if let Some(w) = &stmt.where_ {
2483 let schema_cols = self
2484 .active_catalog()
2485 .get(&stmt.table)
2486 .ok_or_else(|| {
2487 EngineError::Storage(StorageError::TableNotFound {
2488 name: stmt.table.clone(),
2489 })
2490 })?
2491 .schema()
2492 .columns
2493 .clone();
2494 if let Some((col_pos, key)) = try_pk_predicate(w, &schema_cols, stmt.table.as_str())
2495 && let Some(idx_name) = self
2496 .active_catalog()
2497 .get(&stmt.table)
2498 .and_then(|t| t.index_on(col_pos).map(|i| i.name.clone()))
2499 {
2500 cold_shadow_count = self
2501 .active_catalog_mut()
2502 .shadow_cold_row(&stmt.table, &idx_name, &key)
2503 .unwrap_or(0);
2504 }
2505 }
2506
2507 let table = self
2508 .active_catalog_mut()
2509 .get_mut(&stmt.table)
2510 .ok_or_else(|| {
2511 EngineError::Storage(StorageError::TableNotFound {
2512 name: stmt.table.clone(),
2513 })
2514 })?;
2515 let schema_cols: Vec<ColumnSchema> = table.schema().columns.clone();
2516 let ctx = EvalContext::new(&schema_cols, Some(stmt.table.as_str()));
2517 let mut positions: Vec<usize> = Vec::new();
2518 let mut to_delete_rows: Vec<Vec<Value>> = Vec::new();
2522 for (i, row) in table.rows().iter().enumerate() {
2523 if i.is_multiple_of(256) {
2524 cancel.check()?;
2525 }
2526 let keep = if let Some(w) = &stmt.where_ {
2527 let cond = eval::eval_expr(w, row, &ctx)?;
2528 !matches!(cond, Value::Bool(true))
2529 } else {
2530 false
2531 };
2532 if !keep {
2533 positions.push(i);
2534 to_delete_rows.push(row.values.clone());
2535 }
2536 }
2537 let _ = table;
2544 let cascade_plan = plan_fk_parent_deletions(
2545 self.active_catalog(),
2546 &stmt.table,
2547 &positions,
2548 &to_delete_rows,
2549 )?;
2550 for step in &cascade_plan {
2557 apply_fk_child_step(self.active_catalog_mut(), step)?;
2558 }
2559 let table = self
2561 .active_catalog_mut()
2562 .get_mut(&stmt.table)
2563 .ok_or_else(|| {
2564 EngineError::Storage(StorageError::TableNotFound {
2565 name: stmt.table.clone(),
2566 })
2567 })?;
2568 let affected = table.delete_rows(&positions) + cold_shadow_count;
2569 let _ = table;
2570 if !self.in_transaction() && affected > 0 {
2572 self.statistics
2573 .record_modifications(&stmt.table, affected as u64);
2574 }
2575 if let Some(items) = &stmt.returning {
2581 return self.build_returning_rows(
2582 &stmt.table,
2583 items,
2584 to_delete_rows,
2585 );
2586 }
2587 Ok(QueryResult::CommandOk {
2588 affected,
2589 modified_catalog: !self.in_transaction(),
2590 })
2591 }
2592
2593 #[allow(clippy::format_push_string)]
2603 fn exec_explain(
2604 &self,
2605 e: &spg_sql::ast::ExplainStatement,
2606 cancel: CancelToken<'_>,
2607 ) -> Result<QueryResult, EngineError> {
2608 let mut lines = Vec::<String>::new();
2609 explain_select(&e.inner, self, 0, &mut lines);
2610 if e.suggest {
2611 let suggestions = build_index_suggestions(&e.inner, self);
2620 for s in suggestions {
2621 lines.push(s);
2622 }
2623 } else if e.analyze {
2624 let started = self.clock.map(|f| f());
2641 let exec = self.exec_select_cancel(&e.inner, cancel)?;
2642 let elapsed_micros = match (self.clock, started) {
2643 (Some(f), Some(s)) => Some(f().saturating_sub(s)),
2644 _ => None,
2645 };
2646 let row_count = if let QueryResult::Rows { rows, .. } = &exec {
2647 rows.len()
2648 } else {
2649 0
2650 };
2651 annotate_explain_lines(&mut lines, row_count, self);
2652 let mut total = alloc::format!("Total: rows={row_count}");
2653 if let Some(us) = elapsed_micros {
2654 total.push_str(&alloc::format!(" elapsed={us}us"));
2655 }
2656 lines.push(total);
2657 }
2658 let columns = alloc::vec![ColumnSchema::new("QUERY PLAN", DataType::Text, false)];
2659 let rows: Vec<Row> = lines
2660 .into_iter()
2661 .map(|l| Row::new(alloc::vec![Value::Text(l)]))
2662 .collect();
2663 Ok(QueryResult::Rows { columns, rows })
2664 }
2665
2666 fn exec_show_tables(&self) -> QueryResult {
2667 let columns = alloc::vec![ColumnSchema::new("name", DataType::Text, false)];
2668 let rows: Vec<Row> = self
2669 .active_catalog()
2670 .table_names()
2671 .into_iter()
2672 .map(|n| Row::new(alloc::vec![Value::Text(n)]))
2673 .collect();
2674 QueryResult::Rows { columns, rows }
2675 }
2676
2677 fn exec_show_columns(&self, table_name: &str) -> Result<QueryResult, EngineError> {
2680 let table =
2681 self.active_catalog()
2682 .get(table_name)
2683 .ok_or_else(|| StorageError::TableNotFound {
2684 name: table_name.into(),
2685 })?;
2686 let columns = alloc::vec![
2687 ColumnSchema::new("name", DataType::Text, false),
2688 ColumnSchema::new("type", DataType::Text, false),
2689 ColumnSchema::new("nullable", DataType::Bool, false),
2690 ];
2691 let rows: Vec<Row> = table
2692 .schema()
2693 .columns
2694 .iter()
2695 .map(|c| {
2696 Row::new(alloc::vec![
2697 Value::Text(c.name.clone()),
2698 Value::Text(alloc::format!("{}", c.ty)),
2699 Value::Bool(c.nullable),
2700 ])
2701 })
2702 .collect();
2703 Ok(QueryResult::Rows { columns, rows })
2704 }
2705
2706 fn exec_begin(&mut self) -> Result<QueryResult, EngineError> {
2707 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2708 if self.tx_catalogs.contains_key(&tx_id) {
2709 return Err(EngineError::TransactionAlreadyOpen);
2710 }
2711 self.tx_catalogs.insert(
2712 tx_id,
2713 TxState {
2714 catalog: self.catalog.clone(),
2715 savepoints: Vec::new(),
2716 },
2717 );
2718 Ok(QueryResult::CommandOk {
2719 affected: 0,
2720 modified_catalog: false,
2721 })
2722 }
2723
2724 fn exec_commit(&mut self) -> Result<QueryResult, EngineError> {
2725 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2726 let state = self
2727 .tx_catalogs
2728 .remove(&tx_id)
2729 .ok_or(EngineError::NoActiveTransaction)?;
2730 self.catalog = state.catalog;
2731 Ok(QueryResult::CommandOk {
2735 affected: 0,
2736 modified_catalog: true,
2737 })
2738 }
2739
2740 fn exec_rollback(&mut self) -> Result<QueryResult, EngineError> {
2741 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2742 if self.tx_catalogs.remove(&tx_id).is_none() {
2743 return Err(EngineError::NoActiveTransaction);
2744 }
2745 Ok(QueryResult::CommandOk {
2747 affected: 0,
2748 modified_catalog: false,
2749 })
2750 }
2751
2752 fn exec_savepoint(&mut self, name: String) -> Result<QueryResult, EngineError> {
2753 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2754 let state = self
2755 .tx_catalogs
2756 .get_mut(&tx_id)
2757 .ok_or(EngineError::NoActiveTransaction)?;
2758 state.savepoints.retain(|(n, _)| n != &name);
2762 let snapshot = state.catalog.clone();
2763 state.savepoints.push((name, snapshot));
2764 Ok(QueryResult::CommandOk {
2765 affected: 0,
2766 modified_catalog: false,
2767 })
2768 }
2769
2770 fn exec_rollback_to_savepoint(&mut self, name: &str) -> Result<QueryResult, EngineError> {
2771 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2772 let state = self
2773 .tx_catalogs
2774 .get_mut(&tx_id)
2775 .ok_or(EngineError::NoActiveTransaction)?;
2776 let pos = state
2777 .savepoints
2778 .iter()
2779 .rposition(|(n, _)| n == name)
2780 .ok_or_else(|| {
2781 EngineError::Unsupported(alloc::format!("savepoint not found: {name}"))
2782 })?;
2783 let snapshot = state.savepoints[pos].1.clone();
2787 state.savepoints.truncate(pos + 1);
2788 state.catalog = snapshot;
2789 Ok(QueryResult::CommandOk {
2790 affected: 0,
2791 modified_catalog: false,
2792 })
2793 }
2794
2795 fn exec_release_savepoint(&mut self, name: &str) -> Result<QueryResult, EngineError> {
2796 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
2797 let state = self
2798 .tx_catalogs
2799 .get_mut(&tx_id)
2800 .ok_or(EngineError::NoActiveTransaction)?;
2801 let pos = state
2802 .savepoints
2803 .iter()
2804 .rposition(|(n, _)| n == name)
2805 .ok_or_else(|| {
2806 EngineError::Unsupported(alloc::format!("savepoint not found: {name}"))
2807 })?;
2808 state.savepoints.truncate(pos);
2811 Ok(QueryResult::CommandOk {
2812 affected: 0,
2813 modified_catalog: false,
2814 })
2815 }
2816
2817 fn exec_alter_table(
2828 &mut self,
2829 s: spg_sql::ast::AlterTableStatement,
2830 ) -> Result<QueryResult, EngineError> {
2831 match s.target {
2832 spg_sql::ast::AlterTableTarget::SetHotTierBytes(n) => {
2833 let table = self
2834 .active_catalog_mut()
2835 .get_mut(&s.name)
2836 .ok_or_else(|| {
2837 EngineError::Storage(StorageError::TableNotFound {
2838 name: s.name.clone(),
2839 })
2840 })?;
2841 table.schema_mut().hot_tier_bytes = Some(n);
2842 }
2843 spg_sql::ast::AlterTableTarget::AddForeignKey(fk) => {
2844 let cols_snapshot = self
2849 .active_catalog()
2850 .get(&s.name)
2851 .ok_or_else(|| {
2852 EngineError::Storage(StorageError::TableNotFound {
2853 name: s.name.clone(),
2854 })
2855 })?
2856 .schema()
2857 .columns
2858 .clone();
2859 let storage_fk = resolve_foreign_key(
2860 &s.name,
2861 &cols_snapshot,
2862 fk,
2863 self.active_catalog(),
2864 )?;
2865 let existing_rows: Vec<Vec<Value>> = self
2868 .active_catalog()
2869 .get(&s.name)
2870 .expect("checked above")
2871 .rows()
2872 .iter()
2873 .map(|r| r.values.clone())
2874 .collect();
2875 enforce_fk_inserts(
2876 self.active_catalog(),
2877 &s.name,
2878 core::slice::from_ref(&storage_fk),
2879 &existing_rows,
2880 )?;
2881 let table = self
2883 .active_catalog_mut()
2884 .get_mut(&s.name)
2885 .expect("checked above");
2886 if let Some(name) = &storage_fk.name
2887 && table
2888 .schema()
2889 .foreign_keys
2890 .iter()
2891 .any(|f| f.name.as_ref() == Some(name))
2892 {
2893 return Err(EngineError::Unsupported(alloc::format!(
2894 "ALTER TABLE ADD CONSTRAINT: a constraint named {name:?} already exists"
2895 )));
2896 }
2897 table.schema_mut().foreign_keys.push(storage_fk);
2898 }
2899 spg_sql::ast::AlterTableTarget::DropForeignKey(name) => {
2900 let table = self
2901 .active_catalog_mut()
2902 .get_mut(&s.name)
2903 .ok_or_else(|| {
2904 EngineError::Storage(StorageError::TableNotFound {
2905 name: s.name.clone(),
2906 })
2907 })?;
2908 let fks = &mut table.schema_mut().foreign_keys;
2909 let before = fks.len();
2910 fks.retain(|f| f.name.as_ref() != Some(&name));
2911 if fks.len() == before {
2912 return Err(EngineError::Unsupported(alloc::format!(
2913 "ALTER TABLE DROP CONSTRAINT: no FK named {name:?} on {:?}",
2914 s.name
2915 )));
2916 }
2917 }
2918 }
2919 Ok(QueryResult::CommandOk {
2920 affected: 0,
2921 modified_catalog: !self.in_transaction(),
2922 })
2923 }
2924
2925 fn exec_alter_index(
2926 &mut self,
2927 stmt: spg_sql::ast::AlterIndexStatement,
2928 ) -> Result<QueryResult, EngineError> {
2929 let spg_sql::ast::AlterIndexStatement {
2933 name: idx_name,
2934 target,
2935 } = stmt;
2936 let spg_sql::ast::AlterIndexTarget::Rebuild { encoding } = target;
2937 let target = encoding.map(|e| match e {
2938 SqlVecEncoding::F32 => VecEncoding::F32,
2939 SqlVecEncoding::Sq8 => VecEncoding::Sq8,
2940 SqlVecEncoding::F16 => VecEncoding::F16,
2941 });
2942 let table_name = {
2947 let cat = self.active_catalog();
2948 let mut found: Option<String> = None;
2949 for tname in cat.table_names() {
2950 if let Some(t) = cat.get(&tname)
2951 && t.indices().iter().any(|i| i.name == idx_name)
2952 {
2953 found = Some(tname);
2954 break;
2955 }
2956 }
2957 found.ok_or_else(|| {
2958 EngineError::Storage(StorageError::IndexNotFound {
2959 name: idx_name.clone(),
2960 })
2961 })?
2962 };
2963 let table = self
2964 .active_catalog_mut()
2965 .get_mut(&table_name)
2966 .expect("table found above");
2967 table.rebuild_nsw_index(&idx_name, target)?;
2968 self.plan_cache.evict_referencing(&table_name);
2971 Ok(QueryResult::CommandOk {
2972 affected: 0,
2973 modified_catalog: !self.in_transaction(),
2974 })
2975 }
2976
2977 fn exec_create_index(
2978 &mut self,
2979 stmt: CreateIndexStatement,
2980 ) -> Result<QueryResult, EngineError> {
2981 let table = self
2982 .active_catalog_mut()
2983 .get_mut(&stmt.table)
2984 .ok_or_else(|| {
2985 EngineError::Storage(StorageError::TableNotFound {
2986 name: stmt.table.clone(),
2987 })
2988 })?;
2989 if stmt.if_not_exists && table.indices().iter().any(|i| i.name == stmt.name) {
2991 return Ok(QueryResult::CommandOk {
2992 affected: 0,
2993 modified_catalog: false,
2994 });
2995 }
2996 let _ = &stmt.extra_columns; let table_name = stmt.table.clone();
3003 let included_positions: Vec<usize> = if stmt.included_columns.is_empty() {
3007 Vec::new()
3008 } else {
3009 let schema = table.schema();
3010 stmt.included_columns
3011 .iter()
3012 .map(|c| {
3013 schema.column_position(c).ok_or_else(|| {
3014 EngineError::Storage(StorageError::ColumnNotFound {
3015 column: c.clone(),
3016 })
3017 })
3018 })
3019 .collect::<Result<Vec<_>, _>>()?
3020 };
3021 match stmt.method {
3022 IndexMethod::BTree => table.add_index(stmt.name.clone(), &stmt.column)?,
3023 IndexMethod::Hnsw => {
3024 if !included_positions.is_empty() {
3025 return Err(EngineError::Unsupported(
3026 "INCLUDE columns are not supported on HNSW indexes".into(),
3027 ));
3028 }
3029 table.add_nsw_index(stmt.name.clone(), &stmt.column, spg_storage::NSW_DEFAULT_M)?;
3030 }
3031 IndexMethod::Brin => {
3033 if !included_positions.is_empty() {
3034 return Err(EngineError::Unsupported(
3035 "INCLUDE columns are not supported on BRIN indexes".into(),
3036 ));
3037 }
3038 table.add_brin_index(stmt.name.clone(), &stmt.column)?;
3039 }
3040 }
3041 if !included_positions.is_empty()
3042 && let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name)
3043 {
3044 idx.included_columns = included_positions;
3045 }
3046 if let Some(pred_expr) = &stmt.partial_predicate {
3054 let canonical = pred_expr.to_string();
3055 if matches!(stmt.method, IndexMethod::Hnsw | IndexMethod::Brin) {
3056 return Err(EngineError::Unsupported(
3057 "WHERE predicates are not supported on HNSW or BRIN indexes".into(),
3058 ));
3059 }
3060 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3061 idx.partial_predicate = Some(canonical);
3062 }
3063 }
3064 if let Some(key_expr) = &stmt.expression {
3072 if matches!(stmt.method, IndexMethod::Hnsw | IndexMethod::Brin) {
3073 return Err(EngineError::Unsupported(
3074 "Expression keys are not supported on HNSW or BRIN indexes".into(),
3075 ));
3076 }
3077 let canonical = key_expr.to_string();
3078 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3079 idx.expression = Some(canonical);
3080 }
3081 }
3082 if stmt.is_unique {
3091 let mut extra_positions: alloc::vec::Vec<usize> = alloc::vec::Vec::new();
3092 for col_name in &stmt.extra_columns {
3093 let pos = table
3094 .schema()
3095 .columns
3096 .iter()
3097 .position(|c| c.name.eq_ignore_ascii_case(col_name))
3098 .ok_or_else(|| {
3099 EngineError::Unsupported(alloc::format!(
3100 "UNIQUE INDEX {:?}: extra column {col_name:?} not in table {:?}",
3101 stmt.name, stmt.table
3102 ))
3103 })?;
3104 extra_positions.push(pos);
3105 }
3106 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3107 idx.is_unique = true;
3108 idx.extra_column_positions = extra_positions;
3109 }
3110 let snapshot_indices = table.indices().to_vec();
3115 let snapshot_rows: alloc::vec::Vec<spg_storage::Row> =
3116 table.rows().iter().cloned().collect();
3117 let snapshot_schema = table.schema().clone();
3118 let idx_ref = snapshot_indices
3119 .iter()
3120 .find(|i| i.name == stmt.name)
3121 .expect("just-added index");
3122 check_existing_unique_violation(idx_ref, &snapshot_schema, &snapshot_rows)?;
3123 }
3124 self.plan_cache.evict_referencing(&table_name);
3127 Ok(QueryResult::CommandOk {
3128 affected: 0,
3129 modified_catalog: !self.in_transaction(),
3130 })
3131 }
3132
3133 fn exec_create_table(
3134 &mut self,
3135 stmt: CreateTableStatement,
3136 ) -> Result<QueryResult, EngineError> {
3137 if stmt.if_not_exists && self.active_catalog().get(&stmt.name).is_some() {
3138 return Ok(QueryResult::CommandOk {
3139 affected: 0,
3140 modified_catalog: false,
3141 });
3142 }
3143 let table_name = stmt.name.clone();
3144 let inline_pk_columns: Vec<String> = stmt
3148 .columns
3149 .iter()
3150 .filter(|c| c.is_primary_key)
3151 .map(|c| c.name.clone())
3152 .collect();
3153 let cols = stmt
3159 .columns
3160 .into_iter()
3161 .map(column_def_to_schema)
3162 .collect::<Result<Vec<_>, _>>()?;
3163 let mut cols = cols;
3165 for tc in &stmt.table_constraints {
3166 if let spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } = tc {
3167 for col_name in columns {
3168 if let Some(col) = cols.iter_mut().find(|c| c.name == *col_name) {
3169 col.nullable = false;
3170 }
3171 }
3172 }
3173 }
3174 let mut fks: Vec<spg_storage::ForeignKeyConstraint> =
3181 Vec::with_capacity(stmt.foreign_keys.len());
3182 for fk in stmt.foreign_keys {
3183 fks.push(resolve_foreign_key(
3184 &table_name,
3185 &cols,
3186 fk,
3187 self.active_catalog(),
3188 )?);
3189 }
3190 let mut schema = TableSchema::new(table_name.clone(), cols);
3191 schema.foreign_keys = fks;
3192 let mut uc_storage: Vec<spg_storage::UniquenessConstraint> = Vec::new();
3196 for tc in &stmt.table_constraints {
3197 let (is_pk, names) = match tc {
3198 spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } => {
3199 (true, columns.clone())
3200 }
3201 spg_sql::ast::TableConstraint::Unique { columns, .. } => {
3202 (false, columns.clone())
3203 }
3204 };
3205 let mut positions = Vec::with_capacity(names.len());
3206 for n in &names {
3207 let pos = schema
3208 .columns
3209 .iter()
3210 .position(|c| c.name == *n)
3211 .ok_or_else(|| {
3212 EngineError::Unsupported(alloc::format!(
3213 "table constraint references unknown column {n:?}"
3214 ))
3215 })?;
3216 positions.push(pos);
3217 }
3218 uc_storage.push(spg_storage::UniquenessConstraint {
3219 is_primary_key: is_pk,
3220 columns: positions,
3221 });
3222 }
3223 schema.uniqueness_constraints = uc_storage.clone();
3224 self.active_catalog_mut().create_table(schema)?;
3225 let table = self
3229 .active_catalog_mut()
3230 .get_mut(&table_name)
3231 .expect("just created");
3232 for (i, col_name) in inline_pk_columns.iter().enumerate() {
3233 let idx_name = if inline_pk_columns.len() == 1 {
3234 alloc::format!("{table_name}_pkey")
3235 } else {
3236 alloc::format!("{table_name}_pkey_{i}")
3237 };
3238 if let Err(e) = table.add_index(idx_name, col_name) {
3239 return Err(EngineError::Storage(e));
3240 }
3241 }
3242 for (i, tc) in stmt.table_constraints.iter().enumerate() {
3243 let (is_pk, names) = match tc {
3244 spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } => {
3245 (true, columns)
3246 }
3247 spg_sql::ast::TableConstraint::Unique { columns, .. } => {
3248 (false, columns)
3249 }
3250 };
3251 let leading = &names[0];
3252 let already = table
3255 .indices()
3256 .iter()
3257 .any(|idx| {
3258 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
3259 && table.schema().columns[idx.column_position].name == *leading
3260 });
3261 if already {
3262 continue;
3263 }
3264 let suffix = if is_pk { "pkey" } else { "key" };
3265 let idx_name = if names.len() == 1 {
3266 alloc::format!("{table_name}_{leading}_{suffix}")
3267 } else {
3268 alloc::format!("{table_name}_{leading}_{suffix}_{i}")
3269 };
3270 if let Err(e) = table.add_index(idx_name, leading) {
3271 return Err(EngineError::Storage(e));
3272 }
3273 }
3274 Ok(QueryResult::CommandOk {
3275 affected: 0,
3276 modified_catalog: !self.in_transaction(),
3277 })
3278 }
3279
3280 fn exec_insert(&mut self, stmt: InsertStatement) -> Result<QueryResult, EngineError> {
3281 let clock = self.clock;
3285 let table = self
3286 .active_catalog_mut()
3287 .get_mut(&stmt.table)
3288 .ok_or_else(|| {
3289 EngineError::Storage(StorageError::TableNotFound {
3290 name: stmt.table.clone(),
3291 })
3292 })?;
3293 let column_meta: Vec<ColumnSchema> = table.schema().columns.clone();
3299 let schema_cols_len = column_meta.len();
3300 let tuple_pos: Option<Vec<Option<usize>>> = match &stmt.columns {
3304 None => None, Some(cols) => {
3306 let mut map = alloc::vec![None; schema_cols_len];
3307 for (j, name) in cols.iter().enumerate() {
3308 let idx = column_meta
3309 .iter()
3310 .position(|c| c.name == *name)
3311 .ok_or_else(|| {
3312 EngineError::Eval(EvalError::ColumnNotFound { name: name.clone() })
3313 })?;
3314 if map[idx].is_some() {
3315 return Err(EngineError::Storage(StorageError::ArityMismatch {
3316 expected: schema_cols_len,
3317 actual: cols.len(),
3318 }));
3319 }
3320 map[idx] = Some(j);
3321 }
3322 for (i, col) in column_meta.iter().enumerate() {
3326 if map[i].is_none()
3327 && !col.nullable
3328 && col.default.is_none()
3329 && col.runtime_default.is_none()
3330 && !col.auto_increment
3331 {
3332 return Err(EngineError::Storage(StorageError::NullInNotNull {
3333 column: col.name.clone(),
3334 }));
3335 }
3336 }
3337 Some(map)
3338 }
3339 };
3340 let expected_tuple_len = stmt.columns.as_ref().map_or(schema_cols_len, Vec::len);
3341 let fks = table.schema().foreign_keys.clone();
3347 let mut affected = 0usize;
3348 let mut all_values: Vec<Vec<Value>> = Vec::with_capacity(stmt.rows.len());
3351 for tuple in stmt.rows {
3352 if tuple.len() != expected_tuple_len {
3353 return Err(EngineError::Storage(StorageError::ArityMismatch {
3354 expected: expected_tuple_len,
3355 actual: tuple.len(),
3356 }));
3357 }
3358 let values: Vec<Value> = if let Some(map) = &tuple_pos {
3362 let raw_tuple: Vec<Value> = tuple
3364 .into_iter()
3365 .map(literal_expr_to_value)
3366 .collect::<Result<_, _>>()?;
3367 let mut out = Vec::with_capacity(schema_cols_len);
3368 for (i, col) in column_meta.iter().enumerate() {
3369 let mut raw = match map[i] {
3370 Some(j) => raw_tuple[j].clone(),
3371 None => resolve_column_default_free(col, clock)?,
3372 };
3373 if col.auto_increment && raw.is_null() {
3374 let next = table.next_auto_value(i).ok_or_else(|| {
3375 EngineError::Unsupported(alloc::format!(
3376 "AUTO_INCREMENT applies to integer columns only (column `{}`)",
3377 col.name
3378 ))
3379 })?;
3380 raw = Value::BigInt(next);
3381 }
3382 out.push(coerce_value(raw, col.ty, &col.name, i)?);
3383 }
3384 out
3385 } else {
3386 let mut out = Vec::with_capacity(schema_cols_len);
3388 for (i, (col, expr)) in column_meta.iter().zip(tuple).enumerate() {
3389 let mut raw = literal_expr_to_value(expr)?;
3390 if col.auto_increment && raw.is_null() {
3391 let next = table.next_auto_value(i).ok_or_else(|| {
3392 EngineError::Unsupported(alloc::format!(
3393 "AUTO_INCREMENT applies to integer columns only (column `{}`)",
3394 col.name
3395 ))
3396 })?;
3397 raw = Value::BigInt(next);
3398 }
3399 out.push(coerce_value(raw, col.ty, &col.name, i)?);
3400 }
3401 out
3402 };
3403 all_values.push(values);
3404 }
3405 let uniqueness = table.schema().uniqueness_constraints.clone();
3410 let _ = table;
3411 if !fks.is_empty() {
3412 enforce_fk_inserts(self.active_catalog(), &stmt.table, &fks, &all_values)?;
3413 }
3414 enforce_uniqueness_inserts(
3416 self.active_catalog(),
3417 &stmt.table,
3418 &uniqueness,
3419 &all_values,
3420 )?;
3421 enforce_unique_index_inserts(
3428 self.active_catalog(),
3429 &stmt.table,
3430 &all_values,
3431 )?;
3432 let mut pending_updates: Vec<(usize, Vec<Value>)> = Vec::new();
3439 let mut skipped_count = 0usize;
3440 if let Some(clause) = &stmt.on_conflict {
3441 let conflict_cols = resolve_on_conflict_columns(
3442 self.active_catalog(),
3443 &stmt.table,
3444 clause.target_columns.as_slice(),
3445 )?;
3446 let mut kept: Vec<Vec<Value>> = Vec::with_capacity(all_values.len());
3447 let mut seen_keys: Vec<Vec<Value>> = Vec::new();
3448 for values in all_values {
3449 let key_tuple: Vec<&Value> =
3450 conflict_cols.iter().map(|&c| &values[c]).collect();
3451 let has_null_key = key_tuple.iter().any(|v| matches!(v, Value::Null));
3454 let collides_with_table = !has_null_key
3455 && on_conflict_keys_exist(
3456 self.active_catalog(),
3457 &stmt.table,
3458 &conflict_cols,
3459 &key_tuple,
3460 );
3461 let key_tuple_owned: Vec<Value> =
3462 key_tuple.iter().map(|v| (*v).clone()).collect();
3463 let collides_with_batch = !has_null_key
3464 && seen_keys.iter().any(|k| k == &key_tuple_owned);
3465 let collides = collides_with_table || collides_with_batch;
3466 match (&clause.action, collides) {
3467 (_, false) => {
3468 seen_keys.push(key_tuple_owned);
3469 kept.push(values);
3470 }
3471 (spg_sql::ast::OnConflictAction::Nothing, true) => {
3472 skipped_count += 1;
3473 }
3474 (
3475 spg_sql::ast::OnConflictAction::Update {
3476 assignments,
3477 where_,
3478 },
3479 true,
3480 ) => {
3481 if !collides_with_table {
3482 skipped_count += 1;
3483 continue;
3484 }
3485 let target_pos = lookup_row_position_by_keys(
3486 self.active_catalog(),
3487 &stmt.table,
3488 &conflict_cols,
3489 &key_tuple,
3490 )
3491 .ok_or_else(|| {
3492 EngineError::Unsupported(
3493 "ON CONFLICT DO UPDATE: conflict detected but row \
3494 position could not be resolved (cold-tier row?)"
3495 .into(),
3496 )
3497 })?;
3498 let updated = apply_on_conflict_assignments(
3499 self.active_catalog(),
3500 &stmt.table,
3501 target_pos,
3502 &values,
3503 assignments,
3504 where_.as_ref(),
3505 )?;
3506 if let Some(new_row) = updated {
3507 pending_updates.push((target_pos, new_row));
3508 } else {
3509 skipped_count += 1;
3510 }
3511 }
3512 }
3513 }
3514 all_values = kept;
3515 }
3516 let table = self
3518 .active_catalog_mut()
3519 .get_mut(&stmt.table)
3520 .ok_or_else(|| {
3521 EngineError::Storage(StorageError::TableNotFound {
3522 name: stmt.table.clone(),
3523 })
3524 })?;
3525 let mut returning_rows: Vec<Vec<Value>> = Vec::new();
3529 for values in all_values {
3530 if stmt.returning.is_some() {
3531 returning_rows.push(values.clone());
3532 }
3533 table.insert(Row::new(values))?;
3534 affected += 1;
3535 }
3536 for (pos, new_row) in pending_updates {
3540 if stmt.returning.is_some() {
3541 returning_rows.push(new_row.clone());
3542 }
3543 table.update_row(pos, new_row)?;
3544 affected += 1;
3545 }
3546 let _ = skipped_count;
3547 if let Some(items) = &stmt.returning {
3551 let _ = table;
3552 return self.build_returning_rows(
3553 &stmt.table,
3554 items,
3555 returning_rows,
3556 );
3557 }
3558 if !self.in_transaction() && affected > 0 {
3563 self.statistics
3564 .record_modifications(&stmt.table, affected as u64);
3565 }
3566 Ok(QueryResult::CommandOk {
3567 affected,
3568 modified_catalog: !self.in_transaction(),
3569 })
3570 }
3571
3572 fn exec_select_as_of_segment(
3585 &self,
3586 stmt: &SelectStatement,
3587 from: &spg_sql::ast::FromClause,
3588 segment_id: u32,
3589 ) -> Result<QueryResult, EngineError> {
3590 if !from.joins.is_empty()
3593 || stmt.group_by.is_some()
3594 || stmt.having.is_some()
3595 || !stmt.unions.is_empty()
3596 || !stmt.order_by.is_empty()
3597 || stmt.offset.is_some()
3598 || stmt.distinct
3599 || aggregate::uses_aggregate(stmt)
3600 {
3601 return Err(EngineError::Unsupported(
3602 "AS OF SEGMENT supports SELECT projection + WHERE + LIMIT only \
3603 (joins / aggregates / ORDER BY are STABILITY § \"Out of v6.10\")"
3604 .into(),
3605 ));
3606 }
3607 let table = self
3608 .active_catalog()
3609 .get(&from.primary.name)
3610 .ok_or_else(|| StorageError::TableNotFound {
3611 name: from.primary.name.clone(),
3612 })?;
3613 let schema = table.schema().clone();
3614 let schema_cols = &schema.columns;
3615 let alias = from
3616 .primary
3617 .alias
3618 .as_deref()
3619 .unwrap_or(from.primary.name.as_str());
3620 let ctx = EvalContext::new(schema_cols, Some(alias));
3621 let seg = self
3622 .active_catalog()
3623 .cold_segment(segment_id)
3624 .ok_or_else(|| {
3625 EngineError::Unsupported(alloc::format!(
3626 "AS OF SEGMENT: cold segment {segment_id} not registered"
3627 ))
3628 })?;
3629 let mut out_rows: Vec<Row> = Vec::new();
3630 let mut limit_remaining: Option<usize> =
3631 stmt.limit_literal().and_then(|n| usize::try_from(n).ok());
3632 for (_key, body) in seg.scan() {
3633 let (row, _consumed) = spg_storage::decode_row_body_dense(&body, &schema)
3634 .map_err(EngineError::Storage)?;
3635 if let Some(where_expr) = &stmt.where_ {
3636 let cond = self.eval_expr_simple(where_expr, &row, &ctx)?;
3637 if !matches!(cond, Value::Bool(true)) {
3638 continue;
3639 }
3640 }
3641 let projected = self.project_row_simple(&row, &stmt.items, schema_cols, alias)?;
3643 out_rows.push(projected);
3644 if let Some(rem) = limit_remaining.as_mut() {
3645 if *rem == 0 {
3646 out_rows.pop();
3647 break;
3648 }
3649 *rem -= 1;
3650 }
3651 }
3652 let columns = self.derive_output_columns(&stmt.items, schema_cols, alias);
3654 Ok(QueryResult::Rows {
3655 columns,
3656 rows: out_rows,
3657 })
3658 }
3659
3660 fn eval_expr_simple(
3665 &self,
3666 expr: &Expr,
3667 row: &Row,
3668 ctx: &EvalContext,
3669 ) -> Result<Value, EngineError> {
3670 let cancel = CancelToken::none();
3671 self.eval_expr_with_correlated(expr, row, ctx, cancel, None)
3672 }
3673
3674 fn build_returning_rows(
3681 &self,
3682 table_name: &str,
3683 items: &[SelectItem],
3684 mutated_rows: Vec<Vec<Value>>,
3685 ) -> Result<QueryResult, EngineError> {
3686 let table = self.active_catalog().get(table_name).ok_or_else(|| {
3687 EngineError::Storage(StorageError::TableNotFound {
3688 name: table_name.into(),
3689 })
3690 })?;
3691 let schema_cols = table.schema().columns.clone();
3692 let columns = self.derive_output_columns(items, &schema_cols, table_name);
3693 let mut out_rows: Vec<Row> = Vec::with_capacity(mutated_rows.len());
3694 for values in mutated_rows {
3695 let row = Row::new(values);
3696 let projected = self.project_row_simple(&row, items, &schema_cols, table_name)?;
3697 out_rows.push(projected);
3698 }
3699 Ok(QueryResult::Rows {
3700 columns,
3701 rows: out_rows,
3702 })
3703 }
3704
3705 fn project_row_simple(
3709 &self,
3710 row: &Row,
3711 items: &[SelectItem],
3712 schema_cols: &[ColumnSchema],
3713 alias: &str,
3714 ) -> Result<Row, EngineError> {
3715 let ctx = EvalContext::new(schema_cols, Some(alias));
3716 let cancel = CancelToken::none();
3717 let mut out_vals = Vec::new();
3718 for item in items {
3719 match item {
3720 SelectItem::Wildcard => {
3721 out_vals.extend(row.values.iter().cloned());
3722 }
3723 SelectItem::Expr { expr, .. } => {
3724 let v = self.eval_expr_with_correlated(expr, row, &ctx, cancel, None)?;
3725 out_vals.push(v);
3726 }
3727 }
3728 }
3729 Ok(Row::new(out_vals))
3730 }
3731
3732 fn derive_output_columns(
3737 &self,
3738 items: &[SelectItem],
3739 schema_cols: &[ColumnSchema],
3740 _alias: &str,
3741 ) -> Vec<ColumnSchema> {
3742 let mut out = Vec::new();
3743 for item in items {
3744 match item {
3745 SelectItem::Wildcard => {
3746 out.extend(schema_cols.iter().cloned());
3747 }
3748 SelectItem::Expr { alias, .. } => {
3749 let name = alias
3750 .clone()
3751 .unwrap_or_else(|| "?column?".to_string());
3752 out.push(ColumnSchema::new(name, DataType::Text, true));
3755 }
3756 }
3757 }
3758 out
3759 }
3760
3761 fn exec_select_cancel(
3762 &self,
3763 stmt: &SelectStatement,
3764 cancel: CancelToken<'_>,
3765 ) -> Result<QueryResult, EngineError> {
3766 cancel.check()?;
3767 if let Some(from) = &stmt.from
3776 && let Some(seg_id) = from.primary.as_of_segment
3777 {
3778 return self.exec_select_as_of_segment(stmt, from, seg_id);
3779 }
3780 if let Some(from) = &stmt.from
3784 && from.joins.is_empty()
3785 && stmt.where_.is_none()
3786 && stmt.group_by.is_none()
3787 && stmt.having.is_none()
3788 && stmt.unions.is_empty()
3789 && stmt.order_by.is_empty()
3790 && stmt.limit.is_none()
3791 && stmt.offset.is_none()
3792 && !stmt.distinct
3793 && stmt.items.iter().all(|i| matches!(i, SelectItem::Wildcard))
3794 {
3795 let lower = from.primary.name.to_ascii_lowercase();
3796 match lower.as_str() {
3797 "spg_statistic" => return Ok(self.exec_spg_statistic()),
3798 "spg_stat_replication" => return Ok(self.exec_spg_stat_replication()),
3800 "spg_stat_segment" => return Ok(self.exec_spg_stat_segment()),
3801 "spg_stat_query" => return Ok(self.exec_spg_stat_query()),
3802 "spg_stat_activity" => return Ok(self.exec_spg_stat_activity()),
3803 "spg_audit_chain" => return Ok(self.exec_spg_audit_chain()),
3804 "spg_audit_verify" => return Ok(self.exec_spg_audit_verify()),
3805 "spg_table_ddl" => return Ok(self.exec_spg_table_ddl()),
3806 "spg_role_ddl" => return Ok(self.exec_spg_role_ddl()),
3807 "spg_database_ddl" => return Ok(self.exec_spg_database_ddl()),
3808 _ => {}
3809 }
3810 }
3811 if !stmt.ctes.is_empty() {
3819 return self.exec_with_ctes(stmt, cancel);
3820 }
3821 let mut stmt_owned;
3828 let stmt_ref: &SelectStatement = if expr_tree_has_subquery(stmt) {
3829 stmt_owned = stmt.clone();
3830 self.resolve_select_subqueries(&mut stmt_owned, cancel)?;
3831 &stmt_owned
3832 } else {
3833 stmt
3834 };
3835 if stmt_ref.unions.is_empty() {
3836 return self.exec_bare_select_cancel(stmt_ref, cancel);
3837 }
3838 let mut head = stmt_ref.clone();
3843 head.unions = Vec::new();
3844 head.order_by = Vec::new();
3845 head.limit = None;
3846 let QueryResult::Rows { columns, mut rows } =
3847 self.exec_bare_select_cancel(&head, cancel)?
3848 else {
3849 unreachable!("bare SELECT cannot return CommandOk")
3850 };
3851 for (kind, peer) in &stmt_ref.unions {
3852 let QueryResult::Rows {
3853 columns: peer_cols,
3854 rows: peer_rows,
3855 } = self.exec_bare_select_cancel(peer, cancel)?
3856 else {
3857 unreachable!("bare SELECT cannot return CommandOk")
3858 };
3859 if peer_cols.len() != columns.len() {
3860 return Err(EngineError::Unsupported(alloc::format!(
3861 "UNION arity mismatch: head has {} columns, peer has {}",
3862 columns.len(),
3863 peer_cols.len()
3864 )));
3865 }
3866 rows.extend(peer_rows);
3867 if matches!(kind, UnionKind::Distinct) {
3868 rows = dedup_rows(rows);
3869 }
3870 }
3871 if !stmt.order_by.is_empty() {
3874 let synth_ctx = EvalContext::new(&columns, None);
3875 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
3876 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::with_capacity(rows.len());
3877 for r in rows {
3878 let keys = build_order_keys(&stmt.order_by, &r, &synth_ctx)?;
3879 tagged.push((keys, r));
3880 }
3881 sort_by_keys(&mut tagged, &descs);
3882 rows = tagged.into_iter().map(|(_, r)| r).collect();
3883 }
3884 apply_offset_and_limit(&mut rows, stmt.offset_literal(), stmt.limit_literal());
3885 Ok(QueryResult::Rows { columns, rows })
3886 }
3887
3888 #[allow(clippy::too_many_lines)]
3889 #[allow(clippy::too_many_lines)] fn exec_bare_select_cancel(
3891 &self,
3892 stmt: &SelectStatement,
3893 cancel: CancelToken<'_>,
3894 ) -> Result<QueryResult, EngineError> {
3895 if select_has_window(stmt) {
3900 return self.exec_select_with_window(stmt, cancel);
3901 }
3902 let Some(from) = &stmt.from else {
3907 let empty_schema: Vec<ColumnSchema> = Vec::new();
3908 let ctx = EvalContext::new(&empty_schema, None);
3909 let projection = build_projection(&stmt.items, &empty_schema, "")?;
3910 let dummy_row = Row::new(Vec::new());
3911 let mut values = Vec::with_capacity(projection.len());
3912 for p in &projection {
3913 values.push(eval::eval_expr(&p.expr, &dummy_row, &ctx)?);
3914 }
3915 let columns: Vec<ColumnSchema> = projection
3916 .into_iter()
3917 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
3918 .collect();
3919 return Ok(QueryResult::Rows {
3920 columns,
3921 rows: alloc::vec![Row::new(values)],
3922 });
3923 };
3924 if !from.joins.is_empty() {
3928 return self.exec_joined_select(stmt, from);
3929 }
3930 let primary = &from.primary;
3931 let table = self.active_catalog().get(&primary.name).ok_or_else(|| {
3932 StorageError::TableNotFound {
3933 name: primary.name.clone(),
3934 }
3935 })?;
3936 let schema_cols = &table.schema().columns;
3937 let alias = primary.alias.as_deref().unwrap_or(primary.name.as_str());
3940 let ctx = EvalContext::new(schema_cols, Some(alias));
3941
3942 if let Some(nsw_rows) = try_nsw_knn(stmt, table, schema_cols, alias) {
3947 return materialise_in_order(stmt, table, schema_cols, alias, &nsw_rows);
3948 }
3949
3950 let indexed_rows: Option<Vec<Cow<'_, Row>>> = stmt
3958 .where_
3959 .as_ref()
3960 .and_then(|w| try_index_seek(w, schema_cols, self.active_catalog(), table, alias));
3961
3962 if aggregate::uses_aggregate(stmt) {
3965 let mut filtered: Vec<&Row> = Vec::new();
3966 let mut memo = memoize::MemoizeCache::new();
3970 if let Some(rows) = &indexed_rows {
3971 for cow in rows {
3972 let row = cow.as_ref();
3973 if let Some(where_expr) = &stmt.where_ {
3974 let cond = self.eval_expr_with_correlated(
3975 where_expr,
3976 row,
3977 &ctx,
3978 cancel,
3979 Some(&mut memo),
3980 )?;
3981 if !matches!(cond, Value::Bool(true)) {
3982 continue;
3983 }
3984 }
3985 filtered.push(row);
3986 }
3987 } else {
3988 for i in 0..table.row_count() {
3989 let row = &table.rows()[i];
3990 if let Some(where_expr) = &stmt.where_ {
3991 let cond = self.eval_expr_with_correlated(
3992 where_expr,
3993 row,
3994 &ctx,
3995 cancel,
3996 Some(&mut memo),
3997 )?;
3998 if !matches!(cond, Value::Bool(true)) {
3999 continue;
4000 }
4001 }
4002 filtered.push(row);
4003 }
4004 }
4005 let mut agg = aggregate::run(stmt, &filtered, schema_cols, Some(alias))?;
4006 apply_offset_and_limit(&mut agg.rows, stmt.offset_literal(), stmt.limit_literal());
4007 return Ok(QueryResult::Rows {
4008 columns: agg.columns,
4009 rows: agg.rows,
4010 });
4011 }
4012
4013 let projection = build_projection(&stmt.items, schema_cols, alias)?;
4014
4015 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::new();
4018 let mut memo = memoize::MemoizeCache::new();
4020 let mut process_row = |row: &Row, loop_idx: usize| -> Result<(), EngineError> {
4023 if loop_idx.is_multiple_of(256) {
4024 cancel.check()?;
4025 }
4026 if let Some(where_expr) = &stmt.where_ {
4027 let cond = self.eval_expr_with_correlated(
4028 where_expr,
4029 row,
4030 &ctx,
4031 cancel,
4032 Some(&mut memo),
4033 )?;
4034 if !matches!(cond, Value::Bool(true)) {
4035 return Ok(());
4036 }
4037 }
4038 let mut values = Vec::with_capacity(projection.len());
4039 for p in &projection {
4040 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
4041 }
4042 let order_keys = if stmt.order_by.is_empty() {
4043 Vec::new()
4044 } else {
4045 build_order_keys(&stmt.order_by, row, &ctx)?
4046 };
4047 tagged.push((order_keys, Row::new(values)));
4048 Ok(())
4049 };
4050 if let Some(rows) = &indexed_rows {
4051 for (loop_idx, cow) in rows.iter().enumerate() {
4052 process_row(cow.as_ref(), loop_idx)?;
4053 }
4054 } else {
4055 for i in 0..table.row_count() {
4056 process_row(&table.rows()[i], i)?;
4057 }
4058 }
4059
4060 if !stmt.order_by.is_empty() {
4061 let keep = if stmt.distinct {
4066 None
4067 } else {
4068 stmt.limit_literal()
4069 .map(|l| l as usize + stmt.offset_literal().map_or(0, |o| o as usize))
4070 };
4071 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
4072 partial_sort_tagged(&mut tagged, keep, &descs);
4073 }
4074
4075 let mut output_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
4076 if stmt.distinct {
4077 output_rows = dedup_rows(output_rows);
4078 }
4079 apply_offset_and_limit(&mut output_rows, stmt.offset_literal(), stmt.limit_literal());
4080
4081 let columns: Vec<ColumnSchema> = projection
4082 .into_iter()
4083 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
4084 .collect();
4085
4086 Ok(QueryResult::Rows {
4087 columns,
4088 rows: output_rows,
4089 })
4090 }
4091
4092 #[allow(clippy::too_many_lines)]
4099 fn exec_joined_select(
4100 &self,
4101 stmt: &SelectStatement,
4102 from: &FromClause,
4103 ) -> Result<QueryResult, EngineError> {
4104 let primary_table = self
4107 .active_catalog()
4108 .get(&from.primary.name)
4109 .ok_or_else(|| StorageError::TableNotFound {
4110 name: from.primary.name.clone(),
4111 })?;
4112 let primary_alias = from
4113 .primary
4114 .alias
4115 .as_deref()
4116 .unwrap_or(from.primary.name.as_str())
4117 .to_string();
4118 let mut joined_tables: Vec<(&Table, String, JoinKind, Option<&Expr>)> = Vec::new();
4119 for j in &from.joins {
4120 let t = self.active_catalog().get(&j.table.name).ok_or_else(|| {
4121 StorageError::TableNotFound {
4122 name: j.table.name.clone(),
4123 }
4124 })?;
4125 let a = j
4126 .table
4127 .alias
4128 .as_deref()
4129 .unwrap_or(j.table.name.as_str())
4130 .to_string();
4131 joined_tables.push((t, a, j.kind, j.on.as_ref()));
4132 }
4133
4134 let mut combined_schema: Vec<ColumnSchema> = Vec::new();
4137 for col in &primary_table.schema().columns {
4138 combined_schema.push(ColumnSchema::new(
4139 alloc::format!("{primary_alias}.{}", col.name),
4140 col.ty,
4141 col.nullable,
4142 ));
4143 }
4144 for (t, a, _, _) in &joined_tables {
4145 for col in &t.schema().columns {
4146 combined_schema.push(ColumnSchema::new(
4147 alloc::format!("{a}.{}", col.name),
4148 col.ty,
4149 col.nullable,
4150 ));
4151 }
4152 }
4153 let ctx = EvalContext::new(&combined_schema, None);
4154
4155 let mut working: Vec<Row> = primary_table.rows().iter().cloned().collect();
4158 let mut produced_len = primary_table.schema().columns.len();
4159 for (t, _, kind, on) in &joined_tables {
4160 let right_arity = t.schema().columns.len();
4161 let mut next: Vec<Row> = Vec::new();
4162 for left in &working {
4163 let mut left_matched = false;
4164 for right in t.rows() {
4165 let mut combined_vals = left.values.clone();
4166 combined_vals.extend(right.values.iter().cloned());
4167 let combined = Row::new(combined_vals);
4170 let keep = if let Some(on_expr) = on {
4171 let cond = eval::eval_expr(on_expr, &combined, &ctx)?;
4172 matches!(cond, Value::Bool(true))
4173 } else {
4174 true
4176 };
4177 if keep {
4178 next.push(combined);
4179 left_matched = true;
4180 }
4181 }
4182 if !left_matched && matches!(kind, JoinKind::Left) {
4183 let mut combined_vals = left.values.clone();
4186 for _ in 0..right_arity {
4187 combined_vals.push(Value::Null);
4188 }
4189 next.push(Row::new(combined_vals));
4190 }
4191 }
4192 working = next;
4193 produced_len += right_arity;
4194 debug_assert!(produced_len <= combined_schema.len());
4195 }
4196
4197 let mut filtered: Vec<Row> = Vec::new();
4199 for row in working {
4200 if let Some(where_expr) = &stmt.where_ {
4201 let cond = eval::eval_expr(where_expr, &row, &ctx)?;
4202 if !matches!(cond, Value::Bool(true)) {
4203 continue;
4204 }
4205 }
4206 filtered.push(row);
4207 }
4208
4209 if aggregate::uses_aggregate(stmt) {
4212 let refs: Vec<&Row> = filtered.iter().collect();
4213 let mut agg = aggregate::run(stmt, &refs, &combined_schema, None)?;
4214 apply_offset_and_limit(&mut agg.rows, stmt.offset_literal(), stmt.limit_literal());
4215 return Ok(QueryResult::Rows {
4216 columns: agg.columns,
4217 rows: agg.rows,
4218 });
4219 }
4220
4221 let projection = build_projection(&stmt.items, &combined_schema, "")?;
4222 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::new();
4223 for row in &filtered {
4224 let mut values = Vec::with_capacity(projection.len());
4225 for p in &projection {
4226 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
4227 }
4228 let order_keys = if stmt.order_by.is_empty() {
4229 Vec::new()
4230 } else {
4231 build_order_keys(&stmt.order_by, row, &ctx)?
4232 };
4233 tagged.push((order_keys, Row::new(values)));
4234 }
4235 if !stmt.order_by.is_empty() {
4236 let keep = if stmt.distinct {
4237 None
4238 } else {
4239 stmt.limit_literal()
4240 .map(|l| l as usize + stmt.offset_literal().map_or(0, |o| o as usize))
4241 };
4242 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
4243 partial_sort_tagged(&mut tagged, keep, &descs);
4244 }
4245 let mut output_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
4246 if stmt.distinct {
4247 output_rows = dedup_rows(output_rows);
4248 }
4249 apply_offset_and_limit(&mut output_rows, stmt.offset_literal(), stmt.limit_literal());
4250 let columns: Vec<ColumnSchema> = projection
4251 .into_iter()
4252 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
4253 .collect();
4254 Ok(QueryResult::Rows {
4255 columns,
4256 rows: output_rows,
4257 })
4258 }
4259}
4260
4261#[derive(Debug, Clone)]
4264struct ProjectedItem {
4265 expr: Expr,
4266 output_name: String,
4267 ty: DataType,
4268 nullable: bool,
4269}
4270
4271fn dedup_rows(rows: Vec<Row>) -> Vec<Row> {
4277 let mut out: Vec<Row> = Vec::with_capacity(rows.len());
4278 for r in rows {
4279 if !out.iter().any(|seen| seen == &r) {
4280 out.push(r);
4281 }
4282 }
4283 out
4284}
4285
4286fn value_to_order_key(v: &Value) -> Result<f64, EngineError> {
4290 match v {
4291 Value::Null => Ok(f64::INFINITY),
4292 Value::SmallInt(n) => Ok(f64::from(*n)),
4293 Value::Int(n) => Ok(f64::from(*n)),
4294 Value::Date(d) => Ok(f64::from(*d)),
4295 #[allow(clippy::cast_precision_loss)]
4296 Value::Timestamp(t) => Ok(*t as f64),
4297 #[allow(clippy::cast_precision_loss)]
4298 Value::Numeric { scaled, scale } => {
4299 let mut divisor = 1.0_f64;
4305 for _ in 0..*scale {
4306 divisor *= 10.0;
4307 }
4308 Ok((*scaled as f64) / divisor)
4309 }
4310 #[allow(clippy::cast_precision_loss)]
4311 Value::BigInt(n) => Ok(*n as f64),
4312 Value::Float(x) => Ok(*x),
4313 Value::Bool(b) => Ok(if *b { 1.0 } else { 0.0 }),
4314 Value::Text(s) => {
4315 let mut key: u64 = 0;
4319 for &b in s.as_bytes().iter().take(8) {
4320 key = (key << 8) | u64::from(b);
4321 }
4322 #[allow(clippy::cast_precision_loss)]
4323 Ok(key as f64)
4324 }
4325 Value::Vector(_) | Value::Sq8Vector(_) | Value::HalfVector(_) => {
4326 Err(EngineError::Unsupported(
4327 "ORDER BY of a raw vector column is not meaningful — use `<->`".into(),
4328 ))
4329 }
4330 Value::Interval { .. } => Err(EngineError::Unsupported(
4331 "ORDER BY of an INTERVAL is not supported in v2.11 \
4332 (months vs micros has no single canonical ordering)"
4333 .into(),
4334 )),
4335 Value::Json(_) => Err(EngineError::Unsupported(
4336 "ORDER BY of a JSON value is not supported — cast the document to text first".into(),
4337 )),
4338 _ => Err(EngineError::Unsupported(
4342 "ORDER BY of this value type is not supported".into(),
4343 )),
4344 }
4345}
4346
4347fn try_nsw_knn(
4361 stmt: &SelectStatement,
4362 table: &Table,
4363 schema_cols: &[ColumnSchema],
4364 table_alias: &str,
4365) -> Option<Vec<usize>> {
4366 if stmt.distinct {
4367 return None;
4368 }
4369 let limit = usize::try_from(stmt.limit_literal()?).ok()?;
4370 if limit == 0 {
4371 return None;
4372 }
4373 if stmt.order_by.len() != 1 {
4377 return None;
4378 }
4379 let order = &stmt.order_by[0];
4380 if order.desc {
4384 return None;
4385 }
4386 let Expr::Binary { lhs, op, rhs } = &order.expr else {
4387 return None;
4388 };
4389 let metric = match op {
4390 BinOp::L2Distance => spg_storage::NswMetric::L2,
4391 BinOp::InnerProduct => spg_storage::NswMetric::InnerProduct,
4392 BinOp::CosineDistance => spg_storage::NswMetric::Cosine,
4393 _ => return None,
4394 };
4395 let ((Expr::Column(col), literal) | (literal, Expr::Column(col))) =
4397 (lhs.as_ref(), rhs.as_ref())
4398 else {
4399 return None;
4400 };
4401 if let Some(q) = &col.qualifier
4402 && q != table_alias
4403 {
4404 return None;
4405 }
4406 let col_pos = schema_cols.iter().position(|s| s.name == col.name)?;
4407 let query = literal_to_vector(literal)?;
4408 let idx = spg_storage::nsw_index_on(table, col_pos)?;
4409 if let Some(where_expr) = &stmt.where_ {
4410 let over_fetch = limit.saturating_mul(10).max(NSW_OVER_FETCH_FLOOR);
4414 let candidates = spg_storage::nsw_query(table, &idx.name, &query, over_fetch, metric);
4415 let ctx = EvalContext::new(schema_cols, Some(table_alias));
4416 let mut kept: Vec<usize> = Vec::with_capacity(limit);
4417 for i in candidates {
4418 let row = &table.rows()[i];
4419 let cond = eval::eval_expr(where_expr, row, &ctx).ok()?;
4420 if matches!(cond, Value::Bool(true)) {
4421 kept.push(i);
4422 if kept.len() >= limit {
4423 break;
4424 }
4425 }
4426 }
4427 Some(kept)
4428 } else {
4429 Some(spg_storage::nsw_query(
4430 table, &idx.name, &query, limit, metric,
4431 ))
4432 }
4433}
4434
4435const NSW_OVER_FETCH_FLOOR: usize = 32;
4439
4440fn literal_to_vector(e: &Expr) -> Option<Vec<f32>> {
4443 match e {
4444 Expr::Literal(Literal::Vector(v)) => Some(v.clone()),
4445 Expr::Cast { expr, .. } => literal_to_vector(expr),
4446 _ => None,
4447 }
4448}
4449
4450fn materialise_in_order(
4454 stmt: &SelectStatement,
4455 table: &Table,
4456 schema_cols: &[ColumnSchema],
4457 table_alias: &str,
4458 ordered_rows: &[usize],
4459) -> Result<QueryResult, EngineError> {
4460 let ctx = EvalContext::new(schema_cols, Some(table_alias));
4461 let projection = build_projection(&stmt.items, schema_cols, table_alias)?;
4462 let mut output_rows: Vec<Row> = Vec::with_capacity(ordered_rows.len());
4463 for &i in ordered_rows {
4464 let row = &table.rows()[i];
4465 let mut values = Vec::with_capacity(projection.len());
4466 for p in &projection {
4467 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
4468 }
4469 output_rows.push(Row::new(values));
4470 }
4471 apply_offset_and_limit(&mut output_rows, stmt.offset_literal(), stmt.limit_literal());
4472 let columns: Vec<ColumnSchema> = projection
4473 .into_iter()
4474 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
4475 .collect();
4476 Ok(QueryResult::Rows {
4477 columns,
4478 rows: output_rows,
4479 })
4480}
4481
4482fn try_index_seek<'a>(
4483 where_expr: &Expr,
4484 schema_cols: &[ColumnSchema],
4485 catalog: &'a Catalog,
4486 table: &'a Table,
4487 table_alias: &str,
4488) -> Option<Vec<Cow<'a, Row>>> {
4489 let Expr::Binary {
4490 lhs,
4491 op: BinOp::Eq,
4492 rhs,
4493 } = where_expr
4494 else {
4495 return None;
4496 };
4497 let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
4498 .or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
4499 let idx = table.index_on(col_pos)?;
4500 let key = IndexKey::from_value(&value)?;
4501 let locators = idx.lookup_eq(&key);
4502 let table_name = table.schema().name.as_str();
4503 let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(locators.len());
4511 for loc in locators {
4512 match *loc {
4513 spg_storage::RowLocator::Hot(i) => {
4514 if let Some(row) = table.rows().get(i) {
4515 out.push(Cow::Borrowed(row));
4516 }
4517 }
4518 spg_storage::RowLocator::Cold { segment_id, .. } => {
4519 if let Some(row) = catalog.resolve_cold_locator(table_name, segment_id, &key) {
4520 out.push(Cow::Owned(row));
4521 }
4522 }
4523 }
4524 }
4525 Some(out)
4526}
4527
4528fn try_pk_predicate(
4540 where_expr: &Expr,
4541 schema_cols: &[ColumnSchema],
4542 table_alias: &str,
4543) -> Option<(usize, IndexKey)> {
4544 let Expr::Binary {
4545 lhs,
4546 op: BinOp::Eq,
4547 rhs,
4548 } = where_expr
4549 else {
4550 return None;
4551 };
4552 let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
4553 .or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
4554 let key = IndexKey::from_value(&value)?;
4555 Some((col_pos, key))
4556}
4557
4558fn resolve_col_literal_pair(
4559 col_side: &Expr,
4560 lit_side: &Expr,
4561 schema_cols: &[ColumnSchema],
4562 table_alias: &str,
4563) -> Option<(usize, Value)> {
4564 let Expr::Column(c) = col_side else {
4565 return None;
4566 };
4567 if let Some(q) = &c.qualifier
4568 && q != table_alias
4569 {
4570 return None;
4571 }
4572 let pos = schema_cols.iter().position(|s| s.name == c.name)?;
4573 let Expr::Literal(l) = lit_side else {
4574 return None;
4575 };
4576 let v = match l {
4577 Literal::Integer(n) => {
4578 if let Ok(small) = i32::try_from(*n) {
4579 Value::Int(small)
4580 } else {
4581 Value::BigInt(*n)
4582 }
4583 }
4584 Literal::Float(x) => Value::Float(*x),
4585 Literal::String(s) => Value::Text(s.clone()),
4586 Literal::Bool(b) => Value::Bool(*b),
4587 Literal::Null => Value::Null,
4588 Literal::Vector(_) | Literal::Interval { .. } => return None,
4591 };
4592 Some((pos, v))
4593}
4594
4595fn resolve_projection_column<'a>(
4600 c: &ColumnName,
4601 schema_cols: &'a [ColumnSchema],
4602 table_alias: &str,
4603) -> Result<&'a ColumnSchema, EngineError> {
4604 if let Some(q) = &c.qualifier {
4605 let composite = alloc::format!("{q}.{name}", name = c.name);
4606 if let Some(s) = schema_cols.iter().find(|s| s.name == composite) {
4607 return Ok(s);
4608 }
4609 if q == table_alias
4612 && let Some(s) = schema_cols.iter().find(|s| s.name == c.name)
4613 {
4614 return Ok(s);
4615 }
4616 let prefix = alloc::format!("{q}.");
4620 let qualifier_known =
4621 q == table_alias || schema_cols.iter().any(|s| s.name.starts_with(&prefix));
4622 if !qualifier_known {
4623 return Err(EngineError::Eval(EvalError::UnknownQualifier {
4624 qualifier: q.clone(),
4625 }));
4626 }
4627 return Err(EngineError::Eval(EvalError::ColumnNotFound {
4628 name: c.name.clone(),
4629 }));
4630 }
4631 if let Some(s) = schema_cols.iter().find(|s| s.name == c.name) {
4632 return Ok(s);
4633 }
4634 let suffix = alloc::format!(".{name}", name = c.name);
4635 let mut matches = schema_cols.iter().filter(|s| s.name.ends_with(&suffix));
4636 let first = matches.next();
4637 let extra = matches.next();
4638 match (first, extra) {
4639 (Some(s), None) => Ok(s),
4640 (Some(_), Some(_)) => Err(EngineError::Eval(EvalError::TypeMismatch {
4641 detail: alloc::format!("ambiguous column reference: {}", c.name),
4642 })),
4643 _ => Err(EngineError::Eval(EvalError::ColumnNotFound {
4644 name: c.name.clone(),
4645 })),
4646 }
4647}
4648
4649fn build_projection(
4650 items: &[SelectItem],
4651 schema_cols: &[ColumnSchema],
4652 table_alias: &str,
4653) -> Result<Vec<ProjectedItem>, EngineError> {
4654 let mut out = Vec::new();
4655 for item in items {
4656 match item {
4657 SelectItem::Wildcard => {
4658 for col in schema_cols {
4659 out.push(ProjectedItem {
4660 expr: Expr::Column(ColumnName {
4661 qualifier: None,
4662 name: col.name.clone(),
4663 }),
4664 output_name: col.name.clone(),
4665 ty: col.ty,
4666 nullable: col.nullable,
4667 });
4668 }
4669 }
4670 SelectItem::Expr { expr, alias } => {
4671 if let Expr::Column(c) = expr {
4676 let sch = resolve_projection_column(c, schema_cols, table_alias)?;
4677 let output_name = alias.clone().unwrap_or_else(|| c.name.clone());
4678 out.push(ProjectedItem {
4679 expr: expr.clone(),
4680 output_name,
4681 ty: sch.ty,
4682 nullable: sch.nullable,
4683 });
4684 } else {
4685 let output_name = alias.clone().unwrap_or_else(|| expr.to_string());
4686 out.push(ProjectedItem {
4687 expr: expr.clone(),
4688 output_name,
4689 ty: DataType::Text,
4690 nullable: true,
4691 });
4692 }
4693 }
4694 }
4695 }
4696 Ok(out)
4697}
4698
4699fn numeric_from_integer(
4703 n: i128,
4704 precision: u8,
4705 scale: u8,
4706 col_name: &str,
4707) -> Result<Value, EngineError> {
4708 let factor = pow10_i128(scale);
4709 let scaled = n.checked_mul(factor).ok_or_else(|| {
4710 EngineError::Unsupported(alloc::format!(
4711 "integer overflow scaling value for column `{col_name}` to scale {scale}"
4712 ))
4713 })?;
4714 check_precision(scaled, precision, col_name)?;
4715 Ok(Value::Numeric { scaled, scale })
4716}
4717
4718#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
4721fn numeric_from_float(
4722 x: f64,
4723 precision: u8,
4724 scale: u8,
4725 col_name: &str,
4726) -> Result<Value, EngineError> {
4727 if !x.is_finite() {
4728 return Err(EngineError::Unsupported(alloc::format!(
4729 "cannot store non-finite float in NUMERIC column `{col_name}`"
4730 )));
4731 }
4732 let mut factor = 1.0_f64;
4733 for _ in 0..scale {
4734 factor *= 10.0;
4735 }
4736 let shifted = x * factor;
4741 let biased = if shifted >= 0.0 {
4742 shifted + 0.5
4743 } else {
4744 shifted - 0.5
4745 };
4746 if !(-1e38..=1e38).contains(&biased) {
4749 return Err(EngineError::Unsupported(alloc::format!(
4750 "value {x} overflows NUMERIC range for column `{col_name}`"
4751 )));
4752 }
4753 let scaled = biased as i128;
4754 check_precision(scaled, precision, col_name)?;
4755 Ok(Value::Numeric { scaled, scale })
4756}
4757
4758fn numeric_rescale(
4761 scaled: i128,
4762 src_scale: u8,
4763 precision: u8,
4764 dst_scale: u8,
4765 col_name: &str,
4766) -> Result<Value, EngineError> {
4767 let new_scaled = if dst_scale >= src_scale {
4768 let bump = pow10_i128(dst_scale - src_scale);
4769 scaled.checked_mul(bump).ok_or_else(|| {
4770 EngineError::Unsupported(alloc::format!(
4771 "overflow rescaling NUMERIC for column `{col_name}`"
4772 ))
4773 })?
4774 } else {
4775 let drop = pow10_i128(src_scale - dst_scale);
4776 let half = drop / 2;
4777 if scaled >= 0 {
4778 (scaled + half) / drop
4779 } else {
4780 (scaled - half) / drop
4781 }
4782 };
4783 check_precision(new_scaled, precision, col_name)?;
4784 Ok(Value::Numeric {
4785 scaled: new_scaled,
4786 scale: dst_scale,
4787 })
4788}
4789
4790const fn numeric_truncate_to_integer(scaled: i128, scale: u8) -> i128 {
4793 if scale == 0 {
4794 return scaled;
4795 }
4796 let factor = pow10_i128_const(scale);
4797 scaled / factor
4798}
4799
4800fn check_precision(scaled: i128, precision: u8, col_name: &str) -> Result<(), EngineError> {
4804 if precision == 0 {
4805 return Ok(());
4806 }
4807 let limit = pow10_i128(precision);
4808 if scaled.unsigned_abs() >= limit.unsigned_abs() {
4809 return Err(EngineError::Unsupported(alloc::format!(
4810 "NUMERIC value exceeds precision {precision} for column `{col_name}`"
4811 )));
4812 }
4813 Ok(())
4814}
4815
4816const fn pow10_i128_const(p: u8) -> i128 {
4817 let mut acc: i128 = 1;
4818 let mut i = 0;
4819 while i < p {
4820 acc *= 10;
4821 i += 1;
4822 }
4823 acc
4824}
4825
4826fn pow10_i128(p: u8) -> i128 {
4827 pow10_i128_const(p)
4828}
4829
4830impl Engine {
4845 #[allow(
4856 clippy::too_many_lines,
4857 clippy::type_complexity,
4858 clippy::needless_range_loop
4859 )] fn exec_select_with_window(
4861 &self,
4862 stmt: &SelectStatement,
4863 cancel: CancelToken<'_>,
4864 ) -> Result<QueryResult, EngineError> {
4865 let from = stmt.from.as_ref().ok_or_else(|| {
4866 EngineError::Unsupported("window functions require a FROM clause".into())
4867 })?;
4868 if !from.joins.is_empty() {
4871 return Err(EngineError::Unsupported(
4872 "JOIN with window functions not yet supported".into(),
4873 ));
4874 }
4875 let primary = &from.primary;
4876 let table = self.active_catalog().get(&primary.name).ok_or_else(|| {
4877 StorageError::TableNotFound {
4878 name: primary.name.clone(),
4879 }
4880 })?;
4881 let alias = primary.alias.as_deref().unwrap_or(primary.name.as_str());
4882 let schema_cols = &table.schema().columns;
4883 let ctx = EvalContext::new(schema_cols, Some(alias));
4884
4885 let mut filtered: Vec<&Row> = Vec::new();
4887 for (i, row) in table.rows().iter().enumerate() {
4888 if i.is_multiple_of(256) {
4889 cancel.check()?;
4890 }
4891 if let Some(w) = &stmt.where_ {
4892 let cond = eval::eval_expr(w, row, &ctx)?;
4893 if !matches!(cond, Value::Bool(true)) {
4894 continue;
4895 }
4896 }
4897 filtered.push(row);
4898 }
4899 let n_rows = filtered.len();
4900
4901 let mut window_nodes: Vec<Expr> = Vec::new();
4903 for item in &stmt.items {
4904 if let SelectItem::Expr { expr, .. } = item {
4905 collect_window_nodes(expr, &mut window_nodes);
4906 }
4907 }
4908
4909 let mut win_vals: Vec<Vec<Value>> = Vec::with_capacity(window_nodes.len());
4912 for wnode in &window_nodes {
4913 let Expr::WindowFunction {
4914 name,
4915 args,
4916 partition_by,
4917 order_by,
4918 frame,
4919 null_treatment,
4920 } = wnode
4921 else {
4922 unreachable!("collect_window_nodes pushes only WindowFunction");
4923 };
4924 let mut indexed: Vec<(Vec<Value>, Vec<(Value, bool)>, usize)> =
4926 Vec::with_capacity(n_rows);
4927 for (i, row) in filtered.iter().enumerate() {
4928 let pkey: Vec<Value> = partition_by
4929 .iter()
4930 .map(|p| eval::eval_expr(p, row, &ctx))
4931 .collect::<Result<_, _>>()?;
4932 let okey: Vec<(Value, bool)> = order_by
4933 .iter()
4934 .map(|(e, desc)| eval::eval_expr(e, row, &ctx).map(|v| (v, *desc)))
4935 .collect::<Result<_, _>>()?;
4936 indexed.push((pkey, okey, i));
4937 }
4938 indexed.sort_by(|a, b| {
4941 let p_cmp = partition_key_cmp(&a.0, &b.0);
4942 if p_cmp != core::cmp::Ordering::Equal {
4943 return p_cmp;
4944 }
4945 order_key_cmp(&a.1, &b.1)
4946 });
4947 let mut out_vals: Vec<Value> = alloc::vec![Value::Null; n_rows];
4949 let mut p_start = 0;
4950 while p_start < indexed.len() {
4951 let mut p_end = p_start + 1;
4952 while p_end < indexed.len()
4953 && partition_key_cmp(&indexed[p_start].0, &indexed[p_end].0)
4954 == core::cmp::Ordering::Equal
4955 {
4956 p_end += 1;
4957 }
4958 compute_window_partition(
4960 name,
4961 args,
4962 !order_by.is_empty(),
4963 frame.as_ref(),
4964 *null_treatment,
4965 &indexed[p_start..p_end],
4966 &filtered,
4967 &ctx,
4968 &mut out_vals,
4969 )?;
4970 p_start = p_end;
4971 }
4972 win_vals.push(out_vals);
4973 }
4974
4975 let mut ext_cols = schema_cols.clone();
4977 for i in 0..window_nodes.len() {
4978 ext_cols.push(ColumnSchema::new(
4979 alloc::format!("__win_{i}"),
4980 DataType::Text, true,
4982 ));
4983 }
4984 let mut ext_rows: Vec<Row> = Vec::with_capacity(n_rows);
4986 for i in 0..n_rows {
4987 let mut values = filtered[i].values.clone();
4988 for w in 0..window_nodes.len() {
4989 values.push(win_vals[w][i].clone());
4990 }
4991 ext_rows.push(Row::new(values));
4992 }
4993 let mut rewritten_items: Vec<SelectItem> = Vec::with_capacity(stmt.items.len());
4995 for item in &stmt.items {
4996 let new_item = match item {
4997 SelectItem::Wildcard => SelectItem::Wildcard,
4998 SelectItem::Expr { expr, alias } => {
4999 let mut e = expr.clone();
5000 rewrite_window_to_columns(&mut e, &window_nodes);
5001 SelectItem::Expr {
5002 expr: e,
5003 alias: alias.clone(),
5004 }
5005 }
5006 };
5007 rewritten_items.push(new_item);
5008 }
5009
5010 let ext_ctx = EvalContext::new(&ext_cols, Some(alias));
5012 let projection = build_projection(&rewritten_items, &ext_cols, alias)?;
5013 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::with_capacity(n_rows);
5014 for (i, row) in ext_rows.iter().enumerate() {
5015 if i.is_multiple_of(256) {
5016 cancel.check()?;
5017 }
5018 let mut values = Vec::with_capacity(projection.len());
5019 for p in &projection {
5020 values.push(eval::eval_expr(&p.expr, row, &ext_ctx)?);
5021 }
5022 let order_keys = if stmt.order_by.is_empty() {
5023 Vec::new()
5024 } else {
5025 let mut keys = Vec::with_capacity(stmt.order_by.len());
5026 for o in &stmt.order_by {
5027 let mut e = o.expr.clone();
5028 rewrite_window_to_columns(&mut e, &window_nodes);
5029 let key = eval::eval_expr(&e, row, &ext_ctx)?;
5030 keys.push(value_to_order_key(&key)?);
5031 }
5032 keys
5033 };
5034 tagged.push((order_keys, Row::new(values)));
5035 }
5036 if !stmt.order_by.is_empty() {
5038 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
5039 sort_by_keys(&mut tagged, &descs);
5040 }
5041 let mut out_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
5042 apply_offset_and_limit(&mut out_rows, stmt.offset_literal(), stmt.limit_literal());
5043 let final_cols: Vec<ColumnSchema> = projection
5044 .into_iter()
5045 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
5046 .collect();
5047 Ok(QueryResult::Rows {
5048 columns: final_cols,
5049 rows: out_rows,
5050 })
5051 }
5052
5053 fn exec_with_ctes(
5060 &self,
5061 stmt: &SelectStatement,
5062 cancel: CancelToken<'_>,
5063 ) -> Result<QueryResult, EngineError> {
5064 cancel.check()?;
5065 let mut catalog = self.active_catalog().clone();
5066 for cte in &stmt.ctes {
5067 if catalog.get(&cte.name).is_some() {
5068 return Err(EngineError::Unsupported(alloc::format!(
5069 "CTE name {:?} shadows an existing table; rename the CTE",
5070 cte.name
5071 )));
5072 }
5073 let (columns, rows) = if cte.recursive {
5074 self.materialise_recursive_cte(cte, &catalog, cancel)?
5075 } else {
5076 let body_result = self.exec_select_cancel(&cte.body, cancel)?;
5077 let QueryResult::Rows { columns, rows } = body_result else {
5078 return Err(EngineError::Unsupported(alloc::format!(
5079 "CTE {:?} body did not return rows",
5080 cte.name
5081 )));
5082 };
5083 (columns, rows)
5084 };
5085 let inferred = infer_column_types(&columns, &rows);
5090 let mut columns = inferred;
5091 if !cte.column_overrides.is_empty() {
5093 if cte.column_overrides.len() != columns.len() {
5094 return Err(EngineError::Unsupported(alloc::format!(
5095 "CTE {:?} column list has {} names but body returns {} columns",
5096 cte.name,
5097 cte.column_overrides.len(),
5098 columns.len()
5099 )));
5100 }
5101 for (col, name) in columns.iter_mut().zip(cte.column_overrides.iter()) {
5102 col.name.clone_from(name);
5103 }
5104 }
5105 let schema = TableSchema::new(cte.name.clone(), columns);
5106 catalog.create_table(schema).map_err(EngineError::Storage)?;
5107 let table = catalog
5108 .get_mut(&cte.name)
5109 .expect("just-created CTE table must exist");
5110 for row in rows {
5111 table.insert(row).map_err(EngineError::Storage)?;
5112 }
5113 }
5114 let mut body = stmt.clone();
5117 body.ctes = Vec::new();
5118 let mut temp = Engine::restore(catalog);
5119 if let Some(c) = self.clock {
5120 temp = temp.with_clock(c);
5121 }
5122 if let Some(f) = self.salt_fn {
5123 temp = temp.with_salt_fn(f);
5124 }
5125 temp.exec_select_cancel(&body, cancel)
5126 }
5127
5128 #[allow(clippy::too_many_lines)]
5138 fn materialise_recursive_cte(
5139 &self,
5140 cte: &spg_sql::ast::Cte,
5141 base_catalog: &Catalog,
5142 cancel: CancelToken<'_>,
5143 ) -> Result<(Vec<ColumnSchema>, Vec<Row>), EngineError> {
5144 const MAX_TOTAL_ROWS: usize = 1_000_000;
5145 const MAX_ITERATIONS: usize = 100_000;
5146 cancel.check()?;
5147 if cte.body.unions.is_empty() {
5148 return Err(EngineError::Unsupported(alloc::format!(
5149 "WITH RECURSIVE {:?} body must be a UNION of an anchor and a recursive term",
5150 cte.name
5151 )));
5152 }
5153 let mut anchor = cte.body.clone();
5155 let union_terms = core::mem::take(&mut anchor.unions);
5156 anchor.ctes = Vec::new();
5157 if select_refers_to(&anchor, &cte.name) {
5159 return Err(EngineError::Unsupported(alloc::format!(
5160 "WITH RECURSIVE {:?}: the anchor must not reference the CTE itself",
5161 cte.name
5162 )));
5163 }
5164 let anchor_result = self.exec_select_cancel(&anchor, cancel)?;
5165 let QueryResult::Rows {
5166 columns: anchor_cols,
5167 rows: anchor_rows,
5168 } = anchor_result
5169 else {
5170 return Err(EngineError::Unsupported(alloc::format!(
5171 "WITH RECURSIVE {:?}: anchor did not return rows",
5172 cte.name
5173 )));
5174 };
5175 let mut columns = infer_column_types(&anchor_cols, &anchor_rows);
5179 if !cte.column_overrides.is_empty() {
5180 if cte.column_overrides.len() != columns.len() {
5181 return Err(EngineError::Unsupported(alloc::format!(
5182 "CTE {:?} column list has {} names but anchor returns {} columns",
5183 cte.name,
5184 cte.column_overrides.len(),
5185 columns.len()
5186 )));
5187 }
5188 for (col, name) in columns.iter_mut().zip(cte.column_overrides.iter()) {
5189 col.name.clone_from(name);
5190 }
5191 }
5192 let mut all_rows: Vec<Row> = anchor_rows.clone();
5193 let mut working_set: Vec<Row> = anchor_rows;
5194 let mut seen: alloc::collections::BTreeSet<Vec<u8>> = alloc::collections::BTreeSet::new();
5195 let all_union_all = union_terms.iter().all(|(k, _)| matches!(k, UnionKind::All));
5198 if !all_union_all {
5199 for r in &all_rows {
5200 seen.insert(encode_row_key(r));
5201 }
5202 }
5203 for iter in 0..MAX_ITERATIONS {
5204 cancel.check()?;
5205 if working_set.is_empty() {
5206 break;
5207 }
5208 let mut iter_catalog = base_catalog.clone();
5210 let schema = TableSchema::new(cte.name.clone(), columns.clone());
5211 iter_catalog
5212 .create_table(schema)
5213 .map_err(EngineError::Storage)?;
5214 {
5215 let table = iter_catalog.get_mut(&cte.name).expect("just-created");
5216 for row in &working_set {
5217 table.insert(row.clone()).map_err(EngineError::Storage)?;
5218 }
5219 }
5220 let mut iter_engine = Engine::restore(iter_catalog);
5221 if let Some(c) = self.clock {
5222 iter_engine = iter_engine.with_clock(c);
5223 }
5224 if let Some(f) = self.salt_fn {
5225 iter_engine = iter_engine.with_salt_fn(f);
5226 }
5227 let mut next_set: Vec<Row> = Vec::new();
5229 for (_, term) in &union_terms {
5230 let mut term = term.clone();
5231 term.ctes = Vec::new();
5232 let r = iter_engine.exec_select_cancel(&term, cancel)?;
5233 let QueryResult::Rows {
5234 columns: rc,
5235 rows: rs,
5236 } = r
5237 else {
5238 return Err(EngineError::Unsupported(alloc::format!(
5239 "WITH RECURSIVE {:?}: recursive term did not return rows",
5240 cte.name
5241 )));
5242 };
5243 if rc.len() != columns.len() {
5244 return Err(EngineError::Unsupported(alloc::format!(
5245 "WITH RECURSIVE {:?}: column count of recursive term ({}) does not match anchor ({})",
5246 cte.name,
5247 rc.len(),
5248 columns.len()
5249 )));
5250 }
5251 for row in rs {
5252 if !all_union_all {
5253 let key = encode_row_key(&row);
5254 if !seen.insert(key) {
5255 continue;
5256 }
5257 }
5258 next_set.push(row);
5259 }
5260 }
5261 if next_set.is_empty() {
5262 break;
5263 }
5264 all_rows.extend(next_set.iter().cloned());
5265 working_set = next_set;
5266 if all_rows.len() > MAX_TOTAL_ROWS {
5267 return Err(EngineError::Unsupported(alloc::format!(
5268 "WITH RECURSIVE {:?}: produced more than {MAX_TOTAL_ROWS} rows — likely runaway recursion",
5269 cte.name
5270 )));
5271 }
5272 if iter + 1 == MAX_ITERATIONS {
5273 return Err(EngineError::Unsupported(alloc::format!(
5274 "WITH RECURSIVE {:?}: exceeded {MAX_ITERATIONS} iterations",
5275 cte.name
5276 )));
5277 }
5278 }
5279 Ok((columns, all_rows))
5280 }
5281
5282 fn resolve_select_subqueries(
5283 &self,
5284 stmt: &mut SelectStatement,
5285 cancel: CancelToken<'_>,
5286 ) -> Result<(), EngineError> {
5287 for item in &mut stmt.items {
5288 if let SelectItem::Expr { expr, .. } = item {
5289 self.resolve_expr_subqueries(expr, cancel)?;
5290 }
5291 }
5292 if let Some(w) = &mut stmt.where_ {
5293 self.resolve_expr_subqueries(w, cancel)?;
5294 }
5295 if let Some(gs) = &mut stmt.group_by {
5296 for g in gs {
5297 self.resolve_expr_subqueries(g, cancel)?;
5298 }
5299 }
5300 if let Some(h) = &mut stmt.having {
5301 self.resolve_expr_subqueries(h, cancel)?;
5302 }
5303 for o in &mut stmt.order_by {
5304 self.resolve_expr_subqueries(&mut o.expr, cancel)?;
5305 }
5306 for (_, peer) in &mut stmt.unions {
5307 self.resolve_select_subqueries(peer, cancel)?;
5308 }
5309 Ok(())
5310 }
5311
5312 #[allow(clippy::only_used_in_recursion)] fn resolve_expr_subqueries(
5314 &self,
5315 e: &mut Expr,
5316 cancel: CancelToken<'_>,
5317 ) -> Result<(), EngineError> {
5318 if let Some(replacement) = self.subquery_replacement(e, cancel)? {
5320 *e = replacement;
5321 return Ok(());
5322 }
5323 match e {
5324 Expr::Binary { lhs, rhs, .. } => {
5325 self.resolve_expr_subqueries(lhs, cancel)?;
5326 self.resolve_expr_subqueries(rhs, cancel)?;
5327 }
5328 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
5329 self.resolve_expr_subqueries(expr, cancel)?;
5330 }
5331 Expr::FunctionCall { args, .. } => {
5332 for a in args {
5333 self.resolve_expr_subqueries(a, cancel)?;
5334 }
5335 }
5336 Expr::Like { expr, pattern, .. } => {
5337 self.resolve_expr_subqueries(expr, cancel)?;
5338 self.resolve_expr_subqueries(pattern, cancel)?;
5339 }
5340 Expr::Extract { source, .. } => self.resolve_expr_subqueries(source, cancel)?,
5341 Expr::WindowFunction {
5344 args,
5345 partition_by,
5346 order_by,
5347 ..
5348 } => {
5349 for a in args {
5350 self.resolve_expr_subqueries(a, cancel)?;
5351 }
5352 for p in partition_by {
5353 self.resolve_expr_subqueries(p, cancel)?;
5354 }
5355 for (e, _) in order_by {
5356 self.resolve_expr_subqueries(e, cancel)?;
5357 }
5358 }
5359 Expr::ScalarSubquery(_)
5363 | Expr::Exists { .. }
5364 | Expr::InSubquery { .. }
5365 | Expr::Literal(_)
5366 | Expr::Placeholder(_)
5367 | Expr::Column(_) => {}
5368 Expr::Array(items) => {
5370 for elem in items {
5371 self.resolve_expr_subqueries(elem, cancel)?;
5372 }
5373 }
5374 Expr::ArraySubscript { target, index } => {
5375 self.resolve_expr_subqueries(target, cancel)?;
5376 self.resolve_expr_subqueries(index, cancel)?;
5377 }
5378 Expr::AnyAll { expr, array, .. } => {
5379 self.resolve_expr_subqueries(expr, cancel)?;
5380 self.resolve_expr_subqueries(array, cancel)?;
5381 }
5382 }
5383 Ok(())
5384 }
5385
5386 fn eval_expr_with_correlated(
5394 &self,
5395 expr: &Expr,
5396 row: &Row,
5397 ctx: &EvalContext<'_>,
5398 cancel: CancelToken<'_>,
5399 memo: Option<&mut memoize::MemoizeCache>,
5400 ) -> Result<Value, EngineError> {
5401 if !expr_has_subquery(expr) {
5402 return eval::eval_expr(expr, row, ctx).map_err(EngineError::Eval);
5403 }
5404 let mut e = expr.clone();
5405 self.resolve_correlated_in_expr(&mut e, row, ctx, cancel, memo)?;
5406 eval::eval_expr(&e, row, ctx).map_err(EngineError::Eval)
5407 }
5408
5409 fn resolve_correlated_in_expr(
5410 &self,
5411 e: &mut Expr,
5412 row: &Row,
5413 ctx: &EvalContext<'_>,
5414 cancel: CancelToken<'_>,
5415 mut memo: Option<&mut memoize::MemoizeCache>,
5416 ) -> Result<(), EngineError> {
5417 match e {
5418 Expr::ScalarSubquery(inner) => {
5419 let cache_key = memo.as_ref().map(|_| memoize::CacheKey {
5424 subquery_repr: alloc::format!("{}", **inner),
5425 outer_values: row.values.clone(),
5426 });
5427 if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key.as_ref())
5428 && let Some(cached) = cache.get(k)
5429 {
5430 *e = value_to_literal_expr(cached)?;
5431 return Ok(());
5432 }
5433 let mut s = (**inner).clone();
5434 substitute_outer_columns(&mut s, row, ctx);
5435 let r = self.exec_select_cancel(&s, cancel)?;
5436 let QueryResult::Rows { rows, .. } = r else {
5437 return Err(EngineError::Unsupported(
5438 "scalar subquery: inner did not return rows".into(),
5439 ));
5440 };
5441 let value = match rows.as_slice() {
5442 [] => Value::Null,
5443 [r0] => r0.values.first().cloned().unwrap_or(Value::Null),
5444 _ => {
5445 return Err(EngineError::Unsupported(alloc::format!(
5446 "scalar subquery returned {} rows; expected 0 or 1",
5447 rows.len()
5448 )));
5449 }
5450 };
5451 if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key) {
5452 cache.insert(k, value.clone());
5453 }
5454 *e = value_to_literal_expr(value)?;
5455 }
5456 Expr::Exists { subquery, negated } => {
5457 let mut s = (**subquery).clone();
5458 substitute_outer_columns(&mut s, row, ctx);
5459 let r = self.exec_select_cancel(&s, cancel)?;
5460 let exists = matches!(r, QueryResult::Rows { rows, .. } if !rows.is_empty());
5461 let bit = if *negated { !exists } else { exists };
5462 *e = Expr::Literal(Literal::Bool(bit));
5463 }
5464 Expr::InSubquery {
5465 expr: lhs,
5466 subquery,
5467 negated,
5468 } => {
5469 self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
5470 let lhs_val = eval::eval_expr(lhs, row, ctx).map_err(EngineError::Eval)?;
5471 let mut s = (**subquery).clone();
5472 substitute_outer_columns(&mut s, row, ctx);
5473 let r = self.exec_select_cancel(&s, cancel)?;
5474 let QueryResult::Rows { columns, rows, .. } = r else {
5475 return Err(EngineError::Unsupported(
5476 "IN-subquery: inner did not return rows".into(),
5477 ));
5478 };
5479 if columns.len() != 1 {
5480 return Err(EngineError::Unsupported(alloc::format!(
5481 "IN-subquery must project exactly one column; got {}",
5482 columns.len()
5483 )));
5484 }
5485 let mut found = false;
5486 let mut any_null = false;
5487 for r0 in rows {
5488 let v = r0.values.into_iter().next().unwrap_or(Value::Null);
5489 if v.is_null() {
5490 any_null = true;
5491 continue;
5492 }
5493 if value_cmp(&v, &lhs_val) == core::cmp::Ordering::Equal {
5494 found = true;
5495 break;
5496 }
5497 }
5498 let bit = if found {
5499 !*negated
5500 } else if any_null {
5501 return Err(EngineError::Unsupported(
5502 "IN-subquery with NULL in result and no match: NULL semantics not yet implemented".into(),
5503 ));
5504 } else {
5505 *negated
5506 };
5507 *e = Expr::Literal(Literal::Bool(bit));
5508 }
5509 Expr::Binary { lhs, rhs, .. } => {
5510 self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
5511 self.resolve_correlated_in_expr(rhs, row, ctx, cancel, memo.as_deref_mut())?;
5512 }
5513 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
5514 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
5515 }
5516 Expr::Like { expr, pattern, .. } => {
5517 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
5518 self.resolve_correlated_in_expr(pattern, row, ctx, cancel, memo.as_deref_mut())?;
5519 }
5520 Expr::FunctionCall { args, .. } => {
5521 for a in args {
5522 self.resolve_correlated_in_expr(a, row, ctx, cancel, memo.as_deref_mut())?;
5523 }
5524 }
5525 Expr::Extract { source, .. } => {
5526 self.resolve_correlated_in_expr(source, row, ctx, cancel, memo.as_deref_mut())?;
5527 }
5528 Expr::WindowFunction { .. } | Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
5529 Expr::Array(items) => {
5531 for elem in items {
5532 self.resolve_correlated_in_expr(elem, row, ctx, cancel, memo.as_deref_mut())?;
5533 }
5534 }
5535 Expr::ArraySubscript { target, index } => {
5536 self.resolve_correlated_in_expr(target, row, ctx, cancel, memo.as_deref_mut())?;
5537 self.resolve_correlated_in_expr(index, row, ctx, cancel, memo.as_deref_mut())?;
5538 }
5539 Expr::AnyAll { expr, array, .. } => {
5540 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
5541 self.resolve_correlated_in_expr(array, row, ctx, cancel, memo.as_deref_mut())?;
5542 }
5543 }
5544 Ok(())
5545 }
5546
5547 fn subquery_replacement(
5548 &self,
5549 e: &Expr,
5550 cancel: CancelToken<'_>,
5551 ) -> Result<Option<Expr>, EngineError> {
5552 match e {
5553 Expr::ScalarSubquery(inner) => {
5554 let mut s = (**inner).clone();
5555 self.resolve_select_subqueries(&mut s, cancel)?;
5558 let r = match self.exec_bare_select_cancel(&s, cancel) {
5559 Ok(r) => r,
5560 Err(e) if is_correlation_error(&e) => return Ok(None),
5561 Err(e) => return Err(e),
5562 };
5563 let QueryResult::Rows { rows, .. } = r else {
5564 return Err(EngineError::Unsupported(
5565 "scalar subquery: inner statement did not return rows".into(),
5566 ));
5567 };
5568 let value = match rows.as_slice() {
5569 [] => Value::Null,
5570 [row] => row.values.first().cloned().unwrap_or(Value::Null),
5571 _ => {
5572 return Err(EngineError::Unsupported(alloc::format!(
5573 "scalar subquery returned {} rows; expected 0 or 1",
5574 rows.len()
5575 )));
5576 }
5577 };
5578 Ok(Some(value_to_literal_expr(value)?))
5579 }
5580 Expr::Exists { subquery, negated } => {
5581 let mut s = (**subquery).clone();
5582 self.resolve_select_subqueries(&mut s, cancel)?;
5583 let r = match self.exec_bare_select_cancel(&s, cancel) {
5584 Ok(r) => r,
5585 Err(e) if is_correlation_error(&e) => return Ok(None),
5586 Err(e) => return Err(e),
5587 };
5588 let exists = match r {
5589 QueryResult::Rows { rows, .. } => !rows.is_empty(),
5590 QueryResult::CommandOk { .. } => false,
5591 };
5592 let bit = if *negated { !exists } else { exists };
5593 Ok(Some(Expr::Literal(Literal::Bool(bit))))
5594 }
5595 Expr::InSubquery {
5596 expr,
5597 subquery,
5598 negated,
5599 } => {
5600 let mut s = (**subquery).clone();
5601 self.resolve_select_subqueries(&mut s, cancel)?;
5602 let r = match self.exec_bare_select_cancel(&s, cancel) {
5603 Ok(r) => r,
5604 Err(e) if is_correlation_error(&e) => return Ok(None),
5605 Err(e) => return Err(e),
5606 };
5607 let QueryResult::Rows { columns, rows, .. } = r else {
5608 return Err(EngineError::Unsupported(
5609 "IN-subquery: inner statement did not return rows".into(),
5610 ));
5611 };
5612 if columns.len() != 1 {
5613 return Err(EngineError::Unsupported(alloc::format!(
5614 "IN-subquery must project exactly one column; got {}",
5615 columns.len()
5616 )));
5617 }
5618 let mut acc: Option<Expr> = None;
5621 for row in rows {
5622 let v = row.values.into_iter().next().unwrap_or(Value::Null);
5623 let lit = value_to_literal_expr(v)?;
5624 let cmp = Expr::Binary {
5625 lhs: expr.clone(),
5626 op: BinOp::Eq,
5627 rhs: Box::new(lit),
5628 };
5629 acc = Some(match acc {
5630 None => cmp,
5631 Some(prev) => Expr::Binary {
5632 lhs: Box::new(prev),
5633 op: BinOp::Or,
5634 rhs: Box::new(cmp),
5635 },
5636 });
5637 }
5638 let combined = acc.unwrap_or(Expr::Literal(Literal::Bool(false)));
5639 let final_expr = if *negated {
5640 Expr::Unary {
5641 op: UnOp::Not,
5642 expr: Box::new(combined),
5643 }
5644 } else {
5645 combined
5646 };
5647 Ok(Some(final_expr))
5648 }
5649 _ => Ok(None),
5650 }
5651 }
5652}
5653
5654fn select_refers_to(stmt: &SelectStatement, target: &str) -> bool {
5666 if let Some(from) = &stmt.from
5667 && from_refers_to(from, target)
5668 {
5669 return true;
5670 }
5671 for (_, peer) in &stmt.unions {
5672 if select_refers_to(peer, target) {
5673 return true;
5674 }
5675 }
5676 for item in &stmt.items {
5677 if let SelectItem::Expr { expr, .. } = item
5678 && expr_refers_to(expr, target)
5679 {
5680 return true;
5681 }
5682 }
5683 if let Some(w) = &stmt.where_
5684 && expr_refers_to(w, target)
5685 {
5686 return true;
5687 }
5688 false
5689}
5690
5691fn from_refers_to(from: &FromClause, target: &str) -> bool {
5692 if from.primary.name.eq_ignore_ascii_case(target) {
5693 return true;
5694 }
5695 from.joins
5696 .iter()
5697 .any(|j| j.table.name.eq_ignore_ascii_case(target))
5698}
5699
5700fn expr_refers_to(e: &Expr, target: &str) -> bool {
5701 match e {
5702 Expr::ScalarSubquery(s) => select_refers_to(s, target),
5703 Expr::Exists { subquery, .. } | Expr::InSubquery { subquery, .. } => {
5704 select_refers_to(subquery, target)
5705 }
5706 Expr::Binary { lhs, rhs, .. } => expr_refers_to(lhs, target) || expr_refers_to(rhs, target),
5707 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
5708 expr_refers_to(expr, target)
5709 }
5710 Expr::Like { expr, pattern, .. } => {
5711 expr_refers_to(expr, target) || expr_refers_to(pattern, target)
5712 }
5713 Expr::FunctionCall { args, .. } => args.iter().any(|a| expr_refers_to(a, target)),
5714 Expr::Extract { source, .. } => expr_refers_to(source, target),
5715 Expr::WindowFunction {
5716 args,
5717 partition_by,
5718 order_by,
5719 ..
5720 } => {
5721 args.iter().any(|a| expr_refers_to(a, target))
5722 || partition_by.iter().any(|p| expr_refers_to(p, target))
5723 || order_by.iter().any(|(o, _)| expr_refers_to(o, target))
5724 }
5725 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => false,
5726 Expr::Array(items) => items.iter().any(|e| expr_refers_to(e, target)),
5727 Expr::ArraySubscript { target: t, index } => {
5728 expr_refers_to(t, target) || expr_refers_to(index, target)
5729 }
5730 Expr::AnyAll { expr, array, .. } => {
5731 expr_refers_to(expr, target) || expr_refers_to(array, target)
5732 }
5733 }
5734}
5735
5736fn infer_column_types(columns: &[ColumnSchema], rows: &[Row]) -> Vec<ColumnSchema> {
5742 let mut out = columns.to_vec();
5743 for (col_idx, col) in out.iter_mut().enumerate() {
5744 if col.ty != DataType::Text {
5745 continue;
5746 }
5747 let mut inferred: Option<DataType> = None;
5748 let mut all_null = true;
5749 for row in rows {
5750 let Some(v) = row.values.get(col_idx) else {
5751 continue;
5752 };
5753 let ty = match v {
5754 Value::Null => continue,
5755 Value::SmallInt(_) => DataType::SmallInt,
5756 Value::Int(_) => DataType::Int,
5757 Value::BigInt(_) => DataType::BigInt,
5758 Value::Float(_) => DataType::Float,
5759 Value::Bool(_) => DataType::Bool,
5760 Value::Vector(_) => DataType::Vector {
5761 dim: 0,
5762 encoding: VecEncoding::F32,
5763 },
5764 _ => DataType::Text,
5765 };
5766 all_null = false;
5767 inferred = Some(match inferred {
5768 None => ty,
5769 Some(prev) if prev == ty => prev,
5770 Some(_) => DataType::Text,
5771 });
5772 }
5773 if let Some(t) = inferred {
5774 col.ty = t;
5775 col.nullable = true;
5776 } else if all_null {
5777 col.nullable = true;
5778 }
5779 }
5780 out
5781}
5782
5783#[allow(clippy::too_many_lines, clippy::format_push_string)]
5788fn build_index_suggestions(stmt: &SelectStatement, engine: &Engine) -> Vec<String> {
5805 use alloc::collections::BTreeSet;
5806 let mut seen: BTreeSet<(String, String)> = BTreeSet::new();
5807 let mut out: Vec<String> = Vec::new();
5808 let cat = engine.active_catalog();
5809 let Some(from) = &stmt.from else {
5813 return out;
5814 };
5815 let mut tables: Vec<String> = Vec::new();
5816 tables.push(from.primary.name.clone());
5817 for j in &from.joins {
5818 tables.push(j.table.name.clone());
5819 }
5820 let mut col_refs: Vec<spg_sql::ast::ColumnName> = Vec::new();
5823 if let Some(w) = &stmt.where_ {
5824 collect_column_refs(w, &mut col_refs);
5825 }
5826 for j in &from.joins {
5827 if let Some(on) = &j.on {
5828 collect_column_refs(on, &mut col_refs);
5829 }
5830 }
5831 for cn in &col_refs {
5832 let owner: Option<String> = if let Some(q) = &cn.qualifier {
5835 tables.iter().find(|t| t == &q).cloned()
5836 } else {
5837 tables.iter().find_map(|t| {
5838 cat.get(t).and_then(|tbl| {
5839 if tbl.schema().column_position(&cn.name).is_some() {
5840 Some(t.clone())
5841 } else {
5842 None
5843 }
5844 })
5845 })
5846 };
5847 let Some(owner) = owner else {
5848 continue;
5849 };
5850 let Some(tbl) = cat.get(&owner) else {
5851 continue;
5852 };
5853 let Some(col_pos) = tbl.schema().column_position(&cn.name) else {
5854 continue;
5855 };
5856 let already_indexed = tbl.indices().iter().any(|i| {
5859 matches!(i.kind, spg_storage::IndexKind::BTree(_))
5860 && i.column_position == col_pos
5861 && i.expression.is_none()
5862 && i.partial_predicate.is_none()
5863 });
5864 if already_indexed {
5865 continue;
5866 }
5867 if seen.insert((owner.clone(), cn.name.clone())) {
5868 out.push(alloc::format!(
5869 "SUGGEST: CREATE INDEX ix_{}_{} ON {} ({})",
5870 owner,
5871 cn.name,
5872 owner,
5873 cn.name
5874 ));
5875 }
5876 }
5877 out
5878}
5879
5880fn collect_column_refs(expr: &Expr, out: &mut Vec<spg_sql::ast::ColumnName>) {
5883 match expr {
5884 Expr::Column(cn) => out.push(cn.clone()),
5885 Expr::FunctionCall { args, .. } => {
5886 for a in args {
5887 collect_column_refs(a, out);
5888 }
5889 }
5890 Expr::Binary { lhs, rhs, .. } => {
5891 collect_column_refs(lhs, out);
5892 collect_column_refs(rhs, out);
5893 }
5894 Expr::Unary { expr: e, .. } => collect_column_refs(e, out),
5895 _ => {}
5896 }
5897}
5898
5899fn annotate_explain_lines(lines: &mut [String], total_rows: usize, engine: &Engine) {
5900 let catalog = engine.active_catalog();
5901 let cold_ids = catalog.cold_segment_ids_global();
5902 let any_cold = !cold_ids.is_empty();
5903 let cold_ids_repr = if any_cold {
5904 let mut s = alloc::string::String::from("[");
5905 for (i, id) in cold_ids.iter().enumerate() {
5906 if i > 0 {
5907 s.push(',');
5908 }
5909 s.push_str(&alloc::format!("{id}"));
5910 }
5911 s.push(']');
5912 s
5913 } else {
5914 alloc::string::String::new()
5915 };
5916 for (idx, line) in lines.iter_mut().enumerate() {
5917 let trimmed = line.trim_start();
5918 let is_top_level = idx == 0;
5919 if is_top_level {
5920 line.push_str(&alloc::format!(" (rows={total_rows})"));
5921 continue;
5922 }
5923 if let Some(rest) = trimmed.strip_prefix("From: ") {
5924 let (name, scan_kind) = match rest.split_once(" [") {
5925 Some((n, k)) => (n.trim(), k.trim_end_matches(']')),
5926 None => (rest.trim(), ""),
5927 };
5928 let bare = name.split_whitespace().next().unwrap_or(name);
5929 let hot = catalog.get(bare).map(|t| t.rows().len());
5930 let annot = match (hot, scan_kind) {
5935 (Some(h), "full scan") => {
5936 let mut s = alloc::format!(" (hot_rows={h}");
5937 if any_cold {
5938 s.push_str(&alloc::format!(
5939 ", cold_tier=present, cold_segments={cold_ids_repr}"
5940 ));
5941 }
5942 s.push(')');
5943 s
5944 }
5945 (Some(h), "index seek") => {
5946 let mut s = alloc::format!(" (hot_rows≤{h}");
5947 if any_cold {
5948 s.push_str(&alloc::format!(
5949 ", cold_tier=present, cold_segments={cold_ids_repr}"
5950 ));
5951 }
5952 s.push(')');
5953 s
5954 }
5955 _ => " (rows=—)".to_string(),
5956 };
5957 line.push_str(&annot);
5958 continue;
5959 }
5960 line.push_str(" (rows=—)");
5962 }
5963}
5964
5965fn explain_select(stmt: &SelectStatement, engine: &Engine, depth: usize, out: &mut Vec<String>) {
5966 let pad = " ".repeat(depth);
5967 let top = if !stmt.ctes.is_empty() {
5969 if stmt.ctes.iter().any(|c| c.recursive) {
5970 "CTEScan (WITH RECURSIVE)"
5971 } else {
5972 "CTEScan (WITH)"
5973 }
5974 } else if !stmt.unions.is_empty() {
5975 "UnionScan"
5976 } else if select_has_window(stmt) {
5977 "WindowAgg"
5978 } else if aggregate::uses_aggregate(stmt) {
5979 "Aggregate"
5980 } else if stmt.distinct {
5981 "Distinct"
5982 } else if stmt.from.is_some() {
5983 "TableScan"
5984 } else {
5985 "Result"
5986 };
5987 out.push(alloc::format!("{pad}{top}"));
5988 let child = " ".repeat(depth + 1);
5989 for cte in &stmt.ctes {
5991 let head = if cte.recursive {
5992 alloc::format!("{child}CTE (recursive): {}", cte.name)
5993 } else {
5994 alloc::format!("{child}CTE: {}", cte.name)
5995 };
5996 out.push(head);
5997 explain_select(&cte.body, engine, depth + 2, out);
5998 }
5999 if let Some(from) = &stmt.from {
6001 let mut tag = alloc::format!("{child}From: {}", from.primary.name);
6002 if let Some(alias) = &from.primary.alias {
6003 tag.push_str(&alloc::format!(" AS {alias}"));
6004 }
6005 if let Some(w) = &stmt.where_
6008 && let Some(table) = engine.active_catalog().get(&from.primary.name)
6009 {
6010 let alias = from.primary.alias.as_deref().unwrap_or(&from.primary.name);
6011 let cols = &table.schema().columns;
6012 if try_index_seek(w, cols, engine.active_catalog(), table, alias).is_some() {
6013 tag.push_str(" [index seek]");
6014 } else {
6015 tag.push_str(" [full scan]");
6016 }
6017 } else {
6018 tag.push_str(" [full scan]");
6019 }
6020 out.push(tag);
6021 for j in &from.joins {
6022 let kind = match j.kind {
6023 spg_sql::ast::JoinKind::Inner => "INNER JOIN",
6024 spg_sql::ast::JoinKind::Left => "LEFT JOIN",
6025 spg_sql::ast::JoinKind::Cross => "CROSS JOIN",
6026 };
6027 let mut s = alloc::format!("{child}{kind}: {}", j.table.name);
6028 if let Some(alias) = &j.table.alias {
6029 s.push_str(&alloc::format!(" AS {alias}"));
6030 }
6031 if j.on.is_some() {
6032 s.push_str(" (ON …)");
6033 }
6034 out.push(s);
6035 }
6036 }
6037 if let Some(w) = &stmt.where_ {
6039 let mut s = alloc::format!("{child}Filter: {w}");
6040 if expr_has_subquery(w) {
6041 s.push_str(" [subquery]");
6042 }
6043 out.push(s);
6044 }
6045 if let Some(gs) = &stmt.group_by {
6046 let mut parts = Vec::new();
6047 for g in gs {
6048 parts.push(alloc::format!("{g}"));
6049 }
6050 out.push(alloc::format!("{child}GroupBy: {}", parts.join(", ")));
6051 }
6052 if let Some(h) = &stmt.having {
6053 out.push(alloc::format!("{child}Having: {h}"));
6054 }
6055 for o in &stmt.order_by {
6056 let dir = if o.desc { "DESC" } else { "ASC" };
6057 out.push(alloc::format!("{child}OrderBy: {} {dir}", o.expr));
6058 }
6059 if let Some(lim) = stmt.limit {
6060 out.push(alloc::format!("{child}Limit: {lim}"));
6061 }
6062 if let Some(off) = stmt.offset {
6063 out.push(alloc::format!("{child}Offset: {off}"));
6064 }
6065 if stmt
6067 .items
6068 .iter()
6069 .any(|it| matches!(it, SelectItem::Wildcard))
6070 {
6071 out.push(alloc::format!("{child}Project: *"));
6072 } else {
6073 out.push(alloc::format!(
6074 "{child}Project: {} item(s)",
6075 stmt.items.len()
6076 ));
6077 }
6078 for (kind, peer) in &stmt.unions {
6080 let label = match kind {
6081 UnionKind::All => "UNION ALL",
6082 UnionKind::Distinct => "UNION",
6083 };
6084 out.push(alloc::format!("{child}{label}"));
6085 explain_select(peer, engine, depth + 2, out);
6086 }
6087}
6088
6089fn is_correlation_error(e: &EngineError) -> bool {
6094 matches!(
6095 e,
6096 EngineError::Eval(
6097 eval::EvalError::ColumnNotFound { .. } | eval::EvalError::UnknownQualifier { .. }
6098 )
6099 )
6100}
6101
6102fn substitute_outer_columns(stmt: &mut SelectStatement, row: &Row, ctx: &EvalContext<'_>) {
6110 let Some(outer_alias) = ctx.table_alias else {
6111 return;
6112 };
6113 substitute_in_select(stmt, row, ctx, outer_alias);
6114}
6115
6116fn substitute_in_select(
6117 stmt: &mut SelectStatement,
6118 row: &Row,
6119 ctx: &EvalContext<'_>,
6120 outer_alias: &str,
6121) {
6122 for item in &mut stmt.items {
6123 if let SelectItem::Expr { expr, .. } = item {
6124 substitute_in_expr(expr, row, ctx, outer_alias);
6125 }
6126 }
6127 if let Some(w) = &mut stmt.where_ {
6128 substitute_in_expr(w, row, ctx, outer_alias);
6129 }
6130 if let Some(gs) = &mut stmt.group_by {
6131 for g in gs {
6132 substitute_in_expr(g, row, ctx, outer_alias);
6133 }
6134 }
6135 if let Some(h) = &mut stmt.having {
6136 substitute_in_expr(h, row, ctx, outer_alias);
6137 }
6138 for o in &mut stmt.order_by {
6139 substitute_in_expr(&mut o.expr, row, ctx, outer_alias);
6140 }
6141 for (_, peer) in &mut stmt.unions {
6142 substitute_in_select(peer, row, ctx, outer_alias);
6143 }
6144}
6145
6146fn substitute_in_expr(e: &mut Expr, row: &Row, ctx: &EvalContext<'_>, outer_alias: &str) {
6147 if let Expr::Column(c) = e
6148 && let Some(qual) = &c.qualifier
6149 && qual.eq_ignore_ascii_case(outer_alias)
6150 {
6151 if let Some(idx) = ctx
6153 .columns
6154 .iter()
6155 .position(|sc| sc.name.eq_ignore_ascii_case(&c.name))
6156 {
6157 let v = row.values.get(idx).cloned().unwrap_or(Value::Null);
6158 if let Ok(lit) = value_to_literal_expr(v) {
6159 *e = lit;
6160 return;
6161 }
6162 }
6163 }
6164 match e {
6165 Expr::Binary { lhs, rhs, .. } => {
6166 substitute_in_expr(lhs, row, ctx, outer_alias);
6167 substitute_in_expr(rhs, row, ctx, outer_alias);
6168 }
6169 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6170 substitute_in_expr(expr, row, ctx, outer_alias);
6171 }
6172 Expr::Like { expr, pattern, .. } => {
6173 substitute_in_expr(expr, row, ctx, outer_alias);
6174 substitute_in_expr(pattern, row, ctx, outer_alias);
6175 }
6176 Expr::FunctionCall { args, .. } => {
6177 for a in args {
6178 substitute_in_expr(a, row, ctx, outer_alias);
6179 }
6180 }
6181 Expr::Extract { source, .. } => substitute_in_expr(source, row, ctx, outer_alias),
6182 Expr::WindowFunction {
6183 args,
6184 partition_by,
6185 order_by,
6186 ..
6187 } => {
6188 for a in args {
6189 substitute_in_expr(a, row, ctx, outer_alias);
6190 }
6191 for p in partition_by {
6192 substitute_in_expr(p, row, ctx, outer_alias);
6193 }
6194 for (o, _) in order_by {
6195 substitute_in_expr(o, row, ctx, outer_alias);
6196 }
6197 }
6198 Expr::ScalarSubquery(s) => substitute_in_select(s, row, ctx, outer_alias),
6199 Expr::Exists { subquery, .. } | Expr::InSubquery { subquery, .. } => {
6200 substitute_in_select(subquery, row, ctx, outer_alias);
6201 }
6202 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
6203 Expr::Array(items) => {
6204 for elem in items {
6205 substitute_in_expr(elem, row, ctx, outer_alias);
6206 }
6207 }
6208 Expr::ArraySubscript { target, index } => {
6209 substitute_in_expr(target, row, ctx, outer_alias);
6210 substitute_in_expr(index, row, ctx, outer_alias);
6211 }
6212 Expr::AnyAll { expr, array, .. } => {
6213 substitute_in_expr(expr, row, ctx, outer_alias);
6214 substitute_in_expr(array, row, ctx, outer_alias);
6215 }
6216 }
6217}
6218
6219fn encode_row_key(row: &Row) -> Vec<u8> {
6223 let mut out = Vec::new();
6224 for v in &row.values {
6225 let s = alloc::format!("{v:?}|");
6226 out.extend_from_slice(s.as_bytes());
6227 }
6228 out
6229}
6230
6231fn select_has_window(stmt: &SelectStatement) -> bool {
6232 for item in &stmt.items {
6233 if let SelectItem::Expr { expr, .. } = item
6234 && expr_has_window(expr)
6235 {
6236 return true;
6237 }
6238 }
6239 false
6240}
6241
6242fn expr_has_window(e: &Expr) -> bool {
6243 match e {
6244 Expr::WindowFunction { .. } => true,
6245 Expr::Binary { lhs, rhs, .. } => expr_has_window(lhs) || expr_has_window(rhs),
6246 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6247 expr_has_window(expr)
6248 }
6249 Expr::FunctionCall { args, .. } => args.iter().any(expr_has_window),
6250 Expr::Like { expr, pattern, .. } => expr_has_window(expr) || expr_has_window(pattern),
6251 Expr::Extract { source, .. } => expr_has_window(source),
6252 Expr::ScalarSubquery(_)
6253 | Expr::Exists { .. }
6254 | Expr::InSubquery { .. }
6255 | Expr::Literal(_)
6256 | Expr::Placeholder(_)
6257 | Expr::Column(_) => false,
6258 Expr::Array(items) => items.iter().any(expr_has_window),
6259 Expr::ArraySubscript { target, index } => {
6260 expr_has_window(target) || expr_has_window(index)
6261 }
6262 Expr::AnyAll { expr, array, .. } => {
6263 expr_has_window(expr) || expr_has_window(array)
6264 }
6265 }
6266}
6267
6268fn collect_window_nodes(e: &Expr, out: &mut Vec<Expr>) {
6269 if let Expr::WindowFunction { .. } = e {
6270 if !out.iter().any(|x| x == e) {
6275 out.push(e.clone());
6276 }
6277 return;
6278 }
6279 match e {
6280 Expr::WindowFunction { .. } => unreachable!(),
6282 Expr::Binary { lhs, rhs, .. } => {
6283 collect_window_nodes(lhs, out);
6284 collect_window_nodes(rhs, out);
6285 }
6286 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6287 collect_window_nodes(expr, out);
6288 }
6289 Expr::FunctionCall { args, .. } => {
6290 for a in args {
6291 collect_window_nodes(a, out);
6292 }
6293 }
6294 Expr::Like { expr, pattern, .. } => {
6295 collect_window_nodes(expr, out);
6296 collect_window_nodes(pattern, out);
6297 }
6298 Expr::Extract { source, .. } => collect_window_nodes(source, out),
6299 _ => {}
6300 }
6301}
6302
6303fn rewrite_window_to_columns(e: &mut Expr, window_nodes: &[Expr]) {
6304 if let Expr::WindowFunction { .. } = e
6305 && let Some(idx) = window_nodes.iter().position(|w| w == e)
6306 {
6307 *e = Expr::Column(spg_sql::ast::ColumnName {
6308 qualifier: None,
6309 name: alloc::format!("__win_{idx}"),
6310 });
6311 return;
6312 }
6313 match e {
6314 Expr::Binary { lhs, rhs, .. } => {
6315 rewrite_window_to_columns(lhs, window_nodes);
6316 rewrite_window_to_columns(rhs, window_nodes);
6317 }
6318 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6319 rewrite_window_to_columns(expr, window_nodes);
6320 }
6321 Expr::FunctionCall { args, .. } => {
6322 for a in args {
6323 rewrite_window_to_columns(a, window_nodes);
6324 }
6325 }
6326 Expr::Like { expr, pattern, .. } => {
6327 rewrite_window_to_columns(expr, window_nodes);
6328 rewrite_window_to_columns(pattern, window_nodes);
6329 }
6330 Expr::Extract { source, .. } => rewrite_window_to_columns(source, window_nodes),
6331 _ => {}
6332 }
6333}
6334
6335fn partition_key_cmp(a: &[Value], b: &[Value]) -> core::cmp::Ordering {
6339 for (x, y) in a.iter().zip(b.iter()) {
6340 let c = value_cmp(x, y);
6341 if c != core::cmp::Ordering::Equal {
6342 return c;
6343 }
6344 }
6345 a.len().cmp(&b.len())
6346}
6347
6348fn order_key_cmp(a: &[(Value, bool)], b: &[(Value, bool)]) -> core::cmp::Ordering {
6349 for ((va, desc), (vb, _)) in a.iter().zip(b.iter()) {
6350 let c = value_cmp(va, vb);
6351 let c = if *desc { c.reverse() } else { c };
6352 if c != core::cmp::Ordering::Equal {
6353 return c;
6354 }
6355 }
6356 a.len().cmp(&b.len())
6357}
6358
6359#[allow(clippy::match_same_arms)] fn value_cmp(a: &Value, b: &Value) -> core::cmp::Ordering {
6361 use core::cmp::Ordering;
6362 match (a, b) {
6363 (Value::Null, Value::Null) => Ordering::Equal,
6364 (Value::Null, _) => Ordering::Less,
6365 (_, Value::Null) => Ordering::Greater,
6366 (Value::Int(x), Value::Int(y)) => x.cmp(y),
6367 (Value::BigInt(x), Value::BigInt(y)) => x.cmp(y),
6368 (Value::SmallInt(x), Value::SmallInt(y)) => x.cmp(y),
6369 (Value::Text(x), Value::Text(y)) => x.cmp(y),
6370 (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
6371 (Value::Float(x), Value::Float(y)) => x.partial_cmp(y).unwrap_or(Ordering::Equal),
6372 (Value::Date(x), Value::Date(y)) => x.cmp(y),
6373 (Value::Timestamp(x), Value::Timestamp(y)) => x.cmp(y),
6374 _ => alloc::format!("{a:?}").cmp(&alloc::format!("{b:?}")),
6377 }
6378}
6379
6380#[allow(
6386 clippy::too_many_arguments,
6387 clippy::cast_possible_truncation,
6388 clippy::cast_possible_wrap,
6389 clippy::cast_precision_loss,
6390 clippy::cast_sign_loss,
6391 clippy::doc_markdown,
6392 clippy::too_many_lines,
6393 clippy::type_complexity,
6394 clippy::match_same_arms
6395)]
6396fn compute_window_partition(
6397 name: &str,
6398 args: &[Expr],
6399 ordered: bool,
6400 frame: Option<&WindowFrame>,
6401 null_treatment: spg_sql::ast::NullTreatment,
6402 slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)],
6403 filtered_rows: &[&Row],
6404 ctx: &EvalContext<'_>,
6405 out_vals: &mut [Value],
6406) -> Result<(), EngineError> {
6407 let ignore_nulls = matches!(null_treatment, spg_sql::ast::NullTreatment::Ignore);
6408 let lower = name.to_ascii_lowercase();
6409 match lower.as_str() {
6410 "row_number" => {
6411 for (rank, (_, _, idx)) in slice.iter().enumerate() {
6412 out_vals[*idx] = Value::BigInt((rank + 1) as i64);
6413 }
6414 Ok(())
6415 }
6416 "rank" => {
6417 let mut prev_key: Option<&[(Value, bool)]> = None;
6418 let mut current_rank: i64 = 1;
6419 for (i, (_, okey, idx)) in slice.iter().enumerate() {
6420 if let Some(p) = prev_key
6421 && order_key_cmp(p, okey) != core::cmp::Ordering::Equal
6422 {
6423 current_rank = (i + 1) as i64;
6424 }
6425 if prev_key.is_none() {
6426 current_rank = 1;
6427 }
6428 out_vals[*idx] = Value::BigInt(current_rank);
6429 prev_key = Some(okey.as_slice());
6430 }
6431 Ok(())
6432 }
6433 "dense_rank" => {
6434 let mut prev_key: Option<&[(Value, bool)]> = None;
6435 let mut current_rank: i64 = 0;
6436 for (_, okey, idx) in slice {
6437 if prev_key.is_none_or(|p| order_key_cmp(p, okey) != core::cmp::Ordering::Equal) {
6438 current_rank += 1;
6439 }
6440 out_vals[*idx] = Value::BigInt(current_rank);
6441 prev_key = Some(okey.as_slice());
6442 }
6443 Ok(())
6444 }
6445 "sum" | "avg" | "min" | "max" | "count" | "count_star" => {
6446 let arg_values: Vec<Value> = if lower == "count_star" || args.is_empty() {
6449 slice.iter().map(|_| Value::Null).collect()
6450 } else {
6451 slice
6452 .iter()
6453 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
6454 .collect::<Result<_, _>>()
6455 .map_err(EngineError::Eval)?
6456 };
6457 let eff = effective_frame(frame, ordered)?;
6461 #[allow(clippy::needless_range_loop)]
6462 for i in 0..slice.len() {
6463 let (lo, hi) = frame_bounds_for_row(&eff, i, slice);
6464 let mut sum: f64 = 0.0;
6465 let mut count: i64 = 0;
6466 let mut min_v: Option<f64> = None;
6467 let mut max_v: Option<f64> = None;
6468 let mut row_count: i64 = 0;
6469 if lo <= hi {
6470 for j in lo..=hi {
6471 let v = &arg_values[j];
6472 match lower.as_str() {
6473 "count_star" => row_count += 1,
6474 "count" => {
6475 if !v.is_null() {
6476 count += 1;
6477 }
6478 }
6479 _ => {
6480 if let Some(x) = value_to_f64(v) {
6481 sum += x;
6482 count += 1;
6483 min_v = Some(min_v.map_or(x, |m| m.min(x)));
6484 max_v = Some(max_v.map_or(x, |m| m.max(x)));
6485 }
6486 }
6487 }
6488 }
6489 }
6490 let value = match lower.as_str() {
6491 "count_star" => Value::BigInt(row_count),
6492 "count" => Value::BigInt(count),
6493 "sum" => Value::Float(sum),
6494 "avg" => {
6495 if count == 0 {
6496 Value::Null
6497 } else {
6498 Value::Float(sum / count as f64)
6499 }
6500 }
6501 "min" => min_v.map_or(Value::Null, Value::Float),
6502 "max" => max_v.map_or(Value::Null, Value::Float),
6503 _ => unreachable!(),
6504 };
6505 let (_, _, idx) = &slice[i];
6506 out_vals[*idx] = value;
6507 }
6508 Ok(())
6509 }
6510 "lag" | "lead" => {
6511 if args.is_empty() {
6514 return Err(EngineError::Unsupported(alloc::format!(
6515 "{lower}() requires at least one argument"
6516 )));
6517 }
6518 let offset: i64 = if args.len() >= 2 {
6519 let v = eval::eval_expr(&args[1], filtered_rows[slice[0].2], ctx)
6520 .map_err(EngineError::Eval)?;
6521 match v {
6522 Value::SmallInt(n) => i64::from(n),
6523 Value::Int(n) => i64::from(n),
6524 Value::BigInt(n) => n,
6525 _ => {
6526 return Err(EngineError::Unsupported(alloc::format!(
6527 "{lower}() offset must be integer"
6528 )));
6529 }
6530 }
6531 } else {
6532 1
6533 };
6534 let default: Value = if args.len() >= 3 {
6535 eval::eval_expr(&args[2], filtered_rows[slice[0].2], ctx)
6536 .map_err(EngineError::Eval)?
6537 } else {
6538 Value::Null
6539 };
6540 let values: Vec<Value> = slice
6541 .iter()
6542 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
6543 .collect::<Result<_, _>>()
6544 .map_err(EngineError::Eval)?;
6545 let n = slice.len();
6546 for (i, (_, _, idx)) in slice.iter().enumerate() {
6547 let signed_offset = if lower == "lag" { -offset } else { offset };
6548 let v = if ignore_nulls {
6549 let step: i64 = if signed_offset >= 0 { 1 } else { -1 };
6553 let needed: i64 = signed_offset.abs();
6554 if needed == 0 {
6555 values[i].clone()
6556 } else {
6557 let mut j: i64 = i as i64;
6558 let mut hits: i64 = 0;
6559 let mut found: Option<Value> = None;
6560 loop {
6561 j += step;
6562 if j < 0 || j >= n as i64 {
6563 break;
6564 }
6565 #[allow(clippy::cast_sign_loss)]
6566 let v = &values[j as usize];
6567 if !v.is_null() {
6568 hits += 1;
6569 if hits == needed {
6570 found = Some(v.clone());
6571 break;
6572 }
6573 }
6574 }
6575 found.unwrap_or_else(|| default.clone())
6576 }
6577 } else {
6578 let target_signed = i64::try_from(i).unwrap_or(i64::MAX) + signed_offset;
6579 if target_signed < 0
6580 || target_signed >= i64::try_from(n).unwrap_or(i64::MAX)
6581 {
6582 default.clone()
6583 } else {
6584 #[allow(clippy::cast_sign_loss)]
6585 {
6586 values[target_signed as usize].clone()
6587 }
6588 }
6589 };
6590 out_vals[*idx] = v;
6591 }
6592 Ok(())
6593 }
6594 "first_value" | "last_value" | "nth_value" => {
6595 if args.is_empty() {
6596 return Err(EngineError::Unsupported(alloc::format!(
6597 "{lower}() requires at least one argument"
6598 )));
6599 }
6600 let values: Vec<Value> = slice
6601 .iter()
6602 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
6603 .collect::<Result<_, _>>()
6604 .map_err(EngineError::Eval)?;
6605 let nth: usize = if lower == "nth_value" {
6606 if args.len() < 2 {
6607 return Err(EngineError::Unsupported(
6608 "nth_value() requires (expr, n)".into(),
6609 ));
6610 }
6611 let v = eval::eval_expr(&args[1], filtered_rows[slice[0].2], ctx)
6612 .map_err(EngineError::Eval)?;
6613 let raw = match v {
6614 Value::SmallInt(n) => i64::from(n),
6615 Value::Int(n) => i64::from(n),
6616 Value::BigInt(n) => n,
6617 _ => {
6618 return Err(EngineError::Unsupported(
6619 "nth_value() n must be integer".into(),
6620 ));
6621 }
6622 };
6623 if raw < 1 {
6624 return Err(EngineError::Unsupported(
6625 "nth_value() n must be >= 1".into(),
6626 ));
6627 }
6628 #[allow(clippy::cast_sign_loss)]
6629 {
6630 raw as usize
6631 }
6632 } else {
6633 0
6634 };
6635 let eff = effective_frame(frame, ordered)?;
6636 for i in 0..slice.len() {
6637 let (lo, hi) = frame_bounds_for_row(&eff, i, slice);
6638 let (_, _, idx) = &slice[i];
6639 let v = if lo > hi {
6640 Value::Null
6641 } else if ignore_nulls && matches!(lower.as_str(), "first_value" | "last_value") {
6642 if lower == "first_value" {
6645 (lo..=hi)
6646 .find_map(|j| {
6647 let v = &values[j];
6648 (!v.is_null()).then(|| v.clone())
6649 })
6650 .unwrap_or(Value::Null)
6651 } else {
6652 (lo..=hi)
6653 .rev()
6654 .find_map(|j| {
6655 let v = &values[j];
6656 (!v.is_null()).then(|| v.clone())
6657 })
6658 .unwrap_or(Value::Null)
6659 }
6660 } else {
6661 match lower.as_str() {
6662 "first_value" => values[lo].clone(),
6663 "last_value" => values[hi].clone(),
6664 "nth_value" => {
6665 let pos = lo + nth - 1;
6666 if pos > hi {
6667 Value::Null
6668 } else {
6669 values[pos].clone()
6670 }
6671 }
6672 _ => unreachable!(),
6673 }
6674 };
6675 out_vals[*idx] = v;
6676 }
6677 Ok(())
6678 }
6679 "ntile" => {
6680 if args.is_empty() {
6681 return Err(EngineError::Unsupported(
6682 "ntile(n) requires an integer argument".into(),
6683 ));
6684 }
6685 let v = eval::eval_expr(&args[0], filtered_rows[slice[0].2], ctx)
6686 .map_err(EngineError::Eval)?;
6687 let bucket_count: i64 = match v {
6688 Value::SmallInt(n) => i64::from(n),
6689 Value::Int(n) => i64::from(n),
6690 Value::BigInt(n) => n,
6691 _ => {
6692 return Err(EngineError::Unsupported(
6693 "ntile() argument must be integer".into(),
6694 ));
6695 }
6696 };
6697 if bucket_count < 1 {
6698 return Err(EngineError::Unsupported(
6699 "ntile() argument must be >= 1".into(),
6700 ));
6701 }
6702 #[allow(clippy::cast_sign_loss)]
6703 let buckets = bucket_count as usize;
6704 let n = slice.len();
6705 let base = n / buckets;
6708 let extras = n % buckets;
6709 let mut bucket: usize = 1;
6710 let mut remaining_in_bucket = if extras > 0 { base + 1 } else { base };
6711 let mut buckets_with_extra_remaining = extras;
6712 for (_, _, idx) in slice {
6713 if remaining_in_bucket == 0 {
6714 bucket += 1;
6715 buckets_with_extra_remaining = buckets_with_extra_remaining.saturating_sub(1);
6716 remaining_in_bucket = if buckets_with_extra_remaining > 0 {
6717 base + 1
6718 } else {
6719 base
6720 };
6721 if remaining_in_bucket == 0 {
6724 remaining_in_bucket = 1;
6725 }
6726 }
6727 out_vals[*idx] = Value::BigInt(i64::try_from(bucket).unwrap_or(i64::MAX));
6728 remaining_in_bucket -= 1;
6729 }
6730 Ok(())
6731 }
6732 "percent_rank" => {
6733 let n = slice.len();
6736 let mut prev_key: Option<&[(Value, bool)]> = None;
6737 let mut current_rank: i64 = 1;
6738 for (i, (_, okey, idx)) in slice.iter().enumerate() {
6739 if let Some(p) = prev_key
6740 && order_key_cmp(p, okey) != core::cmp::Ordering::Equal
6741 {
6742 current_rank = i64::try_from(i + 1).unwrap_or(i64::MAX);
6743 }
6744 if prev_key.is_none() {
6745 current_rank = 1;
6746 }
6747 #[allow(clippy::cast_precision_loss)]
6748 let pr = if n <= 1 {
6749 0.0
6750 } else {
6751 (current_rank - 1) as f64 / (n - 1) as f64
6752 };
6753 out_vals[*idx] = Value::Float(pr);
6754 prev_key = Some(okey.as_slice());
6755 }
6756 Ok(())
6757 }
6758 "cume_dist" => {
6759 let n = slice.len();
6761 for i in 0..slice.len() {
6763 let peer_end = peer_group_end(slice, i);
6764 #[allow(clippy::cast_precision_loss)]
6765 let cd = (peer_end + 1) as f64 / n as f64;
6766 let (_, _, idx) = &slice[i];
6767 out_vals[*idx] = Value::Float(cd);
6768 }
6769 Ok(())
6770 }
6771 other => Err(EngineError::Unsupported(alloc::format!(
6772 "window function {other:?} not supported (v4.21: row_number/rank/dense_rank/sum/avg/count/min/max/lag/lead/first_value/last_value/nth_value/ntile/percent_rank/cume_dist)"
6773 ))),
6774 }
6775}
6776
6777fn effective_frame(
6784 frame: Option<&WindowFrame>,
6785 ordered: bool,
6786) -> Result<(FrameKind, FrameBound, FrameBound), EngineError> {
6787 match frame {
6788 None => {
6789 if ordered {
6790 Ok((
6791 FrameKind::Range,
6792 FrameBound::UnboundedPreceding,
6793 FrameBound::CurrentRow,
6794 ))
6795 } else {
6796 Ok((
6797 FrameKind::Rows,
6798 FrameBound::UnboundedPreceding,
6799 FrameBound::UnboundedFollowing,
6800 ))
6801 }
6802 }
6803 Some(fr) => {
6804 let end = fr.end.clone().unwrap_or(FrameBound::CurrentRow);
6805 if matches!(fr.start, FrameBound::UnboundedFollowing)
6807 || matches!(end, FrameBound::UnboundedPreceding)
6808 {
6809 return Err(EngineError::Unsupported(alloc::format!(
6810 "invalid frame: start={:?} end={:?}",
6811 fr.start,
6812 end
6813 )));
6814 }
6815 if fr.kind == FrameKind::Range
6820 && (matches!(
6821 fr.start,
6822 FrameBound::OffsetPreceding(_) | FrameBound::OffsetFollowing(_)
6823 ) || matches!(
6824 end,
6825 FrameBound::OffsetPreceding(_) | FrameBound::OffsetFollowing(_)
6826 ))
6827 {
6828 return Err(EngineError::Unsupported(
6829 "RANGE with explicit offset bounds is not supported (v4.20: only UNBOUNDED / CURRENT ROW for RANGE)".into(),
6830 ));
6831 }
6832 Ok((fr.kind, fr.start.clone(), end))
6833 }
6834 }
6835}
6836
6837#[allow(clippy::type_complexity)]
6841fn frame_bounds_for_row(
6842 eff: &(FrameKind, FrameBound, FrameBound),
6843 i: usize,
6844 slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)],
6845) -> (usize, usize) {
6846 let (kind, start, end) = eff;
6847 let n = slice.len();
6848 let last = n.saturating_sub(1);
6849 let (mut lo, mut hi) = match kind {
6850 FrameKind::Rows => {
6851 let lo = match start {
6852 FrameBound::UnboundedPreceding => 0,
6853 FrameBound::OffsetPreceding(k) => {
6854 let k = usize::try_from(*k).unwrap_or(usize::MAX);
6855 i.saturating_sub(k)
6856 }
6857 FrameBound::CurrentRow => i,
6858 FrameBound::OffsetFollowing(k) => {
6859 let k = usize::try_from(*k).unwrap_or(usize::MAX);
6860 i.saturating_add(k).min(last)
6861 }
6862 FrameBound::UnboundedFollowing => last,
6863 };
6864 let hi = match end {
6865 FrameBound::UnboundedPreceding => 0,
6866 FrameBound::OffsetPreceding(k) => {
6867 let k = usize::try_from(*k).unwrap_or(usize::MAX);
6868 i.saturating_sub(k)
6869 }
6870 FrameBound::CurrentRow => i,
6871 FrameBound::OffsetFollowing(k) => {
6872 let k = usize::try_from(*k).unwrap_or(usize::MAX);
6873 i.saturating_add(k).min(last)
6874 }
6875 FrameBound::UnboundedFollowing => last,
6876 };
6877 (lo, hi)
6878 }
6879 FrameKind::Range => {
6880 let lo = match start {
6886 FrameBound::UnboundedPreceding => 0,
6887 FrameBound::CurrentRow => peer_group_start(slice, i),
6888 FrameBound::UnboundedFollowing => last,
6889 _ => unreachable!("offset bounds rejected for RANGE"),
6890 };
6891 let hi = match end {
6892 FrameBound::UnboundedPreceding => 0,
6893 FrameBound::CurrentRow => peer_group_end(slice, i),
6894 FrameBound::UnboundedFollowing => last,
6895 _ => unreachable!("offset bounds rejected for RANGE"),
6896 };
6897 (lo, hi)
6898 }
6899 };
6900 if hi >= n {
6901 hi = last;
6902 }
6903 if lo >= n {
6904 lo = last;
6905 }
6906 (lo, hi)
6907}
6908
6909#[allow(clippy::type_complexity)]
6913fn peer_group_start(slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)], i: usize) -> usize {
6914 let key = &slice[i].1;
6915 let mut j = i;
6916 while j > 0 && order_key_cmp(&slice[j - 1].1, key) == core::cmp::Ordering::Equal {
6917 j -= 1;
6918 }
6919 j
6920}
6921
6922#[allow(clippy::type_complexity)]
6925fn peer_group_end(slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)], i: usize) -> usize {
6926 let key = &slice[i].1;
6927 let mut j = i;
6928 while j + 1 < slice.len() && order_key_cmp(&slice[j + 1].1, key) == core::cmp::Ordering::Equal {
6929 j += 1;
6930 }
6931 j
6932}
6933
6934fn value_to_f64(v: &Value) -> Option<f64> {
6935 match v {
6936 Value::SmallInt(n) => Some(f64::from(*n)),
6937 Value::Int(n) => Some(f64::from(*n)),
6938 #[allow(clippy::cast_precision_loss)]
6939 Value::BigInt(n) => Some(*n as f64),
6940 Value::Float(x) => Some(*x),
6941 _ => None,
6942 }
6943}
6944
6945fn expr_tree_has_subquery(stmt: &SelectStatement) -> bool {
6949 let mut any = false;
6950 for item in &stmt.items {
6951 if let SelectItem::Expr { expr, .. } = item {
6952 any = any || expr_has_subquery(expr);
6953 }
6954 }
6955 if let Some(w) = &stmt.where_ {
6956 any = any || expr_has_subquery(w);
6957 }
6958 if let Some(h) = &stmt.having {
6959 any = any || expr_has_subquery(h);
6960 }
6961 for o in &stmt.order_by {
6962 any = any || expr_has_subquery(&o.expr);
6963 }
6964 for (_, peer) in &stmt.unions {
6965 any = any || expr_tree_has_subquery(peer);
6966 }
6967 any
6968}
6969
6970fn expr_has_subquery(e: &Expr) -> bool {
6971 match e {
6972 Expr::ScalarSubquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. } => true,
6973 Expr::Binary { lhs, rhs, .. } => expr_has_subquery(lhs) || expr_has_subquery(rhs),
6974 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6975 expr_has_subquery(expr)
6976 }
6977 Expr::FunctionCall { args, .. } => args.iter().any(expr_has_subquery),
6978 Expr::Like { expr, pattern, .. } => expr_has_subquery(expr) || expr_has_subquery(pattern),
6979 Expr::Extract { source, .. } => expr_has_subquery(source),
6980 Expr::WindowFunction {
6981 args,
6982 partition_by,
6983 order_by,
6984 ..
6985 } => {
6986 args.iter().any(expr_has_subquery)
6987 || partition_by.iter().any(expr_has_subquery)
6988 || order_by.iter().any(|(e, _)| expr_has_subquery(e))
6989 }
6990 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => false,
6991 Expr::Array(items) => items.iter().any(expr_has_subquery),
6992 Expr::ArraySubscript { target, index } => {
6993 expr_has_subquery(target) || expr_has_subquery(index)
6994 }
6995 Expr::AnyAll { expr, array, .. } => {
6996 expr_has_subquery(expr) || expr_has_subquery(array)
6997 }
6998 }
6999}
7000
7001fn value_to_literal_expr(v: Value) -> Result<Expr, EngineError> {
7008 let lit = match v {
7009 Value::Null => Literal::Null,
7010 Value::SmallInt(n) => Literal::Integer(i64::from(n)),
7011 Value::Int(n) => Literal::Integer(i64::from(n)),
7012 Value::BigInt(n) => Literal::Integer(n),
7013 Value::Float(x) => Literal::Float(x),
7014 Value::Text(s) | Value::Json(s) => Literal::String(s),
7015 Value::Bool(b) => Literal::Bool(b),
7016 other => {
7017 return Err(EngineError::Unsupported(alloc::format!(
7018 "subquery result type {:?} not yet materialisable; cast to text or integer in the inner SELECT",
7019 other.data_type()
7020 )));
7021 }
7022 };
7023 Ok(Expr::Literal(lit))
7024}
7025
7026fn substitute_placeholders(stmt: &mut Statement, params: &[Value]) -> Result<(), EngineError> {
7037 match stmt {
7038 Statement::Select(s) => substitute_select(s, params)?,
7039 Statement::Insert(ins) => {
7040 for row in &mut ins.rows {
7041 for e in row {
7042 substitute_expr(e, params)?;
7043 }
7044 }
7045 }
7046 Statement::Update(u) => {
7047 for (_, e) in &mut u.assignments {
7048 substitute_expr(e, params)?;
7049 }
7050 if let Some(w) = &mut u.where_ {
7051 substitute_expr(w, params)?;
7052 }
7053 }
7054 Statement::Delete(d) => {
7055 if let Some(w) = &mut d.where_ {
7056 substitute_expr(w, params)?;
7057 }
7058 }
7059 Statement::Explain(e) => substitute_select(&mut e.inner, params)?,
7060 _ => {}
7063 }
7064 Ok(())
7065}
7066
7067fn substitute_select(
7068 s: &mut SelectStatement,
7069 params: &[Value],
7070) -> Result<(), EngineError> {
7071 for item in &mut s.items {
7072 if let SelectItem::Expr { expr, .. } = item {
7073 substitute_expr(expr, params)?;
7074 }
7075 }
7076 if let Some(w) = &mut s.where_ {
7077 substitute_expr(w, params)?;
7078 }
7079 if let Some(gs) = &mut s.group_by {
7080 for g in gs {
7081 substitute_expr(g, params)?;
7082 }
7083 }
7084 if let Some(h) = &mut s.having {
7085 substitute_expr(h, params)?;
7086 }
7087 for o in &mut s.order_by {
7088 substitute_expr(&mut o.expr, params)?;
7089 }
7090 for (_, peer) in &mut s.unions {
7091 substitute_select(peer, params)?;
7092 }
7093 if let Some(le) = s.limit {
7098 s.limit = Some(resolve_limit_placeholder(le, params)?);
7099 }
7100 if let Some(le) = s.offset {
7101 s.offset = Some(resolve_limit_placeholder(le, params)?);
7102 }
7103 Ok(())
7104}
7105
7106fn resolve_limit_placeholder(
7107 le: spg_sql::ast::LimitExpr,
7108 params: &[Value],
7109) -> Result<spg_sql::ast::LimitExpr, EngineError> {
7110 use spg_sql::ast::LimitExpr;
7111 match le {
7112 LimitExpr::Literal(_) => Ok(le),
7113 LimitExpr::Placeholder(n) => {
7114 let idx = usize::from(n).saturating_sub(1);
7115 let v = params.get(idx).ok_or_else(|| {
7116 EngineError::Eval(EvalError::PlaceholderOutOfRange {
7117 n,
7118 bound: u16::try_from(params.len()).unwrap_or(u16::MAX),
7119 })
7120 })?;
7121 let int = match v {
7122 Value::SmallInt(x) => Some(i64::from(*x)),
7123 Value::Int(x) => Some(i64::from(*x)),
7124 Value::BigInt(x) => Some(*x),
7125 _ => None,
7126 }
7127 .ok_or_else(|| {
7128 EngineError::Unsupported(alloc::format!(
7129 "LIMIT/OFFSET ${n} bound to non-integer {v:?}"
7130 ))
7131 })?;
7132 if int < 0 {
7133 return Err(EngineError::Unsupported(alloc::format!(
7134 "LIMIT/OFFSET ${n} bound to negative value {int}"
7135 )));
7136 }
7137 let bounded = u32::try_from(int).map_err(|_| {
7138 EngineError::Unsupported(alloc::format!(
7139 "LIMIT/OFFSET ${n} value {int} exceeds u32 range"
7140 ))
7141 })?;
7142 Ok(LimitExpr::Literal(bounded))
7143 }
7144 }
7145}
7146
7147fn substitute_expr(e: &mut Expr, params: &[Value]) -> Result<(), EngineError> {
7148 if let Expr::Placeholder(n) = e {
7149 let idx = usize::from(*n).saturating_sub(1);
7150 let v = params.get(idx).ok_or_else(|| {
7151 EngineError::Eval(EvalError::PlaceholderOutOfRange {
7152 n: *n,
7153 bound: u16::try_from(params.len()).unwrap_or(u16::MAX),
7154 })
7155 })?;
7156 *e = Expr::Literal(value_to_literal(v.clone()));
7157 return Ok(());
7158 }
7159 match e {
7160 Expr::Binary { lhs, rhs, .. } => {
7161 substitute_expr(lhs, params)?;
7162 substitute_expr(rhs, params)?;
7163 }
7164 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7165 substitute_expr(expr, params)?;
7166 }
7167 Expr::FunctionCall { args, .. } => {
7168 for a in args {
7169 substitute_expr(a, params)?;
7170 }
7171 }
7172 Expr::Like { expr, pattern, .. } => {
7173 substitute_expr(expr, params)?;
7174 substitute_expr(pattern, params)?;
7175 }
7176 Expr::Extract { source, .. } => substitute_expr(source, params)?,
7177 Expr::ScalarSubquery(s) => substitute_select(s, params)?,
7178 Expr::Exists { subquery, .. } => substitute_select(subquery, params)?,
7179 Expr::InSubquery { expr, subquery, .. } => {
7180 substitute_expr(expr, params)?;
7181 substitute_select(subquery, params)?;
7182 }
7183 Expr::WindowFunction {
7184 args,
7185 partition_by,
7186 order_by,
7187 ..
7188 } => {
7189 for a in args {
7190 substitute_expr(a, params)?;
7191 }
7192 for p in partition_by {
7193 substitute_expr(p, params)?;
7194 }
7195 for (e, _) in order_by {
7196 substitute_expr(e, params)?;
7197 }
7198 }
7199 Expr::Literal(_) | Expr::Column(_) => {}
7200 Expr::Placeholder(_) => unreachable!("Placeholder handled at top of fn"),
7202 Expr::Array(items) => {
7203 for elem in items {
7204 substitute_expr(elem, params)?;
7205 }
7206 }
7207 Expr::ArraySubscript { target, index } => {
7208 substitute_expr(target, params)?;
7209 substitute_expr(index, params)?;
7210 }
7211 Expr::AnyAll { expr, array, .. } => {
7212 substitute_expr(expr, params)?;
7213 substitute_expr(array, params)?;
7214 }
7215 }
7216 Ok(())
7217}
7218
7219fn sort_values_for_histogram(a: &Value, b: &Value) -> core::cmp::Ordering {
7237 use core::cmp::Ordering;
7238 match (a, b) {
7239 (Value::SmallInt(a), Value::SmallInt(b)) => a.cmp(b),
7240 (Value::Int(a), Value::Int(b)) => a.cmp(b),
7241 (Value::BigInt(a), Value::BigInt(b)) => a.cmp(b),
7242 (Value::SmallInt(a), Value::Int(b)) => i32::from(*a).cmp(b),
7243 (Value::Int(a), Value::SmallInt(b)) => a.cmp(&i32::from(*b)),
7244 (Value::Int(a), Value::BigInt(b)) => i64::from(*a).cmp(b),
7245 (Value::BigInt(a), Value::Int(b)) => a.cmp(&i64::from(*b)),
7246 (Value::SmallInt(a), Value::BigInt(b)) => i64::from(*a).cmp(b),
7247 (Value::BigInt(a), Value::SmallInt(b)) => a.cmp(&i64::from(*b)),
7248 (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(Ordering::Equal),
7249 (Value::Text(a), Value::Text(b)) | (Value::Json(a), Value::Json(b)) => a.cmp(b),
7250 (Value::Bool(a), Value::Bool(b)) => a.cmp(b),
7251 (Value::Date(a), Value::Date(b)) => a.cmp(b),
7252 (Value::Timestamp(a), Value::Timestamp(b)) => a.cmp(b),
7253 (Value::SmallInt(n), Value::Float(x)) => {
7255 (f64::from(*n)).partial_cmp(x).unwrap_or(Ordering::Equal)
7256 }
7257 (Value::Float(x), Value::SmallInt(n)) => {
7258 x.partial_cmp(&f64::from(*n)).unwrap_or(Ordering::Equal)
7259 }
7260 (Value::Int(n), Value::Float(x)) => {
7261 (f64::from(*n)).partial_cmp(x).unwrap_or(Ordering::Equal)
7262 }
7263 (Value::Float(x), Value::Int(n)) => {
7264 x.partial_cmp(&f64::from(*n)).unwrap_or(Ordering::Equal)
7265 }
7266 (Value::BigInt(n), Value::Float(x)) => {
7267 #[allow(clippy::cast_precision_loss)]
7268 let nf = *n as f64;
7269 nf.partial_cmp(x).unwrap_or(Ordering::Equal)
7270 }
7271 (Value::Float(x), Value::BigInt(n)) => {
7272 #[allow(clippy::cast_precision_loss)]
7273 let nf = *n as f64;
7274 x.partial_cmp(&nf).unwrap_or(Ordering::Equal)
7275 }
7276 _ => canonical_value_repr(a).cmp(&canonical_value_repr(b)),
7279 }
7280}
7281
7282fn render_histogram_bounds(bounds: &[alloc::string::String]) -> alloc::string::String {
7289 let mut out = alloc::string::String::with_capacity(bounds.len() * 8 + 2);
7290 out.push('[');
7291 for (i, b) in bounds.iter().enumerate() {
7292 if i > 0 {
7293 out.push_str(", ");
7294 }
7295 let needs_quote = b.contains([',', '[', ']', '"']) || b.is_empty();
7296 if needs_quote {
7297 out.push('"');
7298 for ch in b.chars() {
7299 if ch == '"' || ch == '\\' {
7300 out.push('\\');
7301 }
7302 out.push(ch);
7303 }
7304 out.push('"');
7305 } else {
7306 out.push_str(b);
7307 }
7308 }
7309 out.push(']');
7310 out
7311}
7312
7313pub(crate) fn canonical_value_repr(v: &Value) -> alloc::string::String {
7323 match v {
7324 Value::Null => "NULL".to_string(),
7325 Value::SmallInt(n) => alloc::format!("{n}"),
7326 Value::Int(n) => alloc::format!("{n}"),
7327 Value::BigInt(n) => alloc::format!("{n}"),
7328 Value::Float(x) => alloc::format!("{x:?}"),
7329 Value::Text(s) | Value::Json(s) => s.clone(),
7330 Value::Bool(b) => if *b { "t" } else { "f" }.to_string(),
7331 Value::Date(d) => eval::format_date(*d),
7332 Value::Timestamp(t) => eval::format_timestamp(*t),
7333 Value::Interval { months, micros } => eval::format_interval(*months, *micros),
7334 Value::Numeric { scaled, scale } => eval::format_numeric(*scaled, *scale),
7335 Value::Vector(_) | Value::Sq8Vector(_) | Value::HalfVector(_) => {
7336 alloc::format!("{v:?}")
7340 }
7341 _ => alloc::format!("{v:?}"),
7345 }
7346}
7347
7348const fn is_internal_table_name(_name: &str) -> bool {
7355 false
7356}
7357
7358fn value_to_literal(v: Value) -> Literal {
7359 match v {
7360 Value::Null => Literal::Null,
7361 Value::SmallInt(n) => Literal::Integer(i64::from(n)),
7362 Value::Int(n) => Literal::Integer(i64::from(n)),
7363 Value::BigInt(n) => Literal::Integer(n),
7364 Value::Float(x) => Literal::Float(x),
7365 Value::Text(s) | Value::Json(s) => Literal::String(s),
7366 Value::Bool(b) => Literal::Bool(b),
7367 Value::Vector(v) => Literal::Vector(v),
7368 Value::Numeric { scaled, scale } => {
7369 Literal::String(eval::format_numeric(scaled, scale))
7370 }
7371 Value::Date(d) => Literal::String(eval::format_date(d)),
7372 Value::Timestamp(t) => Literal::String(eval::format_timestamp(t)),
7373 Value::Interval { months, micros } => Literal::Interval {
7374 months,
7375 micros,
7376 text: eval::format_interval(months, micros),
7377 },
7378 Value::Sq8Vector(q) => Literal::Vector(spg_storage::quantize::dequantize(&q)),
7381 Value::HalfVector(h) => Literal::Vector(h.to_f32_vec()),
7382 v => Literal::String(alloc::format!("{v:?}")),
7386 }
7387}
7388
7389fn rewrite_clock_calls(stmt: &mut Statement, now_micros: Option<i64>) {
7390 let Some(now) = now_micros else {
7391 return;
7392 };
7393 match stmt {
7394 Statement::Select(s) => rewrite_select_clock(s, now),
7395 Statement::Insert(ins) => {
7396 for row in &mut ins.rows {
7397 for e in row {
7398 rewrite_expr_clock(e, now);
7399 }
7400 }
7401 }
7402 _ => {}
7403 }
7404}
7405
7406fn rewrite_select_clock(s: &mut SelectStatement, now: i64) {
7407 for item in &mut s.items {
7408 if let SelectItem::Expr { expr, .. } = item {
7409 rewrite_expr_clock(expr, now);
7410 }
7411 }
7412 if let Some(w) = &mut s.where_ {
7413 rewrite_expr_clock(w, now);
7414 }
7415 if let Some(gs) = &mut s.group_by {
7416 for g in gs {
7417 rewrite_expr_clock(g, now);
7418 }
7419 }
7420 if let Some(h) = &mut s.having {
7421 rewrite_expr_clock(h, now);
7422 }
7423 for o in &mut s.order_by {
7424 rewrite_expr_clock(&mut o.expr, now);
7425 }
7426 for (_, peer) in &mut s.unions {
7427 rewrite_select_clock(peer, now);
7428 }
7429}
7430
7431fn rewrite_expr_clock(e: &mut Expr, now: i64) {
7439 if let Some(replacement) = clock_replacement_for(e, now) {
7443 *e = replacement;
7444 return;
7445 }
7446 match e {
7447 Expr::Binary { lhs, rhs, .. } => {
7448 rewrite_expr_clock(lhs, now);
7449 rewrite_expr_clock(rhs, now);
7450 }
7451 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7452 rewrite_expr_clock(expr, now);
7453 }
7454 Expr::FunctionCall { args, .. } => {
7455 for a in args {
7456 rewrite_expr_clock(a, now);
7457 }
7458 }
7459 Expr::Like { expr, pattern, .. } => {
7460 rewrite_expr_clock(expr, now);
7461 rewrite_expr_clock(pattern, now);
7462 }
7463 Expr::Extract { source, .. } => rewrite_expr_clock(source, now),
7464 Expr::ScalarSubquery(s) => rewrite_select_clock(s, now),
7468 Expr::Exists { subquery, .. } => rewrite_select_clock(subquery, now),
7469 Expr::InSubquery { expr, subquery, .. } => {
7470 rewrite_expr_clock(expr, now);
7471 rewrite_select_clock(subquery, now);
7472 }
7473 Expr::WindowFunction {
7476 args,
7477 partition_by,
7478 order_by,
7479 ..
7480 } => {
7481 for a in args {
7482 rewrite_expr_clock(a, now);
7483 }
7484 for p in partition_by {
7485 rewrite_expr_clock(p, now);
7486 }
7487 for (e, _) in order_by {
7488 rewrite_expr_clock(e, now);
7489 }
7490 }
7491 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
7492 Expr::Array(items) => {
7493 for elem in items {
7494 rewrite_expr_clock(elem, now);
7495 }
7496 }
7497 Expr::ArraySubscript { target, index } => {
7498 rewrite_expr_clock(target, now);
7499 rewrite_expr_clock(index, now);
7500 }
7501 Expr::AnyAll { expr, array, .. } => {
7502 rewrite_expr_clock(expr, now);
7503 rewrite_expr_clock(array, now);
7504 }
7505 }
7506}
7507
7508fn clock_replacement_for(e: &Expr, now: i64) -> Option<Expr> {
7515 let (kind, name) = match e {
7516 Expr::FunctionCall { name, args } if args.is_empty() => (ClockSite::Fn, name.as_str()),
7517 Expr::Column(c) if c.qualifier.is_none() => (ClockSite::BareIdent, c.name.as_str()),
7518 _ => return None,
7519 };
7520 let matched = match name.len() {
7523 3 if kind == ClockSite::Fn && name.eq_ignore_ascii_case("now") => Some(true),
7524 12 if name.eq_ignore_ascii_case("current_date") => Some(false),
7525 17 if name.eq_ignore_ascii_case("current_timestamp") => Some(true),
7526 _ => None,
7527 };
7528 let is_timestamp = matched?;
7529 let payload = if is_timestamp {
7530 now
7531 } else {
7532 now.div_euclid(86_400_000_000)
7533 };
7534 let target = if is_timestamp {
7535 spg_sql::ast::CastTarget::Timestamp
7536 } else {
7537 spg_sql::ast::CastTarget::Date
7538 };
7539 Some(Expr::Cast {
7540 expr: alloc::boxed::Box::new(Expr::Literal(spg_sql::ast::Literal::Integer(payload))),
7541 target,
7542 })
7543}
7544
7545#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7546enum ClockSite {
7547 Fn,
7548 BareIdent,
7549}
7550
7551fn expand_group_by_all(s: &mut SelectStatement) {
7562 if !s.group_by_all {
7563 for (_, peer) in &mut s.unions {
7564 expand_group_by_all(peer);
7565 }
7566 return;
7567 }
7568 let mut groups: Vec<Expr> = Vec::new();
7569 for item in &s.items {
7570 if let SelectItem::Expr { expr, .. } = item
7571 && !aggregate::contains_aggregate(expr)
7572 {
7573 groups.push(expr.clone());
7574 }
7575 }
7576 s.group_by = Some(groups);
7577 s.group_by_all = false;
7578 for (_, peer) in &mut s.unions {
7579 expand_group_by_all(peer);
7580 }
7581}
7582
7583fn resolve_order_by_position(s: &mut SelectStatement) {
7584 for order in &mut s.order_by {
7589 match &order.expr {
7590 Expr::Literal(Literal::Integer(n)) if *n >= 1 => {
7591 if let Ok(idx_one_based) = usize::try_from(*n) {
7592 let idx = idx_one_based - 1;
7593 if idx < s.items.len()
7594 && let SelectItem::Expr { expr, .. } = &s.items[idx]
7595 {
7596 order.expr = expr.clone();
7597 }
7598 }
7599 }
7600 Expr::Column(c) if c.qualifier.is_none() => {
7601 for item in &s.items {
7603 if let SelectItem::Expr {
7604 expr,
7605 alias: Some(a),
7606 } = item
7607 && a == &c.name
7608 {
7609 order.expr = expr.clone();
7610 break;
7611 }
7612 }
7613 }
7614 _ => {}
7615 }
7616 }
7617 for (_, peer) in &mut s.unions {
7618 resolve_order_by_position(peer);
7619 }
7620}
7621
7622fn partial_sort_tagged(
7635 tagged: &mut Vec<(Vec<f64>, Row)>,
7636 keep: Option<usize>,
7637 descs: &[bool],
7638) {
7639 let cmp = |a: &(Vec<f64>, Row), b: &(Vec<f64>, Row)| cmp_multi_key(&a.0, &b.0, descs);
7640 match keep {
7641 Some(k) if k < tagged.len() && k > 0 => {
7642 let pivot = k - 1;
7643 tagged.select_nth_unstable_by(pivot, cmp);
7644 tagged[..k].sort_by(cmp);
7645 tagged.truncate(k);
7646 }
7647 _ => {
7648 tagged.sort_by(cmp);
7649 }
7650 }
7651}
7652
7653fn sort_by_keys(tagged: &mut [(Vec<f64>, Row)], descs: &[bool]) {
7654 tagged.sort_by(|a, b| cmp_multi_key(&a.0, &b.0, descs));
7655}
7656
7657fn cmp_multi_key(a: &[f64], b: &[f64], descs: &[bool]) -> core::cmp::Ordering {
7661 use core::cmp::Ordering;
7662 for (i, (ka, kb)) in a.iter().zip(b.iter()).enumerate() {
7663 let ord = ka.partial_cmp(kb).unwrap_or(Ordering::Equal);
7664 let ord = if descs.get(i).copied().unwrap_or(false) {
7665 ord.reverse()
7666 } else {
7667 ord
7668 };
7669 if ord != Ordering::Equal {
7670 return ord;
7671 }
7672 }
7673 Ordering::Equal
7674}
7675
7676fn build_order_keys(
7679 order_by: &[OrderBy],
7680 row: &Row,
7681 ctx: &EvalContext,
7682) -> Result<Vec<f64>, EngineError> {
7683 let mut keys = Vec::with_capacity(order_by.len());
7684 for o in order_by {
7685 let v = eval::eval_expr(&o.expr, row, ctx)?;
7686 keys.push(value_to_order_key(&v)?);
7687 }
7688 Ok(keys)
7689}
7690
7691fn apply_offset_and_limit(rows: &mut Vec<Row>, offset: Option<u32>, limit: Option<u32>) {
7695 if let Some(off) = offset {
7696 let off = off as usize;
7697 if off >= rows.len() {
7698 rows.clear();
7699 } else {
7700 rows.drain(..off);
7701 }
7702 }
7703 if let Some(n) = limit {
7704 rows.truncate(n as usize);
7705 }
7706}
7707
7708fn resolve_foreign_key(
7722 local_table_name: &str,
7723 local_cols: &[ColumnSchema],
7724 fk: spg_sql::ast::ForeignKeyConstraint,
7725 catalog: &Catalog,
7726) -> Result<spg_storage::ForeignKeyConstraint, EngineError> {
7727 let mut local_columns = Vec::with_capacity(fk.columns.len());
7729 for name in &fk.columns {
7730 let pos = local_cols
7731 .iter()
7732 .position(|c| c.name == *name)
7733 .ok_or_else(|| {
7734 EngineError::Unsupported(alloc::format!(
7735 "FOREIGN KEY references unknown local column {name:?}"
7736 ))
7737 })?;
7738 local_columns.push(pos);
7739 }
7740 let is_self_ref = fk.parent_table == local_table_name;
7744 let (parent_cols_for_lookup, parent_table_str): (&[ColumnSchema], &str) = if is_self_ref {
7745 (local_cols, local_table_name)
7746 } else {
7747 let parent_table = catalog.get(&fk.parent_table).ok_or_else(|| {
7748 EngineError::Storage(StorageError::TableNotFound {
7749 name: fk.parent_table.clone(),
7750 })
7751 })?;
7752 (parent_table.schema().columns.as_slice(), fk.parent_table.as_str())
7753 };
7754 let parent_columns: Vec<usize> = if fk.parent_columns.is_empty() {
7759 if fk.columns.len() != 1 {
7760 return Err(EngineError::Unsupported(
7761 "composite FOREIGN KEY without explicit parent column list is not supported \
7762 — list the parent columns explicitly"
7763 .into(),
7764 ));
7765 }
7766 let pos = pick_pk_index_column(catalog, parent_table_str, is_self_ref, local_cols)
7768 .ok_or_else(|| {
7769 EngineError::Unsupported(alloc::format!(
7770 "parent table {parent_table_str:?} has no PRIMARY-key / UNIQUE BTree index \
7771 to default the FOREIGN KEY against"
7772 ))
7773 })?;
7774 alloc::vec![pos]
7775 } else {
7776 let mut out = Vec::with_capacity(fk.parent_columns.len());
7777 for name in &fk.parent_columns {
7778 let pos = parent_cols_for_lookup
7779 .iter()
7780 .position(|c| c.name == *name)
7781 .ok_or_else(|| {
7782 EngineError::Unsupported(alloc::format!(
7783 "FOREIGN KEY references unknown parent column \
7784 {name:?} on table {parent_table_str:?}"
7785 ))
7786 })?;
7787 out.push(pos);
7788 }
7789 out
7790 };
7791 if parent_columns.len() != local_columns.len() {
7792 return Err(EngineError::Unsupported(alloc::format!(
7793 "FOREIGN KEY arity mismatch: {} local columns vs {} parent columns",
7794 local_columns.len(),
7795 parent_columns.len()
7796 )));
7797 }
7798 if !is_self_ref {
7808 let parent_table = catalog
7809 .get(&fk.parent_table)
7810 .expect("checked above");
7811 let primary_parent_col = parent_columns[0];
7812 let has_btree = parent_table.schema().columns.get(primary_parent_col).is_some()
7813 && parent_table
7814 .indices()
7815 .iter()
7816 .any(|idx| {
7817 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
7818 && idx.column_position == primary_parent_col
7819 && idx.partial_predicate.is_none()
7820 });
7821 if !has_btree {
7822 return Err(EngineError::Unsupported(alloc::format!(
7823 "FOREIGN KEY parent column on {:?} is not covered by an unconditional BTree \
7824 index — create one with `CREATE INDEX ... ON {} ({})` first",
7825 parent_table_str,
7826 parent_table_str,
7827 parent_table.schema().columns[primary_parent_col].name,
7828 )));
7829 }
7830 }
7831 let on_delete = fk_action_sql_to_storage(fk.on_delete);
7832 let on_update = fk_action_sql_to_storage(fk.on_update);
7833 Ok(spg_storage::ForeignKeyConstraint {
7834 name: fk.name,
7835 local_columns,
7836 parent_table: fk.parent_table,
7837 parent_columns,
7838 on_delete,
7839 on_update,
7840 })
7841}
7842
7843fn pick_pk_index_column(
7849 catalog: &Catalog,
7850 parent_name: &str,
7851 is_self_ref: bool,
7852 local_cols: &[ColumnSchema],
7853) -> Option<usize> {
7854 if is_self_ref {
7855 let _ = local_cols;
7859 return Some(0);
7860 }
7861 let parent = catalog.get(parent_name)?;
7862 parent.indices().iter().find_map(|idx| {
7863 if matches!(idx.kind, spg_storage::IndexKind::BTree(_))
7864 && idx.partial_predicate.is_none()
7865 && idx.included_columns.is_empty()
7866 && idx.expression.is_none()
7867 {
7868 Some(idx.column_position)
7869 } else {
7870 None
7871 }
7872 })
7873}
7874
7875fn resolve_on_conflict_columns(
7882 catalog: &Catalog,
7883 table_name: &str,
7884 target: &[String],
7885) -> Result<Vec<usize>, EngineError> {
7886 let table = catalog.get(table_name).ok_or_else(|| {
7887 EngineError::Storage(StorageError::TableNotFound {
7888 name: table_name.into(),
7889 })
7890 })?;
7891 if target.is_empty() {
7892 let pos = table
7893 .indices()
7894 .iter()
7895 .find_map(|idx| {
7896 if matches!(idx.kind, spg_storage::IndexKind::BTree(_))
7897 && idx.partial_predicate.is_none()
7898 && idx.included_columns.is_empty()
7899 && idx.expression.is_none()
7900 {
7901 Some(idx.column_position)
7902 } else {
7903 None
7904 }
7905 })
7906 .ok_or_else(|| {
7907 EngineError::Unsupported(alloc::format!(
7908 "ON CONFLICT without target requires a UNIQUE BTree index on {table_name:?}"
7909 ))
7910 })?;
7911 return Ok(alloc::vec![pos]);
7912 }
7913 let mut out = Vec::with_capacity(target.len());
7914 for name in target {
7915 let pos = table
7916 .schema()
7917 .columns
7918 .iter()
7919 .position(|c| c.name == *name)
7920 .ok_or_else(|| {
7921 EngineError::Unsupported(alloc::format!(
7922 "ON CONFLICT target column {name:?} not found on {table_name:?}"
7923 ))
7924 })?;
7925 out.push(pos);
7926 }
7927 Ok(out)
7928}
7929
7930fn on_conflict_key_exists(
7933 catalog: &Catalog,
7934 table_name: &str,
7935 column_pos: usize,
7936 key: &Value,
7937) -> bool {
7938 let Some(table) = catalog.get(table_name) else {
7939 return false;
7940 };
7941 let Some(idx_key) = spg_storage::IndexKey::from_value(key) else {
7942 return false;
7943 };
7944 table.indices().iter().any(|idx| {
7945 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
7946 && idx.column_position == column_pos
7947 && idx.partial_predicate.is_none()
7948 && !idx.lookup_eq(&idx_key).is_empty()
7949 })
7950}
7951
7952fn lookup_row_position_by_keys(
7958 catalog: &Catalog,
7959 table_name: &str,
7960 column_positions: &[usize],
7961 key: &[&Value],
7962) -> Option<usize> {
7963 let table = catalog.get(table_name)?;
7964 table.rows().iter().position(|r| {
7965 column_positions
7966 .iter()
7967 .enumerate()
7968 .all(|(i, &pos)| r.values.get(pos) == Some(key[i]))
7969 })
7970}
7971
7972fn on_conflict_keys_exist(
7977 catalog: &Catalog,
7978 table_name: &str,
7979 column_positions: &[usize],
7980 key: &[&Value],
7981) -> bool {
7982 if column_positions.len() == 1 {
7983 return on_conflict_key_exists(
7984 catalog,
7985 table_name,
7986 column_positions[0],
7987 key[0],
7988 );
7989 }
7990 let Some(table) = catalog.get(table_name) else {
7991 return false;
7992 };
7993 table.rows().iter().any(|r| {
7994 column_positions
7995 .iter()
7996 .enumerate()
7997 .all(|(i, &pos)| r.values.get(pos) == Some(key[i]))
7998 })
7999}
8000
8001fn apply_on_conflict_assignments(
8014 catalog: &Catalog,
8015 table_name: &str,
8016 target_pos: usize,
8017 incoming: &[Value],
8018 assignments: &[(String, Expr)],
8019 where_: Option<&Expr>,
8020) -> Result<Option<Vec<Value>>, EngineError> {
8021 let table = catalog.get(table_name).ok_or_else(|| {
8022 EngineError::Storage(StorageError::TableNotFound {
8023 name: table_name.into(),
8024 })
8025 })?;
8026 let schema_cols = table.schema().columns.clone();
8027 let existing = table
8028 .rows()
8029 .get(target_pos)
8030 .ok_or_else(|| {
8031 EngineError::Unsupported(alloc::format!(
8032 "ON CONFLICT DO UPDATE: row position {target_pos} out of bounds on {table_name:?}"
8033 ))
8034 })?
8035 .clone();
8036 let ctx = eval::EvalContext::new(&schema_cols, Some(table_name));
8037 if let Some(w) = where_ {
8039 let pred = w.clone();
8040 let pred = substitute_excluded_refs(pred, &schema_cols, incoming);
8041 let v = eval::eval_expr(&pred, &existing, &ctx)?;
8042 if !matches!(v, Value::Bool(true)) {
8043 return Ok(None);
8044 }
8045 }
8046 let mut new_values = existing.values.clone();
8047 for (col_name, expr) in assignments {
8048 let target_idx = schema_cols
8049 .iter()
8050 .position(|c| c.name == *col_name)
8051 .ok_or_else(|| {
8052 EngineError::Eval(EvalError::ColumnNotFound {
8053 name: col_name.clone(),
8054 })
8055 })?;
8056 let sub = substitute_excluded_refs(expr.clone(), &schema_cols, incoming);
8057 let v = eval::eval_expr(&sub, &existing, &ctx)?;
8058 new_values[target_idx] =
8059 coerce_value(v, schema_cols[target_idx].ty, col_name, target_idx)?;
8060 }
8061 Ok(Some(new_values))
8062}
8063
8064fn substitute_excluded_refs(
8069 expr: Expr,
8070 schema_cols: &[ColumnSchema],
8071 incoming: &[Value],
8072) -> Expr {
8073 use spg_sql::ast::ColumnName;
8074 match expr {
8075 Expr::Column(ColumnName { qualifier, name })
8076 if qualifier
8077 .as_deref()
8078 .is_some_and(|q| q.eq_ignore_ascii_case("excluded")) =>
8079 {
8080 let pos = schema_cols.iter().position(|c| c.name == name);
8081 match pos {
8082 Some(p) => {
8083 let v = incoming.get(p).cloned().unwrap_or(Value::Null);
8084 value_to_literal_expr(v).unwrap_or_else(|_| {
8085 Expr::Literal(spg_sql::ast::Literal::Null)
8086 })
8087 }
8088 None => Expr::Column(ColumnName { qualifier, name }),
8089 }
8090 }
8091 Expr::Binary { op, lhs, rhs } => Expr::Binary {
8092 op,
8093 lhs: Box::new(substitute_excluded_refs(*lhs, schema_cols, incoming)),
8094 rhs: Box::new(substitute_excluded_refs(*rhs, schema_cols, incoming)),
8095 },
8096 Expr::Unary { op, expr } => Expr::Unary {
8097 op,
8098 expr: Box::new(substitute_excluded_refs(*expr, schema_cols, incoming)),
8099 },
8100 Expr::FunctionCall { name, args } => Expr::FunctionCall {
8101 name,
8102 args: args
8103 .into_iter()
8104 .map(|a| substitute_excluded_refs(a, schema_cols, incoming))
8105 .collect(),
8106 },
8107 other => other,
8108 }
8109}
8110
8111fn enforce_uniqueness_inserts(
8134 catalog: &Catalog,
8135 child_table: &str,
8136 constraints: &[spg_storage::UniquenessConstraint],
8137 rows: &[Vec<Value>],
8138) -> Result<(), EngineError> {
8139 if constraints.is_empty() {
8140 return Ok(());
8141 }
8142 let table = catalog.get(child_table).ok_or_else(|| {
8143 EngineError::Storage(StorageError::TableNotFound {
8144 name: child_table.into(),
8145 })
8146 })?;
8147 for uc in constraints {
8148 for (batch_idx, row_values) in rows.iter().enumerate() {
8149 let key: Vec<&Value> = uc.columns.iter().map(|&i| &row_values[i]).collect();
8150 let has_null = key.iter().any(|v| matches!(v, Value::Null));
8151 if has_null {
8152 continue;
8153 }
8154 let collides_in_table = table.rows().iter().any(|prow| {
8156 uc.columns
8157 .iter()
8158 .enumerate()
8159 .all(|(i, &p)| prow.values.get(p) == Some(key[i]))
8160 });
8161 let collides_in_batch = rows[..batch_idx].iter().any(|earlier| {
8163 uc.columns
8164 .iter()
8165 .enumerate()
8166 .all(|(i, &p)| earlier.get(p) == Some(key[i]))
8167 });
8168 if collides_in_table || collides_in_batch {
8169 let kind = if uc.is_primary_key { "PRIMARY KEY" } else { "UNIQUE" };
8170 let col_names: Vec<String> = uc
8171 .columns
8172 .iter()
8173 .map(|&i| table.schema().columns[i].name.clone())
8174 .collect();
8175 return Err(EngineError::Unsupported(alloc::format!(
8176 "{kind} violation on {child_table:?} columns {col_names:?}: \
8177 row #{batch_idx} duplicates an existing key"
8178 )));
8179 }
8180 }
8181 }
8182 Ok(())
8183}
8184
8185fn predicate_truthy(v: &spg_storage::Value) -> bool {
8193 use spg_storage::Value as V;
8194 match v {
8195 V::Bool(b) => *b,
8196 V::Int(n) => *n != 0,
8197 V::BigInt(n) => *n != 0,
8198 V::SmallInt(n) => *n != 0,
8199 _ => false,
8200 }
8201}
8202
8203fn check_existing_unique_violation(
8208 idx: &spg_storage::Index,
8209 schema: &spg_storage::TableSchema,
8210 rows: &[spg_storage::Row],
8211) -> Result<(), EngineError> {
8212 let predicate_expr = match idx.partial_predicate.as_deref() {
8213 Some(s) => Some(spg_sql::parser::parse_expression(s).map_err(|e| {
8214 EngineError::Unsupported(alloc::format!(
8215 "stored partial predicate {s:?} failed to re-parse: {e:?}"
8216 ))
8217 })?),
8218 None => None,
8219 };
8220 let ctx = eval::EvalContext::new(&schema.columns, None);
8221 let key_positions = unique_key_positions(idx);
8222 let mut seen: alloc::vec::Vec<alloc::vec::Vec<spg_storage::Value>> = alloc::vec::Vec::new();
8223 for row in rows {
8224 if let Some(expr) = &predicate_expr {
8225 let v = eval::eval_expr(expr, row, &ctx).map_err(|e| {
8226 EngineError::Unsupported(alloc::format!(
8227 "evaluating UNIQUE INDEX predicate against existing row: {e:?}"
8228 ))
8229 })?;
8230 if !predicate_truthy(&v) {
8231 continue;
8232 }
8233 }
8234 let key: alloc::vec::Vec<spg_storage::Value> = key_positions
8235 .iter()
8236 .map(|&p| {
8237 row.values
8238 .get(p)
8239 .cloned()
8240 .unwrap_or(spg_storage::Value::Null)
8241 })
8242 .collect();
8243 if key.iter().any(|v| matches!(v, spg_storage::Value::Null)) {
8244 continue;
8245 }
8246 if seen.iter().any(|other| *other == key) {
8247 return Err(EngineError::Unsupported(alloc::format!(
8248 "CREATE UNIQUE INDEX {:?}: existing rows already violate the constraint",
8249 idx.name
8250 )));
8251 }
8252 seen.push(key);
8253 }
8254 Ok(())
8255}
8256
8257fn unique_key_positions(idx: &spg_storage::Index) -> alloc::vec::Vec<usize> {
8261 let mut out = alloc::vec::Vec::with_capacity(1 + idx.extra_column_positions.len());
8262 out.push(idx.column_position);
8263 out.extend_from_slice(&idx.extra_column_positions);
8264 out
8265}
8266
8267fn enforce_unique_index_inserts(
8275 catalog: &Catalog,
8276 table_name: &str,
8277 rows: &[alloc::vec::Vec<spg_storage::Value>],
8278) -> Result<(), EngineError> {
8279 let table = catalog.get(table_name).ok_or_else(|| {
8280 EngineError::Storage(StorageError::TableNotFound {
8281 name: table_name.into(),
8282 })
8283 })?;
8284 let schema = table.schema();
8285 let ctx = eval::EvalContext::new(&schema.columns, None);
8286 for idx in table.indices() {
8287 if !idx.is_unique {
8288 continue;
8289 }
8290 let predicate_expr = match idx.partial_predicate.as_deref() {
8292 Some(s) => Some(spg_sql::parser::parse_expression(s).map_err(|e| {
8293 EngineError::Unsupported(alloc::format!(
8294 "UNIQUE INDEX {:?} predicate {s:?} failed to re-parse: {e:?}",
8295 idx.name
8296 ))
8297 })?),
8298 None => None,
8299 };
8300 let key_positions = unique_key_positions(idx);
8301 let key_of = |values: &[spg_storage::Value]| -> alloc::vec::Vec<spg_storage::Value> {
8302 key_positions
8303 .iter()
8304 .map(|&p| {
8305 values
8306 .get(p)
8307 .cloned()
8308 .unwrap_or(spg_storage::Value::Null)
8309 })
8310 .collect()
8311 };
8312 let participates = |values: &[spg_storage::Value]| -> Result<bool, EngineError> {
8316 let Some(expr) = &predicate_expr else {
8317 return Ok(true);
8318 };
8319 let tmp_row = spg_storage::Row {
8320 values: values.to_vec(),
8321 };
8322 let v = eval::eval_expr(expr, &tmp_row, &ctx).map_err(|e| {
8323 EngineError::Unsupported(alloc::format!(
8324 "UNIQUE INDEX {:?} predicate eval: {e:?}",
8325 idx.name
8326 ))
8327 })?;
8328 Ok(predicate_truthy(&v))
8329 };
8330 for (batch_idx, row_values) in rows.iter().enumerate() {
8331 if !participates(row_values)? {
8332 continue;
8333 }
8334 let key = key_of(row_values);
8335 if key.iter().any(|v| matches!(v, spg_storage::Value::Null)) {
8336 continue;
8337 }
8338 for prow in table.rows() {
8340 if !participates(&prow.values)? {
8341 continue;
8342 }
8343 if key_of(&prow.values) == key {
8344 return Err(EngineError::Unsupported(alloc::format!(
8345 "UNIQUE INDEX {:?} violation on {table_name:?}: \
8346 row #{batch_idx} duplicates an existing key",
8347 idx.name
8348 )));
8349 }
8350 }
8351 for earlier in &rows[..batch_idx] {
8353 if !participates(earlier)? {
8354 continue;
8355 }
8356 if key_of(earlier) == key {
8357 return Err(EngineError::Unsupported(alloc::format!(
8358 "UNIQUE INDEX {:?} violation on {table_name:?}: \
8359 row #{batch_idx} duplicates an earlier row in the same batch",
8360 idx.name
8361 )));
8362 }
8363 }
8364 }
8365 }
8366 Ok(())
8367}
8368
8369fn enforce_fk_inserts(
8370 catalog: &Catalog,
8371 child_table: &str,
8372 fks: &[spg_storage::ForeignKeyConstraint],
8373 rows: &[Vec<Value>],
8374) -> Result<(), EngineError> {
8375 for fk in fks {
8376 let parent_is_self = fk.parent_table == child_table;
8377 let parent = if parent_is_self {
8378 catalog.get(child_table).ok_or_else(|| {
8381 EngineError::Storage(StorageError::TableNotFound {
8382 name: child_table.into(),
8383 })
8384 })?
8385 } else {
8386 catalog.get(&fk.parent_table).ok_or_else(|| {
8387 EngineError::Storage(StorageError::TableNotFound {
8388 name: fk.parent_table.clone(),
8389 })
8390 })?
8391 };
8392 for (batch_idx, row_values) in rows.iter().enumerate() {
8393 if fk.local_columns.len() == 1 {
8397 let v = &row_values[fk.local_columns[0]];
8398 if matches!(v, Value::Null) {
8399 continue;
8400 }
8401 let parent_col = fk.parent_columns[0];
8402 let key = spg_storage::IndexKey::from_value(v).ok_or_else(|| {
8403 EngineError::Unsupported(alloc::format!(
8404 "FOREIGN KEY column value of type {:?} is not index-eligible",
8405 v.data_type()
8406 ))
8407 })?;
8408 let present_committed = parent.indices().iter().any(|idx| {
8409 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
8410 && idx.column_position == parent_col
8411 && idx.partial_predicate.is_none()
8412 && !idx.lookup_eq(&key).is_empty()
8413 });
8414 let present_in_batch = parent_is_self
8418 && rows[..batch_idx].iter().any(|earlier| {
8419 earlier.get(parent_col) == Some(v)
8420 });
8421 if !(present_committed || present_in_batch) {
8422 return Err(EngineError::Unsupported(alloc::format!(
8423 "FOREIGN KEY violation: no parent row in {:?} where {} = {:?}",
8424 fk.parent_table,
8425 parent
8426 .schema()
8427 .columns
8428 .get(parent_col)
8429 .map_or("?", |c| c.name.as_str()),
8430 v,
8431 )));
8432 }
8433 } else {
8434 if fk.local_columns
8438 .iter()
8439 .all(|&i| matches!(row_values.get(i), Some(Value::Null)))
8440 {
8441 continue;
8442 }
8443 let local: Vec<&Value> = fk.local_columns.iter().map(|&i| &row_values[i]).collect();
8444 let parent_match_committed = parent.rows().iter().any(|prow| {
8445 fk.parent_columns
8446 .iter()
8447 .enumerate()
8448 .all(|(i, &pi)| prow.values.get(pi) == Some(local[i]))
8449 });
8450 let parent_match_in_batch = parent_is_self
8451 && rows[..batch_idx].iter().any(|earlier| {
8452 fk.parent_columns
8453 .iter()
8454 .enumerate()
8455 .all(|(i, &pi)| earlier.get(pi) == Some(local[i]))
8456 });
8457 if !(parent_match_committed || parent_match_in_batch) {
8458 return Err(EngineError::Unsupported(alloc::format!(
8459 "FOREIGN KEY violation: no parent row in {:?} matching composite key",
8460 fk.parent_table,
8461 )));
8462 }
8463 }
8464 }
8465 }
8466 Ok(())
8467}
8468
8469#[derive(Debug, Clone)]
8473struct FkChildStep {
8474 child_table: String,
8475 action: FkChildAction,
8476}
8477
8478#[derive(Debug, Clone)]
8479enum FkChildAction {
8480 Delete { positions: Vec<usize> },
8482 SetNull {
8486 positions: Vec<usize>,
8487 columns: Vec<usize>,
8488 },
8489 SetDefault {
8493 positions: Vec<usize>,
8494 columns: Vec<usize>,
8495 defaults: Vec<Value>,
8496 },
8497}
8498
8499fn plan_fk_parent_deletions(
8515 catalog: &Catalog,
8516 parent_table_name: &str,
8517 to_delete_positions: &[usize],
8518 to_delete_rows: &[Vec<Value>],
8519) -> Result<Vec<FkChildStep>, EngineError> {
8520 use alloc::collections::{BTreeMap, BTreeSet};
8521 if to_delete_rows.is_empty() {
8522 return Ok(Vec::new());
8523 }
8524 let mut delete_plan: BTreeMap<String, BTreeSet<usize>> = BTreeMap::new();
8525 let mut setnull_plan: BTreeMap<String, BTreeSet<(usize, usize)>> = BTreeMap::new();
8527 let mut setdefault_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> =
8528 BTreeMap::new();
8529 let mut visited: BTreeSet<(String, usize)> = BTreeSet::new();
8530 for &p in to_delete_positions {
8531 visited.insert((parent_table_name.to_string(), p));
8532 }
8533 let mut work: Vec<(String, Vec<Value>)> = to_delete_rows
8534 .iter()
8535 .map(|r| (parent_table_name.to_string(), r.clone()))
8536 .collect();
8537 while let Some((cur_parent, parent_row)) = work.pop() {
8538 for child_name in catalog.table_names() {
8539 let child = catalog
8540 .get(&child_name)
8541 .expect("table_names → catalog.get round-trip is total");
8542 for fk in &child.schema().foreign_keys {
8543 if fk.parent_table != cur_parent {
8544 continue;
8545 }
8546 let parent_key: Vec<&Value> = fk
8547 .parent_columns
8548 .iter()
8549 .map(|&pi| &parent_row[pi])
8550 .collect();
8551 if parent_key.iter().any(|v| matches!(v, Value::Null)) {
8552 continue;
8553 }
8554 for (child_row_idx, child_row) in child.rows().iter().enumerate() {
8555 if child_name == cur_parent
8556 && visited.contains(&(child_name.clone(), child_row_idx))
8557 {
8558 continue;
8559 }
8560 let matches_key = fk
8561 .local_columns
8562 .iter()
8563 .enumerate()
8564 .all(|(i, &li)| child_row.values.get(li) == Some(parent_key[i]));
8565 if !matches_key {
8566 continue;
8567 }
8568 match fk.on_delete {
8569 spg_storage::FkAction::Restrict
8570 | spg_storage::FkAction::NoAction => {
8571 return Err(EngineError::Unsupported(alloc::format!(
8572 "FOREIGN KEY violation: DELETE on {cur_parent:?} is \
8573 restricted by FK from {child_name:?}.{:?}",
8574 fk.local_columns,
8575 )));
8576 }
8577 spg_storage::FkAction::Cascade => {
8578 if visited.insert((child_name.clone(), child_row_idx)) {
8579 delete_plan
8580 .entry(child_name.clone())
8581 .or_default()
8582 .insert(child_row_idx);
8583 work.push((child_name.clone(), child_row.values.clone()));
8584 }
8585 }
8586 spg_storage::FkAction::SetNull => {
8587 for &li in &fk.local_columns {
8589 let col = child.schema().columns.get(li).ok_or_else(|| {
8590 EngineError::Unsupported(alloc::format!(
8591 "FK local column {li} missing in {child_name:?}"
8592 ))
8593 })?;
8594 if !col.nullable {
8595 return Err(EngineError::Unsupported(alloc::format!(
8596 "FOREIGN KEY ON DELETE SET NULL: column \
8597 {child_name:?}.{:?} is NOT NULL — cannot SET NULL",
8598 col.name,
8599 )));
8600 }
8601 }
8602 let entry = setnull_plan.entry(child_name.clone()).or_default();
8603 for &li in &fk.local_columns {
8604 entry.insert((child_row_idx, li));
8605 }
8606 }
8607 spg_storage::FkAction::SetDefault => {
8608 let entry =
8610 setdefault_plan.entry(child_name.clone()).or_default();
8611 for &li in &fk.local_columns {
8612 let col = child.schema().columns.get(li).ok_or_else(|| {
8613 EngineError::Unsupported(alloc::format!(
8614 "FK local column {li} missing in {child_name:?}"
8615 ))
8616 })?;
8617 let default = col.default.clone().ok_or_else(|| {
8618 EngineError::Unsupported(alloc::format!(
8619 "FOREIGN KEY ON DELETE SET DEFAULT: column \
8620 {child_name:?}.{:?} has no DEFAULT declared",
8621 col.name,
8622 ))
8623 })?;
8624 entry.insert((child_row_idx, li), default);
8625 }
8626 }
8627 }
8628 }
8629 }
8630 }
8631 }
8632 let mut steps: Vec<FkChildStep> = Vec::new();
8640 for (child_table, entries) in setnull_plan {
8641 let (positions, columns): (Vec<usize>, Vec<usize>) = entries.into_iter().unzip();
8642 steps.push(FkChildStep {
8643 child_table,
8644 action: FkChildAction::SetNull { positions, columns },
8645 });
8646 }
8647 for (child_table, entries) in setdefault_plan {
8648 let mut positions = Vec::with_capacity(entries.len());
8649 let mut columns = Vec::with_capacity(entries.len());
8650 let mut defaults = Vec::with_capacity(entries.len());
8651 for ((p, c), v) in entries {
8652 positions.push(p);
8653 columns.push(c);
8654 defaults.push(v);
8655 }
8656 steps.push(FkChildStep {
8657 child_table,
8658 action: FkChildAction::SetDefault {
8659 positions,
8660 columns,
8661 defaults,
8662 },
8663 });
8664 }
8665 for (child_table, positions) in delete_plan {
8666 steps.push(FkChildStep {
8667 child_table,
8668 action: FkChildAction::Delete {
8669 positions: positions.into_iter().collect(),
8670 },
8671 });
8672 }
8673 Ok(steps)
8674}
8675
8676fn plan_fk_parent_updates(
8693 catalog: &Catalog,
8694 parent_table_name: &str,
8695 plan_with_old: &[(usize, Vec<Value>, Vec<Value>)],
8696) -> Result<Vec<FkChildStep>, EngineError> {
8697 use alloc::collections::BTreeMap;
8698 if plan_with_old.is_empty() {
8699 return Ok(Vec::new());
8700 }
8701 let delete_plan: BTreeMap<String, alloc::collections::BTreeSet<usize>> = BTreeMap::new();
8706 let mut setnull_plan: BTreeMap<
8707 String,
8708 alloc::collections::BTreeSet<(usize, usize)>,
8709 > = BTreeMap::new();
8710 let mut setdefault_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> =
8711 BTreeMap::new();
8712 let mut cascade_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> = BTreeMap::new();
8714
8715 for child_name in catalog.table_names() {
8716 let child = catalog
8717 .get(&child_name)
8718 .expect("table_names → catalog.get total");
8719 for fk in &child.schema().foreign_keys {
8720 if fk.parent_table != parent_table_name {
8721 continue;
8722 }
8723 for (_pos, old_row, new_row) in plan_with_old {
8724 let key_changed = fk
8726 .parent_columns
8727 .iter()
8728 .any(|&pi| old_row.get(pi) != new_row.get(pi));
8729 if !key_changed {
8730 continue;
8731 }
8732 let old_key: Vec<&Value> = fk
8734 .parent_columns
8735 .iter()
8736 .map(|&pi| &old_row[pi])
8737 .collect();
8738 if old_key.iter().any(|v| matches!(v, Value::Null)) {
8739 continue;
8741 }
8742 let new_key: Vec<&Value> = fk
8743 .parent_columns
8744 .iter()
8745 .map(|&pi| &new_row[pi])
8746 .collect();
8747 for (child_row_idx, child_row) in child.rows().iter().enumerate() {
8748 if child_name == parent_table_name
8751 && plan_with_old
8752 .iter()
8753 .any(|(p, _, _)| *p == child_row_idx)
8754 {
8755 continue;
8756 }
8757 let matches_key = fk
8758 .local_columns
8759 .iter()
8760 .enumerate()
8761 .all(|(i, &li)| child_row.values.get(li) == Some(old_key[i]));
8762 if !matches_key {
8763 continue;
8764 }
8765 match fk.on_update {
8766 spg_storage::FkAction::Restrict
8767 | spg_storage::FkAction::NoAction => {
8768 return Err(EngineError::Unsupported(alloc::format!(
8769 "FOREIGN KEY violation: UPDATE on {parent_table_name:?} PK is \
8770 restricted by FK from {child_name:?}.{:?}",
8771 fk.local_columns,
8772 )));
8773 }
8774 spg_storage::FkAction::Cascade => {
8775 let entry = cascade_plan.entry(child_name.clone()).or_default();
8777 for (i, &li) in fk.local_columns.iter().enumerate() {
8778 entry.insert((child_row_idx, li), new_key[i].clone());
8779 }
8780 }
8781 spg_storage::FkAction::SetNull => {
8782 for &li in &fk.local_columns {
8783 let col = child.schema().columns.get(li).ok_or_else(|| {
8784 EngineError::Unsupported(alloc::format!(
8785 "FK local column {li} missing in {child_name:?}"
8786 ))
8787 })?;
8788 if !col.nullable {
8789 return Err(EngineError::Unsupported(alloc::format!(
8790 "FOREIGN KEY ON UPDATE SET NULL: column \
8791 {child_name:?}.{:?} is NOT NULL",
8792 col.name,
8793 )));
8794 }
8795 }
8796 let entry = setnull_plan.entry(child_name.clone()).or_default();
8797 for &li in &fk.local_columns {
8798 entry.insert((child_row_idx, li));
8799 }
8800 }
8801 spg_storage::FkAction::SetDefault => {
8802 let entry =
8803 setdefault_plan.entry(child_name.clone()).or_default();
8804 for &li in &fk.local_columns {
8805 let col = child.schema().columns.get(li).ok_or_else(|| {
8806 EngineError::Unsupported(alloc::format!(
8807 "FK local column {li} missing in {child_name:?}"
8808 ))
8809 })?;
8810 let default = col.default.clone().ok_or_else(|| {
8811 EngineError::Unsupported(alloc::format!(
8812 "FOREIGN KEY ON UPDATE SET DEFAULT: column \
8813 {child_name:?}.{:?} has no DEFAULT",
8814 col.name,
8815 ))
8816 })?;
8817 entry.insert((child_row_idx, li), default);
8818 }
8819 }
8820 }
8821 }
8822 }
8823 }
8824 }
8825 let mut steps: Vec<FkChildStep> = Vec::new();
8828 for (child_table, entries) in cascade_plan {
8829 let mut positions = Vec::with_capacity(entries.len());
8830 let mut columns = Vec::with_capacity(entries.len());
8831 let mut defaults = Vec::with_capacity(entries.len());
8832 for ((p, c), v) in entries {
8833 positions.push(p);
8834 columns.push(c);
8835 defaults.push(v);
8836 }
8837 steps.push(FkChildStep {
8842 child_table,
8843 action: FkChildAction::SetDefault {
8844 positions,
8845 columns,
8846 defaults,
8847 },
8848 });
8849 }
8850 for (child_table, entries) in setnull_plan {
8851 let (positions, columns): (Vec<usize>, Vec<usize>) = entries.into_iter().unzip();
8852 steps.push(FkChildStep {
8853 child_table,
8854 action: FkChildAction::SetNull { positions, columns },
8855 });
8856 }
8857 for (child_table, entries) in setdefault_plan {
8858 let mut positions = Vec::with_capacity(entries.len());
8859 let mut columns = Vec::with_capacity(entries.len());
8860 let mut defaults = Vec::with_capacity(entries.len());
8861 for ((p, c), v) in entries {
8862 positions.push(p);
8863 columns.push(c);
8864 defaults.push(v);
8865 }
8866 steps.push(FkChildStep {
8867 child_table,
8868 action: FkChildAction::SetDefault {
8869 positions,
8870 columns,
8871 defaults,
8872 },
8873 });
8874 }
8875 let _ = delete_plan; Ok(steps)
8877}
8878
8879fn apply_fk_child_step(
8883 catalog: &mut Catalog,
8884 step: &FkChildStep,
8885) -> Result<(), EngineError> {
8886 let child = catalog.get_mut(&step.child_table).ok_or_else(|| {
8887 EngineError::Storage(StorageError::TableNotFound {
8888 name: step.child_table.clone(),
8889 })
8890 })?;
8891 match &step.action {
8892 FkChildAction::Delete { positions } => {
8893 let _ = child.delete_rows(positions);
8894 }
8895 FkChildAction::SetNull { positions, columns } => {
8896 apply_per_cell_writes(child, positions, columns, |_| Value::Null)?;
8897 }
8898 FkChildAction::SetDefault {
8899 positions,
8900 columns,
8901 defaults,
8902 } => {
8903 apply_per_cell_writes(child, positions, columns, |i| defaults[i].clone())?;
8904 }
8905 }
8906 Ok(())
8907}
8908
8909fn apply_per_cell_writes(
8915 child: &mut spg_storage::Table,
8916 positions: &[usize],
8917 columns: &[usize],
8918 mut value_for: impl FnMut(usize) -> Value,
8919) -> Result<(), EngineError> {
8920 use alloc::collections::BTreeMap;
8921 let mut by_row: BTreeMap<usize, Vec<(usize, Value)>> = BTreeMap::new();
8922 for i in 0..positions.len() {
8923 by_row
8924 .entry(positions[i])
8925 .or_default()
8926 .push((columns[i], value_for(i)));
8927 }
8928 for (pos, mutations) in by_row {
8929 let mut new_values = child.rows()[pos].values.clone();
8930 for (col, v) in mutations {
8931 if let Some(slot) = new_values.get_mut(col) {
8932 *slot = v;
8933 }
8934 }
8935 child
8936 .update_row(pos, new_values)
8937 .map_err(EngineError::Storage)?;
8938 }
8939 Ok(())
8940}
8941
8942fn fk_action_sql_to_storage(a: spg_sql::ast::FkAction) -> spg_storage::FkAction {
8943 match a {
8944 spg_sql::ast::FkAction::Restrict => spg_storage::FkAction::Restrict,
8945 spg_sql::ast::FkAction::Cascade => spg_storage::FkAction::Cascade,
8946 spg_sql::ast::FkAction::SetNull => spg_storage::FkAction::SetNull,
8947 spg_sql::ast::FkAction::SetDefault => spg_storage::FkAction::SetDefault,
8948 spg_sql::ast::FkAction::NoAction => spg_storage::FkAction::NoAction,
8949 }
8950}
8951
8952fn resolve_column_default_free(
8958 col: &ColumnSchema,
8959 clock_fn: Option<ClockFn>,
8960) -> Result<Value, EngineError> {
8961 if let Some(rt) = &col.runtime_default {
8962 return eval_runtime_default_free(rt, col.ty, clock_fn);
8963 }
8964 Ok(col.default.clone().unwrap_or(Value::Null))
8965}
8966
8967fn eval_runtime_default_free(
8968 rt: &str,
8969 ty: DataType,
8970 clock_fn: Option<ClockFn>,
8971) -> Result<Value, EngineError> {
8972 let s = rt.trim().to_ascii_lowercase();
8973 let canonical = s.trim_end_matches("()");
8974 let now_us = match clock_fn {
8975 Some(f) => f(),
8976 None => 0,
8977 };
8978 let v = match canonical {
8979 "now" | "current_timestamp" | "localtimestamp" => {
8980 Value::Timestamp(now_us)
8981 }
8982 "current_date" => Value::Date((now_us / 86_400_000_000) as i32),
8983 "current_time" | "localtime" => Value::Timestamp(now_us),
8984 other => {
8985 return Err(EngineError::Unsupported(alloc::format!(
8986 "runtime DEFAULT expression {other:?} not supported \
8987 (v7.9.21 whitelist: now() / current_timestamp / \
8988 current_date / current_time / localtimestamp / \
8989 localtime)"
8990 )));
8991 }
8992 };
8993 coerce_value(v, ty, "DEFAULT", 0)
8994}
8995
8996fn is_runtime_default_expr(expr: &Expr) -> bool {
9002 match expr {
9003 Expr::FunctionCall { .. } => true,
9004 Expr::Unary { expr, .. } => is_runtime_default_expr(expr),
9005 _ => false,
9006 }
9007}
9008
9009fn column_def_to_schema(c: ColumnDef) -> Result<ColumnSchema, EngineError> {
9010 let ty = column_type_to_data_type(c.ty);
9011 let mut schema = ColumnSchema::new(c.name.clone(), ty, c.nullable);
9012 if let Some(default_expr) = c.default {
9013 if is_runtime_default_expr(&default_expr) {
9019 let display = alloc::format!("{default_expr}");
9020 schema = schema.with_runtime_default(display);
9021 } else {
9022 let raw = literal_expr_to_value(default_expr)?;
9023 let coerced = coerce_value(raw, ty, &c.name, 0)?;
9024 schema = schema.with_default(coerced);
9025 }
9026 }
9027 if c.auto_increment {
9028 if !matches!(ty, DataType::SmallInt | DataType::Int | DataType::BigInt) {
9030 return Err(EngineError::Unsupported(alloc::format!(
9031 "AUTO_INCREMENT requires an integer column type, got {ty:?}"
9032 )));
9033 }
9034 schema = schema.with_auto_increment();
9035 }
9036 Ok(schema)
9037}
9038
9039fn decode_bytea_literal(s: &str) -> Result<alloc::vec::Vec<u8>, &'static str> {
9044 let s = s.trim();
9045 if let Some(hex) = s.strip_prefix("\\x").or_else(|| s.strip_prefix("\\X")) {
9046 let cleaned: alloc::string::String = hex.chars().filter(|c| !c.is_whitespace()).collect();
9048 if cleaned.len() % 2 != 0 {
9049 return Err("odd-length hex literal");
9050 }
9051 let mut out = alloc::vec::Vec::with_capacity(cleaned.len() / 2);
9052 let cleaned_bytes = cleaned.as_bytes();
9053 for i in (0..cleaned_bytes.len()).step_by(2) {
9054 let hi = hex_nibble(cleaned_bytes[i])?;
9055 let lo = hex_nibble(cleaned_bytes[i + 1])?;
9056 out.push((hi << 4) | lo);
9057 }
9058 return Ok(out);
9059 }
9060 let bytes = s.as_bytes();
9063 let mut out = alloc::vec::Vec::with_capacity(bytes.len());
9064 let mut i = 0;
9065 while i < bytes.len() {
9066 let b = bytes[i];
9067 if b == b'\\' && i + 1 < bytes.len() {
9068 let n = bytes[i + 1];
9069 if n == b'\\' {
9070 out.push(b'\\');
9071 i += 2;
9072 continue;
9073 }
9074 if n.is_ascii_digit() && i + 3 < bytes.len() && bytes[i + 2].is_ascii_digit()
9075 && bytes[i + 3].is_ascii_digit()
9076 {
9077 let oct = |x: u8| (x - b'0') as u32;
9078 let v = oct(n) * 64 + oct(bytes[i + 2]) * 8 + oct(bytes[i + 3]);
9079 if v <= 0xFF {
9080 out.push(v as u8);
9081 i += 4;
9082 continue;
9083 }
9084 }
9085 }
9086 out.push(b);
9087 i += 1;
9088 }
9089 Ok(out)
9090}
9091
9092fn hex_nibble(b: u8) -> Result<u8, &'static str> {
9093 match b {
9094 b'0'..=b'9' => Ok(b - b'0'),
9095 b'a'..=b'f' => Ok(b - b'a' + 10),
9096 b'A'..=b'F' => Ok(b - b'A' + 10),
9097 _ => Err("invalid hex digit"),
9098 }
9099}
9100
9101fn decode_text_array_literal(
9108 s: &str,
9109) -> Result<alloc::vec::Vec<Option<alloc::string::String>>, &'static str> {
9110 let trimmed = s.trim();
9111 let inner = trimmed
9112 .strip_prefix('{')
9113 .and_then(|x| x.strip_suffix('}'))
9114 .ok_or("TEXT[] literal must be enclosed in '{...}'")?;
9115 let mut out: alloc::vec::Vec<Option<alloc::string::String>> = alloc::vec::Vec::new();
9116 if inner.trim().is_empty() {
9117 return Ok(out);
9118 }
9119 let bytes = inner.as_bytes();
9120 let mut i = 0;
9121 while i <= bytes.len() {
9122 while i < bytes.len() && (bytes[i] == b' ' || bytes[i] == b'\t') {
9124 i += 1;
9125 }
9126 if i < bytes.len() && bytes[i] == b'"' {
9128 i += 1; let mut buf = alloc::string::String::new();
9130 while i < bytes.len() && bytes[i] != b'"' {
9131 if bytes[i] == b'\\' && i + 1 < bytes.len() {
9132 buf.push(bytes[i + 1] as char);
9133 i += 2;
9134 } else {
9135 buf.push(bytes[i] as char);
9136 i += 1;
9137 }
9138 }
9139 if i >= bytes.len() {
9140 return Err("unterminated quoted element");
9141 }
9142 i += 1; out.push(Some(buf));
9144 } else {
9145 let start = i;
9147 while i < bytes.len() && bytes[i] != b',' {
9148 i += 1;
9149 }
9150 let raw = inner[start..i].trim();
9151 if raw.eq_ignore_ascii_case("NULL") {
9152 out.push(None);
9153 } else {
9154 out.push(Some(alloc::string::ToString::to_string(raw)));
9155 }
9156 }
9157 while i < bytes.len() && (bytes[i] == b' ' || bytes[i] == b'\t') {
9159 i += 1;
9160 }
9161 if i >= bytes.len() {
9162 break;
9163 }
9164 if bytes[i] != b',' {
9165 return Err("expected ',' between TEXT[] elements");
9166 }
9167 i += 1;
9168 }
9169 Ok(out)
9170}
9171
9172fn encode_text_array(items: &[Option<alloc::string::String>]) -> alloc::string::String {
9177 let mut out = alloc::string::String::with_capacity(2 + items.len() * 8);
9178 out.push('{');
9179 for (i, item) in items.iter().enumerate() {
9180 if i > 0 {
9181 out.push(',');
9182 }
9183 match item {
9184 None => out.push_str("NULL"),
9185 Some(s) => {
9186 let needs_quote = s.is_empty()
9187 || s.eq_ignore_ascii_case("NULL")
9188 || s.chars().any(|c| {
9189 matches!(c, ',' | '{' | '}' | '"' | '\\' | ' ' | '\t')
9190 });
9191 if needs_quote {
9192 out.push('"');
9193 for c in s.chars() {
9194 if c == '"' || c == '\\' {
9195 out.push('\\');
9196 }
9197 out.push(c);
9198 }
9199 out.push('"');
9200 } else {
9201 out.push_str(s);
9202 }
9203 }
9204 }
9205 }
9206 out.push('}');
9207 out
9208}
9209
9210fn encode_bytea_hex(b: &[u8]) -> alloc::string::String {
9214 let mut out = alloc::string::String::with_capacity(2 + 2 * b.len());
9215 out.push_str("\\x");
9216 for byte in b {
9217 let hi = byte >> 4;
9218 let lo = byte & 0x0F;
9219 out.push(hex_digit(hi));
9220 out.push(hex_digit(lo));
9221 }
9222 out
9223}
9224
9225const fn hex_digit(n: u8) -> char {
9226 match n {
9227 0..=9 => (b'0' + n) as char,
9228 10..=15 => (b'a' + n - 10) as char,
9229 _ => '?',
9230 }
9231}
9232
9233const fn column_type_to_data_type(t: ColumnTypeName) -> DataType {
9234 match t {
9235 ColumnTypeName::SmallInt => DataType::SmallInt,
9236 ColumnTypeName::Int => DataType::Int,
9237 ColumnTypeName::BigInt => DataType::BigInt,
9238 ColumnTypeName::Float => DataType::Float,
9239 ColumnTypeName::Text => DataType::Text,
9240 ColumnTypeName::Varchar(n) => DataType::Varchar(n),
9241 ColumnTypeName::Char(n) => DataType::Char(n),
9242 ColumnTypeName::Bool => DataType::Bool,
9243 ColumnTypeName::Vector { dim, encoding } => DataType::Vector {
9244 dim,
9245 encoding: match encoding {
9246 SqlVecEncoding::F32 => VecEncoding::F32,
9247 SqlVecEncoding::Sq8 => VecEncoding::Sq8,
9248 SqlVecEncoding::F16 => VecEncoding::F16,
9249 },
9250 },
9251 ColumnTypeName::Numeric(precision, scale) => DataType::Numeric { precision, scale },
9252 ColumnTypeName::Date => DataType::Date,
9253 ColumnTypeName::Timestamp => DataType::Timestamp,
9254 ColumnTypeName::Timestamptz => DataType::Timestamptz,
9255 ColumnTypeName::Json => DataType::Json,
9256 ColumnTypeName::Jsonb => DataType::Jsonb,
9257 ColumnTypeName::Bytes => DataType::Bytes,
9258 ColumnTypeName::TextArray => DataType::TextArray,
9259 }
9260}
9261
9262fn literal_expr_to_value(expr: Expr) -> Result<Value, EngineError> {
9266 match expr {
9267 Expr::Literal(l) => Ok(literal_to_value(l)),
9268 Expr::Cast { expr, target } => {
9269 let inner_value = literal_expr_to_value(*expr)?;
9270 crate::eval::cast_value(inner_value, target).map_err(EngineError::Eval)
9271 }
9272 Expr::Unary {
9273 op: UnOp::Neg,
9274 expr,
9275 } => match *expr {
9276 Expr::Literal(Literal::Integer(n)) => {
9277 let neg = n.checked_neg().ok_or_else(|| {
9280 EngineError::Unsupported("integer literal overflow on negation".into())
9281 })?;
9282 Ok(int_value_for(neg))
9283 }
9284 Expr::Literal(Literal::Float(x)) => Ok(Value::Float(-x)),
9285 other => Err(EngineError::Unsupported(alloc::format!(
9286 "unary minus over non-literal expression: {other:?}"
9287 ))),
9288 },
9289 Expr::Array(items) => {
9295 let mut out: alloc::vec::Vec<Option<alloc::string::String>> =
9296 alloc::vec::Vec::with_capacity(items.len());
9297 for elem in items {
9298 match literal_expr_to_value(elem)? {
9299 Value::Null => out.push(None),
9300 Value::Text(s) => out.push(Some(s)),
9301 other => out.push(Some(alloc::format!("{other:?}"))),
9302 }
9303 }
9304 Ok(Value::TextArray(out))
9305 }
9306 other => Err(EngineError::Unsupported(alloc::format!(
9307 "non-literal INSERT value expression: {other:?}"
9308 ))),
9309 }
9310}
9311
9312fn literal_to_value(l: Literal) -> Value {
9313 match l {
9314 Literal::Integer(n) => int_value_for(n),
9315 Literal::Float(x) => Value::Float(x),
9316 Literal::String(s) => Value::Text(s),
9317 Literal::Bool(b) => Value::Bool(b),
9318 Literal::Null => Value::Null,
9319 Literal::Vector(v) => Value::Vector(v),
9320 Literal::Interval { months, micros, .. } => Value::Interval { months, micros },
9321 }
9322}
9323
9324fn int_value_for(n: i64) -> Value {
9328 if let Ok(small) = i32::try_from(n) {
9329 Value::Int(small)
9330 } else {
9331 Value::BigInt(n)
9332 }
9333}
9334
9335#[allow(clippy::too_many_lines)]
9341fn coerce_value(
9342 v: Value,
9343 expected: DataType,
9344 col_name: &str,
9345 position: usize,
9346) -> Result<Value, EngineError> {
9347 if v.is_null() {
9348 return Ok(Value::Null);
9349 }
9350 let actual = v.data_type().expect("non-null");
9351 if actual == expected {
9352 return Ok(v);
9353 }
9354 let coerced =
9355 match (v, expected) {
9356 (Value::Int(n), DataType::BigInt) => Some(Value::BigInt(i64::from(n))),
9357 (Value::Int(n), DataType::Float) => Some(Value::Float(f64::from(n))),
9358 (Value::Int(n), DataType::SmallInt) => i16::try_from(n).ok().map(Value::SmallInt),
9359 (Value::Int(n), DataType::Numeric { precision, scale }) => Some(numeric_from_integer(
9360 i128::from(n),
9361 precision,
9362 scale,
9363 col_name,
9364 )?),
9365 (Value::SmallInt(n), DataType::Int) => Some(Value::Int(i32::from(n))),
9366 (Value::SmallInt(n), DataType::BigInt) => Some(Value::BigInt(i64::from(n))),
9367 (Value::SmallInt(n), DataType::Float) => Some(Value::Float(f64::from(n))),
9368 (Value::SmallInt(n), DataType::Numeric { precision, scale }) => Some(
9369 numeric_from_integer(i128::from(n), precision, scale, col_name)?,
9370 ),
9371 (Value::BigInt(n), DataType::Int) => i32::try_from(n).ok().map(Value::Int),
9372 (Value::BigInt(n), DataType::SmallInt) => i16::try_from(n).ok().map(Value::SmallInt),
9373 #[allow(clippy::cast_precision_loss)]
9374 (Value::BigInt(n), DataType::Float) => Some(Value::Float(n as f64)),
9375 (Value::BigInt(n), DataType::Numeric { precision, scale }) => Some(
9376 numeric_from_integer(i128::from(n), precision, scale, col_name)?,
9377 ),
9378 (Value::Float(x), DataType::Numeric { precision, scale }) => {
9379 Some(numeric_from_float(x, precision, scale, col_name)?)
9380 }
9381 (Value::Text(s), DataType::Date) => {
9383 let d = eval::parse_date_literal(&s).ok_or_else(|| {
9384 EngineError::Eval(EvalError::TypeMismatch {
9385 detail: alloc::format!(
9386 "cannot parse {s:?} as DATE for column `{col_name}`"
9387 ),
9388 })
9389 })?;
9390 Some(Value::Date(d))
9391 }
9392 (Value::Text(s), DataType::Json | DataType::Jsonb) => Some(Value::Json(s)),
9396 (Value::Json(s), DataType::Text) => Some(Value::Text(s)),
9397 (Value::Text(s), DataType::Bytes) => {
9404 let bytes = decode_bytea_literal(&s).map_err(|e| {
9405 EngineError::Eval(EvalError::TypeMismatch {
9406 detail: alloc::format!(
9407 "cannot parse {s:?} as BYTEA for column `{col_name}`: {e}"
9408 ),
9409 })
9410 })?;
9411 Some(Value::Bytes(bytes))
9412 }
9413 (Value::Bytes(b), DataType::Text) => Some(Value::Text(encode_bytea_hex(&b))),
9417 (Value::Text(s), DataType::TextArray) => {
9422 let arr = decode_text_array_literal(&s).map_err(|e| {
9423 EngineError::Eval(EvalError::TypeMismatch {
9424 detail: alloc::format!(
9425 "cannot parse {s:?} as TEXT[] for column `{col_name}`: {e}"
9426 ),
9427 })
9428 })?;
9429 Some(Value::TextArray(arr))
9430 }
9431 (Value::TextArray(items), DataType::Text) => {
9435 Some(Value::Text(encode_text_array(&items)))
9436 }
9437 (Value::Text(s), DataType::Timestamp | DataType::Timestamptz) => {
9438 let t = eval::parse_timestamp_literal(&s).ok_or_else(|| {
9439 EngineError::Eval(EvalError::TypeMismatch {
9440 detail: alloc::format!(
9441 "cannot parse {s:?} as TIMESTAMP for column `{col_name}`"
9442 ),
9443 })
9444 })?;
9445 Some(Value::Timestamp(t))
9446 }
9447 (Value::Date(d), DataType::Timestamp | DataType::Timestamptz) => {
9450 Some(Value::Timestamp(i64::from(d) * 86_400_000_000))
9451 }
9452 (Value::Timestamp(t), DataType::Timestamptz) => Some(Value::Timestamp(t)),
9456 (Value::Timestamp(t), DataType::Date) => {
9457 let days = t.div_euclid(86_400_000_000);
9458 i32::try_from(days).ok().map(Value::Date)
9459 }
9460 (
9461 Value::Numeric {
9462 scaled,
9463 scale: src_scale,
9464 },
9465 DataType::Numeric { precision, scale },
9466 ) => Some(numeric_rescale(
9467 scaled, src_scale, precision, scale, col_name,
9468 )?),
9469 #[allow(clippy::cast_precision_loss)]
9470 (Value::Numeric { scaled, scale }, DataType::Float) => {
9471 let mut div = 1.0_f64;
9472 for _ in 0..scale {
9473 div *= 10.0;
9474 }
9475 Some(Value::Float((scaled as f64) / div))
9476 }
9477 (Value::Numeric { scaled, scale }, DataType::Int) => {
9478 let truncated = numeric_truncate_to_integer(scaled, scale);
9479 i32::try_from(truncated).ok().map(Value::Int)
9480 }
9481 (Value::Numeric { scaled, scale }, DataType::BigInt) => {
9482 let truncated = numeric_truncate_to_integer(scaled, scale);
9483 i64::try_from(truncated).ok().map(Value::BigInt)
9484 }
9485 (Value::Numeric { scaled, scale }, DataType::SmallInt) => {
9486 let truncated = numeric_truncate_to_integer(scaled, scale);
9487 i16::try_from(truncated).ok().map(Value::SmallInt)
9488 }
9489 (Value::Text(s), DataType::Varchar(max)) => {
9491 if u32::try_from(s.chars().count()).unwrap_or(u32::MAX) <= max {
9492 Some(Value::Text(s))
9493 } else {
9494 return Err(EngineError::Unsupported(alloc::format!(
9495 "value for VARCHAR({max}) column `{col_name}` exceeds length: \
9496 {} chars",
9497 s.chars().count()
9498 )));
9499 }
9500 }
9501 (
9509 Value::Vector(v),
9510 DataType::Vector {
9511 dim,
9512 encoding: VecEncoding::Sq8,
9513 },
9514 ) if v.len() == dim as usize => {
9515 Some(Value::Sq8Vector(spg_storage::quantize::quantize(&v)))
9516 }
9517 (
9522 Value::Vector(v),
9523 DataType::Vector {
9524 dim,
9525 encoding: VecEncoding::F16,
9526 },
9527 ) if v.len() == dim as usize => Some(Value::HalfVector(
9528 spg_storage::halfvec::HalfVector::from_f32_slice(&v),
9529 )),
9530 (Value::Text(s), DataType::Char(size)) => {
9534 let len = u32::try_from(s.chars().count()).unwrap_or(u32::MAX);
9535 if len > size {
9536 return Err(EngineError::Unsupported(alloc::format!(
9537 "value for CHAR({size}) column `{col_name}` exceeds length: \
9538 {len} chars"
9539 )));
9540 }
9541 let need = (size - len) as usize;
9542 let mut padded = s;
9543 padded.reserve(need);
9544 for _ in 0..need {
9545 padded.push(' ');
9546 }
9547 Some(Value::Text(padded))
9548 }
9549 _ => None,
9550 };
9551 coerced.ok_or(EngineError::Storage(StorageError::TypeMismatch {
9552 column: col_name.into(),
9553 expected,
9554 actual,
9555 position,
9556 }))
9557}
9558
9559#[cfg(test)]
9560mod tests {
9561 use super::*;
9562 use alloc::vec;
9563
9564 fn unwrap_command_ok(r: &QueryResult) -> usize {
9565 match r {
9566 QueryResult::CommandOk { affected, .. } => *affected,
9567 QueryResult::Rows { .. } => panic!("expected CommandOk, got Rows"),
9568 }
9569 }
9570
9571 #[test]
9572 fn create_table_registers_schema() {
9573 let mut e = Engine::new();
9574 e.execute("CREATE TABLE foo (a INT NOT NULL, b TEXT)")
9575 .unwrap();
9576 assert_eq!(e.catalog().table_count(), 1);
9577 let t = e.catalog().get("foo").unwrap();
9578 assert_eq!(t.schema().columns.len(), 2);
9579 assert_eq!(t.schema().columns[0].ty, DataType::Int);
9580 assert!(!t.schema().columns[0].nullable);
9581 assert_eq!(t.schema().columns[1].ty, DataType::Text);
9582 }
9583
9584 #[test]
9585 fn create_table_vector_default_is_f32_encoded() {
9586 let mut e = Engine::new();
9587 e.execute("CREATE TABLE t (v VECTOR(8))").unwrap();
9588 let t = e.catalog().get("t").unwrap();
9589 assert_eq!(
9590 t.schema().columns[0].ty,
9591 DataType::Vector {
9592 dim: 8,
9593 encoding: VecEncoding::F32,
9594 },
9595 );
9596 }
9597
9598 #[test]
9599 fn create_table_vector_using_sq8_succeeds() {
9600 let mut e = Engine::new();
9604 e.execute("CREATE TABLE t (v VECTOR(8) USING SQ8)").unwrap();
9605 let t = e.catalog().get("t").unwrap();
9606 assert_eq!(
9607 t.schema().columns[0].ty,
9608 DataType::Vector {
9609 dim: 8,
9610 encoding: VecEncoding::Sq8,
9611 },
9612 );
9613 }
9614
9615 #[test]
9616 fn insert_into_sq8_column_quantises_f32_payload() {
9617 let mut e = Engine::new();
9624 e.execute("CREATE TABLE t (v VECTOR(4) USING SQ8)").unwrap();
9625 e.execute("INSERT INTO t VALUES ([0.0, 0.25, 0.5, 1.0])")
9626 .unwrap();
9627 let t = e.catalog().get("t").unwrap();
9628 assert_eq!(t.rows().len(), 1);
9629 match &t.rows()[0].values[0] {
9630 Value::Sq8Vector(q) => {
9631 assert_eq!(q.bytes.len(), 4);
9632 assert!((q.min - 0.0).abs() < 1e-6);
9634 assert!((q.max - 1.0).abs() < 1e-6);
9635 }
9636 other => panic!("expected Sq8Vector cell, got {other:?}"),
9637 }
9638 }
9639
9640 #[test]
9641 fn create_table_vector_using_half_succeeds_and_insert_converts_to_f16() {
9642 let mut e = Engine::new();
9649 e.execute("CREATE TABLE t (v VECTOR(4) USING HALF)")
9650 .unwrap();
9651 e.execute("INSERT INTO t VALUES ([0.0, 0.25, 0.5, 1.0])")
9652 .unwrap();
9653 let t = e.catalog().get("t").unwrap();
9654 assert_eq!(t.rows().len(), 1);
9655 match &t.rows()[0].values[0] {
9656 Value::HalfVector(h) => {
9657 assert_eq!(h.dim(), 4);
9658 let back = h.to_f32_vec();
9659 let expected = alloc::vec![0.0_f32, 0.25, 0.5, 1.0];
9660 for (g, e) in back.iter().zip(expected.iter()) {
9661 assert!(
9662 (g - e).abs() < 1e-6,
9663 "{g} vs {e} should be exact on f16 grid"
9664 );
9665 }
9666 }
9667 other => panic!("expected HalfVector cell, got {other:?}"),
9668 }
9669 }
9670
9671 #[test]
9672 fn alter_index_rebuild_in_place_succeeds() {
9673 let mut e = Engine::new();
9678 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(3) NOT NULL)")
9679 .unwrap();
9680 for i in 0..8_i32 {
9681 #[allow(clippy::cast_precision_loss)]
9682 let base = (i as f32) * 0.1;
9683 e.execute(&alloc::format!(
9684 "INSERT INTO t VALUES ({i}, [{base}, {b1}, {b2}])",
9685 b1 = base + 0.01,
9686 b2 = base + 0.02,
9687 ))
9688 .unwrap();
9689 }
9690 e.execute("CREATE INDEX t_idx ON t USING hnsw (v)").unwrap();
9691 e.execute("ALTER INDEX t_idx REBUILD").unwrap();
9692 assert_eq!(
9694 e.catalog().get("t").unwrap().schema().columns[1].ty,
9695 DataType::Vector {
9696 dim: 3,
9697 encoding: VecEncoding::F32,
9698 },
9699 );
9700 }
9701
9702 #[test]
9703 fn alter_index_rebuild_with_encoding_switches_cell_type() {
9704 let mut e = Engine::new();
9709 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(4) NOT NULL)")
9710 .unwrap();
9711 e.execute("INSERT INTO t VALUES (1, [0.0, 0.25, 0.5, 1.0])")
9712 .unwrap();
9713 e.execute("CREATE INDEX t_idx ON t USING hnsw (v)").unwrap();
9714 e.execute("ALTER INDEX t_idx REBUILD WITH (encoding = SQ8)")
9715 .unwrap();
9716 let t = e.catalog().get("t").unwrap();
9717 assert_eq!(
9718 t.schema().columns[1].ty,
9719 DataType::Vector {
9720 dim: 4,
9721 encoding: VecEncoding::Sq8,
9722 },
9723 );
9724 assert!(matches!(t.rows()[0].values[1], Value::Sq8Vector(_)));
9725 }
9726
9727 #[test]
9728 fn alter_index_rebuild_unknown_index_errors() {
9729 let mut e = Engine::new();
9730 let err = e.execute("ALTER INDEX nope REBUILD").unwrap_err();
9731 assert!(
9732 matches!(
9733 &err,
9734 EngineError::Storage(StorageError::IndexNotFound { name }) if name == "nope"
9735 ),
9736 "got: {err}"
9737 );
9738 }
9739
9740 #[test]
9741 fn alter_index_rebuild_on_btree_index_errors() {
9742 let mut e = Engine::new();
9745 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
9746 e.execute("INSERT INTO t VALUES (1)").unwrap();
9747 e.execute("CREATE INDEX t_idx ON t (id)").unwrap();
9748 let err = e.execute("ALTER INDEX t_idx REBUILD").unwrap_err();
9749 assert!(
9750 matches!(&err, EngineError::Storage(StorageError::Unsupported(_))),
9751 "got: {err}"
9752 );
9753 }
9754
9755 #[test]
9756 fn prepared_insert_substitutes_placeholders() {
9757 let mut e = Engine::new();
9763 e.execute("CREATE TABLE t (id INT NOT NULL, name TEXT NOT NULL)")
9764 .unwrap();
9765 let stmt = e.prepare("INSERT INTO t VALUES ($1, $2)").unwrap();
9766 for (id, name) in [(1, "alice"), (2, "bob"), (3, "carol")] {
9767 e.execute_prepared(
9768 stmt.clone(),
9769 &[Value::Int(id), Value::Text(name.into())],
9770 )
9771 .unwrap();
9772 }
9773 let rows_result = e.execute("SELECT id, name FROM t").unwrap();
9775 let QueryResult::Rows { rows, .. } = rows_result else {
9776 panic!("expected Rows")
9777 };
9778 assert_eq!(rows.len(), 3);
9779 }
9780
9781 #[test]
9782 fn prepared_select_with_placeholder_filters_rows() {
9783 let mut e = Engine::new();
9784 e.execute("CREATE TABLE t (id INT NOT NULL, v INT NOT NULL)")
9785 .unwrap();
9786 for i in 0..10_i32 {
9787 e.execute(&alloc::format!("INSERT INTO t VALUES ({i}, {})", i * 7))
9788 .unwrap();
9789 }
9790 let stmt = e
9791 .prepare("SELECT id FROM t WHERE v = $1")
9792 .unwrap();
9793 let QueryResult::Rows { rows, .. } = e
9794 .execute_prepared(stmt, &[Value::Int(35)])
9795 .unwrap()
9796 else {
9797 panic!("expected Rows")
9798 };
9799 assert_eq!(rows.len(), 1);
9801 assert_eq!(rows[0].values[0], Value::Int(5));
9802 }
9803
9804 #[test]
9805 fn prepared_too_few_params_errors() {
9806 let mut e = Engine::new();
9807 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
9808 let stmt = e.prepare("INSERT INTO t VALUES ($1)").unwrap();
9809 let err = e.execute_prepared(stmt, &[]).unwrap_err();
9810 assert!(
9811 matches!(
9812 &err,
9813 EngineError::Eval(EvalError::PlaceholderOutOfRange { n: 1, bound: 0 })
9814 ),
9815 "got: {err}"
9816 );
9817 }
9818
9819 #[test]
9820 fn insert_into_half_column_dim_mismatch_errors() {
9821 let mut e = Engine::new();
9822 e.execute("CREATE TABLE t (v VECTOR(4) USING HALF)")
9823 .unwrap();
9824 let err = e.execute("INSERT INTO t VALUES ([1.0, 2.0])").unwrap_err();
9825 assert!(matches!(
9826 &err,
9827 EngineError::Storage(StorageError::TypeMismatch { .. })
9828 ));
9829 }
9830
9831 #[test]
9832 fn insert_into_sq8_column_dim_mismatch_errors() {
9833 let mut e = Engine::new();
9838 e.execute("CREATE TABLE t (v VECTOR(4) USING SQ8)").unwrap();
9839 let err = e.execute("INSERT INTO t VALUES ([1.0, 2.0])").unwrap_err();
9840 assert!(
9841 matches!(
9842 &err,
9843 EngineError::Storage(StorageError::TypeMismatch { .. })
9844 ),
9845 "got: {err}",
9846 );
9847 }
9848
9849 #[test]
9850 fn create_table_duplicate_errors() {
9851 let mut e = Engine::new();
9852 e.execute("CREATE TABLE foo (a INT)").unwrap();
9853 let err = e.execute("CREATE TABLE foo (a INT)").unwrap_err();
9854 assert!(matches!(
9855 err,
9856 EngineError::Storage(StorageError::DuplicateTable { ref name }) if name == "foo"
9857 ));
9858 }
9859
9860 #[test]
9861 fn insert_into_unknown_table_errors() {
9862 let mut e = Engine::new();
9863 let err = e.execute("INSERT INTO ghost VALUES (1)").unwrap_err();
9864 assert!(matches!(
9865 err,
9866 EngineError::Storage(StorageError::TableNotFound { ref name }) if name == "ghost"
9867 ));
9868 }
9869
9870 #[test]
9871 fn insert_happy_path_reports_one_affected() {
9872 let mut e = Engine::new();
9873 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
9874 let r = e.execute("INSERT INTO foo VALUES (42)").unwrap();
9875 assert_eq!(unwrap_command_ok(&r), 1);
9876 assert_eq!(e.catalog().get("foo").unwrap().row_count(), 1);
9877 }
9878
9879 #[test]
9880 fn insert_arity_mismatch_propagates() {
9881 let mut e = Engine::new();
9882 e.execute("CREATE TABLE foo (a INT, b TEXT)").unwrap();
9883 let err = e.execute("INSERT INTO foo VALUES (1)").unwrap_err();
9884 assert!(matches!(
9885 err,
9886 EngineError::Storage(StorageError::ArityMismatch { .. })
9887 ));
9888 }
9889
9890 #[test]
9891 fn insert_negative_integer_via_unary_minus() {
9892 let mut e = Engine::new();
9893 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
9894 e.execute("INSERT INTO foo VALUES (-7)").unwrap();
9895 let rows = e.catalog().get("foo").unwrap().rows();
9896 assert_eq!(rows[0].values[0], Value::Int(-7));
9897 }
9898
9899 #[test]
9900 fn insert_non_literal_expr_unsupported() {
9901 let mut e = Engine::new();
9902 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
9903 let err = e.execute("INSERT INTO foo VALUES (1 + 2)").unwrap_err();
9904 assert!(matches!(err, EngineError::Unsupported(_)));
9905 }
9906
9907 #[test]
9908 fn select_star_returns_all_rows_in_insertion_order() {
9909 let mut e = Engine::new();
9910 e.execute("CREATE TABLE foo (a INT NOT NULL, b TEXT NOT NULL)")
9911 .unwrap();
9912 e.execute("INSERT INTO foo VALUES (1, 'one')").unwrap();
9913 e.execute("INSERT INTO foo VALUES (2, 'two')").unwrap();
9914 e.execute("INSERT INTO foo VALUES (3, 'three')").unwrap();
9915
9916 let r = e.execute("SELECT * FROM foo").unwrap();
9917 let QueryResult::Rows { columns, rows } = r else {
9918 panic!("expected Rows")
9919 };
9920 assert_eq!(columns.len(), 2);
9921 assert_eq!(columns[0].name, "a");
9922 assert_eq!(rows.len(), 3);
9923 assert_eq!(
9924 rows[1].values,
9925 vec![Value::Int(2), Value::Text("two".into())]
9926 );
9927 }
9928
9929 #[test]
9930 fn select_star_on_empty_table_returns_zero_rows() {
9931 let mut e = Engine::new();
9932 e.execute("CREATE TABLE foo (a INT)").unwrap();
9933 let r = e.execute("SELECT * FROM foo").unwrap();
9934 match r {
9935 QueryResult::Rows { rows, .. } => assert!(rows.is_empty()),
9936 QueryResult::CommandOk { .. } => panic!("expected Rows"),
9937 }
9938 }
9939
9940 fn make_three_row_users(e: &mut Engine) {
9943 e.execute("CREATE TABLE users (id INT NOT NULL, name TEXT NOT NULL, score INT)")
9944 .unwrap();
9945 e.execute("INSERT INTO users VALUES (1, 'alice', 90)")
9946 .unwrap();
9947 e.execute("INSERT INTO users VALUES (2, 'bob', NULL)")
9948 .unwrap();
9949 e.execute("INSERT INTO users VALUES (3, 'cara', 70)")
9950 .unwrap();
9951 }
9952
9953 fn unwrap_rows(r: QueryResult) -> (Vec<ColumnSchema>, Vec<Row>) {
9954 match r {
9955 QueryResult::Rows { columns, rows } => (columns, rows),
9956 QueryResult::CommandOk { .. } => panic!("expected Rows"),
9957 }
9958 }
9959
9960 #[test]
9961 fn where_filter_passes_only_true_rows() {
9962 let mut e = Engine::new();
9963 make_three_row_users(&mut e);
9964 let r = e.execute("SELECT * FROM users WHERE id > 1").unwrap();
9965 let (_, rows) = unwrap_rows(r);
9966 assert_eq!(rows.len(), 2);
9967 assert_eq!(rows[0].values[0], Value::Int(2));
9968 assert_eq!(rows[1].values[0], Value::Int(3));
9969 }
9970
9971 #[test]
9972 fn where_with_null_result_filters_out_row() {
9973 let mut e = Engine::new();
9974 make_three_row_users(&mut e);
9975 let r = e.execute("SELECT * FROM users WHERE score > 80").unwrap();
9977 let (_, rows) = unwrap_rows(r);
9978 assert_eq!(rows.len(), 1);
9979 assert_eq!(rows[0].values[1], Value::Text("alice".into()));
9980 }
9981
9982 #[test]
9983 fn projection_named_columns() {
9984 let mut e = Engine::new();
9985 make_three_row_users(&mut e);
9986 let r = e.execute("SELECT name, score FROM users").unwrap();
9987 let (cols, rows) = unwrap_rows(r);
9988 assert_eq!(cols.len(), 2);
9989 assert_eq!(cols[0].name, "name");
9990 assert_eq!(cols[1].name, "score");
9991 assert_eq!(rows.len(), 3);
9992 assert_eq!(
9993 rows[0].values,
9994 vec![Value::Text("alice".into()), Value::Int(90)]
9995 );
9996 }
9997
9998 #[test]
9999 fn projection_with_column_alias() {
10000 let mut e = Engine::new();
10001 make_three_row_users(&mut e);
10002 let r = e
10003 .execute("SELECT name AS who FROM users WHERE id = 1")
10004 .unwrap();
10005 let (cols, rows) = unwrap_rows(r);
10006 assert_eq!(cols[0].name, "who");
10007 assert_eq!(rows.len(), 1);
10008 assert_eq!(rows[0].values[0], Value::Text("alice".into()));
10009 }
10010
10011 #[test]
10012 fn qualified_column_with_table_alias_resolves() {
10013 let mut e = Engine::new();
10014 make_three_row_users(&mut e);
10015 let r = e
10016 .execute("SELECT u.id, u.name FROM users AS u WHERE u.id < 3")
10017 .unwrap();
10018 let (cols, rows) = unwrap_rows(r);
10019 assert_eq!(cols.len(), 2);
10020 assert_eq!(rows.len(), 2);
10021 }
10022
10023 #[test]
10024 fn qualified_column_with_wrong_alias_errors() {
10025 let mut e = Engine::new();
10026 make_three_row_users(&mut e);
10027 let err = e.execute("SELECT x.id FROM users AS u").unwrap_err();
10028 assert!(matches!(
10029 err,
10030 EngineError::Eval(EvalError::UnknownQualifier { ref qualifier }) if qualifier == "x"
10031 ));
10032 }
10033
10034 #[test]
10035 fn select_unknown_column_errors_in_projection() {
10036 let mut e = Engine::new();
10037 make_three_row_users(&mut e);
10038 let err = e.execute("SELECT ghost FROM users").unwrap_err();
10039 assert!(matches!(
10040 err,
10041 EngineError::Eval(EvalError::ColumnNotFound { ref name }) if name == "ghost"
10042 ));
10043 }
10044
10045 #[test]
10046 fn where_unknown_column_errors() {
10047 let mut e = Engine::new();
10048 make_three_row_users(&mut e);
10049 let err = e
10050 .execute("SELECT * FROM users WHERE ghost = 1")
10051 .unwrap_err();
10052 assert!(matches!(
10053 err,
10054 EngineError::Eval(EvalError::ColumnNotFound { .. })
10055 ));
10056 }
10057
10058 #[test]
10059 fn expression_projection_evaluates_and_renders() {
10060 let mut e = Engine::new();
10063 e.execute("CREATE TABLE t (a INT NOT NULL)").unwrap();
10064 e.execute("INSERT INTO t VALUES (3)").unwrap();
10065 let (_, rows) = unwrap_rows(e.execute("SELECT 1 + 2 FROM t").unwrap());
10066 assert_eq!(rows.len(), 1);
10067 assert_eq!(rows[0].values[0], Value::Int(3));
10070 }
10071
10072 #[test]
10073 fn select_unknown_table_errors() {
10074 let mut e = Engine::new();
10075 let err = e.execute("SELECT * FROM ghost").unwrap_err();
10076 assert!(matches!(
10077 err,
10078 EngineError::Storage(StorageError::TableNotFound { .. })
10079 ));
10080 }
10081
10082 #[test]
10083 fn invalid_sql_returns_parse_error() {
10084 let mut e = Engine::new();
10087 let err = e.execute("THIS_IS_NOT_A_KEYWORD foo bar baz").unwrap_err();
10088 assert!(matches!(err, EngineError::Parse(_)));
10089 }
10090
10091 #[test]
10094 fn create_index_registers_on_table() {
10095 let mut e = Engine::new();
10096 make_three_row_users(&mut e);
10097 e.execute("CREATE INDEX by_name ON users (name)").unwrap();
10098 let t = e.catalog().get("users").unwrap();
10099 assert_eq!(t.indices().len(), 1);
10100 assert_eq!(t.indices()[0].name, "by_name");
10101 }
10102
10103 #[test]
10104 fn create_index_on_unknown_table_errors() {
10105 let mut e = Engine::new();
10106 let err = e.execute("CREATE INDEX i ON ghost (a)").unwrap_err();
10107 assert!(matches!(
10108 err,
10109 EngineError::Storage(StorageError::TableNotFound { .. })
10110 ));
10111 }
10112
10113 #[test]
10114 fn create_index_on_unknown_column_errors() {
10115 let mut e = Engine::new();
10116 make_three_row_users(&mut e);
10117 let err = e.execute("CREATE INDEX i ON users (ghost)").unwrap_err();
10118 assert!(matches!(
10119 err,
10120 EngineError::Storage(StorageError::ColumnNotFound { .. })
10121 ));
10122 }
10123
10124 #[test]
10125 fn select_eq_uses_index_returns_same_rows_as_scan() {
10126 let mut without = Engine::new();
10130 make_three_row_users(&mut without);
10131 let mut with = Engine::new();
10132 make_three_row_users(&mut with);
10133 with.execute("CREATE INDEX by_id ON users (id)").unwrap();
10134
10135 let q = "SELECT * FROM users WHERE id = 2";
10136 let (_, no_idx_rows) = unwrap_rows(without.execute(q).unwrap());
10137 let (_, idx_rows) = unwrap_rows(with.execute(q).unwrap());
10138 assert_eq!(no_idx_rows, idx_rows);
10139 assert_eq!(idx_rows.len(), 1);
10140 }
10141
10142 #[test]
10143 fn select_eq_with_no_matching_index_value_returns_empty() {
10144 let mut e = Engine::new();
10145 make_three_row_users(&mut e);
10146 e.execute("CREATE INDEX by_id ON users (id)").unwrap();
10147 let (_, rows) = unwrap_rows(e.execute("SELECT * FROM users WHERE id = 999").unwrap());
10148 assert_eq!(rows.len(), 0);
10149 }
10150
10151 #[test]
10154 fn begin_sets_in_transaction_flag() {
10155 let mut e = Engine::new();
10156 assert!(!e.in_transaction());
10157 e.execute("BEGIN").unwrap();
10158 assert!(e.in_transaction());
10159 }
10160
10161 #[test]
10162 fn double_begin_errors() {
10163 let mut e = Engine::new();
10164 e.execute("BEGIN").unwrap();
10165 let err = e.execute("BEGIN").unwrap_err();
10166 assert_eq!(err, EngineError::TransactionAlreadyOpen);
10167 }
10168
10169 #[test]
10170 fn commit_without_begin_errors() {
10171 let mut e = Engine::new();
10172 let err = e.execute("COMMIT").unwrap_err();
10173 assert_eq!(err, EngineError::NoActiveTransaction);
10174 }
10175
10176 #[test]
10177 fn rollback_without_begin_errors() {
10178 let mut e = Engine::new();
10179 let err = e.execute("ROLLBACK").unwrap_err();
10180 assert_eq!(err, EngineError::NoActiveTransaction);
10181 }
10182
10183 #[test]
10184 fn commit_applies_shadow_to_committed_catalog() {
10185 let mut e = Engine::new();
10186 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
10187 e.execute("BEGIN").unwrap();
10188 e.execute("INSERT INTO t VALUES (1)").unwrap();
10189 e.execute("INSERT INTO t VALUES (2)").unwrap();
10190 e.execute("COMMIT").unwrap();
10191 assert!(!e.in_transaction());
10192 assert_eq!(e.catalog().get("t").unwrap().row_count(), 2);
10193 }
10194
10195 #[test]
10196 fn rollback_discards_shadow() {
10197 let mut e = Engine::new();
10198 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
10199 e.execute("BEGIN").unwrap();
10200 e.execute("INSERT INTO t VALUES (1)").unwrap();
10201 e.execute("INSERT INTO t VALUES (2)").unwrap();
10202 e.execute("ROLLBACK").unwrap();
10203 assert!(!e.in_transaction());
10204 assert_eq!(e.catalog().get("t").unwrap().row_count(), 0);
10205 }
10206
10207 #[test]
10208 fn select_during_tx_sees_uncommitted_writes_own_session() {
10209 let mut e = Engine::new();
10212 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
10213 e.execute("BEGIN").unwrap();
10214 e.execute("INSERT INTO t VALUES (42)").unwrap();
10215 let (_, rows) = unwrap_rows(e.execute("SELECT * FROM t").unwrap());
10216 assert_eq!(rows.len(), 1);
10217 assert_eq!(rows[0].values[0], Value::Int(42));
10218 }
10219
10220 #[test]
10221 fn snapshot_with_no_users_is_bare_catalog_format() {
10222 let mut e = Engine::new();
10223 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10224 let bytes = e.snapshot();
10225 assert_eq!(
10226 &bytes[..8],
10227 b"SPGDB001",
10228 "must be the bare v3.x catalog magic"
10229 );
10230 let e2 = Engine::restore_envelope(&bytes).unwrap();
10231 assert!(e2.users().is_empty());
10232 assert_eq!(e2.catalog().table_count(), 1);
10233 }
10234
10235 #[test]
10236 fn snapshot_with_users_round_trips_both_via_envelope() {
10237 let mut e = Engine::new();
10238 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10239 e.create_user("alice", "pw1", Role::Admin, [9; 16]).unwrap();
10240 e.create_user("bob", "pw2", Role::ReadOnly, [5; 16])
10241 .unwrap();
10242 let bytes = e.snapshot();
10243 assert_eq!(&bytes[..8], b"SPGENV01", "must be the v4.1 envelope magic");
10244 let e2 = Engine::restore_envelope(&bytes).unwrap();
10245 assert_eq!(e2.users().len(), 2);
10246 assert_eq!(e2.verify_user("alice", "pw1"), Some(Role::Admin));
10247 assert_eq!(e2.verify_user("bob", "pw2"), Some(Role::ReadOnly));
10248 assert_eq!(e2.verify_user("alice", "wrong"), None);
10249 assert_eq!(e2.catalog().table_count(), 1);
10250 }
10251
10252 #[test]
10253 fn ddl_inside_tx_also_rolled_back() {
10254 let mut e = Engine::new();
10255 e.execute("BEGIN").unwrap();
10256 e.execute("CREATE TABLE t (v INT)").unwrap();
10257 e.execute("SELECT * FROM t").unwrap();
10259 e.execute("ROLLBACK").unwrap();
10260 let err = e.execute("SELECT * FROM t").unwrap_err();
10262 assert!(matches!(
10263 err,
10264 EngineError::Storage(StorageError::TableNotFound { .. })
10265 ));
10266 }
10267
10268 #[test]
10271 fn create_publication_lands_in_catalog() {
10272 let mut e = Engine::new();
10273 assert!(e.publications().is_empty());
10274 e.execute("CREATE PUBLICATION pub_a").unwrap();
10275 assert_eq!(e.publications().len(), 1);
10276 assert!(e.publications().contains("pub_a"));
10277 }
10278
10279 #[test]
10280 fn create_publication_duplicate_errors() {
10281 let mut e = Engine::new();
10282 e.execute("CREATE PUBLICATION pub_a").unwrap();
10283 let err = e.execute("CREATE PUBLICATION pub_a").unwrap_err();
10284 assert!(
10285 alloc::format!("{err:?}").contains("DuplicateName"),
10286 "got {err:?}"
10287 );
10288 }
10289
10290 #[test]
10291 fn drop_publication_silent_when_absent() {
10292 let mut e = Engine::new();
10293 let r = e.execute("DROP PUBLICATION nope").unwrap();
10296 match r {
10297 QueryResult::CommandOk { affected, .. } => assert_eq!(affected, 0),
10298 other => panic!("expected CommandOk, got {other:?}"),
10299 }
10300 }
10301
10302 #[test]
10303 fn drop_publication_present_reports_one_affected() {
10304 let mut e = Engine::new();
10305 e.execute("CREATE PUBLICATION pub_a").unwrap();
10306 let r = e.execute("DROP PUBLICATION pub_a").unwrap();
10307 match r {
10308 QueryResult::CommandOk {
10309 affected,
10310 modified_catalog,
10311 } => {
10312 assert_eq!(affected, 1);
10313 assert!(modified_catalog);
10314 }
10315 other => panic!("expected CommandOk, got {other:?}"),
10316 }
10317 assert!(e.publications().is_empty());
10318 }
10319
10320 #[test]
10321 fn publications_persist_across_snapshot_restore() {
10322 let mut e = Engine::new();
10327 e.execute("CREATE PUBLICATION pub_a").unwrap();
10328 e.execute("CREATE PUBLICATION pub_b FOR ALL TABLES").unwrap();
10329 let snap = e.snapshot();
10330 let e2 = Engine::restore_envelope(&snap).unwrap();
10331 assert_eq!(e2.publications().len(), 2);
10332 assert!(e2.publications().contains("pub_a"));
10333 assert!(e2.publications().contains("pub_b"));
10334 }
10335
10336 #[test]
10337 fn create_publication_allowed_inside_transaction() {
10338 let mut e = Engine::new();
10342 e.execute("BEGIN").unwrap();
10343 e.execute("CREATE PUBLICATION pub_a").unwrap();
10344 e.execute("COMMIT").unwrap();
10345 assert!(e.publications().contains("pub_a"));
10346 }
10347
10348 #[test]
10351 fn create_publication_for_table_list_lands_with_scope() {
10352 let mut e = Engine::new();
10353 e.execute("CREATE TABLE t1 (id INT NOT NULL)").unwrap();
10354 e.execute("CREATE TABLE t2 (id INT NOT NULL)").unwrap();
10355 e.execute("CREATE PUBLICATION pub_a FOR TABLE t1, t2")
10356 .unwrap();
10357 let scope = e.publications().get("pub_a").cloned();
10358 let Some(spg_sql::ast::PublicationScope::ForTables(ts)) = scope else {
10359 panic!("expected ForTables scope, got {scope:?}")
10360 };
10361 assert_eq!(ts, alloc::vec!["t1".to_string(), "t2".to_string()]);
10362 }
10363
10364 #[test]
10365 fn create_publication_all_tables_except_lands_with_scope() {
10366 let mut e = Engine::new();
10367 e.execute("CREATE PUBLICATION pub_a FOR ALL TABLES EXCEPT t3")
10368 .unwrap();
10369 let scope = e.publications().get("pub_a").cloned();
10370 let Some(spg_sql::ast::PublicationScope::AllTablesExcept(ts)) = scope else {
10371 panic!("expected AllTablesExcept scope, got {scope:?}")
10372 };
10373 assert_eq!(ts, alloc::vec!["t3".to_string()]);
10374 }
10375
10376 #[test]
10377 fn show_publications_empty_returns_zero_rows() {
10378 let e = Engine::new();
10379 let r = e.execute_readonly("SHOW PUBLICATIONS").unwrap();
10380 let QueryResult::Rows { rows, columns } = r else {
10381 panic!()
10382 };
10383 assert!(rows.is_empty());
10384 assert_eq!(columns.len(), 3);
10385 assert_eq!(columns[0].name, "name");
10386 assert_eq!(columns[1].name, "scope");
10387 assert_eq!(columns[2].name, "table_count");
10388 }
10389
10390 #[test]
10391 fn show_publications_returns_one_row_per_publication_ordered_by_name() {
10392 let mut e = Engine::new();
10393 e.execute("CREATE PUBLICATION z_pub").unwrap();
10394 e.execute("CREATE PUBLICATION a_pub FOR TABLE t1, t2")
10395 .unwrap();
10396 e.execute("CREATE PUBLICATION m_pub FOR ALL TABLES EXCEPT bad")
10397 .unwrap();
10398 let r = e.execute_readonly("SHOW PUBLICATIONS").unwrap();
10399 let QueryResult::Rows { rows, .. } = r else {
10400 panic!()
10401 };
10402 assert_eq!(rows.len(), 3);
10403 let names: Vec<&str> = rows
10405 .iter()
10406 .map(|r| {
10407 if let Value::Text(s) = &r.values[0] {
10408 s.as_str()
10409 } else {
10410 panic!()
10411 }
10412 })
10413 .collect();
10414 assert_eq!(names, alloc::vec!["a_pub", "m_pub", "z_pub"]);
10415 match &rows[0].values[1] {
10417 Value::Text(s) => assert_eq!(s, "FOR TABLE t1, t2"),
10418 other => panic!("expected Text, got {other:?}"),
10419 }
10420 assert_eq!(rows[0].values[2], Value::Int(2));
10421 match &rows[1].values[1] {
10423 Value::Text(s) => assert_eq!(s, "FOR ALL TABLES EXCEPT bad"),
10424 other => panic!("expected Text, got {other:?}"),
10425 }
10426 assert_eq!(rows[1].values[2], Value::Int(1));
10427 match &rows[2].values[1] {
10429 Value::Text(s) => assert_eq!(s, "FOR ALL TABLES"),
10430 other => panic!("expected Text, got {other:?}"),
10431 }
10432 assert_eq!(rows[2].values[2], Value::Null);
10433 }
10434
10435 #[test]
10436 fn for_list_scopes_persist_across_snapshot() {
10437 let mut e = Engine::new();
10440 e.execute("CREATE PUBLICATION p1 FOR TABLE t1, t2").unwrap();
10441 e.execute("CREATE PUBLICATION p2 FOR ALL TABLES EXCEPT bad, worse")
10442 .unwrap();
10443 let snap = e.snapshot();
10444 let e2 = Engine::restore_envelope(&snap).unwrap();
10445 assert_eq!(e2.publications().len(), 2);
10446 let p1 = e2.publications().get("p1").cloned();
10447 let Some(spg_sql::ast::PublicationScope::ForTables(ts)) = p1 else {
10448 panic!("p1 scope lost: {p1:?}")
10449 };
10450 assert_eq!(ts, alloc::vec!["t1".to_string(), "t2".to_string()]);
10451 let p2 = e2.publications().get("p2").cloned();
10452 let Some(spg_sql::ast::PublicationScope::AllTablesExcept(ts)) = p2 else {
10453 panic!("p2 scope lost: {p2:?}")
10454 };
10455 assert_eq!(ts, alloc::vec!["bad".to_string(), "worse".to_string()]);
10456 }
10457
10458 #[test]
10461 fn create_subscription_lands_in_catalog_with_defaults() {
10462 let mut e = Engine::new();
10463 e.execute(
10464 "CREATE SUBSCRIPTION sub_a CONNECTION 'host=127.0.0.1 port=20002' PUBLICATION pub_a",
10465 )
10466 .unwrap();
10467 let s = e.subscriptions().get("sub_a").cloned().expect("present");
10468 assert_eq!(s.conn_str, "host=127.0.0.1 port=20002");
10469 assert_eq!(s.publications, alloc::vec!["pub_a".to_string()]);
10470 assert!(s.enabled);
10471 assert_eq!(s.last_received_pos, 0);
10472 }
10473
10474 #[test]
10475 fn create_subscription_duplicate_name_errors() {
10476 let mut e = Engine::new();
10477 e.execute("CREATE SUBSCRIPTION s CONNECTION 'host=x' PUBLICATION p")
10478 .unwrap();
10479 let err = e
10480 .execute("CREATE SUBSCRIPTION s CONNECTION 'host=y' PUBLICATION p")
10481 .unwrap_err();
10482 assert!(
10483 alloc::format!("{err:?}").contains("DuplicateName"),
10484 "got {err:?}"
10485 );
10486 }
10487
10488 #[test]
10489 fn drop_subscription_silent_when_absent() {
10490 let mut e = Engine::new();
10491 let r = e.execute("DROP SUBSCRIPTION never").unwrap();
10492 match r {
10493 QueryResult::CommandOk { affected, .. } => assert_eq!(affected, 0),
10494 other => panic!("expected CommandOk, got {other:?}"),
10495 }
10496 }
10497
10498 #[test]
10499 fn subscription_advance_updates_last_pos_monotone() {
10500 let mut e = Engine::new();
10501 e.execute("CREATE SUBSCRIPTION s CONNECTION 'h=x' PUBLICATION p")
10502 .unwrap();
10503 assert!(e.subscription_advance("s", 100));
10504 assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 100);
10505 assert!(e.subscription_advance("s", 50)); assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 100);
10507 assert!(e.subscription_advance("s", 200));
10508 assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 200);
10509 assert!(!e.subscription_advance("missing", 1));
10510 }
10511
10512 #[test]
10513 fn show_subscriptions_returns_rows_ordered_by_name() {
10514 let mut e = Engine::new();
10515 e.execute("CREATE SUBSCRIPTION z_sub CONNECTION 'h=x' PUBLICATION p1, p2")
10516 .unwrap();
10517 e.execute("CREATE SUBSCRIPTION a_sub CONNECTION 'h=y' PUBLICATION p3")
10518 .unwrap();
10519 let r = e.execute_readonly("SHOW SUBSCRIPTIONS").unwrap();
10520 let QueryResult::Rows { rows, columns } = r else {
10521 panic!()
10522 };
10523 assert_eq!(rows.len(), 2);
10524 assert_eq!(columns.len(), 5);
10525 assert_eq!(columns[0].name, "name");
10526 assert_eq!(columns[4].name, "last_received_pos");
10527 let names: Vec<&str> = rows
10529 .iter()
10530 .map(|r| {
10531 if let Value::Text(s) = &r.values[0] {
10532 s.as_str()
10533 } else {
10534 panic!()
10535 }
10536 })
10537 .collect();
10538 assert_eq!(names, alloc::vec!["a_sub", "z_sub"]);
10539 assert_eq!(rows[0].values[1], Value::Text("h=y".to_string()));
10541 assert_eq!(rows[0].values[2], Value::Text("p3".to_string()));
10542 assert_eq!(rows[0].values[3], Value::Bool(true));
10543 assert_eq!(rows[0].values[4], Value::BigInt(0));
10544 assert_eq!(rows[1].values[2], Value::Text("p1, p2".to_string()));
10546 }
10547
10548 #[test]
10549 fn subscriptions_persist_across_snapshot_envelope_v4() {
10550 let mut e = Engine::new();
10551 e.execute("CREATE SUBSCRIPTION s1 CONNECTION 'h=A' PUBLICATION p1, p2")
10552 .unwrap();
10553 e.execute("CREATE SUBSCRIPTION s2 CONNECTION 'h=B' PUBLICATION p3")
10554 .unwrap();
10555 e.subscription_advance("s2", 42);
10556 let snap = e.snapshot();
10557 let e2 = Engine::restore_envelope(&snap).unwrap();
10558 assert_eq!(e2.subscriptions().len(), 2);
10559 let s1 = e2.subscriptions().get("s1").unwrap();
10560 assert_eq!(s1.conn_str, "h=A");
10561 assert_eq!(s1.publications, alloc::vec!["p1".to_string(), "p2".to_string()]);
10562 assert_eq!(s1.last_received_pos, 0);
10563 let s2 = e2.subscriptions().get("s2").unwrap();
10564 assert_eq!(s2.last_received_pos, 42);
10565 }
10566
10567 #[test]
10568 fn v3_envelope_loads_with_empty_subscriptions() {
10569 let mut e = Engine::new();
10573 e.execute("CREATE PUBLICATION pub_legacy").unwrap();
10574 let catalog = e.catalog.serialize();
10575 let users = crate::users::serialize_users(&e.users);
10576 let pubs = e.publications.serialize();
10577 let mut buf = Vec::new();
10578 buf.extend_from_slice(b"SPGENV01");
10579 buf.push(3u8); buf.extend_from_slice(&u32::try_from(catalog.len()).unwrap().to_le_bytes());
10581 buf.extend_from_slice(&catalog);
10582 buf.extend_from_slice(&u32::try_from(users.len()).unwrap().to_le_bytes());
10583 buf.extend_from_slice(&users);
10584 buf.extend_from_slice(&u32::try_from(pubs.len()).unwrap().to_le_bytes());
10585 buf.extend_from_slice(&pubs);
10586 let crc = spg_crypto::crc32::crc32(&buf);
10587 buf.extend_from_slice(&crc.to_le_bytes());
10588
10589 let e2 = Engine::restore_envelope(&buf).expect("v3 envelope restores under v4 reader");
10590 assert!(e2.subscriptions().is_empty());
10591 assert!(e2.publications().contains("pub_legacy"));
10592 }
10593
10594 #[test]
10595 fn create_subscription_allowed_inside_transaction() {
10596 let mut e = Engine::new();
10597 e.execute("BEGIN").unwrap();
10598 e.execute("CREATE SUBSCRIPTION s CONNECTION 'h=x' PUBLICATION p")
10599 .unwrap();
10600 e.execute("COMMIT").unwrap();
10601 assert!(e.subscriptions().contains("s"));
10602 }
10603
10604 #[test]
10605 #[test]
10608 fn analyze_populates_histogram_bounds() {
10609 let mut e = Engine::new();
10610 e.execute("CREATE TABLE t (id INT NOT NULL, name TEXT)").unwrap();
10611 for i in 0..50 {
10612 e.execute(&alloc::format!(
10613 "INSERT INTO t VALUES ({i}, 'name{i}')"
10614 ))
10615 .unwrap();
10616 }
10617 e.execute("ANALYZE t").unwrap();
10618 let stats = e.statistics();
10619 let id_stats = stats.get("t", "id").unwrap();
10620 assert!(id_stats.histogram_bounds.len() >= 2);
10621 assert_eq!(id_stats.histogram_bounds.first().unwrap(), "0");
10622 assert_eq!(id_stats.histogram_bounds.last().unwrap(), "49");
10623 assert!((id_stats.null_frac - 0.0).abs() < 1e-6);
10624 assert_eq!(id_stats.n_distinct, 50);
10625 }
10626
10627 #[test]
10628 fn reanalyze_overwrites_prior_stats() {
10629 let mut e = Engine::new();
10630 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10631 for i in 0..10 {
10632 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10633 }
10634 e.execute("ANALYZE t").unwrap();
10635 let n1 = e.statistics().get("t", "id").unwrap().n_distinct;
10636 assert_eq!(n1, 10);
10637 for i in 10..30 {
10638 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10639 }
10640 e.execute("ANALYZE t").unwrap();
10641 let n2 = e.statistics().get("t", "id").unwrap().n_distinct;
10642 assert_eq!(n2, 30);
10643 }
10644
10645 #[test]
10646 fn analyze_unknown_table_errors() {
10647 let mut e = Engine::new();
10648 let err = e.execute("ANALYZE nonexistent").unwrap_err();
10649 assert!(matches!(err, EngineError::Storage(StorageError::TableNotFound { .. })));
10650 }
10651
10652 #[test]
10653 fn bare_analyze_covers_all_user_tables() {
10654 let mut e = Engine::new();
10655 e.execute("CREATE TABLE t1 (id INT NOT NULL)").unwrap();
10656 e.execute("CREATE TABLE t2 (name TEXT NOT NULL)").unwrap();
10657 e.execute("INSERT INTO t1 VALUES (1)").unwrap();
10658 e.execute("INSERT INTO t2 VALUES ('alice')").unwrap();
10659 let r = e.execute("ANALYZE").unwrap();
10660 match r {
10661 QueryResult::CommandOk { affected, modified_catalog } => {
10662 assert_eq!(affected, 2);
10663 assert!(modified_catalog);
10664 }
10665 other => panic!("expected CommandOk, got {other:?}"),
10666 }
10667 assert!(e.statistics().get("t1", "id").is_some());
10668 assert!(e.statistics().get("t2", "name").is_some());
10669 }
10670
10671 #[test]
10672 fn select_from_spg_statistic_returns_rows_per_column() {
10673 let mut e = Engine::new();
10674 e.execute("CREATE TABLE t (id INT NOT NULL, label TEXT)")
10675 .unwrap();
10676 e.execute("INSERT INTO t VALUES (1, 'a')").unwrap();
10677 e.execute("INSERT INTO t VALUES (2, 'b')").unwrap();
10678 e.execute("ANALYZE t").unwrap();
10679 let r = e.execute_readonly("SELECT * FROM spg_statistic").unwrap();
10680 let QueryResult::Rows { rows, columns } = r else {
10681 panic!()
10682 };
10683 assert_eq!(columns.len(), 6);
10685 assert_eq!(columns[0].name, "table_name");
10686 assert_eq!(columns[4].name, "histogram_bounds");
10687 assert_eq!(columns[5].name, "cold_row_count");
10688 assert_eq!(rows.len(), 2, "one row per column of t");
10689 match (&rows[0].values[0], &rows[0].values[1]) {
10691 (Value::Text(t), Value::Text(c)) => {
10692 assert_eq!(t, "t");
10693 assert_eq!(c, "id");
10695 }
10696 _ => panic!(),
10697 }
10698 }
10699
10700 #[test]
10701 fn analyze_skips_vector_columns() {
10702 let mut e = Engine::new();
10705 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(3) NOT NULL)")
10706 .unwrap();
10707 e.execute("INSERT INTO t VALUES (1, [1, 2, 3])").unwrap();
10708 e.execute("ANALYZE t").unwrap();
10709 assert!(e.statistics().get("t", "id").is_some());
10710 assert!(e.statistics().get("t", "v").is_none());
10711 }
10712
10713 #[test]
10714 fn statistics_persist_across_envelope_v5_round_trip() {
10715 let mut e = Engine::new();
10716 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10717 for i in 0..20 {
10718 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10719 }
10720 e.execute("ANALYZE").unwrap();
10721 let snap = e.snapshot();
10722 let e2 = Engine::restore_envelope(&snap).unwrap();
10723 let s = e2.statistics().get("t", "id").unwrap();
10724 assert_eq!(s.n_distinct, 20);
10725 }
10726
10727 #[test]
10730 fn auto_analyze_threshold_fires_after_10pct_of_min_rows_on_small_table() {
10731 let mut e = Engine::new();
10735 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10736 for i in 0..9 {
10737 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10738 }
10739 assert!(e.tables_needing_analyze().is_empty(), "9 < threshold");
10740 e.execute("INSERT INTO t VALUES (9)").unwrap();
10741 let needs = e.tables_needing_analyze();
10742 assert_eq!(needs, alloc::vec!["t".to_string()]);
10743 }
10744
10745 #[test]
10746 fn auto_analyze_threshold_uses_10pct_of_row_count_for_large_tables() {
10747 let mut e = Engine::new();
10753 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10754 for i in 0..1000 {
10755 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10756 }
10757 e.execute("ANALYZE t").unwrap();
10758 assert!(e.tables_needing_analyze().is_empty(), "fresh ANALYZE");
10759 for i in 1000..1050 {
10760 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10761 }
10762 assert!(
10763 e.tables_needing_analyze().is_empty(),
10764 "50 inserts < threshold of ~105"
10765 );
10766 for i in 1050..1200 {
10767 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10768 }
10769 assert_eq!(
10770 e.tables_needing_analyze(),
10771 alloc::vec!["t".to_string()],
10772 "200 inserts > 0.1 × 1200 threshold"
10773 );
10774 }
10775
10776 #[test]
10777 fn auto_analyze_threshold_resets_after_analyze() {
10778 let mut e = Engine::new();
10779 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
10780 for i in 0..200 {
10781 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})")).unwrap();
10782 }
10783 assert!(!e.tables_needing_analyze().is_empty());
10784 e.execute("ANALYZE").unwrap();
10785 assert!(
10786 e.tables_needing_analyze().is_empty(),
10787 "ANALYZE must reset the counter"
10788 );
10789 }
10790
10791 #[test]
10792 fn auto_analyze_threshold_tracks_updates_and_deletes() {
10793 let mut e = Engine::new();
10794 e.execute("CREATE TABLE t (id INT NOT NULL, label TEXT)").unwrap();
10795 for i in 0..50 {
10796 e.execute(&alloc::format!("INSERT INTO t VALUES ({i}, 'x')"))
10797 .unwrap();
10798 }
10799 e.execute("ANALYZE t").unwrap();
10800 e.execute("UPDATE t SET label = 'y' WHERE id < 20").unwrap();
10803 e.execute("DELETE FROM t WHERE id >= 45").unwrap();
10804 assert_eq!(
10805 e.tables_needing_analyze(),
10806 alloc::vec!["t".to_string()]
10807 );
10808 }
10809
10810 #[test]
10811 fn v4_envelope_loads_with_empty_statistics() {
10812 let mut e = Engine::new();
10816 e.create_user("alice", "secret", crate::users::Role::ReadOnly, [0u8; 16])
10817 .unwrap();
10818 let catalog = e.catalog.serialize();
10819 let users = crate::users::serialize_users(&e.users);
10820 let pubs = e.publications.serialize();
10821 let subs = e.subscriptions.serialize();
10822 let mut buf = Vec::new();
10823 buf.extend_from_slice(b"SPGENV01");
10824 buf.push(4u8);
10825 buf.extend_from_slice(&u32::try_from(catalog.len()).unwrap().to_le_bytes());
10826 buf.extend_from_slice(&catalog);
10827 buf.extend_from_slice(&u32::try_from(users.len()).unwrap().to_le_bytes());
10828 buf.extend_from_slice(&users);
10829 buf.extend_from_slice(&u32::try_from(pubs.len()).unwrap().to_le_bytes());
10830 buf.extend_from_slice(&pubs);
10831 buf.extend_from_slice(&u32::try_from(subs.len()).unwrap().to_le_bytes());
10832 buf.extend_from_slice(&subs);
10833 let crc = spg_crypto::crc32::crc32(&buf);
10834 buf.extend_from_slice(&crc.to_le_bytes());
10835 let e2 = Engine::restore_envelope(&buf).expect("v4 envelope restores");
10836 assert!(e2.statistics().is_empty());
10837 }
10838
10839 #[test]
10840 fn v1_v2_envelope_loads_with_empty_publications() {
10841 let mut e = Engine::new();
10848 e.create_user(
10851 "alice",
10852 "secret",
10853 crate::users::Role::ReadOnly,
10854 [0u8; 16],
10855 )
10856 .unwrap();
10857
10858 let catalog = e.catalog.serialize();
10860 let users = crate::users::serialize_users(&e.users);
10861 let mut buf = Vec::new();
10862 buf.extend_from_slice(b"SPGENV01");
10863 buf.push(2u8); buf.extend_from_slice(
10865 &u32::try_from(catalog.len()).unwrap().to_le_bytes(),
10866 );
10867 buf.extend_from_slice(&catalog);
10868 buf.extend_from_slice(
10869 &u32::try_from(users.len()).unwrap().to_le_bytes(),
10870 );
10871 buf.extend_from_slice(&users);
10872 let crc = spg_crypto::crc32::crc32(&buf);
10873 buf.extend_from_slice(&crc.to_le_bytes());
10874
10875 let e2 = Engine::restore_envelope(&buf).expect("v2 envelope restores");
10876 assert!(e2.publications().is_empty());
10877 }
10878}