1#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
2use crate::errors::{DynoxideError, Result};
3#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
4use crate::storage_backend::clock::{Clock, SystemClock};
5#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
6use crate::storage_backend::sql_builders::{self, escape_table_name};
7use crate::types::AttributeValue;
8#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
9use rusqlite::{Connection, params};
10#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
11use std::{cell::RefCell, collections::HashMap, sync::Arc};
12
13#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
17const SCHEMA_VERSION: &str = sql_builders::SCHEMA_VERSION;
18
19pub(crate) const HASH_BUCKETS: u32 = 4096;
22
23pub fn compute_hash_prefix(pk_value: &AttributeValue) -> String {
37 let key_bytes = match pk_value {
38 AttributeValue::S(s) => s.as_bytes().to_vec(),
39 AttributeValue::N(n) => num_to_buffer(n),
40 AttributeValue::B(b) => b.clone(),
41 _ => vec![], };
43
44 let digest = md5::compute([b"Outliers" as &[u8], &key_bytes].concat());
45 format!("{:032x}", digest)[..6].to_string()
46}
47
48pub fn hash_bucket(hash_prefix: &str) -> u32 {
50 let prefix_3 = &hash_prefix[..3.min(hash_prefix.len())];
51 u32::from_str_radix(prefix_3, 16).unwrap_or(0)
52}
53
54fn num_to_buffer(num_str: &str) -> Vec<u8> {
60 let trimmed = num_str.trim();
61 if trimmed.is_empty() {
62 return vec![0x80];
63 }
64
65 use bigdecimal::BigDecimal;
66 use std::str::FromStr;
67
68 let bd = match BigDecimal::from_str(trimmed) {
69 Ok(v) => v,
70 Err(_) => return vec![0x80],
71 };
72
73 if bd.sign() == bigdecimal::num_bigint::Sign::NoSign {
74 return vec![0x80];
75 }
76
77 let is_negative = bd.sign() == bigdecimal::num_bigint::Sign::Minus;
78 let bd_abs = if is_negative { -&bd } else { bd.clone() };
79
80 let (mantissa, exponent) = extract_mantissa_and_exponent(&bd_abs);
81 if mantissa.is_empty() {
82 return vec![0x80];
83 }
84
85 let append_zero: i64 = if exponent % 2 != 0 { 1 } else { 0 };
88 let byte_len_no_exp = ((mantissa.len() as i64 + append_zero + 1) / 2) as usize;
89
90 let mut byte_array: Vec<u8>;
91 if byte_len_no_exp < 20 && is_negative {
92 byte_array = vec![0u8; byte_len_no_exp + 2];
93 byte_array[byte_len_no_exp + 1] = 102;
94 } else {
95 byte_array = vec![0u8; byte_len_no_exp + 1];
96 }
97
98 let exp_sum = exponent + append_zero;
102 let exp_byte_val = floor_div(exp_sum, 2) - 64;
103 if is_negative {
104 byte_array[0] = (exp_byte_val ^ !0i64) as u8;
107 } else {
108 byte_array[0] = exp_byte_val as u8;
109 }
110
111 let mut mi: i64 = 0; let mlen = mantissa.len() as i64;
115 let mut appended_zero = false;
116
117 while mi < mlen {
118 let bai = ((mi + append_zero) / 2 + 1) as usize; if append_zero != 0 && mi == 0 && !appended_zero {
120 byte_array[bai] = 0;
121 appended_zero = true;
122 mi -= 1; } else if (mi + append_zero) % 2 == 0 {
124 byte_array[bai] = mantissa[mi as usize] * 10;
125 } else {
126 byte_array[bai] += mantissa[mi as usize];
127 }
128
129 if ((mi + append_zero) % 2 != 0) || (mi == mlen - 1) {
131 if is_negative {
132 byte_array[bai] = 101u8.wrapping_sub(byte_array[bai]);
133 } else {
134 byte_array[bai] = byte_array[bai].wrapping_add(1);
135 }
136 }
137
138 mi += 1; }
140
141 byte_array
142}
143
144fn floor_div(a: i64, b: i64) -> i64 {
146 let d = a / b;
147 let r = a % b;
148 if (r != 0) && ((r ^ b) < 0) { d - 1 } else { d }
149}
150
151fn extract_mantissa_and_exponent(bd: &bigdecimal::BigDecimal) -> (Vec<u8>, i64) {
158 let normalized = bd.normalized();
160
161 let (bigint, scale) = normalized.as_bigint_and_exponent();
164 let digits_str = bigint.to_string();
165 let digits_str = digits_str.trim_start_matches('-');
166
167 let digits: Vec<u8> = digits_str
168 .chars()
169 .map(|c| c.to_digit(10).unwrap() as u8)
170 .collect();
171
172 let exponent = digits.len() as i64 - scale;
175
176 (digits, exponent)
177}
178
179pub fn hash_in_segment(hash_prefix: &str, segment: u32, total_segments: u32) -> bool {
186 let bucket = hash_bucket(hash_prefix);
187 let start = ceiling_div(HASH_BUCKETS * segment, total_segments);
188 let end = ceiling_div(HASH_BUCKETS * (segment + 1), total_segments) - 1;
189 bucket >= start && bucket <= end
190}
191
192pub(crate) fn ceiling_div(a: u32, b: u32) -> u32 {
193 a.div_ceil(b)
194}
195
196#[derive(Debug, Default)]
198pub struct ScanParams<'a> {
199 pub limit: Option<usize>,
200 pub exclusive_start_pk: Option<&'a str>,
201 pub exclusive_start_sk: Option<&'a str>,
202 pub segment: Option<u32>,
203 pub total_segments: Option<u32>,
204 pub exclusive_start_base_pk: Option<&'a str>,
206 pub exclusive_start_base_sk: Option<&'a str>,
208}
209
210#[derive(Debug, Default)]
212pub struct CreateTableMetadata<'a> {
213 pub table_name: &'a str,
214 pub key_schema: &'a str,
215 pub attribute_definitions: &'a str,
216 pub gsi_definitions: Option<&'a str>,
217 pub lsi_definitions: Option<&'a str>,
218 pub provisioned_throughput: Option<&'a str>,
219 pub created_at: i64,
220 pub sse_specification: Option<&'a str>,
221 pub table_class: Option<&'a str>,
222 pub deletion_protection_enabled: bool,
223 pub billing_mode: Option<&'a str>,
224 pub on_demand_throughput: Option<&'a str>,
225 pub vector_index_definitions: Option<&'a str>,
226}
227
228#[derive(Debug, Default)]
230pub struct QueryParams<'a> {
231 pub sk_condition: Option<&'a str>,
232 pub sk_params: &'a [&'a str],
233 pub forward: bool,
234 pub limit: Option<usize>,
235 pub exclusive_start_sk: Option<&'a str>,
236 pub exclusive_start_base_pk: Option<&'a str>,
238 pub exclusive_start_base_sk: Option<&'a str>,
240}
241
242#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
250pub struct Storage {
251 conn: Connection,
252 metadata_cache: RefCell<HashMap<String, TableMetadata>>,
255 clock: Arc<dyn Clock>,
257 #[cfg(test)]
261 gsi_insert_fault: RefCell<Option<u32>>,
262}
263
264#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
265impl Storage {
266 pub fn new(path: &str) -> Result<Self> {
268 let conn = Connection::open(path)?;
269 let mut storage = Self {
270 conn,
271 metadata_cache: RefCell::new(HashMap::new()),
272 #[cfg(test)]
273 gsi_insert_fault: RefCell::new(None),
274 clock: Arc::new(SystemClock),
275 };
276 storage.initialize().map_err(Self::maybe_encrypted_error)?;
277 Ok(storage)
278 }
279
280 #[cfg(test)]
285 pub(crate) fn fail_gsi_insert_after(&self, successes: u32) {
286 *self.gsi_insert_fault.borrow_mut() = Some(successes);
287 }
288
289 #[cfg(test)]
291 pub(crate) fn check_gsi_insert_fault(&self) -> Result<()> {
292 let mut fault = self.gsi_insert_fault.borrow_mut();
293 match fault.as_mut() {
294 Some(0) => Err(crate::errors::DynoxideError::InternalServerError(
295 "injected GSI insert failure".to_string(),
296 )),
297 Some(remaining) => {
298 *remaining -= 1;
299 Ok(())
300 }
301 None => Ok(()),
302 }
303 }
304
305 pub fn with_clock(mut self, clock: Arc<dyn Clock>) -> Self {
311 self.clock = clock;
312 self
313 }
314
315 pub(crate) fn clock(&self) -> &dyn Clock {
317 self.clock.as_ref()
318 }
319
320 fn maybe_encrypted_error(err: DynoxideError) -> DynoxideError {
323 if let DynoxideError::SqliteError(ref sqlite_err) = err
324 && let Some(rusqlite::ErrorCode::NotADatabase) = sqlite_err.sqlite_error_code()
325 {
326 return DynoxideError::InternalServerError(
327 "Database file is encrypted or not a valid SQLite database. \
328 If encrypted, enable the `encryption` or `encryption-cc` feature \
329 and use Database::new_encrypted() with the correct key."
330 .to_string(),
331 );
332 }
333 err
334 }
335
336 #[cfg(feature = "_has-encryption")]
342 pub fn new_encrypted(path: &str, key: &str) -> Result<Self> {
343 use zeroize::Zeroize;
344
345 let conn = Connection::open(path)?;
346 let mut pragma_val = format!("x'{key}'");
351 conn.pragma_update(None, "key", &pragma_val)?;
352 pragma_val.zeroize();
353 conn.execute_batch("SELECT count(*) FROM sqlite_master;")?;
355 let mut storage = Self {
356 conn,
357 metadata_cache: RefCell::new(HashMap::new()),
358 #[cfg(test)]
359 gsi_insert_fault: RefCell::new(None),
360 clock: Arc::new(SystemClock),
361 };
362 storage.initialize()?;
363 Ok(storage)
364 }
365
366 pub fn memory() -> Result<Self> {
368 let conn = Connection::open_in_memory()?;
369 let mut storage = Self {
370 conn,
371 metadata_cache: RefCell::new(HashMap::new()),
372 #[cfg(test)]
373 gsi_insert_fault: RefCell::new(None),
374 clock: Arc::new(SystemClock),
375 };
376 storage.initialize()?;
377 Ok(storage)
378 }
379
380 fn initialize(&mut self) -> Result<()> {
382 self.conn.pragma_update(None, "journal_mode", "WAL")?;
384
385 self.conn.create_scalar_function(
388 "fnv1a_hash",
389 1,
390 rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC
391 | rusqlite::functions::FunctionFlags::SQLITE_UTF8,
392 |ctx: &rusqlite::functions::Context| -> rusqlite::Result<i64> {
393 let pk_ref = ctx.get_raw(0);
394 let pk_bytes = match pk_ref {
395 rusqlite::types::ValueRef::Text(bytes) => bytes,
396 _ => {
397 return Err(rusqlite::Error::InvalidFunctionParameterType(
398 0,
399 rusqlite::types::Type::Text,
400 ));
401 }
402 };
403 let mut hash: u32 = 2166136261;
404 for &byte in pk_bytes {
405 hash ^= byte as u32;
406 hash = hash.wrapping_mul(16777619);
407 }
408 Ok(hash as i64)
409 },
410 )?;
411
412 self.conn.execute_batch(sql_builders::INIT_SCHEMA)?;
414
415 let _ = self
417 .conn
418 .execute_batch("ALTER TABLE _stream_records ADD COLUMN user_identity TEXT");
419
420 let (version_sql, version_params) = sql_builders::init_schema_version(SCHEMA_VERSION);
422 self.conn.execute(
423 &version_sql,
424 rusqlite::params_from_iter(version_params.iter()),
425 )?;
426
427 let version: i32 = self
429 .conn
430 .query_row(
431 "SELECT value FROM _config WHERE key = 'schema_version'",
432 [],
433 |r| r.get::<_, String>(0),
434 )
435 .unwrap_or_else(|_| "1".to_string())
436 .parse()
437 .unwrap_or(1);
438
439 if version < 2 {
440 self.migrate_v1_to_v2()?;
441 }
442 if version < 3 {
443 self.migrate_v2_to_v3()?;
444 }
445 if version < 4 {
446 self.migrate_v3_to_v4()?;
447 }
448 if version < 5 {
449 self.migrate_v4_to_v5()?;
450 }
451 if version < 6 {
452 self.migrate_v5_to_v6()?;
453 }
454 if version < 7 {
455 self.migrate_v6_to_v7()?;
456 }
457 if version < 8 {
458 self.migrate_v7_to_v8()?;
459 }
460 if version < 9 {
461 self.migrate_v8_to_v9()?;
462 }
463
464 Ok(())
465 }
466
467 fn migrate_v1_to_v2(&self) -> Result<()> {
469 let mut stmt = self.conn.prepare("SELECT table_name FROM _tables")?;
470 let table_names: Vec<String> = stmt
471 .query_map([], |row| row.get(0))?
472 .collect::<std::result::Result<Vec<_>, _>>()?;
473
474 for table_name in &table_names {
475 let escaped = format!("\"{}\"", table_name.replace('"', "\"\""));
476 let _ = self.conn.execute(
477 &format!("ALTER TABLE {escaped} ADD COLUMN cached_at REAL"),
478 [],
479 );
480 }
481
482 self.conn.execute(
483 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '2')",
484 [],
485 )?;
486
487 Ok(())
488 }
489
490 fn migrate_v2_to_v3(&self) -> Result<()> {
492 let _ = self
493 .conn
494 .execute("ALTER TABLE _tables ADD COLUMN tags TEXT", []);
495
496 self.conn.execute(
497 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '3')",
498 [],
499 )?;
500
501 Ok(())
502 }
503
504 fn migrate_v3_to_v4(&self) -> Result<()> {
506 let _ = self
507 .conn
508 .execute("ALTER TABLE _tables ADD COLUMN sse_specification TEXT", []);
509 let _ = self
510 .conn
511 .execute("ALTER TABLE _tables ADD COLUMN table_class TEXT", []);
512 let _ = self.conn.execute(
513 "ALTER TABLE _tables ADD COLUMN deletion_protection_enabled INTEGER DEFAULT 0",
514 [],
515 );
516
517 self.conn.execute(
518 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '4')",
519 [],
520 )?;
521
522 Ok(())
523 }
524
525 fn migrate_v4_to_v5(&self) -> Result<()> {
527 let mut stmt = self
528 .conn
529 .prepare("SELECT table_name, gsi_definitions, lsi_definitions FROM _tables")?;
530 let tables: Vec<(String, Option<String>, Option<String>)> = stmt
531 .query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))?
532 .collect::<std::result::Result<Vec<_>, _>>()?;
533
534 for (table_name, gsi_json, lsi_json) in &tables {
535 if let Some(json) = gsi_json
536 && let Ok(gsis) = serde_json::from_str::<Vec<serde_json::Value>>(json)
537 {
538 for gsi in &gsis {
539 if let Some(idx) = gsi.get("IndexName").and_then(|v| v.as_str()) {
540 let gsi_table = escape_table_name(&format!("{table_name}::gsi::{idx}"));
541 let idx_name =
542 escape_table_name(&format!("{table_name}::gsi::{idx}::base_key"));
543 let _ = self.conn.execute_batch(&format!(
544 "CREATE INDEX IF NOT EXISTS \"{idx_name}\" ON \"{gsi_table}\" (table_pk, table_sk)"
545 ));
546 }
547 }
548 }
549 if let Some(json) = lsi_json
550 && let Ok(lsis) = serde_json::from_str::<Vec<serde_json::Value>>(json)
551 {
552 for lsi in &lsis {
553 if let Some(idx) = lsi.get("IndexName").and_then(|v| v.as_str()) {
554 let lsi_table = escape_table_name(&format!("{table_name}::lsi::{idx}"));
555 let idx_name =
556 escape_table_name(&format!("{table_name}::lsi::{idx}::base_key"));
557 let _ = self.conn.execute_batch(&format!(
558 "CREATE INDEX IF NOT EXISTS \"{idx_name}\" ON \"{lsi_table}\" (base_pk, base_sk)"
559 ));
560 }
561 }
562 }
563 }
564
565 self.conn.execute(
566 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '5')",
567 [],
568 )?;
569
570 Ok(())
571 }
572
573 fn migrate_v5_to_v6(&self) -> Result<()> {
575 let mut stmt = self.conn.prepare("SELECT table_name FROM _tables")?;
576 let table_names: Vec<String> = stmt
577 .query_map([], |row| row.get(0))?
578 .collect::<std::result::Result<Vec<_>, _>>()?;
579
580 for table_name in &table_names {
581 let escaped = escape_table_name(table_name);
582 let _ = self.conn.execute(
583 &format!(
584 "ALTER TABLE \"{escaped}\" ADD COLUMN hash_prefix TEXT NOT NULL DEFAULT ''"
585 ),
586 [],
587 );
588 }
589
590 self.conn.execute(
591 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '6')",
592 [],
593 )?;
594
595 Ok(())
596 }
597
598 fn migrate_v6_to_v7(&self) -> Result<()> {
600 let _ = self.conn.execute(
601 "ALTER TABLE _tables ADD COLUMN on_demand_throughput TEXT",
602 [],
603 );
604
605 self.conn.execute(
606 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '7')",
607 [],
608 )?;
609
610 Ok(())
611 }
612
613 fn migrate_v7_to_v8(&self) -> Result<()> {
624 let _ = self
625 .conn
626 .execute("ALTER TABLE _tables ADD COLUMN table_id TEXT", []);
627
628 let names: Vec<String> = {
629 let mut stmt = self
630 .conn
631 .prepare("SELECT table_name FROM _tables WHERE table_id IS NULL")?;
632 let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
633 rows.collect::<rusqlite::Result<Vec<String>>>()?
634 };
635 for name in names {
636 let id = uuid::Uuid::new_v4().to_string();
637 self.conn.execute(
638 "UPDATE _tables SET table_id = ?1 WHERE table_name = ?2",
639 params![id, name],
640 )?;
641 }
642
643 self.conn.execute(
644 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '8')",
645 [],
646 )?;
647
648 Ok(())
649 }
650
651 fn migrate_v8_to_v9(&self) -> Result<()> {
659 if let Err(e) = self
660 .conn
661 .execute(sql_builders::ADD_VECTOR_INDEX_DEFINITIONS_COLUMN, [])
662 && !sql_builders::is_duplicate_column_error(&e.to_string())
663 {
664 return Err(e.into());
665 }
666
667 self.conn.execute(
668 "INSERT OR REPLACE INTO _config (key, value) VALUES ('schema_version', '9')",
669 [],
670 )?;
671
672 Ok(())
673 }
674
675 pub fn conn(&self) -> &Connection {
677 &self.conn
678 }
679
680 pub fn conn_mut(&mut self) -> &mut Connection {
682 &mut self.conn
683 }
684
685 pub fn insert_table_metadata(&self, m: &CreateTableMetadata) -> Result<()> {
691 let table_name = m.table_name;
692 let (sql, params) = sql_builders::insert_table_metadata(m);
693 self.conn
694 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
695 self.metadata_cache.borrow_mut().remove(table_name);
696 Ok(())
697 }
698
699 pub fn get_table_metadata(&self, table_name: &str) -> Result<Option<TableMetadata>> {
705 if let Some(cached) = self.metadata_cache.borrow().get(table_name) {
707 return Ok(Some(cached.clone()));
708 }
709
710 let (sql, params) = sql_builders::get_table_metadata(table_name);
711 let mut stmt = self.conn.prepare(&sql)?;
712
713 let result = stmt.query_row(rusqlite::params_from_iter(params.iter()), row_to_metadata);
714
715 match result {
716 Ok(meta) => {
717 self.metadata_cache
718 .borrow_mut()
719 .insert(table_name.to_string(), meta.clone());
720 Ok(Some(meta))
721 }
722 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
723 Err(e) => Err(DynoxideError::from(e)),
724 }
725 }
726
727 pub fn delete_table_metadata(&self, table_name: &str) -> Result<bool> {
729 let (sql, params) = sql_builders::delete_table_metadata(table_name);
730 let affected = self
731 .conn
732 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
733 self.metadata_cache.borrow_mut().remove(table_name);
734 Ok(affected > 0)
735 }
736
737 pub fn update_table_metadata(
739 &self,
740 table_name: &str,
741 attribute_definitions: &str,
742 gsi_definitions: Option<&str>,
743 ) -> Result<()> {
744 let (sql, params) =
745 sql_builders::update_table_metadata(table_name, attribute_definitions, gsi_definitions);
746 self.conn
747 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
748 self.metadata_cache.borrow_mut().remove(table_name);
749 Ok(())
750 }
751
752 pub fn update_vector_index_definitions(
754 &self,
755 table_name: &str,
756 vector_index_definitions: Option<&str>,
757 ) -> Result<()> {
758 let (sql, params) =
759 sql_builders::update_vector_index_definitions(table_name, vector_index_definitions);
760 self.conn
761 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
762 self.metadata_cache.borrow_mut().remove(table_name);
763 Ok(())
764 }
765
766 pub fn update_provisioned_throughput(
768 &self,
769 table_name: &str,
770 provisioned_throughput: &str,
771 ) -> Result<()> {
772 let (sql, params) =
773 sql_builders::update_provisioned_throughput(table_name, provisioned_throughput);
774 self.conn
775 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
776 self.metadata_cache.borrow_mut().remove(table_name);
777 Ok(())
778 }
779
780 pub fn clear_provisioned_throughput(&self, table_name: &str) -> Result<()> {
782 let (sql, params) = sql_builders::clear_provisioned_throughput(table_name);
783 self.conn
784 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
785 self.metadata_cache.borrow_mut().remove(table_name);
786 Ok(())
787 }
788
789 pub fn clear_on_demand_throughput(&self, table_name: &str) -> Result<()> {
791 let (sql, params) = sql_builders::clear_on_demand_throughput(table_name);
792 self.conn
793 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
794 self.metadata_cache.borrow_mut().remove(table_name);
795 Ok(())
796 }
797
798 pub fn update_billing_mode(&self, table_name: &str, billing_mode: &str) -> Result<()> {
800 let (sql, params) = sql_builders::update_billing_mode(table_name, billing_mode);
801 self.conn
802 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
803 self.metadata_cache.borrow_mut().remove(table_name);
804 Ok(())
805 }
806
807 pub fn update_table_class(&self, table_name: &str, table_class: &str) -> Result<()> {
809 let (sql, params) = sql_builders::update_table_class(table_name, table_class);
810 self.conn
811 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
812 self.metadata_cache.borrow_mut().remove(table_name);
813 Ok(())
814 }
815
816 pub fn update_on_demand_throughput(
818 &self,
819 table_name: &str,
820 on_demand_throughput: &str,
821 ) -> Result<()> {
822 let (sql, params) =
823 sql_builders::update_on_demand_throughput(table_name, on_demand_throughput);
824 self.conn
825 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
826 self.metadata_cache.borrow_mut().remove(table_name);
827 Ok(())
828 }
829
830 pub fn get_tags(&self, table_name: &str) -> Result<Vec<crate::types::Tag>> {
836 let tags_json: Option<String> = self.conn.query_row(
837 "SELECT tags FROM _tables WHERE table_name = ?1",
838 params![table_name],
839 |row| row.get(0),
840 )?;
841
842 match tags_json {
843 Some(json) => serde_json::from_str(&json)
844 .map_err(|e| DynoxideError::InternalServerError(format!("Bad tags JSON: {e}"))),
845 None => Ok(Vec::new()),
846 }
847 }
848
849 pub fn set_tags(&self, table_name: &str, new_tags: &[crate::types::Tag]) -> Result<()> {
851 use std::collections::BTreeMap;
852
853 let existing = self.get_tags(table_name)?;
854 let mut tag_map: BTreeMap<String, String> =
855 existing.into_iter().map(|t| (t.key, t.value)).collect();
856
857 for tag in new_tags {
858 tag_map.insert(tag.key.clone(), tag.value.clone());
859 }
860
861 if tag_map.len() > 50 {
862 return Err(DynoxideError::ValidationException(
863 "One or more parameter values were invalid: \
864 Too many tags: tag limit is 50"
865 .to_string(),
866 ));
867 }
868
869 let merged: Vec<crate::types::Tag> = tag_map
870 .into_iter()
871 .map(|(k, v)| crate::types::Tag { key: k, value: v })
872 .collect();
873
874 let json = serde_json::to_string(&merged)
875 .map_err(|e| DynoxideError::InternalServerError(e.to_string()))?;
876
877 self.conn.execute(
878 "UPDATE _tables SET tags = ?1 WHERE table_name = ?2",
879 params![json, table_name],
880 )?;
881 Ok(())
882 }
883
884 pub fn update_deletion_protection(&self, table_name: &str, enabled: bool) -> Result<()> {
886 let (sql, params) = sql_builders::update_deletion_protection(table_name, enabled);
887 self.conn
888 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
889 self.metadata_cache.borrow_mut().remove(table_name);
890 Ok(())
891 }
892
893 pub fn remove_tags(&self, table_name: &str, keys: &[String]) -> Result<()> {
895 let mut tags = self.get_tags(table_name)?;
896 tags.retain(|t| !keys.contains(&t.key));
897
898 let json = if tags.is_empty() {
899 None
900 } else {
901 Some(
902 serde_json::to_string(&tags)
903 .map_err(|e| DynoxideError::InternalServerError(e.to_string()))?,
904 )
905 };
906
907 self.conn.execute(
908 "UPDATE _tables SET tags = ?1 WHERE table_name = ?2",
909 params![json, table_name],
910 )?;
911 Ok(())
912 }
913
914 pub fn list_table_names(&self) -> Result<Vec<String>> {
916 let (sql, params) = sql_builders::list_table_names();
917 let mut stmt = self.conn.prepare(&sql)?;
918 let names = stmt
919 .query_map(rusqlite::params_from_iter(params.iter()), |row| row.get(0))?
920 .collect::<std::result::Result<Vec<String>, _>>()?;
921 Ok(names)
922 }
923
924 pub fn table_exists(&self, table_name: &str) -> Result<bool> {
926 let (sql, params) = sql_builders::table_exists(table_name);
927 let count: i32 =
928 self.conn
929 .query_row(&sql, rusqlite::params_from_iter(params.iter()), |row| {
930 row.get(0)
931 })?;
932 Ok(count > 0)
933 }
934
935 #[allow(dead_code)]
937 pub(crate) fn invalidate_metadata_cache(&self, table_name: &str) {
938 self.metadata_cache.borrow_mut().remove(table_name);
939 }
940
941 pub fn create_data_table(&self, table_name: &str) -> Result<()> {
947 let (sql, params) = sql_builders::create_data_table(table_name);
948 self.conn
949 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
950 Ok(())
951 }
952
953 pub fn drop_data_table(&self, table_name: &str) -> Result<()> {
955 let (sql, params) = sql_builders::drop_data_table(table_name);
956 self.conn
957 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
958 Ok(())
959 }
960
961 pub fn create_gsi_table(&self, table_name: &str, index_name: &str) -> Result<()> {
963 let (sql, _) = sql_builders::create_gsi_table(table_name, index_name);
964 self.conn.execute_batch(&sql)?;
965 Ok(())
966 }
967
968 pub fn drop_gsi_table(&self, table_name: &str, index_name: &str) -> Result<()> {
970 let (sql, params) = sql_builders::drop_gsi_table(table_name, index_name);
971 self.conn
972 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
973 Ok(())
974 }
975
976 #[allow(clippy::too_many_arguments)]
982 pub fn insert_gsi_item(
983 &self,
984 table_name: &str,
985 index_name: &str,
986 gsi_pk: &str,
987 gsi_sk: &str,
988 table_pk: &str,
989 table_sk: &str,
990 item_json: &str,
991 ) -> Result<()> {
992 let sql = sql_builders::gsi_insert_sql(table_name, index_name);
993 let params = sql_builders::gsi_insert_params(gsi_pk, gsi_sk, table_pk, table_sk, item_json);
994 self.conn
995 .prepare_cached(&sql)?
996 .execute(rusqlite::params_from_iter(params.iter()))?;
997 Ok(())
998 }
999
1000 pub fn insert_gsi_items(
1004 &self,
1005 table_name: &str,
1006 index_name: &str,
1007 rows: &[crate::storage_backend::GsiItemRow],
1008 ) -> Result<()> {
1009 let sql = sql_builders::gsi_insert_sql(table_name, index_name);
1010 let mut stmt = self.conn.prepare_cached(&sql)?;
1011 for row in rows {
1012 let params = sql_builders::gsi_insert_params(
1013 &row.gsi_pk,
1014 &row.gsi_sk,
1015 &row.table_pk,
1016 &row.table_sk,
1017 &row.item_json,
1018 );
1019 stmt.execute(rusqlite::params_from_iter(params.iter()))?;
1020 }
1021 Ok(())
1022 }
1023
1024 pub fn delete_gsi_item(
1026 &self,
1027 table_name: &str,
1028 index_name: &str,
1029 table_pk: &str,
1030 table_sk: &str,
1031 ) -> Result<()> {
1032 let (sql, params) =
1033 sql_builders::delete_gsi_item(table_name, index_name, table_pk, table_sk);
1034 self.conn
1035 .prepare_cached(&sql)?
1036 .execute(rusqlite::params_from_iter(params.iter()))?;
1037 Ok(())
1038 }
1039
1040 pub fn query_gsi_items(
1042 &self,
1043 table_name: &str,
1044 index_name: &str,
1045 gsi_pk: &str,
1046 params: &QueryParams,
1047 ) -> Result<Vec<(String, String, String)>> {
1048 let (sql, params_vec) =
1049 sql_builders::query_gsi_items(table_name, index_name, gsi_pk, params);
1050 let mut stmt = self.conn.prepare(&sql)?;
1051 let rows = stmt
1052 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1053 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1054 })?
1055 .collect::<std::result::Result<Vec<_>, _>>()?;
1056
1057 Ok(rows)
1058 }
1059
1060 pub fn scan_gsi_items(
1062 &self,
1063 table_name: &str,
1064 index_name: &str,
1065 params: &ScanParams,
1066 ) -> Result<Vec<(String, String, String)>> {
1067 let (sql, params_vec) = sql_builders::scan_gsi_items(table_name, index_name, params);
1068 let mut stmt = self.conn.prepare(&sql)?;
1069 let rows = stmt
1070 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1071 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1072 })?
1073 .collect::<std::result::Result<Vec<_>, _>>()?;
1074
1075 Ok(rows)
1076 }
1077
1078 pub fn create_lsi_table(&self, table_name: &str, index_name: &str) -> Result<()> {
1084 let (sql, _) = sql_builders::create_lsi_table(table_name, index_name);
1085 self.conn.execute_batch(&sql)?;
1086 Ok(())
1087 }
1088
1089 pub fn create_vector_table(&self, table_name: &str, index_name: &str) -> Result<()> {
1091 let (sql, _) = sql_builders::create_vector_table(table_name, index_name);
1092 self.conn.execute_batch(&sql)?;
1093 Ok(())
1094 }
1095
1096 pub fn insert_vector_items(
1099 &self,
1100 table_name: &str,
1101 index_name: &str,
1102 rows: &[crate::storage_backend::VectorItemRow],
1103 ) -> Result<()> {
1104 let sql = sql_builders::vector_insert_sql(table_name, index_name);
1105 let mut stmt = self.conn.prepare_cached(&sql)?;
1106 for row in rows {
1107 let params = sql_builders::vector_insert_params(
1108 &row.table_pk,
1109 &row.table_sk,
1110 &row.hash_value,
1111 &row.vector_json,
1112 &row.filter_json,
1113 &row.item_json,
1114 row.entry_bytes,
1115 );
1116 stmt.execute(rusqlite::params_from_iter(params.iter()))?;
1117 }
1118 Ok(())
1119 }
1120
1121 pub fn delete_vector_item(
1124 &self,
1125 table_name: &str,
1126 index_name: &str,
1127 table_pk: &str,
1128 table_sk: &str,
1129 ) -> Result<()> {
1130 let (sql, params) =
1131 sql_builders::delete_vector_item(table_name, index_name, table_pk, table_sk);
1132 self.conn
1133 .prepare_cached(&sql)?
1134 .execute(rusqlite::params_from_iter(params.iter()))?;
1135 Ok(())
1136 }
1137
1138 pub fn query_vector_candidates(
1142 &self,
1143 table_name: &str,
1144 index_name: &str,
1145 hash_value: Option<&str>,
1146 ) -> Result<Vec<crate::storage_backend::VectorCandidateRow>> {
1147 let (sql, params) =
1148 sql_builders::query_vector_candidates(table_name, index_name, hash_value);
1149 let mut stmt = self.conn.prepare(&sql)?;
1150 let rows = stmt
1151 .query_map(rusqlite::params_from_iter(params.iter()), |row| {
1152 Ok(crate::storage_backend::VectorCandidateRow {
1153 table_pk: row.get(0)?,
1154 table_sk: row.get(1)?,
1155 vector_json: row.get(2)?,
1156 filter_json: row.get(3)?,
1157 entry_bytes: row.get(4)?,
1158 })
1159 })?
1160 .collect::<std::result::Result<Vec<_>, _>>()?;
1161 Ok(rows)
1162 }
1163
1164 pub fn vector_items_for_keys(
1166 &self,
1167 table_name: &str,
1168 index_name: &str,
1169 keys: &[(String, String)],
1170 ) -> Result<Vec<(String, String, String)>> {
1171 if keys.is_empty() {
1172 return Ok(Vec::new());
1173 }
1174 let (sql, params) = sql_builders::vector_items_for_keys(table_name, index_name, keys);
1175 let mut stmt = self.conn.prepare(&sql)?;
1176 let rows = stmt
1177 .query_map(rusqlite::params_from_iter(params.iter()), |row| {
1178 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1179 })?
1180 .collect::<std::result::Result<Vec<_>, _>>()?;
1181 Ok(rows)
1182 }
1183
1184 pub fn drop_vector_table(&self, table_name: &str, index_name: &str) -> Result<()> {
1186 let (sql, params) = sql_builders::drop_vector_table(table_name, index_name);
1187 self.conn
1188 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
1189 Ok(())
1190 }
1191
1192 pub fn drop_lsi_table(&self, table_name: &str, index_name: &str) -> Result<()> {
1194 let (sql, params) = sql_builders::drop_lsi_table(table_name, index_name);
1195 self.conn
1196 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
1197 Ok(())
1198 }
1199
1200 #[allow(clippy::too_many_arguments)]
1206 pub fn insert_lsi_item(
1207 &self,
1208 table_name: &str,
1209 index_name: &str,
1210 pk: &str,
1211 sk: &str,
1212 base_pk: &str,
1213 base_sk: &str,
1214 item_json: &str,
1215 ) -> Result<()> {
1216 let sql = sql_builders::lsi_insert_sql(table_name, index_name);
1217 let params = sql_builders::lsi_insert_params(pk, sk, base_pk, base_sk, item_json);
1218 self.conn
1219 .prepare_cached(&sql)?
1220 .execute(rusqlite::params_from_iter(params.iter()))?;
1221 Ok(())
1222 }
1223
1224 pub fn delete_lsi_item(
1226 &self,
1227 table_name: &str,
1228 index_name: &str,
1229 base_pk: &str,
1230 base_sk: &str,
1231 ) -> Result<()> {
1232 let (sql, params) = sql_builders::delete_lsi_item(table_name, index_name, base_pk, base_sk);
1233 self.conn
1234 .prepare_cached(&sql)?
1235 .execute(rusqlite::params_from_iter(params.iter()))?;
1236 Ok(())
1237 }
1238
1239 pub fn query_lsi_items(
1241 &self,
1242 table_name: &str,
1243 index_name: &str,
1244 pk: &str,
1245 params: &QueryParams,
1246 ) -> Result<Vec<(String, String, String)>> {
1247 let (sql, params_vec) = sql_builders::query_lsi_items(table_name, index_name, pk, params);
1248 let mut stmt = self.conn.prepare(&sql)?;
1249 let rows = stmt
1250 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1251 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1252 })?
1253 .collect::<std::result::Result<Vec<_>, _>>()?;
1254
1255 Ok(rows)
1256 }
1257
1258 pub fn scan_lsi_items(
1260 &self,
1261 table_name: &str,
1262 index_name: &str,
1263 params: &ScanParams,
1264 ) -> Result<Vec<(String, String, String)>> {
1265 let (sql, params_vec) = sql_builders::scan_lsi_items(table_name, index_name, params);
1266 let mut stmt = self.conn.prepare(&sql)?;
1267 let rows = stmt
1268 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1269 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1270 })?
1271 .collect::<std::result::Result<Vec<_>, _>>()?;
1272
1273 Ok(rows)
1274 }
1275
1276 pub fn begin_transaction(&self) -> Result<()> {
1282 self.conn.execute_batch(sql_builders::BEGIN)?;
1283 Ok(())
1284 }
1285
1286 pub fn commit(&self) -> Result<()> {
1288 self.conn.execute_batch(sql_builders::COMMIT)?;
1289 Ok(())
1290 }
1291
1292 pub fn rollback(&self) -> Result<()> {
1294 self.conn.execute_batch(sql_builders::ROLLBACK)?;
1295 Ok(())
1296 }
1297
1298 pub fn enable_bulk_loading(&self) -> Result<()> {
1304 self.conn.execute_batch(
1305 "PRAGMA synchronous = OFF;
1306 PRAGMA cache_size = -64000;
1307 PRAGMA temp_store = MEMORY;
1308 PRAGMA mmap_size = 268435456;",
1309 )?;
1310 Ok(())
1311 }
1312
1313 pub fn disable_bulk_loading(&self) -> Result<()> {
1315 self.conn.execute_batch(
1316 "PRAGMA synchronous = NORMAL;
1317 PRAGMA cache_size = -2000;
1318 PRAGMA temp_store = DEFAULT;
1319 PRAGMA mmap_size = 0;",
1320 )?;
1321 Ok(())
1322 }
1323
1324 pub fn put_item(
1330 &self,
1331 table_name: &str,
1332 pk: &str,
1333 sk: &str,
1334 item_json: &str,
1335 item_size: usize,
1336 ) -> Result<Option<String>> {
1337 self.put_item_with_hash(table_name, pk, sk, item_json, item_size, "")
1338 }
1339
1340 pub fn put_item_with_hash(
1342 &self,
1343 table_name: &str,
1344 pk: &str,
1345 sk: &str,
1346 item_json: &str,
1347 item_size: usize,
1348 hash_prefix: &str,
1349 ) -> Result<Option<String>> {
1350 let old_item = self.get_item(table_name, pk, sk)?;
1352
1353 let (sql, params) =
1354 sql_builders::put_item_with_hash(table_name, pk, sk, item_json, item_size, hash_prefix);
1355 self.conn
1356 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
1357
1358 Ok(old_item)
1359 }
1360
1361 pub fn put_base_items(
1366 &self,
1367 table_name: &str,
1368 rows: &[crate::storage_backend::BaseItemRow],
1369 ) -> Result<()> {
1370 let escaped = escape_table_name(table_name);
1371 let sql = format!(
1372 "INSERT OR REPLACE INTO \"{escaped}\" (pk, sk, item_json, item_size, cached_at, hash_prefix) \
1373 VALUES (?1, ?2, ?3, ?4, ?5, ?6)"
1374 );
1375 let mut stmt = self.conn.prepare_cached(&sql)?;
1376 for row in rows {
1377 stmt.execute(params![
1378 row.pk,
1379 row.sk,
1380 row.item_json,
1381 row.item_size as i64,
1382 row.cached_at,
1383 row.hash_prefix
1384 ])?;
1385 }
1386 Ok(())
1387 }
1388
1389 pub fn get_item(&self, table_name: &str, pk: &str, sk: &str) -> Result<Option<String>> {
1391 let (sql, params) = sql_builders::get_item(table_name, pk, sk);
1392 let result = self
1393 .conn
1394 .query_row(&sql, rusqlite::params_from_iter(params.iter()), |row| {
1395 row.get(0)
1396 });
1397
1398 match result {
1399 Ok(json) => Ok(Some(json)),
1400 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
1401 Err(e) => Err(DynoxideError::from(e)),
1402 }
1403 }
1404
1405 pub fn get_partition_size(&self, table_name: &str, pk: &str) -> Result<i64> {
1407 let (sql, params) = sql_builders::get_partition_size(table_name, pk);
1408 let size: i64 =
1409 self.conn
1410 .query_row(&sql, rusqlite::params_from_iter(params.iter()), |row| {
1411 row.get(0)
1412 })?;
1413 Ok(size)
1414 }
1415
1416 pub fn get_lsi_partition_size(
1419 &self,
1420 table_name: &str,
1421 index_name: &str,
1422 pk: &str,
1423 ) -> Result<i64> {
1424 let (sql, params) = sql_builders::get_lsi_partition_size(table_name, index_name, pk);
1425 let size: i64 =
1426 self.conn
1427 .query_row(&sql, rusqlite::params_from_iter(params.iter()), |row| {
1428 row.get(0)
1429 })?;
1430 Ok(size)
1431 }
1432
1433 pub fn delete_item(&self, table_name: &str, pk: &str, sk: &str) -> Result<Option<String>> {
1435 let old_item = self.get_item(table_name, pk, sk)?;
1436
1437 let (sql, params) = sql_builders::delete_item(table_name, pk, sk);
1438 self.conn
1439 .execute(&sql, rusqlite::params_from_iter(params.iter()))?;
1440
1441 Ok(old_item)
1442 }
1443
1444 pub fn query_items(
1449 &self,
1450 table_name: &str,
1451 pk: &str,
1452 params: &QueryParams,
1453 ) -> Result<Vec<(String, String, String)>> {
1454 let (sql, params_vec) = sql_builders::query_items(table_name, pk, params);
1455 let mut stmt = self.conn.prepare(&sql)?;
1456 let rows = stmt
1457 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1458 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1459 })?
1460 .collect::<std::result::Result<Vec<_>, _>>()?;
1461
1462 Ok(rows)
1463 }
1464
1465 pub fn scan_items(
1470 &self,
1471 table_name: &str,
1472 params: &ScanParams,
1473 ) -> Result<Vec<(String, String, String)>> {
1474 let (sql, params_vec) = sql_builders::scan_items(table_name, params);
1475 let mut stmt = self.conn.prepare(&sql)?;
1476 let rows = stmt
1477 .query_map(rusqlite::params_from_iter(params_vec.iter()), |row| {
1478 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1479 })?
1480 .collect::<std::result::Result<Vec<_>, _>>()?;
1481
1482 Ok(rows)
1483 }
1484
1485 pub fn count_items(&self, table_name: &str) -> Result<i64> {
1487 let (sql, params) = sql_builders::count_items(table_name);
1488 let count: i64 =
1489 self.conn
1490 .query_row(&sql, rusqlite::params_from_iter(params.iter()), |row| {
1491 row.get(0)
1492 })?;
1493 Ok(count)
1494 }
1495
1496 pub fn db_path(&self) -> Option<String> {
1502 self.conn
1503 .path()
1504 .filter(|p| !p.is_empty())
1505 .map(|p| p.to_owned())
1506 }
1507
1508 pub fn db_size_bytes(&self) -> Result<u64> {
1510 let size: i64 = self.conn.query_row(
1511 "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()",
1512 [],
1513 |row| row.get(0),
1514 )?;
1515 Ok(size as u64)
1516 }
1517
1518 pub fn table_count(&self) -> Result<usize> {
1520 let count: i64 = self
1521 .conn
1522 .query_row("SELECT COUNT(*) FROM _tables", [], |row| row.get(0))?;
1523 Ok(count as usize)
1524 }
1525
1526 pub fn table_stats(&self) -> Result<Vec<TableStats>> {
1530 let table_names = self.list_table_names()?;
1531 let mut stats = Vec::with_capacity(table_names.len());
1532 for name in table_names {
1533 let sql = format!(
1534 "SELECT COUNT(*), COALESCE(SUM(item_size), 0) FROM \"{}\"",
1535 escape_table_name(&name)
1536 );
1537 let (item_count, size_bytes): (i64, i64) = self
1538 .conn
1539 .query_row(&sql, [], |row| Ok((row.get(0)?, row.get(1)?)))?;
1540 stats.push(TableStats {
1541 table_name: name,
1542 item_count,
1543 size_bytes: size_bytes as u64,
1544 });
1545 }
1546 Ok(stats)
1547 }
1548
1549 pub fn database_info(&self) -> Result<DatabaseInfo> {
1554 let path = self.db_path();
1555 let size_bytes = self.db_size_bytes()?;
1556 let table_count = self.table_count()?;
1557 let stats = self.table_stats()?;
1558
1559 let mut table_details = Vec::with_capacity(stats.len());
1560 for s in stats {
1561 let metadata = self.get_table_metadata(&s.table_name)?;
1562 table_details.push(TableInfoEntry { stats: s, metadata });
1563 }
1564
1565 Ok(DatabaseInfo {
1566 path,
1567 size_bytes,
1568 table_count,
1569 tables: table_details,
1570 })
1571 }
1572
1573 pub fn vacuum_into(&self, path: &str) -> Result<()> {
1580 if path.contains('\0') {
1581 return Err(DynoxideError::ValidationException(
1582 "path contains null byte".to_string(),
1583 ));
1584 }
1585 self.conn
1586 .execute_batch(&format!("VACUUM INTO '{}'", path.replace('\'', "''")))?;
1587 Ok(())
1588 }
1589
1590 pub fn vacuum(&self) -> Result<()> {
1592 self.conn.execute_batch("VACUUM")?;
1593 Ok(())
1594 }
1595
1596 pub fn restore_from(&mut self, path: &str) -> Result<()> {
1600 let source = Connection::open(path)?;
1601 self.restore_from_connection(&source)
1602 }
1603
1604 pub fn backup_to_memory(&self) -> Result<Connection> {
1609 let mut dest = Connection::open_in_memory()?;
1610 {
1611 let backup = rusqlite::backup::Backup::new(&self.conn, &mut dest)?;
1612 backup.run_to_completion(100, std::time::Duration::from_millis(0), None)?;
1613 }
1614 Ok(dest)
1615 }
1616
1617 pub fn restore_from_connection(&mut self, source: &Connection) -> Result<()> {
1622 let backup = rusqlite::backup::Backup::new(source, &mut self.conn)?;
1623 backup.run_to_completion(100, std::time::Duration::from_millis(0), None)?;
1624 self.metadata_cache.borrow_mut().clear();
1625 Ok(())
1626 }
1627
1628 pub fn connection_size_bytes(conn: &Connection) -> Result<u64> {
1630 let size: i64 = conn.query_row(
1631 "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()",
1632 [],
1633 |row| row.get(0),
1634 )?;
1635 Ok(size as u64)
1636 }
1637
1638 pub fn enable_stream(&self, table_name: &str, view_type: &str, label: &str) -> Result<()> {
1644 self.conn.execute(
1645 "UPDATE _tables SET stream_enabled = 1, stream_view_type = ?1, stream_label = ?2 WHERE table_name = ?3",
1646 params![view_type, label, table_name],
1647 )?;
1648 self.metadata_cache.borrow_mut().remove(table_name);
1649 Ok(())
1650 }
1651
1652 pub fn disable_stream(&self, table_name: &str) -> Result<()> {
1654 self.conn.execute(
1655 "UPDATE _tables SET stream_enabled = 0 WHERE table_name = ?1",
1656 params![table_name],
1657 )?;
1658 self.metadata_cache.borrow_mut().remove(table_name);
1659 Ok(())
1660 }
1661
1662 #[allow(clippy::too_many_arguments)]
1664 pub fn insert_stream_record(
1665 &self,
1666 table_name: &str,
1667 event_name: &str,
1668 keys_json: &str,
1669 new_image: Option<&str>,
1670 old_image: Option<&str>,
1671 sequence_number: &str,
1672 shard_id: &str,
1673 created_at: i64,
1674 ) -> Result<()> {
1675 self.insert_stream_record_with_identity(
1676 table_name,
1677 event_name,
1678 keys_json,
1679 new_image,
1680 old_image,
1681 sequence_number,
1682 shard_id,
1683 created_at,
1684 None,
1685 )
1686 }
1687
1688 #[allow(clippy::too_many_arguments)]
1690 pub fn insert_stream_record_with_identity(
1691 &self,
1692 table_name: &str,
1693 event_name: &str,
1694 keys_json: &str,
1695 new_image: Option<&str>,
1696 old_image: Option<&str>,
1697 sequence_number: &str,
1698 shard_id: &str,
1699 created_at: i64,
1700 user_identity: Option<&str>,
1701 ) -> Result<()> {
1702 self.conn.execute(
1703 "INSERT INTO _stream_records (table_name, event_name, keys_json, new_image, old_image, sequence_number, shard_id, created_at, user_identity)
1704 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
1705 params![table_name, event_name, keys_json, new_image, old_image, sequence_number, shard_id, created_at, user_identity],
1706 )?;
1707 Ok(())
1708 }
1709
1710 pub fn next_stream_sequence_number(&self, table_name: &str) -> Result<i64> {
1712 let result: std::result::Result<i64, _> = self.conn.query_row(
1713 "SELECT COALESCE(MAX(CAST(sequence_number AS INTEGER)), 0) + 1 FROM _stream_records WHERE table_name = ?1",
1714 params![table_name],
1715 |row| row.get(0),
1716 );
1717 match result {
1718 Ok(n) => Ok(n),
1719 Err(_) => Ok(1),
1720 }
1721 }
1722
1723 pub fn get_stream_records(
1725 &self,
1726 table_name: &str,
1727 shard_id: &str,
1728 after_sequence: i64,
1729 limit: usize,
1730 ) -> Result<Vec<StreamRecord>> {
1731 let mut stmt = self.conn.prepare(
1732 "SELECT event_name, keys_json, new_image, old_image, sequence_number, created_at, user_identity
1733 FROM _stream_records
1734 WHERE table_name = ?1 AND shard_id = ?2 AND CAST(sequence_number AS INTEGER) > ?3
1735 ORDER BY CAST(sequence_number AS INTEGER) ASC
1736 LIMIT ?4",
1737 )?;
1738 let rows = stmt
1739 .query_map(
1740 params![table_name, shard_id, after_sequence, limit as i64],
1741 |row| {
1742 Ok(StreamRecord {
1743 event_name: row.get(0)?,
1744 keys_json: row.get(1)?,
1745 new_image: row.get(2)?,
1746 old_image: row.get(3)?,
1747 sequence_number: row.get(4)?,
1748 created_at: row.get(5)?,
1749 user_identity: row.get(6)?,
1750 })
1751 },
1752 )?
1753 .collect::<std::result::Result<Vec<_>, _>>()?;
1754 Ok(rows)
1755 }
1756
1757 pub fn list_stream_enabled_tables(&self) -> Result<Vec<TableMetadata>> {
1759 let sql = format!(
1760 "SELECT {} FROM _tables WHERE stream_enabled = 1 ORDER BY table_name",
1761 sql_builders::TABLE_METADATA_COLUMNS
1762 );
1763 let mut stmt = self.conn.prepare(&sql)?;
1764 let rows = stmt
1765 .query_map([], row_to_metadata)?
1766 .collect::<std::result::Result<Vec<_>, _>>()?;
1767 Ok(rows)
1768 }
1769
1770 pub fn update_ttl_config(
1776 &self,
1777 table_name: &str,
1778 attribute_name: Option<&str>,
1779 enabled: bool,
1780 ) -> Result<()> {
1781 self.conn.execute(
1782 "UPDATE _tables SET ttl_attribute = ?1, ttl_enabled = ?2 WHERE table_name = ?3",
1783 params![attribute_name, enabled as i32, table_name],
1784 )?;
1785 self.metadata_cache.borrow_mut().remove(table_name);
1786 Ok(())
1787 }
1788
1789 pub fn list_ttl_enabled_tables(&self) -> Result<Vec<TableMetadata>> {
1791 let sql = format!(
1792 "SELECT {} FROM _tables WHERE ttl_enabled = 1 ORDER BY table_name",
1793 sql_builders::TABLE_METADATA_COLUMNS
1794 );
1795 let mut stmt = self.conn.prepare(&sql)?;
1796 let rows = stmt
1797 .query_map([], row_to_metadata)?
1798 .collect::<std::result::Result<Vec<_>, _>>()?;
1799 Ok(rows)
1800 }
1801
1802 pub fn get_shard_sequence_range(
1804 &self,
1805 table_name: &str,
1806 shard_id: &str,
1807 ) -> Result<(Option<String>, Option<String>)> {
1808 let result: std::result::Result<(Option<String>, Option<String>), _> = self.conn.query_row(
1809 "SELECT MIN(sequence_number), MAX(sequence_number) FROM _stream_records WHERE table_name = ?1 AND shard_id = ?2",
1810 params![table_name, shard_id],
1811 |row| Ok((row.get(0)?, row.get(1)?)),
1812 );
1813 match result {
1814 Ok(range) => Ok(range),
1815 Err(_) => Ok((None, None)),
1816 }
1817 }
1818
1819 pub fn touch_cached_at(
1825 &self,
1826 table_name: &str,
1827 pk: &str,
1828 sk: &str,
1829 timestamp: f64,
1830 ) -> Result<()> {
1831 let sql = format!(
1832 "UPDATE \"{}\" SET cached_at = ?1 WHERE pk = ?2 AND sk = ?3",
1833 escape_table_name(table_name)
1834 );
1835 self.conn.execute(&sql, params![timestamp, pk, sk])?;
1836 Ok(())
1837 }
1838
1839 pub fn get_lru_items(
1844 &self,
1845 table_name: &str,
1846 limit: usize,
1847 ) -> Result<Vec<(String, String, i64)>> {
1848 let sql = format!(
1849 "SELECT pk, sk, item_size FROM \"{}\" WHERE cached_at IS NOT NULL ORDER BY cached_at ASC LIMIT ?1",
1850 escape_table_name(table_name)
1851 );
1852 let mut stmt = self.conn.prepare(&sql)?;
1853 let rows = stmt
1854 .query_map(params![limit as i64], |row| {
1855 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1856 })?
1857 .collect::<std::result::Result<Vec<_>, _>>()?;
1858 Ok(rows)
1859 }
1860}
1861
1862#[derive(Debug, Clone)]
1864pub struct StreamRecord {
1865 pub event_name: String,
1866 pub keys_json: String,
1867 pub new_image: Option<String>,
1868 pub old_image: Option<String>,
1869 pub sequence_number: String,
1870 pub created_at: i64,
1871 pub user_identity: Option<String>,
1872}
1873
1874#[derive(Debug, Clone)]
1876pub struct TableStats {
1877 pub table_name: String,
1878 pub item_count: i64,
1879 pub size_bytes: u64,
1880}
1881
1882#[derive(Debug, Clone)]
1884pub struct DatabaseInfo {
1885 pub path: Option<String>,
1886 pub size_bytes: u64,
1887 pub table_count: usize,
1888 pub tables: Vec<TableInfoEntry>,
1889}
1890
1891#[derive(Debug, Clone)]
1893pub struct TableInfoEntry {
1894 pub stats: TableStats,
1895 pub metadata: Option<TableMetadata>,
1896}
1897
1898#[derive(Debug, Clone)]
1904pub struct TableMetadata {
1905 pub table_name: String,
1906 pub key_schema: String,
1907 pub attribute_definitions: String,
1908 pub gsi_definitions: Option<String>,
1909 pub lsi_definitions: Option<String>,
1910 pub stream_enabled: bool,
1911 pub stream_view_type: Option<String>,
1912 pub stream_label: Option<String>,
1913 pub ttl_attribute: Option<String>,
1914 pub ttl_enabled: bool,
1915 pub created_at: i64,
1916 pub table_status: String,
1917 pub billing_mode: Option<String>,
1918 pub provisioned_throughput: Option<String>,
1919 pub sse_specification: Option<String>,
1920 pub table_class: Option<String>,
1921 pub deletion_protection_enabled: bool,
1922 pub on_demand_throughput: Option<String>,
1923 pub table_id: Option<String>,
1926 pub vector_index_definitions: Option<String>,
1929}
1930
1931#[cfg(any(feature = "native-sqlite", feature = "_has-encryption"))]
1933fn row_to_metadata(row: &rusqlite::Row) -> rusqlite::Result<TableMetadata> {
1934 Ok(TableMetadata {
1935 table_name: row.get(0)?,
1936 key_schema: row.get(1)?,
1937 attribute_definitions: row.get(2)?,
1938 gsi_definitions: row.get(3)?,
1939 lsi_definitions: row.get(4)?,
1940 stream_enabled: row.get::<_, i32>(5)? != 0,
1941 stream_view_type: row.get(6)?,
1942 stream_label: row.get(7)?,
1943 ttl_attribute: row.get(8)?,
1944 ttl_enabled: row.get::<_, i32>(9)? != 0,
1945 created_at: row.get(10)?,
1946 table_status: row.get(11)?,
1947 billing_mode: row.get(12)?,
1948 provisioned_throughput: row.get(13)?,
1949 sse_specification: row.get(14)?,
1950 table_class: row.get(15)?,
1951 deletion_protection_enabled: row.get::<_, i32>(16).unwrap_or(0) != 0,
1952 on_demand_throughput: row.get(17)?,
1953 table_id: row.get(18)?,
1954 vector_index_definitions: row.get(19)?,
1955 })
1956}
1957
1958#[cfg(all(test, any(feature = "native-sqlite", feature = "_has-encryption")))]
1959mod tests {
1960 use super::*;
1961
1962 fn test_storage() -> Storage {
1963 Storage::memory().expect("Failed to create in-memory storage")
1964 }
1965
1966 #[test]
1967 fn test_initialize_creates_metadata_tables() {
1968 let storage = test_storage();
1969 let version: String = storage
1971 .conn()
1972 .query_row(
1973 "SELECT value FROM _config WHERE key = 'schema_version'",
1974 [],
1975 |row| row.get(0),
1976 )
1977 .unwrap();
1978 assert_eq!(version, SCHEMA_VERSION);
1979 }
1980
1981 #[test]
1982 fn test_migrate_v6_to_v7_adds_on_demand_throughput_column() {
1983 let tmp = tempfile::NamedTempFile::new().unwrap();
1987 let path = tmp.path().to_str().unwrap().to_string();
1988
1989 {
1992 let conn = Connection::open(&path).unwrap();
1993 conn.execute_batch(
1994 "CREATE TABLE _config (key TEXT PRIMARY KEY, value TEXT NOT NULL);
1995 CREATE TABLE _tables (
1996 table_name TEXT PRIMARY KEY,
1997 key_schema TEXT NOT NULL,
1998 attribute_definitions TEXT NOT NULL,
1999 gsi_definitions TEXT,
2000 lsi_definitions TEXT,
2001 stream_enabled INTEGER DEFAULT 0,
2002 stream_view_type TEXT,
2003 stream_label TEXT,
2004 ttl_attribute TEXT,
2005 ttl_enabled INTEGER DEFAULT 0,
2006 created_at INTEGER NOT NULL,
2007 table_status TEXT NOT NULL DEFAULT 'ACTIVE',
2008 billing_mode TEXT DEFAULT 'PAY_PER_REQUEST',
2009 provisioned_throughput TEXT,
2010 tags TEXT,
2011 sse_specification TEXT,
2012 table_class TEXT,
2013 deletion_protection_enabled INTEGER DEFAULT 0
2014 );",
2015 )
2016 .unwrap();
2017 conn.execute(
2018 "INSERT INTO _config (key, value) VALUES ('schema_version', '6')",
2019 [],
2020 )
2021 .unwrap();
2022 conn.execute(
2023 "INSERT INTO _tables (table_name, key_schema, attribute_definitions, created_at) \
2024 VALUES ('LegacyTable', ?1, ?2, 0)",
2025 params![
2026 r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2027 r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2028 ],
2029 )
2030 .unwrap();
2031 }
2032
2033 let storage = Storage::new(&path).unwrap();
2036 let version: String = storage
2037 .conn()
2038 .query_row(
2039 "SELECT value FROM _config WHERE key = 'schema_version'",
2040 [],
2041 |r| r.get(0),
2042 )
2043 .unwrap();
2044 assert_eq!(version, SCHEMA_VERSION);
2045
2046 let meta = storage.get_table_metadata("LegacyTable").unwrap().unwrap();
2049 assert_eq!(meta.table_name, "LegacyTable");
2050 assert!(meta.on_demand_throughput.is_none());
2051
2052 let col: Option<String> = storage
2054 .conn()
2055 .query_row(
2056 "SELECT on_demand_throughput FROM _tables WHERE table_name = 'LegacyTable'",
2057 [],
2058 |r| r.get(0),
2059 )
2060 .unwrap();
2061 assert!(col.is_none());
2062 }
2063
2064 #[test]
2068 fn test_migrate_v7_to_v8_backfills_table_id() {
2069 let tmp = tempfile::NamedTempFile::new().unwrap();
2070 let path = tmp.path().to_str().unwrap().to_string();
2071
2072 {
2075 let conn = Connection::open(&path).unwrap();
2076 conn.execute_batch(
2077 "CREATE TABLE _config (key TEXT PRIMARY KEY, value TEXT NOT NULL);
2078 CREATE TABLE _tables (
2079 table_name TEXT PRIMARY KEY,
2080 key_schema TEXT NOT NULL,
2081 attribute_definitions TEXT NOT NULL,
2082 gsi_definitions TEXT,
2083 lsi_definitions TEXT,
2084 stream_enabled INTEGER DEFAULT 0,
2085 stream_view_type TEXT,
2086 stream_label TEXT,
2087 ttl_attribute TEXT,
2088 ttl_enabled INTEGER DEFAULT 0,
2089 created_at INTEGER NOT NULL,
2090 table_status TEXT NOT NULL DEFAULT 'ACTIVE',
2091 billing_mode TEXT DEFAULT 'PAY_PER_REQUEST',
2092 provisioned_throughput TEXT,
2093 tags TEXT,
2094 sse_specification TEXT,
2095 table_class TEXT,
2096 deletion_protection_enabled INTEGER DEFAULT 0,
2097 on_demand_throughput TEXT
2098 );",
2099 )
2100 .unwrap();
2101 conn.execute(
2102 "INSERT INTO _config (key, value) VALUES ('schema_version', '7')",
2103 [],
2104 )
2105 .unwrap();
2106 conn.execute(
2107 "INSERT INTO _tables (table_name, key_schema, attribute_definitions, created_at) \
2108 VALUES ('LegacyTable', ?1, ?2, 0)",
2109 params![
2110 r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2111 r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2112 ],
2113 )
2114 .unwrap();
2115 }
2116
2117 let storage = Storage::new(&path).unwrap();
2119 let version: String = storage
2120 .conn()
2121 .query_row(
2122 "SELECT value FROM _config WHERE key = 'schema_version'",
2123 [],
2124 |r| r.get(0),
2125 )
2126 .unwrap();
2127 assert_eq!(version, SCHEMA_VERSION);
2128
2129 let meta = storage.get_table_metadata("LegacyTable").unwrap().unwrap();
2132 let id = meta.table_id.expect("legacy table should be backfilled");
2133 assert!(!id.is_empty());
2134
2135 drop(storage);
2137 let storage2 = Storage::new(&path).unwrap();
2138 let meta2 = storage2.get_table_metadata("LegacyTable").unwrap().unwrap();
2139 assert_eq!(meta2.table_id.as_deref(), Some(id.as_str()));
2140 }
2141
2142 #[test]
2146 fn test_migrate_v8_to_v9_adds_vector_index_definitions_column() {
2147 let tmp = tempfile::NamedTempFile::new().unwrap();
2148 let path = tmp.path().to_str().unwrap().to_string();
2149
2150 {
2153 let conn = Connection::open(&path).unwrap();
2154 conn.execute_batch(
2155 "CREATE TABLE _config (key TEXT PRIMARY KEY, value TEXT NOT NULL);
2156 CREATE TABLE _tables (
2157 table_name TEXT PRIMARY KEY,
2158 key_schema TEXT NOT NULL,
2159 attribute_definitions TEXT NOT NULL,
2160 gsi_definitions TEXT,
2161 lsi_definitions TEXT,
2162 stream_enabled INTEGER DEFAULT 0,
2163 stream_view_type TEXT,
2164 stream_label TEXT,
2165 ttl_attribute TEXT,
2166 ttl_enabled INTEGER DEFAULT 0,
2167 created_at INTEGER NOT NULL,
2168 table_status TEXT NOT NULL DEFAULT 'ACTIVE',
2169 billing_mode TEXT DEFAULT 'PAY_PER_REQUEST',
2170 provisioned_throughput TEXT,
2171 tags TEXT,
2172 sse_specification TEXT,
2173 table_class TEXT,
2174 deletion_protection_enabled INTEGER DEFAULT 0,
2175 on_demand_throughput TEXT,
2176 table_id TEXT
2177 );",
2178 )
2179 .unwrap();
2180 conn.execute(
2181 "INSERT INTO _config (key, value) VALUES ('schema_version', '8')",
2182 [],
2183 )
2184 .unwrap();
2185 conn.execute(
2186 "INSERT INTO _tables (table_name, key_schema, attribute_definitions, created_at, table_id) \
2187 VALUES ('LegacyTable', ?1, ?2, 0, 'fixed-id')",
2188 params![
2189 r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2190 r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2191 ],
2192 )
2193 .unwrap();
2194 }
2195
2196 let storage = Storage::new(&path).unwrap();
2198 let version: String = storage
2199 .conn()
2200 .query_row(
2201 "SELECT value FROM _config WHERE key = 'schema_version'",
2202 [],
2203 |r| r.get(0),
2204 )
2205 .unwrap();
2206 assert_eq!(version, SCHEMA_VERSION);
2207
2208 let meta = storage.get_table_metadata("LegacyTable").unwrap().unwrap();
2211 assert_eq!(meta.table_name, "LegacyTable");
2212 assert!(meta.vector_index_definitions.is_none());
2213
2214 let col: Option<String> = storage
2216 .conn()
2217 .query_row(
2218 "SELECT vector_index_definitions FROM _tables WHERE table_name = 'LegacyTable'",
2219 [],
2220 |r| r.get(0),
2221 )
2222 .unwrap();
2223 assert!(col.is_none());
2224 }
2225
2226 #[test]
2227 fn test_wal_mode_enabled() {
2228 let storage = test_storage();
2229 let mode: String = storage
2230 .conn()
2231 .query_row("PRAGMA journal_mode", [], |row| row.get(0))
2232 .unwrap();
2233 assert!(mode == "wal" || mode == "memory", "Got mode: {mode}");
2235 }
2236
2237 #[test]
2243 fn fnv1a_hash_matches_known_vectors() {
2244 let storage = test_storage();
2245 let cases: [(&str, i64); 6] = [
2246 ("", 2166136261),
2247 ("a", 3826002220),
2248 ("u#1", 2199603432),
2249 ("artist#42", 2385694177),
2250 ("café", 2821410889),
2251 ("tenant#9007199254740993", 2022216178),
2252 ];
2253 for (input, expected) in cases {
2254 let got: i64 = storage
2255 .conn()
2256 .query_row("SELECT fnv1a_hash(?1)", [input], |row| row.get(0))
2257 .unwrap();
2258 assert_eq!(got, expected, "fnv1a_hash({input:?})");
2259 }
2260 }
2261
2262 #[test]
2263 fn test_table_metadata_crud() {
2264 let storage = test_storage();
2265
2266 assert!(!storage.table_exists("TestTable").unwrap());
2268 assert!(storage.list_table_names().unwrap().is_empty());
2269
2270 storage
2272 .insert_table_metadata(&CreateTableMetadata {
2273 table_name: "TestTable",
2274 key_schema: r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2275 attribute_definitions: r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2276 created_at: 1000000,
2277 ..Default::default()
2278 })
2279 .unwrap();
2280
2281 assert!(storage.table_exists("TestTable").unwrap());
2282 assert_eq!(storage.list_table_names().unwrap(), vec!["TestTable"]);
2283
2284 let meta = storage.get_table_metadata("TestTable").unwrap().unwrap();
2286 assert_eq!(meta.table_name, "TestTable");
2287 assert_eq!(meta.table_status, "ACTIVE");
2288 assert_eq!(meta.created_at, 1000000);
2289
2290 assert!(storage.delete_table_metadata("TestTable").unwrap());
2292 assert!(!storage.table_exists("TestTable").unwrap());
2293 }
2294
2295 #[test]
2296 fn test_create_and_drop_data_table() {
2297 let storage = test_storage();
2298 storage.create_data_table("MyTable").unwrap();
2299
2300 storage
2302 .put_item("MyTable", "pk1", "", r#"{"pk":{"S":"pk1"}}"#, 10)
2303 .unwrap();
2304
2305 let item = storage.get_item("MyTable", "pk1", "").unwrap();
2306 assert!(item.is_some());
2307
2308 storage.drop_data_table("MyTable").unwrap();
2309 }
2310
2311 #[test]
2312 fn test_item_crud() {
2313 let storage = test_storage();
2314 storage.create_data_table("Items").unwrap();
2315
2316 let old = storage
2318 .put_item(
2319 "Items",
2320 "user#1",
2321 "profile",
2322 r#"{"name":{"S":"Alice"}}"#,
2323 20,
2324 )
2325 .unwrap();
2326 assert!(old.is_none()); let item = storage.get_item("Items", "user#1", "profile").unwrap();
2330 assert_eq!(item.unwrap(), r#"{"name":{"S":"Alice"}}"#);
2331
2332 let old = storage
2334 .put_item("Items", "user#1", "profile", r#"{"name":{"S":"Bob"}}"#, 18)
2335 .unwrap();
2336 assert_eq!(old.unwrap(), r#"{"name":{"S":"Alice"}}"#);
2337
2338 let deleted = storage.delete_item("Items", "user#1", "profile").unwrap();
2340 assert_eq!(deleted.unwrap(), r#"{"name":{"S":"Bob"}}"#);
2341
2342 assert!(
2344 storage
2345 .get_item("Items", "user#1", "profile")
2346 .unwrap()
2347 .is_none()
2348 );
2349 }
2350
2351 #[test]
2352 fn test_query_items() {
2353 let storage = test_storage();
2354 storage.create_data_table("Orders").unwrap();
2355
2356 for i in 1..=5 {
2358 let sk = format!("order#{i:03}");
2359 let json = format!(r#"{{"id":{{"N":"{i}"}}}}"#);
2360 storage
2361 .put_item("Orders", "user#1", &sk, &json, 10)
2362 .unwrap();
2363 }
2364
2365 let results = storage
2367 .query_items(
2368 "Orders",
2369 "user#1",
2370 &QueryParams {
2371 forward: true,
2372 ..Default::default()
2373 },
2374 )
2375 .unwrap();
2376 assert_eq!(results.len(), 5);
2377 assert_eq!(results[0].1, "order#001"); let results = storage
2381 .query_items(
2382 "Orders",
2383 "user#1",
2384 &QueryParams {
2385 forward: true,
2386 limit: Some(2),
2387 ..Default::default()
2388 },
2389 )
2390 .unwrap();
2391 assert_eq!(results.len(), 2);
2392
2393 let results = storage
2395 .query_items(
2396 "Orders",
2397 "user#1",
2398 &QueryParams {
2399 forward: false,
2400 limit: Some(2),
2401 ..Default::default()
2402 },
2403 )
2404 .unwrap();
2405 assert_eq!(results.len(), 2);
2406 assert_eq!(results[0].1, "order#005"); }
2408
2409 #[test]
2410 fn test_scan_items() {
2411 let storage = test_storage();
2412 storage.create_data_table("ScanTest").unwrap();
2413
2414 storage.put_item("ScanTest", "a", "1", r#"{}"#, 2).unwrap();
2415 storage.put_item("ScanTest", "b", "2", r#"{}"#, 2).unwrap();
2416 storage.put_item("ScanTest", "c", "3", r#"{}"#, 2).unwrap();
2417
2418 let results = storage.scan_items("ScanTest", &Default::default()).unwrap();
2419 assert_eq!(results.len(), 3);
2420
2421 let results = storage
2423 .scan_items(
2424 "ScanTest",
2425 &ScanParams {
2426 limit: Some(2),
2427 ..Default::default()
2428 },
2429 )
2430 .unwrap();
2431 assert_eq!(results.len(), 2);
2432
2433 let results = storage
2435 .scan_items(
2436 "ScanTest",
2437 &ScanParams {
2438 limit: Some(2),
2439 exclusive_start_pk: Some("a"),
2440 exclusive_start_sk: Some("1"),
2441 ..Default::default()
2442 },
2443 )
2444 .unwrap();
2445 assert_eq!(results.len(), 2);
2446 assert_eq!(results[0].0, "b"); }
2448
2449 #[test]
2450 fn test_count_items() {
2451 let storage = test_storage();
2452 storage.create_data_table("CountTest").unwrap();
2453
2454 assert_eq!(storage.count_items("CountTest").unwrap(), 0);
2455
2456 storage.put_item("CountTest", "a", "", r#"{}"#, 2).unwrap();
2457 storage.put_item("CountTest", "b", "", r#"{}"#, 2).unwrap();
2458
2459 assert_eq!(storage.count_items("CountTest").unwrap(), 2);
2460 }
2461
2462 #[test]
2463 fn test_gsi_table_lifecycle() {
2464 let storage = test_storage();
2465 storage.create_gsi_table("Orders", "ByDate").unwrap();
2466
2467 let gsi_name = "Orders::gsi::ByDate";
2469 let sql = format!(
2470 "INSERT INTO \"{}\" (gsi_pk, gsi_sk, table_pk, table_sk, item_json) VALUES (?1, ?2, ?3, ?4, ?5)",
2471 gsi_name.replace('"', "\"\"")
2472 );
2473 storage
2474 .conn()
2475 .execute(
2476 &sql,
2477 params!["2024-01-01", "001", "user#1", "order#001", r#"{}"#],
2478 )
2479 .unwrap();
2480
2481 storage.drop_gsi_table("Orders", "ByDate").unwrap();
2482 }
2483
2484 #[test]
2485 fn test_nonexistent_table_metadata() {
2486 let storage = test_storage();
2487 assert!(storage.get_table_metadata("Nonexistent").unwrap().is_none());
2488 assert!(!storage.delete_table_metadata("Nonexistent").unwrap());
2489 }
2490
2491 #[test]
2492 fn test_metadata_cache_hit() {
2493 let storage = test_storage();
2494 storage
2495 .insert_table_metadata(&CreateTableMetadata {
2496 table_name: "CacheTest",
2497 key_schema: r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2498 attribute_definitions: r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2499 created_at: 1000000,
2500 ..Default::default()
2501 })
2502 .unwrap();
2503
2504 let meta1 = storage.get_table_metadata("CacheTest").unwrap().unwrap();
2506 assert_eq!(meta1.table_name, "CacheTest");
2507
2508 let meta2 = storage.get_table_metadata("CacheTest").unwrap().unwrap();
2510 assert_eq!(meta2.table_name, "CacheTest");
2511 assert_eq!(meta1.created_at, meta2.created_at);
2512
2513 assert!(storage.metadata_cache.borrow().contains_key("CacheTest"));
2515 }
2516
2517 #[test]
2518 fn test_metadata_cache_invalidated_on_delete() {
2519 let storage = test_storage();
2520 storage
2521 .insert_table_metadata(&CreateTableMetadata {
2522 table_name: "DelCache",
2523 key_schema: r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2524 attribute_definitions: r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2525 created_at: 1000000,
2526 ..Default::default()
2527 })
2528 .unwrap();
2529
2530 storage.get_table_metadata("DelCache").unwrap();
2532 assert!(storage.metadata_cache.borrow().contains_key("DelCache"));
2533
2534 storage.delete_table_metadata("DelCache").unwrap();
2536 assert!(!storage.metadata_cache.borrow().contains_key("DelCache"));
2537 }
2538
2539 #[test]
2540 fn test_metadata_cache_invalidated_on_stream_enable() {
2541 let storage = test_storage();
2542 storage
2543 .insert_table_metadata(&CreateTableMetadata {
2544 table_name: "StreamCache",
2545 key_schema: r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2546 attribute_definitions: r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2547 created_at: 1000000,
2548 ..Default::default()
2549 })
2550 .unwrap();
2551
2552 let meta = storage.get_table_metadata("StreamCache").unwrap().unwrap();
2554 assert!(!meta.stream_enabled);
2555
2556 storage
2558 .enable_stream("StreamCache", "NEW_AND_OLD_IMAGES", "2024-01-01T00:00:00")
2559 .unwrap();
2560 assert!(!storage.metadata_cache.borrow().contains_key("StreamCache"));
2561
2562 let meta = storage.get_table_metadata("StreamCache").unwrap().unwrap();
2564 assert!(meta.stream_enabled);
2565 }
2566
2567 #[test]
2568 fn test_metadata_cache_invalidated_on_ttl_update() {
2569 let storage = test_storage();
2570 storage
2571 .insert_table_metadata(&CreateTableMetadata {
2572 table_name: "TtlCache",
2573 key_schema: r#"[{"AttributeName":"pk","KeyType":"HASH"}]"#,
2574 attribute_definitions: r#"[{"AttributeName":"pk","AttributeType":"S"}]"#,
2575 created_at: 1000000,
2576 ..Default::default()
2577 })
2578 .unwrap();
2579
2580 let meta = storage.get_table_metadata("TtlCache").unwrap().unwrap();
2582 assert!(!meta.ttl_enabled);
2583
2584 storage
2586 .update_ttl_config("TtlCache", Some("expires_at"), true)
2587 .unwrap();
2588 assert!(!storage.metadata_cache.borrow().contains_key("TtlCache"));
2589
2590 let meta = storage.get_table_metadata("TtlCache").unwrap().unwrap();
2592 assert!(meta.ttl_enabled);
2593 assert_eq!(meta.ttl_attribute, Some("expires_at".to_string()));
2594 }
2595
2596 #[test]
2597 fn test_num_to_buffer_zero() {
2598 assert_eq!(num_to_buffer("0"), vec![0x80]);
2600 assert_eq!(num_to_buffer("-0"), vec![0x80]);
2601 }
2602
2603 #[test]
2604 fn test_hash_prefix_string_keys() {
2605 let h1 = compute_hash_prefix(&AttributeValue::S("3635".into()));
2608 let h2 = compute_hash_prefix(&AttributeValue::S("228".into()));
2609 let h3 = compute_hash_prefix(&AttributeValue::S("1668".into()));
2610 let h4 = compute_hash_prefix(&AttributeValue::S("3435".into()));
2611
2612 assert_eq!(
2615 hash_bucket(&h1),
2616 0,
2617 "3635 should be bucket 0, got hash {h1}"
2618 );
2619 assert_eq!(hash_bucket(&h2), 0, "228 should be bucket 0, got hash {h2}");
2620
2621 assert_eq!(
2623 hash_bucket(&h3),
2624 1,
2625 "1668 should be bucket 1, got hash {h3}"
2626 );
2627
2628 assert_eq!(
2630 hash_bucket(&h4),
2631 4,
2632 "3435 should be bucket 4, got hash {h4}"
2633 );
2634 }
2635
2636 #[test]
2637 fn test_hash_prefix_number_keys() {
2638 let h1 = compute_hash_prefix(&AttributeValue::N("251".into()));
2641 assert_eq!(hash_bucket(&h1), 1, "251 should be bucket 1, got hash {h1}");
2642
2643 let h2 = compute_hash_prefix(&AttributeValue::N("2388".into()));
2645 assert_eq!(
2646 hash_bucket(&h2),
2647 4095,
2648 "2388 should be bucket 4095, got hash {h2}"
2649 );
2650 }
2651
2652 #[test]
2653 fn test_hash_in_segment() {
2654 assert!(hash_in_segment("000000", 0, 4096));
2656 assert!(!hash_in_segment("000000", 1, 4096));
2657
2658 assert!(hash_in_segment("001000", 1, 4096));
2660 assert!(!hash_in_segment("001000", 0, 4096));
2661
2662 assert!(hash_in_segment("fff000", 4095, 4096));
2664 assert!(!hash_in_segment("fff000", 0, 4096));
2665
2666 assert!(hash_in_segment("000000", 0, 2));
2668 assert!(hash_in_segment("7ff000", 0, 2));
2669 assert!(hash_in_segment("800000", 1, 2));
2670 assert!(hash_in_segment("fff000", 1, 2));
2671 }
2672}