1#![no_std]
6
7extern crate alloc;
8
9pub mod aggregate;
10pub mod describe;
11pub mod eval;
12pub mod fts;
13pub mod json;
14pub mod memoize;
15pub mod plan_cache;
16pub mod publications;
17pub mod query_stats;
18pub mod reorder;
19pub mod selectivity;
20pub mod statistics;
21pub mod subscriptions;
22pub mod triggers;
23pub mod users;
24
25pub use crate::users::{Role, ScramSecrets, UserError, UserStore};
26
27use alloc::borrow::Cow;
28use alloc::boxed::Box;
29use alloc::collections::BTreeMap;
30use alloc::string::{String, ToString};
31use alloc::vec::Vec;
32use core::fmt;
33
34use spg_sql::ast::{
35 BinOp, ColumnDef, ColumnName, ColumnTypeName, CreateIndexStatement, CreatePublicationStatement,
36 CreateSubscriptionStatement, CreateTableStatement, CreateUserStatement, Expr, FrameBound,
37 FrameKind, FromClause, IndexMethod, InsertStatement, JoinKind, Literal, OrderBy, SelectItem,
38 SelectStatement, Statement, TableRef, UnOp, UnionKind, VecEncoding as SqlVecEncoding,
39 WindowFrame,
40};
41use spg_sql::parser::{self, ParseError};
42use spg_storage::{
43 Catalog, ColumnSchema, CompactReport, DataType, IndexKey, IndexKind, Row, StorageError, Table,
44 TableSchema, Value, VecEncoding,
45};
46
47use crate::eval::{EvalContext, EvalError};
48
49#[derive(Debug, Clone, PartialEq)]
51#[non_exhaustive]
52pub enum QueryResult {
53 CommandOk {
62 affected: usize,
63 modified_catalog: bool,
64 },
65 Rows {
67 columns: Vec<ColumnSchema>,
68 rows: Vec<Row>,
69 },
70}
71
72#[derive(Debug, Clone, PartialEq)]
78#[non_exhaustive]
79pub enum EngineError {
80 Parse(ParseError),
81 Storage(StorageError),
82 Eval(EvalError),
83 Unsupported(String),
85 TransactionAlreadyOpen,
87 NoActiveTransaction,
89 WriteRequired,
94 RowLimitExceeded(usize),
97 Cancelled,
103}
104
105impl fmt::Display for EngineError {
106 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107 match self {
108 Self::Parse(e) => write!(f, "parse: {e}"),
109 Self::Storage(e) => write!(f, "storage: {e}"),
110 Self::Eval(e) => write!(f, "eval: {e}"),
111 Self::Unsupported(s) => write!(f, "unsupported: {s}"),
112 Self::TransactionAlreadyOpen => f.write_str("a transaction is already open"),
113 Self::NoActiveTransaction => f.write_str("no active transaction"),
114 Self::WriteRequired => {
115 f.write_str("statement requires a write lock (use execute, not execute_readonly)")
116 }
117 Self::RowLimitExceeded(n) => {
118 write!(f, "query exceeded max_query_rows={n}")
119 }
120 Self::Cancelled => f.write_str("query cancelled (timeout or client request)"),
121 }
122 }
123}
124
125impl From<ParseError> for EngineError {
126 fn from(e: ParseError) -> Self {
127 Self::Parse(e)
128 }
129}
130impl From<StorageError> for EngineError {
131 fn from(e: StorageError) -> Self {
132 Self::Storage(e)
133 }
134}
135impl From<EvalError> for EngineError {
136 fn from(e: EvalError) -> Self {
137 Self::Eval(e)
138 }
139}
140
141pub type ClockFn = fn() -> i64;
150
151pub type SaltFn = fn() -> [u8; 16];
158
159#[derive(Debug, Clone, Copy)]
170pub struct CancelToken<'a> {
171 flag: Option<&'a core::sync::atomic::AtomicBool>,
172}
173
174impl<'a> CancelToken<'a> {
175 #[must_use]
176 pub const fn none() -> Self {
177 Self { flag: None }
178 }
179
180 #[must_use]
181 pub const fn from_flag(f: &'a core::sync::atomic::AtomicBool) -> Self {
182 Self { flag: Some(f) }
183 }
184
185 #[must_use]
186 pub fn is_cancelled(self) -> bool {
187 self.flag
188 .is_some_and(|f| f.load(core::sync::atomic::Ordering::Relaxed))
189 }
190
191 #[inline]
195 pub fn check(self) -> Result<(), EngineError> {
196 if self.is_cancelled() {
197 Err(EngineError::Cancelled)
198 } else {
199 Ok(())
200 }
201 }
202}
203
204const ENVELOPE_MAGIC: &[u8; 8] = b"SPGENV01";
262const ENVELOPE_VERSION_V1: u8 = 1;
263const ENVELOPE_VERSION_V2: u8 = 2;
264const ENVELOPE_VERSION_V3: u8 = 3;
265const ENVELOPE_VERSION_V4: u8 = 4;
266const ENVELOPE_VERSION_V5: u8 = 5;
267
268fn build_envelope(catalog: &[u8], users: &[u8], pubs: &[u8], subs: &[u8], stats: &[u8]) -> Vec<u8> {
269 let mut out = Vec::with_capacity(
270 8 + 1
271 + 4
272 + catalog.len()
273 + 4
274 + users.len()
275 + 4
276 + pubs.len()
277 + 4
278 + subs.len()
279 + 4
280 + stats.len()
281 + 4,
282 );
283 out.extend_from_slice(ENVELOPE_MAGIC);
284 out.push(ENVELOPE_VERSION_V5);
285 out.extend_from_slice(
286 &u32::try_from(catalog.len())
287 .expect("≤ 4G catalog")
288 .to_le_bytes(),
289 );
290 out.extend_from_slice(catalog);
291 out.extend_from_slice(
292 &u32::try_from(users.len())
293 .expect("≤ 4G users")
294 .to_le_bytes(),
295 );
296 out.extend_from_slice(users);
297 out.extend_from_slice(
298 &u32::try_from(pubs.len())
299 .expect("≤ 4G publications")
300 .to_le_bytes(),
301 );
302 out.extend_from_slice(pubs);
303 out.extend_from_slice(
304 &u32::try_from(subs.len())
305 .expect("≤ 4G subscriptions")
306 .to_le_bytes(),
307 );
308 out.extend_from_slice(subs);
309 out.extend_from_slice(
310 &u32::try_from(stats.len())
311 .expect("≤ 4G statistics")
312 .to_le_bytes(),
313 );
314 out.extend_from_slice(stats);
315 let crc = spg_crypto::crc32::crc32(&out);
316 out.extend_from_slice(&crc.to_le_bytes());
317 out
318}
319
320enum EnvelopeParse<'a> {
327 Bare,
328 Pair {
329 catalog: &'a [u8],
330 users: &'a [u8],
331 publications: Option<&'a [u8]>,
332 subscriptions: Option<&'a [u8]>,
333 statistics: Option<&'a [u8]>,
334 },
335 CrcMismatch {
336 expected: u32,
337 computed: u32,
338 },
339}
340
341fn split_envelope(buf: &[u8]) -> EnvelopeParse<'_> {
346 if buf.len() < 8 + 1 + 4 || &buf[..8] != ENVELOPE_MAGIC {
347 return EnvelopeParse::Bare;
348 }
349 let version = buf[8];
350 if !matches!(
351 version,
352 ENVELOPE_VERSION_V1
353 | ENVELOPE_VERSION_V2
354 | ENVELOPE_VERSION_V3
355 | ENVELOPE_VERSION_V4
356 | ENVELOPE_VERSION_V5
357 ) {
358 return EnvelopeParse::Bare;
359 }
360 let mut p = 9usize;
361 let Some(cat_len_bytes) = buf.get(p..p + 4) else {
362 return EnvelopeParse::Bare;
363 };
364 let Ok(cat_len_arr) = cat_len_bytes.try_into() else {
365 return EnvelopeParse::Bare;
366 };
367 let cat_len = u32::from_le_bytes(cat_len_arr) as usize;
368 p += 4;
369 if p + cat_len + 4 > buf.len() {
370 return EnvelopeParse::Bare;
371 }
372 let catalog = &buf[p..p + cat_len];
373 p += cat_len;
374 let Some(user_len_bytes) = buf.get(p..p + 4) else {
375 return EnvelopeParse::Bare;
376 };
377 let Ok(user_len_arr) = user_len_bytes.try_into() else {
378 return EnvelopeParse::Bare;
379 };
380 let user_len = u32::from_le_bytes(user_len_arr) as usize;
381 p += 4;
382 if p + user_len > buf.len() {
383 return EnvelopeParse::Bare;
384 }
385 let users = &buf[p..p + user_len];
386 p += user_len;
387 let publications = if matches!(
388 version,
389 ENVELOPE_VERSION_V3 | ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5
390 ) {
391 let Some(pubs_len_bytes) = buf.get(p..p + 4) else {
393 return EnvelopeParse::Bare;
394 };
395 let Ok(pubs_len_arr) = pubs_len_bytes.try_into() else {
396 return EnvelopeParse::Bare;
397 };
398 let pubs_len = u32::from_le_bytes(pubs_len_arr) as usize;
399 p += 4;
400 if p + pubs_len > buf.len() {
401 return EnvelopeParse::Bare;
402 }
403 let pubs_slice = &buf[p..p + pubs_len];
404 p += pubs_len;
405 Some(pubs_slice)
406 } else {
407 None
408 };
409 let subscriptions = if matches!(version, ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5) {
410 let Some(subs_len_bytes) = buf.get(p..p + 4) else {
412 return EnvelopeParse::Bare;
413 };
414 let Ok(subs_len_arr) = subs_len_bytes.try_into() else {
415 return EnvelopeParse::Bare;
416 };
417 let subs_len = u32::from_le_bytes(subs_len_arr) as usize;
418 p += 4;
419 if p + subs_len > buf.len() {
420 return EnvelopeParse::Bare;
421 }
422 let subs_slice = &buf[p..p + subs_len];
423 p += subs_len;
424 Some(subs_slice)
425 } else {
426 None
427 };
428 let statistics = if version == ENVELOPE_VERSION_V5 {
429 let Some(stats_len_bytes) = buf.get(p..p + 4) else {
431 return EnvelopeParse::Bare;
432 };
433 let Ok(stats_len_arr) = stats_len_bytes.try_into() else {
434 return EnvelopeParse::Bare;
435 };
436 let stats_len = u32::from_le_bytes(stats_len_arr) as usize;
437 p += 4;
438 if p + stats_len > buf.len() {
439 return EnvelopeParse::Bare;
440 }
441 let stats_slice = &buf[p..p + stats_len];
442 p += stats_len;
443 Some(stats_slice)
444 } else {
445 None
446 };
447 if matches!(
448 version,
449 ENVELOPE_VERSION_V2 | ENVELOPE_VERSION_V3 | ENVELOPE_VERSION_V4 | ENVELOPE_VERSION_V5
450 ) {
451 if p + 4 != buf.len() {
452 return EnvelopeParse::Bare;
453 }
454 let Ok(crc_arr) = buf[p..p + 4].try_into() else {
455 return EnvelopeParse::Bare;
456 };
457 let expected = u32::from_le_bytes(crc_arr);
458 let computed = spg_crypto::crc32::crc32(&buf[..p]);
459 if expected != computed {
460 return EnvelopeParse::CrcMismatch { expected, computed };
461 }
462 } else if p != buf.len() {
463 return EnvelopeParse::Bare;
465 }
466 EnvelopeParse::Pair {
467 catalog,
468 users,
469 publications,
470 subscriptions,
471 statistics,
472 }
473}
474
475#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
485pub struct TxId(pub u64);
486
487pub const IMPLICIT_TX: TxId = TxId(0);
490
491pub const COMPACTION_TARGET_DEFAULT_BYTES: u64 = 4 * 1024 * 1024;
497
498#[derive(Debug, Default, Clone)]
503struct TxState {
504 catalog: Catalog,
509 savepoints: Vec<(String, Catalog)>,
515}
516
517#[derive(Debug, Clone)]
531pub struct CatalogSnapshot {
532 catalog: Catalog,
533 statistics: statistics::Statistics,
534 clock: Option<ClockFn>,
535 max_query_rows: Option<usize>,
536}
537
538#[derive(Debug, Default)]
539pub struct Engine {
540 catalog: Catalog,
543 tx_catalogs: BTreeMap<TxId, TxState>,
548 current_tx: Option<TxId>,
553 next_tx_id: u64,
556 clock: Option<ClockFn>,
559 salt_fn: Option<SaltFn>,
563 max_query_rows: Option<usize>,
569 users: UserStore,
575 publications: publications::Publications,
579 subscriptions: subscriptions::Subscriptions,
583 statistics: statistics::Statistics,
587 plan_cache: plan_cache::PlanCache,
591 query_stats: query_stats::QueryStats,
595 activity_provider: Option<ActivityProvider>,
602 audit_chain_provider: Option<AuditChainProvider>,
607 audit_verifier: Option<AuditVerifier>,
608 slow_query_threshold_us: Option<u64>,
614 slow_query_logger: Option<SlowQueryLogger>,
615 session_params: BTreeMap<String, String>,
622 trigger_recursion_depth: u32,
630}
631
632const MAX_TRIGGER_RECURSION: u32 = 16;
636
637pub type SlowQueryLogger = fn(&str, u64);
641
642fn render_create_table(name: &str, columns: &[ColumnSchema]) -> String {
647 let mut out = alloc::format!("CREATE TABLE {name} (");
648 for (i, col) in columns.iter().enumerate() {
649 if i > 0 {
650 out.push_str(", ");
651 }
652 out.push_str(&col.name);
653 out.push(' ');
654 out.push_str(&render_data_type(col.ty));
655 if !col.nullable {
656 out.push_str(" NOT NULL");
657 }
658 if col.auto_increment {
659 out.push_str(" AUTO_INCREMENT");
660 }
661 }
662 out.push(')');
663 out
664}
665
666fn render_data_type(ty: DataType) -> String {
667 match ty {
668 DataType::SmallInt => "SMALLINT".into(),
669 DataType::Int => "INT".into(),
670 DataType::BigInt => "BIGINT".into(),
671 DataType::Float => "FLOAT".into(),
672 DataType::Text => "TEXT".into(),
673 DataType::Varchar(n) => alloc::format!("VARCHAR({n})"),
674 DataType::Char(n) => alloc::format!("CHAR({n})"),
675 DataType::Bool => "BOOL".into(),
676 DataType::Vector { dim, encoding } => match encoding {
677 spg_storage::VecEncoding::F32 => alloc::format!("VECTOR({dim})"),
678 spg_storage::VecEncoding::Sq8 => alloc::format!("VECTOR({dim}) USING SQ8"),
679 spg_storage::VecEncoding::F16 => alloc::format!("VECTOR({dim}) USING HALF"),
680 },
681 DataType::Numeric { precision, scale } => {
682 alloc::format!("NUMERIC({precision},{scale})")
683 }
684 DataType::Date => "DATE".into(),
685 DataType::Timestamp => "TIMESTAMP".into(),
686 DataType::Interval => "INTERVAL".into(),
687 DataType::Json => "JSON".into(),
688 DataType::Jsonb => "JSONB".into(),
689 DataType::Timestamptz => "TIMESTAMPTZ".into(),
690 DataType::Bytes => "BYTEA".into(),
691 DataType::TextArray => "TEXT[]".into(),
692 DataType::IntArray => "INT[]".into(),
693 DataType::BigIntArray => "BIGINT[]".into(),
694 DataType::TsVector => "TSVECTOR".into(),
695 DataType::TsQuery => "TSQUERY".into(),
696 }
697}
698
699#[derive(Debug, Clone)]
703pub struct ActivityRow {
704 pub pid: u32,
705 pub user: String,
706 pub started_at_us: i64,
707 pub current_sql: String,
708 pub wait_event: String,
709 pub elapsed_us: i64,
710 pub in_transaction: bool,
711}
712
713pub type ActivityProvider = fn() -> Vec<ActivityRow>;
716
717#[derive(Debug, Clone)]
720pub struct AuditRow {
721 pub seq: i64,
722 pub ts_ms: i64,
723 pub prev_hash_hex: String,
724 pub entry_hash_hex: String,
725 pub sql: String,
726}
727
728pub type AuditChainProvider = fn() -> Vec<AuditRow>;
733pub type AuditVerifier = fn() -> (i64, i64);
734
735impl Engine {
736 pub fn new() -> Self {
737 Self {
738 catalog: Catalog::new(),
739 tx_catalogs: BTreeMap::new(),
740 current_tx: None,
741 next_tx_id: 1,
742 clock: None,
743 salt_fn: None,
744 max_query_rows: None,
745 users: UserStore::new(),
746 publications: publications::Publications::new(),
747 subscriptions: subscriptions::Subscriptions::new(),
748 statistics: statistics::Statistics::new(),
749 plan_cache: plan_cache::PlanCache::new(),
750 query_stats: query_stats::QueryStats::new(),
751 activity_provider: None,
752 audit_chain_provider: None,
753 audit_verifier: None,
754 slow_query_threshold_us: None,
755 slow_query_logger: None,
756 session_params: BTreeMap::new(),
757 trigger_recursion_depth: 0,
758 }
759 }
760
761 #[must_use]
770 pub fn clone_snapshot(&self) -> CatalogSnapshot {
771 CatalogSnapshot {
772 catalog: self.active_catalog().clone(),
773 statistics: self.statistics.clone(),
774 clock: self.clock,
775 max_query_rows: self.max_query_rows,
776 }
777 }
778
779 pub fn execute_readonly_on_snapshot(
787 snapshot: &CatalogSnapshot,
788 sql: &str,
789 ) -> Result<QueryResult, EngineError> {
790 Self::execute_readonly_on_snapshot_with_cancel(snapshot, sql, CancelToken::none())
791 }
792
793 pub fn execute_readonly_on_snapshot_with_cancel(
800 snapshot: &CatalogSnapshot,
801 sql: &str,
802 cancel: CancelToken<'_>,
803 ) -> Result<QueryResult, EngineError> {
804 let transient = Engine {
805 catalog: snapshot.catalog.clone(),
806 statistics: snapshot.statistics.clone(),
807 clock: snapshot.clock,
808 max_query_rows: snapshot.max_query_rows,
809 ..Engine::default()
810 };
811 transient.execute_readonly_with_cancel(sql, cancel)
812 }
813
814 pub fn restore(catalog: Catalog) -> Self {
817 Self {
818 catalog,
819 tx_catalogs: BTreeMap::new(),
820 current_tx: None,
821 next_tx_id: 1,
822 clock: None,
823 salt_fn: None,
824 max_query_rows: None,
825 users: UserStore::new(),
826 publications: publications::Publications::new(),
827 subscriptions: subscriptions::Subscriptions::new(),
828 statistics: statistics::Statistics::new(),
829 plan_cache: plan_cache::PlanCache::new(),
830 query_stats: query_stats::QueryStats::new(),
831 activity_provider: None,
832 audit_chain_provider: None,
833 audit_verifier: None,
834 slow_query_threshold_us: None,
835 slow_query_logger: None,
836 session_params: BTreeMap::new(),
837 trigger_recursion_depth: 0,
838 }
839 }
840
841 pub fn restore_envelope(buf: &[u8]) -> Result<Self, EngineError> {
848 match split_envelope(buf) {
849 EnvelopeParse::Pair {
850 catalog: catalog_bytes,
851 users: user_bytes,
852 publications: pub_bytes,
853 subscriptions: sub_bytes,
854 statistics: stats_bytes,
855 } => {
856 let catalog = Catalog::deserialize(catalog_bytes).map_err(EngineError::Storage)?;
857 let users = users::deserialize_users(user_bytes)
858 .map_err(|e| EngineError::Unsupported(alloc::format!("users restore: {e}")))?;
859 let publications = match pub_bytes {
860 Some(b) => publications::Publications::deserialize(b).map_err(|e| {
861 EngineError::Unsupported(alloc::format!("publications restore: {e:?}"))
862 })?,
863 None => publications::Publications::new(),
864 };
865 let subscriptions = match sub_bytes {
866 Some(b) => subscriptions::Subscriptions::deserialize(b).map_err(|e| {
867 EngineError::Unsupported(alloc::format!("subscriptions restore: {e:?}"))
868 })?,
869 None => subscriptions::Subscriptions::new(),
870 };
871 let statistics = match stats_bytes {
872 Some(b) => statistics::Statistics::deserialize(b).map_err(|e| {
873 EngineError::Unsupported(alloc::format!("statistics restore: {e:?}"))
874 })?,
875 None => statistics::Statistics::new(),
876 };
877 Ok(Self {
878 catalog,
879 tx_catalogs: BTreeMap::new(),
880 current_tx: None,
881 next_tx_id: 1,
882 clock: None,
883 salt_fn: None,
884 max_query_rows: None,
885 users,
886 publications,
887 subscriptions,
888 statistics,
889 plan_cache: plan_cache::PlanCache::new(),
890 query_stats: query_stats::QueryStats::new(),
891 activity_provider: None,
892 audit_chain_provider: None,
893 audit_verifier: None,
894 slow_query_threshold_us: None,
895 slow_query_logger: None,
896 session_params: BTreeMap::new(),
897 trigger_recursion_depth: 0,
898 })
899 }
900 EnvelopeParse::CrcMismatch { expected, computed } => {
901 Err(EngineError::Storage(StorageError::Corrupt(alloc::format!(
902 "snapshot envelope CRC32 mismatch (expected={expected:#010x}, computed={computed:#010x})"
903 ))))
904 }
905 EnvelopeParse::Bare => {
906 let catalog = Catalog::deserialize(buf).map_err(EngineError::Storage)?;
907 Ok(Self::restore(catalog))
908 }
909 }
910 }
911
912 pub const fn users(&self) -> &UserStore {
913 &self.users
914 }
915
916 pub fn create_user(
920 &mut self,
921 name: &str,
922 password: &str,
923 role: Role,
924 salt: [u8; 16],
925 ) -> Result<(), UserError> {
926 self.users.create(name, password, role, salt)?;
927 let scram_salt = self.salt_fn.map_or_else(
933 || {
934 let mut s = [0u8; users::SCRAM_SALT_LEN];
935 let digest = spg_crypto::hash(name.as_bytes());
936 s.copy_from_slice(&digest[16..32]);
939 s
940 },
941 |f| f(),
942 );
943 self.users
944 .enable_scram(name, password, scram_salt, users::SCRAM_DEFAULT_ITERS)?;
945 Ok(())
946 }
947
948 pub fn drop_user(&mut self, name: &str) -> Result<(), UserError> {
949 self.users.drop(name)
950 }
951
952 pub fn verify_user(&self, name: &str, password: &str) -> Option<Role> {
953 self.users.verify(name, password)
954 }
955
956 #[must_use]
959 pub const fn with_clock(mut self, clock: ClockFn) -> Self {
960 self.clock = Some(clock);
961 self
962 }
963
964 #[must_use]
967 pub const fn with_salt_fn(mut self, f: SaltFn) -> Self {
968 self.salt_fn = Some(f);
969 self
970 }
971
972 #[must_use]
978 pub const fn with_max_query_rows(mut self, n: usize) -> Self {
979 self.max_query_rows = Some(n);
980 self
981 }
982
983 pub const fn catalog(&self) -> &Catalog {
987 &self.catalog
988 }
989
990 pub fn snapshot(&self) -> Vec<u8> {
998 if self.users.is_empty()
999 && self.publications.is_empty()
1000 && self.subscriptions.is_empty()
1001 && self.statistics.is_empty()
1002 {
1003 self.catalog.serialize()
1004 } else {
1005 build_envelope(
1006 &self.catalog.serialize(),
1007 &users::serialize_users(&self.users),
1008 &self.publications.serialize(),
1009 &self.subscriptions.serialize(),
1010 &self.statistics.serialize(),
1011 )
1012 }
1013 }
1014
1015 pub fn in_transaction(&self) -> bool {
1020 !self.tx_catalogs.is_empty()
1021 }
1022
1023 pub fn alloc_tx_id(&mut self) -> TxId {
1032 let id = TxId(self.next_tx_id);
1033 self.next_tx_id = self.next_tx_id.saturating_add(1);
1034 id
1035 }
1036
1037 pub fn replace_catalog(&mut self, catalog: Catalog) {
1057 self.catalog = catalog;
1058 }
1059
1060 pub fn freeze_oldest_to_cold(
1068 &mut self,
1069 table_name: &str,
1070 index_name: &str,
1071 max_rows: usize,
1072 ) -> Result<spg_storage::FreezeReport, EngineError> {
1073 let report = self
1074 .active_catalog_mut()
1075 .freeze_oldest_to_cold(table_name, index_name, max_rows)
1076 .map_err(EngineError::Storage)?;
1077 if let Some(t) = self.active_catalog_mut().get_mut(table_name) {
1078 t.mark_cold_row_count_stale();
1079 }
1080 Ok(report)
1081 }
1082
1083 pub fn receive_cold_segment(
1097 &mut self,
1098 segment_id: u32,
1099 bytes: Vec<u8>,
1100 ) -> Result<(), EngineError> {
1101 let mut new_cat = self.catalog.clone();
1102 match new_cat.load_segment_bytes_at(segment_id, bytes) {
1103 Ok(()) => {
1104 self.replace_catalog(new_cat);
1105 Ok(())
1106 }
1107 Err(StorageError::Corrupt(msg)) if msg.contains("already occupied") => Ok(()),
1108 Err(e) => Err(EngineError::Storage(e)),
1109 }
1110 }
1111
1112 pub fn compact_cold_segments_with_target(
1126 &mut self,
1127 target_segment_bytes: u64,
1128 ) -> Result<Vec<(String, String, CompactReport)>, EngineError> {
1129 let table_names = self.active_catalog().table_names();
1130 let mut reports: Vec<(String, String, CompactReport)> = Vec::new();
1131 for tname in table_names {
1132 if is_internal_table_name(&tname) {
1133 continue;
1134 }
1135 let idx_names: Vec<String> = {
1136 let Some(t) = self.active_catalog().get(&tname) else {
1137 continue;
1138 };
1139 t.indices()
1140 .iter()
1141 .filter(|i| matches!(i.kind, IndexKind::BTree(_)))
1142 .map(|i| i.name.clone())
1143 .collect()
1144 };
1145 for iname in idx_names {
1146 let report = self
1147 .active_catalog_mut()
1148 .compact_cold_segments(&tname, &iname, target_segment_bytes)
1149 .map_err(EngineError::Storage)?;
1150 if report.merged_segment_id.is_some() {
1151 if let Some(t) = self.active_catalog_mut().get_mut(&tname) {
1152 t.mark_cold_row_count_stale();
1153 }
1154 reports.push((tname.clone(), iname, report));
1155 }
1156 }
1157 }
1158 Ok(reports)
1159 }
1160
1161 fn active_catalog(&self) -> &Catalog {
1162 match self.current_tx {
1163 Some(t) => self
1164 .tx_catalogs
1165 .get(&t)
1166 .map_or(&self.catalog, |s| &s.catalog),
1167 None => &self.catalog,
1168 }
1169 }
1170
1171 fn snapshot_row_triggers(
1179 &self,
1180 table: &str,
1181 event: &str,
1182 timing: &str,
1183 ) -> Vec<spg_storage::FunctionDef> {
1184 let cat = self.active_catalog();
1185 cat.triggers()
1186 .iter()
1187 .filter(|t| {
1188 t.table == table
1189 && t.timing.eq_ignore_ascii_case(timing)
1190 && t.for_each.eq_ignore_ascii_case("row")
1191 && t.events.iter().any(|e| e.eq_ignore_ascii_case(event))
1192 })
1193 .filter_map(|t| cat.functions().get(&t.function).cloned())
1194 .collect()
1195 }
1196
1197 fn snapshot_update_row_triggers(
1202 &self,
1203 table: &str,
1204 timing: &str,
1205 ) -> Vec<(spg_storage::FunctionDef, Vec<String>)> {
1206 let cat = self.active_catalog();
1207 cat.triggers()
1208 .iter()
1209 .filter(|t| {
1210 t.table == table
1211 && t.timing.eq_ignore_ascii_case(timing)
1212 && t.for_each.eq_ignore_ascii_case("row")
1213 && t.events.iter().any(|e| e.eq_ignore_ascii_case("UPDATE"))
1214 })
1215 .filter_map(|t| {
1216 cat.functions()
1217 .get(&t.function)
1218 .cloned()
1219 .map(|fd| (fd, t.update_columns.clone()))
1220 })
1221 .collect()
1222 }
1223
1224 fn execute_deferred_trigger_stmts(
1233 &mut self,
1234 deferred: Vec<triggers::DeferredEmbeddedStmt>,
1235 cancel: CancelToken<'_>,
1236 ) -> Result<(), EngineError> {
1237 for d in deferred {
1238 if self.trigger_recursion_depth >= MAX_TRIGGER_RECURSION {
1239 return Err(EngineError::Storage(StorageError::Corrupt(alloc::format!(
1240 "trigger embedded SQL recursion depth {} exceeded (trigger function \
1241 {:?} would push past the {} cap — check for trigger cycles)",
1242 self.trigger_recursion_depth,
1243 d.function,
1244 MAX_TRIGGER_RECURSION,
1245 ))));
1246 }
1247 self.trigger_recursion_depth += 1;
1248 let res = self.execute_stmt_with_cancel(d.stmt, cancel);
1249 self.trigger_recursion_depth -= 1;
1250 res?;
1251 }
1252 Ok(())
1253 }
1254
1255 fn active_catalog_mut(&mut self) -> &mut Catalog {
1256 let tx = self.current_tx;
1257 match tx {
1258 Some(t) => match self.tx_catalogs.get_mut(&t) {
1259 Some(s) => &mut s.catalog,
1260 None => &mut self.catalog,
1261 },
1262 None => &mut self.catalog,
1263 }
1264 }
1265
1266 pub fn execute_readonly(&self, sql: &str) -> Result<QueryResult, EngineError> {
1278 self.execute_readonly_with_cancel(sql, CancelToken::none())
1279 }
1280
1281 pub fn execute_readonly_with_cancel(
1287 &self,
1288 sql: &str,
1289 cancel: CancelToken<'_>,
1290 ) -> Result<QueryResult, EngineError> {
1291 cancel.check()?;
1292 let mut stmt = parser::parse_statement(sql)?;
1293 let now_micros = self.clock.map(|f| f());
1294 rewrite_clock_calls(&mut stmt, now_micros);
1295 if let Statement::Select(s) = &mut stmt {
1296 resolve_order_by_position(s);
1297 reorder::reorder_joins(s, &self.catalog, &self.statistics);
1299 }
1300 let result = match stmt {
1301 Statement::Select(s) => self.exec_select_cancel(&s, cancel),
1302 Statement::ShowTables => Ok(self.exec_show_tables()),
1303 Statement::ShowColumns(table) => self.exec_show_columns(&table),
1304 Statement::ShowUsers => Ok(self.exec_show_users()),
1305 Statement::ShowPublications => Ok(self.exec_show_publications()),
1306 Statement::ShowSubscriptions => Ok(self.exec_show_subscriptions()),
1307 Statement::WaitForWalPosition { .. } => Err(EngineError::Unsupported(
1308 "WAIT FOR WAL POSITION must be handled by the server layer".into(),
1309 )),
1310 Statement::Explain(e) => self.exec_explain(&e, cancel),
1311 _ => Err(EngineError::WriteRequired),
1312 };
1313 self.enforce_row_limit(result)
1314 }
1315
1316 fn enforce_row_limit(
1320 &self,
1321 result: Result<QueryResult, EngineError>,
1322 ) -> Result<QueryResult, EngineError> {
1323 if let (Ok(QueryResult::Rows { rows, .. }), Some(cap)) = (&result, self.max_query_rows)
1324 && rows.len() > cap
1325 {
1326 return Err(EngineError::RowLimitExceeded(cap));
1327 }
1328 result
1329 }
1330
1331 pub fn execute(&mut self, sql: &str) -> Result<QueryResult, EngineError> {
1332 self.execute_in_with_cancel(sql, IMPLICIT_TX, CancelToken::none())
1333 }
1334
1335 pub fn execute_with_cancel(
1340 &mut self,
1341 sql: &str,
1342 cancel: CancelToken<'_>,
1343 ) -> Result<QueryResult, EngineError> {
1344 self.execute_in_with_cancel(sql, IMPLICIT_TX, cancel)
1345 }
1346
1347 pub fn execute_in(&mut self, sql: &str, tx_id: TxId) -> Result<QueryResult, EngineError> {
1354 self.execute_in_with_cancel(sql, tx_id, CancelToken::none())
1355 }
1356
1357 pub fn execute_in_with_cancel(
1363 &mut self,
1364 sql: &str,
1365 tx_id: TxId,
1366 cancel: CancelToken<'_>,
1367 ) -> Result<QueryResult, EngineError> {
1368 let saved = self.current_tx;
1369 self.current_tx = Some(tx_id);
1370 let result = self.execute_inner_with_cancel(sql, cancel);
1371 self.current_tx = saved;
1372 result
1373 }
1374
1375 pub fn prepare(&self, sql: &str) -> Result<Statement, ParseError> {
1387 let mut stmt = parser::parse_statement(sql)?;
1388 let now_micros = self.clock.map(|f| f());
1389 rewrite_clock_calls(&mut stmt, now_micros);
1390 if let Statement::Select(s) = &mut stmt {
1391 expand_group_by_all(s);
1395 resolve_order_by_position(s);
1396 reorder::reorder_joins(s, &self.catalog, &self.statistics);
1399 }
1400 Ok(stmt)
1401 }
1402
1403 pub fn prepare_cached(&mut self, sql: &str) -> Result<Statement, ParseError> {
1415 let current_version = self.statistics.version();
1418 if let Some(plan) = self.plan_cache.get(sql) {
1419 if plan.statistics_version == current_version {
1420 return Ok(plan.stmt.clone());
1421 }
1422 }
1424 self.plan_cache.evict(sql);
1425 let stmt = self.prepare(sql)?;
1426 let source_tables = plan_cache::collect_source_tables(&stmt);
1427 let plan = plan_cache::PreparedPlan {
1428 stmt: stmt.clone(),
1429 statistics_version: current_version,
1430 source_tables,
1431 describe_columns: alloc::vec::Vec::new(),
1432 };
1433 self.plan_cache.insert(String::from(sql), plan);
1434 Ok(stmt)
1435 }
1436
1437 pub fn plan_cache(&self) -> &plan_cache::PlanCache {
1439 &self.plan_cache
1440 }
1441
1442 pub fn plan_cache_mut(&mut self) -> &mut plan_cache::PlanCache {
1444 &mut self.plan_cache
1445 }
1446
1447 pub fn describe_prepared(&self, stmt: &Statement) -> (Vec<u32>, Vec<ColumnSchema>) {
1453 describe::describe_prepared(stmt, self.active_catalog())
1454 }
1455
1456 pub fn execute_prepared(
1466 &mut self,
1467 mut stmt: Statement,
1468 params: &[Value],
1469 ) -> Result<QueryResult, EngineError> {
1470 substitute_placeholders(&mut stmt, params)?;
1471 self.execute_stmt_with_cancel(stmt, CancelToken::none())
1472 }
1473
1474 fn execute_inner_with_cancel(
1475 &mut self,
1476 sql: &str,
1477 cancel: CancelToken<'_>,
1478 ) -> Result<QueryResult, EngineError> {
1479 cancel.check()?;
1480 let stmt = self.prepare(sql)?;
1481 let start_us = self.clock.map(|f| f());
1485 let result = self.execute_stmt_with_cancel(stmt, cancel);
1486 if let (Some(t0), Ok(_)) = (start_us, &result) {
1487 let now = self.clock.map_or(t0, |f| f());
1488 let elapsed = now.saturating_sub(t0).max(0) as u64;
1489 self.query_stats.record(sql, elapsed, now as u64);
1490 if let (Some(threshold), Some(logger)) =
1493 (self.slow_query_threshold_us, self.slow_query_logger)
1494 && elapsed >= threshold
1495 {
1496 logger(sql, elapsed);
1497 }
1498 }
1499 result
1500 }
1501
1502 fn execute_stmt_with_cancel(
1503 &mut self,
1504 stmt: Statement,
1505 cancel: CancelToken<'_>,
1506 ) -> Result<QueryResult, EngineError> {
1507 cancel.check()?;
1508 let result = match stmt {
1509 Statement::CreateTable(s) => self.exec_create_table(s),
1510 Statement::CreateExtension(_) => Ok(QueryResult::CommandOk {
1514 affected: 0,
1515 modified_catalog: false,
1516 }),
1517 Statement::DoBlock => Ok(QueryResult::CommandOk {
1520 affected: 0,
1521 modified_catalog: false,
1522 }),
1523 Statement::CreateIndex(s) => self.exec_create_index(s),
1524 Statement::Insert(s) => self.exec_insert(s),
1525 Statement::Update(s) => self.exec_update_cancel(&s, cancel),
1526 Statement::Delete(s) => self.exec_delete_cancel(&s, cancel),
1527 Statement::Select(s) => self.exec_select_cancel(&s, cancel),
1528 Statement::Begin => self.exec_begin(),
1529 Statement::Commit => self.exec_commit(),
1530 Statement::Rollback => self.exec_rollback(),
1531 Statement::Savepoint(name) => self.exec_savepoint(name),
1532 Statement::RollbackToSavepoint(name) => self.exec_rollback_to_savepoint(&name),
1533 Statement::ReleaseSavepoint(name) => self.exec_release_savepoint(&name),
1534 Statement::ShowTables => Ok(self.exec_show_tables()),
1535 Statement::ShowColumns(table) => self.exec_show_columns(&table),
1536 Statement::ShowUsers => Ok(self.exec_show_users()),
1537 Statement::ShowPublications => Ok(self.exec_show_publications()),
1538 Statement::ShowSubscriptions => Ok(self.exec_show_subscriptions()),
1539 Statement::CreateUser(s) => self.exec_create_user(&s),
1540 Statement::DropUser(name) => self.exec_drop_user(&name),
1541 Statement::Explain(e) => self.exec_explain(&e, cancel),
1542 Statement::AlterIndex(s) => self.exec_alter_index(s),
1543 Statement::AlterTable(s) => self.exec_alter_table(s),
1544 Statement::CreatePublication(s) => self.exec_create_publication(s),
1545 Statement::DropPublication(name) => self.exec_drop_publication(&name),
1546 Statement::CreateSubscription(s) => self.exec_create_subscription(s),
1547 Statement::DropSubscription(name) => self.exec_drop_subscription(&name),
1548 Statement::WaitForWalPosition { .. } => Err(EngineError::Unsupported(
1555 "WAIT FOR WAL POSITION must be handled by the server layer".into(),
1556 )),
1557 Statement::Analyze(target) => self.exec_analyze(target.as_deref()),
1559 Statement::CompactColdSegments => self.exec_compact_cold_segments(),
1561 Statement::SetParameter { name, value } => {
1566 self.set_session_param(name, value);
1567 Ok(QueryResult::CommandOk {
1568 affected: 0,
1569 modified_catalog: false,
1570 })
1571 }
1572 Statement::CreateFunction(s) => self.exec_create_function(s),
1576 Statement::CreateTrigger(s) => self.exec_create_trigger(s),
1577 Statement::DropTrigger {
1578 name,
1579 table,
1580 if_exists,
1581 } => self.exec_drop_trigger(&name, &table, if_exists),
1582 Statement::DropFunction { name, if_exists } => {
1583 self.exec_drop_function(&name, if_exists)
1584 }
1585 Statement::ResetParameter(target) => {
1586 match target {
1587 None => self.session_params.clear(),
1588 Some(name) => {
1589 self.session_params.remove(&name.to_ascii_lowercase());
1590 }
1591 }
1592 Ok(QueryResult::CommandOk {
1593 affected: 0,
1594 modified_catalog: false,
1595 })
1596 }
1597 };
1598 self.enforce_row_limit(result)
1599 }
1600
1601 fn exec_create_publication(
1609 &mut self,
1610 s: CreatePublicationStatement,
1611 ) -> Result<QueryResult, EngineError> {
1612 self.publications
1618 .create(s.name, s.scope)
1619 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE PUBLICATION: {e:?}")))?;
1620 Ok(QueryResult::CommandOk {
1621 affected: 1,
1622 modified_catalog: true,
1623 })
1624 }
1625
1626 fn exec_drop_publication(&mut self, name: &str) -> Result<QueryResult, EngineError> {
1631 let removed = self.publications.drop(name);
1632 Ok(QueryResult::CommandOk {
1633 affected: usize::from(removed),
1634 modified_catalog: removed,
1635 })
1636 }
1637
1638 pub const fn publications(&self) -> &publications::Publications {
1643 &self.publications
1644 }
1645
1646 fn exec_create_subscription(
1651 &mut self,
1652 s: CreateSubscriptionStatement,
1653 ) -> Result<QueryResult, EngineError> {
1654 let sub = subscriptions::Subscription {
1658 conn_str: s.conn_str,
1659 publications: s.publications,
1660 enabled: true,
1661 last_received_pos: 0,
1662 };
1663 self.subscriptions
1664 .create(s.name, sub)
1665 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE SUBSCRIPTION: {e:?}")))?;
1666 Ok(QueryResult::CommandOk {
1667 affected: 1,
1668 modified_catalog: true,
1669 })
1670 }
1671
1672 fn exec_drop_subscription(&mut self, name: &str) -> Result<QueryResult, EngineError> {
1680 let removed = self.subscriptions.drop(name);
1681 Ok(QueryResult::CommandOk {
1682 affected: usize::from(removed),
1683 modified_catalog: removed,
1684 })
1685 }
1686
1687 pub const fn subscriptions(&self) -> &subscriptions::Subscriptions {
1692 &self.subscriptions
1693 }
1694
1695 pub fn subscription_advance(&mut self, name: &str, pos: u64) -> bool {
1701 self.subscriptions.update_last_received_pos(name, pos)
1702 }
1703
1704 fn exec_show_subscriptions(&self) -> QueryResult {
1710 let columns = alloc::vec![
1711 ColumnSchema::new("name", DataType::Text, false),
1712 ColumnSchema::new("conn_str", DataType::Text, false),
1713 ColumnSchema::new("publications", DataType::Text, false),
1714 ColumnSchema::new("enabled", DataType::Bool, false),
1715 ColumnSchema::new("last_received_pos", DataType::BigInt, false),
1716 ];
1717 let rows: Vec<Row> = self
1718 .subscriptions
1719 .iter()
1720 .map(|(name, sub)| {
1721 Row::new(alloc::vec![
1722 Value::Text(name.clone()),
1723 Value::Text(sub.conn_str.clone()),
1724 Value::Text(sub.publications.join(", ")),
1725 Value::Bool(sub.enabled),
1726 Value::BigInt(i64::try_from(sub.last_received_pos).unwrap_or(i64::MAX)),
1727 ])
1728 })
1729 .collect();
1730 QueryResult::Rows { columns, rows }
1731 }
1732
1733 fn exec_spg_statistic(&self) -> QueryResult {
1738 let columns = alloc::vec![
1739 ColumnSchema::new("table_name", DataType::Text, false),
1740 ColumnSchema::new("column_name", DataType::Text, false),
1741 ColumnSchema::new("null_frac", DataType::Float, false),
1742 ColumnSchema::new("n_distinct", DataType::BigInt, false),
1743 ColumnSchema::new("histogram_bounds", DataType::Text, false),
1744 ColumnSchema::new("cold_row_count", DataType::BigInt, false),
1749 ];
1750 let rows: Vec<Row> = self
1751 .statistics
1752 .iter()
1753 .map(|((t, c), s)| {
1754 let cold = self
1755 .catalog
1756 .get(t)
1757 .map_or(0, |table| table.cold_row_count());
1758 Row::new(alloc::vec![
1759 Value::Text(t.clone()),
1760 Value::Text(c.clone()),
1761 Value::Float(f64::from(s.null_frac)),
1762 Value::BigInt(i64::try_from(s.n_distinct).unwrap_or(i64::MAX)),
1763 Value::Text(render_histogram_bounds(&s.histogram_bounds)),
1764 Value::BigInt(i64::try_from(cold).unwrap_or(i64::MAX)),
1765 ])
1766 })
1767 .collect();
1768 QueryResult::Rows { columns, rows }
1769 }
1770
1771 fn exec_spg_stat_replication(&self) -> QueryResult {
1778 let columns = alloc::vec![
1779 ColumnSchema::new("name", DataType::Text, false),
1780 ColumnSchema::new("conn_str", DataType::Text, false),
1781 ColumnSchema::new("publications", DataType::Text, false),
1782 ColumnSchema::new("last_received_pos", DataType::BigInt, false),
1783 ColumnSchema::new("enabled", DataType::Bool, false),
1784 ];
1785 let rows: Vec<Row> = self
1786 .subscriptions
1787 .iter()
1788 .map(|(name, sub)| {
1789 Row::new(alloc::vec![
1790 Value::Text(name.clone()),
1791 Value::Text(sub.conn_str.clone()),
1792 Value::Text(sub.publications.join(",")),
1793 Value::BigInt(i64::try_from(sub.last_received_pos).unwrap_or(i64::MAX)),
1794 Value::Bool(sub.enabled),
1795 ])
1796 })
1797 .collect();
1798 QueryResult::Rows { columns, rows }
1799 }
1800
1801 fn exec_spg_stat_segment(&self) -> QueryResult {
1813 let columns = alloc::vec![
1814 ColumnSchema::new("segment_id", DataType::BigInt, false),
1815 ColumnSchema::new("table_name", DataType::Text, false),
1816 ColumnSchema::new("num_rows", DataType::BigInt, false),
1817 ColumnSchema::new("num_pages", DataType::BigInt, false),
1818 ColumnSchema::new("total_bytes", DataType::BigInt, false),
1819 ];
1820 let mut segment_owners: alloc::collections::BTreeMap<u32, String> = BTreeMap::new();
1826 for tname in self.catalog.table_names() {
1827 if is_internal_table_name(&tname) {
1828 continue;
1829 }
1830 let Some(t) = self.catalog.get(&tname) else {
1831 continue;
1832 };
1833 for idx in t.indices() {
1834 if let spg_storage::IndexKind::BTree(map) = &idx.kind {
1835 for (_, locs) in map.iter() {
1836 for loc in locs {
1837 if let spg_storage::RowLocator::Cold { segment_id, .. } = loc {
1838 segment_owners
1839 .entry(*segment_id)
1840 .or_insert_with(|| tname.clone());
1841 }
1842 }
1843 }
1844 }
1845 }
1846 }
1847 let rows: Vec<Row> = self
1848 .catalog
1849 .cold_segment_ids_global()
1850 .iter()
1851 .filter_map(|&id| {
1852 let seg = self.catalog.cold_segment(id)?;
1853 let meta = seg.meta();
1854 let owner = segment_owners.get(&id).cloned().unwrap_or_default();
1855 Some(Row::new(alloc::vec![
1856 Value::BigInt(i64::from(id)),
1857 Value::Text(owner),
1858 Value::BigInt(i64::try_from(meta.num_rows).unwrap_or(i64::MAX)),
1859 Value::BigInt(i64::from(meta.num_pages)),
1860 Value::BigInt(i64::try_from(meta.total_bytes).unwrap_or(i64::MAX)),
1861 ]))
1862 })
1863 .collect();
1864 QueryResult::Rows { columns, rows }
1865 }
1866
1867 fn exec_spg_stat_query(&self) -> QueryResult {
1873 let columns = alloc::vec![
1874 ColumnSchema::new("sql", DataType::Text, false),
1875 ColumnSchema::new("exec_count", DataType::BigInt, false),
1876 ColumnSchema::new("total_us", DataType::BigInt, false),
1877 ColumnSchema::new("mean_us", DataType::BigInt, false),
1878 ColumnSchema::new("max_us", DataType::BigInt, false),
1879 ColumnSchema::new("last_seen_us", DataType::BigInt, false),
1880 ];
1881 let rows: Vec<Row> = self
1882 .query_stats
1883 .snapshot()
1884 .into_iter()
1885 .map(|(sql, s)| {
1886 let mean = if s.exec_count == 0 {
1887 0
1888 } else {
1889 s.total_us / s.exec_count
1890 };
1891 Row::new(alloc::vec![
1892 Value::Text(sql),
1893 Value::BigInt(i64::try_from(s.exec_count).unwrap_or(i64::MAX)),
1894 Value::BigInt(i64::try_from(s.total_us).unwrap_or(i64::MAX)),
1895 Value::BigInt(i64::try_from(mean).unwrap_or(i64::MAX)),
1896 Value::BigInt(i64::try_from(s.max_us).unwrap_or(i64::MAX)),
1897 Value::BigInt(i64::try_from(s.last_seen_us).unwrap_or(i64::MAX)),
1898 ])
1899 })
1900 .collect();
1901 QueryResult::Rows { columns, rows }
1902 }
1903
1904 #[must_use]
1909 pub const fn with_activity_provider(mut self, f: ActivityProvider) -> Self {
1910 self.activity_provider = Some(f);
1911 self
1912 }
1913
1914 #[must_use]
1916 pub const fn with_audit_providers(
1917 mut self,
1918 chain: AuditChainProvider,
1919 verify: AuditVerifier,
1920 ) -> Self {
1921 self.audit_chain_provider = Some(chain);
1922 self.audit_verifier = Some(verify);
1923 self
1924 }
1925
1926 #[must_use]
1931 pub const fn with_slow_query_log(mut self, threshold_us: u64, logger: SlowQueryLogger) -> Self {
1932 self.slow_query_threshold_us = Some(threshold_us);
1933 self.slow_query_logger = Some(logger);
1934 self
1935 }
1936
1937 pub fn set_plan_cache_max(&mut self, n: usize) {
1941 self.plan_cache.set_max_entries(n);
1942 }
1943
1944 fn exec_spg_stat_activity(&self) -> QueryResult {
1949 let columns = alloc::vec![
1950 ColumnSchema::new("pid", DataType::Int, false),
1951 ColumnSchema::new("user", DataType::Text, false),
1952 ColumnSchema::new("started_at_us", DataType::BigInt, false),
1953 ColumnSchema::new("current_sql", DataType::Text, false),
1954 ColumnSchema::new("wait_event", DataType::Text, false),
1955 ColumnSchema::new("elapsed_us", DataType::BigInt, false),
1956 ColumnSchema::new("in_transaction", DataType::Bool, false),
1957 ];
1958 let rows: Vec<Row> = self
1959 .activity_provider
1960 .map(|f| f())
1961 .unwrap_or_default()
1962 .into_iter()
1963 .map(|r| {
1964 Row::new(alloc::vec![
1965 Value::Int(i32::try_from(r.pid).unwrap_or(i32::MAX)),
1966 Value::Text(r.user),
1967 Value::BigInt(r.started_at_us),
1968 Value::Text(r.current_sql),
1969 Value::Text(r.wait_event),
1970 Value::BigInt(r.elapsed_us),
1971 Value::Bool(r.in_transaction),
1972 ])
1973 })
1974 .collect();
1975 QueryResult::Rows { columns, rows }
1976 }
1977
1978 fn exec_spg_table_ddl(&self) -> QueryResult {
1982 let columns = alloc::vec![
1983 ColumnSchema::new("table_name", DataType::Text, false),
1984 ColumnSchema::new("ddl", DataType::Text, false),
1985 ];
1986 let rows: Vec<Row> = self
1987 .catalog
1988 .table_names()
1989 .into_iter()
1990 .filter(|n| !is_internal_table_name(n))
1991 .filter_map(|name| {
1992 let table = self.catalog.get(&name)?;
1993 let ddl = render_create_table(&name, &table.schema().columns);
1994 Some(Row::new(alloc::vec![Value::Text(name), Value::Text(ddl),]))
1995 })
1996 .collect();
1997 QueryResult::Rows { columns, rows }
1998 }
1999
2000 fn exec_spg_role_ddl(&self) -> QueryResult {
2004 let columns = alloc::vec![
2005 ColumnSchema::new("role_name", DataType::Text, false),
2006 ColumnSchema::new("ddl", DataType::Text, false),
2007 ];
2008 let rows: Vec<Row> = self
2009 .users
2010 .iter()
2011 .map(|(name, rec)| {
2012 let ddl = alloc::format!(
2013 "CREATE USER {name} WITH PASSWORD '<redacted>' ROLE '{}'",
2014 rec.role.as_str(),
2015 );
2016 Row::new(alloc::vec![
2017 Value::Text(String::from(name)),
2018 Value::Text(ddl)
2019 ])
2020 })
2021 .collect();
2022 QueryResult::Rows { columns, rows }
2023 }
2024
2025 fn exec_spg_database_ddl(&self) -> QueryResult {
2031 let columns = alloc::vec![ColumnSchema::new("ddl", DataType::Text, false)];
2032 let mut out = String::new();
2033 for (name, rec) in self.users.iter() {
2034 out.push_str(&alloc::format!(
2035 "CREATE USER {name} WITH PASSWORD '<redacted>' ROLE '{}';\n",
2036 rec.role.as_str(),
2037 ));
2038 }
2039 for name in self.catalog.table_names() {
2040 if is_internal_table_name(&name) {
2041 continue;
2042 }
2043 if let Some(table) = self.catalog.get(&name) {
2044 out.push_str(&render_create_table(&name, &table.schema().columns));
2045 out.push_str(";\n");
2046 }
2047 }
2048 QueryResult::Rows {
2049 columns,
2050 rows: alloc::vec![Row::new(alloc::vec![Value::Text(out)])],
2051 }
2052 }
2053
2054 fn exec_spg_audit_chain(&self) -> QueryResult {
2058 let columns = alloc::vec![
2059 ColumnSchema::new("seq", DataType::BigInt, false),
2060 ColumnSchema::new("ts_ms", DataType::BigInt, false),
2061 ColumnSchema::new("prev_hash", DataType::Text, false),
2062 ColumnSchema::new("entry_hash", DataType::Text, false),
2063 ColumnSchema::new("sql", DataType::Text, false),
2064 ];
2065 let rows: Vec<Row> = self
2066 .audit_chain_provider
2067 .map(|f| f())
2068 .unwrap_or_default()
2069 .into_iter()
2070 .map(|r| {
2071 Row::new(alloc::vec![
2072 Value::BigInt(r.seq),
2073 Value::BigInt(r.ts_ms),
2074 Value::Text(r.prev_hash_hex),
2075 Value::Text(r.entry_hash_hex),
2076 Value::Text(r.sql),
2077 ])
2078 })
2079 .collect();
2080 QueryResult::Rows { columns, rows }
2081 }
2082
2083 fn exec_spg_audit_verify(&self) -> QueryResult {
2089 let columns = alloc::vec![
2090 ColumnSchema::new("verified_count", DataType::BigInt, false),
2091 ColumnSchema::new("broken_at_seq", DataType::BigInt, false),
2092 ];
2093 let (verified, broken) = self.audit_verifier.map(|f| f()).unwrap_or((0, -1));
2094 let row = Row::new(alloc::vec![Value::BigInt(verified), Value::BigInt(broken),]);
2095 QueryResult::Rows {
2096 columns,
2097 rows: alloc::vec![row],
2098 }
2099 }
2100
2101 pub fn query_stats(&self) -> &query_stats::QueryStats {
2103 &self.query_stats
2104 }
2105
2106 pub fn query_stats_mut(&mut self) -> &mut query_stats::QueryStats {
2108 &mut self.query_stats
2109 }
2110
2111 pub const fn statistics(&self) -> &statistics::Statistics {
2115 &self.statistics
2116 }
2117
2118 pub fn tables_needing_analyze(&self) -> Vec<String> {
2131 const MIN_ROWS: u64 = 100;
2132 let mut out = Vec::new();
2133 for name in self.catalog.table_names() {
2134 if is_internal_table_name(&name) {
2135 continue;
2136 }
2137 let Some(table) = self.catalog.get(&name) else {
2138 continue;
2139 };
2140 let row_count = table.rows().len() as u64;
2141 let modified = self.statistics.modified_since_last_analyze(&name);
2142 let base = row_count.max(MIN_ROWS);
2147 let threshold = base.saturating_add(9) / 10;
2148 if modified >= threshold {
2149 out.push(name);
2150 }
2151 }
2152 out
2153 }
2154
2155 fn exec_analyze(&mut self, target: Option<&str>) -> Result<QueryResult, EngineError> {
2166 let names: Vec<String> = if let Some(name) = target {
2167 if self.catalog.get(name).is_none() {
2169 return Err(EngineError::Storage(StorageError::TableNotFound {
2170 name: name.to_string(),
2171 }));
2172 }
2173 alloc::vec![name.to_string()]
2174 } else {
2175 self.catalog
2176 .table_names()
2177 .into_iter()
2178 .filter(|n| !is_internal_table_name(n))
2179 .collect()
2180 };
2181 let mut analysed = 0usize;
2182 for table_name in &names {
2183 self.analyze_one_table(table_name)?;
2184 analysed += 1;
2185 }
2186 if analysed > 0 {
2192 self.statistics.bump_version();
2193 if target.is_some() {
2194 for t in &names {
2195 self.plan_cache.evict_referencing(t);
2196 }
2197 } else {
2198 self.plan_cache.clear();
2199 }
2200 }
2201 Ok(QueryResult::CommandOk {
2202 affected: analysed,
2203 modified_catalog: true,
2204 })
2205 }
2206
2207 fn set_session_param(&mut self, name: String, value: spg_sql::ast::SetValue) {
2220 let normalised = match value {
2221 spg_sql::ast::SetValue::String(s) => s,
2222 spg_sql::ast::SetValue::Ident(s) => s,
2223 spg_sql::ast::SetValue::Number(s) => s,
2224 spg_sql::ast::SetValue::Default => String::new(),
2225 };
2226 self.session_params
2227 .insert(name.to_ascii_lowercase(), normalised);
2228 }
2229
2230 #[must_use]
2234 pub fn session_param(&self, name: &str) -> Option<&str> {
2235 self.session_params
2236 .get(&name.to_ascii_lowercase())
2237 .map(String::as_str)
2238 }
2239
2240 fn ev_ctx<'a>(
2245 &'a self,
2246 columns: &'a [ColumnSchema],
2247 alias: Option<&'a str>,
2248 ) -> EvalContext<'a> {
2249 EvalContext::new(columns, alias)
2250 .with_default_text_search_config(self.session_param("default_text_search_config"))
2251 }
2252
2253 fn exec_compact_cold_segments(&mut self) -> Result<QueryResult, EngineError> {
2257 let target = COMPACTION_TARGET_DEFAULT_BYTES;
2258 let reports = self.compact_cold_segments_with_target(target)?;
2259 let columns = alloc::vec![
2260 ColumnSchema::new("table_name", DataType::Text, false),
2261 ColumnSchema::new("index_name", DataType::Text, false),
2262 ColumnSchema::new("sources_merged", DataType::BigInt, false),
2263 ColumnSchema::new("merged_segment_id", DataType::BigInt, false),
2264 ColumnSchema::new("merged_rows", DataType::BigInt, false),
2265 ColumnSchema::new("deleted_rows_pruned", DataType::BigInt, false),
2266 ColumnSchema::new("bytes_reclaimed_estimate", DataType::BigInt, false),
2267 ];
2268 let rows: Vec<Row> = reports
2269 .into_iter()
2270 .map(|(tname, iname, report)| {
2271 Row::new(alloc::vec![
2272 Value::Text(tname),
2273 Value::Text(iname),
2274 Value::BigInt(i64::try_from(report.sources.len()).unwrap_or(i64::MAX)),
2275 Value::BigInt(i64::from(report.merged_segment_id.unwrap_or(0))),
2276 Value::BigInt(i64::try_from(report.merged_rows).unwrap_or(i64::MAX)),
2277 Value::BigInt(i64::try_from(report.deleted_rows_pruned).unwrap_or(i64::MAX),),
2278 Value::BigInt(
2279 i64::try_from(report.bytes_reclaimed_estimate).unwrap_or(i64::MAX),
2280 ),
2281 ])
2282 })
2283 .collect();
2284 Ok(QueryResult::Rows { columns, rows })
2285 }
2286
2287 fn analyze_one_table(&mut self, table_name: &str) -> Result<(), EngineError> {
2292 let table = self.catalog.get(table_name).ok_or_else(|| {
2293 EngineError::Storage(StorageError::TableNotFound {
2294 name: table_name.to_string(),
2295 })
2296 })?;
2297 let schema = table.schema().clone();
2298 let row_count = table.rows().len();
2299 self.statistics.clear_table(table_name);
2304 for (col_pos, col_schema) in schema.columns.iter().enumerate() {
2305 if matches!(col_schema.ty, DataType::Vector { .. }) {
2308 continue;
2309 }
2310 let mut non_null_values: Vec<Value> = Vec::with_capacity(row_count);
2311 let mut nulls: u64 = 0;
2312 for row in table.rows() {
2313 match row.values.get(col_pos) {
2314 Some(Value::Null) | None => nulls += 1,
2315 Some(v) => non_null_values.push(v.clone()),
2316 }
2317 }
2318 non_null_values.sort_by(|a, b| sort_values_for_histogram(a, b));
2323 let non_null: Vec<String> = non_null_values.iter().map(canonical_value_repr).collect();
2324 let null_frac = if row_count == 0 {
2325 0.0
2326 } else {
2327 #[allow(clippy::cast_precision_loss)]
2328 let f = nulls as f32 / row_count as f32;
2329 f
2330 };
2331 let n_distinct = statistics::estimate_n_distinct(&non_null);
2332 let histogram_bounds = statistics::build_histogram(&non_null);
2333 self.statistics.set(
2334 table_name.to_string(),
2335 col_schema.name.clone(),
2336 statistics::ColumnStats {
2337 null_frac,
2338 n_distinct,
2339 histogram_bounds,
2340 },
2341 );
2342 }
2343 self.statistics.reset_modified(table_name);
2344 let cold_count = {
2350 let table = self
2351 .active_catalog()
2352 .get(table_name)
2353 .expect("table still present");
2354 table.count_cold_locators()
2355 };
2356 let table_mut = self
2357 .active_catalog_mut()
2358 .get_mut(table_name)
2359 .expect("table still present");
2360 table_mut.set_cold_row_count(cold_count);
2361 Ok(())
2362 }
2363
2364 fn exec_show_publications(&self) -> QueryResult {
2376 let columns = alloc::vec![
2377 ColumnSchema::new("name", DataType::Text, false),
2378 ColumnSchema::new("scope", DataType::Text, false),
2379 ColumnSchema::new("table_count", DataType::Int, true),
2380 ];
2381 let rows: Vec<Row> = self
2382 .publications
2383 .iter()
2384 .map(|(name, scope)| {
2385 let (scope_str, count_val) = match scope {
2386 spg_sql::ast::PublicationScope::AllTables => {
2387 ("FOR ALL TABLES".to_string(), Value::Null)
2388 }
2389 spg_sql::ast::PublicationScope::ForTables(ts) => (
2390 alloc::format!("FOR TABLE {}", ts.join(", ")),
2391 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
2392 ),
2393 spg_sql::ast::PublicationScope::AllTablesExcept(ts) => (
2394 alloc::format!("FOR ALL TABLES EXCEPT {}", ts.join(", ")),
2395 Value::Int(i32::try_from(ts.len()).unwrap_or(i32::MAX)),
2396 ),
2397 };
2398 Row::new(alloc::vec![
2399 Value::Text(name.clone()),
2400 Value::Text(scope_str),
2401 count_val,
2402 ])
2403 })
2404 .collect();
2405 QueryResult::Rows { columns, rows }
2406 }
2407
2408 fn exec_show_users(&self) -> QueryResult {
2410 let columns = alloc::vec![
2411 ColumnSchema::new("name", DataType::Text, false),
2412 ColumnSchema::new("role", DataType::Text, false),
2413 ];
2414 let rows: Vec<Row> = self
2415 .users
2416 .iter()
2417 .map(|(name, rec)| {
2418 Row::new(alloc::vec![
2419 Value::Text(name.to_string()),
2420 Value::Text(rec.role.as_str().to_string()),
2421 ])
2422 })
2423 .collect();
2424 QueryResult::Rows { columns, rows }
2425 }
2426
2427 fn exec_create_user(&mut self, s: &CreateUserStatement) -> Result<QueryResult, EngineError> {
2428 if self.in_transaction() {
2429 return Err(EngineError::Unsupported(
2430 "CREATE USER is not allowed inside a transaction".into(),
2431 ));
2432 }
2433 let role = users::Role::parse(&s.role).ok_or_else(|| {
2434 EngineError::Unsupported(alloc::format!("invalid role: {:?}", s.role))
2435 })?;
2436 let salt = self.salt_fn.map_or_else(
2440 || {
2441 let mut s_bytes = [0u8; 16];
2442 let digest = spg_crypto::hash(s.name.as_bytes());
2443 s_bytes.copy_from_slice(&digest[..16]);
2444 s_bytes
2445 },
2446 |f| f(),
2447 );
2448 self.users
2449 .create(&s.name, &s.password, role, salt)
2450 .map_err(|e| EngineError::Unsupported(alloc::format!("CREATE USER: {e}")))?;
2451 Ok(QueryResult::CommandOk {
2452 affected: 1,
2453 modified_catalog: true,
2454 })
2455 }
2456
2457 fn exec_drop_user(&mut self, name: &str) -> Result<QueryResult, EngineError> {
2458 if self.in_transaction() {
2459 return Err(EngineError::Unsupported(
2460 "DROP USER is not allowed inside a transaction".into(),
2461 ));
2462 }
2463 self.users
2464 .drop(name)
2465 .map_err(|e| EngineError::Unsupported(alloc::format!("DROP USER: {e}")))?;
2466 Ok(QueryResult::CommandOk {
2467 affected: 1,
2468 modified_catalog: true,
2469 })
2470 }
2471
2472 fn exec_create_function(
2478 &mut self,
2479 s: spg_sql::ast::CreateFunctionStatement,
2480 ) -> Result<QueryResult, EngineError> {
2481 let args_repr = render_function_args(&s.args);
2482 let returns = match &s.returns {
2483 spg_sql::ast::FunctionReturn::Trigger => alloc::string::String::from("TRIGGER"),
2484 spg_sql::ast::FunctionReturn::Void => alloc::string::String::from("VOID"),
2485 spg_sql::ast::FunctionReturn::Type(t) => alloc::format!("{t}"),
2486 spg_sql::ast::FunctionReturn::Other(s) => s.clone(),
2487 };
2488 let body_text = match &s.body {
2489 spg_sql::ast::FunctionBody::PlPgSql(b) => alloc::format!("{b}"),
2490 spg_sql::ast::FunctionBody::Raw(s) => s.clone(),
2491 };
2492 let def = spg_storage::FunctionDef {
2493 name: s.name.clone(),
2494 args_repr,
2495 returns,
2496 language: s.language.clone(),
2497 body: body_text,
2498 };
2499 self.active_catalog_mut()
2500 .create_function(def, s.or_replace)
2501 .map_err(EngineError::Storage)?;
2502 Ok(QueryResult::CommandOk {
2503 affected: 0,
2504 modified_catalog: true,
2505 })
2506 }
2507
2508 fn exec_create_trigger(
2513 &mut self,
2514 s: spg_sql::ast::CreateTriggerStatement,
2515 ) -> Result<QueryResult, EngineError> {
2516 let timing = match s.timing {
2517 spg_sql::ast::TriggerTiming::Before => "BEFORE",
2518 spg_sql::ast::TriggerTiming::After => "AFTER",
2519 spg_sql::ast::TriggerTiming::InsteadOf => "INSTEAD OF",
2520 };
2521 let events: Vec<alloc::string::String> = s
2522 .events
2523 .iter()
2524 .map(|e| match e {
2525 spg_sql::ast::TriggerEvent::Insert => alloc::string::String::from("INSERT"),
2526 spg_sql::ast::TriggerEvent::Update => alloc::string::String::from("UPDATE"),
2527 spg_sql::ast::TriggerEvent::Delete => alloc::string::String::from("DELETE"),
2528 spg_sql::ast::TriggerEvent::Truncate => alloc::string::String::from("TRUNCATE"),
2529 })
2530 .collect();
2531 let for_each = match s.for_each {
2532 spg_sql::ast::TriggerForEach::Row => "ROW",
2533 spg_sql::ast::TriggerForEach::Statement => "STATEMENT",
2534 };
2535 let def = spg_storage::TriggerDef {
2536 name: s.name.clone(),
2537 table: s.table.clone(),
2538 timing: alloc::string::String::from(timing),
2539 events,
2540 for_each: alloc::string::String::from(for_each),
2541 function: s.function.clone(),
2542 update_columns: s.update_columns.clone(),
2543 };
2544 self.active_catalog_mut()
2545 .create_trigger(def, s.or_replace)
2546 .map_err(EngineError::Storage)?;
2547 Ok(QueryResult::CommandOk {
2548 affected: 0,
2549 modified_catalog: true,
2550 })
2551 }
2552
2553 fn exec_drop_trigger(
2554 &mut self,
2555 name: &str,
2556 table: &str,
2557 if_exists: bool,
2558 ) -> Result<QueryResult, EngineError> {
2559 let removed = self.active_catalog_mut().drop_trigger(name, table);
2560 if !removed && !if_exists {
2561 return Err(EngineError::Storage(spg_storage::StorageError::Corrupt(
2562 alloc::format!("trigger {name:?} on {table:?} does not exist"),
2563 )));
2564 }
2565 Ok(QueryResult::CommandOk {
2566 affected: usize::from(removed),
2567 modified_catalog: removed,
2568 })
2569 }
2570
2571 fn exec_drop_function(
2572 &mut self,
2573 name: &str,
2574 if_exists: bool,
2575 ) -> Result<QueryResult, EngineError> {
2576 let removed = self.active_catalog_mut().drop_function(name);
2577 if !removed && !if_exists {
2578 return Err(EngineError::Storage(spg_storage::StorageError::Corrupt(
2579 alloc::format!("function {name:?} does not exist"),
2580 )));
2581 }
2582 Ok(QueryResult::CommandOk {
2583 affected: usize::from(removed),
2584 modified_catalog: removed,
2585 })
2586 }
2587
2588 fn exec_update_cancel(
2595 &mut self,
2596 stmt: &spg_sql::ast::UpdateStatement,
2597 cancel: CancelToken<'_>,
2598 ) -> Result<QueryResult, EngineError> {
2599 let before_update_triggers = self.snapshot_update_row_triggers(&stmt.table, "BEFORE");
2608 let after_update_triggers = self.snapshot_update_row_triggers(&stmt.table, "AFTER");
2609 let trigger_session_cfg: Option<String> = self
2610 .session_params
2611 .get("default_text_search_config")
2612 .cloned();
2613 if let Some(w) = &stmt.where_ {
2621 let schema_cols = self
2622 .active_catalog()
2623 .get(&stmt.table)
2624 .ok_or_else(|| {
2625 EngineError::Storage(StorageError::TableNotFound {
2626 name: stmt.table.clone(),
2627 })
2628 })?
2629 .schema()
2630 .columns
2631 .clone();
2632 if let Some((col_pos, key)) = try_pk_predicate(w, &schema_cols, stmt.table.as_str())
2633 && let Some(idx_name) = self
2634 .active_catalog()
2635 .get(&stmt.table)
2636 .and_then(|t| t.index_on(col_pos).map(|i| i.name.clone()))
2637 {
2638 let _ = self
2642 .active_catalog_mut()
2643 .promote_cold_row(&stmt.table, &idx_name, &key);
2644 }
2645 }
2646
2647 let ts_cfg: Option<String> = self
2650 .session_param("default_text_search_config")
2651 .map(String::from);
2652 let table = self
2653 .active_catalog_mut()
2654 .get_mut(&stmt.table)
2655 .ok_or_else(|| {
2656 EngineError::Storage(StorageError::TableNotFound {
2657 name: stmt.table.clone(),
2658 })
2659 })?;
2660 let schema_cols: Vec<ColumnSchema> = table.schema().columns.clone();
2661 let mut targets: Vec<(usize, &Expr)> = Vec::with_capacity(stmt.assignments.len());
2665 for (col, expr) in &stmt.assignments {
2666 let pos = schema_cols
2667 .iter()
2668 .position(|c| c.name == *col)
2669 .ok_or_else(|| {
2670 EngineError::Eval(EvalError::ColumnNotFound { name: col.clone() })
2671 })?;
2672 targets.push((pos, expr));
2673 }
2674 let ctx = EvalContext::new(&schema_cols, Some(stmt.table.as_str()))
2675 .with_default_text_search_config(ts_cfg.as_deref());
2676 let mut planned: Vec<(usize, Vec<Value>)> = Vec::new();
2682 for (i, row) in table.rows().iter().enumerate() {
2683 if i.is_multiple_of(256) {
2687 cancel.check()?;
2688 }
2689 if let Some(w) = &stmt.where_ {
2690 let cond = eval::eval_expr(w, row, &ctx)?;
2691 if !matches!(cond, Value::Bool(true)) {
2692 continue;
2693 }
2694 }
2695 let mut new_vals = row.values.clone();
2696 for (pos, expr) in &targets {
2697 let v = eval::eval_expr(expr, row, &ctx)?;
2698 new_vals[*pos] =
2699 coerce_value(v, schema_cols[*pos].ty, &schema_cols[*pos].name, *pos)?;
2700 }
2701 planned.push((i, new_vals));
2702 }
2703 let plan_with_old: Vec<(usize, Vec<Value>, Vec<Value>)> = planned
2707 .iter()
2708 .map(|(pos, new_vals)| (*pos, table.rows()[*pos].values.clone(), new_vals.clone()))
2709 .collect();
2710 let self_fks = table.schema().foreign_keys.clone();
2711 let _ = table;
2716 if !self_fks.is_empty() {
2720 let new_rows: Vec<Vec<Value>> = planned
2721 .iter()
2722 .map(|(_pos, new_vals)| new_vals.clone())
2723 .collect();
2724 enforce_fk_inserts(self.active_catalog(), &stmt.table, &self_fks, &new_rows)?;
2725 }
2726 {
2730 let new_rows: Vec<Vec<Value>> = planned
2731 .iter()
2732 .map(|(_pos, new_vals)| new_vals.clone())
2733 .collect();
2734 enforce_check_constraints(self.active_catalog(), &stmt.table, &new_rows)?;
2735 }
2736 let child_plan =
2740 plan_fk_parent_updates(self.active_catalog(), &stmt.table, &plan_with_old)?;
2741 for step in &child_plan {
2743 apply_fk_child_step(self.active_catalog_mut(), step)?;
2744 }
2745 let table = self
2747 .active_catalog_mut()
2748 .get_mut(&stmt.table)
2749 .ok_or_else(|| {
2750 EngineError::Storage(StorageError::TableNotFound {
2751 name: stmt.table.clone(),
2752 })
2753 })?;
2754 let mut applied_after_before: Vec<(usize, Row, Row)> = Vec::with_capacity(planned.len());
2764 let mut deferred_embedded: Vec<triggers::DeferredEmbeddedStmt> = Vec::new();
2766 for (pos, new_vals) in &planned {
2767 let old_row = table.rows()[*pos].clone();
2768 let mut new_row = Row::new(new_vals.clone());
2769 let mut skip = false;
2770 for (fd, filter) in &before_update_triggers {
2771 if !filter.is_empty()
2776 && !any_column_changed(filter, &schema_cols, &old_row, &new_row)
2777 {
2778 continue;
2779 }
2780 let (outcome, deferred) = triggers::fire_row_trigger(
2781 fd,
2782 Some(new_row.clone()),
2783 Some(&old_row),
2784 &stmt.table,
2785 &schema_cols,
2786 &[],
2787 trigger_session_cfg.as_deref(),
2788 false,
2789 )
2790 .map_err(|e| EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}"))))?;
2791 deferred_embedded.extend(deferred);
2792 match outcome {
2793 triggers::TriggerOutcome::Row(r) => new_row = r,
2794 triggers::TriggerOutcome::Skip => {
2795 skip = true;
2796 break;
2797 }
2798 }
2799 }
2800 if !skip {
2801 applied_after_before.push((*pos, new_row, old_row));
2802 }
2803 }
2804 let updated_for_returning: Vec<Vec<Value>> = if stmt.returning.is_some() {
2807 applied_after_before
2808 .iter()
2809 .map(|(_pos, new_row, _old)| new_row.values.clone())
2810 .collect()
2811 } else {
2812 Vec::new()
2813 };
2814 let affected = applied_after_before.len();
2815 for (pos, new_row, old_row) in applied_after_before {
2819 table.update_row(pos, new_row.values.clone())?;
2820 for (fd, filter) in &after_update_triggers {
2821 if !filter.is_empty()
2822 && !any_column_changed(filter, &schema_cols, &old_row, &new_row)
2823 {
2824 continue;
2825 }
2826 let (_outcome, deferred) = triggers::fire_row_trigger(
2827 fd,
2828 Some(new_row.clone()),
2829 Some(&old_row),
2830 &stmt.table,
2831 &schema_cols,
2832 &[],
2833 trigger_session_cfg.as_deref(),
2834 true,
2835 )
2836 .map_err(|e| EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}"))))?;
2837 deferred_embedded.extend(deferred);
2838 }
2839 }
2840 let _ = table;
2841 self.execute_deferred_trigger_stmts(deferred_embedded, cancel)?;
2843 if !self.in_transaction() && affected > 0 {
2845 self.statistics
2846 .record_modifications(&stmt.table, affected as u64);
2847 }
2848 if let Some(items) = &stmt.returning {
2850 return self.build_returning_rows(&stmt.table, items, updated_for_returning);
2851 }
2852 Ok(QueryResult::CommandOk {
2853 affected,
2854 modified_catalog: !self.in_transaction(),
2855 })
2856 }
2857
2858 fn exec_delete_cancel(
2862 &mut self,
2863 stmt: &spg_sql::ast::DeleteStatement,
2864 cancel: CancelToken<'_>,
2865 ) -> Result<QueryResult, EngineError> {
2866 let before_delete_triggers = self.snapshot_row_triggers(&stmt.table, "DELETE", "BEFORE");
2870 let after_delete_triggers = self.snapshot_row_triggers(&stmt.table, "DELETE", "AFTER");
2871 let trigger_session_cfg: Option<String> = self
2872 .session_params
2873 .get("default_text_search_config")
2874 .cloned();
2875 let mut cold_shadow_count: usize = 0;
2883 if let Some(w) = &stmt.where_ {
2884 let schema_cols = self
2885 .active_catalog()
2886 .get(&stmt.table)
2887 .ok_or_else(|| {
2888 EngineError::Storage(StorageError::TableNotFound {
2889 name: stmt.table.clone(),
2890 })
2891 })?
2892 .schema()
2893 .columns
2894 .clone();
2895 if let Some((col_pos, key)) = try_pk_predicate(w, &schema_cols, stmt.table.as_str())
2896 && let Some(idx_name) = self
2897 .active_catalog()
2898 .get(&stmt.table)
2899 .and_then(|t| t.index_on(col_pos).map(|i| i.name.clone()))
2900 {
2901 cold_shadow_count = self
2902 .active_catalog_mut()
2903 .shadow_cold_row(&stmt.table, &idx_name, &key)
2904 .unwrap_or(0);
2905 }
2906 }
2907
2908 let ts_cfg: Option<String> = self
2914 .session_param("default_text_search_config")
2915 .map(String::from);
2916 let table = self
2917 .active_catalog_mut()
2918 .get_mut(&stmt.table)
2919 .ok_or_else(|| {
2920 EngineError::Storage(StorageError::TableNotFound {
2921 name: stmt.table.clone(),
2922 })
2923 })?;
2924 let schema_cols: Vec<ColumnSchema> = table.schema().columns.clone();
2925 let ctx = EvalContext::new(&schema_cols, Some(stmt.table.as_str()))
2926 .with_default_text_search_config(ts_cfg.as_deref());
2927 let mut positions: Vec<usize> = Vec::new();
2928 let mut to_delete_rows: Vec<Vec<Value>> = Vec::new();
2932 for (i, row) in table.rows().iter().enumerate() {
2933 if i.is_multiple_of(256) {
2934 cancel.check()?;
2935 }
2936 let keep = if let Some(w) = &stmt.where_ {
2937 let cond = eval::eval_expr(w, row, &ctx)?;
2938 !matches!(cond, Value::Bool(true))
2939 } else {
2940 false
2941 };
2942 if !keep {
2943 positions.push(i);
2944 to_delete_rows.push(row.values.clone());
2945 }
2946 }
2947 let _ = table;
2954 let mut deferred_embedded: Vec<triggers::DeferredEmbeddedStmt> = Vec::new();
2962 if !before_delete_triggers.is_empty() {
2963 let mut filtered_positions: Vec<usize> = Vec::with_capacity(positions.len());
2964 let mut filtered_old_rows: Vec<Vec<Value>> = Vec::with_capacity(to_delete_rows.len());
2965 for (pos, old_vals) in positions.iter().zip(to_delete_rows.iter()) {
2966 let old_row = Row::new(old_vals.clone());
2967 let mut cancel_this = false;
2968 for fd in &before_delete_triggers {
2969 let (outcome, deferred) = triggers::fire_row_trigger(
2970 fd,
2971 None,
2972 Some(&old_row),
2973 &stmt.table,
2974 &schema_cols,
2975 &[],
2976 trigger_session_cfg.as_deref(),
2977 false,
2978 )
2979 .map_err(|e| {
2980 EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}")))
2981 })?;
2982 deferred_embedded.extend(deferred);
2983 if matches!(outcome, triggers::TriggerOutcome::Skip) {
2984 cancel_this = true;
2985 break;
2986 }
2987 }
2988 if !cancel_this {
2989 filtered_positions.push(*pos);
2990 filtered_old_rows.push(old_vals.clone());
2991 }
2992 }
2993 positions = filtered_positions;
2994 to_delete_rows = filtered_old_rows;
2995 }
2996 let cascade_plan = plan_fk_parent_deletions(
2997 self.active_catalog(),
2998 &stmt.table,
2999 &positions,
3000 &to_delete_rows,
3001 )?;
3002 for step in &cascade_plan {
3009 apply_fk_child_step(self.active_catalog_mut(), step)?;
3010 }
3011 let table = self
3013 .active_catalog_mut()
3014 .get_mut(&stmt.table)
3015 .ok_or_else(|| {
3016 EngineError::Storage(StorageError::TableNotFound {
3017 name: stmt.table.clone(),
3018 })
3019 })?;
3020 let affected = table.delete_rows(&positions) + cold_shadow_count;
3021 let _ = table;
3022 if !after_delete_triggers.is_empty() {
3027 for old_vals in &to_delete_rows {
3028 let old_row = Row::new(old_vals.clone());
3029 for fd in &after_delete_triggers {
3030 let (_outcome, deferred) = triggers::fire_row_trigger(
3031 fd,
3032 None,
3033 Some(&old_row),
3034 &stmt.table,
3035 &schema_cols,
3036 &[],
3037 trigger_session_cfg.as_deref(),
3038 true,
3039 )
3040 .map_err(|e| {
3041 EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}")))
3042 })?;
3043 deferred_embedded.extend(deferred);
3044 }
3045 }
3046 }
3047 self.execute_deferred_trigger_stmts(deferred_embedded, cancel)?;
3049 if !self.in_transaction() && affected > 0 {
3051 self.statistics
3052 .record_modifications(&stmt.table, affected as u64);
3053 }
3054 if let Some(items) = &stmt.returning {
3060 return self.build_returning_rows(&stmt.table, items, to_delete_rows);
3061 }
3062 Ok(QueryResult::CommandOk {
3063 affected,
3064 modified_catalog: !self.in_transaction(),
3065 })
3066 }
3067
3068 #[allow(clippy::format_push_string)]
3078 fn exec_explain(
3079 &self,
3080 e: &spg_sql::ast::ExplainStatement,
3081 cancel: CancelToken<'_>,
3082 ) -> Result<QueryResult, EngineError> {
3083 let mut lines = Vec::<String>::new();
3084 explain_select(&e.inner, self, 0, &mut lines);
3085 if e.suggest {
3086 let suggestions = build_index_suggestions(&e.inner, self);
3095 for s in suggestions {
3096 lines.push(s);
3097 }
3098 } else if e.analyze {
3099 let started = self.clock.map(|f| f());
3116 let exec = self.exec_select_cancel(&e.inner, cancel)?;
3117 let elapsed_micros = match (self.clock, started) {
3118 (Some(f), Some(s)) => Some(f().saturating_sub(s)),
3119 _ => None,
3120 };
3121 let row_count = if let QueryResult::Rows { rows, .. } = &exec {
3122 rows.len()
3123 } else {
3124 0
3125 };
3126 annotate_explain_lines(&mut lines, row_count, self);
3127 let mut total = alloc::format!("Total: rows={row_count}");
3128 if let Some(us) = elapsed_micros {
3129 total.push_str(&alloc::format!(" elapsed={us}us"));
3130 }
3131 lines.push(total);
3132 }
3133 let columns = alloc::vec![ColumnSchema::new("QUERY PLAN", DataType::Text, false)];
3134 let rows: Vec<Row> = lines
3135 .into_iter()
3136 .map(|l| Row::new(alloc::vec![Value::Text(l)]))
3137 .collect();
3138 Ok(QueryResult::Rows { columns, rows })
3139 }
3140
3141 fn exec_show_tables(&self) -> QueryResult {
3142 let columns = alloc::vec![ColumnSchema::new("name", DataType::Text, false)];
3143 let rows: Vec<Row> = self
3144 .active_catalog()
3145 .table_names()
3146 .into_iter()
3147 .map(|n| Row::new(alloc::vec![Value::Text(n)]))
3148 .collect();
3149 QueryResult::Rows { columns, rows }
3150 }
3151
3152 fn exec_show_columns(&self, table_name: &str) -> Result<QueryResult, EngineError> {
3155 let table =
3156 self.active_catalog()
3157 .get(table_name)
3158 .ok_or_else(|| StorageError::TableNotFound {
3159 name: table_name.into(),
3160 })?;
3161 let columns = alloc::vec![
3162 ColumnSchema::new("name", DataType::Text, false),
3163 ColumnSchema::new("type", DataType::Text, false),
3164 ColumnSchema::new("nullable", DataType::Bool, false),
3165 ];
3166 let rows: Vec<Row> = table
3167 .schema()
3168 .columns
3169 .iter()
3170 .map(|c| {
3171 Row::new(alloc::vec![
3172 Value::Text(c.name.clone()),
3173 Value::Text(alloc::format!("{}", c.ty)),
3174 Value::Bool(c.nullable),
3175 ])
3176 })
3177 .collect();
3178 Ok(QueryResult::Rows { columns, rows })
3179 }
3180
3181 fn exec_begin(&mut self) -> Result<QueryResult, EngineError> {
3182 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3183 if self.tx_catalogs.contains_key(&tx_id) {
3184 return Err(EngineError::TransactionAlreadyOpen);
3185 }
3186 self.tx_catalogs.insert(
3187 tx_id,
3188 TxState {
3189 catalog: self.catalog.clone(),
3190 savepoints: Vec::new(),
3191 },
3192 );
3193 Ok(QueryResult::CommandOk {
3194 affected: 0,
3195 modified_catalog: false,
3196 })
3197 }
3198
3199 fn exec_commit(&mut self) -> Result<QueryResult, EngineError> {
3200 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3201 let state = self
3202 .tx_catalogs
3203 .remove(&tx_id)
3204 .ok_or(EngineError::NoActiveTransaction)?;
3205 self.catalog = state.catalog;
3206 Ok(QueryResult::CommandOk {
3210 affected: 0,
3211 modified_catalog: true,
3212 })
3213 }
3214
3215 fn exec_rollback(&mut self) -> Result<QueryResult, EngineError> {
3216 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3217 if self.tx_catalogs.remove(&tx_id).is_none() {
3218 return Err(EngineError::NoActiveTransaction);
3219 }
3220 Ok(QueryResult::CommandOk {
3222 affected: 0,
3223 modified_catalog: false,
3224 })
3225 }
3226
3227 fn exec_savepoint(&mut self, name: String) -> Result<QueryResult, EngineError> {
3228 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3229 let state = self
3230 .tx_catalogs
3231 .get_mut(&tx_id)
3232 .ok_or(EngineError::NoActiveTransaction)?;
3233 state.savepoints.retain(|(n, _)| n != &name);
3237 let snapshot = state.catalog.clone();
3238 state.savepoints.push((name, snapshot));
3239 Ok(QueryResult::CommandOk {
3240 affected: 0,
3241 modified_catalog: false,
3242 })
3243 }
3244
3245 fn exec_rollback_to_savepoint(&mut self, name: &str) -> Result<QueryResult, EngineError> {
3246 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3247 let state = self
3248 .tx_catalogs
3249 .get_mut(&tx_id)
3250 .ok_or(EngineError::NoActiveTransaction)?;
3251 let pos = state
3252 .savepoints
3253 .iter()
3254 .rposition(|(n, _)| n == name)
3255 .ok_or_else(|| {
3256 EngineError::Unsupported(alloc::format!("savepoint not found: {name}"))
3257 })?;
3258 let snapshot = state.savepoints[pos].1.clone();
3262 state.savepoints.truncate(pos + 1);
3263 state.catalog = snapshot;
3264 Ok(QueryResult::CommandOk {
3265 affected: 0,
3266 modified_catalog: false,
3267 })
3268 }
3269
3270 fn exec_release_savepoint(&mut self, name: &str) -> Result<QueryResult, EngineError> {
3271 let tx_id = self.current_tx.ok_or(EngineError::NoActiveTransaction)?;
3272 let state = self
3273 .tx_catalogs
3274 .get_mut(&tx_id)
3275 .ok_or(EngineError::NoActiveTransaction)?;
3276 let pos = state
3277 .savepoints
3278 .iter()
3279 .rposition(|(n, _)| n == name)
3280 .ok_or_else(|| {
3281 EngineError::Unsupported(alloc::format!("savepoint not found: {name}"))
3282 })?;
3283 state.savepoints.truncate(pos);
3286 Ok(QueryResult::CommandOk {
3287 affected: 0,
3288 modified_catalog: false,
3289 })
3290 }
3291
3292 fn exec_alter_table(
3303 &mut self,
3304 s: spg_sql::ast::AlterTableStatement,
3305 ) -> Result<QueryResult, EngineError> {
3306 let table_name = s.name.clone();
3311 for target in s.targets {
3312 self.exec_alter_table_subaction(&table_name, target)?;
3313 }
3314 Ok(QueryResult::CommandOk {
3315 affected: 0,
3316 modified_catalog: !self.in_transaction(),
3317 })
3318 }
3319
3320 fn exec_alter_table_subaction(
3321 &mut self,
3322 table_name_outer: &str,
3323 target: spg_sql::ast::AlterTableTarget,
3324 ) -> Result<(), EngineError> {
3325 struct S<'a> {
3328 name: &'a str,
3329 }
3330 let s = S {
3331 name: table_name_outer,
3332 };
3333 match target {
3334 spg_sql::ast::AlterTableTarget::SetHotTierBytes(n) => {
3335 let table = self.active_catalog_mut().get_mut(s.name).ok_or_else(|| {
3336 EngineError::Storage(StorageError::TableNotFound {
3337 name: s.name.into(),
3338 })
3339 })?;
3340 table.schema_mut().hot_tier_bytes = Some(n);
3341 }
3342 spg_sql::ast::AlterTableTarget::AddForeignKey(fk) => {
3343 let cols_snapshot = self
3348 .active_catalog()
3349 .get(s.name)
3350 .ok_or_else(|| {
3351 EngineError::Storage(StorageError::TableNotFound {
3352 name: s.name.into(),
3353 })
3354 })?
3355 .schema()
3356 .columns
3357 .clone();
3358 let storage_fk =
3359 resolve_foreign_key(s.name, &cols_snapshot, fk, self.active_catalog())?;
3360 let existing_rows: Vec<Vec<Value>> = self
3363 .active_catalog()
3364 .get(&s.name)
3365 .expect("checked above")
3366 .rows()
3367 .iter()
3368 .map(|r| r.values.clone())
3369 .collect();
3370 enforce_fk_inserts(
3371 self.active_catalog(),
3372 s.name,
3373 core::slice::from_ref(&storage_fk),
3374 &existing_rows,
3375 )?;
3376 let table = self
3378 .active_catalog_mut()
3379 .get_mut(s.name)
3380 .expect("checked above");
3381 if let Some(name) = &storage_fk.name
3382 && table
3383 .schema()
3384 .foreign_keys
3385 .iter()
3386 .any(|f| f.name.as_ref() == Some(name))
3387 {
3388 return Err(EngineError::Unsupported(alloc::format!(
3389 "ALTER TABLE ADD CONSTRAINT: a constraint named {name:?} already exists"
3390 )));
3391 }
3392 table.schema_mut().foreign_keys.push(storage_fk);
3393 }
3394 spg_sql::ast::AlterTableTarget::DropForeignKey { name, if_exists } => {
3395 let table = self.active_catalog_mut().get_mut(s.name).ok_or_else(|| {
3396 EngineError::Storage(StorageError::TableNotFound {
3397 name: s.name.into(),
3398 })
3399 })?;
3400 let fks = &mut table.schema_mut().foreign_keys;
3401 let before = fks.len();
3402 fks.retain(|f| f.name.as_ref() != Some(&name));
3403 if fks.len() == before && !if_exists {
3404 return Err(EngineError::Unsupported(alloc::format!(
3405 "ALTER TABLE DROP CONSTRAINT: no FK named {name:?} on {:?}",
3406 s.name
3407 )));
3408 }
3409 }
3411 spg_sql::ast::AlterTableTarget::AddColumn {
3412 column,
3413 if_not_exists,
3414 } => {
3415 let clock = self.clock;
3420 let table = self.active_catalog_mut().get_mut(s.name).ok_or_else(|| {
3421 EngineError::Storage(StorageError::TableNotFound {
3422 name: s.name.into(),
3423 })
3424 })?;
3425 if table
3426 .schema()
3427 .columns
3428 .iter()
3429 .any(|c| c.name.eq_ignore_ascii_case(&column.name))
3430 {
3431 if if_not_exists {
3432 return Ok(());
3433 }
3434 return Err(EngineError::Unsupported(alloc::format!(
3435 "ALTER TABLE ADD COLUMN: column {:?} already exists on {:?}",
3436 column.name,
3437 s.name
3438 )));
3439 }
3440 let col_name = column.name.clone();
3441 let nullable = column.nullable;
3442 let has_default =
3443 column.default.is_some() || column.auto_increment;
3444 let col_schema = column_def_to_schema(column)?;
3445 let row_count = table.row_count();
3446 let fill_value: Value = if has_default
3453 || col_schema.runtime_default.is_some()
3454 {
3455 resolve_column_default_free(&col_schema, clock)?
3456 } else if nullable || row_count == 0 {
3457 Value::Null
3458 } else {
3459 return Err(EngineError::Unsupported(alloc::format!(
3460 "ALTER TABLE ADD COLUMN {col_name:?}: NOT NULL column requires DEFAULT \
3461 when the table has existing rows"
3462 )));
3463 };
3464 table.add_column(col_schema, fill_value);
3465 }
3466 spg_sql::ast::AlterTableTarget::AlterColumnType {
3467 column,
3468 new_type,
3469 using,
3470 } => {
3471 let new_data_type = column_type_to_data_type(new_type);
3477 let table = self.active_catalog_mut().get_mut(s.name).ok_or_else(|| {
3478 EngineError::Storage(StorageError::TableNotFound {
3479 name: s.name.into(),
3480 })
3481 })?;
3482 let col_pos = table
3483 .schema()
3484 .columns
3485 .iter()
3486 .position(|c| c.name.eq_ignore_ascii_case(&column))
3487 .ok_or_else(|| {
3488 EngineError::Unsupported(alloc::format!(
3489 "ALTER COLUMN TYPE: column {column:?} not found on {:?}",
3490 s.name
3491 ))
3492 })?;
3493 let schema_cols = table.schema().columns.clone();
3494 let ctx = eval::EvalContext::new(&schema_cols, None);
3495 let mut new_values: alloc::vec::Vec<Value> =
3496 alloc::vec::Vec::with_capacity(table.row_count());
3497 for row in table.rows().iter() {
3498 let raw = match &using {
3499 Some(expr) => eval::eval_expr(expr, row, &ctx).map_err(|e| {
3500 EngineError::Unsupported(alloc::format!(
3501 "ALTER COLUMN TYPE: USING expression failed: {e:?}"
3502 ))
3503 })?,
3504 None => row.values.get(col_pos).cloned().unwrap_or(Value::Null),
3505 };
3506 let coerced = coerce_value(raw, new_data_type, &column, col_pos)?;
3507 new_values.push(coerced);
3508 }
3509 table.schema_mut().columns[col_pos].ty = new_data_type;
3510 for (i, v) in new_values.into_iter().enumerate() {
3511 let mut row_values = table
3512 .rows()
3513 .get(i)
3514 .expect("bounds-checked above")
3515 .values
3516 .clone();
3517 row_values[col_pos] = v;
3518 table.update_row(i, row_values)?;
3519 }
3520 }
3521 }
3522 Ok(())
3523 }
3524
3525 fn exec_alter_index(
3526 &mut self,
3527 stmt: spg_sql::ast::AlterIndexStatement,
3528 ) -> Result<QueryResult, EngineError> {
3529 let spg_sql::ast::AlterIndexStatement {
3533 name: idx_name,
3534 target,
3535 } = stmt;
3536 let spg_sql::ast::AlterIndexTarget::Rebuild { encoding } = target;
3537 let target = encoding.map(|e| match e {
3538 SqlVecEncoding::F32 => VecEncoding::F32,
3539 SqlVecEncoding::Sq8 => VecEncoding::Sq8,
3540 SqlVecEncoding::F16 => VecEncoding::F16,
3541 });
3542 let table_name = {
3547 let cat = self.active_catalog();
3548 let mut found: Option<String> = None;
3549 for tname in cat.table_names() {
3550 if let Some(t) = cat.get(&tname)
3551 && t.indices().iter().any(|i| i.name == idx_name)
3552 {
3553 found = Some(tname);
3554 break;
3555 }
3556 }
3557 found.ok_or_else(|| {
3558 EngineError::Storage(StorageError::IndexNotFound {
3559 name: idx_name.clone(),
3560 })
3561 })?
3562 };
3563 let table = self
3564 .active_catalog_mut()
3565 .get_mut(&table_name)
3566 .expect("table found above");
3567 table.rebuild_nsw_index(&idx_name, target)?;
3568 self.plan_cache.evict_referencing(&table_name);
3571 Ok(QueryResult::CommandOk {
3572 affected: 0,
3573 modified_catalog: !self.in_transaction(),
3574 })
3575 }
3576
3577 fn exec_create_index(
3578 &mut self,
3579 stmt: CreateIndexStatement,
3580 ) -> Result<QueryResult, EngineError> {
3581 let table = self
3582 .active_catalog_mut()
3583 .get_mut(&stmt.table)
3584 .ok_or_else(|| {
3585 EngineError::Storage(StorageError::TableNotFound {
3586 name: stmt.table.clone(),
3587 })
3588 })?;
3589 if stmt.if_not_exists && table.indices().iter().any(|i| i.name == stmt.name) {
3591 return Ok(QueryResult::CommandOk {
3592 affected: 0,
3593 modified_catalog: false,
3594 });
3595 }
3596 let _ = &stmt.extra_columns; let table_name = stmt.table.clone();
3603 let included_positions: Vec<usize> = if stmt.included_columns.is_empty() {
3607 Vec::new()
3608 } else {
3609 let schema = table.schema();
3610 stmt.included_columns
3611 .iter()
3612 .map(|c| {
3613 schema.column_position(c).ok_or_else(|| {
3614 EngineError::Storage(StorageError::ColumnNotFound { column: c.clone() })
3615 })
3616 })
3617 .collect::<Result<Vec<_>, _>>()?
3618 };
3619 match stmt.method {
3620 IndexMethod::BTree => table.add_index(stmt.name.clone(), &stmt.column)?,
3621 IndexMethod::Hnsw => {
3622 if !included_positions.is_empty() {
3623 return Err(EngineError::Unsupported(
3624 "INCLUDE columns are not supported on HNSW indexes".into(),
3625 ));
3626 }
3627 table.add_nsw_index(stmt.name.clone(), &stmt.column, spg_storage::NSW_DEFAULT_M)?;
3628 }
3629 IndexMethod::Brin => {
3631 if !included_positions.is_empty() {
3632 return Err(EngineError::Unsupported(
3633 "INCLUDE columns are not supported on BRIN indexes".into(),
3634 ));
3635 }
3636 table.add_brin_index(stmt.name.clone(), &stmt.column)?;
3637 }
3638 IndexMethod::Gin => {
3646 if !included_positions.is_empty() {
3647 return Err(EngineError::Unsupported(
3648 "INCLUDE columns are not supported on GIN indexes".into(),
3649 ));
3650 }
3651 let col_pos = table
3652 .schema()
3653 .column_position(&stmt.column)
3654 .ok_or_else(|| {
3655 EngineError::Storage(StorageError::ColumnNotFound {
3656 column: stmt.column.clone(),
3657 })
3658 })?;
3659 if table.schema().columns[col_pos].ty == spg_storage::DataType::TsVector {
3660 table
3661 .add_gin_index(stmt.name.clone(), &stmt.column)
3662 .map_err(EngineError::Storage)?;
3663 } else {
3664 table.add_index(stmt.name.clone(), &stmt.column)?;
3670 }
3671 }
3672 }
3673 if !included_positions.is_empty()
3674 && let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name)
3675 {
3676 idx.included_columns = included_positions;
3677 }
3678 if let Some(pred_expr) = &stmt.partial_predicate {
3686 let canonical = pred_expr.to_string();
3687 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3699 idx.partial_predicate = Some(canonical);
3700 }
3701 }
3702 if let Some(key_expr) = &stmt.expression {
3710 if matches!(
3711 stmt.method,
3712 IndexMethod::Hnsw | IndexMethod::Brin | IndexMethod::Gin
3713 ) {
3714 return Err(EngineError::Unsupported(
3715 "Expression keys are not supported on HNSW or BRIN indexes".into(),
3716 ));
3717 }
3718 let canonical = key_expr.to_string();
3719 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3720 idx.expression = Some(canonical);
3721 }
3722 }
3723 if stmt.is_unique {
3732 let mut extra_positions: alloc::vec::Vec<usize> = alloc::vec::Vec::new();
3733 for col_name in &stmt.extra_columns {
3734 let pos = table
3735 .schema()
3736 .columns
3737 .iter()
3738 .position(|c| c.name.eq_ignore_ascii_case(col_name))
3739 .ok_or_else(|| {
3740 EngineError::Unsupported(alloc::format!(
3741 "UNIQUE INDEX {:?}: extra column {col_name:?} not in table {:?}",
3742 stmt.name,
3743 stmt.table
3744 ))
3745 })?;
3746 extra_positions.push(pos);
3747 }
3748 if let Some(idx) = table.indices_mut().iter_mut().find(|i| i.name == stmt.name) {
3749 idx.is_unique = true;
3750 idx.extra_column_positions = extra_positions;
3751 }
3752 let snapshot_indices = table.indices().to_vec();
3757 let snapshot_rows: alloc::vec::Vec<spg_storage::Row> =
3758 table.rows().iter().cloned().collect();
3759 let snapshot_schema = table.schema().clone();
3760 let idx_ref = snapshot_indices
3761 .iter()
3762 .find(|i| i.name == stmt.name)
3763 .expect("just-added index");
3764 check_existing_unique_violation(idx_ref, &snapshot_schema, &snapshot_rows)?;
3765 }
3766 self.plan_cache.evict_referencing(&table_name);
3769 Ok(QueryResult::CommandOk {
3770 affected: 0,
3771 modified_catalog: !self.in_transaction(),
3772 })
3773 }
3774
3775 fn exec_create_table(
3776 &mut self,
3777 stmt: CreateTableStatement,
3778 ) -> Result<QueryResult, EngineError> {
3779 if stmt.if_not_exists && self.active_catalog().get(&stmt.name).is_some() {
3780 return Ok(QueryResult::CommandOk {
3781 affected: 0,
3782 modified_catalog: false,
3783 });
3784 }
3785 let table_name = stmt.name.clone();
3786 let inline_pk_columns: Vec<String> = stmt
3790 .columns
3791 .iter()
3792 .filter(|c| c.is_primary_key)
3793 .map(|c| c.name.clone())
3794 .collect();
3795 let cols = stmt
3801 .columns
3802 .into_iter()
3803 .map(column_def_to_schema)
3804 .collect::<Result<Vec<_>, _>>()?;
3805 let mut cols = cols;
3807 for tc in &stmt.table_constraints {
3808 if let spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } = tc {
3809 for col_name in columns {
3810 if let Some(col) = cols.iter_mut().find(|c| c.name == *col_name) {
3811 col.nullable = false;
3812 }
3813 }
3814 }
3815 }
3816 let mut fks: Vec<spg_storage::ForeignKeyConstraint> =
3823 Vec::with_capacity(stmt.foreign_keys.len());
3824 for fk in stmt.foreign_keys {
3825 fks.push(resolve_foreign_key(
3826 &table_name,
3827 &cols,
3828 fk,
3829 self.active_catalog(),
3830 )?);
3831 }
3832 let mut schema = TableSchema::new(table_name.clone(), cols);
3833 schema.foreign_keys = fks;
3834 let mut uc_storage: Vec<spg_storage::UniquenessConstraint> = Vec::new();
3838 let mut check_exprs: Vec<String> = Vec::new();
3839 for tc in &stmt.table_constraints {
3840 let (is_pk, names, nnd) = match tc {
3841 spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } => {
3842 (true, columns.clone(), false)
3843 }
3844 spg_sql::ast::TableConstraint::Unique {
3845 columns,
3846 nulls_not_distinct,
3847 ..
3848 } => (false, columns.clone(), *nulls_not_distinct),
3849 spg_sql::ast::TableConstraint::Check { expr, .. } => {
3850 check_exprs.push(alloc::format!("{expr}"));
3853 continue;
3854 }
3855 };
3856 let mut positions = Vec::with_capacity(names.len());
3857 for n in &names {
3858 let pos = schema
3859 .columns
3860 .iter()
3861 .position(|c| c.name == *n)
3862 .ok_or_else(|| {
3863 EngineError::Unsupported(alloc::format!(
3864 "table constraint references unknown column {n:?}"
3865 ))
3866 })?;
3867 positions.push(pos);
3868 }
3869 uc_storage.push(spg_storage::UniquenessConstraint {
3870 is_primary_key: is_pk,
3871 columns: positions,
3872 nulls_not_distinct: nnd,
3873 });
3874 }
3875 schema.uniqueness_constraints = uc_storage.clone();
3876 schema.checks = check_exprs;
3877 self.active_catalog_mut().create_table(schema)?;
3878 let table = self
3882 .active_catalog_mut()
3883 .get_mut(&table_name)
3884 .expect("just created");
3885 for (i, col_name) in inline_pk_columns.iter().enumerate() {
3886 let idx_name = if inline_pk_columns.len() == 1 {
3887 alloc::format!("{table_name}_pkey")
3888 } else {
3889 alloc::format!("{table_name}_pkey_{i}")
3890 };
3891 if let Err(e) = table.add_index(idx_name, col_name) {
3892 return Err(EngineError::Storage(e));
3893 }
3894 }
3895 for (i, tc) in stmt.table_constraints.iter().enumerate() {
3896 let (is_pk, names) = match tc {
3897 spg_sql::ast::TableConstraint::PrimaryKey { columns, .. } => (true, columns),
3898 spg_sql::ast::TableConstraint::Unique { columns, .. } => (false, columns),
3899 spg_sql::ast::TableConstraint::Check { .. } => continue,
3900 };
3901 let leading = &names[0];
3902 let already = table.indices().iter().any(|idx| {
3905 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
3906 && table.schema().columns[idx.column_position].name == *leading
3907 });
3908 if already {
3909 continue;
3910 }
3911 let suffix = if is_pk { "pkey" } else { "key" };
3912 let idx_name = if names.len() == 1 {
3913 alloc::format!("{table_name}_{leading}_{suffix}")
3914 } else {
3915 alloc::format!("{table_name}_{leading}_{suffix}_{i}")
3916 };
3917 if let Err(e) = table.add_index(idx_name, leading) {
3918 return Err(EngineError::Storage(e));
3919 }
3920 }
3921 Ok(QueryResult::CommandOk {
3922 affected: 0,
3923 modified_catalog: !self.in_transaction(),
3924 })
3925 }
3926
3927 fn exec_insert(&mut self, stmt: InsertStatement) -> Result<QueryResult, EngineError> {
3928 if let Some(select) = stmt.select_source.clone() {
3933 let select_result = self.exec_select_cancel(&select, CancelToken::none())?;
3934 let rows = match select_result {
3935 QueryResult::Rows { rows, .. } => rows,
3936 other => {
3937 return Err(EngineError::Unsupported(alloc::format!(
3938 "INSERT … SELECT: inner statement produced {other:?} instead of a row set"
3939 )));
3940 }
3941 };
3942 let mut materialised: Vec<Vec<Expr>> = Vec::with_capacity(rows.len());
3943 for row in rows {
3944 let mut tuple: Vec<Expr> = Vec::with_capacity(row.values.len());
3945 for v in row.values {
3946 tuple.push(value_to_literal_expr_permissive(v)?);
3947 }
3948 materialised.push(tuple);
3949 }
3950 let recurse = InsertStatement {
3951 table: stmt.table,
3952 columns: stmt.columns,
3953 rows: materialised,
3954 select_source: None,
3955 on_conflict: stmt.on_conflict,
3956 returning: stmt.returning,
3957 };
3958 return self.exec_insert(recurse);
3959 }
3960 let clock = self.clock;
3964 let before_insert_triggers = self.snapshot_row_triggers(&stmt.table, "INSERT", "BEFORE");
3970 let after_insert_triggers = self.snapshot_row_triggers(&stmt.table, "INSERT", "AFTER");
3971 let trigger_session_cfg: Option<alloc::string::String> = self
3972 .session_params
3973 .get("default_text_search_config")
3974 .cloned();
3975 let table = self
3976 .active_catalog_mut()
3977 .get_mut(&stmt.table)
3978 .ok_or_else(|| {
3979 EngineError::Storage(StorageError::TableNotFound {
3980 name: stmt.table.clone(),
3981 })
3982 })?;
3983 let column_meta: Vec<ColumnSchema> = table.schema().columns.clone();
3989 let schema_cols_len = column_meta.len();
3990 let tuple_pos: Option<Vec<Option<usize>>> = match &stmt.columns {
3994 None => None, Some(cols) => {
3996 let mut map = alloc::vec![None; schema_cols_len];
3997 for (j, name) in cols.iter().enumerate() {
3998 let idx = column_meta
3999 .iter()
4000 .position(|c| c.name == *name)
4001 .ok_or_else(|| {
4002 EngineError::Eval(EvalError::ColumnNotFound { name: name.clone() })
4003 })?;
4004 if map[idx].is_some() {
4005 return Err(EngineError::Storage(StorageError::ArityMismatch {
4006 expected: schema_cols_len,
4007 actual: cols.len(),
4008 }));
4009 }
4010 map[idx] = Some(j);
4011 }
4012 for (i, col) in column_meta.iter().enumerate() {
4016 if map[i].is_none()
4017 && !col.nullable
4018 && col.default.is_none()
4019 && col.runtime_default.is_none()
4020 && !col.auto_increment
4021 {
4022 return Err(EngineError::Storage(StorageError::NullInNotNull {
4023 column: col.name.clone(),
4024 }));
4025 }
4026 }
4027 Some(map)
4028 }
4029 };
4030 let expected_tuple_len = stmt.columns.as_ref().map_or(schema_cols_len, Vec::len);
4031 let fks = table.schema().foreign_keys.clone();
4037 let mut affected = 0usize;
4038 let mut all_values: Vec<Vec<Value>> = Vec::with_capacity(stmt.rows.len());
4041 for tuple in stmt.rows {
4042 if tuple.len() != expected_tuple_len {
4043 return Err(EngineError::Storage(StorageError::ArityMismatch {
4044 expected: expected_tuple_len,
4045 actual: tuple.len(),
4046 }));
4047 }
4048 let values: Vec<Value> = if let Some(map) = &tuple_pos {
4052 let raw_tuple: Vec<Value> = tuple
4054 .into_iter()
4055 .map(literal_expr_to_value)
4056 .collect::<Result<_, _>>()?;
4057 let mut out = Vec::with_capacity(schema_cols_len);
4058 for (i, col) in column_meta.iter().enumerate() {
4059 let mut raw = match map[i] {
4060 Some(j) => raw_tuple[j].clone(),
4061 None => resolve_column_default_free(col, clock)?,
4062 };
4063 if col.auto_increment && raw.is_null() {
4064 let next = table.next_auto_value(i).ok_or_else(|| {
4065 EngineError::Unsupported(alloc::format!(
4066 "AUTO_INCREMENT applies to integer columns only (column `{}`)",
4067 col.name
4068 ))
4069 })?;
4070 raw = Value::BigInt(next);
4071 }
4072 out.push(coerce_value(raw, col.ty, &col.name, i)?);
4073 }
4074 out
4075 } else {
4076 let mut out = Vec::with_capacity(schema_cols_len);
4078 for (i, (col, expr)) in column_meta.iter().zip(tuple).enumerate() {
4079 let mut raw = literal_expr_to_value(expr)?;
4080 if col.auto_increment && raw.is_null() {
4081 let next = table.next_auto_value(i).ok_or_else(|| {
4082 EngineError::Unsupported(alloc::format!(
4083 "AUTO_INCREMENT applies to integer columns only (column `{}`)",
4084 col.name
4085 ))
4086 })?;
4087 raw = Value::BigInt(next);
4088 }
4089 out.push(coerce_value(raw, col.ty, &col.name, i)?);
4090 }
4091 out
4092 };
4093 all_values.push(values);
4094 }
4095 let uniqueness = table.schema().uniqueness_constraints.clone();
4100 let _ = table;
4101 if !fks.is_empty() {
4102 enforce_fk_inserts(self.active_catalog(), &stmt.table, &fks, &all_values)?;
4103 }
4104 enforce_check_constraints(self.active_catalog(), &stmt.table, &all_values)?;
4106 enforce_uniqueness_inserts(self.active_catalog(), &stmt.table, &uniqueness, &all_values)?;
4108 enforce_unique_index_inserts(self.active_catalog(), &stmt.table, &all_values)?;
4115 let mut pending_updates: Vec<(usize, Vec<Value>)> = Vec::new();
4122 let mut skipped_count = 0usize;
4123 if let Some(clause) = &stmt.on_conflict {
4124 let conflict_cols = resolve_on_conflict_columns(
4125 self.active_catalog(),
4126 &stmt.table,
4127 clause.target_columns.as_slice(),
4128 )?;
4129 let mut kept: Vec<Vec<Value>> = Vec::with_capacity(all_values.len());
4130 let mut seen_keys: Vec<Vec<Value>> = Vec::new();
4131 for values in all_values {
4132 let key_tuple: Vec<&Value> = conflict_cols.iter().map(|&c| &values[c]).collect();
4133 let has_null_key = key_tuple.iter().any(|v| matches!(v, Value::Null));
4136 let collides_with_table = !has_null_key
4137 && on_conflict_keys_exist(
4138 self.active_catalog(),
4139 &stmt.table,
4140 &conflict_cols,
4141 &key_tuple,
4142 );
4143 let key_tuple_owned: Vec<Value> = key_tuple.iter().map(|v| (*v).clone()).collect();
4144 let collides_with_batch =
4145 !has_null_key && seen_keys.iter().any(|k| k == &key_tuple_owned);
4146 let collides = collides_with_table || collides_with_batch;
4147 match (&clause.action, collides) {
4148 (_, false) => {
4149 seen_keys.push(key_tuple_owned);
4150 kept.push(values);
4151 }
4152 (spg_sql::ast::OnConflictAction::Nothing, true) => {
4153 skipped_count += 1;
4154 }
4155 (
4156 spg_sql::ast::OnConflictAction::Update {
4157 assignments,
4158 where_,
4159 },
4160 true,
4161 ) => {
4162 if !collides_with_table {
4163 skipped_count += 1;
4164 continue;
4165 }
4166 let target_pos = lookup_row_position_by_keys(
4167 self.active_catalog(),
4168 &stmt.table,
4169 &conflict_cols,
4170 &key_tuple,
4171 )
4172 .ok_or_else(|| {
4173 EngineError::Unsupported(
4174 "ON CONFLICT DO UPDATE: conflict detected but row \
4175 position could not be resolved (cold-tier row?)"
4176 .into(),
4177 )
4178 })?;
4179 let updated = apply_on_conflict_assignments(
4180 self.active_catalog(),
4181 &stmt.table,
4182 target_pos,
4183 &values,
4184 assignments,
4185 where_.as_ref(),
4186 )?;
4187 if let Some(new_row) = updated {
4188 pending_updates.push((target_pos, new_row));
4189 } else {
4190 skipped_count += 1;
4191 }
4192 }
4193 }
4194 }
4195 all_values = kept;
4196 }
4197 let table = self
4199 .active_catalog_mut()
4200 .get_mut(&stmt.table)
4201 .ok_or_else(|| {
4202 EngineError::Storage(StorageError::TableNotFound {
4203 name: stmt.table.clone(),
4204 })
4205 })?;
4206 let mut returning_rows: Vec<Vec<Value>> = Vec::new();
4210 let mut deferred_embedded: Vec<triggers::DeferredEmbeddedStmt> = Vec::new();
4214 'rowloop: for values in all_values {
4215 let mut row = Row::new(values);
4216 for fd in &before_insert_triggers {
4221 let (outcome, deferred) = triggers::fire_row_trigger(
4222 fd,
4223 Some(row.clone()),
4224 None,
4225 &stmt.table,
4226 &column_meta,
4227 &[],
4228 trigger_session_cfg.as_deref(),
4229 false,
4230 )
4231 .map_err(|e| EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}"))))?;
4232 deferred_embedded.extend(deferred);
4233 match outcome {
4234 triggers::TriggerOutcome::Row(r) => row = r,
4235 triggers::TriggerOutcome::Skip => continue 'rowloop,
4236 }
4237 }
4238 if stmt.returning.is_some() {
4239 returning_rows.push(row.values.clone());
4240 }
4241 let inserted = row.clone();
4244 table.insert(row)?;
4245 affected += 1;
4246 for fd in &after_insert_triggers {
4250 let (_outcome, deferred) = triggers::fire_row_trigger(
4251 fd,
4252 Some(inserted.clone()),
4253 None,
4254 &stmt.table,
4255 &column_meta,
4256 &[],
4257 trigger_session_cfg.as_deref(),
4258 true,
4259 )
4260 .map_err(|e| EngineError::Storage(StorageError::Corrupt(alloc::format!("{e}"))))?;
4261 deferred_embedded.extend(deferred);
4262 }
4263 }
4264 for (pos, new_row) in pending_updates {
4268 if stmt.returning.is_some() {
4269 returning_rows.push(new_row.clone());
4270 }
4271 table.update_row(pos, new_row)?;
4272 affected += 1;
4273 }
4274 let _ = skipped_count;
4275 let _ = table;
4281 self.execute_deferred_trigger_stmts(deferred_embedded, CancelToken::none())?;
4282 if let Some(items) = &stmt.returning {
4286 return self.build_returning_rows(&stmt.table, items, returning_rows);
4287 }
4288 if !self.in_transaction() && affected > 0 {
4293 self.statistics
4294 .record_modifications(&stmt.table, affected as u64);
4295 }
4296 Ok(QueryResult::CommandOk {
4297 affected,
4298 modified_catalog: !self.in_transaction(),
4299 })
4300 }
4301
4302 fn exec_select_as_of_segment(
4315 &self,
4316 stmt: &SelectStatement,
4317 from: &spg_sql::ast::FromClause,
4318 segment_id: u32,
4319 ) -> Result<QueryResult, EngineError> {
4320 if !from.joins.is_empty()
4323 || stmt.group_by.is_some()
4324 || stmt.having.is_some()
4325 || !stmt.unions.is_empty()
4326 || !stmt.order_by.is_empty()
4327 || stmt.offset.is_some()
4328 || stmt.distinct
4329 || aggregate::uses_aggregate(stmt)
4330 {
4331 return Err(EngineError::Unsupported(
4332 "AS OF SEGMENT supports SELECT projection + WHERE + LIMIT only \
4333 (joins / aggregates / ORDER BY are STABILITY § \"Out of v6.10\")"
4334 .into(),
4335 ));
4336 }
4337 let table = self
4338 .active_catalog()
4339 .get(&from.primary.name)
4340 .ok_or_else(|| StorageError::TableNotFound {
4341 name: from.primary.name.clone(),
4342 })?;
4343 let schema = table.schema().clone();
4344 let schema_cols = &schema.columns;
4345 let alias = from
4346 .primary
4347 .alias
4348 .as_deref()
4349 .unwrap_or(from.primary.name.as_str());
4350 let ctx = EvalContext::new(schema_cols, Some(alias));
4351 let seg = self
4352 .active_catalog()
4353 .cold_segment(segment_id)
4354 .ok_or_else(|| {
4355 EngineError::Unsupported(alloc::format!(
4356 "AS OF SEGMENT: cold segment {segment_id} not registered"
4357 ))
4358 })?;
4359 let mut out_rows: Vec<Row> = Vec::new();
4360 let mut limit_remaining: Option<usize> =
4361 stmt.limit_literal().and_then(|n| usize::try_from(n).ok());
4362 for (_key, body) in seg.scan() {
4363 let (row, _consumed) =
4364 spg_storage::decode_row_body_dense(&body, &schema).map_err(EngineError::Storage)?;
4365 if let Some(where_expr) = &stmt.where_ {
4366 let cond = self.eval_expr_simple(where_expr, &row, &ctx)?;
4367 if !matches!(cond, Value::Bool(true)) {
4368 continue;
4369 }
4370 }
4371 let projected = self.project_row_simple(&row, &stmt.items, schema_cols, alias)?;
4373 out_rows.push(projected);
4374 if let Some(rem) = limit_remaining.as_mut() {
4375 if *rem == 0 {
4376 out_rows.pop();
4377 break;
4378 }
4379 *rem -= 1;
4380 }
4381 }
4382 let columns = self.derive_output_columns(&stmt.items, schema_cols, alias);
4384 Ok(QueryResult::Rows {
4385 columns,
4386 rows: out_rows,
4387 })
4388 }
4389
4390 fn eval_expr_simple(
4395 &self,
4396 expr: &Expr,
4397 row: &Row,
4398 ctx: &EvalContext,
4399 ) -> Result<Value, EngineError> {
4400 let cancel = CancelToken::none();
4401 self.eval_expr_with_correlated(expr, row, ctx, cancel, None)
4402 }
4403
4404 fn build_returning_rows(
4411 &self,
4412 table_name: &str,
4413 items: &[SelectItem],
4414 mutated_rows: Vec<Vec<Value>>,
4415 ) -> Result<QueryResult, EngineError> {
4416 let table = self.active_catalog().get(table_name).ok_or_else(|| {
4417 EngineError::Storage(StorageError::TableNotFound {
4418 name: table_name.into(),
4419 })
4420 })?;
4421 let schema_cols = table.schema().columns.clone();
4422 let columns = self.derive_output_columns(items, &schema_cols, table_name);
4423 let mut out_rows: Vec<Row> = Vec::with_capacity(mutated_rows.len());
4424 for values in mutated_rows {
4425 let row = Row::new(values);
4426 let projected = self.project_row_simple(&row, items, &schema_cols, table_name)?;
4427 out_rows.push(projected);
4428 }
4429 Ok(QueryResult::Rows {
4430 columns,
4431 rows: out_rows,
4432 })
4433 }
4434
4435 fn project_row_simple(
4439 &self,
4440 row: &Row,
4441 items: &[SelectItem],
4442 schema_cols: &[ColumnSchema],
4443 alias: &str,
4444 ) -> Result<Row, EngineError> {
4445 let ctx = EvalContext::new(schema_cols, Some(alias));
4446 let cancel = CancelToken::none();
4447 let mut out_vals = Vec::new();
4448 for item in items {
4449 match item {
4450 SelectItem::Wildcard => {
4451 out_vals.extend(row.values.iter().cloned());
4452 }
4453 SelectItem::Expr { expr, .. } => {
4454 let v = self.eval_expr_with_correlated(expr, row, &ctx, cancel, None)?;
4455 out_vals.push(v);
4456 }
4457 }
4458 }
4459 Ok(Row::new(out_vals))
4460 }
4461
4462 fn derive_output_columns(
4467 &self,
4468 items: &[SelectItem],
4469 schema_cols: &[ColumnSchema],
4470 _alias: &str,
4471 ) -> Vec<ColumnSchema> {
4472 let mut out = Vec::new();
4473 for item in items {
4474 match item {
4475 SelectItem::Wildcard => {
4476 out.extend(schema_cols.iter().cloned());
4477 }
4478 SelectItem::Expr { alias, .. } => {
4479 let name = alias.clone().unwrap_or_else(|| "?column?".to_string());
4480 out.push(ColumnSchema::new(name, DataType::Text, true));
4483 }
4484 }
4485 }
4486 out
4487 }
4488
4489 fn exec_select_cancel(
4490 &self,
4491 stmt: &SelectStatement,
4492 cancel: CancelToken<'_>,
4493 ) -> Result<QueryResult, EngineError> {
4494 cancel.check()?;
4495 if let Some(from) = &stmt.from
4504 && let Some(seg_id) = from.primary.as_of_segment
4505 {
4506 return self.exec_select_as_of_segment(stmt, from, seg_id);
4507 }
4508 if let Some(from) = &stmt.from
4512 && from.joins.is_empty()
4513 && stmt.where_.is_none()
4514 && stmt.group_by.is_none()
4515 && stmt.having.is_none()
4516 && stmt.unions.is_empty()
4517 && stmt.order_by.is_empty()
4518 && stmt.limit.is_none()
4519 && stmt.offset.is_none()
4520 && !stmt.distinct
4521 && stmt.items.iter().all(|i| matches!(i, SelectItem::Wildcard))
4522 {
4523 let lower = from.primary.name.to_ascii_lowercase();
4524 match lower.as_str() {
4525 "spg_statistic" => return Ok(self.exec_spg_statistic()),
4526 "spg_stat_replication" => return Ok(self.exec_spg_stat_replication()),
4528 "spg_stat_segment" => return Ok(self.exec_spg_stat_segment()),
4529 "spg_stat_query" => return Ok(self.exec_spg_stat_query()),
4530 "spg_stat_activity" => return Ok(self.exec_spg_stat_activity()),
4531 "spg_audit_chain" => return Ok(self.exec_spg_audit_chain()),
4532 "spg_audit_verify" => return Ok(self.exec_spg_audit_verify()),
4533 "spg_table_ddl" => return Ok(self.exec_spg_table_ddl()),
4534 "spg_role_ddl" => return Ok(self.exec_spg_role_ddl()),
4535 "spg_database_ddl" => return Ok(self.exec_spg_database_ddl()),
4536 _ => {}
4537 }
4538 }
4539 if !stmt.ctes.is_empty() {
4547 return self.exec_with_ctes(stmt, cancel);
4548 }
4549 let mut stmt_owned;
4556 let stmt_ref: &SelectStatement = if expr_tree_has_subquery(stmt) {
4557 stmt_owned = stmt.clone();
4558 self.resolve_select_subqueries(&mut stmt_owned, cancel)?;
4559 &stmt_owned
4560 } else {
4561 stmt
4562 };
4563 if stmt_ref.unions.is_empty() {
4564 return self.exec_bare_select_cancel(stmt_ref, cancel);
4565 }
4566 let mut head = stmt_ref.clone();
4571 head.unions = Vec::new();
4572 head.order_by = Vec::new();
4573 head.limit = None;
4574 let QueryResult::Rows { columns, mut rows } =
4575 self.exec_bare_select_cancel(&head, cancel)?
4576 else {
4577 unreachable!("bare SELECT cannot return CommandOk")
4578 };
4579 for (kind, peer) in &stmt_ref.unions {
4580 let QueryResult::Rows {
4581 columns: peer_cols,
4582 rows: peer_rows,
4583 } = self.exec_bare_select_cancel(peer, cancel)?
4584 else {
4585 unreachable!("bare SELECT cannot return CommandOk")
4586 };
4587 if peer_cols.len() != columns.len() {
4588 return Err(EngineError::Unsupported(alloc::format!(
4589 "UNION arity mismatch: head has {} columns, peer has {}",
4590 columns.len(),
4591 peer_cols.len()
4592 )));
4593 }
4594 rows.extend(peer_rows);
4595 if matches!(kind, UnionKind::Distinct) {
4596 rows = dedup_rows(rows);
4597 }
4598 }
4599 if !stmt.order_by.is_empty() {
4602 let synth_ctx = EvalContext::new(&columns, None);
4603 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
4604 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::with_capacity(rows.len());
4605 for r in rows {
4606 let keys = build_order_keys(&stmt.order_by, &r, &synth_ctx)?;
4607 tagged.push((keys, r));
4608 }
4609 sort_by_keys(&mut tagged, &descs);
4610 rows = tagged.into_iter().map(|(_, r)| r).collect();
4611 }
4612 apply_offset_and_limit(&mut rows, stmt.offset_literal(), stmt.limit_literal());
4613 Ok(QueryResult::Rows { columns, rows })
4614 }
4615
4616 #[allow(clippy::too_many_lines)]
4617 #[allow(clippy::too_many_lines)] fn exec_select_unnest(
4625 &self,
4626 stmt: &SelectStatement,
4627 primary: &TableRef,
4628 cancel: CancelToken<'_>,
4629 ) -> Result<QueryResult, EngineError> {
4630 let expr = primary
4631 .unnest_expr
4632 .as_deref()
4633 .expect("caller guards unnest_expr.is_some()");
4634 let empty_schema: alloc::vec::Vec<ColumnSchema> = alloc::vec::Vec::new();
4637 let ctx = EvalContext::new(&empty_schema, None);
4638 let dummy_row = Row::new(alloc::vec::Vec::new());
4639 let (elem_dtype, rows): (DataType, alloc::vec::Vec<Row>) =
4642 match eval::eval_expr(expr, &dummy_row, &ctx).map_err(EngineError::Eval)? {
4643 Value::Null => (DataType::Text, alloc::vec::Vec::new()),
4644 Value::TextArray(items) => {
4645 let rows = items
4646 .into_iter()
4647 .map(|item| {
4648 Row::new(alloc::vec![match item {
4649 Some(s) => Value::Text(s),
4650 None => Value::Null,
4651 }])
4652 })
4653 .collect();
4654 (DataType::Text, rows)
4655 }
4656 Value::IntArray(items) => {
4657 let rows = items
4658 .into_iter()
4659 .map(|item| {
4660 Row::new(alloc::vec![match item {
4661 Some(n) => Value::Int(n),
4662 None => Value::Null,
4663 }])
4664 })
4665 .collect();
4666 (DataType::Int, rows)
4667 }
4668 Value::BigIntArray(items) => {
4669 let rows = items
4670 .into_iter()
4671 .map(|item| {
4672 Row::new(alloc::vec![match item {
4673 Some(n) => Value::BigInt(n),
4674 None => Value::Null,
4675 }])
4676 })
4677 .collect();
4678 (DataType::BigInt, rows)
4679 }
4680 other => {
4681 return Err(EngineError::Unsupported(alloc::format!(
4682 "unnest() expects an array argument, got {:?}",
4683 other.data_type()
4684 )));
4685 }
4686 };
4687 let alias = primary
4688 .alias
4689 .clone()
4690 .unwrap_or_else(|| "unnest".to_string());
4691 let col_name = primary
4697 .unnest_column_aliases
4698 .first()
4699 .cloned()
4700 .unwrap_or_else(|| alias.clone());
4701 let col_schema = ColumnSchema::new(col_name, elem_dtype, true);
4702 let schema_cols = alloc::vec![col_schema.clone()];
4703 let scan_ctx = EvalContext::new(&schema_cols, Some(&alias));
4704 let filtered: alloc::vec::Vec<Row> = if let Some(w) = &stmt.where_ {
4706 let mut out = alloc::vec::Vec::with_capacity(rows.len());
4707 for row in rows {
4708 cancel.check()?;
4709 let v = eval::eval_expr(w, &row, &scan_ctx).map_err(EngineError::Eval)?;
4710 if matches!(v, Value::Bool(true)) {
4711 out.push(row);
4712 }
4713 }
4714 out
4715 } else {
4716 rows
4717 };
4718 let projection = build_projection(&stmt.items, &schema_cols, &alias)?;
4720 let mut projected_rows: alloc::vec::Vec<Row> =
4721 alloc::vec::Vec::with_capacity(filtered.len());
4722 for row in &filtered {
4723 let mut vals = alloc::vec::Vec::with_capacity(projection.len());
4724 for p in &projection {
4725 vals.push(eval::eval_expr(&p.expr, row, &scan_ctx).map_err(EngineError::Eval)?);
4726 }
4727 projected_rows.push(Row::new(vals));
4728 }
4729 let columns: alloc::vec::Vec<ColumnSchema> = projection
4732 .iter()
4733 .map(|p| ColumnSchema::new(p.output_name.clone(), p.ty, p.nullable))
4734 .collect();
4735 if !stmt.order_by.is_empty() {
4738 let mut indexed: alloc::vec::Vec<(usize, Vec<Value>)> = filtered
4739 .iter()
4740 .enumerate()
4741 .map(|(i, r)| -> Result<_, EngineError> {
4742 let keys: Result<Vec<Value>, EngineError> = stmt
4743 .order_by
4744 .iter()
4745 .map(|ob| {
4746 eval::eval_expr(&ob.expr, r, &scan_ctx).map_err(EngineError::Eval)
4747 })
4748 .collect();
4749 Ok((i, keys?))
4750 })
4751 .collect::<Result<_, _>>()?;
4752 indexed.sort_by(|a, b| {
4753 for (idx, (ka, kb)) in a.1.iter().zip(b.1.iter()).enumerate() {
4754 let mut cmp = value_cmp(ka, kb);
4755 if stmt.order_by[idx].desc {
4756 cmp = cmp.reverse();
4757 }
4758 if cmp != core::cmp::Ordering::Equal {
4759 return cmp;
4760 }
4761 }
4762 core::cmp::Ordering::Equal
4763 });
4764 projected_rows = indexed
4765 .into_iter()
4766 .map(|(i, _)| projected_rows[i].clone())
4767 .collect();
4768 }
4769 if let Some(offset) = stmt.offset_literal() {
4771 let off = (offset as usize).min(projected_rows.len());
4772 projected_rows.drain(..off);
4773 }
4774 if let Some(limit) = stmt.limit_literal() {
4775 projected_rows.truncate(limit as usize);
4776 }
4777 Ok(QueryResult::Rows {
4778 columns,
4779 rows: projected_rows,
4780 })
4781 }
4782
4783 fn exec_bare_select_cancel(
4784 &self,
4785 stmt: &SelectStatement,
4786 cancel: CancelToken<'_>,
4787 ) -> Result<QueryResult, EngineError> {
4788 if select_has_window(stmt) {
4793 return self.exec_select_with_window(stmt, cancel);
4794 }
4795 let Some(from) = &stmt.from else {
4800 let empty_schema: Vec<ColumnSchema> = Vec::new();
4801 let ctx = self.ev_ctx(&empty_schema, None);
4802 let projection = build_projection(&stmt.items, &empty_schema, "")?;
4803 let dummy_row = Row::new(Vec::new());
4804 let mut values = Vec::with_capacity(projection.len());
4805 for p in &projection {
4806 values.push(eval::eval_expr(&p.expr, &dummy_row, &ctx)?);
4807 }
4808 let columns: Vec<ColumnSchema> = projection
4809 .into_iter()
4810 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
4811 .collect();
4812 return Ok(QueryResult::Rows {
4813 columns,
4814 rows: alloc::vec![Row::new(values)],
4815 });
4816 };
4817 if !from.joins.is_empty() {
4821 return self.exec_joined_select(stmt, from);
4822 }
4823 if from.primary.unnest_expr.is_some() {
4830 return self.exec_select_unnest(stmt, &from.primary, cancel);
4831 }
4832 let primary = &from.primary;
4833 let table = self.active_catalog().get(&primary.name).ok_or_else(|| {
4834 StorageError::TableNotFound {
4835 name: primary.name.clone(),
4836 }
4837 })?;
4838 let schema_cols = &table.schema().columns;
4839 let alias = primary.alias.as_deref().unwrap_or(primary.name.as_str());
4842 let ctx = self.ev_ctx(schema_cols, Some(alias));
4843
4844 if let Some(nsw_rows) = try_nsw_knn(stmt, table, schema_cols, alias) {
4849 return materialise_in_order(stmt, table, schema_cols, alias, &nsw_rows);
4850 }
4851
4852 let indexed_rows: Option<Vec<Cow<'_, Row>>> = stmt.where_.as_ref().and_then(|w| {
4860 try_index_seek(w, schema_cols, self.active_catalog(), table, alias).or_else(|| {
4863 try_gin_seek(w, schema_cols, self.active_catalog(), table, alias, &ctx)
4868 })
4869 });
4870
4871 if aggregate::uses_aggregate(stmt) {
4874 let mut filtered: Vec<&Row> = Vec::new();
4875 let mut memo = memoize::MemoizeCache::new();
4879 if let Some(rows) = &indexed_rows {
4880 for cow in rows {
4881 let row = cow.as_ref();
4882 if let Some(where_expr) = &stmt.where_ {
4883 let cond = self.eval_expr_with_correlated(
4884 where_expr,
4885 row,
4886 &ctx,
4887 cancel,
4888 Some(&mut memo),
4889 )?;
4890 if !matches!(cond, Value::Bool(true)) {
4891 continue;
4892 }
4893 }
4894 filtered.push(row);
4895 }
4896 } else {
4897 for i in 0..table.row_count() {
4898 let row = &table.rows()[i];
4899 if let Some(where_expr) = &stmt.where_ {
4900 let cond = self.eval_expr_with_correlated(
4901 where_expr,
4902 row,
4903 &ctx,
4904 cancel,
4905 Some(&mut memo),
4906 )?;
4907 if !matches!(cond, Value::Bool(true)) {
4908 continue;
4909 }
4910 }
4911 filtered.push(row);
4912 }
4913 }
4914 let mut agg = aggregate::run(stmt, &filtered, schema_cols, Some(alias))?;
4915 apply_offset_and_limit(&mut agg.rows, stmt.offset_literal(), stmt.limit_literal());
4916 return Ok(QueryResult::Rows {
4917 columns: agg.columns,
4918 rows: agg.rows,
4919 });
4920 }
4921
4922 let projection = build_projection(&stmt.items, schema_cols, alias)?;
4923
4924 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::new();
4927 let mut memo = memoize::MemoizeCache::new();
4929 let mut process_row = |row: &Row, loop_idx: usize| -> Result<(), EngineError> {
4932 if loop_idx.is_multiple_of(256) {
4933 cancel.check()?;
4934 }
4935 if let Some(where_expr) = &stmt.where_ {
4936 let cond =
4937 self.eval_expr_with_correlated(where_expr, row, &ctx, cancel, Some(&mut memo))?;
4938 if !matches!(cond, Value::Bool(true)) {
4939 return Ok(());
4940 }
4941 }
4942 let mut values = Vec::with_capacity(projection.len());
4943 for p in &projection {
4944 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
4945 }
4946 let order_keys = if stmt.order_by.is_empty() {
4947 Vec::new()
4948 } else {
4949 build_order_keys(&stmt.order_by, row, &ctx)?
4950 };
4951 tagged.push((order_keys, Row::new(values)));
4952 Ok(())
4953 };
4954 if let Some(rows) = &indexed_rows {
4955 for (loop_idx, cow) in rows.iter().enumerate() {
4956 process_row(cow.as_ref(), loop_idx)?;
4957 }
4958 } else {
4959 for i in 0..table.row_count() {
4960 process_row(&table.rows()[i], i)?;
4961 }
4962 }
4963
4964 if !stmt.order_by.is_empty() {
4965 let keep = if stmt.distinct {
4970 None
4971 } else {
4972 stmt.limit_literal()
4973 .map(|l| l as usize + stmt.offset_literal().map_or(0, |o| o as usize))
4974 };
4975 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
4976 partial_sort_tagged(&mut tagged, keep, &descs);
4977 }
4978
4979 let mut output_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
4980 if stmt.distinct {
4981 output_rows = dedup_rows(output_rows);
4982 }
4983 apply_offset_and_limit(
4984 &mut output_rows,
4985 stmt.offset_literal(),
4986 stmt.limit_literal(),
4987 );
4988
4989 let columns: Vec<ColumnSchema> = projection
4990 .into_iter()
4991 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
4992 .collect();
4993
4994 Ok(QueryResult::Rows {
4995 columns,
4996 rows: output_rows,
4997 })
4998 }
4999
5000 #[allow(clippy::too_many_lines)]
5007 fn materialise_table_ref(
5015 &self,
5016 tref: &TableRef,
5017 ) -> Result<(Vec<Row>, Vec<ColumnSchema>), EngineError> {
5018 if let Some(expr) = tref.unnest_expr.as_deref() {
5019 let empty_schema: Vec<ColumnSchema> = Vec::new();
5020 let ctx = EvalContext::new(&empty_schema, None);
5021 let dummy_row = Row::new(Vec::new());
5022 let (elem_dtype, rows) =
5023 match eval::eval_expr(expr, &dummy_row, &ctx).map_err(EngineError::Eval)? {
5024 Value::Null => (DataType::Text, Vec::new()),
5025 Value::TextArray(items) => (
5026 DataType::Text,
5027 items
5028 .into_iter()
5029 .map(|item| {
5030 Row::new(alloc::vec![match item {
5031 Some(s) => Value::Text(s),
5032 None => Value::Null,
5033 }])
5034 })
5035 .collect(),
5036 ),
5037 Value::IntArray(items) => (
5038 DataType::Int,
5039 items
5040 .into_iter()
5041 .map(|item| {
5042 Row::new(alloc::vec![match item {
5043 Some(n) => Value::Int(n),
5044 None => Value::Null,
5045 }])
5046 })
5047 .collect(),
5048 ),
5049 Value::BigIntArray(items) => (
5050 DataType::BigInt,
5051 items
5052 .into_iter()
5053 .map(|item| {
5054 Row::new(alloc::vec![match item {
5055 Some(n) => Value::BigInt(n),
5056 None => Value::Null,
5057 }])
5058 })
5059 .collect(),
5060 ),
5061 other => {
5062 return Err(EngineError::Unsupported(alloc::format!(
5063 "unnest() expects an array argument, got {:?}",
5064 other.data_type()
5065 )));
5066 }
5067 };
5068 let alias = tref.alias.clone().unwrap_or_else(|| "unnest".to_string());
5069 let col_name = tref
5070 .unnest_column_aliases
5071 .first()
5072 .cloned()
5073 .unwrap_or(alias);
5074 return Ok((rows, alloc::vec![ColumnSchema::new(col_name, elem_dtype, true)]));
5075 }
5076 let table = self
5077 .active_catalog()
5078 .get(&tref.name)
5079 .ok_or_else(|| StorageError::TableNotFound {
5080 name: tref.name.clone(),
5081 })?;
5082 let rows: Vec<Row> = table.rows().iter().cloned().collect();
5083 let cols = table.schema().columns.clone();
5084 Ok((rows, cols))
5085 }
5086
5087 fn exec_joined_select(
5088 &self,
5089 stmt: &SelectStatement,
5090 from: &FromClause,
5091 ) -> Result<QueryResult, EngineError> {
5092 let (primary_rows, primary_cols) = self.materialise_table_ref(&from.primary)?;
5098 let primary_alias = from
5099 .primary
5100 .alias
5101 .as_deref()
5102 .unwrap_or(from.primary.name.as_str())
5103 .to_string();
5104 let mut joined: Vec<(Vec<Row>, Vec<ColumnSchema>, String, JoinKind, Option<&Expr>)> =
5107 Vec::new();
5108 for j in &from.joins {
5109 let (rows, cols) = self.materialise_table_ref(&j.table)?;
5110 let a = j
5111 .table
5112 .alias
5113 .as_deref()
5114 .unwrap_or(j.table.name.as_str())
5115 .to_string();
5116 joined.push((rows, cols, a, j.kind, j.on.as_ref()));
5117 }
5118
5119 let mut combined_schema: Vec<ColumnSchema> = Vec::new();
5122 for col in &primary_cols {
5123 combined_schema.push(ColumnSchema::new(
5124 alloc::format!("{primary_alias}.{}", col.name),
5125 col.ty,
5126 col.nullable,
5127 ));
5128 }
5129 for (_, cols, a, _, _) in &joined {
5130 for col in cols {
5131 combined_schema.push(ColumnSchema::new(
5132 alloc::format!("{a}.{}", col.name),
5133 col.ty,
5134 col.nullable,
5135 ));
5136 }
5137 }
5138 let ctx = EvalContext::new(&combined_schema, None);
5139
5140 let mut working: Vec<Row> = primary_rows;
5142 let mut produced_len = primary_cols.len();
5143 for (rrows, rcols, _, kind, on) in &joined {
5144 let right_arity = rcols.len();
5145 let mut next: Vec<Row> = Vec::new();
5146 for left in &working {
5147 let mut left_matched = false;
5148 for right in rrows {
5149 let mut combined_vals = left.values.clone();
5150 combined_vals.extend(right.values.iter().cloned());
5151 let combined = Row::new(combined_vals);
5154 let keep = if let Some(on_expr) = on {
5155 let cond = eval::eval_expr(on_expr, &combined, &ctx)?;
5156 matches!(cond, Value::Bool(true))
5157 } else {
5158 true
5160 };
5161 if keep {
5162 next.push(combined);
5163 left_matched = true;
5164 }
5165 }
5166 if !left_matched && matches!(kind, JoinKind::Left) {
5167 let mut combined_vals = left.values.clone();
5170 for _ in 0..right_arity {
5171 combined_vals.push(Value::Null);
5172 }
5173 next.push(Row::new(combined_vals));
5174 }
5175 }
5176 working = next;
5177 produced_len += right_arity;
5178 debug_assert!(produced_len <= combined_schema.len());
5179 }
5180
5181 let mut filtered: Vec<Row> = Vec::new();
5183 for row in working {
5184 if let Some(where_expr) = &stmt.where_ {
5185 let cond = eval::eval_expr(where_expr, &row, &ctx)?;
5186 if !matches!(cond, Value::Bool(true)) {
5187 continue;
5188 }
5189 }
5190 filtered.push(row);
5191 }
5192
5193 if aggregate::uses_aggregate(stmt) {
5196 let refs: Vec<&Row> = filtered.iter().collect();
5197 let mut agg = aggregate::run(stmt, &refs, &combined_schema, None)?;
5198 apply_offset_and_limit(&mut agg.rows, stmt.offset_literal(), stmt.limit_literal());
5199 return Ok(QueryResult::Rows {
5200 columns: agg.columns,
5201 rows: agg.rows,
5202 });
5203 }
5204
5205 let projection = build_projection(&stmt.items, &combined_schema, "")?;
5206 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::new();
5207 for row in &filtered {
5208 let mut values = Vec::with_capacity(projection.len());
5209 for p in &projection {
5210 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
5211 }
5212 let order_keys = if stmt.order_by.is_empty() {
5213 Vec::new()
5214 } else {
5215 build_order_keys(&stmt.order_by, row, &ctx)?
5216 };
5217 tagged.push((order_keys, Row::new(values)));
5218 }
5219 if !stmt.order_by.is_empty() {
5220 let keep = if stmt.distinct {
5221 None
5222 } else {
5223 stmt.limit_literal()
5224 .map(|l| l as usize + stmt.offset_literal().map_or(0, |o| o as usize))
5225 };
5226 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
5227 partial_sort_tagged(&mut tagged, keep, &descs);
5228 }
5229 let mut output_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
5230 if stmt.distinct {
5231 output_rows = dedup_rows(output_rows);
5232 }
5233 apply_offset_and_limit(
5234 &mut output_rows,
5235 stmt.offset_literal(),
5236 stmt.limit_literal(),
5237 );
5238 let columns: Vec<ColumnSchema> = projection
5239 .into_iter()
5240 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
5241 .collect();
5242 Ok(QueryResult::Rows {
5243 columns,
5244 rows: output_rows,
5245 })
5246 }
5247}
5248
5249#[derive(Debug, Clone)]
5252struct ProjectedItem {
5253 expr: Expr,
5254 output_name: String,
5255 ty: DataType,
5256 nullable: bool,
5257}
5258
5259fn dedup_rows(rows: Vec<Row>) -> Vec<Row> {
5265 let mut out: Vec<Row> = Vec::with_capacity(rows.len());
5266 for r in rows {
5267 if !out.iter().any(|seen| seen == &r) {
5268 out.push(r);
5269 }
5270 }
5271 out
5272}
5273
5274fn value_to_order_key(v: &Value) -> Result<f64, EngineError> {
5278 match v {
5279 Value::Null => Ok(f64::INFINITY),
5280 Value::SmallInt(n) => Ok(f64::from(*n)),
5281 Value::Int(n) => Ok(f64::from(*n)),
5282 Value::Date(d) => Ok(f64::from(*d)),
5283 #[allow(clippy::cast_precision_loss)]
5284 Value::Timestamp(t) => Ok(*t as f64),
5285 #[allow(clippy::cast_precision_loss)]
5286 Value::Numeric { scaled, scale } => {
5287 let mut divisor = 1.0_f64;
5293 for _ in 0..*scale {
5294 divisor *= 10.0;
5295 }
5296 Ok((*scaled as f64) / divisor)
5297 }
5298 #[allow(clippy::cast_precision_loss)]
5299 Value::BigInt(n) => Ok(*n as f64),
5300 Value::Float(x) => Ok(*x),
5301 Value::Bool(b) => Ok(if *b { 1.0 } else { 0.0 }),
5302 Value::Text(s) => {
5303 let mut key: u64 = 0;
5307 for &b in s.as_bytes().iter().take(8) {
5308 key = (key << 8) | u64::from(b);
5309 }
5310 #[allow(clippy::cast_precision_loss)]
5311 Ok(key as f64)
5312 }
5313 Value::Vector(_) | Value::Sq8Vector(_) | Value::HalfVector(_) => {
5314 Err(EngineError::Unsupported(
5315 "ORDER BY of a raw vector column is not meaningful — use `<->`".into(),
5316 ))
5317 }
5318 Value::Interval { .. } => Err(EngineError::Unsupported(
5319 "ORDER BY of an INTERVAL is not supported in v2.11 \
5320 (months vs micros has no single canonical ordering)"
5321 .into(),
5322 )),
5323 Value::Json(_) => Err(EngineError::Unsupported(
5324 "ORDER BY of a JSON value is not supported — cast the document to text first".into(),
5325 )),
5326 _ => Err(EngineError::Unsupported(
5330 "ORDER BY of this value type is not supported".into(),
5331 )),
5332 }
5333}
5334
5335fn try_nsw_knn(
5349 stmt: &SelectStatement,
5350 table: &Table,
5351 schema_cols: &[ColumnSchema],
5352 table_alias: &str,
5353) -> Option<Vec<usize>> {
5354 if stmt.distinct {
5355 return None;
5356 }
5357 let limit = usize::try_from(stmt.limit_literal()?).ok()?;
5358 if limit == 0 {
5359 return None;
5360 }
5361 if stmt.order_by.len() != 1 {
5365 return None;
5366 }
5367 let order = &stmt.order_by[0];
5368 if order.desc {
5372 return None;
5373 }
5374 let Expr::Binary { lhs, op, rhs } = &order.expr else {
5375 return None;
5376 };
5377 let metric = match op {
5378 BinOp::L2Distance => spg_storage::NswMetric::L2,
5379 BinOp::InnerProduct => spg_storage::NswMetric::InnerProduct,
5380 BinOp::CosineDistance => spg_storage::NswMetric::Cosine,
5381 _ => return None,
5382 };
5383 let ((Expr::Column(col), literal) | (literal, Expr::Column(col))) =
5385 (lhs.as_ref(), rhs.as_ref())
5386 else {
5387 return None;
5388 };
5389 if let Some(q) = &col.qualifier
5390 && q != table_alias
5391 {
5392 return None;
5393 }
5394 let col_pos = schema_cols.iter().position(|s| s.name == col.name)?;
5395 let query = literal_to_vector(literal)?;
5396 let idx = spg_storage::nsw_index_on(table, col_pos)?;
5397 if let Some(where_expr) = &stmt.where_ {
5398 let over_fetch = limit.saturating_mul(10).max(NSW_OVER_FETCH_FLOOR);
5402 let candidates = spg_storage::nsw_query(table, &idx.name, &query, over_fetch, metric);
5403 let ctx = EvalContext::new(schema_cols, Some(table_alias));
5404 let mut kept: Vec<usize> = Vec::with_capacity(limit);
5405 for i in candidates {
5406 let row = &table.rows()[i];
5407 let cond = eval::eval_expr(where_expr, row, &ctx).ok()?;
5408 if matches!(cond, Value::Bool(true)) {
5409 kept.push(i);
5410 if kept.len() >= limit {
5411 break;
5412 }
5413 }
5414 }
5415 Some(kept)
5416 } else {
5417 Some(spg_storage::nsw_query(
5418 table, &idx.name, &query, limit, metric,
5419 ))
5420 }
5421}
5422
5423const NSW_OVER_FETCH_FLOOR: usize = 32;
5427
5428fn literal_to_vector(e: &Expr) -> Option<Vec<f32>> {
5431 match e {
5432 Expr::Literal(Literal::Vector(v)) => Some(v.clone()),
5433 Expr::Cast { expr, .. } => literal_to_vector(expr),
5434 _ => None,
5435 }
5436}
5437
5438fn materialise_in_order(
5442 stmt: &SelectStatement,
5443 table: &Table,
5444 schema_cols: &[ColumnSchema],
5445 table_alias: &str,
5446 ordered_rows: &[usize],
5447) -> Result<QueryResult, EngineError> {
5448 let ctx = EvalContext::new(schema_cols, Some(table_alias));
5449 let projection = build_projection(&stmt.items, schema_cols, table_alias)?;
5450 let mut output_rows: Vec<Row> = Vec::with_capacity(ordered_rows.len());
5451 for &i in ordered_rows {
5452 let row = &table.rows()[i];
5453 let mut values = Vec::with_capacity(projection.len());
5454 for p in &projection {
5455 values.push(eval::eval_expr(&p.expr, row, &ctx)?);
5456 }
5457 output_rows.push(Row::new(values));
5458 }
5459 apply_offset_and_limit(
5460 &mut output_rows,
5461 stmt.offset_literal(),
5462 stmt.limit_literal(),
5463 );
5464 let columns: Vec<ColumnSchema> = projection
5465 .into_iter()
5466 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
5467 .collect();
5468 Ok(QueryResult::Rows {
5469 columns,
5470 rows: output_rows,
5471 })
5472}
5473
5474fn try_index_seek<'a>(
5475 where_expr: &Expr,
5476 schema_cols: &[ColumnSchema],
5477 catalog: &'a Catalog,
5478 table: &'a Table,
5479 table_alias: &str,
5480) -> Option<Vec<Cow<'a, Row>>> {
5481 if let Expr::Binary {
5488 lhs,
5489 op: BinOp::And,
5490 rhs,
5491 } = where_expr
5492 {
5493 if let Some(rows) = try_index_seek(lhs, schema_cols, catalog, table, table_alias) {
5496 return Some(rows);
5497 }
5498 return try_index_seek(rhs, schema_cols, catalog, table, table_alias);
5499 }
5500 let Expr::Binary {
5501 lhs,
5502 op: BinOp::Eq,
5503 rhs,
5504 } = where_expr
5505 else {
5506 return None;
5507 };
5508 let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
5509 .or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
5510 let idx = table.index_on(col_pos)?;
5511 let key = IndexKey::from_value(&value)?;
5512 let locators = idx.lookup_eq(&key);
5513 let table_name = table.schema().name.as_str();
5514 let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(locators.len());
5522 for loc in locators {
5523 match *loc {
5524 spg_storage::RowLocator::Hot(i) => {
5525 if let Some(row) = table.rows().get(i) {
5526 out.push(Cow::Borrowed(row));
5527 }
5528 }
5529 spg_storage::RowLocator::Cold { segment_id, .. } => {
5530 if let Some(row) = catalog.resolve_cold_locator(table_name, segment_id, &key) {
5531 out.push(Cow::Owned(row));
5532 }
5533 }
5534 }
5535 }
5536 Some(out)
5537}
5538
5539fn try_gin_seek<'a>(
5558 where_expr: &Expr,
5559 schema_cols: &[ColumnSchema],
5560 catalog: &'a Catalog,
5561 table: &'a Table,
5562 table_alias: &str,
5563 ctx: &eval::EvalContext<'_>,
5564) -> Option<Vec<Cow<'a, Row>>> {
5565 if let Expr::Binary {
5566 lhs,
5567 op: BinOp::And,
5568 rhs,
5569 } = where_expr
5570 {
5571 if let Some(rows) = try_gin_seek(lhs, schema_cols, catalog, table, table_alias, ctx) {
5572 return Some(rows);
5573 }
5574 return try_gin_seek(rhs, schema_cols, catalog, table, table_alias, ctx);
5575 }
5576 let Expr::Binary {
5577 lhs,
5578 op: BinOp::TsMatch,
5579 rhs,
5580 } = where_expr
5581 else {
5582 return None;
5583 };
5584 let (col_pos, query) = resolve_gin_col_query(lhs, rhs, schema_cols, table_alias, ctx)
5589 .or_else(|| resolve_gin_col_query(rhs, lhs, schema_cols, table_alias, ctx))?;
5590 let idx = table
5591 .indices()
5592 .iter()
5593 .find(|i| i.column_position == col_pos && i.is_gin())?;
5594 let candidates = gin_query_candidates(idx, &query)?;
5595 let _ = catalog; let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(candidates.len());
5597 for loc in candidates {
5598 match loc {
5599 spg_storage::RowLocator::Hot(i) => {
5600 if let Some(row) = table.rows().get(i) {
5601 out.push(Cow::Borrowed(row));
5602 }
5603 }
5604 spg_storage::RowLocator::Cold { .. } => {}
5611 }
5612 }
5613 Some(out)
5614}
5615
5616fn resolve_gin_col_query(
5622 col_side: &Expr,
5623 query_side: &Expr,
5624 schema_cols: &[ColumnSchema],
5625 table_alias: &str,
5626 ctx: &eval::EvalContext<'_>,
5627) -> Option<(usize, spg_storage::TsQueryAst)> {
5628 let Expr::Column(c) = col_side else {
5629 return None;
5630 };
5631 if let Some(q) = &c.qualifier
5632 && q != table_alias
5633 {
5634 return None;
5635 }
5636 let pos = schema_cols.iter().position(|s| s.name == c.name)?;
5637 let empty_row = Row::new(Vec::new());
5641 let v = eval::eval_expr(query_side, &empty_row, ctx).ok()?;
5642 let Value::TsQuery(q) = v else { return None };
5643 Some((pos, q))
5644}
5645
5646fn gin_query_candidates(
5657 idx: &spg_storage::Index,
5658 query: &spg_storage::TsQueryAst,
5659) -> Option<Vec<spg_storage::RowLocator>> {
5660 use spg_storage::TsQueryAst;
5661 match query {
5662 TsQueryAst::Term { word, .. } => {
5663 let mut v: Vec<spg_storage::RowLocator> = idx.gin_lookup_word(word).to_vec();
5664 v.sort_by_key(locator_sort_key);
5665 v.dedup_by_key(|l| locator_sort_key(l));
5666 Some(v)
5667 }
5668 TsQueryAst::And(l, r) => {
5669 let mut left = gin_query_candidates(idx, l)?;
5670 let mut right = gin_query_candidates(idx, r)?;
5671 left.sort_by_key(locator_sort_key);
5672 right.sort_by_key(locator_sort_key);
5673 let mut out: Vec<spg_storage::RowLocator> = Vec::new();
5675 let (mut i, mut j) = (0usize, 0usize);
5676 while i < left.len() && j < right.len() {
5677 let lk = locator_sort_key(&left[i]);
5678 let rk = locator_sort_key(&right[j]);
5679 match lk.cmp(&rk) {
5680 core::cmp::Ordering::Less => i += 1,
5681 core::cmp::Ordering::Greater => j += 1,
5682 core::cmp::Ordering::Equal => {
5683 out.push(left[i]);
5684 i += 1;
5685 j += 1;
5686 }
5687 }
5688 }
5689 Some(out)
5690 }
5691 TsQueryAst::Or(l, r) => {
5692 let mut out = gin_query_candidates(idx, l)?;
5693 out.extend(gin_query_candidates(idx, r)?);
5694 out.sort_by_key(locator_sort_key);
5695 out.dedup_by_key(|l| locator_sort_key(l));
5696 Some(out)
5697 }
5698 TsQueryAst::Not(_) | TsQueryAst::Phrase { .. } => None,
5703 }
5704}
5705
5706fn locator_sort_key(l: &spg_storage::RowLocator) -> (u8, u64, u64) {
5711 match *l {
5712 spg_storage::RowLocator::Hot(i) => (0, i as u64, 0),
5713 spg_storage::RowLocator::Cold {
5714 segment_id,
5715 page_offset,
5716 } => (1, u64::from(segment_id), u64::from(page_offset)),
5717 }
5718}
5719
5720fn try_pk_predicate(
5732 where_expr: &Expr,
5733 schema_cols: &[ColumnSchema],
5734 table_alias: &str,
5735) -> Option<(usize, IndexKey)> {
5736 let Expr::Binary {
5737 lhs,
5738 op: BinOp::Eq,
5739 rhs,
5740 } = where_expr
5741 else {
5742 return None;
5743 };
5744 let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
5745 .or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
5746 let key = IndexKey::from_value(&value)?;
5747 Some((col_pos, key))
5748}
5749
5750fn resolve_col_literal_pair(
5751 col_side: &Expr,
5752 lit_side: &Expr,
5753 schema_cols: &[ColumnSchema],
5754 table_alias: &str,
5755) -> Option<(usize, Value)> {
5756 let Expr::Column(c) = col_side else {
5757 return None;
5758 };
5759 if let Some(q) = &c.qualifier
5760 && q != table_alias
5761 {
5762 return None;
5763 }
5764 let pos = schema_cols.iter().position(|s| s.name == c.name)?;
5765 let Expr::Literal(l) = lit_side else {
5766 return None;
5767 };
5768 let v = match l {
5769 Literal::Integer(n) => {
5770 if let Ok(small) = i32::try_from(*n) {
5771 Value::Int(small)
5772 } else {
5773 Value::BigInt(*n)
5774 }
5775 }
5776 Literal::Float(x) => Value::Float(*x),
5777 Literal::String(s) => Value::Text(s.clone()),
5778 Literal::Bool(b) => Value::Bool(*b),
5779 Literal::Null => Value::Null,
5780 Literal::Vector(_) | Literal::Interval { .. } => return None,
5783 };
5784 Some((pos, v))
5785}
5786
5787fn resolve_projection_column<'a>(
5792 c: &ColumnName,
5793 schema_cols: &'a [ColumnSchema],
5794 table_alias: &str,
5795) -> Result<&'a ColumnSchema, EngineError> {
5796 if let Some(q) = &c.qualifier {
5797 let composite = alloc::format!("{q}.{name}", name = c.name);
5798 if let Some(s) = schema_cols.iter().find(|s| s.name == composite) {
5799 return Ok(s);
5800 }
5801 if q == table_alias
5804 && let Some(s) = schema_cols.iter().find(|s| s.name == c.name)
5805 {
5806 return Ok(s);
5807 }
5808 let prefix = alloc::format!("{q}.");
5812 let qualifier_known =
5813 q == table_alias || schema_cols.iter().any(|s| s.name.starts_with(&prefix));
5814 if !qualifier_known {
5815 return Err(EngineError::Eval(EvalError::UnknownQualifier {
5816 qualifier: q.clone(),
5817 }));
5818 }
5819 return Err(EngineError::Eval(EvalError::ColumnNotFound {
5820 name: c.name.clone(),
5821 }));
5822 }
5823 if let Some(s) = schema_cols.iter().find(|s| s.name == c.name) {
5824 return Ok(s);
5825 }
5826 let suffix = alloc::format!(".{name}", name = c.name);
5827 let mut matches = schema_cols.iter().filter(|s| s.name.ends_with(&suffix));
5828 let first = matches.next();
5829 let extra = matches.next();
5830 match (first, extra) {
5831 (Some(s), None) => Ok(s),
5832 (Some(_), Some(_)) => Err(EngineError::Eval(EvalError::TypeMismatch {
5833 detail: alloc::format!("ambiguous column reference: {}", c.name),
5834 })),
5835 _ => Err(EngineError::Eval(EvalError::ColumnNotFound {
5836 name: c.name.clone(),
5837 })),
5838 }
5839}
5840
5841fn build_projection(
5842 items: &[SelectItem],
5843 schema_cols: &[ColumnSchema],
5844 table_alias: &str,
5845) -> Result<Vec<ProjectedItem>, EngineError> {
5846 let mut out = Vec::new();
5847 for item in items {
5848 match item {
5849 SelectItem::Wildcard => {
5850 for col in schema_cols {
5851 out.push(ProjectedItem {
5852 expr: Expr::Column(ColumnName {
5853 qualifier: None,
5854 name: col.name.clone(),
5855 }),
5856 output_name: col.name.clone(),
5857 ty: col.ty,
5858 nullable: col.nullable,
5859 });
5860 }
5861 }
5862 SelectItem::Expr { expr, alias } => {
5863 if let Expr::Column(c) = expr {
5868 let sch = resolve_projection_column(c, schema_cols, table_alias)?;
5869 let output_name = alias.clone().unwrap_or_else(|| c.name.clone());
5870 out.push(ProjectedItem {
5871 expr: expr.clone(),
5872 output_name,
5873 ty: sch.ty,
5874 nullable: sch.nullable,
5875 });
5876 } else {
5877 let output_name = alias.clone().unwrap_or_else(|| expr.to_string());
5878 out.push(ProjectedItem {
5879 expr: expr.clone(),
5880 output_name,
5881 ty: DataType::Text,
5882 nullable: true,
5883 });
5884 }
5885 }
5886 }
5887 }
5888 Ok(out)
5889}
5890
5891fn numeric_from_integer(
5895 n: i128,
5896 precision: u8,
5897 scale: u8,
5898 col_name: &str,
5899) -> Result<Value, EngineError> {
5900 let factor = pow10_i128(scale);
5901 let scaled = n.checked_mul(factor).ok_or_else(|| {
5902 EngineError::Unsupported(alloc::format!(
5903 "integer overflow scaling value for column `{col_name}` to scale {scale}"
5904 ))
5905 })?;
5906 check_precision(scaled, precision, col_name)?;
5907 Ok(Value::Numeric { scaled, scale })
5908}
5909
5910#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
5913fn numeric_from_float(
5914 x: f64,
5915 precision: u8,
5916 scale: u8,
5917 col_name: &str,
5918) -> Result<Value, EngineError> {
5919 if !x.is_finite() {
5920 return Err(EngineError::Unsupported(alloc::format!(
5921 "cannot store non-finite float in NUMERIC column `{col_name}`"
5922 )));
5923 }
5924 let mut factor = 1.0_f64;
5925 for _ in 0..scale {
5926 factor *= 10.0;
5927 }
5928 let shifted = x * factor;
5933 let biased = if shifted >= 0.0 {
5934 shifted + 0.5
5935 } else {
5936 shifted - 0.5
5937 };
5938 if !(-1e38..=1e38).contains(&biased) {
5941 return Err(EngineError::Unsupported(alloc::format!(
5942 "value {x} overflows NUMERIC range for column `{col_name}`"
5943 )));
5944 }
5945 let scaled = biased as i128;
5946 check_precision(scaled, precision, col_name)?;
5947 Ok(Value::Numeric { scaled, scale })
5948}
5949
5950fn numeric_rescale(
5953 scaled: i128,
5954 src_scale: u8,
5955 precision: u8,
5956 dst_scale: u8,
5957 col_name: &str,
5958) -> Result<Value, EngineError> {
5959 let new_scaled = if dst_scale >= src_scale {
5960 let bump = pow10_i128(dst_scale - src_scale);
5961 scaled.checked_mul(bump).ok_or_else(|| {
5962 EngineError::Unsupported(alloc::format!(
5963 "overflow rescaling NUMERIC for column `{col_name}`"
5964 ))
5965 })?
5966 } else {
5967 let drop = pow10_i128(src_scale - dst_scale);
5968 let half = drop / 2;
5969 if scaled >= 0 {
5970 (scaled + half) / drop
5971 } else {
5972 (scaled - half) / drop
5973 }
5974 };
5975 check_precision(new_scaled, precision, col_name)?;
5976 Ok(Value::Numeric {
5977 scaled: new_scaled,
5978 scale: dst_scale,
5979 })
5980}
5981
5982const fn numeric_truncate_to_integer(scaled: i128, scale: u8) -> i128 {
5985 if scale == 0 {
5986 return scaled;
5987 }
5988 let factor = pow10_i128_const(scale);
5989 scaled / factor
5990}
5991
5992fn check_precision(scaled: i128, precision: u8, col_name: &str) -> Result<(), EngineError> {
5996 if precision == 0 {
5997 return Ok(());
5998 }
5999 let limit = pow10_i128(precision);
6000 if scaled.unsigned_abs() >= limit.unsigned_abs() {
6001 return Err(EngineError::Unsupported(alloc::format!(
6002 "NUMERIC value exceeds precision {precision} for column `{col_name}`"
6003 )));
6004 }
6005 Ok(())
6006}
6007
6008const fn pow10_i128_const(p: u8) -> i128 {
6009 let mut acc: i128 = 1;
6010 let mut i = 0;
6011 while i < p {
6012 acc *= 10;
6013 i += 1;
6014 }
6015 acc
6016}
6017
6018fn pow10_i128(p: u8) -> i128 {
6019 pow10_i128_const(p)
6020}
6021
6022impl Engine {
6037 #[allow(
6048 clippy::too_many_lines,
6049 clippy::type_complexity,
6050 clippy::needless_range_loop
6051 )] fn exec_select_with_window(
6053 &self,
6054 stmt: &SelectStatement,
6055 cancel: CancelToken<'_>,
6056 ) -> Result<QueryResult, EngineError> {
6057 let from = stmt.from.as_ref().ok_or_else(|| {
6058 EngineError::Unsupported("window functions require a FROM clause".into())
6059 })?;
6060 if !from.joins.is_empty() {
6063 return Err(EngineError::Unsupported(
6064 "JOIN with window functions not yet supported".into(),
6065 ));
6066 }
6067 let primary = &from.primary;
6068 let table = self.active_catalog().get(&primary.name).ok_or_else(|| {
6069 StorageError::TableNotFound {
6070 name: primary.name.clone(),
6071 }
6072 })?;
6073 let alias = primary.alias.as_deref().unwrap_or(primary.name.as_str());
6074 let schema_cols = &table.schema().columns;
6075 let ctx = self.ev_ctx(schema_cols, Some(alias));
6076
6077 let mut filtered: Vec<&Row> = Vec::new();
6079 for (i, row) in table.rows().iter().enumerate() {
6080 if i.is_multiple_of(256) {
6081 cancel.check()?;
6082 }
6083 if let Some(w) = &stmt.where_ {
6084 let cond = eval::eval_expr(w, row, &ctx)?;
6085 if !matches!(cond, Value::Bool(true)) {
6086 continue;
6087 }
6088 }
6089 filtered.push(row);
6090 }
6091 let n_rows = filtered.len();
6092
6093 let mut window_nodes: Vec<Expr> = Vec::new();
6095 for item in &stmt.items {
6096 if let SelectItem::Expr { expr, .. } = item {
6097 collect_window_nodes(expr, &mut window_nodes);
6098 }
6099 }
6100
6101 let mut win_vals: Vec<Vec<Value>> = Vec::with_capacity(window_nodes.len());
6104 for wnode in &window_nodes {
6105 let Expr::WindowFunction {
6106 name,
6107 args,
6108 partition_by,
6109 order_by,
6110 frame,
6111 null_treatment,
6112 } = wnode
6113 else {
6114 unreachable!("collect_window_nodes pushes only WindowFunction");
6115 };
6116 let mut indexed: Vec<(Vec<Value>, Vec<(Value, bool)>, usize)> =
6118 Vec::with_capacity(n_rows);
6119 for (i, row) in filtered.iter().enumerate() {
6120 let pkey: Vec<Value> = partition_by
6121 .iter()
6122 .map(|p| eval::eval_expr(p, row, &ctx))
6123 .collect::<Result<_, _>>()?;
6124 let okey: Vec<(Value, bool)> = order_by
6125 .iter()
6126 .map(|(e, desc)| eval::eval_expr(e, row, &ctx).map(|v| (v, *desc)))
6127 .collect::<Result<_, _>>()?;
6128 indexed.push((pkey, okey, i));
6129 }
6130 indexed.sort_by(|a, b| {
6133 let p_cmp = partition_key_cmp(&a.0, &b.0);
6134 if p_cmp != core::cmp::Ordering::Equal {
6135 return p_cmp;
6136 }
6137 order_key_cmp(&a.1, &b.1)
6138 });
6139 let mut out_vals: Vec<Value> = alloc::vec![Value::Null; n_rows];
6141 let mut p_start = 0;
6142 while p_start < indexed.len() {
6143 let mut p_end = p_start + 1;
6144 while p_end < indexed.len()
6145 && partition_key_cmp(&indexed[p_start].0, &indexed[p_end].0)
6146 == core::cmp::Ordering::Equal
6147 {
6148 p_end += 1;
6149 }
6150 compute_window_partition(
6152 name,
6153 args,
6154 !order_by.is_empty(),
6155 frame.as_ref(),
6156 *null_treatment,
6157 &indexed[p_start..p_end],
6158 &filtered,
6159 &ctx,
6160 &mut out_vals,
6161 )?;
6162 p_start = p_end;
6163 }
6164 win_vals.push(out_vals);
6165 }
6166
6167 let mut ext_cols = schema_cols.clone();
6169 for i in 0..window_nodes.len() {
6170 ext_cols.push(ColumnSchema::new(
6171 alloc::format!("__win_{i}"),
6172 DataType::Text, true,
6174 ));
6175 }
6176 let mut ext_rows: Vec<Row> = Vec::with_capacity(n_rows);
6178 for i in 0..n_rows {
6179 let mut values = filtered[i].values.clone();
6180 for w in 0..window_nodes.len() {
6181 values.push(win_vals[w][i].clone());
6182 }
6183 ext_rows.push(Row::new(values));
6184 }
6185 let mut rewritten_items: Vec<SelectItem> = Vec::with_capacity(stmt.items.len());
6187 for item in &stmt.items {
6188 let new_item = match item {
6189 SelectItem::Wildcard => SelectItem::Wildcard,
6190 SelectItem::Expr { expr, alias } => {
6191 let mut e = expr.clone();
6192 rewrite_window_to_columns(&mut e, &window_nodes);
6193 SelectItem::Expr {
6194 expr: e,
6195 alias: alias.clone(),
6196 }
6197 }
6198 };
6199 rewritten_items.push(new_item);
6200 }
6201
6202 let ext_ctx = EvalContext::new(&ext_cols, Some(alias));
6204 let projection = build_projection(&rewritten_items, &ext_cols, alias)?;
6205 let mut tagged: Vec<(Vec<f64>, Row)> = Vec::with_capacity(n_rows);
6206 for (i, row) in ext_rows.iter().enumerate() {
6207 if i.is_multiple_of(256) {
6208 cancel.check()?;
6209 }
6210 let mut values = Vec::with_capacity(projection.len());
6211 for p in &projection {
6212 values.push(eval::eval_expr(&p.expr, row, &ext_ctx)?);
6213 }
6214 let order_keys = if stmt.order_by.is_empty() {
6215 Vec::new()
6216 } else {
6217 let mut keys = Vec::with_capacity(stmt.order_by.len());
6218 for o in &stmt.order_by {
6219 let mut e = o.expr.clone();
6220 rewrite_window_to_columns(&mut e, &window_nodes);
6221 let key = eval::eval_expr(&e, row, &ext_ctx)?;
6222 keys.push(value_to_order_key(&key)?);
6223 }
6224 keys
6225 };
6226 tagged.push((order_keys, Row::new(values)));
6227 }
6228 if !stmt.order_by.is_empty() {
6230 let descs: Vec<bool> = stmt.order_by.iter().map(|o| o.desc).collect();
6231 sort_by_keys(&mut tagged, &descs);
6232 }
6233 let mut out_rows: Vec<Row> = tagged.into_iter().map(|(_, r)| r).collect();
6234 apply_offset_and_limit(&mut out_rows, stmt.offset_literal(), stmt.limit_literal());
6235 let final_cols: Vec<ColumnSchema> = projection
6236 .into_iter()
6237 .map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
6238 .collect();
6239 Ok(QueryResult::Rows {
6240 columns: final_cols,
6241 rows: out_rows,
6242 })
6243 }
6244
6245 fn exec_with_ctes(
6252 &self,
6253 stmt: &SelectStatement,
6254 cancel: CancelToken<'_>,
6255 ) -> Result<QueryResult, EngineError> {
6256 cancel.check()?;
6257 let mut catalog = self.active_catalog().clone();
6258 for cte in &stmt.ctes {
6259 if catalog.get(&cte.name).is_some() {
6260 return Err(EngineError::Unsupported(alloc::format!(
6261 "CTE name {:?} shadows an existing table; rename the CTE",
6262 cte.name
6263 )));
6264 }
6265 let (columns, rows) = if cte.recursive {
6266 self.materialise_recursive_cte(cte, &catalog, cancel)?
6267 } else {
6268 let body_result = self.exec_select_cancel(&cte.body, cancel)?;
6269 let QueryResult::Rows { columns, rows } = body_result else {
6270 return Err(EngineError::Unsupported(alloc::format!(
6271 "CTE {:?} body did not return rows",
6272 cte.name
6273 )));
6274 };
6275 (columns, rows)
6276 };
6277 let inferred = infer_column_types(&columns, &rows);
6282 let mut columns = inferred;
6283 if !cte.column_overrides.is_empty() {
6285 if cte.column_overrides.len() != columns.len() {
6286 return Err(EngineError::Unsupported(alloc::format!(
6287 "CTE {:?} column list has {} names but body returns {} columns",
6288 cte.name,
6289 cte.column_overrides.len(),
6290 columns.len()
6291 )));
6292 }
6293 for (col, name) in columns.iter_mut().zip(cte.column_overrides.iter()) {
6294 col.name.clone_from(name);
6295 }
6296 }
6297 let schema = TableSchema::new(cte.name.clone(), columns);
6298 catalog.create_table(schema).map_err(EngineError::Storage)?;
6299 let table = catalog
6300 .get_mut(&cte.name)
6301 .expect("just-created CTE table must exist");
6302 for row in rows {
6303 table.insert(row).map_err(EngineError::Storage)?;
6304 }
6305 }
6306 let mut body = stmt.clone();
6309 body.ctes = Vec::new();
6310 let mut temp = Engine::restore(catalog);
6311 if let Some(c) = self.clock {
6312 temp = temp.with_clock(c);
6313 }
6314 if let Some(f) = self.salt_fn {
6315 temp = temp.with_salt_fn(f);
6316 }
6317 temp.exec_select_cancel(&body, cancel)
6318 }
6319
6320 #[allow(clippy::too_many_lines)]
6330 fn materialise_recursive_cte(
6331 &self,
6332 cte: &spg_sql::ast::Cte,
6333 base_catalog: &Catalog,
6334 cancel: CancelToken<'_>,
6335 ) -> Result<(Vec<ColumnSchema>, Vec<Row>), EngineError> {
6336 const MAX_TOTAL_ROWS: usize = 1_000_000;
6337 const MAX_ITERATIONS: usize = 100_000;
6338 cancel.check()?;
6339 if cte.body.unions.is_empty() {
6340 return Err(EngineError::Unsupported(alloc::format!(
6341 "WITH RECURSIVE {:?} body must be a UNION of an anchor and a recursive term",
6342 cte.name
6343 )));
6344 }
6345 let mut anchor = cte.body.clone();
6347 let union_terms = core::mem::take(&mut anchor.unions);
6348 anchor.ctes = Vec::new();
6349 if select_refers_to(&anchor, &cte.name) {
6351 return Err(EngineError::Unsupported(alloc::format!(
6352 "WITH RECURSIVE {:?}: the anchor must not reference the CTE itself",
6353 cte.name
6354 )));
6355 }
6356 let anchor_result = self.exec_select_cancel(&anchor, cancel)?;
6357 let QueryResult::Rows {
6358 columns: anchor_cols,
6359 rows: anchor_rows,
6360 } = anchor_result
6361 else {
6362 return Err(EngineError::Unsupported(alloc::format!(
6363 "WITH RECURSIVE {:?}: anchor did not return rows",
6364 cte.name
6365 )));
6366 };
6367 let mut columns = infer_column_types(&anchor_cols, &anchor_rows);
6371 if !cte.column_overrides.is_empty() {
6372 if cte.column_overrides.len() != columns.len() {
6373 return Err(EngineError::Unsupported(alloc::format!(
6374 "CTE {:?} column list has {} names but anchor returns {} columns",
6375 cte.name,
6376 cte.column_overrides.len(),
6377 columns.len()
6378 )));
6379 }
6380 for (col, name) in columns.iter_mut().zip(cte.column_overrides.iter()) {
6381 col.name.clone_from(name);
6382 }
6383 }
6384 let mut all_rows: Vec<Row> = anchor_rows.clone();
6385 let mut working_set: Vec<Row> = anchor_rows;
6386 let mut seen: alloc::collections::BTreeSet<Vec<u8>> = alloc::collections::BTreeSet::new();
6387 let all_union_all = union_terms.iter().all(|(k, _)| matches!(k, UnionKind::All));
6390 if !all_union_all {
6391 for r in &all_rows {
6392 seen.insert(encode_row_key(r));
6393 }
6394 }
6395 for iter in 0..MAX_ITERATIONS {
6396 cancel.check()?;
6397 if working_set.is_empty() {
6398 break;
6399 }
6400 let mut iter_catalog = base_catalog.clone();
6402 let schema = TableSchema::new(cte.name.clone(), columns.clone());
6403 iter_catalog
6404 .create_table(schema)
6405 .map_err(EngineError::Storage)?;
6406 {
6407 let table = iter_catalog.get_mut(&cte.name).expect("just-created");
6408 for row in &working_set {
6409 table.insert(row.clone()).map_err(EngineError::Storage)?;
6410 }
6411 }
6412 let mut iter_engine = Engine::restore(iter_catalog);
6413 if let Some(c) = self.clock {
6414 iter_engine = iter_engine.with_clock(c);
6415 }
6416 if let Some(f) = self.salt_fn {
6417 iter_engine = iter_engine.with_salt_fn(f);
6418 }
6419 let mut next_set: Vec<Row> = Vec::new();
6421 for (_, term) in &union_terms {
6422 let mut term = term.clone();
6423 term.ctes = Vec::new();
6424 let r = iter_engine.exec_select_cancel(&term, cancel)?;
6425 let QueryResult::Rows {
6426 columns: rc,
6427 rows: rs,
6428 } = r
6429 else {
6430 return Err(EngineError::Unsupported(alloc::format!(
6431 "WITH RECURSIVE {:?}: recursive term did not return rows",
6432 cte.name
6433 )));
6434 };
6435 if rc.len() != columns.len() {
6436 return Err(EngineError::Unsupported(alloc::format!(
6437 "WITH RECURSIVE {:?}: column count of recursive term ({}) does not match anchor ({})",
6438 cte.name,
6439 rc.len(),
6440 columns.len()
6441 )));
6442 }
6443 for row in rs {
6444 if !all_union_all {
6445 let key = encode_row_key(&row);
6446 if !seen.insert(key) {
6447 continue;
6448 }
6449 }
6450 next_set.push(row);
6451 }
6452 }
6453 if next_set.is_empty() {
6454 break;
6455 }
6456 all_rows.extend(next_set.iter().cloned());
6457 working_set = next_set;
6458 if all_rows.len() > MAX_TOTAL_ROWS {
6459 return Err(EngineError::Unsupported(alloc::format!(
6460 "WITH RECURSIVE {:?}: produced more than {MAX_TOTAL_ROWS} rows — likely runaway recursion",
6461 cte.name
6462 )));
6463 }
6464 if iter + 1 == MAX_ITERATIONS {
6465 return Err(EngineError::Unsupported(alloc::format!(
6466 "WITH RECURSIVE {:?}: exceeded {MAX_ITERATIONS} iterations",
6467 cte.name
6468 )));
6469 }
6470 }
6471 Ok((columns, all_rows))
6472 }
6473
6474 fn resolve_select_subqueries(
6475 &self,
6476 stmt: &mut SelectStatement,
6477 cancel: CancelToken<'_>,
6478 ) -> Result<(), EngineError> {
6479 for item in &mut stmt.items {
6480 if let SelectItem::Expr { expr, .. } = item {
6481 self.resolve_expr_subqueries(expr, cancel)?;
6482 }
6483 }
6484 if let Some(w) = &mut stmt.where_ {
6485 self.resolve_expr_subqueries(w, cancel)?;
6486 }
6487 if let Some(gs) = &mut stmt.group_by {
6488 for g in gs {
6489 self.resolve_expr_subqueries(g, cancel)?;
6490 }
6491 }
6492 if let Some(h) = &mut stmt.having {
6493 self.resolve_expr_subqueries(h, cancel)?;
6494 }
6495 for o in &mut stmt.order_by {
6496 self.resolve_expr_subqueries(&mut o.expr, cancel)?;
6497 }
6498 for (_, peer) in &mut stmt.unions {
6499 self.resolve_select_subqueries(peer, cancel)?;
6500 }
6501 Ok(())
6502 }
6503
6504 #[allow(clippy::only_used_in_recursion)] fn resolve_expr_subqueries(
6506 &self,
6507 e: &mut Expr,
6508 cancel: CancelToken<'_>,
6509 ) -> Result<(), EngineError> {
6510 if let Some(replacement) = self.subquery_replacement(e, cancel)? {
6512 *e = replacement;
6513 return Ok(());
6514 }
6515 match e {
6516 Expr::Binary { lhs, rhs, .. } => {
6517 self.resolve_expr_subqueries(lhs, cancel)?;
6518 self.resolve_expr_subqueries(rhs, cancel)?;
6519 }
6520 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6521 self.resolve_expr_subqueries(expr, cancel)?;
6522 }
6523 Expr::FunctionCall { args, .. } => {
6524 for a in args {
6525 self.resolve_expr_subqueries(a, cancel)?;
6526 }
6527 }
6528 Expr::Like { expr, pattern, .. } => {
6529 self.resolve_expr_subqueries(expr, cancel)?;
6530 self.resolve_expr_subqueries(pattern, cancel)?;
6531 }
6532 Expr::Extract { source, .. } => self.resolve_expr_subqueries(source, cancel)?,
6533 Expr::WindowFunction {
6536 args,
6537 partition_by,
6538 order_by,
6539 ..
6540 } => {
6541 for a in args {
6542 self.resolve_expr_subqueries(a, cancel)?;
6543 }
6544 for p in partition_by {
6545 self.resolve_expr_subqueries(p, cancel)?;
6546 }
6547 for (e, _) in order_by {
6548 self.resolve_expr_subqueries(e, cancel)?;
6549 }
6550 }
6551 Expr::ScalarSubquery(_)
6555 | Expr::Exists { .. }
6556 | Expr::InSubquery { .. }
6557 | Expr::Literal(_)
6558 | Expr::Placeholder(_)
6559 | Expr::Column(_) => {}
6560 Expr::Array(items) => {
6562 for elem in items {
6563 self.resolve_expr_subqueries(elem, cancel)?;
6564 }
6565 }
6566 Expr::ArraySubscript { target, index } => {
6567 self.resolve_expr_subqueries(target, cancel)?;
6568 self.resolve_expr_subqueries(index, cancel)?;
6569 }
6570 Expr::AnyAll { expr, array, .. } => {
6571 self.resolve_expr_subqueries(expr, cancel)?;
6572 self.resolve_expr_subqueries(array, cancel)?;
6573 }
6574 Expr::Case {
6575 operand,
6576 branches,
6577 else_branch,
6578 } => {
6579 if let Some(o) = operand {
6580 self.resolve_expr_subqueries(o, cancel)?;
6581 }
6582 for (w, t) in branches {
6583 self.resolve_expr_subqueries(w, cancel)?;
6584 self.resolve_expr_subqueries(t, cancel)?;
6585 }
6586 if let Some(e) = else_branch {
6587 self.resolve_expr_subqueries(e, cancel)?;
6588 }
6589 }
6590 }
6591 Ok(())
6592 }
6593
6594 fn eval_expr_with_correlated(
6602 &self,
6603 expr: &Expr,
6604 row: &Row,
6605 ctx: &EvalContext<'_>,
6606 cancel: CancelToken<'_>,
6607 memo: Option<&mut memoize::MemoizeCache>,
6608 ) -> Result<Value, EngineError> {
6609 if !expr_has_subquery(expr) {
6610 return eval::eval_expr(expr, row, ctx).map_err(EngineError::Eval);
6611 }
6612 let mut e = expr.clone();
6613 self.resolve_correlated_in_expr(&mut e, row, ctx, cancel, memo)?;
6614 eval::eval_expr(&e, row, ctx).map_err(EngineError::Eval)
6615 }
6616
6617 fn resolve_correlated_in_expr(
6618 &self,
6619 e: &mut Expr,
6620 row: &Row,
6621 ctx: &EvalContext<'_>,
6622 cancel: CancelToken<'_>,
6623 mut memo: Option<&mut memoize::MemoizeCache>,
6624 ) -> Result<(), EngineError> {
6625 match e {
6626 Expr::ScalarSubquery(inner) => {
6627 let cache_key = memo.as_ref().map(|_| memoize::CacheKey {
6632 subquery_repr: alloc::format!("{}", **inner),
6633 outer_values: row.values.clone(),
6634 });
6635 if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key.as_ref())
6636 && let Some(cached) = cache.get(k)
6637 {
6638 *e = value_to_literal_expr(cached)?;
6639 return Ok(());
6640 }
6641 let mut s = (**inner).clone();
6642 substitute_outer_columns(&mut s, row, ctx);
6643 let r = self.exec_select_cancel(&s, cancel)?;
6644 let QueryResult::Rows { rows, .. } = r else {
6645 return Err(EngineError::Unsupported(
6646 "scalar subquery: inner did not return rows".into(),
6647 ));
6648 };
6649 let value = match rows.as_slice() {
6650 [] => Value::Null,
6651 [r0] => r0.values.first().cloned().unwrap_or(Value::Null),
6652 _ => {
6653 return Err(EngineError::Unsupported(alloc::format!(
6654 "scalar subquery returned {} rows; expected 0 or 1",
6655 rows.len()
6656 )));
6657 }
6658 };
6659 if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key) {
6660 cache.insert(k, value.clone());
6661 }
6662 *e = value_to_literal_expr(value)?;
6663 }
6664 Expr::Exists { subquery, negated } => {
6665 let mut s = (**subquery).clone();
6666 substitute_outer_columns(&mut s, row, ctx);
6667 let r = self.exec_select_cancel(&s, cancel)?;
6668 let exists = matches!(r, QueryResult::Rows { rows, .. } if !rows.is_empty());
6669 let bit = if *negated { !exists } else { exists };
6670 *e = Expr::Literal(Literal::Bool(bit));
6671 }
6672 Expr::InSubquery {
6673 expr: lhs,
6674 subquery,
6675 negated,
6676 } => {
6677 self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
6678 let lhs_val = eval::eval_expr(lhs, row, ctx).map_err(EngineError::Eval)?;
6679 let mut s = (**subquery).clone();
6680 substitute_outer_columns(&mut s, row, ctx);
6681 let r = self.exec_select_cancel(&s, cancel)?;
6682 let QueryResult::Rows { columns, rows, .. } = r else {
6683 return Err(EngineError::Unsupported(
6684 "IN-subquery: inner did not return rows".into(),
6685 ));
6686 };
6687 if columns.len() != 1 {
6688 return Err(EngineError::Unsupported(alloc::format!(
6689 "IN-subquery must project exactly one column; got {}",
6690 columns.len()
6691 )));
6692 }
6693 let mut found = false;
6694 let mut any_null = false;
6695 for r0 in rows {
6696 let v = r0.values.into_iter().next().unwrap_or(Value::Null);
6697 if v.is_null() {
6698 any_null = true;
6699 continue;
6700 }
6701 if value_cmp(&v, &lhs_val) == core::cmp::Ordering::Equal {
6702 found = true;
6703 break;
6704 }
6705 }
6706 let bit = if found {
6707 !*negated
6708 } else if any_null {
6709 return Err(EngineError::Unsupported(
6710 "IN-subquery with NULL in result and no match: NULL semantics not yet implemented".into(),
6711 ));
6712 } else {
6713 *negated
6714 };
6715 *e = Expr::Literal(Literal::Bool(bit));
6716 }
6717 Expr::Binary { lhs, rhs, .. } => {
6718 self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
6719 self.resolve_correlated_in_expr(rhs, row, ctx, cancel, memo.as_deref_mut())?;
6720 }
6721 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6722 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
6723 }
6724 Expr::Like { expr, pattern, .. } => {
6725 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
6726 self.resolve_correlated_in_expr(pattern, row, ctx, cancel, memo.as_deref_mut())?;
6727 }
6728 Expr::FunctionCall { args, .. } => {
6729 for a in args {
6730 self.resolve_correlated_in_expr(a, row, ctx, cancel, memo.as_deref_mut())?;
6731 }
6732 }
6733 Expr::Extract { source, .. } => {
6734 self.resolve_correlated_in_expr(source, row, ctx, cancel, memo.as_deref_mut())?;
6735 }
6736 Expr::WindowFunction { .. }
6737 | Expr::Literal(_)
6738 | Expr::Placeholder(_)
6739 | Expr::Column(_) => {}
6740 Expr::Array(items) => {
6742 for elem in items {
6743 self.resolve_correlated_in_expr(elem, row, ctx, cancel, memo.as_deref_mut())?;
6744 }
6745 }
6746 Expr::ArraySubscript { target, index } => {
6747 self.resolve_correlated_in_expr(target, row, ctx, cancel, memo.as_deref_mut())?;
6748 self.resolve_correlated_in_expr(index, row, ctx, cancel, memo.as_deref_mut())?;
6749 }
6750 Expr::AnyAll { expr, array, .. } => {
6751 self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
6752 self.resolve_correlated_in_expr(array, row, ctx, cancel, memo.as_deref_mut())?;
6753 }
6754 Expr::Case {
6755 operand,
6756 branches,
6757 else_branch,
6758 } => {
6759 if let Some(o) = operand {
6760 self.resolve_correlated_in_expr(o, row, ctx, cancel, memo.as_deref_mut())?;
6761 }
6762 for (w, t) in branches {
6763 self.resolve_correlated_in_expr(w, row, ctx, cancel, memo.as_deref_mut())?;
6764 self.resolve_correlated_in_expr(t, row, ctx, cancel, memo.as_deref_mut())?;
6765 }
6766 if let Some(e) = else_branch {
6767 self.resolve_correlated_in_expr(e, row, ctx, cancel, memo.as_deref_mut())?;
6768 }
6769 }
6770 }
6771 Ok(())
6772 }
6773
6774 fn subquery_replacement(
6775 &self,
6776 e: &Expr,
6777 cancel: CancelToken<'_>,
6778 ) -> Result<Option<Expr>, EngineError> {
6779 match e {
6780 Expr::ScalarSubquery(inner) => {
6781 let mut s = (**inner).clone();
6782 self.resolve_select_subqueries(&mut s, cancel)?;
6785 let r = match self.exec_bare_select_cancel(&s, cancel) {
6786 Ok(r) => r,
6787 Err(e) if is_correlation_error(&e) => return Ok(None),
6788 Err(e) => return Err(e),
6789 };
6790 let QueryResult::Rows { rows, .. } = r else {
6791 return Err(EngineError::Unsupported(
6792 "scalar subquery: inner statement did not return rows".into(),
6793 ));
6794 };
6795 let value = match rows.as_slice() {
6796 [] => Value::Null,
6797 [row] => row.values.first().cloned().unwrap_or(Value::Null),
6798 _ => {
6799 return Err(EngineError::Unsupported(alloc::format!(
6800 "scalar subquery returned {} rows; expected 0 or 1",
6801 rows.len()
6802 )));
6803 }
6804 };
6805 Ok(Some(value_to_literal_expr(value)?))
6806 }
6807 Expr::Exists { subquery, negated } => {
6808 let mut s = (**subquery).clone();
6809 self.resolve_select_subqueries(&mut s, cancel)?;
6810 let r = match self.exec_bare_select_cancel(&s, cancel) {
6811 Ok(r) => r,
6812 Err(e) if is_correlation_error(&e) => return Ok(None),
6813 Err(e) => return Err(e),
6814 };
6815 let exists = match r {
6816 QueryResult::Rows { rows, .. } => !rows.is_empty(),
6817 QueryResult::CommandOk { .. } => false,
6818 };
6819 let bit = if *negated { !exists } else { exists };
6820 Ok(Some(Expr::Literal(Literal::Bool(bit))))
6821 }
6822 Expr::InSubquery {
6823 expr,
6824 subquery,
6825 negated,
6826 } => {
6827 let mut s = (**subquery).clone();
6828 self.resolve_select_subqueries(&mut s, cancel)?;
6829 let r = match self.exec_bare_select_cancel(&s, cancel) {
6830 Ok(r) => r,
6831 Err(e) if is_correlation_error(&e) => return Ok(None),
6832 Err(e) => return Err(e),
6833 };
6834 let QueryResult::Rows { columns, rows, .. } = r else {
6835 return Err(EngineError::Unsupported(
6836 "IN-subquery: inner statement did not return rows".into(),
6837 ));
6838 };
6839 if columns.len() != 1 {
6840 return Err(EngineError::Unsupported(alloc::format!(
6841 "IN-subquery must project exactly one column; got {}",
6842 columns.len()
6843 )));
6844 }
6845 let mut acc: Option<Expr> = None;
6848 for row in rows {
6849 let v = row.values.into_iter().next().unwrap_or(Value::Null);
6850 let lit = value_to_literal_expr(v)?;
6851 let cmp = Expr::Binary {
6852 lhs: expr.clone(),
6853 op: BinOp::Eq,
6854 rhs: Box::new(lit),
6855 };
6856 acc = Some(match acc {
6857 None => cmp,
6858 Some(prev) => Expr::Binary {
6859 lhs: Box::new(prev),
6860 op: BinOp::Or,
6861 rhs: Box::new(cmp),
6862 },
6863 });
6864 }
6865 let combined = acc.unwrap_or(Expr::Literal(Literal::Bool(false)));
6866 let final_expr = if *negated {
6867 Expr::Unary {
6868 op: UnOp::Not,
6869 expr: Box::new(combined),
6870 }
6871 } else {
6872 combined
6873 };
6874 Ok(Some(final_expr))
6875 }
6876 _ => Ok(None),
6877 }
6878 }
6879}
6880
6881fn select_refers_to(stmt: &SelectStatement, target: &str) -> bool {
6893 if let Some(from) = &stmt.from
6894 && from_refers_to(from, target)
6895 {
6896 return true;
6897 }
6898 for (_, peer) in &stmt.unions {
6899 if select_refers_to(peer, target) {
6900 return true;
6901 }
6902 }
6903 for item in &stmt.items {
6904 if let SelectItem::Expr { expr, .. } = item
6905 && expr_refers_to(expr, target)
6906 {
6907 return true;
6908 }
6909 }
6910 if let Some(w) = &stmt.where_
6911 && expr_refers_to(w, target)
6912 {
6913 return true;
6914 }
6915 false
6916}
6917
6918fn from_refers_to(from: &FromClause, target: &str) -> bool {
6919 if from.primary.name.eq_ignore_ascii_case(target) {
6920 return true;
6921 }
6922 from.joins
6923 .iter()
6924 .any(|j| j.table.name.eq_ignore_ascii_case(target))
6925}
6926
6927fn expr_refers_to(e: &Expr, target: &str) -> bool {
6928 match e {
6929 Expr::ScalarSubquery(s) => select_refers_to(s, target),
6930 Expr::Exists { subquery, .. } | Expr::InSubquery { subquery, .. } => {
6931 select_refers_to(subquery, target)
6932 }
6933 Expr::Binary { lhs, rhs, .. } => expr_refers_to(lhs, target) || expr_refers_to(rhs, target),
6934 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
6935 expr_refers_to(expr, target)
6936 }
6937 Expr::Like { expr, pattern, .. } => {
6938 expr_refers_to(expr, target) || expr_refers_to(pattern, target)
6939 }
6940 Expr::FunctionCall { args, .. } => args.iter().any(|a| expr_refers_to(a, target)),
6941 Expr::Extract { source, .. } => expr_refers_to(source, target),
6942 Expr::WindowFunction {
6943 args,
6944 partition_by,
6945 order_by,
6946 ..
6947 } => {
6948 args.iter().any(|a| expr_refers_to(a, target))
6949 || partition_by.iter().any(|p| expr_refers_to(p, target))
6950 || order_by.iter().any(|(o, _)| expr_refers_to(o, target))
6951 }
6952 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => false,
6953 Expr::Array(items) => items.iter().any(|e| expr_refers_to(e, target)),
6954 Expr::ArraySubscript { target: t, index } => {
6955 expr_refers_to(t, target) || expr_refers_to(index, target)
6956 }
6957 Expr::AnyAll { expr, array, .. } => {
6958 expr_refers_to(expr, target) || expr_refers_to(array, target)
6959 }
6960 Expr::Case {
6961 operand,
6962 branches,
6963 else_branch,
6964 } => {
6965 operand.as_deref().is_some_and(|o| expr_refers_to(o, target))
6966 || branches
6967 .iter()
6968 .any(|(w, t)| expr_refers_to(w, target) || expr_refers_to(t, target))
6969 || else_branch
6970 .as_deref()
6971 .is_some_and(|e| expr_refers_to(e, target))
6972 }
6973 }
6974}
6975
6976fn infer_column_types(columns: &[ColumnSchema], rows: &[Row]) -> Vec<ColumnSchema> {
6982 let mut out = columns.to_vec();
6983 for (col_idx, col) in out.iter_mut().enumerate() {
6984 if col.ty != DataType::Text {
6985 continue;
6986 }
6987 let mut inferred: Option<DataType> = None;
6988 let mut all_null = true;
6989 for row in rows {
6990 let Some(v) = row.values.get(col_idx) else {
6991 continue;
6992 };
6993 let ty = match v {
6994 Value::Null => continue,
6995 Value::SmallInt(_) => DataType::SmallInt,
6996 Value::Int(_) => DataType::Int,
6997 Value::BigInt(_) => DataType::BigInt,
6998 Value::Float(_) => DataType::Float,
6999 Value::Bool(_) => DataType::Bool,
7000 Value::Vector(_) => DataType::Vector {
7001 dim: 0,
7002 encoding: VecEncoding::F32,
7003 },
7004 _ => DataType::Text,
7005 };
7006 all_null = false;
7007 inferred = Some(match inferred {
7008 None => ty,
7009 Some(prev) if prev == ty => prev,
7010 Some(_) => DataType::Text,
7011 });
7012 }
7013 if let Some(t) = inferred {
7014 col.ty = t;
7015 col.nullable = true;
7016 } else if all_null {
7017 col.nullable = true;
7018 }
7019 }
7020 out
7021}
7022
7023#[allow(clippy::too_many_lines, clippy::format_push_string)]
7028fn build_index_suggestions(stmt: &SelectStatement, engine: &Engine) -> Vec<String> {
7045 use alloc::collections::BTreeSet;
7046 let mut seen: BTreeSet<(String, String)> = BTreeSet::new();
7047 let mut out: Vec<String> = Vec::new();
7048 let cat = engine.active_catalog();
7049 let Some(from) = &stmt.from else {
7053 return out;
7054 };
7055 let mut tables: Vec<String> = Vec::new();
7056 tables.push(from.primary.name.clone());
7057 for j in &from.joins {
7058 tables.push(j.table.name.clone());
7059 }
7060 let mut col_refs: Vec<spg_sql::ast::ColumnName> = Vec::new();
7063 if let Some(w) = &stmt.where_ {
7064 collect_column_refs(w, &mut col_refs);
7065 }
7066 for j in &from.joins {
7067 if let Some(on) = &j.on {
7068 collect_column_refs(on, &mut col_refs);
7069 }
7070 }
7071 for cn in &col_refs {
7072 let owner: Option<String> = if let Some(q) = &cn.qualifier {
7075 tables.iter().find(|t| t == &q).cloned()
7076 } else {
7077 tables.iter().find_map(|t| {
7078 cat.get(t).and_then(|tbl| {
7079 if tbl.schema().column_position(&cn.name).is_some() {
7080 Some(t.clone())
7081 } else {
7082 None
7083 }
7084 })
7085 })
7086 };
7087 let Some(owner) = owner else {
7088 continue;
7089 };
7090 let Some(tbl) = cat.get(&owner) else {
7091 continue;
7092 };
7093 let Some(col_pos) = tbl.schema().column_position(&cn.name) else {
7094 continue;
7095 };
7096 let already_indexed = tbl.indices().iter().any(|i| {
7099 matches!(i.kind, spg_storage::IndexKind::BTree(_))
7100 && i.column_position == col_pos
7101 && i.expression.is_none()
7102 && i.partial_predicate.is_none()
7103 });
7104 if already_indexed {
7105 continue;
7106 }
7107 if seen.insert((owner.clone(), cn.name.clone())) {
7108 out.push(alloc::format!(
7109 "SUGGEST: CREATE INDEX ix_{}_{} ON {} ({})",
7110 owner,
7111 cn.name,
7112 owner,
7113 cn.name
7114 ));
7115 }
7116 }
7117 out
7118}
7119
7120fn collect_column_refs(expr: &Expr, out: &mut Vec<spg_sql::ast::ColumnName>) {
7123 match expr {
7124 Expr::Column(cn) => out.push(cn.clone()),
7125 Expr::FunctionCall { args, .. } => {
7126 for a in args {
7127 collect_column_refs(a, out);
7128 }
7129 }
7130 Expr::Binary { lhs, rhs, .. } => {
7131 collect_column_refs(lhs, out);
7132 collect_column_refs(rhs, out);
7133 }
7134 Expr::Unary { expr: e, .. } => collect_column_refs(e, out),
7135 _ => {}
7136 }
7137}
7138
7139fn annotate_explain_lines(lines: &mut [String], total_rows: usize, engine: &Engine) {
7140 let catalog = engine.active_catalog();
7141 let cold_ids = catalog.cold_segment_ids_global();
7142 let any_cold = !cold_ids.is_empty();
7143 let cold_ids_repr = if any_cold {
7144 let mut s = alloc::string::String::from("[");
7145 for (i, id) in cold_ids.iter().enumerate() {
7146 if i > 0 {
7147 s.push(',');
7148 }
7149 s.push_str(&alloc::format!("{id}"));
7150 }
7151 s.push(']');
7152 s
7153 } else {
7154 alloc::string::String::new()
7155 };
7156 for (idx, line) in lines.iter_mut().enumerate() {
7157 let trimmed = line.trim_start();
7158 let is_top_level = idx == 0;
7159 if is_top_level {
7160 line.push_str(&alloc::format!(" (rows={total_rows})"));
7161 continue;
7162 }
7163 if let Some(rest) = trimmed.strip_prefix("From: ") {
7164 let (name, scan_kind) = match rest.split_once(" [") {
7165 Some((n, k)) => (n.trim(), k.trim_end_matches(']')),
7166 None => (rest.trim(), ""),
7167 };
7168 let bare = name.split_whitespace().next().unwrap_or(name);
7169 let hot = catalog.get(bare).map(|t| t.rows().len());
7170 let annot = match (hot, scan_kind) {
7175 (Some(h), "full scan") => {
7176 let mut s = alloc::format!(" (hot_rows={h}");
7177 if any_cold {
7178 s.push_str(&alloc::format!(
7179 ", cold_tier=present, cold_segments={cold_ids_repr}"
7180 ));
7181 }
7182 s.push(')');
7183 s
7184 }
7185 (Some(h), "index seek") => {
7186 let mut s = alloc::format!(" (hot_rows≤{h}");
7187 if any_cold {
7188 s.push_str(&alloc::format!(
7189 ", cold_tier=present, cold_segments={cold_ids_repr}"
7190 ));
7191 }
7192 s.push(')');
7193 s
7194 }
7195 _ => " (rows=—)".to_string(),
7196 };
7197 line.push_str(&annot);
7198 continue;
7199 }
7200 line.push_str(" (rows=—)");
7202 }
7203}
7204
7205fn explain_select(stmt: &SelectStatement, engine: &Engine, depth: usize, out: &mut Vec<String>) {
7206 let pad = " ".repeat(depth);
7207 let top = if !stmt.ctes.is_empty() {
7209 if stmt.ctes.iter().any(|c| c.recursive) {
7210 "CTEScan (WITH RECURSIVE)"
7211 } else {
7212 "CTEScan (WITH)"
7213 }
7214 } else if !stmt.unions.is_empty() {
7215 "UnionScan"
7216 } else if select_has_window(stmt) {
7217 "WindowAgg"
7218 } else if aggregate::uses_aggregate(stmt) {
7219 "Aggregate"
7220 } else if stmt.distinct {
7221 "Distinct"
7222 } else if stmt.from.is_some() {
7223 "TableScan"
7224 } else {
7225 "Result"
7226 };
7227 out.push(alloc::format!("{pad}{top}"));
7228 let child = " ".repeat(depth + 1);
7229 for cte in &stmt.ctes {
7231 let head = if cte.recursive {
7232 alloc::format!("{child}CTE (recursive): {}", cte.name)
7233 } else {
7234 alloc::format!("{child}CTE: {}", cte.name)
7235 };
7236 out.push(head);
7237 explain_select(&cte.body, engine, depth + 2, out);
7238 }
7239 if let Some(from) = &stmt.from {
7241 let mut tag = alloc::format!("{child}From: {}", from.primary.name);
7242 if let Some(alias) = &from.primary.alias {
7243 tag.push_str(&alloc::format!(" AS {alias}"));
7244 }
7245 if let Some(w) = &stmt.where_
7248 && let Some(table) = engine.active_catalog().get(&from.primary.name)
7249 {
7250 let alias = from.primary.alias.as_deref().unwrap_or(&from.primary.name);
7251 let cols = &table.schema().columns;
7252 if try_index_seek(w, cols, engine.active_catalog(), table, alias).is_some() {
7253 tag.push_str(" [index seek]");
7254 } else {
7255 tag.push_str(" [full scan]");
7256 }
7257 } else {
7258 tag.push_str(" [full scan]");
7259 }
7260 out.push(tag);
7261 for j in &from.joins {
7262 let kind = match j.kind {
7263 spg_sql::ast::JoinKind::Inner => "INNER JOIN",
7264 spg_sql::ast::JoinKind::Left => "LEFT JOIN",
7265 spg_sql::ast::JoinKind::Cross => "CROSS JOIN",
7266 };
7267 let mut s = alloc::format!("{child}{kind}: {}", j.table.name);
7268 if let Some(alias) = &j.table.alias {
7269 s.push_str(&alloc::format!(" AS {alias}"));
7270 }
7271 if j.on.is_some() {
7272 s.push_str(" (ON …)");
7273 }
7274 out.push(s);
7275 }
7276 }
7277 if let Some(w) = &stmt.where_ {
7279 let mut s = alloc::format!("{child}Filter: {w}");
7280 if expr_has_subquery(w) {
7281 s.push_str(" [subquery]");
7282 }
7283 out.push(s);
7284 }
7285 if let Some(gs) = &stmt.group_by {
7286 let mut parts = Vec::new();
7287 for g in gs {
7288 parts.push(alloc::format!("{g}"));
7289 }
7290 out.push(alloc::format!("{child}GroupBy: {}", parts.join(", ")));
7291 }
7292 if let Some(h) = &stmt.having {
7293 out.push(alloc::format!("{child}Having: {h}"));
7294 }
7295 for o in &stmt.order_by {
7296 let dir = if o.desc { "DESC" } else { "ASC" };
7297 out.push(alloc::format!("{child}OrderBy: {} {dir}", o.expr));
7298 }
7299 if let Some(lim) = stmt.limit {
7300 out.push(alloc::format!("{child}Limit: {lim}"));
7301 }
7302 if let Some(off) = stmt.offset {
7303 out.push(alloc::format!("{child}Offset: {off}"));
7304 }
7305 if stmt
7307 .items
7308 .iter()
7309 .any(|it| matches!(it, SelectItem::Wildcard))
7310 {
7311 out.push(alloc::format!("{child}Project: *"));
7312 } else {
7313 out.push(alloc::format!(
7314 "{child}Project: {} item(s)",
7315 stmt.items.len()
7316 ));
7317 }
7318 for (kind, peer) in &stmt.unions {
7320 let label = match kind {
7321 UnionKind::All => "UNION ALL",
7322 UnionKind::Distinct => "UNION",
7323 };
7324 out.push(alloc::format!("{child}{label}"));
7325 explain_select(peer, engine, depth + 2, out);
7326 }
7327}
7328
7329fn is_correlation_error(e: &EngineError) -> bool {
7334 matches!(
7335 e,
7336 EngineError::Eval(
7337 eval::EvalError::ColumnNotFound { .. } | eval::EvalError::UnknownQualifier { .. }
7338 )
7339 )
7340}
7341
7342fn substitute_outer_columns(stmt: &mut SelectStatement, row: &Row, ctx: &EvalContext<'_>) {
7350 let Some(outer_alias) = ctx.table_alias else {
7351 return;
7352 };
7353 substitute_in_select(stmt, row, ctx, outer_alias);
7354}
7355
7356fn substitute_in_select(
7357 stmt: &mut SelectStatement,
7358 row: &Row,
7359 ctx: &EvalContext<'_>,
7360 outer_alias: &str,
7361) {
7362 for item in &mut stmt.items {
7363 if let SelectItem::Expr { expr, .. } = item {
7364 substitute_in_expr(expr, row, ctx, outer_alias);
7365 }
7366 }
7367 if let Some(w) = &mut stmt.where_ {
7368 substitute_in_expr(w, row, ctx, outer_alias);
7369 }
7370 if let Some(gs) = &mut stmt.group_by {
7371 for g in gs {
7372 substitute_in_expr(g, row, ctx, outer_alias);
7373 }
7374 }
7375 if let Some(h) = &mut stmt.having {
7376 substitute_in_expr(h, row, ctx, outer_alias);
7377 }
7378 for o in &mut stmt.order_by {
7379 substitute_in_expr(&mut o.expr, row, ctx, outer_alias);
7380 }
7381 for (_, peer) in &mut stmt.unions {
7382 substitute_in_select(peer, row, ctx, outer_alias);
7383 }
7384}
7385
7386fn substitute_in_expr(e: &mut Expr, row: &Row, ctx: &EvalContext<'_>, outer_alias: &str) {
7387 if let Expr::Column(c) = e
7388 && let Some(qual) = &c.qualifier
7389 && qual.eq_ignore_ascii_case(outer_alias)
7390 {
7391 if let Some(idx) = ctx
7393 .columns
7394 .iter()
7395 .position(|sc| sc.name.eq_ignore_ascii_case(&c.name))
7396 {
7397 let v = row.values.get(idx).cloned().unwrap_or(Value::Null);
7398 if let Ok(lit) = value_to_literal_expr(v) {
7399 *e = lit;
7400 return;
7401 }
7402 }
7403 }
7404 match e {
7405 Expr::Binary { lhs, rhs, .. } => {
7406 substitute_in_expr(lhs, row, ctx, outer_alias);
7407 substitute_in_expr(rhs, row, ctx, outer_alias);
7408 }
7409 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7410 substitute_in_expr(expr, row, ctx, outer_alias);
7411 }
7412 Expr::Like { expr, pattern, .. } => {
7413 substitute_in_expr(expr, row, ctx, outer_alias);
7414 substitute_in_expr(pattern, row, ctx, outer_alias);
7415 }
7416 Expr::FunctionCall { args, .. } => {
7417 for a in args {
7418 substitute_in_expr(a, row, ctx, outer_alias);
7419 }
7420 }
7421 Expr::Extract { source, .. } => substitute_in_expr(source, row, ctx, outer_alias),
7422 Expr::WindowFunction {
7423 args,
7424 partition_by,
7425 order_by,
7426 ..
7427 } => {
7428 for a in args {
7429 substitute_in_expr(a, row, ctx, outer_alias);
7430 }
7431 for p in partition_by {
7432 substitute_in_expr(p, row, ctx, outer_alias);
7433 }
7434 for (o, _) in order_by {
7435 substitute_in_expr(o, row, ctx, outer_alias);
7436 }
7437 }
7438 Expr::ScalarSubquery(s) => substitute_in_select(s, row, ctx, outer_alias),
7439 Expr::Exists { subquery, .. } | Expr::InSubquery { subquery, .. } => {
7440 substitute_in_select(subquery, row, ctx, outer_alias);
7441 }
7442 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
7443 Expr::Array(items) => {
7444 for elem in items {
7445 substitute_in_expr(elem, row, ctx, outer_alias);
7446 }
7447 }
7448 Expr::ArraySubscript { target, index } => {
7449 substitute_in_expr(target, row, ctx, outer_alias);
7450 substitute_in_expr(index, row, ctx, outer_alias);
7451 }
7452 Expr::AnyAll { expr, array, .. } => {
7453 substitute_in_expr(expr, row, ctx, outer_alias);
7454 substitute_in_expr(array, row, ctx, outer_alias);
7455 }
7456 Expr::Case {
7457 operand,
7458 branches,
7459 else_branch,
7460 } => {
7461 if let Some(o) = operand {
7462 substitute_in_expr(o, row, ctx, outer_alias);
7463 }
7464 for (w, t) in branches {
7465 substitute_in_expr(w, row, ctx, outer_alias);
7466 substitute_in_expr(t, row, ctx, outer_alias);
7467 }
7468 if let Some(e) = else_branch {
7469 substitute_in_expr(e, row, ctx, outer_alias);
7470 }
7471 }
7472 }
7473}
7474
7475fn encode_row_key(row: &Row) -> Vec<u8> {
7479 let mut out = Vec::new();
7480 for v in &row.values {
7481 let s = alloc::format!("{v:?}|");
7482 out.extend_from_slice(s.as_bytes());
7483 }
7484 out
7485}
7486
7487fn select_has_window(stmt: &SelectStatement) -> bool {
7488 for item in &stmt.items {
7489 if let SelectItem::Expr { expr, .. } = item
7490 && expr_has_window(expr)
7491 {
7492 return true;
7493 }
7494 }
7495 false
7496}
7497
7498fn expr_has_window(e: &Expr) -> bool {
7499 match e {
7500 Expr::WindowFunction { .. } => true,
7501 Expr::Binary { lhs, rhs, .. } => expr_has_window(lhs) || expr_has_window(rhs),
7502 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7503 expr_has_window(expr)
7504 }
7505 Expr::FunctionCall { args, .. } => args.iter().any(expr_has_window),
7506 Expr::Like { expr, pattern, .. } => expr_has_window(expr) || expr_has_window(pattern),
7507 Expr::Extract { source, .. } => expr_has_window(source),
7508 Expr::ScalarSubquery(_)
7509 | Expr::Exists { .. }
7510 | Expr::InSubquery { .. }
7511 | Expr::Literal(_)
7512 | Expr::Placeholder(_)
7513 | Expr::Column(_) => false,
7514 Expr::Array(items) => items.iter().any(expr_has_window),
7515 Expr::ArraySubscript { target, index } => expr_has_window(target) || expr_has_window(index),
7516 Expr::AnyAll { expr, array, .. } => expr_has_window(expr) || expr_has_window(array),
7517 Expr::Case {
7518 operand,
7519 branches,
7520 else_branch,
7521 } => {
7522 operand.as_deref().is_some_and(expr_has_window)
7523 || branches
7524 .iter()
7525 .any(|(w, t)| expr_has_window(w) || expr_has_window(t))
7526 || else_branch.as_deref().is_some_and(expr_has_window)
7527 }
7528 }
7529}
7530
7531fn collect_window_nodes(e: &Expr, out: &mut Vec<Expr>) {
7532 if let Expr::WindowFunction { .. } = e {
7533 if !out.iter().any(|x| x == e) {
7538 out.push(e.clone());
7539 }
7540 return;
7541 }
7542 match e {
7543 Expr::WindowFunction { .. } => unreachable!(),
7545 Expr::Binary { lhs, rhs, .. } => {
7546 collect_window_nodes(lhs, out);
7547 collect_window_nodes(rhs, out);
7548 }
7549 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7550 collect_window_nodes(expr, out);
7551 }
7552 Expr::FunctionCall { args, .. } => {
7553 for a in args {
7554 collect_window_nodes(a, out);
7555 }
7556 }
7557 Expr::Like { expr, pattern, .. } => {
7558 collect_window_nodes(expr, out);
7559 collect_window_nodes(pattern, out);
7560 }
7561 Expr::Extract { source, .. } => collect_window_nodes(source, out),
7562 _ => {}
7563 }
7564}
7565
7566fn rewrite_window_to_columns(e: &mut Expr, window_nodes: &[Expr]) {
7567 if let Expr::WindowFunction { .. } = e
7568 && let Some(idx) = window_nodes.iter().position(|w| w == e)
7569 {
7570 *e = Expr::Column(spg_sql::ast::ColumnName {
7571 qualifier: None,
7572 name: alloc::format!("__win_{idx}"),
7573 });
7574 return;
7575 }
7576 match e {
7577 Expr::Binary { lhs, rhs, .. } => {
7578 rewrite_window_to_columns(lhs, window_nodes);
7579 rewrite_window_to_columns(rhs, window_nodes);
7580 }
7581 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
7582 rewrite_window_to_columns(expr, window_nodes);
7583 }
7584 Expr::FunctionCall { args, .. } => {
7585 for a in args {
7586 rewrite_window_to_columns(a, window_nodes);
7587 }
7588 }
7589 Expr::Like { expr, pattern, .. } => {
7590 rewrite_window_to_columns(expr, window_nodes);
7591 rewrite_window_to_columns(pattern, window_nodes);
7592 }
7593 Expr::Extract { source, .. } => rewrite_window_to_columns(source, window_nodes),
7594 _ => {}
7595 }
7596}
7597
7598fn partition_key_cmp(a: &[Value], b: &[Value]) -> core::cmp::Ordering {
7602 for (x, y) in a.iter().zip(b.iter()) {
7603 let c = value_cmp(x, y);
7604 if c != core::cmp::Ordering::Equal {
7605 return c;
7606 }
7607 }
7608 a.len().cmp(&b.len())
7609}
7610
7611fn order_key_cmp(a: &[(Value, bool)], b: &[(Value, bool)]) -> core::cmp::Ordering {
7612 for ((va, desc), (vb, _)) in a.iter().zip(b.iter()) {
7613 let c = value_cmp(va, vb);
7614 let c = if *desc { c.reverse() } else { c };
7615 if c != core::cmp::Ordering::Equal {
7616 return c;
7617 }
7618 }
7619 a.len().cmp(&b.len())
7620}
7621
7622#[allow(clippy::match_same_arms)] fn value_cmp(a: &Value, b: &Value) -> core::cmp::Ordering {
7624 use core::cmp::Ordering;
7625 match (a, b) {
7626 (Value::Null, Value::Null) => Ordering::Equal,
7627 (Value::Null, _) => Ordering::Less,
7628 (_, Value::Null) => Ordering::Greater,
7629 (Value::Int(x), Value::Int(y)) => x.cmp(y),
7630 (Value::BigInt(x), Value::BigInt(y)) => x.cmp(y),
7631 (Value::SmallInt(x), Value::SmallInt(y)) => x.cmp(y),
7632 (Value::Text(x), Value::Text(y)) => x.cmp(y),
7633 (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
7634 (Value::Float(x), Value::Float(y)) => x.partial_cmp(y).unwrap_or(Ordering::Equal),
7635 (Value::Date(x), Value::Date(y)) => x.cmp(y),
7636 (Value::Timestamp(x), Value::Timestamp(y)) => x.cmp(y),
7637 _ => alloc::format!("{a:?}").cmp(&alloc::format!("{b:?}")),
7640 }
7641}
7642
7643#[allow(
7649 clippy::too_many_arguments,
7650 clippy::cast_possible_truncation,
7651 clippy::cast_possible_wrap,
7652 clippy::cast_precision_loss,
7653 clippy::cast_sign_loss,
7654 clippy::doc_markdown,
7655 clippy::too_many_lines,
7656 clippy::type_complexity,
7657 clippy::match_same_arms
7658)]
7659fn compute_window_partition(
7660 name: &str,
7661 args: &[Expr],
7662 ordered: bool,
7663 frame: Option<&WindowFrame>,
7664 null_treatment: spg_sql::ast::NullTreatment,
7665 slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)],
7666 filtered_rows: &[&Row],
7667 ctx: &EvalContext<'_>,
7668 out_vals: &mut [Value],
7669) -> Result<(), EngineError> {
7670 let ignore_nulls = matches!(null_treatment, spg_sql::ast::NullTreatment::Ignore);
7671 let lower = name.to_ascii_lowercase();
7672 match lower.as_str() {
7673 "row_number" => {
7674 for (rank, (_, _, idx)) in slice.iter().enumerate() {
7675 out_vals[*idx] = Value::BigInt((rank + 1) as i64);
7676 }
7677 Ok(())
7678 }
7679 "rank" => {
7680 let mut prev_key: Option<&[(Value, bool)]> = None;
7681 let mut current_rank: i64 = 1;
7682 for (i, (_, okey, idx)) in slice.iter().enumerate() {
7683 if let Some(p) = prev_key
7684 && order_key_cmp(p, okey) != core::cmp::Ordering::Equal
7685 {
7686 current_rank = (i + 1) as i64;
7687 }
7688 if prev_key.is_none() {
7689 current_rank = 1;
7690 }
7691 out_vals[*idx] = Value::BigInt(current_rank);
7692 prev_key = Some(okey.as_slice());
7693 }
7694 Ok(())
7695 }
7696 "dense_rank" => {
7697 let mut prev_key: Option<&[(Value, bool)]> = None;
7698 let mut current_rank: i64 = 0;
7699 for (_, okey, idx) in slice {
7700 if prev_key.is_none_or(|p| order_key_cmp(p, okey) != core::cmp::Ordering::Equal) {
7701 current_rank += 1;
7702 }
7703 out_vals[*idx] = Value::BigInt(current_rank);
7704 prev_key = Some(okey.as_slice());
7705 }
7706 Ok(())
7707 }
7708 "sum" | "avg" | "min" | "max" | "count" | "count_star" => {
7709 let arg_values: Vec<Value> = if lower == "count_star" || args.is_empty() {
7712 slice.iter().map(|_| Value::Null).collect()
7713 } else {
7714 slice
7715 .iter()
7716 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
7717 .collect::<Result<_, _>>()
7718 .map_err(EngineError::Eval)?
7719 };
7720 let eff = effective_frame(frame, ordered)?;
7724 #[allow(clippy::needless_range_loop)]
7725 for i in 0..slice.len() {
7726 let (lo, hi) = frame_bounds_for_row(&eff, i, slice);
7727 let mut sum: f64 = 0.0;
7728 let mut count: i64 = 0;
7729 let mut min_v: Option<f64> = None;
7730 let mut max_v: Option<f64> = None;
7731 let mut row_count: i64 = 0;
7732 if lo <= hi {
7733 for j in lo..=hi {
7734 let v = &arg_values[j];
7735 match lower.as_str() {
7736 "count_star" => row_count += 1,
7737 "count" => {
7738 if !v.is_null() {
7739 count += 1;
7740 }
7741 }
7742 _ => {
7743 if let Some(x) = value_to_f64(v) {
7744 sum += x;
7745 count += 1;
7746 min_v = Some(min_v.map_or(x, |m| m.min(x)));
7747 max_v = Some(max_v.map_or(x, |m| m.max(x)));
7748 }
7749 }
7750 }
7751 }
7752 }
7753 let value = match lower.as_str() {
7754 "count_star" => Value::BigInt(row_count),
7755 "count" => Value::BigInt(count),
7756 "sum" => Value::Float(sum),
7757 "avg" => {
7758 if count == 0 {
7759 Value::Null
7760 } else {
7761 Value::Float(sum / count as f64)
7762 }
7763 }
7764 "min" => min_v.map_or(Value::Null, Value::Float),
7765 "max" => max_v.map_or(Value::Null, Value::Float),
7766 _ => unreachable!(),
7767 };
7768 let (_, _, idx) = &slice[i];
7769 out_vals[*idx] = value;
7770 }
7771 Ok(())
7772 }
7773 "lag" | "lead" => {
7774 if args.is_empty() {
7777 return Err(EngineError::Unsupported(alloc::format!(
7778 "{lower}() requires at least one argument"
7779 )));
7780 }
7781 let offset: i64 = if args.len() >= 2 {
7782 let v = eval::eval_expr(&args[1], filtered_rows[slice[0].2], ctx)
7783 .map_err(EngineError::Eval)?;
7784 match v {
7785 Value::SmallInt(n) => i64::from(n),
7786 Value::Int(n) => i64::from(n),
7787 Value::BigInt(n) => n,
7788 _ => {
7789 return Err(EngineError::Unsupported(alloc::format!(
7790 "{lower}() offset must be integer"
7791 )));
7792 }
7793 }
7794 } else {
7795 1
7796 };
7797 let default: Value = if args.len() >= 3 {
7798 eval::eval_expr(&args[2], filtered_rows[slice[0].2], ctx)
7799 .map_err(EngineError::Eval)?
7800 } else {
7801 Value::Null
7802 };
7803 let values: Vec<Value> = slice
7804 .iter()
7805 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
7806 .collect::<Result<_, _>>()
7807 .map_err(EngineError::Eval)?;
7808 let n = slice.len();
7809 for (i, (_, _, idx)) in slice.iter().enumerate() {
7810 let signed_offset = if lower == "lag" { -offset } else { offset };
7811 let v = if ignore_nulls {
7812 let step: i64 = if signed_offset >= 0 { 1 } else { -1 };
7816 let needed: i64 = signed_offset.abs();
7817 if needed == 0 {
7818 values[i].clone()
7819 } else {
7820 let mut j: i64 = i as i64;
7821 let mut hits: i64 = 0;
7822 let mut found: Option<Value> = None;
7823 loop {
7824 j += step;
7825 if j < 0 || j >= n as i64 {
7826 break;
7827 }
7828 #[allow(clippy::cast_sign_loss)]
7829 let v = &values[j as usize];
7830 if !v.is_null() {
7831 hits += 1;
7832 if hits == needed {
7833 found = Some(v.clone());
7834 break;
7835 }
7836 }
7837 }
7838 found.unwrap_or_else(|| default.clone())
7839 }
7840 } else {
7841 let target_signed = i64::try_from(i).unwrap_or(i64::MAX) + signed_offset;
7842 if target_signed < 0 || target_signed >= i64::try_from(n).unwrap_or(i64::MAX) {
7843 default.clone()
7844 } else {
7845 #[allow(clippy::cast_sign_loss)]
7846 {
7847 values[target_signed as usize].clone()
7848 }
7849 }
7850 };
7851 out_vals[*idx] = v;
7852 }
7853 Ok(())
7854 }
7855 "first_value" | "last_value" | "nth_value" => {
7856 if args.is_empty() {
7857 return Err(EngineError::Unsupported(alloc::format!(
7858 "{lower}() requires at least one argument"
7859 )));
7860 }
7861 let values: Vec<Value> = slice
7862 .iter()
7863 .map(|(_, _, idx)| eval::eval_expr(&args[0], filtered_rows[*idx], ctx))
7864 .collect::<Result<_, _>>()
7865 .map_err(EngineError::Eval)?;
7866 let nth: usize = if lower == "nth_value" {
7867 if args.len() < 2 {
7868 return Err(EngineError::Unsupported(
7869 "nth_value() requires (expr, n)".into(),
7870 ));
7871 }
7872 let v = eval::eval_expr(&args[1], filtered_rows[slice[0].2], ctx)
7873 .map_err(EngineError::Eval)?;
7874 let raw = match v {
7875 Value::SmallInt(n) => i64::from(n),
7876 Value::Int(n) => i64::from(n),
7877 Value::BigInt(n) => n,
7878 _ => {
7879 return Err(EngineError::Unsupported(
7880 "nth_value() n must be integer".into(),
7881 ));
7882 }
7883 };
7884 if raw < 1 {
7885 return Err(EngineError::Unsupported(
7886 "nth_value() n must be >= 1".into(),
7887 ));
7888 }
7889 #[allow(clippy::cast_sign_loss)]
7890 {
7891 raw as usize
7892 }
7893 } else {
7894 0
7895 };
7896 let eff = effective_frame(frame, ordered)?;
7897 for i in 0..slice.len() {
7898 let (lo, hi) = frame_bounds_for_row(&eff, i, slice);
7899 let (_, _, idx) = &slice[i];
7900 let v = if lo > hi {
7901 Value::Null
7902 } else if ignore_nulls && matches!(lower.as_str(), "first_value" | "last_value") {
7903 if lower == "first_value" {
7906 (lo..=hi)
7907 .find_map(|j| {
7908 let v = &values[j];
7909 (!v.is_null()).then(|| v.clone())
7910 })
7911 .unwrap_or(Value::Null)
7912 } else {
7913 (lo..=hi)
7914 .rev()
7915 .find_map(|j| {
7916 let v = &values[j];
7917 (!v.is_null()).then(|| v.clone())
7918 })
7919 .unwrap_or(Value::Null)
7920 }
7921 } else {
7922 match lower.as_str() {
7923 "first_value" => values[lo].clone(),
7924 "last_value" => values[hi].clone(),
7925 "nth_value" => {
7926 let pos = lo + nth - 1;
7927 if pos > hi {
7928 Value::Null
7929 } else {
7930 values[pos].clone()
7931 }
7932 }
7933 _ => unreachable!(),
7934 }
7935 };
7936 out_vals[*idx] = v;
7937 }
7938 Ok(())
7939 }
7940 "ntile" => {
7941 if args.is_empty() {
7942 return Err(EngineError::Unsupported(
7943 "ntile(n) requires an integer argument".into(),
7944 ));
7945 }
7946 let v = eval::eval_expr(&args[0], filtered_rows[slice[0].2], ctx)
7947 .map_err(EngineError::Eval)?;
7948 let bucket_count: i64 = match v {
7949 Value::SmallInt(n) => i64::from(n),
7950 Value::Int(n) => i64::from(n),
7951 Value::BigInt(n) => n,
7952 _ => {
7953 return Err(EngineError::Unsupported(
7954 "ntile() argument must be integer".into(),
7955 ));
7956 }
7957 };
7958 if bucket_count < 1 {
7959 return Err(EngineError::Unsupported(
7960 "ntile() argument must be >= 1".into(),
7961 ));
7962 }
7963 #[allow(clippy::cast_sign_loss)]
7964 let buckets = bucket_count as usize;
7965 let n = slice.len();
7966 let base = n / buckets;
7969 let extras = n % buckets;
7970 let mut bucket: usize = 1;
7971 let mut remaining_in_bucket = if extras > 0 { base + 1 } else { base };
7972 let mut buckets_with_extra_remaining = extras;
7973 for (_, _, idx) in slice {
7974 if remaining_in_bucket == 0 {
7975 bucket += 1;
7976 buckets_with_extra_remaining = buckets_with_extra_remaining.saturating_sub(1);
7977 remaining_in_bucket = if buckets_with_extra_remaining > 0 {
7978 base + 1
7979 } else {
7980 base
7981 };
7982 if remaining_in_bucket == 0 {
7985 remaining_in_bucket = 1;
7986 }
7987 }
7988 out_vals[*idx] = Value::BigInt(i64::try_from(bucket).unwrap_or(i64::MAX));
7989 remaining_in_bucket -= 1;
7990 }
7991 Ok(())
7992 }
7993 "percent_rank" => {
7994 let n = slice.len();
7997 let mut prev_key: Option<&[(Value, bool)]> = None;
7998 let mut current_rank: i64 = 1;
7999 for (i, (_, okey, idx)) in slice.iter().enumerate() {
8000 if let Some(p) = prev_key
8001 && order_key_cmp(p, okey) != core::cmp::Ordering::Equal
8002 {
8003 current_rank = i64::try_from(i + 1).unwrap_or(i64::MAX);
8004 }
8005 if prev_key.is_none() {
8006 current_rank = 1;
8007 }
8008 #[allow(clippy::cast_precision_loss)]
8009 let pr = if n <= 1 {
8010 0.0
8011 } else {
8012 (current_rank - 1) as f64 / (n - 1) as f64
8013 };
8014 out_vals[*idx] = Value::Float(pr);
8015 prev_key = Some(okey.as_slice());
8016 }
8017 Ok(())
8018 }
8019 "cume_dist" => {
8020 let n = slice.len();
8022 for i in 0..slice.len() {
8024 let peer_end = peer_group_end(slice, i);
8025 #[allow(clippy::cast_precision_loss)]
8026 let cd = (peer_end + 1) as f64 / n as f64;
8027 let (_, _, idx) = &slice[i];
8028 out_vals[*idx] = Value::Float(cd);
8029 }
8030 Ok(())
8031 }
8032 other => Err(EngineError::Unsupported(alloc::format!(
8033 "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)"
8034 ))),
8035 }
8036}
8037
8038fn effective_frame(
8045 frame: Option<&WindowFrame>,
8046 ordered: bool,
8047) -> Result<(FrameKind, FrameBound, FrameBound), EngineError> {
8048 match frame {
8049 None => {
8050 if ordered {
8051 Ok((
8052 FrameKind::Range,
8053 FrameBound::UnboundedPreceding,
8054 FrameBound::CurrentRow,
8055 ))
8056 } else {
8057 Ok((
8058 FrameKind::Rows,
8059 FrameBound::UnboundedPreceding,
8060 FrameBound::UnboundedFollowing,
8061 ))
8062 }
8063 }
8064 Some(fr) => {
8065 let end = fr.end.clone().unwrap_or(FrameBound::CurrentRow);
8066 if matches!(fr.start, FrameBound::UnboundedFollowing)
8068 || matches!(end, FrameBound::UnboundedPreceding)
8069 {
8070 return Err(EngineError::Unsupported(alloc::format!(
8071 "invalid frame: start={:?} end={:?}",
8072 fr.start,
8073 end
8074 )));
8075 }
8076 if fr.kind == FrameKind::Range
8081 && (matches!(
8082 fr.start,
8083 FrameBound::OffsetPreceding(_) | FrameBound::OffsetFollowing(_)
8084 ) || matches!(
8085 end,
8086 FrameBound::OffsetPreceding(_) | FrameBound::OffsetFollowing(_)
8087 ))
8088 {
8089 return Err(EngineError::Unsupported(
8090 "RANGE with explicit offset bounds is not supported (v4.20: only UNBOUNDED / CURRENT ROW for RANGE)".into(),
8091 ));
8092 }
8093 Ok((fr.kind, fr.start.clone(), end))
8094 }
8095 }
8096}
8097
8098#[allow(clippy::type_complexity)]
8102fn frame_bounds_for_row(
8103 eff: &(FrameKind, FrameBound, FrameBound),
8104 i: usize,
8105 slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)],
8106) -> (usize, usize) {
8107 let (kind, start, end) = eff;
8108 let n = slice.len();
8109 let last = n.saturating_sub(1);
8110 let (mut lo, mut hi) = match kind {
8111 FrameKind::Rows => {
8112 let lo = match start {
8113 FrameBound::UnboundedPreceding => 0,
8114 FrameBound::OffsetPreceding(k) => {
8115 let k = usize::try_from(*k).unwrap_or(usize::MAX);
8116 i.saturating_sub(k)
8117 }
8118 FrameBound::CurrentRow => i,
8119 FrameBound::OffsetFollowing(k) => {
8120 let k = usize::try_from(*k).unwrap_or(usize::MAX);
8121 i.saturating_add(k).min(last)
8122 }
8123 FrameBound::UnboundedFollowing => last,
8124 };
8125 let hi = match end {
8126 FrameBound::UnboundedPreceding => 0,
8127 FrameBound::OffsetPreceding(k) => {
8128 let k = usize::try_from(*k).unwrap_or(usize::MAX);
8129 i.saturating_sub(k)
8130 }
8131 FrameBound::CurrentRow => i,
8132 FrameBound::OffsetFollowing(k) => {
8133 let k = usize::try_from(*k).unwrap_or(usize::MAX);
8134 i.saturating_add(k).min(last)
8135 }
8136 FrameBound::UnboundedFollowing => last,
8137 };
8138 (lo, hi)
8139 }
8140 FrameKind::Range => {
8141 let lo = match start {
8147 FrameBound::UnboundedPreceding => 0,
8148 FrameBound::CurrentRow => peer_group_start(slice, i),
8149 FrameBound::UnboundedFollowing => last,
8150 _ => unreachable!("offset bounds rejected for RANGE"),
8151 };
8152 let hi = match end {
8153 FrameBound::UnboundedPreceding => 0,
8154 FrameBound::CurrentRow => peer_group_end(slice, i),
8155 FrameBound::UnboundedFollowing => last,
8156 _ => unreachable!("offset bounds rejected for RANGE"),
8157 };
8158 (lo, hi)
8159 }
8160 };
8161 if hi >= n {
8162 hi = last;
8163 }
8164 if lo >= n {
8165 lo = last;
8166 }
8167 (lo, hi)
8168}
8169
8170#[allow(clippy::type_complexity)]
8174fn peer_group_start(slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)], i: usize) -> usize {
8175 let key = &slice[i].1;
8176 let mut j = i;
8177 while j > 0 && order_key_cmp(&slice[j - 1].1, key) == core::cmp::Ordering::Equal {
8178 j -= 1;
8179 }
8180 j
8181}
8182
8183#[allow(clippy::type_complexity)]
8186fn peer_group_end(slice: &[(Vec<Value>, Vec<(Value, bool)>, usize)], i: usize) -> usize {
8187 let key = &slice[i].1;
8188 let mut j = i;
8189 while j + 1 < slice.len() && order_key_cmp(&slice[j + 1].1, key) == core::cmp::Ordering::Equal {
8190 j += 1;
8191 }
8192 j
8193}
8194
8195fn value_to_f64(v: &Value) -> Option<f64> {
8196 match v {
8197 Value::SmallInt(n) => Some(f64::from(*n)),
8198 Value::Int(n) => Some(f64::from(*n)),
8199 #[allow(clippy::cast_precision_loss)]
8200 Value::BigInt(n) => Some(*n as f64),
8201 Value::Float(x) => Some(*x),
8202 _ => None,
8203 }
8204}
8205
8206fn expr_tree_has_subquery(stmt: &SelectStatement) -> bool {
8210 let mut any = false;
8211 for item in &stmt.items {
8212 if let SelectItem::Expr { expr, .. } = item {
8213 any = any || expr_has_subquery(expr);
8214 }
8215 }
8216 if let Some(w) = &stmt.where_ {
8217 any = any || expr_has_subquery(w);
8218 }
8219 if let Some(h) = &stmt.having {
8220 any = any || expr_has_subquery(h);
8221 }
8222 for o in &stmt.order_by {
8223 any = any || expr_has_subquery(&o.expr);
8224 }
8225 for (_, peer) in &stmt.unions {
8226 any = any || expr_tree_has_subquery(peer);
8227 }
8228 any
8229}
8230
8231fn expr_has_subquery(e: &Expr) -> bool {
8232 match e {
8233 Expr::ScalarSubquery(_) | Expr::Exists { .. } | Expr::InSubquery { .. } => true,
8234 Expr::Binary { lhs, rhs, .. } => expr_has_subquery(lhs) || expr_has_subquery(rhs),
8235 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
8236 expr_has_subquery(expr)
8237 }
8238 Expr::FunctionCall { args, .. } => args.iter().any(expr_has_subquery),
8239 Expr::Like { expr, pattern, .. } => expr_has_subquery(expr) || expr_has_subquery(pattern),
8240 Expr::Extract { source, .. } => expr_has_subquery(source),
8241 Expr::WindowFunction {
8242 args,
8243 partition_by,
8244 order_by,
8245 ..
8246 } => {
8247 args.iter().any(expr_has_subquery)
8248 || partition_by.iter().any(expr_has_subquery)
8249 || order_by.iter().any(|(e, _)| expr_has_subquery(e))
8250 }
8251 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => false,
8252 Expr::Array(items) => items.iter().any(expr_has_subquery),
8253 Expr::ArraySubscript { target, index } => {
8254 expr_has_subquery(target) || expr_has_subquery(index)
8255 }
8256 Expr::AnyAll { expr, array, .. } => expr_has_subquery(expr) || expr_has_subquery(array),
8257 Expr::Case {
8258 operand,
8259 branches,
8260 else_branch,
8261 } => {
8262 operand.as_deref().is_some_and(expr_has_subquery)
8263 || branches
8264 .iter()
8265 .any(|(w, t)| expr_has_subquery(w) || expr_has_subquery(t))
8266 || else_branch.as_deref().is_some_and(expr_has_subquery)
8267 }
8268 }
8269}
8270
8271fn value_to_literal_expr(v: Value) -> Result<Expr, EngineError> {
8278 let lit = match v {
8279 Value::Null => Literal::Null,
8280 Value::SmallInt(n) => Literal::Integer(i64::from(n)),
8281 Value::Int(n) => Literal::Integer(i64::from(n)),
8282 Value::BigInt(n) => Literal::Integer(n),
8283 Value::Float(x) => Literal::Float(x),
8284 Value::Text(s) | Value::Json(s) => Literal::String(s),
8285 Value::Bool(b) => Literal::Bool(b),
8286 other => {
8287 return Err(EngineError::Unsupported(alloc::format!(
8288 "subquery result type {:?} not yet materialisable; cast to text or integer in the inner SELECT",
8289 other.data_type()
8290 )));
8291 }
8292 };
8293 Ok(Expr::Literal(lit))
8294}
8295
8296fn value_to_literal_expr_permissive(v: Value) -> Result<Expr, EngineError> {
8302 let lit = match v {
8303 Value::Null => Literal::Null,
8304 Value::SmallInt(n) => Literal::Integer(i64::from(n)),
8305 Value::Int(n) => Literal::Integer(i64::from(n)),
8306 Value::BigInt(n) => Literal::Integer(n),
8307 Value::Float(x) => Literal::Float(x),
8308 Value::Text(s) | Value::Json(s) => Literal::String(s),
8309 Value::Bool(b) => Literal::Bool(b),
8310 Value::Vector(xs) => Literal::Vector(xs),
8311 Value::Date(days) => {
8315 let micros = (i64::from(days)) * 86_400_000_000;
8316 Literal::String(format_timestamp_micros_as_date(micros))
8317 }
8318 Value::Timestamp(us) => Literal::String(format_timestamp_micros(us)),
8319 Value::Numeric { scaled, scale } => {
8320 Literal::String(format_numeric(scaled, scale))
8321 }
8322 other => {
8323 return Err(EngineError::Unsupported(alloc::format!(
8324 "INSERT … SELECT cannot materialise value of type {:?}; \
8325 add an explicit CAST in the inner SELECT",
8326 other.data_type()
8327 )));
8328 }
8329 };
8330 Ok(Expr::Literal(lit))
8331}
8332
8333fn format_timestamp_micros(us: i64) -> String {
8334 let days = us.div_euclid(86_400_000_000);
8336 let intra_day = us.rem_euclid(86_400_000_000);
8337 let date = format_timestamp_micros_as_date(days * 86_400_000_000);
8338 let secs = intra_day / 1_000_000;
8339 let us_rem = intra_day % 1_000_000;
8340 let h = (secs / 3600) % 24;
8341 let m = (secs / 60) % 60;
8342 let s = secs % 60;
8343 if us_rem == 0 {
8344 alloc::format!("{date} {h:02}:{m:02}:{s:02}")
8345 } else {
8346 alloc::format!("{date} {h:02}:{m:02}:{s:02}.{us_rem:06}")
8347 }
8348}
8349
8350fn format_timestamp_micros_as_date(us: i64) -> String {
8351 let days = us.div_euclid(86_400_000_000);
8354 let jdn = days + 2_440_588;
8356 let (y, mo, d) = jdn_to_ymd(jdn);
8357 alloc::format!("{y:04}-{mo:02}-{d:02}")
8358}
8359
8360fn jdn_to_ymd(jdn: i64) -> (i64, u32, u32) {
8361 let l = jdn + 68569;
8363 let n = (4 * l) / 146_097;
8364 let l = l - (146_097 * n + 3) / 4;
8365 let i = (4000 * (l + 1)) / 1_461_001;
8366 let l = l - (1461 * i) / 4 + 31;
8367 let j = (80 * l) / 2447;
8368 let day = (l - (2447 * j) / 80) as u32;
8369 let l = j / 11;
8370 let month = (j + 2 - 12 * l) as u32;
8371 let year = 100 * (n - 49) + i + l;
8372 (year, month, day)
8373}
8374
8375fn format_numeric(scaled: i128, scale: u8) -> String {
8376 if scale == 0 {
8377 return alloc::format!("{scaled}");
8378 }
8379 let abs = scaled.unsigned_abs();
8380 let divisor = 10u128.pow(u32::from(scale));
8381 let whole = abs / divisor;
8382 let frac = abs % divisor;
8383 let sign = if scaled < 0 { "-" } else { "" };
8384 alloc::format!(
8385 "{sign}{whole}.{frac:0width$}",
8386 width = usize::from(scale)
8387 )
8388}
8389
8390fn substitute_placeholders(stmt: &mut Statement, params: &[Value]) -> Result<(), EngineError> {
8401 match stmt {
8402 Statement::Select(s) => substitute_select(s, params)?,
8403 Statement::Insert(ins) => {
8404 for row in &mut ins.rows {
8405 for e in row {
8406 substitute_expr(e, params)?;
8407 }
8408 }
8409 }
8410 Statement::Update(u) => {
8411 for (_, e) in &mut u.assignments {
8412 substitute_expr(e, params)?;
8413 }
8414 if let Some(w) = &mut u.where_ {
8415 substitute_expr(w, params)?;
8416 }
8417 }
8418 Statement::Delete(d) => {
8419 if let Some(w) = &mut d.where_ {
8420 substitute_expr(w, params)?;
8421 }
8422 }
8423 Statement::Explain(e) => substitute_select(&mut e.inner, params)?,
8424 _ => {}
8427 }
8428 Ok(())
8429}
8430
8431fn substitute_select(s: &mut SelectStatement, params: &[Value]) -> Result<(), EngineError> {
8432 for item in &mut s.items {
8433 if let SelectItem::Expr { expr, .. } = item {
8434 substitute_expr(expr, params)?;
8435 }
8436 }
8437 if let Some(w) = &mut s.where_ {
8438 substitute_expr(w, params)?;
8439 }
8440 if let Some(gs) = &mut s.group_by {
8441 for g in gs {
8442 substitute_expr(g, params)?;
8443 }
8444 }
8445 if let Some(h) = &mut s.having {
8446 substitute_expr(h, params)?;
8447 }
8448 for o in &mut s.order_by {
8449 substitute_expr(&mut o.expr, params)?;
8450 }
8451 for (_, peer) in &mut s.unions {
8452 substitute_select(peer, params)?;
8453 }
8454 if let Some(le) = s.limit {
8459 s.limit = Some(resolve_limit_placeholder(le, params)?);
8460 }
8461 if let Some(le) = s.offset {
8462 s.offset = Some(resolve_limit_placeholder(le, params)?);
8463 }
8464 Ok(())
8465}
8466
8467fn resolve_limit_placeholder(
8468 le: spg_sql::ast::LimitExpr,
8469 params: &[Value],
8470) -> Result<spg_sql::ast::LimitExpr, EngineError> {
8471 use spg_sql::ast::LimitExpr;
8472 match le {
8473 LimitExpr::Literal(_) => Ok(le),
8474 LimitExpr::Placeholder(n) => {
8475 let idx = usize::from(n).saturating_sub(1);
8476 let v = params.get(idx).ok_or_else(|| {
8477 EngineError::Eval(EvalError::PlaceholderOutOfRange {
8478 n,
8479 bound: u16::try_from(params.len()).unwrap_or(u16::MAX),
8480 })
8481 })?;
8482 let int = match v {
8483 Value::SmallInt(x) => Some(i64::from(*x)),
8484 Value::Int(x) => Some(i64::from(*x)),
8485 Value::BigInt(x) => Some(*x),
8486 _ => None,
8487 }
8488 .ok_or_else(|| {
8489 EngineError::Unsupported(alloc::format!(
8490 "LIMIT/OFFSET ${n} bound to non-integer {v:?}"
8491 ))
8492 })?;
8493 if int < 0 {
8494 return Err(EngineError::Unsupported(alloc::format!(
8495 "LIMIT/OFFSET ${n} bound to negative value {int}"
8496 )));
8497 }
8498 let bounded = u32::try_from(int).map_err(|_| {
8499 EngineError::Unsupported(alloc::format!(
8500 "LIMIT/OFFSET ${n} value {int} exceeds u32 range"
8501 ))
8502 })?;
8503 Ok(LimitExpr::Literal(bounded))
8504 }
8505 }
8506}
8507
8508fn substitute_expr(e: &mut Expr, params: &[Value]) -> Result<(), EngineError> {
8509 if let Expr::Placeholder(n) = e {
8510 let idx = usize::from(*n).saturating_sub(1);
8511 let v = params.get(idx).ok_or_else(|| {
8512 EngineError::Eval(EvalError::PlaceholderOutOfRange {
8513 n: *n,
8514 bound: u16::try_from(params.len()).unwrap_or(u16::MAX),
8515 })
8516 })?;
8517 *e = Expr::Literal(value_to_literal(v.clone()));
8518 return Ok(());
8519 }
8520 match e {
8521 Expr::Binary { lhs, rhs, .. } => {
8522 substitute_expr(lhs, params)?;
8523 substitute_expr(rhs, params)?;
8524 }
8525 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
8526 substitute_expr(expr, params)?;
8527 }
8528 Expr::FunctionCall { args, .. } => {
8529 for a in args {
8530 substitute_expr(a, params)?;
8531 }
8532 }
8533 Expr::Like { expr, pattern, .. } => {
8534 substitute_expr(expr, params)?;
8535 substitute_expr(pattern, params)?;
8536 }
8537 Expr::Extract { source, .. } => substitute_expr(source, params)?,
8538 Expr::ScalarSubquery(s) => substitute_select(s, params)?,
8539 Expr::Exists { subquery, .. } => substitute_select(subquery, params)?,
8540 Expr::InSubquery { expr, subquery, .. } => {
8541 substitute_expr(expr, params)?;
8542 substitute_select(subquery, params)?;
8543 }
8544 Expr::WindowFunction {
8545 args,
8546 partition_by,
8547 order_by,
8548 ..
8549 } => {
8550 for a in args {
8551 substitute_expr(a, params)?;
8552 }
8553 for p in partition_by {
8554 substitute_expr(p, params)?;
8555 }
8556 for (e, _) in order_by {
8557 substitute_expr(e, params)?;
8558 }
8559 }
8560 Expr::Literal(_) | Expr::Column(_) => {}
8561 Expr::Placeholder(_) => unreachable!("Placeholder handled at top of fn"),
8563 Expr::Array(items) => {
8564 for elem in items {
8565 substitute_expr(elem, params)?;
8566 }
8567 }
8568 Expr::ArraySubscript { target, index } => {
8569 substitute_expr(target, params)?;
8570 substitute_expr(index, params)?;
8571 }
8572 Expr::AnyAll { expr, array, .. } => {
8573 substitute_expr(expr, params)?;
8574 substitute_expr(array, params)?;
8575 }
8576 Expr::Case {
8577 operand,
8578 branches,
8579 else_branch,
8580 } => {
8581 if let Some(o) = operand {
8582 substitute_expr(o, params)?;
8583 }
8584 for (w, t) in branches {
8585 substitute_expr(w, params)?;
8586 substitute_expr(t, params)?;
8587 }
8588 if let Some(e) = else_branch {
8589 substitute_expr(e, params)?;
8590 }
8591 }
8592 }
8593 Ok(())
8594}
8595
8596fn sort_values_for_histogram(a: &Value, b: &Value) -> core::cmp::Ordering {
8614 use core::cmp::Ordering;
8615 match (a, b) {
8616 (Value::SmallInt(a), Value::SmallInt(b)) => a.cmp(b),
8617 (Value::Int(a), Value::Int(b)) => a.cmp(b),
8618 (Value::BigInt(a), Value::BigInt(b)) => a.cmp(b),
8619 (Value::SmallInt(a), Value::Int(b)) => i32::from(*a).cmp(b),
8620 (Value::Int(a), Value::SmallInt(b)) => a.cmp(&i32::from(*b)),
8621 (Value::Int(a), Value::BigInt(b)) => i64::from(*a).cmp(b),
8622 (Value::BigInt(a), Value::Int(b)) => a.cmp(&i64::from(*b)),
8623 (Value::SmallInt(a), Value::BigInt(b)) => i64::from(*a).cmp(b),
8624 (Value::BigInt(a), Value::SmallInt(b)) => a.cmp(&i64::from(*b)),
8625 (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(Ordering::Equal),
8626 (Value::Text(a), Value::Text(b)) | (Value::Json(a), Value::Json(b)) => a.cmp(b),
8627 (Value::Bool(a), Value::Bool(b)) => a.cmp(b),
8628 (Value::Date(a), Value::Date(b)) => a.cmp(b),
8629 (Value::Timestamp(a), Value::Timestamp(b)) => a.cmp(b),
8630 (Value::SmallInt(n), Value::Float(x)) => {
8632 (f64::from(*n)).partial_cmp(x).unwrap_or(Ordering::Equal)
8633 }
8634 (Value::Float(x), Value::SmallInt(n)) => {
8635 x.partial_cmp(&f64::from(*n)).unwrap_or(Ordering::Equal)
8636 }
8637 (Value::Int(n), Value::Float(x)) => {
8638 (f64::from(*n)).partial_cmp(x).unwrap_or(Ordering::Equal)
8639 }
8640 (Value::Float(x), Value::Int(n)) => {
8641 x.partial_cmp(&f64::from(*n)).unwrap_or(Ordering::Equal)
8642 }
8643 (Value::BigInt(n), Value::Float(x)) => {
8644 #[allow(clippy::cast_precision_loss)]
8645 let nf = *n as f64;
8646 nf.partial_cmp(x).unwrap_or(Ordering::Equal)
8647 }
8648 (Value::Float(x), Value::BigInt(n)) => {
8649 #[allow(clippy::cast_precision_loss)]
8650 let nf = *n as f64;
8651 x.partial_cmp(&nf).unwrap_or(Ordering::Equal)
8652 }
8653 _ => canonical_value_repr(a).cmp(&canonical_value_repr(b)),
8656 }
8657}
8658
8659fn render_histogram_bounds(bounds: &[alloc::string::String]) -> alloc::string::String {
8666 let mut out = alloc::string::String::with_capacity(bounds.len() * 8 + 2);
8667 out.push('[');
8668 for (i, b) in bounds.iter().enumerate() {
8669 if i > 0 {
8670 out.push_str(", ");
8671 }
8672 let needs_quote = b.contains([',', '[', ']', '"']) || b.is_empty();
8673 if needs_quote {
8674 out.push('"');
8675 for ch in b.chars() {
8676 if ch == '"' || ch == '\\' {
8677 out.push('\\');
8678 }
8679 out.push(ch);
8680 }
8681 out.push('"');
8682 } else {
8683 out.push_str(b);
8684 }
8685 }
8686 out.push(']');
8687 out
8688}
8689
8690pub(crate) fn canonical_value_repr(v: &Value) -> alloc::string::String {
8700 match v {
8701 Value::Null => "NULL".to_string(),
8702 Value::SmallInt(n) => alloc::format!("{n}"),
8703 Value::Int(n) => alloc::format!("{n}"),
8704 Value::BigInt(n) => alloc::format!("{n}"),
8705 Value::Float(x) => alloc::format!("{x:?}"),
8706 Value::Text(s) | Value::Json(s) => s.clone(),
8707 Value::Bool(b) => if *b { "t" } else { "f" }.to_string(),
8708 Value::Date(d) => eval::format_date(*d),
8709 Value::Timestamp(t) => eval::format_timestamp(*t),
8710 Value::Interval { months, micros } => eval::format_interval(*months, *micros),
8711 Value::Numeric { scaled, scale } => eval::format_numeric(*scaled, *scale),
8712 Value::Vector(_) | Value::Sq8Vector(_) | Value::HalfVector(_) => {
8713 alloc::format!("{v:?}")
8717 }
8718 _ => alloc::format!("{v:?}"),
8722 }
8723}
8724
8725const fn is_internal_table_name(_name: &str) -> bool {
8732 false
8733}
8734
8735fn value_to_literal(v: Value) -> Literal {
8736 match v {
8737 Value::Null => Literal::Null,
8738 Value::SmallInt(n) => Literal::Integer(i64::from(n)),
8739 Value::Int(n) => Literal::Integer(i64::from(n)),
8740 Value::BigInt(n) => Literal::Integer(n),
8741 Value::Float(x) => Literal::Float(x),
8742 Value::Text(s) | Value::Json(s) => Literal::String(s),
8743 Value::Bool(b) => Literal::Bool(b),
8744 Value::Vector(v) => Literal::Vector(v),
8745 Value::Numeric { scaled, scale } => Literal::String(eval::format_numeric(scaled, scale)),
8746 Value::Date(d) => Literal::String(eval::format_date(d)),
8747 Value::Timestamp(t) => Literal::String(eval::format_timestamp(t)),
8748 Value::Interval { months, micros } => Literal::Interval {
8749 months,
8750 micros,
8751 text: eval::format_interval(months, micros),
8752 },
8753 Value::Sq8Vector(q) => Literal::Vector(spg_storage::quantize::dequantize(&q)),
8756 Value::HalfVector(h) => Literal::Vector(h.to_f32_vec()),
8757 v => Literal::String(alloc::format!("{v:?}")),
8761 }
8762}
8763
8764fn rewrite_clock_calls(stmt: &mut Statement, now_micros: Option<i64>) {
8765 let Some(now) = now_micros else {
8766 return;
8767 };
8768 match stmt {
8769 Statement::Select(s) => rewrite_select_clock(s, now),
8770 Statement::Insert(ins) => {
8771 for row in &mut ins.rows {
8772 for e in row {
8773 rewrite_expr_clock(e, now);
8774 }
8775 }
8776 }
8777 _ => {}
8778 }
8779}
8780
8781fn rewrite_select_clock(s: &mut SelectStatement, now: i64) {
8782 for item in &mut s.items {
8783 if let SelectItem::Expr { expr, .. } = item {
8784 rewrite_expr_clock(expr, now);
8785 }
8786 }
8787 if let Some(w) = &mut s.where_ {
8788 rewrite_expr_clock(w, now);
8789 }
8790 if let Some(gs) = &mut s.group_by {
8791 for g in gs {
8792 rewrite_expr_clock(g, now);
8793 }
8794 }
8795 if let Some(h) = &mut s.having {
8796 rewrite_expr_clock(h, now);
8797 }
8798 for o in &mut s.order_by {
8799 rewrite_expr_clock(&mut o.expr, now);
8800 }
8801 for (_, peer) in &mut s.unions {
8802 rewrite_select_clock(peer, now);
8803 }
8804}
8805
8806fn rewrite_expr_clock(e: &mut Expr, now: i64) {
8814 if let Some(replacement) = clock_replacement_for(e, now) {
8818 *e = replacement;
8819 return;
8820 }
8821 match e {
8822 Expr::Binary { lhs, rhs, .. } => {
8823 rewrite_expr_clock(lhs, now);
8824 rewrite_expr_clock(rhs, now);
8825 }
8826 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
8827 rewrite_expr_clock(expr, now);
8828 }
8829 Expr::FunctionCall { args, .. } => {
8830 for a in args {
8831 rewrite_expr_clock(a, now);
8832 }
8833 }
8834 Expr::Like { expr, pattern, .. } => {
8835 rewrite_expr_clock(expr, now);
8836 rewrite_expr_clock(pattern, now);
8837 }
8838 Expr::Extract { source, .. } => rewrite_expr_clock(source, now),
8839 Expr::ScalarSubquery(s) => rewrite_select_clock(s, now),
8843 Expr::Exists { subquery, .. } => rewrite_select_clock(subquery, now),
8844 Expr::InSubquery { expr, subquery, .. } => {
8845 rewrite_expr_clock(expr, now);
8846 rewrite_select_clock(subquery, now);
8847 }
8848 Expr::WindowFunction {
8851 args,
8852 partition_by,
8853 order_by,
8854 ..
8855 } => {
8856 for a in args {
8857 rewrite_expr_clock(a, now);
8858 }
8859 for p in partition_by {
8860 rewrite_expr_clock(p, now);
8861 }
8862 for (e, _) in order_by {
8863 rewrite_expr_clock(e, now);
8864 }
8865 }
8866 Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
8867 Expr::Array(items) => {
8868 for elem in items {
8869 rewrite_expr_clock(elem, now);
8870 }
8871 }
8872 Expr::ArraySubscript { target, index } => {
8873 rewrite_expr_clock(target, now);
8874 rewrite_expr_clock(index, now);
8875 }
8876 Expr::AnyAll { expr, array, .. } => {
8877 rewrite_expr_clock(expr, now);
8878 rewrite_expr_clock(array, now);
8879 }
8880 Expr::Case {
8881 operand,
8882 branches,
8883 else_branch,
8884 } => {
8885 if let Some(o) = operand {
8886 rewrite_expr_clock(o, now);
8887 }
8888 for (w, t) in branches {
8889 rewrite_expr_clock(w, now);
8890 rewrite_expr_clock(t, now);
8891 }
8892 if let Some(e) = else_branch {
8893 rewrite_expr_clock(e, now);
8894 }
8895 }
8896 }
8897}
8898
8899fn clock_replacement_for(e: &Expr, now: i64) -> Option<Expr> {
8906 let (kind, name) = match e {
8907 Expr::FunctionCall { name, args } if args.is_empty() => (ClockSite::Fn, name.as_str()),
8908 Expr::Column(c) if c.qualifier.is_none() => (ClockSite::BareIdent, c.name.as_str()),
8909 _ => return None,
8910 };
8911 let matched = match name.len() {
8914 3 if kind == ClockSite::Fn && name.eq_ignore_ascii_case("now") => Some(true),
8915 12 if name.eq_ignore_ascii_case("current_date") => Some(false),
8916 17 if name.eq_ignore_ascii_case("current_timestamp") => Some(true),
8917 _ => None,
8918 };
8919 let is_timestamp = matched?;
8920 let payload = if is_timestamp {
8921 now
8922 } else {
8923 now.div_euclid(86_400_000_000)
8924 };
8925 let target = if is_timestamp {
8926 spg_sql::ast::CastTarget::Timestamp
8927 } else {
8928 spg_sql::ast::CastTarget::Date
8929 };
8930 Some(Expr::Cast {
8931 expr: alloc::boxed::Box::new(Expr::Literal(spg_sql::ast::Literal::Integer(payload))),
8932 target,
8933 })
8934}
8935
8936#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8937enum ClockSite {
8938 Fn,
8939 BareIdent,
8940}
8941
8942fn expand_group_by_all(s: &mut SelectStatement) {
8953 if !s.group_by_all {
8954 for (_, peer) in &mut s.unions {
8955 expand_group_by_all(peer);
8956 }
8957 return;
8958 }
8959 let mut groups: Vec<Expr> = Vec::new();
8960 for item in &s.items {
8961 if let SelectItem::Expr { expr, .. } = item
8962 && !aggregate::contains_aggregate(expr)
8963 {
8964 groups.push(expr.clone());
8965 }
8966 }
8967 s.group_by = Some(groups);
8968 s.group_by_all = false;
8969 for (_, peer) in &mut s.unions {
8970 expand_group_by_all(peer);
8971 }
8972}
8973
8974fn resolve_order_by_position(s: &mut SelectStatement) {
8975 for order in &mut s.order_by {
8980 match &order.expr {
8981 Expr::Literal(Literal::Integer(n)) if *n >= 1 => {
8982 if let Ok(idx_one_based) = usize::try_from(*n) {
8983 let idx = idx_one_based - 1;
8984 if idx < s.items.len()
8985 && let SelectItem::Expr { expr, .. } = &s.items[idx]
8986 {
8987 order.expr = expr.clone();
8988 }
8989 }
8990 }
8991 Expr::Column(c) if c.qualifier.is_none() => {
8992 for item in &s.items {
8994 if let SelectItem::Expr {
8995 expr,
8996 alias: Some(a),
8997 } = item
8998 && a == &c.name
8999 {
9000 order.expr = expr.clone();
9001 break;
9002 }
9003 }
9004 }
9005 _ => {}
9006 }
9007 }
9008 for (_, peer) in &mut s.unions {
9009 resolve_order_by_position(peer);
9010 }
9011}
9012
9013fn partial_sort_tagged(tagged: &mut Vec<(Vec<f64>, Row)>, keep: Option<usize>, descs: &[bool]) {
9026 let cmp = |a: &(Vec<f64>, Row), b: &(Vec<f64>, Row)| cmp_multi_key(&a.0, &b.0, descs);
9027 match keep {
9028 Some(k) if k < tagged.len() && k > 0 => {
9029 let pivot = k - 1;
9030 tagged.select_nth_unstable_by(pivot, cmp);
9031 tagged[..k].sort_by(cmp);
9032 tagged.truncate(k);
9033 }
9034 _ => {
9035 tagged.sort_by(cmp);
9036 }
9037 }
9038}
9039
9040fn sort_by_keys(tagged: &mut [(Vec<f64>, Row)], descs: &[bool]) {
9041 tagged.sort_by(|a, b| cmp_multi_key(&a.0, &b.0, descs));
9042}
9043
9044fn cmp_multi_key(a: &[f64], b: &[f64], descs: &[bool]) -> core::cmp::Ordering {
9048 use core::cmp::Ordering;
9049 for (i, (ka, kb)) in a.iter().zip(b.iter()).enumerate() {
9050 let ord = ka.partial_cmp(kb).unwrap_or(Ordering::Equal);
9051 let ord = if descs.get(i).copied().unwrap_or(false) {
9052 ord.reverse()
9053 } else {
9054 ord
9055 };
9056 if ord != Ordering::Equal {
9057 return ord;
9058 }
9059 }
9060 Ordering::Equal
9061}
9062
9063fn build_order_keys(
9066 order_by: &[OrderBy],
9067 row: &Row,
9068 ctx: &EvalContext,
9069) -> Result<Vec<f64>, EngineError> {
9070 let mut keys = Vec::with_capacity(order_by.len());
9071 for o in order_by {
9072 let v = eval::eval_expr(&o.expr, row, ctx)?;
9073 keys.push(value_to_order_key(&v)?);
9074 }
9075 Ok(keys)
9076}
9077
9078fn apply_offset_and_limit(rows: &mut Vec<Row>, offset: Option<u32>, limit: Option<u32>) {
9082 if let Some(off) = offset {
9083 let off = off as usize;
9084 if off >= rows.len() {
9085 rows.clear();
9086 } else {
9087 rows.drain(..off);
9088 }
9089 }
9090 if let Some(n) = limit {
9091 rows.truncate(n as usize);
9092 }
9093}
9094
9095fn resolve_foreign_key(
9109 local_table_name: &str,
9110 local_cols: &[ColumnSchema],
9111 fk: spg_sql::ast::ForeignKeyConstraint,
9112 catalog: &Catalog,
9113) -> Result<spg_storage::ForeignKeyConstraint, EngineError> {
9114 let mut local_columns = Vec::with_capacity(fk.columns.len());
9116 for name in &fk.columns {
9117 let pos = local_cols
9118 .iter()
9119 .position(|c| c.name == *name)
9120 .ok_or_else(|| {
9121 EngineError::Unsupported(alloc::format!(
9122 "FOREIGN KEY references unknown local column {name:?}"
9123 ))
9124 })?;
9125 local_columns.push(pos);
9126 }
9127 let is_self_ref = fk.parent_table == local_table_name;
9131 let (parent_cols_for_lookup, parent_table_str): (&[ColumnSchema], &str) = if is_self_ref {
9132 (local_cols, local_table_name)
9133 } else {
9134 let parent_table = catalog.get(&fk.parent_table).ok_or_else(|| {
9135 EngineError::Storage(StorageError::TableNotFound {
9136 name: fk.parent_table.clone(),
9137 })
9138 })?;
9139 (
9140 parent_table.schema().columns.as_slice(),
9141 fk.parent_table.as_str(),
9142 )
9143 };
9144 let parent_columns: Vec<usize> = if fk.parent_columns.is_empty() {
9149 if fk.columns.len() != 1 {
9150 return Err(EngineError::Unsupported(
9151 "composite FOREIGN KEY without explicit parent column list is not supported \
9152 — list the parent columns explicitly"
9153 .into(),
9154 ));
9155 }
9156 let pos = pick_pk_index_column(catalog, parent_table_str, is_self_ref, local_cols)
9158 .ok_or_else(|| {
9159 EngineError::Unsupported(alloc::format!(
9160 "parent table {parent_table_str:?} has no PRIMARY-key / UNIQUE BTree index \
9161 to default the FOREIGN KEY against"
9162 ))
9163 })?;
9164 alloc::vec![pos]
9165 } else {
9166 let mut out = Vec::with_capacity(fk.parent_columns.len());
9167 for name in &fk.parent_columns {
9168 let pos = parent_cols_for_lookup
9169 .iter()
9170 .position(|c| c.name == *name)
9171 .ok_or_else(|| {
9172 EngineError::Unsupported(alloc::format!(
9173 "FOREIGN KEY references unknown parent column \
9174 {name:?} on table {parent_table_str:?}"
9175 ))
9176 })?;
9177 out.push(pos);
9178 }
9179 out
9180 };
9181 if parent_columns.len() != local_columns.len() {
9182 return Err(EngineError::Unsupported(alloc::format!(
9183 "FOREIGN KEY arity mismatch: {} local columns vs {} parent columns",
9184 local_columns.len(),
9185 parent_columns.len()
9186 )));
9187 }
9188 if !is_self_ref {
9198 let parent_table = catalog.get(&fk.parent_table).expect("checked above");
9199 let primary_parent_col = parent_columns[0];
9200 let has_btree = parent_table
9201 .schema()
9202 .columns
9203 .get(primary_parent_col)
9204 .is_some()
9205 && parent_table.indices().iter().any(|idx| {
9206 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
9207 && idx.column_position == primary_parent_col
9208 && idx.partial_predicate.is_none()
9209 });
9210 if !has_btree {
9211 return Err(EngineError::Unsupported(alloc::format!(
9212 "FOREIGN KEY parent column on {:?} is not covered by an unconditional BTree \
9213 index — create one with `CREATE INDEX ... ON {} ({})` first",
9214 parent_table_str,
9215 parent_table_str,
9216 parent_table.schema().columns[primary_parent_col].name,
9217 )));
9218 }
9219 }
9220 let on_delete = fk_action_sql_to_storage(fk.on_delete);
9221 let on_update = fk_action_sql_to_storage(fk.on_update);
9222 Ok(spg_storage::ForeignKeyConstraint {
9223 name: fk.name,
9224 local_columns,
9225 parent_table: fk.parent_table,
9226 parent_columns,
9227 on_delete,
9228 on_update,
9229 })
9230}
9231
9232fn pick_pk_index_column(
9238 catalog: &Catalog,
9239 parent_name: &str,
9240 is_self_ref: bool,
9241 local_cols: &[ColumnSchema],
9242) -> Option<usize> {
9243 if is_self_ref {
9244 let _ = local_cols;
9248 return Some(0);
9249 }
9250 let parent = catalog.get(parent_name)?;
9251 parent.indices().iter().find_map(|idx| {
9252 if matches!(idx.kind, spg_storage::IndexKind::BTree(_))
9253 && idx.partial_predicate.is_none()
9254 && idx.included_columns.is_empty()
9255 && idx.expression.is_none()
9256 {
9257 Some(idx.column_position)
9258 } else {
9259 None
9260 }
9261 })
9262}
9263
9264fn resolve_on_conflict_columns(
9271 catalog: &Catalog,
9272 table_name: &str,
9273 target: &[String],
9274) -> Result<Vec<usize>, EngineError> {
9275 let table = catalog.get(table_name).ok_or_else(|| {
9276 EngineError::Storage(StorageError::TableNotFound {
9277 name: table_name.into(),
9278 })
9279 })?;
9280 if target.is_empty() {
9281 if let Some(uc) = table.schema().uniqueness_constraints.first() {
9291 return Ok(uc.columns.clone());
9292 }
9293 let pos = table
9294 .indices()
9295 .iter()
9296 .find_map(|idx| {
9297 if matches!(idx.kind, spg_storage::IndexKind::BTree(_))
9298 && idx.partial_predicate.is_none()
9299 && idx.included_columns.is_empty()
9300 && idx.expression.is_none()
9301 {
9302 Some(idx.column_position)
9303 } else {
9304 None
9305 }
9306 })
9307 .ok_or_else(|| {
9308 EngineError::Unsupported(alloc::format!(
9309 "ON CONFLICT without target requires a UNIQUE BTree index on {table_name:?}"
9310 ))
9311 })?;
9312 return Ok(alloc::vec![pos]);
9313 }
9314 let mut out = Vec::with_capacity(target.len());
9315 for name in target {
9316 let pos = table
9317 .schema()
9318 .columns
9319 .iter()
9320 .position(|c| c.name == *name)
9321 .ok_or_else(|| {
9322 EngineError::Unsupported(alloc::format!(
9323 "ON CONFLICT target column {name:?} not found on {table_name:?}"
9324 ))
9325 })?;
9326 out.push(pos);
9327 }
9328 Ok(out)
9329}
9330
9331fn on_conflict_key_exists(
9334 catalog: &Catalog,
9335 table_name: &str,
9336 column_pos: usize,
9337 key: &Value,
9338) -> bool {
9339 let Some(table) = catalog.get(table_name) else {
9340 return false;
9341 };
9342 let Some(idx_key) = spg_storage::IndexKey::from_value(key) else {
9343 return false;
9344 };
9345 table.indices().iter().any(|idx| {
9346 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
9347 && idx.column_position == column_pos
9348 && idx.partial_predicate.is_none()
9349 && !idx.lookup_eq(&idx_key).is_empty()
9350 })
9351}
9352
9353fn lookup_row_position_by_keys(
9359 catalog: &Catalog,
9360 table_name: &str,
9361 column_positions: &[usize],
9362 key: &[&Value],
9363) -> Option<usize> {
9364 let table = catalog.get(table_name)?;
9365 table.rows().iter().position(|r| {
9366 column_positions
9367 .iter()
9368 .enumerate()
9369 .all(|(i, &pos)| r.values.get(pos) == Some(key[i]))
9370 })
9371}
9372
9373fn on_conflict_keys_exist(
9378 catalog: &Catalog,
9379 table_name: &str,
9380 column_positions: &[usize],
9381 key: &[&Value],
9382) -> bool {
9383 if column_positions.len() == 1 {
9384 return on_conflict_key_exists(catalog, table_name, column_positions[0], key[0]);
9385 }
9386 let Some(table) = catalog.get(table_name) else {
9387 return false;
9388 };
9389 table.rows().iter().any(|r| {
9390 column_positions
9391 .iter()
9392 .enumerate()
9393 .all(|(i, &pos)| r.values.get(pos) == Some(key[i]))
9394 })
9395}
9396
9397fn apply_on_conflict_assignments(
9410 catalog: &Catalog,
9411 table_name: &str,
9412 target_pos: usize,
9413 incoming: &[Value],
9414 assignments: &[(String, Expr)],
9415 where_: Option<&Expr>,
9416) -> Result<Option<Vec<Value>>, EngineError> {
9417 let table = catalog.get(table_name).ok_or_else(|| {
9418 EngineError::Storage(StorageError::TableNotFound {
9419 name: table_name.into(),
9420 })
9421 })?;
9422 let schema_cols = table.schema().columns.clone();
9423 let existing = table
9424 .rows()
9425 .get(target_pos)
9426 .ok_or_else(|| {
9427 EngineError::Unsupported(alloc::format!(
9428 "ON CONFLICT DO UPDATE: row position {target_pos} out of bounds on {table_name:?}"
9429 ))
9430 })?
9431 .clone();
9432 let ctx = eval::EvalContext::new(&schema_cols, Some(table_name));
9433 if let Some(w) = where_ {
9435 let pred = w.clone();
9436 let pred = substitute_excluded_refs(pred, &schema_cols, incoming);
9437 let v = eval::eval_expr(&pred, &existing, &ctx)?;
9438 if !matches!(v, Value::Bool(true)) {
9439 return Ok(None);
9440 }
9441 }
9442 let mut new_values = existing.values.clone();
9443 for (col_name, expr) in assignments {
9444 let target_idx = schema_cols
9445 .iter()
9446 .position(|c| c.name == *col_name)
9447 .ok_or_else(|| {
9448 EngineError::Eval(EvalError::ColumnNotFound {
9449 name: col_name.clone(),
9450 })
9451 })?;
9452 let sub = substitute_excluded_refs(expr.clone(), &schema_cols, incoming);
9453 let v = eval::eval_expr(&sub, &existing, &ctx)?;
9454 new_values[target_idx] = coerce_value(v, schema_cols[target_idx].ty, col_name, target_idx)?;
9455 }
9456 Ok(Some(new_values))
9457}
9458
9459fn substitute_excluded_refs(expr: Expr, schema_cols: &[ColumnSchema], incoming: &[Value]) -> Expr {
9464 use spg_sql::ast::ColumnName;
9465 match expr {
9466 Expr::Column(ColumnName { qualifier, name })
9467 if qualifier
9468 .as_deref()
9469 .is_some_and(|q| q.eq_ignore_ascii_case("excluded")) =>
9470 {
9471 let pos = schema_cols.iter().position(|c| c.name == name);
9472 match pos {
9473 Some(p) => {
9474 let v = incoming.get(p).cloned().unwrap_or(Value::Null);
9475 value_to_literal_expr(v)
9476 .unwrap_or_else(|_| Expr::Literal(spg_sql::ast::Literal::Null))
9477 }
9478 None => Expr::Column(ColumnName { qualifier, name }),
9479 }
9480 }
9481 Expr::Binary { op, lhs, rhs } => Expr::Binary {
9482 op,
9483 lhs: Box::new(substitute_excluded_refs(*lhs, schema_cols, incoming)),
9484 rhs: Box::new(substitute_excluded_refs(*rhs, schema_cols, incoming)),
9485 },
9486 Expr::Unary { op, expr } => Expr::Unary {
9487 op,
9488 expr: Box::new(substitute_excluded_refs(*expr, schema_cols, incoming)),
9489 },
9490 Expr::FunctionCall { name, args } => Expr::FunctionCall {
9491 name,
9492 args: args
9493 .into_iter()
9494 .map(|a| substitute_excluded_refs(a, schema_cols, incoming))
9495 .collect(),
9496 },
9497 other => other,
9498 }
9499}
9500
9501fn enforce_uniqueness_inserts(
9524 catalog: &Catalog,
9525 child_table: &str,
9526 constraints: &[spg_storage::UniquenessConstraint],
9527 rows: &[Vec<Value>],
9528) -> Result<(), EngineError> {
9529 if constraints.is_empty() {
9530 return Ok(());
9531 }
9532 let table = catalog.get(child_table).ok_or_else(|| {
9533 EngineError::Storage(StorageError::TableNotFound {
9534 name: child_table.into(),
9535 })
9536 })?;
9537 for uc in constraints {
9538 for (batch_idx, row_values) in rows.iter().enumerate() {
9539 let key: Vec<&Value> = uc.columns.iter().map(|&i| &row_values[i]).collect();
9540 let has_null = key.iter().any(|v| matches!(v, Value::Null));
9541 if has_null && !uc.nulls_not_distinct {
9546 continue;
9547 }
9548 let collides_in_table = table.rows().iter().any(|prow| {
9550 uc.columns
9551 .iter()
9552 .enumerate()
9553 .all(|(i, &p)| prow.values.get(p) == Some(key[i]))
9554 });
9555 let collides_in_batch = rows[..batch_idx].iter().any(|earlier| {
9557 uc.columns
9558 .iter()
9559 .enumerate()
9560 .all(|(i, &p)| earlier.get(p) == Some(key[i]))
9561 });
9562 if collides_in_table || collides_in_batch {
9563 let kind = if uc.is_primary_key {
9564 "PRIMARY KEY"
9565 } else {
9566 "UNIQUE"
9567 };
9568 let col_names: Vec<String> = uc
9569 .columns
9570 .iter()
9571 .map(|&i| table.schema().columns[i].name.clone())
9572 .collect();
9573 return Err(EngineError::Unsupported(alloc::format!(
9574 "{kind} violation on {child_table:?} columns {col_names:?}: \
9575 row #{batch_idx} duplicates an existing key"
9576 )));
9577 }
9578 }
9579 }
9580 Ok(())
9581}
9582
9583fn predicate_truthy(v: &spg_storage::Value) -> bool {
9591 use spg_storage::Value as V;
9592 match v {
9593 V::Bool(b) => *b,
9594 V::Int(n) => *n != 0,
9595 V::BigInt(n) => *n != 0,
9596 V::SmallInt(n) => *n != 0,
9597 _ => false,
9598 }
9599}
9600
9601fn check_existing_unique_violation(
9606 idx: &spg_storage::Index,
9607 schema: &spg_storage::TableSchema,
9608 rows: &[spg_storage::Row],
9609) -> Result<(), EngineError> {
9610 let predicate_expr = match idx.partial_predicate.as_deref() {
9611 Some(s) => Some(spg_sql::parser::parse_expression(s).map_err(|e| {
9612 EngineError::Unsupported(alloc::format!(
9613 "stored partial predicate {s:?} failed to re-parse: {e:?}"
9614 ))
9615 })?),
9616 None => None,
9617 };
9618 let ctx = eval::EvalContext::new(&schema.columns, None);
9619 let key_positions = unique_key_positions(idx);
9620 let mut seen: alloc::vec::Vec<alloc::vec::Vec<spg_storage::Value>> = alloc::vec::Vec::new();
9621 for row in rows {
9622 if let Some(expr) = &predicate_expr {
9623 let v = eval::eval_expr(expr, row, &ctx).map_err(|e| {
9624 EngineError::Unsupported(alloc::format!(
9625 "evaluating UNIQUE INDEX predicate against existing row: {e:?}"
9626 ))
9627 })?;
9628 if !predicate_truthy(&v) {
9629 continue;
9630 }
9631 }
9632 let key: alloc::vec::Vec<spg_storage::Value> = key_positions
9633 .iter()
9634 .map(|&p| {
9635 row.values
9636 .get(p)
9637 .cloned()
9638 .unwrap_or(spg_storage::Value::Null)
9639 })
9640 .collect();
9641 if key.iter().any(|v| matches!(v, spg_storage::Value::Null)) {
9642 continue;
9643 }
9644 if seen.iter().any(|other| *other == key) {
9645 return Err(EngineError::Unsupported(alloc::format!(
9646 "CREATE UNIQUE INDEX {:?}: existing rows already violate the constraint",
9647 idx.name
9648 )));
9649 }
9650 seen.push(key);
9651 }
9652 Ok(())
9653}
9654
9655fn unique_key_positions(idx: &spg_storage::Index) -> alloc::vec::Vec<usize> {
9659 let mut out = alloc::vec::Vec::with_capacity(1 + idx.extra_column_positions.len());
9660 out.push(idx.column_position);
9661 out.extend_from_slice(&idx.extra_column_positions);
9662 out
9663}
9664
9665fn enforce_unique_index_inserts(
9673 catalog: &Catalog,
9674 table_name: &str,
9675 rows: &[alloc::vec::Vec<spg_storage::Value>],
9676) -> Result<(), EngineError> {
9677 let table = catalog.get(table_name).ok_or_else(|| {
9678 EngineError::Storage(StorageError::TableNotFound {
9679 name: table_name.into(),
9680 })
9681 })?;
9682 let schema = table.schema();
9683 let ctx = eval::EvalContext::new(&schema.columns, None);
9684 for idx in table.indices() {
9685 if !idx.is_unique {
9686 continue;
9687 }
9688 let predicate_expr = match idx.partial_predicate.as_deref() {
9690 Some(s) => Some(spg_sql::parser::parse_expression(s).map_err(|e| {
9691 EngineError::Unsupported(alloc::format!(
9692 "UNIQUE INDEX {:?} predicate {s:?} failed to re-parse: {e:?}",
9693 idx.name
9694 ))
9695 })?),
9696 None => None,
9697 };
9698 let key_positions = unique_key_positions(idx);
9699 let key_of = |values: &[spg_storage::Value]| -> alloc::vec::Vec<spg_storage::Value> {
9700 key_positions
9701 .iter()
9702 .map(|&p| values.get(p).cloned().unwrap_or(spg_storage::Value::Null))
9703 .collect()
9704 };
9705 let participates = |values: &[spg_storage::Value]| -> Result<bool, EngineError> {
9709 let Some(expr) = &predicate_expr else {
9710 return Ok(true);
9711 };
9712 let tmp_row = spg_storage::Row {
9713 values: values.to_vec(),
9714 };
9715 let v = eval::eval_expr(expr, &tmp_row, &ctx).map_err(|e| {
9716 EngineError::Unsupported(alloc::format!(
9717 "UNIQUE INDEX {:?} predicate eval: {e:?}",
9718 idx.name
9719 ))
9720 })?;
9721 Ok(predicate_truthy(&v))
9722 };
9723 for (batch_idx, row_values) in rows.iter().enumerate() {
9724 if !participates(row_values)? {
9725 continue;
9726 }
9727 let key = key_of(row_values);
9728 if key.iter().any(|v| matches!(v, spg_storage::Value::Null)) {
9729 continue;
9730 }
9731 for prow in table.rows() {
9733 if !participates(&prow.values)? {
9734 continue;
9735 }
9736 if key_of(&prow.values) == key {
9737 return Err(EngineError::Unsupported(alloc::format!(
9738 "UNIQUE INDEX {:?} violation on {table_name:?}: \
9739 row #{batch_idx} duplicates an existing key",
9740 idx.name
9741 )));
9742 }
9743 }
9744 for earlier in &rows[..batch_idx] {
9746 if !participates(earlier)? {
9747 continue;
9748 }
9749 if key_of(earlier) == key {
9750 return Err(EngineError::Unsupported(alloc::format!(
9751 "UNIQUE INDEX {:?} violation on {table_name:?}: \
9752 row #{batch_idx} duplicates an earlier row in the same batch",
9753 idx.name
9754 )));
9755 }
9756 }
9757 }
9758 }
9759 Ok(())
9760}
9761
9762fn any_column_changed(
9770 filter_cols: &[String],
9771 schema_cols: &[ColumnSchema],
9772 old_row: &Row,
9773 new_row: &Row,
9774) -> bool {
9775 for col_name in filter_cols {
9776 let Some(pos) = schema_cols
9777 .iter()
9778 .position(|c| c.name.eq_ignore_ascii_case(col_name))
9779 else {
9780 continue;
9781 };
9782 let old_v = old_row.values.get(pos);
9783 let new_v = new_row.values.get(pos);
9784 if old_v != new_v {
9785 return true;
9786 }
9787 }
9788 false
9789}
9790
9791fn enforce_check_constraints(
9796 catalog: &Catalog,
9797 table_name: &str,
9798 rows: &[alloc::vec::Vec<spg_storage::Value>],
9799) -> Result<(), EngineError> {
9800 let table = catalog.get(table_name).ok_or_else(|| {
9801 EngineError::Storage(StorageError::TableNotFound {
9802 name: table_name.into(),
9803 })
9804 })?;
9805 let schema = table.schema();
9806 if schema.checks.is_empty() {
9807 return Ok(());
9808 }
9809 let ctx = eval::EvalContext::new(&schema.columns, None);
9810 let mut parsed: alloc::vec::Vec<(usize, Expr)> = alloc::vec::Vec::new();
9811 for (i, src) in schema.checks.iter().enumerate() {
9812 let expr = spg_sql::parser::parse_expression(src).map_err(|e| {
9813 EngineError::Unsupported(alloc::format!(
9814 "CHECK constraint #{i} on {table_name:?} ({src:?}) failed to re-parse: {e:?}"
9815 ))
9816 })?;
9817 parsed.push((i, expr));
9818 }
9819 for (batch_idx, row_values) in rows.iter().enumerate() {
9820 let tmp_row = spg_storage::Row {
9821 values: row_values.clone(),
9822 };
9823 for (i, expr) in &parsed {
9824 let v = eval::eval_expr(expr, &tmp_row, &ctx).map_err(|e| {
9825 EngineError::Unsupported(alloc::format!(
9826 "CHECK constraint #{i} on {table_name:?} eval at row #{batch_idx}: {e:?}"
9827 ))
9828 })?;
9829 if matches!(v, spg_storage::Value::Bool(false)) {
9831 return Err(EngineError::Unsupported(alloc::format!(
9832 "CHECK constraint violation on {table_name:?} (row #{batch_idx}): {:?}",
9833 schema.checks[*i]
9834 )));
9835 }
9836 }
9837 }
9838 Ok(())
9839}
9840
9841fn enforce_fk_inserts(
9842 catalog: &Catalog,
9843 child_table: &str,
9844 fks: &[spg_storage::ForeignKeyConstraint],
9845 rows: &[Vec<Value>],
9846) -> Result<(), EngineError> {
9847 for fk in fks {
9848 let parent_is_self = fk.parent_table == child_table;
9849 let parent = if parent_is_self {
9850 catalog.get(child_table).ok_or_else(|| {
9853 EngineError::Storage(StorageError::TableNotFound {
9854 name: child_table.into(),
9855 })
9856 })?
9857 } else {
9858 catalog.get(&fk.parent_table).ok_or_else(|| {
9859 EngineError::Storage(StorageError::TableNotFound {
9860 name: fk.parent_table.clone(),
9861 })
9862 })?
9863 };
9864 for (batch_idx, row_values) in rows.iter().enumerate() {
9865 if fk.local_columns.len() == 1 {
9869 let v = &row_values[fk.local_columns[0]];
9870 if matches!(v, Value::Null) {
9871 continue;
9872 }
9873 let parent_col = fk.parent_columns[0];
9874 let key = spg_storage::IndexKey::from_value(v).ok_or_else(|| {
9875 EngineError::Unsupported(alloc::format!(
9876 "FOREIGN KEY column value of type {:?} is not index-eligible",
9877 v.data_type()
9878 ))
9879 })?;
9880 let present_committed = parent.indices().iter().any(|idx| {
9881 matches!(idx.kind, spg_storage::IndexKind::BTree(_))
9882 && idx.column_position == parent_col
9883 && idx.partial_predicate.is_none()
9884 && !idx.lookup_eq(&key).is_empty()
9885 });
9886 let present_in_batch = parent_is_self
9890 && rows[..batch_idx]
9891 .iter()
9892 .any(|earlier| earlier.get(parent_col) == Some(v));
9893 if !(present_committed || present_in_batch) {
9894 return Err(EngineError::Unsupported(alloc::format!(
9895 "FOREIGN KEY violation: no parent row in {:?} where {} = {:?}",
9896 fk.parent_table,
9897 parent
9898 .schema()
9899 .columns
9900 .get(parent_col)
9901 .map_or("?", |c| c.name.as_str()),
9902 v,
9903 )));
9904 }
9905 } else {
9906 if fk
9910 .local_columns
9911 .iter()
9912 .all(|&i| matches!(row_values.get(i), Some(Value::Null)))
9913 {
9914 continue;
9915 }
9916 let local: Vec<&Value> = fk.local_columns.iter().map(|&i| &row_values[i]).collect();
9917 let parent_match_committed = parent.rows().iter().any(|prow| {
9918 fk.parent_columns
9919 .iter()
9920 .enumerate()
9921 .all(|(i, &pi)| prow.values.get(pi) == Some(local[i]))
9922 });
9923 let parent_match_in_batch = parent_is_self
9924 && rows[..batch_idx].iter().any(|earlier| {
9925 fk.parent_columns
9926 .iter()
9927 .enumerate()
9928 .all(|(i, &pi)| earlier.get(pi) == Some(local[i]))
9929 });
9930 if !(parent_match_committed || parent_match_in_batch) {
9931 return Err(EngineError::Unsupported(alloc::format!(
9932 "FOREIGN KEY violation: no parent row in {:?} matching composite key",
9933 fk.parent_table,
9934 )));
9935 }
9936 }
9937 }
9938 }
9939 Ok(())
9940}
9941
9942#[derive(Debug, Clone)]
9946struct FkChildStep {
9947 child_table: String,
9948 action: FkChildAction,
9949}
9950
9951#[derive(Debug, Clone)]
9952enum FkChildAction {
9953 Delete { positions: Vec<usize> },
9955 SetNull {
9959 positions: Vec<usize>,
9960 columns: Vec<usize>,
9961 },
9962 SetDefault {
9966 positions: Vec<usize>,
9967 columns: Vec<usize>,
9968 defaults: Vec<Value>,
9969 },
9970}
9971
9972fn plan_fk_parent_deletions(
9988 catalog: &Catalog,
9989 parent_table_name: &str,
9990 to_delete_positions: &[usize],
9991 to_delete_rows: &[Vec<Value>],
9992) -> Result<Vec<FkChildStep>, EngineError> {
9993 use alloc::collections::{BTreeMap, BTreeSet};
9994 if to_delete_rows.is_empty() {
9995 return Ok(Vec::new());
9996 }
9997 let mut delete_plan: BTreeMap<String, BTreeSet<usize>> = BTreeMap::new();
9998 let mut setnull_plan: BTreeMap<String, BTreeSet<(usize, usize)>> = BTreeMap::new();
10000 let mut setdefault_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> = BTreeMap::new();
10001 let mut visited: BTreeSet<(String, usize)> = BTreeSet::new();
10002 for &p in to_delete_positions {
10003 visited.insert((parent_table_name.to_string(), p));
10004 }
10005 let mut work: Vec<(String, Vec<Value>)> = to_delete_rows
10006 .iter()
10007 .map(|r| (parent_table_name.to_string(), r.clone()))
10008 .collect();
10009 while let Some((cur_parent, parent_row)) = work.pop() {
10010 for child_name in catalog.table_names() {
10011 let child = catalog
10012 .get(&child_name)
10013 .expect("table_names → catalog.get round-trip is total");
10014 for fk in &child.schema().foreign_keys {
10015 if fk.parent_table != cur_parent {
10016 continue;
10017 }
10018 let parent_key: Vec<&Value> = fk
10019 .parent_columns
10020 .iter()
10021 .map(|&pi| &parent_row[pi])
10022 .collect();
10023 if parent_key.iter().any(|v| matches!(v, Value::Null)) {
10024 continue;
10025 }
10026 for (child_row_idx, child_row) in child.rows().iter().enumerate() {
10027 if child_name == cur_parent
10028 && visited.contains(&(child_name.clone(), child_row_idx))
10029 {
10030 continue;
10031 }
10032 let matches_key = fk
10033 .local_columns
10034 .iter()
10035 .enumerate()
10036 .all(|(i, &li)| child_row.values.get(li) == Some(parent_key[i]));
10037 if !matches_key {
10038 continue;
10039 }
10040 match fk.on_delete {
10041 spg_storage::FkAction::Restrict | spg_storage::FkAction::NoAction => {
10042 return Err(EngineError::Unsupported(alloc::format!(
10043 "FOREIGN KEY violation: DELETE on {cur_parent:?} is \
10044 restricted by FK from {child_name:?}.{:?}",
10045 fk.local_columns,
10046 )));
10047 }
10048 spg_storage::FkAction::Cascade => {
10049 if visited.insert((child_name.clone(), child_row_idx)) {
10050 delete_plan
10051 .entry(child_name.clone())
10052 .or_default()
10053 .insert(child_row_idx);
10054 work.push((child_name.clone(), child_row.values.clone()));
10055 }
10056 }
10057 spg_storage::FkAction::SetNull => {
10058 for &li in &fk.local_columns {
10060 let col = child.schema().columns.get(li).ok_or_else(|| {
10061 EngineError::Unsupported(alloc::format!(
10062 "FK local column {li} missing in {child_name:?}"
10063 ))
10064 })?;
10065 if !col.nullable {
10066 return Err(EngineError::Unsupported(alloc::format!(
10067 "FOREIGN KEY ON DELETE SET NULL: column \
10068 {child_name:?}.{:?} is NOT NULL — cannot SET NULL",
10069 col.name,
10070 )));
10071 }
10072 }
10073 let entry = setnull_plan.entry(child_name.clone()).or_default();
10074 for &li in &fk.local_columns {
10075 entry.insert((child_row_idx, li));
10076 }
10077 }
10078 spg_storage::FkAction::SetDefault => {
10079 let entry = setdefault_plan.entry(child_name.clone()).or_default();
10081 for &li in &fk.local_columns {
10082 let col = child.schema().columns.get(li).ok_or_else(|| {
10083 EngineError::Unsupported(alloc::format!(
10084 "FK local column {li} missing in {child_name:?}"
10085 ))
10086 })?;
10087 let default = col.default.clone().ok_or_else(|| {
10088 EngineError::Unsupported(alloc::format!(
10089 "FOREIGN KEY ON DELETE SET DEFAULT: column \
10090 {child_name:?}.{:?} has no DEFAULT declared",
10091 col.name,
10092 ))
10093 })?;
10094 entry.insert((child_row_idx, li), default);
10095 }
10096 }
10097 }
10098 }
10099 }
10100 }
10101 }
10102 let mut steps: Vec<FkChildStep> = Vec::new();
10110 for (child_table, entries) in setnull_plan {
10111 let (positions, columns): (Vec<usize>, Vec<usize>) = entries.into_iter().unzip();
10112 steps.push(FkChildStep {
10113 child_table,
10114 action: FkChildAction::SetNull { positions, columns },
10115 });
10116 }
10117 for (child_table, entries) in setdefault_plan {
10118 let mut positions = Vec::with_capacity(entries.len());
10119 let mut columns = Vec::with_capacity(entries.len());
10120 let mut defaults = Vec::with_capacity(entries.len());
10121 for ((p, c), v) in entries {
10122 positions.push(p);
10123 columns.push(c);
10124 defaults.push(v);
10125 }
10126 steps.push(FkChildStep {
10127 child_table,
10128 action: FkChildAction::SetDefault {
10129 positions,
10130 columns,
10131 defaults,
10132 },
10133 });
10134 }
10135 for (child_table, positions) in delete_plan {
10136 steps.push(FkChildStep {
10137 child_table,
10138 action: FkChildAction::Delete {
10139 positions: positions.into_iter().collect(),
10140 },
10141 });
10142 }
10143 Ok(steps)
10144}
10145
10146fn plan_fk_parent_updates(
10163 catalog: &Catalog,
10164 parent_table_name: &str,
10165 plan_with_old: &[(usize, Vec<Value>, Vec<Value>)],
10166) -> Result<Vec<FkChildStep>, EngineError> {
10167 use alloc::collections::BTreeMap;
10168 if plan_with_old.is_empty() {
10169 return Ok(Vec::new());
10170 }
10171 let delete_plan: BTreeMap<String, alloc::collections::BTreeSet<usize>> = BTreeMap::new();
10176 let mut setnull_plan: BTreeMap<String, alloc::collections::BTreeSet<(usize, usize)>> =
10177 BTreeMap::new();
10178 let mut setdefault_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> = BTreeMap::new();
10179 let mut cascade_plan: BTreeMap<String, BTreeMap<(usize, usize), Value>> = BTreeMap::new();
10181
10182 for child_name in catalog.table_names() {
10183 let child = catalog
10184 .get(&child_name)
10185 .expect("table_names → catalog.get total");
10186 for fk in &child.schema().foreign_keys {
10187 if fk.parent_table != parent_table_name {
10188 continue;
10189 }
10190 for (_pos, old_row, new_row) in plan_with_old {
10191 let key_changed = fk
10193 .parent_columns
10194 .iter()
10195 .any(|&pi| old_row.get(pi) != new_row.get(pi));
10196 if !key_changed {
10197 continue;
10198 }
10199 let old_key: Vec<&Value> =
10201 fk.parent_columns.iter().map(|&pi| &old_row[pi]).collect();
10202 if old_key.iter().any(|v| matches!(v, Value::Null)) {
10203 continue;
10205 }
10206 let new_key: Vec<&Value> =
10207 fk.parent_columns.iter().map(|&pi| &new_row[pi]).collect();
10208 for (child_row_idx, child_row) in child.rows().iter().enumerate() {
10209 if child_name == parent_table_name
10212 && plan_with_old.iter().any(|(p, _, _)| *p == child_row_idx)
10213 {
10214 continue;
10215 }
10216 let matches_key = fk
10217 .local_columns
10218 .iter()
10219 .enumerate()
10220 .all(|(i, &li)| child_row.values.get(li) == Some(old_key[i]));
10221 if !matches_key {
10222 continue;
10223 }
10224 match fk.on_update {
10225 spg_storage::FkAction::Restrict | spg_storage::FkAction::NoAction => {
10226 return Err(EngineError::Unsupported(alloc::format!(
10227 "FOREIGN KEY violation: UPDATE on {parent_table_name:?} PK is \
10228 restricted by FK from {child_name:?}.{:?}",
10229 fk.local_columns,
10230 )));
10231 }
10232 spg_storage::FkAction::Cascade => {
10233 let entry = cascade_plan.entry(child_name.clone()).or_default();
10235 for (i, &li) in fk.local_columns.iter().enumerate() {
10236 entry.insert((child_row_idx, li), new_key[i].clone());
10237 }
10238 }
10239 spg_storage::FkAction::SetNull => {
10240 for &li in &fk.local_columns {
10241 let col = child.schema().columns.get(li).ok_or_else(|| {
10242 EngineError::Unsupported(alloc::format!(
10243 "FK local column {li} missing in {child_name:?}"
10244 ))
10245 })?;
10246 if !col.nullable {
10247 return Err(EngineError::Unsupported(alloc::format!(
10248 "FOREIGN KEY ON UPDATE SET NULL: column \
10249 {child_name:?}.{:?} is NOT NULL",
10250 col.name,
10251 )));
10252 }
10253 }
10254 let entry = setnull_plan.entry(child_name.clone()).or_default();
10255 for &li in &fk.local_columns {
10256 entry.insert((child_row_idx, li));
10257 }
10258 }
10259 spg_storage::FkAction::SetDefault => {
10260 let entry = setdefault_plan.entry(child_name.clone()).or_default();
10261 for &li in &fk.local_columns {
10262 let col = child.schema().columns.get(li).ok_or_else(|| {
10263 EngineError::Unsupported(alloc::format!(
10264 "FK local column {li} missing in {child_name:?}"
10265 ))
10266 })?;
10267 let default = col.default.clone().ok_or_else(|| {
10268 EngineError::Unsupported(alloc::format!(
10269 "FOREIGN KEY ON UPDATE SET DEFAULT: column \
10270 {child_name:?}.{:?} has no DEFAULT",
10271 col.name,
10272 ))
10273 })?;
10274 entry.insert((child_row_idx, li), default);
10275 }
10276 }
10277 }
10278 }
10279 }
10280 }
10281 }
10282 let mut steps: Vec<FkChildStep> = Vec::new();
10285 for (child_table, entries) in cascade_plan {
10286 let mut positions = Vec::with_capacity(entries.len());
10287 let mut columns = Vec::with_capacity(entries.len());
10288 let mut defaults = Vec::with_capacity(entries.len());
10289 for ((p, c), v) in entries {
10290 positions.push(p);
10291 columns.push(c);
10292 defaults.push(v);
10293 }
10294 steps.push(FkChildStep {
10299 child_table,
10300 action: FkChildAction::SetDefault {
10301 positions,
10302 columns,
10303 defaults,
10304 },
10305 });
10306 }
10307 for (child_table, entries) in setnull_plan {
10308 let (positions, columns): (Vec<usize>, Vec<usize>) = entries.into_iter().unzip();
10309 steps.push(FkChildStep {
10310 child_table,
10311 action: FkChildAction::SetNull { positions, columns },
10312 });
10313 }
10314 for (child_table, entries) in setdefault_plan {
10315 let mut positions = Vec::with_capacity(entries.len());
10316 let mut columns = Vec::with_capacity(entries.len());
10317 let mut defaults = Vec::with_capacity(entries.len());
10318 for ((p, c), v) in entries {
10319 positions.push(p);
10320 columns.push(c);
10321 defaults.push(v);
10322 }
10323 steps.push(FkChildStep {
10324 child_table,
10325 action: FkChildAction::SetDefault {
10326 positions,
10327 columns,
10328 defaults,
10329 },
10330 });
10331 }
10332 let _ = delete_plan; Ok(steps)
10334}
10335
10336fn apply_fk_child_step(catalog: &mut Catalog, step: &FkChildStep) -> Result<(), EngineError> {
10340 let child = catalog.get_mut(&step.child_table).ok_or_else(|| {
10341 EngineError::Storage(StorageError::TableNotFound {
10342 name: step.child_table.clone(),
10343 })
10344 })?;
10345 match &step.action {
10346 FkChildAction::Delete { positions } => {
10347 let _ = child.delete_rows(positions);
10348 }
10349 FkChildAction::SetNull { positions, columns } => {
10350 apply_per_cell_writes(child, positions, columns, |_| Value::Null)?;
10351 }
10352 FkChildAction::SetDefault {
10353 positions,
10354 columns,
10355 defaults,
10356 } => {
10357 apply_per_cell_writes(child, positions, columns, |i| defaults[i].clone())?;
10358 }
10359 }
10360 Ok(())
10361}
10362
10363fn apply_per_cell_writes(
10369 child: &mut spg_storage::Table,
10370 positions: &[usize],
10371 columns: &[usize],
10372 mut value_for: impl FnMut(usize) -> Value,
10373) -> Result<(), EngineError> {
10374 use alloc::collections::BTreeMap;
10375 let mut by_row: BTreeMap<usize, Vec<(usize, Value)>> = BTreeMap::new();
10376 for i in 0..positions.len() {
10377 by_row
10378 .entry(positions[i])
10379 .or_default()
10380 .push((columns[i], value_for(i)));
10381 }
10382 for (pos, mutations) in by_row {
10383 let mut new_values = child.rows()[pos].values.clone();
10384 for (col, v) in mutations {
10385 if let Some(slot) = new_values.get_mut(col) {
10386 *slot = v;
10387 }
10388 }
10389 child
10390 .update_row(pos, new_values)
10391 .map_err(EngineError::Storage)?;
10392 }
10393 Ok(())
10394}
10395
10396fn fk_action_sql_to_storage(a: spg_sql::ast::FkAction) -> spg_storage::FkAction {
10397 match a {
10398 spg_sql::ast::FkAction::Restrict => spg_storage::FkAction::Restrict,
10399 spg_sql::ast::FkAction::Cascade => spg_storage::FkAction::Cascade,
10400 spg_sql::ast::FkAction::SetNull => spg_storage::FkAction::SetNull,
10401 spg_sql::ast::FkAction::SetDefault => spg_storage::FkAction::SetDefault,
10402 spg_sql::ast::FkAction::NoAction => spg_storage::FkAction::NoAction,
10403 }
10404}
10405
10406fn resolve_column_default_free(
10412 col: &ColumnSchema,
10413 clock_fn: Option<ClockFn>,
10414) -> Result<Value, EngineError> {
10415 if let Some(rt) = &col.runtime_default {
10416 return eval_runtime_default_free(rt, col.ty, clock_fn);
10417 }
10418 Ok(col.default.clone().unwrap_or(Value::Null))
10419}
10420
10421fn eval_runtime_default_free(
10422 rt: &str,
10423 ty: DataType,
10424 clock_fn: Option<ClockFn>,
10425) -> Result<Value, EngineError> {
10426 let s = rt.trim().to_ascii_lowercase();
10427 let canonical = s.trim_end_matches("()");
10428 let now_us = match clock_fn {
10429 Some(f) => f(),
10430 None => 0,
10431 };
10432 let v = match canonical {
10433 "now" | "current_timestamp" | "localtimestamp" => Value::Timestamp(now_us),
10434 "current_date" => Value::Date((now_us / 86_400_000_000) as i32),
10435 "current_time" | "localtime" => Value::Timestamp(now_us),
10436 other => {
10437 return Err(EngineError::Unsupported(alloc::format!(
10438 "runtime DEFAULT expression {other:?} not supported \
10439 (v7.9.21 whitelist: now() / current_timestamp / \
10440 current_date / current_time / localtimestamp / \
10441 localtime)"
10442 )));
10443 }
10444 };
10445 coerce_value(v, ty, "DEFAULT", 0)
10446}
10447
10448fn is_runtime_default_expr(expr: &Expr) -> bool {
10454 match expr {
10455 Expr::FunctionCall { .. } => true,
10456 Expr::Unary { expr, .. } => is_runtime_default_expr(expr),
10457 _ => false,
10458 }
10459}
10460
10461fn column_def_to_schema(c: ColumnDef) -> Result<ColumnSchema, EngineError> {
10462 let ty = column_type_to_data_type(c.ty);
10463 let mut schema = ColumnSchema::new(c.name.clone(), ty, c.nullable);
10464 if let Some(default_expr) = c.default {
10465 if is_runtime_default_expr(&default_expr) {
10471 let display = alloc::format!("{default_expr}");
10472 schema = schema.with_runtime_default(display);
10473 } else {
10474 let raw = literal_expr_to_value(default_expr)?;
10475 let coerced = coerce_value(raw, ty, &c.name, 0)?;
10476 schema = schema.with_default(coerced);
10477 }
10478 }
10479 if c.auto_increment {
10480 if !matches!(ty, DataType::SmallInt | DataType::Int | DataType::BigInt) {
10482 return Err(EngineError::Unsupported(alloc::format!(
10483 "AUTO_INCREMENT requires an integer column type, got {ty:?}"
10484 )));
10485 }
10486 schema = schema.with_auto_increment();
10487 }
10488 Ok(schema)
10489}
10490
10491fn decode_bytea_literal(s: &str) -> Result<alloc::vec::Vec<u8>, &'static str> {
10496 let s = s.trim();
10497 if let Some(hex) = s.strip_prefix("\\x").or_else(|| s.strip_prefix("\\X")) {
10498 let cleaned: alloc::string::String = hex.chars().filter(|c| !c.is_whitespace()).collect();
10500 if cleaned.len() % 2 != 0 {
10501 return Err("odd-length hex literal");
10502 }
10503 let mut out = alloc::vec::Vec::with_capacity(cleaned.len() / 2);
10504 let cleaned_bytes = cleaned.as_bytes();
10505 for i in (0..cleaned_bytes.len()).step_by(2) {
10506 let hi = hex_nibble(cleaned_bytes[i])?;
10507 let lo = hex_nibble(cleaned_bytes[i + 1])?;
10508 out.push((hi << 4) | lo);
10509 }
10510 return Ok(out);
10511 }
10512 let bytes = s.as_bytes();
10515 let mut out = alloc::vec::Vec::with_capacity(bytes.len());
10516 let mut i = 0;
10517 while i < bytes.len() {
10518 let b = bytes[i];
10519 if b == b'\\' && i + 1 < bytes.len() {
10520 let n = bytes[i + 1];
10521 if n == b'\\' {
10522 out.push(b'\\');
10523 i += 2;
10524 continue;
10525 }
10526 if n.is_ascii_digit()
10527 && i + 3 < bytes.len()
10528 && bytes[i + 2].is_ascii_digit()
10529 && bytes[i + 3].is_ascii_digit()
10530 {
10531 let oct = |x: u8| (x - b'0') as u32;
10532 let v = oct(n) * 64 + oct(bytes[i + 2]) * 8 + oct(bytes[i + 3]);
10533 if v <= 0xFF {
10534 out.push(v as u8);
10535 i += 4;
10536 continue;
10537 }
10538 }
10539 }
10540 out.push(b);
10541 i += 1;
10542 }
10543 Ok(out)
10544}
10545
10546fn hex_nibble(b: u8) -> Result<u8, &'static str> {
10547 match b {
10548 b'0'..=b'9' => Ok(b - b'0'),
10549 b'a'..=b'f' => Ok(b - b'a' + 10),
10550 b'A'..=b'F' => Ok(b - b'A' + 10),
10551 _ => Err("invalid hex digit"),
10552 }
10553}
10554
10555fn array_literal_widen(items: alloc::vec::Vec<Value>) -> Value {
10569 let mut has_text = false;
10570 let mut has_bigint = false;
10571 let mut has_int = false;
10572 for v in &items {
10573 match v {
10574 Value::Null => {}
10575 Value::Text(_) | Value::Json(_) => has_text = true,
10576 Value::BigInt(_) => has_bigint = true,
10577 Value::Int(_) | Value::SmallInt(_) => has_int = true,
10578 _ => has_text = true,
10579 }
10580 }
10581 if has_text || (!has_bigint && !has_int) {
10582 let out: alloc::vec::Vec<Option<alloc::string::String>> = items
10583 .into_iter()
10584 .map(|v| match v {
10585 Value::Null => None,
10586 Value::Text(s) | Value::Json(s) => Some(s),
10587 other => Some(alloc::format!("{other:?}")),
10588 })
10589 .collect();
10590 return Value::TextArray(out);
10591 }
10592 if has_bigint {
10593 let out: alloc::vec::Vec<Option<i64>> = items
10594 .into_iter()
10595 .map(|v| match v {
10596 Value::Null => None,
10597 Value::Int(n) => Some(i64::from(n)),
10598 Value::SmallInt(n) => Some(i64::from(n)),
10599 Value::BigInt(n) => Some(n),
10600 _ => unreachable!("widen: unexpected non-integer in BigInt path"),
10601 })
10602 .collect();
10603 return Value::BigIntArray(out);
10604 }
10605 let out: alloc::vec::Vec<Option<i32>> = items
10606 .into_iter()
10607 .map(|v| match v {
10608 Value::Null => None,
10609 Value::Int(n) => Some(n),
10610 Value::SmallInt(n) => Some(i32::from(n)),
10611 _ => unreachable!("widen: unexpected non-i32-compatible in Int path"),
10612 })
10613 .collect();
10614 Value::IntArray(out)
10615}
10616
10617fn decode_text_array_literal(
10618 s: &str,
10619) -> Result<alloc::vec::Vec<Option<alloc::string::String>>, &'static str> {
10620 let trimmed = s.trim();
10621 let inner = trimmed
10622 .strip_prefix('{')
10623 .and_then(|x| x.strip_suffix('}'))
10624 .ok_or("TEXT[] literal must be enclosed in '{...}'")?;
10625 let mut out: alloc::vec::Vec<Option<alloc::string::String>> = alloc::vec::Vec::new();
10626 if inner.trim().is_empty() {
10627 return Ok(out);
10628 }
10629 let bytes = inner.as_bytes();
10630 let mut i = 0;
10631 while i <= bytes.len() {
10632 while i < bytes.len() && (bytes[i] == b' ' || bytes[i] == b'\t') {
10634 i += 1;
10635 }
10636 if i < bytes.len() && bytes[i] == b'"' {
10638 i += 1; let mut buf = alloc::string::String::new();
10640 while i < bytes.len() && bytes[i] != b'"' {
10641 if bytes[i] == b'\\' && i + 1 < bytes.len() {
10642 buf.push(bytes[i + 1] as char);
10643 i += 2;
10644 } else {
10645 buf.push(bytes[i] as char);
10646 i += 1;
10647 }
10648 }
10649 if i >= bytes.len() {
10650 return Err("unterminated quoted element");
10651 }
10652 i += 1; out.push(Some(buf));
10654 } else {
10655 let start = i;
10657 while i < bytes.len() && bytes[i] != b',' {
10658 i += 1;
10659 }
10660 let raw = inner[start..i].trim();
10661 if raw.eq_ignore_ascii_case("NULL") {
10662 out.push(None);
10663 } else {
10664 out.push(Some(alloc::string::ToString::to_string(raw)));
10665 }
10666 }
10667 while i < bytes.len() && (bytes[i] == b' ' || bytes[i] == b'\t') {
10669 i += 1;
10670 }
10671 if i >= bytes.len() {
10672 break;
10673 }
10674 if bytes[i] != b',' {
10675 return Err("expected ',' between TEXT[] elements");
10676 }
10677 i += 1;
10678 }
10679 Ok(out)
10680}
10681
10682fn encode_text_array(items: &[Option<alloc::string::String>]) -> alloc::string::String {
10687 let mut out = alloc::string::String::with_capacity(2 + items.len() * 8);
10688 out.push('{');
10689 for (i, item) in items.iter().enumerate() {
10690 if i > 0 {
10691 out.push(',');
10692 }
10693 match item {
10694 None => out.push_str("NULL"),
10695 Some(s) => {
10696 let needs_quote = s.is_empty()
10697 || s.eq_ignore_ascii_case("NULL")
10698 || s.chars()
10699 .any(|c| matches!(c, ',' | '{' | '}' | '"' | '\\' | ' ' | '\t'));
10700 if needs_quote {
10701 out.push('"');
10702 for c in s.chars() {
10703 if c == '"' || c == '\\' {
10704 out.push('\\');
10705 }
10706 out.push(c);
10707 }
10708 out.push('"');
10709 } else {
10710 out.push_str(s);
10711 }
10712 }
10713 }
10714 }
10715 out.push('}');
10716 out
10717}
10718
10719fn encode_bytea_hex(b: &[u8]) -> alloc::string::String {
10723 let mut out = alloc::string::String::with_capacity(2 + 2 * b.len());
10724 out.push_str("\\x");
10725 for byte in b {
10726 let hi = byte >> 4;
10727 let lo = byte & 0x0F;
10728 out.push(hex_digit(hi));
10729 out.push(hex_digit(lo));
10730 }
10731 out
10732}
10733
10734const fn hex_digit(n: u8) -> char {
10735 match n {
10736 0..=9 => (b'0' + n) as char,
10737 10..=15 => (b'a' + n - 10) as char,
10738 _ => '?',
10739 }
10740}
10741
10742const fn column_type_to_data_type(t: ColumnTypeName) -> DataType {
10743 match t {
10744 ColumnTypeName::SmallInt => DataType::SmallInt,
10745 ColumnTypeName::Int => DataType::Int,
10746 ColumnTypeName::BigInt => DataType::BigInt,
10747 ColumnTypeName::Float => DataType::Float,
10748 ColumnTypeName::Text => DataType::Text,
10749 ColumnTypeName::Varchar(n) => DataType::Varchar(n),
10750 ColumnTypeName::Char(n) => DataType::Char(n),
10751 ColumnTypeName::Bool => DataType::Bool,
10752 ColumnTypeName::Vector { dim, encoding } => DataType::Vector {
10753 dim,
10754 encoding: match encoding {
10755 SqlVecEncoding::F32 => VecEncoding::F32,
10756 SqlVecEncoding::Sq8 => VecEncoding::Sq8,
10757 SqlVecEncoding::F16 => VecEncoding::F16,
10758 },
10759 },
10760 ColumnTypeName::Numeric(precision, scale) => DataType::Numeric { precision, scale },
10761 ColumnTypeName::Date => DataType::Date,
10762 ColumnTypeName::Timestamp => DataType::Timestamp,
10763 ColumnTypeName::Timestamptz => DataType::Timestamptz,
10764 ColumnTypeName::Json => DataType::Json,
10765 ColumnTypeName::Jsonb => DataType::Jsonb,
10766 ColumnTypeName::Bytes => DataType::Bytes,
10767 ColumnTypeName::TextArray => DataType::TextArray,
10768 ColumnTypeName::IntArray => DataType::IntArray,
10769 ColumnTypeName::BigIntArray => DataType::BigIntArray,
10770 ColumnTypeName::TsVector => DataType::TsVector,
10771 ColumnTypeName::TsQuery => DataType::TsQuery,
10772 }
10773}
10774
10775fn literal_expr_to_value(expr: Expr) -> Result<Value, EngineError> {
10779 match expr {
10780 Expr::Literal(l) => Ok(literal_to_value(l)),
10781 Expr::Cast { expr, target } => {
10782 let inner_value = literal_expr_to_value(*expr)?;
10783 crate::eval::cast_value(inner_value, target).map_err(EngineError::Eval)
10784 }
10785 Expr::Unary {
10786 op: UnOp::Neg,
10787 expr,
10788 } => match *expr {
10789 Expr::Literal(Literal::Integer(n)) => {
10790 let neg = n.checked_neg().ok_or_else(|| {
10793 EngineError::Unsupported("integer literal overflow on negation".into())
10794 })?;
10795 Ok(int_value_for(neg))
10796 }
10797 Expr::Literal(Literal::Float(x)) => Ok(Value::Float(-x)),
10798 other => Err(EngineError::Unsupported(alloc::format!(
10799 "unary minus over non-literal expression: {other:?}"
10800 ))),
10801 },
10802 Expr::Array(items) => {
10810 let mut materialised: alloc::vec::Vec<Value> =
10811 alloc::vec::Vec::with_capacity(items.len());
10812 for elem in items {
10813 materialised.push(literal_expr_to_value(elem)?);
10814 }
10815 Ok(array_literal_widen(materialised))
10816 }
10817 other => Err(EngineError::Unsupported(alloc::format!(
10818 "non-literal INSERT value expression: {other:?}"
10819 ))),
10820 }
10821}
10822
10823fn literal_to_value(l: Literal) -> Value {
10824 match l {
10825 Literal::Integer(n) => int_value_for(n),
10826 Literal::Float(x) => Value::Float(x),
10827 Literal::String(s) => Value::Text(s),
10828 Literal::Bool(b) => Value::Bool(b),
10829 Literal::Null => Value::Null,
10830 Literal::Vector(v) => Value::Vector(v),
10831 Literal::Interval { months, micros, .. } => Value::Interval { months, micros },
10832 }
10833}
10834
10835fn int_value_for(n: i64) -> Value {
10839 if let Ok(small) = i32::try_from(n) {
10840 Value::Int(small)
10841 } else {
10842 Value::BigInt(n)
10843 }
10844}
10845
10846#[allow(clippy::too_many_lines)]
10852fn coerce_value(
10853 v: Value,
10854 expected: DataType,
10855 col_name: &str,
10856 position: usize,
10857) -> Result<Value, EngineError> {
10858 if v.is_null() {
10859 return Ok(Value::Null);
10860 }
10861 let actual = v.data_type().expect("non-null");
10862 if actual == expected {
10863 return Ok(v);
10864 }
10865 let coerced = match (v, expected) {
10866 (Value::Int(n), DataType::BigInt) => Some(Value::BigInt(i64::from(n))),
10867 (Value::Int(n), DataType::Float) => Some(Value::Float(f64::from(n))),
10868 (Value::Int(n), DataType::SmallInt) => i16::try_from(n).ok().map(Value::SmallInt),
10869 (Value::Int(n), DataType::Numeric { precision, scale }) => Some(numeric_from_integer(
10870 i128::from(n),
10871 precision,
10872 scale,
10873 col_name,
10874 )?),
10875 (Value::SmallInt(n), DataType::Int) => Some(Value::Int(i32::from(n))),
10876 (Value::SmallInt(n), DataType::BigInt) => Some(Value::BigInt(i64::from(n))),
10877 (Value::SmallInt(n), DataType::Float) => Some(Value::Float(f64::from(n))),
10878 (Value::SmallInt(n), DataType::Numeric { precision, scale }) => Some(numeric_from_integer(
10879 i128::from(n),
10880 precision,
10881 scale,
10882 col_name,
10883 )?),
10884 (Value::BigInt(n), DataType::Int) => i32::try_from(n).ok().map(Value::Int),
10885 (Value::BigInt(n), DataType::SmallInt) => i16::try_from(n).ok().map(Value::SmallInt),
10886 #[allow(clippy::cast_precision_loss)]
10887 (Value::BigInt(n), DataType::Float) => Some(Value::Float(n as f64)),
10888 (Value::BigInt(n), DataType::Numeric { precision, scale }) => Some(numeric_from_integer(
10889 i128::from(n),
10890 precision,
10891 scale,
10892 col_name,
10893 )?),
10894 (Value::Float(x), DataType::Numeric { precision, scale }) => {
10895 Some(numeric_from_float(x, precision, scale, col_name)?)
10896 }
10897 (Value::Text(s), DataType::Date) => {
10899 let d = eval::parse_date_literal(&s).ok_or_else(|| {
10900 EngineError::Eval(EvalError::TypeMismatch {
10901 detail: alloc::format!("cannot parse {s:?} as DATE for column `{col_name}`"),
10902 })
10903 })?;
10904 Some(Value::Date(d))
10905 }
10906 (Value::Text(s), DataType::Json | DataType::Jsonb) => Some(Value::Json(s)),
10910 (Value::Json(s), DataType::Text) => Some(Value::Text(s)),
10911 (Value::Text(s), DataType::Bytes) => {
10918 let bytes = decode_bytea_literal(&s).map_err(|e| {
10919 EngineError::Eval(EvalError::TypeMismatch {
10920 detail: alloc::format!(
10921 "cannot parse {s:?} as BYTEA for column `{col_name}`: {e}"
10922 ),
10923 })
10924 })?;
10925 Some(Value::Bytes(bytes))
10926 }
10927 (Value::Bytes(b), DataType::Text) => Some(Value::Text(encode_bytea_hex(&b))),
10931 (Value::Text(s), DataType::TextArray) => {
10936 let arr = decode_text_array_literal(&s).map_err(|e| {
10937 EngineError::Eval(EvalError::TypeMismatch {
10938 detail: alloc::format!(
10939 "cannot parse {s:?} as TEXT[] for column `{col_name}`: {e}"
10940 ),
10941 })
10942 })?;
10943 Some(Value::TextArray(arr))
10944 }
10945 (Value::TextArray(items), DataType::Text) => Some(Value::Text(encode_text_array(&items))),
10949 (Value::Text(s), DataType::Timestamp | DataType::Timestamptz) => {
10950 let t = eval::parse_timestamp_literal(&s).ok_or_else(|| {
10951 EngineError::Eval(EvalError::TypeMismatch {
10952 detail: alloc::format!(
10953 "cannot parse {s:?} as TIMESTAMP for column `{col_name}`"
10954 ),
10955 })
10956 })?;
10957 Some(Value::Timestamp(t))
10958 }
10959 (Value::Date(d), DataType::Timestamp | DataType::Timestamptz) => {
10962 Some(Value::Timestamp(i64::from(d) * 86_400_000_000))
10963 }
10964 (Value::Timestamp(t), DataType::Timestamptz) => Some(Value::Timestamp(t)),
10968 (Value::Timestamp(t), DataType::Date) => {
10969 let days = t.div_euclid(86_400_000_000);
10970 i32::try_from(days).ok().map(Value::Date)
10971 }
10972 (
10973 Value::Numeric {
10974 scaled,
10975 scale: src_scale,
10976 },
10977 DataType::Numeric { precision, scale },
10978 ) => Some(numeric_rescale(
10979 scaled, src_scale, precision, scale, col_name,
10980 )?),
10981 #[allow(clippy::cast_precision_loss)]
10982 (Value::Numeric { scaled, scale }, DataType::Float) => {
10983 let mut div = 1.0_f64;
10984 for _ in 0..scale {
10985 div *= 10.0;
10986 }
10987 Some(Value::Float((scaled as f64) / div))
10988 }
10989 (Value::Numeric { scaled, scale }, DataType::Int) => {
10990 let truncated = numeric_truncate_to_integer(scaled, scale);
10991 i32::try_from(truncated).ok().map(Value::Int)
10992 }
10993 (Value::Numeric { scaled, scale }, DataType::BigInt) => {
10994 let truncated = numeric_truncate_to_integer(scaled, scale);
10995 i64::try_from(truncated).ok().map(Value::BigInt)
10996 }
10997 (Value::Numeric { scaled, scale }, DataType::SmallInt) => {
10998 let truncated = numeric_truncate_to_integer(scaled, scale);
10999 i16::try_from(truncated).ok().map(Value::SmallInt)
11000 }
11001 (Value::Text(s), DataType::Varchar(max)) => {
11003 if u32::try_from(s.chars().count()).unwrap_or(u32::MAX) <= max {
11004 Some(Value::Text(s))
11005 } else {
11006 return Err(EngineError::Unsupported(alloc::format!(
11007 "value for VARCHAR({max}) column `{col_name}` exceeds length: \
11008 {} chars",
11009 s.chars().count()
11010 )));
11011 }
11012 }
11013 (
11021 Value::Vector(v),
11022 DataType::Vector {
11023 dim,
11024 encoding: VecEncoding::Sq8,
11025 },
11026 ) if v.len() == dim as usize => Some(Value::Sq8Vector(spg_storage::quantize::quantize(&v))),
11027 (
11032 Value::Vector(v),
11033 DataType::Vector {
11034 dim,
11035 encoding: VecEncoding::F16,
11036 },
11037 ) if v.len() == dim as usize => Some(Value::HalfVector(
11038 spg_storage::halfvec::HalfVector::from_f32_slice(&v),
11039 )),
11040 (Value::Text(s), DataType::Char(size)) => {
11044 let len = u32::try_from(s.chars().count()).unwrap_or(u32::MAX);
11045 if len > size {
11046 return Err(EngineError::Unsupported(alloc::format!(
11047 "value for CHAR({size}) column `{col_name}` exceeds length: \
11048 {len} chars"
11049 )));
11050 }
11051 let need = (size - len) as usize;
11052 let mut padded = s;
11053 padded.reserve(need);
11054 for _ in 0..need {
11055 padded.push(' ');
11056 }
11057 Some(Value::Text(padded))
11058 }
11059 _ => None,
11060 };
11061 coerced.ok_or(EngineError::Storage(StorageError::TypeMismatch {
11062 column: col_name.into(),
11063 expected,
11064 actual,
11065 position,
11066 }))
11067}
11068
11069fn render_function_args(args: &[spg_sql::ast::FunctionArg]) -> alloc::string::String {
11075 use core::fmt::Write;
11076 let mut out = alloc::string::String::from("(");
11077 for (i, a) in args.iter().enumerate() {
11078 if i > 0 {
11079 out.push_str(", ");
11080 }
11081 match a.mode {
11082 spg_sql::ast::FunctionArgMode::In => {}
11083 spg_sql::ast::FunctionArgMode::Out => out.push_str("OUT "),
11084 spg_sql::ast::FunctionArgMode::InOut => out.push_str("INOUT "),
11085 }
11086 if let Some(n) = &a.name {
11087 out.push_str(n);
11088 out.push(' ');
11089 }
11090 match &a.ty {
11091 spg_sql::ast::FunctionArgType::Typed(t) => {
11092 let _ = write!(out, "{t}");
11093 }
11094 spg_sql::ast::FunctionArgType::Raw(s) => out.push_str(s),
11095 }
11096 }
11097 out.push(')');
11098 out
11099}
11100
11101#[cfg(test)]
11102mod tests {
11103 use super::*;
11104 use alloc::vec;
11105
11106 fn unwrap_command_ok(r: &QueryResult) -> usize {
11107 match r {
11108 QueryResult::CommandOk { affected, .. } => *affected,
11109 QueryResult::Rows { .. } => panic!("expected CommandOk, got Rows"),
11110 }
11111 }
11112
11113 #[test]
11114 fn create_table_registers_schema() {
11115 let mut e = Engine::new();
11116 e.execute("CREATE TABLE foo (a INT NOT NULL, b TEXT)")
11117 .unwrap();
11118 assert_eq!(e.catalog().table_count(), 1);
11119 let t = e.catalog().get("foo").unwrap();
11120 assert_eq!(t.schema().columns.len(), 2);
11121 assert_eq!(t.schema().columns[0].ty, DataType::Int);
11122 assert!(!t.schema().columns[0].nullable);
11123 assert_eq!(t.schema().columns[1].ty, DataType::Text);
11124 }
11125
11126 #[test]
11127 fn create_table_vector_default_is_f32_encoded() {
11128 let mut e = Engine::new();
11129 e.execute("CREATE TABLE t (v VECTOR(8))").unwrap();
11130 let t = e.catalog().get("t").unwrap();
11131 assert_eq!(
11132 t.schema().columns[0].ty,
11133 DataType::Vector {
11134 dim: 8,
11135 encoding: VecEncoding::F32,
11136 },
11137 );
11138 }
11139
11140 #[test]
11141 fn create_table_vector_using_sq8_succeeds() {
11142 let mut e = Engine::new();
11146 e.execute("CREATE TABLE t (v VECTOR(8) USING SQ8)").unwrap();
11147 let t = e.catalog().get("t").unwrap();
11148 assert_eq!(
11149 t.schema().columns[0].ty,
11150 DataType::Vector {
11151 dim: 8,
11152 encoding: VecEncoding::Sq8,
11153 },
11154 );
11155 }
11156
11157 #[test]
11158 fn insert_into_sq8_column_quantises_f32_payload() {
11159 let mut e = Engine::new();
11166 e.execute("CREATE TABLE t (v VECTOR(4) USING SQ8)").unwrap();
11167 e.execute("INSERT INTO t VALUES ([0.0, 0.25, 0.5, 1.0])")
11168 .unwrap();
11169 let t = e.catalog().get("t").unwrap();
11170 assert_eq!(t.rows().len(), 1);
11171 match &t.rows()[0].values[0] {
11172 Value::Sq8Vector(q) => {
11173 assert_eq!(q.bytes.len(), 4);
11174 assert!((q.min - 0.0).abs() < 1e-6);
11176 assert!((q.max - 1.0).abs() < 1e-6);
11177 }
11178 other => panic!("expected Sq8Vector cell, got {other:?}"),
11179 }
11180 }
11181
11182 #[test]
11183 fn create_table_vector_using_half_succeeds_and_insert_converts_to_f16() {
11184 let mut e = Engine::new();
11191 e.execute("CREATE TABLE t (v VECTOR(4) USING HALF)")
11192 .unwrap();
11193 e.execute("INSERT INTO t VALUES ([0.0, 0.25, 0.5, 1.0])")
11194 .unwrap();
11195 let t = e.catalog().get("t").unwrap();
11196 assert_eq!(t.rows().len(), 1);
11197 match &t.rows()[0].values[0] {
11198 Value::HalfVector(h) => {
11199 assert_eq!(h.dim(), 4);
11200 let back = h.to_f32_vec();
11201 let expected = alloc::vec![0.0_f32, 0.25, 0.5, 1.0];
11202 for (g, e) in back.iter().zip(expected.iter()) {
11203 assert!(
11204 (g - e).abs() < 1e-6,
11205 "{g} vs {e} should be exact on f16 grid"
11206 );
11207 }
11208 }
11209 other => panic!("expected HalfVector cell, got {other:?}"),
11210 }
11211 }
11212
11213 #[test]
11214 fn alter_index_rebuild_in_place_succeeds() {
11215 let mut e = Engine::new();
11220 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(3) NOT NULL)")
11221 .unwrap();
11222 for i in 0..8_i32 {
11223 #[allow(clippy::cast_precision_loss)]
11224 let base = (i as f32) * 0.1;
11225 e.execute(&alloc::format!(
11226 "INSERT INTO t VALUES ({i}, [{base}, {b1}, {b2}])",
11227 b1 = base + 0.01,
11228 b2 = base + 0.02,
11229 ))
11230 .unwrap();
11231 }
11232 e.execute("CREATE INDEX t_idx ON t USING hnsw (v)").unwrap();
11233 e.execute("ALTER INDEX t_idx REBUILD").unwrap();
11234 assert_eq!(
11236 e.catalog().get("t").unwrap().schema().columns[1].ty,
11237 DataType::Vector {
11238 dim: 3,
11239 encoding: VecEncoding::F32,
11240 },
11241 );
11242 }
11243
11244 #[test]
11245 fn alter_index_rebuild_with_encoding_switches_cell_type() {
11246 let mut e = Engine::new();
11251 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(4) NOT NULL)")
11252 .unwrap();
11253 e.execute("INSERT INTO t VALUES (1, [0.0, 0.25, 0.5, 1.0])")
11254 .unwrap();
11255 e.execute("CREATE INDEX t_idx ON t USING hnsw (v)").unwrap();
11256 e.execute("ALTER INDEX t_idx REBUILD WITH (encoding = SQ8)")
11257 .unwrap();
11258 let t = e.catalog().get("t").unwrap();
11259 assert_eq!(
11260 t.schema().columns[1].ty,
11261 DataType::Vector {
11262 dim: 4,
11263 encoding: VecEncoding::Sq8,
11264 },
11265 );
11266 assert!(matches!(t.rows()[0].values[1], Value::Sq8Vector(_)));
11267 }
11268
11269 #[test]
11270 fn alter_index_rebuild_unknown_index_errors() {
11271 let mut e = Engine::new();
11272 let err = e.execute("ALTER INDEX nope REBUILD").unwrap_err();
11273 assert!(
11274 matches!(
11275 &err,
11276 EngineError::Storage(StorageError::IndexNotFound { name }) if name == "nope"
11277 ),
11278 "got: {err}"
11279 );
11280 }
11281
11282 #[test]
11283 fn alter_index_rebuild_on_btree_index_errors() {
11284 let mut e = Engine::new();
11287 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
11288 e.execute("INSERT INTO t VALUES (1)").unwrap();
11289 e.execute("CREATE INDEX t_idx ON t (id)").unwrap();
11290 let err = e.execute("ALTER INDEX t_idx REBUILD").unwrap_err();
11291 assert!(
11292 matches!(&err, EngineError::Storage(StorageError::Unsupported(_))),
11293 "got: {err}"
11294 );
11295 }
11296
11297 #[test]
11298 fn prepared_insert_substitutes_placeholders() {
11299 let mut e = Engine::new();
11305 e.execute("CREATE TABLE t (id INT NOT NULL, name TEXT NOT NULL)")
11306 .unwrap();
11307 let stmt = e.prepare("INSERT INTO t VALUES ($1, $2)").unwrap();
11308 for (id, name) in [(1, "alice"), (2, "bob"), (3, "carol")] {
11309 e.execute_prepared(stmt.clone(), &[Value::Int(id), Value::Text(name.into())])
11310 .unwrap();
11311 }
11312 let rows_result = e.execute("SELECT id, name FROM t").unwrap();
11314 let QueryResult::Rows { rows, .. } = rows_result else {
11315 panic!("expected Rows")
11316 };
11317 assert_eq!(rows.len(), 3);
11318 }
11319
11320 #[test]
11321 fn prepared_select_with_placeholder_filters_rows() {
11322 let mut e = Engine::new();
11323 e.execute("CREATE TABLE t (id INT NOT NULL, v INT NOT NULL)")
11324 .unwrap();
11325 for i in 0..10_i32 {
11326 e.execute(&alloc::format!("INSERT INTO t VALUES ({i}, {})", i * 7))
11327 .unwrap();
11328 }
11329 let stmt = e.prepare("SELECT id FROM t WHERE v = $1").unwrap();
11330 let QueryResult::Rows { rows, .. } = e.execute_prepared(stmt, &[Value::Int(35)]).unwrap()
11331 else {
11332 panic!("expected Rows")
11333 };
11334 assert_eq!(rows.len(), 1);
11336 assert_eq!(rows[0].values[0], Value::Int(5));
11337 }
11338
11339 #[test]
11340 fn prepared_too_few_params_errors() {
11341 let mut e = Engine::new();
11342 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
11343 let stmt = e.prepare("INSERT INTO t VALUES ($1)").unwrap();
11344 let err = e.execute_prepared(stmt, &[]).unwrap_err();
11345 assert!(
11346 matches!(
11347 &err,
11348 EngineError::Eval(EvalError::PlaceholderOutOfRange { n: 1, bound: 0 })
11349 ),
11350 "got: {err}"
11351 );
11352 }
11353
11354 #[test]
11355 fn insert_into_half_column_dim_mismatch_errors() {
11356 let mut e = Engine::new();
11357 e.execute("CREATE TABLE t (v VECTOR(4) USING HALF)")
11358 .unwrap();
11359 let err = e.execute("INSERT INTO t VALUES ([1.0, 2.0])").unwrap_err();
11360 assert!(matches!(
11361 &err,
11362 EngineError::Storage(StorageError::TypeMismatch { .. })
11363 ));
11364 }
11365
11366 #[test]
11367 fn insert_into_sq8_column_dim_mismatch_errors() {
11368 let mut e = Engine::new();
11373 e.execute("CREATE TABLE t (v VECTOR(4) USING SQ8)").unwrap();
11374 let err = e.execute("INSERT INTO t VALUES ([1.0, 2.0])").unwrap_err();
11375 assert!(
11376 matches!(
11377 &err,
11378 EngineError::Storage(StorageError::TypeMismatch { .. })
11379 ),
11380 "got: {err}",
11381 );
11382 }
11383
11384 #[test]
11385 fn create_table_duplicate_errors() {
11386 let mut e = Engine::new();
11387 e.execute("CREATE TABLE foo (a INT)").unwrap();
11388 let err = e.execute("CREATE TABLE foo (a INT)").unwrap_err();
11389 assert!(matches!(
11390 err,
11391 EngineError::Storage(StorageError::DuplicateTable { ref name }) if name == "foo"
11392 ));
11393 }
11394
11395 #[test]
11396 fn insert_into_unknown_table_errors() {
11397 let mut e = Engine::new();
11398 let err = e.execute("INSERT INTO ghost VALUES (1)").unwrap_err();
11399 assert!(matches!(
11400 err,
11401 EngineError::Storage(StorageError::TableNotFound { ref name }) if name == "ghost"
11402 ));
11403 }
11404
11405 #[test]
11406 fn insert_happy_path_reports_one_affected() {
11407 let mut e = Engine::new();
11408 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
11409 let r = e.execute("INSERT INTO foo VALUES (42)").unwrap();
11410 assert_eq!(unwrap_command_ok(&r), 1);
11411 assert_eq!(e.catalog().get("foo").unwrap().row_count(), 1);
11412 }
11413
11414 #[test]
11415 fn insert_arity_mismatch_propagates() {
11416 let mut e = Engine::new();
11417 e.execute("CREATE TABLE foo (a INT, b TEXT)").unwrap();
11418 let err = e.execute("INSERT INTO foo VALUES (1)").unwrap_err();
11419 assert!(matches!(
11420 err,
11421 EngineError::Storage(StorageError::ArityMismatch { .. })
11422 ));
11423 }
11424
11425 #[test]
11426 fn insert_negative_integer_via_unary_minus() {
11427 let mut e = Engine::new();
11428 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
11429 e.execute("INSERT INTO foo VALUES (-7)").unwrap();
11430 let rows = e.catalog().get("foo").unwrap().rows();
11431 assert_eq!(rows[0].values[0], Value::Int(-7));
11432 }
11433
11434 #[test]
11435 fn insert_non_literal_expr_unsupported() {
11436 let mut e = Engine::new();
11437 e.execute("CREATE TABLE foo (a INT NOT NULL)").unwrap();
11438 let err = e.execute("INSERT INTO foo VALUES (1 + 2)").unwrap_err();
11439 assert!(matches!(err, EngineError::Unsupported(_)));
11440 }
11441
11442 #[test]
11443 fn select_star_returns_all_rows_in_insertion_order() {
11444 let mut e = Engine::new();
11445 e.execute("CREATE TABLE foo (a INT NOT NULL, b TEXT NOT NULL)")
11446 .unwrap();
11447 e.execute("INSERT INTO foo VALUES (1, 'one')").unwrap();
11448 e.execute("INSERT INTO foo VALUES (2, 'two')").unwrap();
11449 e.execute("INSERT INTO foo VALUES (3, 'three')").unwrap();
11450
11451 let r = e.execute("SELECT * FROM foo").unwrap();
11452 let QueryResult::Rows { columns, rows } = r else {
11453 panic!("expected Rows")
11454 };
11455 assert_eq!(columns.len(), 2);
11456 assert_eq!(columns[0].name, "a");
11457 assert_eq!(rows.len(), 3);
11458 assert_eq!(
11459 rows[1].values,
11460 vec![Value::Int(2), Value::Text("two".into())]
11461 );
11462 }
11463
11464 #[test]
11465 fn select_star_on_empty_table_returns_zero_rows() {
11466 let mut e = Engine::new();
11467 e.execute("CREATE TABLE foo (a INT)").unwrap();
11468 let r = e.execute("SELECT * FROM foo").unwrap();
11469 match r {
11470 QueryResult::Rows { rows, .. } => assert!(rows.is_empty()),
11471 QueryResult::CommandOk { .. } => panic!("expected Rows"),
11472 }
11473 }
11474
11475 fn make_three_row_users(e: &mut Engine) {
11478 e.execute("CREATE TABLE users (id INT NOT NULL, name TEXT NOT NULL, score INT)")
11479 .unwrap();
11480 e.execute("INSERT INTO users VALUES (1, 'alice', 90)")
11481 .unwrap();
11482 e.execute("INSERT INTO users VALUES (2, 'bob', NULL)")
11483 .unwrap();
11484 e.execute("INSERT INTO users VALUES (3, 'cara', 70)")
11485 .unwrap();
11486 }
11487
11488 fn unwrap_rows(r: QueryResult) -> (Vec<ColumnSchema>, Vec<Row>) {
11489 match r {
11490 QueryResult::Rows { columns, rows } => (columns, rows),
11491 QueryResult::CommandOk { .. } => panic!("expected Rows"),
11492 }
11493 }
11494
11495 #[test]
11496 fn where_filter_passes_only_true_rows() {
11497 let mut e = Engine::new();
11498 make_three_row_users(&mut e);
11499 let r = e.execute("SELECT * FROM users WHERE id > 1").unwrap();
11500 let (_, rows) = unwrap_rows(r);
11501 assert_eq!(rows.len(), 2);
11502 assert_eq!(rows[0].values[0], Value::Int(2));
11503 assert_eq!(rows[1].values[0], Value::Int(3));
11504 }
11505
11506 #[test]
11507 fn where_with_null_result_filters_out_row() {
11508 let mut e = Engine::new();
11509 make_three_row_users(&mut e);
11510 let r = e.execute("SELECT * FROM users WHERE score > 80").unwrap();
11512 let (_, rows) = unwrap_rows(r);
11513 assert_eq!(rows.len(), 1);
11514 assert_eq!(rows[0].values[1], Value::Text("alice".into()));
11515 }
11516
11517 #[test]
11518 fn projection_named_columns() {
11519 let mut e = Engine::new();
11520 make_three_row_users(&mut e);
11521 let r = e.execute("SELECT name, score FROM users").unwrap();
11522 let (cols, rows) = unwrap_rows(r);
11523 assert_eq!(cols.len(), 2);
11524 assert_eq!(cols[0].name, "name");
11525 assert_eq!(cols[1].name, "score");
11526 assert_eq!(rows.len(), 3);
11527 assert_eq!(
11528 rows[0].values,
11529 vec![Value::Text("alice".into()), Value::Int(90)]
11530 );
11531 }
11532
11533 #[test]
11534 fn projection_with_column_alias() {
11535 let mut e = Engine::new();
11536 make_three_row_users(&mut e);
11537 let r = e
11538 .execute("SELECT name AS who FROM users WHERE id = 1")
11539 .unwrap();
11540 let (cols, rows) = unwrap_rows(r);
11541 assert_eq!(cols[0].name, "who");
11542 assert_eq!(rows.len(), 1);
11543 assert_eq!(rows[0].values[0], Value::Text("alice".into()));
11544 }
11545
11546 #[test]
11547 fn qualified_column_with_table_alias_resolves() {
11548 let mut e = Engine::new();
11549 make_three_row_users(&mut e);
11550 let r = e
11551 .execute("SELECT u.id, u.name FROM users AS u WHERE u.id < 3")
11552 .unwrap();
11553 let (cols, rows) = unwrap_rows(r);
11554 assert_eq!(cols.len(), 2);
11555 assert_eq!(rows.len(), 2);
11556 }
11557
11558 #[test]
11559 fn qualified_column_with_wrong_alias_errors() {
11560 let mut e = Engine::new();
11561 make_three_row_users(&mut e);
11562 let err = e.execute("SELECT x.id FROM users AS u").unwrap_err();
11563 assert!(matches!(
11564 err,
11565 EngineError::Eval(EvalError::UnknownQualifier { ref qualifier }) if qualifier == "x"
11566 ));
11567 }
11568
11569 #[test]
11570 fn select_unknown_column_errors_in_projection() {
11571 let mut e = Engine::new();
11572 make_three_row_users(&mut e);
11573 let err = e.execute("SELECT ghost FROM users").unwrap_err();
11574 assert!(matches!(
11575 err,
11576 EngineError::Eval(EvalError::ColumnNotFound { ref name }) if name == "ghost"
11577 ));
11578 }
11579
11580 #[test]
11581 fn where_unknown_column_errors() {
11582 let mut e = Engine::new();
11583 make_three_row_users(&mut e);
11584 let err = e
11585 .execute("SELECT * FROM users WHERE ghost = 1")
11586 .unwrap_err();
11587 assert!(matches!(
11588 err,
11589 EngineError::Eval(EvalError::ColumnNotFound { .. })
11590 ));
11591 }
11592
11593 #[test]
11594 fn expression_projection_evaluates_and_renders() {
11595 let mut e = Engine::new();
11598 e.execute("CREATE TABLE t (a INT NOT NULL)").unwrap();
11599 e.execute("INSERT INTO t VALUES (3)").unwrap();
11600 let (_, rows) = unwrap_rows(e.execute("SELECT 1 + 2 FROM t").unwrap());
11601 assert_eq!(rows.len(), 1);
11602 assert_eq!(rows[0].values[0], Value::Int(3));
11605 }
11606
11607 #[test]
11608 fn select_unknown_table_errors() {
11609 let mut e = Engine::new();
11610 let err = e.execute("SELECT * FROM ghost").unwrap_err();
11611 assert!(matches!(
11612 err,
11613 EngineError::Storage(StorageError::TableNotFound { .. })
11614 ));
11615 }
11616
11617 #[test]
11618 fn invalid_sql_returns_parse_error() {
11619 let mut e = Engine::new();
11622 let err = e.execute("THIS_IS_NOT_A_KEYWORD foo bar baz").unwrap_err();
11623 assert!(matches!(err, EngineError::Parse(_)));
11624 }
11625
11626 #[test]
11629 fn create_index_registers_on_table() {
11630 let mut e = Engine::new();
11631 make_three_row_users(&mut e);
11632 e.execute("CREATE INDEX by_name ON users (name)").unwrap();
11633 let t = e.catalog().get("users").unwrap();
11634 assert_eq!(t.indices().len(), 1);
11635 assert_eq!(t.indices()[0].name, "by_name");
11636 }
11637
11638 #[test]
11639 fn create_index_on_unknown_table_errors() {
11640 let mut e = Engine::new();
11641 let err = e.execute("CREATE INDEX i ON ghost (a)").unwrap_err();
11642 assert!(matches!(
11643 err,
11644 EngineError::Storage(StorageError::TableNotFound { .. })
11645 ));
11646 }
11647
11648 #[test]
11649 fn create_index_on_unknown_column_errors() {
11650 let mut e = Engine::new();
11651 make_three_row_users(&mut e);
11652 let err = e.execute("CREATE INDEX i ON users (ghost)").unwrap_err();
11653 assert!(matches!(
11654 err,
11655 EngineError::Storage(StorageError::ColumnNotFound { .. })
11656 ));
11657 }
11658
11659 #[test]
11660 fn select_eq_uses_index_returns_same_rows_as_scan() {
11661 let mut without = Engine::new();
11665 make_three_row_users(&mut without);
11666 let mut with = Engine::new();
11667 make_three_row_users(&mut with);
11668 with.execute("CREATE INDEX by_id ON users (id)").unwrap();
11669
11670 let q = "SELECT * FROM users WHERE id = 2";
11671 let (_, no_idx_rows) = unwrap_rows(without.execute(q).unwrap());
11672 let (_, idx_rows) = unwrap_rows(with.execute(q).unwrap());
11673 assert_eq!(no_idx_rows, idx_rows);
11674 assert_eq!(idx_rows.len(), 1);
11675 }
11676
11677 #[test]
11678 fn select_eq_with_no_matching_index_value_returns_empty() {
11679 let mut e = Engine::new();
11680 make_three_row_users(&mut e);
11681 e.execute("CREATE INDEX by_id ON users (id)").unwrap();
11682 let (_, rows) = unwrap_rows(e.execute("SELECT * FROM users WHERE id = 999").unwrap());
11683 assert_eq!(rows.len(), 0);
11684 }
11685
11686 #[test]
11689 fn begin_sets_in_transaction_flag() {
11690 let mut e = Engine::new();
11691 assert!(!e.in_transaction());
11692 e.execute("BEGIN").unwrap();
11693 assert!(e.in_transaction());
11694 }
11695
11696 #[test]
11697 fn double_begin_errors() {
11698 let mut e = Engine::new();
11699 e.execute("BEGIN").unwrap();
11700 let err = e.execute("BEGIN").unwrap_err();
11701 assert_eq!(err, EngineError::TransactionAlreadyOpen);
11702 }
11703
11704 #[test]
11705 fn commit_without_begin_errors() {
11706 let mut e = Engine::new();
11707 let err = e.execute("COMMIT").unwrap_err();
11708 assert_eq!(err, EngineError::NoActiveTransaction);
11709 }
11710
11711 #[test]
11712 fn rollback_without_begin_errors() {
11713 let mut e = Engine::new();
11714 let err = e.execute("ROLLBACK").unwrap_err();
11715 assert_eq!(err, EngineError::NoActiveTransaction);
11716 }
11717
11718 #[test]
11719 fn commit_applies_shadow_to_committed_catalog() {
11720 let mut e = Engine::new();
11721 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
11722 e.execute("BEGIN").unwrap();
11723 e.execute("INSERT INTO t VALUES (1)").unwrap();
11724 e.execute("INSERT INTO t VALUES (2)").unwrap();
11725 e.execute("COMMIT").unwrap();
11726 assert!(!e.in_transaction());
11727 assert_eq!(e.catalog().get("t").unwrap().row_count(), 2);
11728 }
11729
11730 #[test]
11731 fn rollback_discards_shadow() {
11732 let mut e = Engine::new();
11733 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
11734 e.execute("BEGIN").unwrap();
11735 e.execute("INSERT INTO t VALUES (1)").unwrap();
11736 e.execute("INSERT INTO t VALUES (2)").unwrap();
11737 e.execute("ROLLBACK").unwrap();
11738 assert!(!e.in_transaction());
11739 assert_eq!(e.catalog().get("t").unwrap().row_count(), 0);
11740 }
11741
11742 #[test]
11743 fn select_during_tx_sees_uncommitted_writes_own_session() {
11744 let mut e = Engine::new();
11747 e.execute("CREATE TABLE t (v INT NOT NULL)").unwrap();
11748 e.execute("BEGIN").unwrap();
11749 e.execute("INSERT INTO t VALUES (42)").unwrap();
11750 let (_, rows) = unwrap_rows(e.execute("SELECT * FROM t").unwrap());
11751 assert_eq!(rows.len(), 1);
11752 assert_eq!(rows[0].values[0], Value::Int(42));
11753 }
11754
11755 #[test]
11756 fn snapshot_with_no_users_is_bare_catalog_format() {
11757 let mut e = Engine::new();
11758 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
11759 let bytes = e.snapshot();
11760 assert_eq!(
11761 &bytes[..8],
11762 b"SPGDB001",
11763 "must be the bare v3.x catalog magic"
11764 );
11765 let e2 = Engine::restore_envelope(&bytes).unwrap();
11766 assert!(e2.users().is_empty());
11767 assert_eq!(e2.catalog().table_count(), 1);
11768 }
11769
11770 #[test]
11771 fn snapshot_with_users_round_trips_both_via_envelope() {
11772 let mut e = Engine::new();
11773 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
11774 e.create_user("alice", "pw1", Role::Admin, [9; 16]).unwrap();
11775 e.create_user("bob", "pw2", Role::ReadOnly, [5; 16])
11776 .unwrap();
11777 let bytes = e.snapshot();
11778 assert_eq!(&bytes[..8], b"SPGENV01", "must be the v4.1 envelope magic");
11779 let e2 = Engine::restore_envelope(&bytes).unwrap();
11780 assert_eq!(e2.users().len(), 2);
11781 assert_eq!(e2.verify_user("alice", "pw1"), Some(Role::Admin));
11782 assert_eq!(e2.verify_user("bob", "pw2"), Some(Role::ReadOnly));
11783 assert_eq!(e2.verify_user("alice", "wrong"), None);
11784 assert_eq!(e2.catalog().table_count(), 1);
11785 }
11786
11787 #[test]
11788 fn ddl_inside_tx_also_rolled_back() {
11789 let mut e = Engine::new();
11790 e.execute("BEGIN").unwrap();
11791 e.execute("CREATE TABLE t (v INT)").unwrap();
11792 e.execute("SELECT * FROM t").unwrap();
11794 e.execute("ROLLBACK").unwrap();
11795 let err = e.execute("SELECT * FROM t").unwrap_err();
11797 assert!(matches!(
11798 err,
11799 EngineError::Storage(StorageError::TableNotFound { .. })
11800 ));
11801 }
11802
11803 #[test]
11806 fn create_publication_lands_in_catalog() {
11807 let mut e = Engine::new();
11808 assert!(e.publications().is_empty());
11809 e.execute("CREATE PUBLICATION pub_a").unwrap();
11810 assert_eq!(e.publications().len(), 1);
11811 assert!(e.publications().contains("pub_a"));
11812 }
11813
11814 #[test]
11815 fn create_publication_duplicate_errors() {
11816 let mut e = Engine::new();
11817 e.execute("CREATE PUBLICATION pub_a").unwrap();
11818 let err = e.execute("CREATE PUBLICATION pub_a").unwrap_err();
11819 assert!(
11820 alloc::format!("{err:?}").contains("DuplicateName"),
11821 "got {err:?}"
11822 );
11823 }
11824
11825 #[test]
11826 fn drop_publication_silent_when_absent() {
11827 let mut e = Engine::new();
11828 let r = e.execute("DROP PUBLICATION nope").unwrap();
11831 match r {
11832 QueryResult::CommandOk { affected, .. } => assert_eq!(affected, 0),
11833 other => panic!("expected CommandOk, got {other:?}"),
11834 }
11835 }
11836
11837 #[test]
11838 fn drop_publication_present_reports_one_affected() {
11839 let mut e = Engine::new();
11840 e.execute("CREATE PUBLICATION pub_a").unwrap();
11841 let r = e.execute("DROP PUBLICATION pub_a").unwrap();
11842 match r {
11843 QueryResult::CommandOk {
11844 affected,
11845 modified_catalog,
11846 } => {
11847 assert_eq!(affected, 1);
11848 assert!(modified_catalog);
11849 }
11850 other => panic!("expected CommandOk, got {other:?}"),
11851 }
11852 assert!(e.publications().is_empty());
11853 }
11854
11855 #[test]
11856 fn publications_persist_across_snapshot_restore() {
11857 let mut e = Engine::new();
11862 e.execute("CREATE PUBLICATION pub_a").unwrap();
11863 e.execute("CREATE PUBLICATION pub_b FOR ALL TABLES")
11864 .unwrap();
11865 let snap = e.snapshot();
11866 let e2 = Engine::restore_envelope(&snap).unwrap();
11867 assert_eq!(e2.publications().len(), 2);
11868 assert!(e2.publications().contains("pub_a"));
11869 assert!(e2.publications().contains("pub_b"));
11870 }
11871
11872 #[test]
11873 fn create_publication_allowed_inside_transaction() {
11874 let mut e = Engine::new();
11878 e.execute("BEGIN").unwrap();
11879 e.execute("CREATE PUBLICATION pub_a").unwrap();
11880 e.execute("COMMIT").unwrap();
11881 assert!(e.publications().contains("pub_a"));
11882 }
11883
11884 #[test]
11887 fn create_publication_for_table_list_lands_with_scope() {
11888 let mut e = Engine::new();
11889 e.execute("CREATE TABLE t1 (id INT NOT NULL)").unwrap();
11890 e.execute("CREATE TABLE t2 (id INT NOT NULL)").unwrap();
11891 e.execute("CREATE PUBLICATION pub_a FOR TABLE t1, t2")
11892 .unwrap();
11893 let scope = e.publications().get("pub_a").cloned();
11894 let Some(spg_sql::ast::PublicationScope::ForTables(ts)) = scope else {
11895 panic!("expected ForTables scope, got {scope:?}")
11896 };
11897 assert_eq!(ts, alloc::vec!["t1".to_string(), "t2".to_string()]);
11898 }
11899
11900 #[test]
11901 fn create_publication_all_tables_except_lands_with_scope() {
11902 let mut e = Engine::new();
11903 e.execute("CREATE PUBLICATION pub_a FOR ALL TABLES EXCEPT t3")
11904 .unwrap();
11905 let scope = e.publications().get("pub_a").cloned();
11906 let Some(spg_sql::ast::PublicationScope::AllTablesExcept(ts)) = scope else {
11907 panic!("expected AllTablesExcept scope, got {scope:?}")
11908 };
11909 assert_eq!(ts, alloc::vec!["t3".to_string()]);
11910 }
11911
11912 #[test]
11913 fn show_publications_empty_returns_zero_rows() {
11914 let e = Engine::new();
11915 let r = e.execute_readonly("SHOW PUBLICATIONS").unwrap();
11916 let QueryResult::Rows { rows, columns } = r else {
11917 panic!()
11918 };
11919 assert!(rows.is_empty());
11920 assert_eq!(columns.len(), 3);
11921 assert_eq!(columns[0].name, "name");
11922 assert_eq!(columns[1].name, "scope");
11923 assert_eq!(columns[2].name, "table_count");
11924 }
11925
11926 #[test]
11927 fn show_publications_returns_one_row_per_publication_ordered_by_name() {
11928 let mut e = Engine::new();
11929 e.execute("CREATE PUBLICATION z_pub").unwrap();
11930 e.execute("CREATE PUBLICATION a_pub FOR TABLE t1, t2")
11931 .unwrap();
11932 e.execute("CREATE PUBLICATION m_pub FOR ALL TABLES EXCEPT bad")
11933 .unwrap();
11934 let r = e.execute_readonly("SHOW PUBLICATIONS").unwrap();
11935 let QueryResult::Rows { rows, .. } = r else {
11936 panic!()
11937 };
11938 assert_eq!(rows.len(), 3);
11939 let names: Vec<&str> = rows
11941 .iter()
11942 .map(|r| {
11943 if let Value::Text(s) = &r.values[0] {
11944 s.as_str()
11945 } else {
11946 panic!()
11947 }
11948 })
11949 .collect();
11950 assert_eq!(names, alloc::vec!["a_pub", "m_pub", "z_pub"]);
11951 match &rows[0].values[1] {
11953 Value::Text(s) => assert_eq!(s, "FOR TABLE t1, t2"),
11954 other => panic!("expected Text, got {other:?}"),
11955 }
11956 assert_eq!(rows[0].values[2], Value::Int(2));
11957 match &rows[1].values[1] {
11959 Value::Text(s) => assert_eq!(s, "FOR ALL TABLES EXCEPT bad"),
11960 other => panic!("expected Text, got {other:?}"),
11961 }
11962 assert_eq!(rows[1].values[2], Value::Int(1));
11963 match &rows[2].values[1] {
11965 Value::Text(s) => assert_eq!(s, "FOR ALL TABLES"),
11966 other => panic!("expected Text, got {other:?}"),
11967 }
11968 assert_eq!(rows[2].values[2], Value::Null);
11969 }
11970
11971 #[test]
11972 fn for_list_scopes_persist_across_snapshot() {
11973 let mut e = Engine::new();
11976 e.execute("CREATE PUBLICATION p1 FOR TABLE t1, t2").unwrap();
11977 e.execute("CREATE PUBLICATION p2 FOR ALL TABLES EXCEPT bad, worse")
11978 .unwrap();
11979 let snap = e.snapshot();
11980 let e2 = Engine::restore_envelope(&snap).unwrap();
11981 assert_eq!(e2.publications().len(), 2);
11982 let p1 = e2.publications().get("p1").cloned();
11983 let Some(spg_sql::ast::PublicationScope::ForTables(ts)) = p1 else {
11984 panic!("p1 scope lost: {p1:?}")
11985 };
11986 assert_eq!(ts, alloc::vec!["t1".to_string(), "t2".to_string()]);
11987 let p2 = e2.publications().get("p2").cloned();
11988 let Some(spg_sql::ast::PublicationScope::AllTablesExcept(ts)) = p2 else {
11989 panic!("p2 scope lost: {p2:?}")
11990 };
11991 assert_eq!(ts, alloc::vec!["bad".to_string(), "worse".to_string()]);
11992 }
11993
11994 #[test]
11997 fn create_subscription_lands_in_catalog_with_defaults() {
11998 let mut e = Engine::new();
11999 e.execute(
12000 "CREATE SUBSCRIPTION sub_a CONNECTION 'host=127.0.0.1 port=20002' PUBLICATION pub_a",
12001 )
12002 .unwrap();
12003 let s = e.subscriptions().get("sub_a").cloned().expect("present");
12004 assert_eq!(s.conn_str, "host=127.0.0.1 port=20002");
12005 assert_eq!(s.publications, alloc::vec!["pub_a".to_string()]);
12006 assert!(s.enabled);
12007 assert_eq!(s.last_received_pos, 0);
12008 }
12009
12010 #[test]
12011 fn create_subscription_duplicate_name_errors() {
12012 let mut e = Engine::new();
12013 e.execute("CREATE SUBSCRIPTION s CONNECTION 'host=x' PUBLICATION p")
12014 .unwrap();
12015 let err = e
12016 .execute("CREATE SUBSCRIPTION s CONNECTION 'host=y' PUBLICATION p")
12017 .unwrap_err();
12018 assert!(
12019 alloc::format!("{err:?}").contains("DuplicateName"),
12020 "got {err:?}"
12021 );
12022 }
12023
12024 #[test]
12025 fn drop_subscription_silent_when_absent() {
12026 let mut e = Engine::new();
12027 let r = e.execute("DROP SUBSCRIPTION never").unwrap();
12028 match r {
12029 QueryResult::CommandOk { affected, .. } => assert_eq!(affected, 0),
12030 other => panic!("expected CommandOk, got {other:?}"),
12031 }
12032 }
12033
12034 #[test]
12035 fn subscription_advance_updates_last_pos_monotone() {
12036 let mut e = Engine::new();
12037 e.execute("CREATE SUBSCRIPTION s CONNECTION 'h=x' PUBLICATION p")
12038 .unwrap();
12039 assert!(e.subscription_advance("s", 100));
12040 assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 100);
12041 assert!(e.subscription_advance("s", 50)); assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 100);
12043 assert!(e.subscription_advance("s", 200));
12044 assert_eq!(e.subscriptions().get("s").unwrap().last_received_pos, 200);
12045 assert!(!e.subscription_advance("missing", 1));
12046 }
12047
12048 #[test]
12049 fn show_subscriptions_returns_rows_ordered_by_name() {
12050 let mut e = Engine::new();
12051 e.execute("CREATE SUBSCRIPTION z_sub CONNECTION 'h=x' PUBLICATION p1, p2")
12052 .unwrap();
12053 e.execute("CREATE SUBSCRIPTION a_sub CONNECTION 'h=y' PUBLICATION p3")
12054 .unwrap();
12055 let r = e.execute_readonly("SHOW SUBSCRIPTIONS").unwrap();
12056 let QueryResult::Rows { rows, columns } = r else {
12057 panic!()
12058 };
12059 assert_eq!(rows.len(), 2);
12060 assert_eq!(columns.len(), 5);
12061 assert_eq!(columns[0].name, "name");
12062 assert_eq!(columns[4].name, "last_received_pos");
12063 let names: Vec<&str> = rows
12065 .iter()
12066 .map(|r| {
12067 if let Value::Text(s) = &r.values[0] {
12068 s.as_str()
12069 } else {
12070 panic!()
12071 }
12072 })
12073 .collect();
12074 assert_eq!(names, alloc::vec!["a_sub", "z_sub"]);
12075 assert_eq!(rows[0].values[1], Value::Text("h=y".to_string()));
12077 assert_eq!(rows[0].values[2], Value::Text("p3".to_string()));
12078 assert_eq!(rows[0].values[3], Value::Bool(true));
12079 assert_eq!(rows[0].values[4], Value::BigInt(0));
12080 assert_eq!(rows[1].values[2], Value::Text("p1, p2".to_string()));
12082 }
12083
12084 #[test]
12085 fn subscriptions_persist_across_snapshot_envelope_v4() {
12086 let mut e = Engine::new();
12087 e.execute("CREATE SUBSCRIPTION s1 CONNECTION 'h=A' PUBLICATION p1, p2")
12088 .unwrap();
12089 e.execute("CREATE SUBSCRIPTION s2 CONNECTION 'h=B' PUBLICATION p3")
12090 .unwrap();
12091 e.subscription_advance("s2", 42);
12092 let snap = e.snapshot();
12093 let e2 = Engine::restore_envelope(&snap).unwrap();
12094 assert_eq!(e2.subscriptions().len(), 2);
12095 let s1 = e2.subscriptions().get("s1").unwrap();
12096 assert_eq!(s1.conn_str, "h=A");
12097 assert_eq!(
12098 s1.publications,
12099 alloc::vec!["p1".to_string(), "p2".to_string()]
12100 );
12101 assert_eq!(s1.last_received_pos, 0);
12102 let s2 = e2.subscriptions().get("s2").unwrap();
12103 assert_eq!(s2.last_received_pos, 42);
12104 }
12105
12106 #[test]
12107 fn v3_envelope_loads_with_empty_subscriptions() {
12108 let mut e = Engine::new();
12112 e.execute("CREATE PUBLICATION pub_legacy").unwrap();
12113 let catalog = e.catalog.serialize();
12114 let users = crate::users::serialize_users(&e.users);
12115 let pubs = e.publications.serialize();
12116 let mut buf = Vec::new();
12117 buf.extend_from_slice(b"SPGENV01");
12118 buf.push(3u8); buf.extend_from_slice(&u32::try_from(catalog.len()).unwrap().to_le_bytes());
12120 buf.extend_from_slice(&catalog);
12121 buf.extend_from_slice(&u32::try_from(users.len()).unwrap().to_le_bytes());
12122 buf.extend_from_slice(&users);
12123 buf.extend_from_slice(&u32::try_from(pubs.len()).unwrap().to_le_bytes());
12124 buf.extend_from_slice(&pubs);
12125 let crc = spg_crypto::crc32::crc32(&buf);
12126 buf.extend_from_slice(&crc.to_le_bytes());
12127
12128 let e2 = Engine::restore_envelope(&buf).expect("v3 envelope restores under v4 reader");
12129 assert!(e2.subscriptions().is_empty());
12130 assert!(e2.publications().contains("pub_legacy"));
12131 }
12132
12133 #[test]
12134 fn create_subscription_allowed_inside_transaction() {
12135 let mut e = Engine::new();
12136 e.execute("BEGIN").unwrap();
12137 e.execute("CREATE SUBSCRIPTION s CONNECTION 'h=x' PUBLICATION p")
12138 .unwrap();
12139 e.execute("COMMIT").unwrap();
12140 assert!(e.subscriptions().contains("s"));
12141 }
12142
12143 #[test]
12145 fn analyze_populates_histogram_bounds() {
12146 let mut e = Engine::new();
12147 e.execute("CREATE TABLE t (id INT NOT NULL, name TEXT)")
12148 .unwrap();
12149 for i in 0..50 {
12150 e.execute(&alloc::format!("INSERT INTO t VALUES ({i}, 'name{i}')"))
12151 .unwrap();
12152 }
12153 e.execute("ANALYZE t").unwrap();
12154 let stats = e.statistics();
12155 let id_stats = stats.get("t", "id").unwrap();
12156 assert!(id_stats.histogram_bounds.len() >= 2);
12157 assert_eq!(id_stats.histogram_bounds.first().unwrap(), "0");
12158 assert_eq!(id_stats.histogram_bounds.last().unwrap(), "49");
12159 assert!((id_stats.null_frac - 0.0).abs() < 1e-6);
12160 assert_eq!(id_stats.n_distinct, 50);
12161 }
12162
12163 #[test]
12164 fn reanalyze_overwrites_prior_stats() {
12165 let mut e = Engine::new();
12166 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
12167 for i in 0..10 {
12168 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12169 .unwrap();
12170 }
12171 e.execute("ANALYZE t").unwrap();
12172 let n1 = e.statistics().get("t", "id").unwrap().n_distinct;
12173 assert_eq!(n1, 10);
12174 for i in 10..30 {
12175 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12176 .unwrap();
12177 }
12178 e.execute("ANALYZE t").unwrap();
12179 let n2 = e.statistics().get("t", "id").unwrap().n_distinct;
12180 assert_eq!(n2, 30);
12181 }
12182
12183 #[test]
12184 fn analyze_unknown_table_errors() {
12185 let mut e = Engine::new();
12186 let err = e.execute("ANALYZE nonexistent").unwrap_err();
12187 assert!(matches!(
12188 err,
12189 EngineError::Storage(StorageError::TableNotFound { .. })
12190 ));
12191 }
12192
12193 #[test]
12194 fn bare_analyze_covers_all_user_tables() {
12195 let mut e = Engine::new();
12196 e.execute("CREATE TABLE t1 (id INT NOT NULL)").unwrap();
12197 e.execute("CREATE TABLE t2 (name TEXT NOT NULL)").unwrap();
12198 e.execute("INSERT INTO t1 VALUES (1)").unwrap();
12199 e.execute("INSERT INTO t2 VALUES ('alice')").unwrap();
12200 let r = e.execute("ANALYZE").unwrap();
12201 match r {
12202 QueryResult::CommandOk {
12203 affected,
12204 modified_catalog,
12205 } => {
12206 assert_eq!(affected, 2);
12207 assert!(modified_catalog);
12208 }
12209 other => panic!("expected CommandOk, got {other:?}"),
12210 }
12211 assert!(e.statistics().get("t1", "id").is_some());
12212 assert!(e.statistics().get("t2", "name").is_some());
12213 }
12214
12215 #[test]
12216 fn select_from_spg_statistic_returns_rows_per_column() {
12217 let mut e = Engine::new();
12218 e.execute("CREATE TABLE t (id INT NOT NULL, label TEXT)")
12219 .unwrap();
12220 e.execute("INSERT INTO t VALUES (1, 'a')").unwrap();
12221 e.execute("INSERT INTO t VALUES (2, 'b')").unwrap();
12222 e.execute("ANALYZE t").unwrap();
12223 let r = e.execute_readonly("SELECT * FROM spg_statistic").unwrap();
12224 let QueryResult::Rows { rows, columns } = r else {
12225 panic!()
12226 };
12227 assert_eq!(columns.len(), 6);
12229 assert_eq!(columns[0].name, "table_name");
12230 assert_eq!(columns[4].name, "histogram_bounds");
12231 assert_eq!(columns[5].name, "cold_row_count");
12232 assert_eq!(rows.len(), 2, "one row per column of t");
12233 match (&rows[0].values[0], &rows[0].values[1]) {
12235 (Value::Text(t), Value::Text(c)) => {
12236 assert_eq!(t, "t");
12237 assert_eq!(c, "id");
12239 }
12240 _ => panic!(),
12241 }
12242 }
12243
12244 #[test]
12245 fn analyze_skips_vector_columns() {
12246 let mut e = Engine::new();
12249 e.execute("CREATE TABLE t (id INT NOT NULL, v VECTOR(3) NOT NULL)")
12250 .unwrap();
12251 e.execute("INSERT INTO t VALUES (1, [1, 2, 3])").unwrap();
12252 e.execute("ANALYZE t").unwrap();
12253 assert!(e.statistics().get("t", "id").is_some());
12254 assert!(e.statistics().get("t", "v").is_none());
12255 }
12256
12257 #[test]
12258 fn statistics_persist_across_envelope_v5_round_trip() {
12259 let mut e = Engine::new();
12260 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
12261 for i in 0..20 {
12262 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12263 .unwrap();
12264 }
12265 e.execute("ANALYZE").unwrap();
12266 let snap = e.snapshot();
12267 let e2 = Engine::restore_envelope(&snap).unwrap();
12268 let s = e2.statistics().get("t", "id").unwrap();
12269 assert_eq!(s.n_distinct, 20);
12270 }
12271
12272 #[test]
12275 fn auto_analyze_threshold_fires_after_10pct_of_min_rows_on_small_table() {
12276 let mut e = Engine::new();
12280 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
12281 for i in 0..9 {
12282 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12283 .unwrap();
12284 }
12285 assert!(e.tables_needing_analyze().is_empty(), "9 < threshold");
12286 e.execute("INSERT INTO t VALUES (9)").unwrap();
12287 let needs = e.tables_needing_analyze();
12288 assert_eq!(needs, alloc::vec!["t".to_string()]);
12289 }
12290
12291 #[test]
12292 fn auto_analyze_threshold_uses_10pct_of_row_count_for_large_tables() {
12293 let mut e = Engine::new();
12299 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
12300 for i in 0..1000 {
12301 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12302 .unwrap();
12303 }
12304 e.execute("ANALYZE t").unwrap();
12305 assert!(e.tables_needing_analyze().is_empty(), "fresh ANALYZE");
12306 for i in 1000..1050 {
12307 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12308 .unwrap();
12309 }
12310 assert!(
12311 e.tables_needing_analyze().is_empty(),
12312 "50 inserts < threshold of ~105"
12313 );
12314 for i in 1050..1200 {
12315 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12316 .unwrap();
12317 }
12318 assert_eq!(
12319 e.tables_needing_analyze(),
12320 alloc::vec!["t".to_string()],
12321 "200 inserts > 0.1 × 1200 threshold"
12322 );
12323 }
12324
12325 #[test]
12326 fn auto_analyze_threshold_resets_after_analyze() {
12327 let mut e = Engine::new();
12328 e.execute("CREATE TABLE t (id INT NOT NULL)").unwrap();
12329 for i in 0..200 {
12330 e.execute(&alloc::format!("INSERT INTO t VALUES ({i})"))
12331 .unwrap();
12332 }
12333 assert!(!e.tables_needing_analyze().is_empty());
12334 e.execute("ANALYZE").unwrap();
12335 assert!(
12336 e.tables_needing_analyze().is_empty(),
12337 "ANALYZE must reset the counter"
12338 );
12339 }
12340
12341 #[test]
12342 fn auto_analyze_threshold_tracks_updates_and_deletes() {
12343 let mut e = Engine::new();
12344 e.execute("CREATE TABLE t (id INT NOT NULL, label TEXT)")
12345 .unwrap();
12346 for i in 0..50 {
12347 e.execute(&alloc::format!("INSERT INTO t VALUES ({i}, 'x')"))
12348 .unwrap();
12349 }
12350 e.execute("ANALYZE t").unwrap();
12351 e.execute("UPDATE t SET label = 'y' WHERE id < 20").unwrap();
12354 e.execute("DELETE FROM t WHERE id >= 45").unwrap();
12355 assert_eq!(e.tables_needing_analyze(), alloc::vec!["t".to_string()]);
12356 }
12357
12358 #[test]
12359 fn v4_envelope_loads_with_empty_statistics() {
12360 let mut e = Engine::new();
12364 e.create_user("alice", "secret", crate::users::Role::ReadOnly, [0u8; 16])
12365 .unwrap();
12366 let catalog = e.catalog.serialize();
12367 let users = crate::users::serialize_users(&e.users);
12368 let pubs = e.publications.serialize();
12369 let subs = e.subscriptions.serialize();
12370 let mut buf = Vec::new();
12371 buf.extend_from_slice(b"SPGENV01");
12372 buf.push(4u8);
12373 buf.extend_from_slice(&u32::try_from(catalog.len()).unwrap().to_le_bytes());
12374 buf.extend_from_slice(&catalog);
12375 buf.extend_from_slice(&u32::try_from(users.len()).unwrap().to_le_bytes());
12376 buf.extend_from_slice(&users);
12377 buf.extend_from_slice(&u32::try_from(pubs.len()).unwrap().to_le_bytes());
12378 buf.extend_from_slice(&pubs);
12379 buf.extend_from_slice(&u32::try_from(subs.len()).unwrap().to_le_bytes());
12380 buf.extend_from_slice(&subs);
12381 let crc = spg_crypto::crc32::crc32(&buf);
12382 buf.extend_from_slice(&crc.to_le_bytes());
12383 let e2 = Engine::restore_envelope(&buf).expect("v4 envelope restores");
12384 assert!(e2.statistics().is_empty());
12385 }
12386
12387 #[test]
12388 fn v1_v2_envelope_loads_with_empty_publications() {
12389 let mut e = Engine::new();
12396 e.create_user("alice", "secret", crate::users::Role::ReadOnly, [0u8; 16])
12399 .unwrap();
12400
12401 let catalog = e.catalog.serialize();
12403 let users = crate::users::serialize_users(&e.users);
12404 let mut buf = Vec::new();
12405 buf.extend_from_slice(b"SPGENV01");
12406 buf.push(2u8); buf.extend_from_slice(&u32::try_from(catalog.len()).unwrap().to_le_bytes());
12408 buf.extend_from_slice(&catalog);
12409 buf.extend_from_slice(&u32::try_from(users.len()).unwrap().to_le_bytes());
12410 buf.extend_from_slice(&users);
12411 let crc = spg_crypto::crc32::crc32(&buf);
12412 buf.extend_from_slice(&crc.to_le_bytes());
12413
12414 let e2 = Engine::restore_envelope(&buf).expect("v2 envelope restores");
12415 assert!(e2.publications().is_empty());
12416 }
12417}