alopex_embedded/
owned_session.rs

1//! Owned embedded-session factories for long-lived local consumers.
2//!
3//! The existing [`crate::Database::begin`] API remains a borrowed Rust facade. This module is the
4//! separate boundary used by Python and asynchronous stream work: it clones the database's
5//! storage `Arc` and returns only core-owned session state.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use alopex_core::kv::{
11    AnyKV, OwnedKVTransactionAdapter, OwnedReadOptions, OwnedReadSession,
12    OwnedSessionFactory as CoreOwnedSessionFactory, OwnedTransactionSession,
13};
14use alopex_core::vector::hnsw::{HnswIndex, HnswTransactionState};
15use alopex_core::TxnMode;
16use alopex_sql::catalog::CatalogOverlay;
17use alopex_sql::storage::LocalRangeChangeJournal;
18
19use crate::{Database, Error, Result};
20
21/// Factory for owned local read and transaction sessions from one embedded database.
22///
23/// The factory owns the same `Arc<AnyKV>` as its source [`Database`]. Consequently, a session
24/// and every cursor it opens keep the backend alive without borrowing the database or promoting
25/// a legacy [`crate::Transaction`] lifetime.
26#[derive(Clone)]
27pub struct EmbeddedOwnedSessionFactory {
28    store: Arc<AnyKV>,
29}
30
31impl EmbeddedOwnedSessionFactory {
32    pub(crate) fn new(store: Arc<AnyKV>) -> Self {
33        Self { store }
34    }
35
36    /// Begin one owned read-only session.
37    pub fn begin_read(&self, options: OwnedReadOptions) -> Result<OwnedReadSession> {
38        self.store
39            .clone()
40            .begin_owned_read(options)
41            .map_err(Error::Core)
42    }
43
44    /// Begin one owned transaction session.
45    ///
46    /// A transaction's active lease, terminal effect, and exactly-once commit/rollback are
47    /// enforced by the returned core session. This factory never performs an implicit commit.
48    pub fn begin_transaction(&self, mode: TxnMode) -> Result<OwnedTransactionSession> {
49        self.store
50            .clone()
51            .begin_owned_transaction(mode)
52            .map_err(Error::Core)
53    }
54}
55
56impl Database {
57    /// Return an owned-session factory bound to this embedded-local database.
58    pub fn owned_session_factory(&self) -> EmbeddedOwnedSessionFactory {
59        EmbeddedOwnedSessionFactory::new(self.store.clone())
60    }
61
62    /// Begin an owned read-only session without changing the borrowed [`Self::begin`] API.
63    pub fn begin_owned_read(&self, options: OwnedReadOptions) -> Result<OwnedReadSession> {
64        self.owned_session_factory().begin_read(options)
65    }
66
67    /// Begin an owned transaction session without changing the borrowed [`Self::begin`] API.
68    pub fn begin_owned_transaction(&self, mode: TxnMode) -> Result<OwnedTransactionSession> {
69        self.owned_session_factory().begin_transaction(mode)
70    }
71
72    /// Begin an owned embedded transaction from an `Arc` database handle.
73    ///
74    /// This is the safe replacement for foreign bindings that formerly extended the lifetime of
75    /// [`crate::Transaction`].  The database `Arc`, core-owned transaction, catalog overlay, and
76    /// commit bookkeeping remain together until one explicit terminal transition.
77    pub fn begin_owned_embedded_transaction(
78        self: Arc<Self>,
79        mode: TxnMode,
80    ) -> Result<OwnedEmbeddedTransaction> {
81        let session = self.begin_owned_transaction(mode)?;
82        let journal = if mode == TxnMode::ReadWrite
83            && self.store.range_change_journal_capability()
84                == alopex_core::kv::RangeChangeJournalCapability::Supported
85        {
86            let scope = {
87                let catalog = self.sql_catalog.read().expect("catalog lock poisoned");
88                crate::sql_api::local_journal_scope(&*catalog)
89            };
90            Some(
91                session
92                    .with_transaction(|transaction| {
93                        let mut transaction = OwnedKVTransactionAdapter::new(transaction);
94                        LocalRangeChangeJournal::capture(&mut transaction, scope)
95                    })
96                    .map_err(Error::Core)?,
97            )
98        } else {
99            None
100        };
101        Ok(OwnedEmbeddedTransaction {
102            db: self,
103            session,
104            overlay: CatalogOverlay::new(),
105            catalog_modified: false,
106            journal,
107            hnsw_indices: HashMap::new(),
108            vector_cache_invalidated: false,
109        })
110    }
111}
112
113/// An embedded-local transaction that owns all state required by a Python or async handle.
114///
115/// The type deliberately has no borrowed lifetime.  Finite compatibility operations borrow the
116/// owned KV transaction only for their duration; public stream leases clone `session` and retain
117/// ownership through the core state machine.
118pub struct OwnedEmbeddedTransaction {
119    pub(crate) db: Arc<Database>,
120    pub(crate) session: OwnedTransactionSession,
121    pub(crate) overlay: CatalogOverlay,
122    pub(crate) catalog_modified: bool,
123    pub(crate) journal: Option<LocalRangeChangeJournal>,
124    pub(crate) hnsw_indices: HashMap<String, (HnswIndex, HnswTransactionState)>,
125    pub(crate) vector_cache_invalidated: bool,
126}
127
128impl OwnedEmbeddedTransaction {
129    /// Return the core session shared with a transaction-owned stream lease.
130    pub fn session(&self) -> OwnedTransactionSession {
131        self.session.clone()
132    }
133
134    /// Read one key inside this owned transaction.
135    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
136        self.session
137            .with_transaction(|transaction| transaction.get(&key.to_vec()))
138            .map_err(Error::Core)
139    }
140
141    /// Stage one key/value pair inside this owned transaction.
142    pub fn put(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
143        self.vector_cache_invalidated = true;
144        self.session
145            .with_transaction(|transaction| transaction.put(key.to_vec(), value.to_vec()))
146            .map_err(Error::Core)
147    }
148
149    /// Stage one deletion inside this owned transaction.
150    pub fn delete(&mut self, key: &[u8]) -> Result<()> {
151        self.vector_cache_invalidated = true;
152        self.session
153            .with_transaction(|transaction| transaction.delete(key.to_vec()))
154            .map_err(Error::Core)
155    }
156
157    /// Execute SQL without committing this transaction.
158    ///
159    /// The implementation is in `sql_api.rs` so it uses the same planner, catalog overlay, and
160    /// error mapping as the borrowed compatibility transaction.
161    pub fn execute_sql(&mut self, sql: &str) -> Result<crate::SqlResult> {
162        crate::sql_api::execute_sql_owned(self, sql)
163    }
164
165    /// Preflight a streamable local SELECT against this transaction's catalog overlay.
166    ///
167    /// The plan copies its required catalog metadata before returning, so the resulting stream
168    /// lease retains no borrow of this transaction or its overlay.
169    pub fn preflight_sql_stream(&self, sql: &str) -> Result<crate::OwnedSqlStreamPlan> {
170        crate::OwnedSqlStreamPlan::preflight_in_transaction(&self.db, &self.overlay, sql)
171    }
172
173    /// Commit the owned transaction after staging catalog and range-change metadata.
174    pub fn commit(&mut self) -> Result<()> {
175        let mut preparation = Ok(());
176        let journal = self.journal.take();
177        self.session
178            .with_transaction(|transaction| {
179                let mut transaction = alopex_core::kv::any::AnyKVTransaction::Owned(
180                    OwnedKVTransactionAdapter::new(transaction),
181                );
182                for (index, state) in self.hnsw_indices.values_mut() {
183                    if preparation.is_ok() {
184                        preparation = index
185                            .commit_staged(&mut transaction, state)
186                            .map_err(Error::Core);
187                    }
188                }
189                let mut catalog = self.db.sql_catalog.write().expect("catalog lock poisoned");
190                if preparation.is_ok() {
191                    preparation = catalog
192                        .persist_overlay(&mut transaction, &self.overlay)
193                        .map_err(|error| Error::Sql(error.into()));
194                }
195                if preparation.is_ok() {
196                    if let Some(journal) = journal {
197                        preparation = journal
198                            .stage(&mut transaction)
199                            .map(|_| ())
200                            .map_err(Error::Core);
201                    }
202                }
203                Ok(())
204            })
205            .map_err(Error::Core)?;
206        preparation?;
207
208        self.session.commit().map_err(Error::Core)?;
209        let overlay = std::mem::take(&mut self.overlay);
210        let mut catalog = self.db.sql_catalog.write().expect("catalog lock poisoned");
211        catalog.apply_overlay(overlay);
212        drop(catalog);
213        if self.catalog_modified {
214            self.db.invalidate_table_info_cache();
215        }
216        if self.vector_cache_invalidated {
217            let mut cache = self
218                .db
219                .vector_cache
220                .write()
221                .expect("vector cache lock poisoned");
222            *cache = None;
223        }
224        if !self.hnsw_indices.is_empty() {
225            let hnsw_indices = std::mem::take(&mut self.hnsw_indices);
226            let mut cache = self
227                .db
228                .hnsw_cache
229                .write()
230                .expect("hnsw cache lock poisoned");
231            for (name, (index, _)) in hnsw_indices {
232                cache.insert(name, Arc::new(index));
233            }
234        }
235        Ok(())
236    }
237
238    /// Roll back the owned transaction once.
239    pub fn rollback(&mut self) -> Result<()> {
240        self.session.rollback().map_err(Error::Core)?;
241        for (index, state) in self.hnsw_indices.values_mut() {
242            let _ = index.rollback(state);
243        }
244        self.hnsw_indices.clear();
245        self.overlay = CatalogOverlay::default();
246        Ok(())
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use alopex_core::kv::OwnedReadOptions;
254    use alopex_core::txn::OwnedLeaseOutcome;
255    use alopex_core::TxnMode;
256    use std::sync::Arc;
257
258    #[test]
259    fn embedded_factory_uses_owned_lifecycle_without_changing_borrowed_transactions() {
260        let database = Database::new();
261        let session = database
262            .begin_owned_transaction(TxnMode::ReadWrite)
263            .unwrap();
264        let lease = session.acquire_lease().unwrap();
265        lease
266            .with_transaction(|transaction| transaction.put(b"owned".to_vec(), b"value".to_vec()))
267            .unwrap();
268        assert!(session.commit().is_err());
269        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
270        session.commit().unwrap();
271
272        let mut borrowed = database.begin(TxnMode::ReadOnly).unwrap();
273        assert_eq!(
274            borrowed.get(b"owned".as_ref()).unwrap(),
275            Some(b"value".to_vec())
276        );
277        borrowed.commit().unwrap();
278
279        let read = database
280            .owned_session_factory()
281            .begin_read(OwnedReadOptions::default())
282            .unwrap();
283        let lease = read.acquire_lease().unwrap();
284        let mut cursor = lease
285            .with_transaction(|transaction| transaction.scan_prefix(b"own"))
286            .unwrap();
287        assert_eq!(
288            cursor.next_entry().unwrap(),
289            Some((b"owned".to_vec(), b"value".to_vec()))
290        );
291        assert_eq!(cursor.next_entry().unwrap(), None);
292        cursor.close().unwrap();
293        drop(cursor);
294        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
295    }
296
297    #[test]
298    fn dropped_embedded_owned_transaction_rolls_back_staged_writes() {
299        let database = Database::new();
300        let session = database
301            .begin_owned_transaction(TxnMode::ReadWrite)
302            .unwrap();
303        let lease = session.acquire_lease().unwrap();
304        lease
305            .with_transaction(|transaction| transaction.put(b"discard".to_vec(), b"value".to_vec()))
306            .unwrap();
307        lease.finish(OwnedLeaseOutcome::Exhausted).unwrap();
308        drop(session);
309
310        let mut borrowed = database.begin(TxnMode::ReadOnly).unwrap();
311        assert_eq!(borrowed.get(b"discard".as_ref()).unwrap(), None);
312        borrowed.commit().unwrap();
313    }
314
315    #[test]
316    fn owned_embedded_transaction_preserves_sql_visibility_until_explicit_commit() {
317        let database = Arc::new(Database::new());
318        database
319            .execute_sql("CREATE TABLE owned_sql (id INTEGER PRIMARY KEY, value TEXT)")
320            .unwrap();
321
322        let mut transaction = Arc::clone(&database)
323            .begin_owned_embedded_transaction(TxnMode::ReadWrite)
324            .unwrap();
325        transaction
326            .execute_sql("INSERT INTO owned_sql (id, value) VALUES (1, 'staged')")
327            .unwrap();
328        let alopex_sql::ExecutionResult::Query(query) = transaction
329            .execute_sql("SELECT value FROM owned_sql WHERE id = 1")
330            .unwrap()
331        else {
332            panic!("owned transaction select must return a query")
333        };
334        assert_eq!(query.rows.len(), 1);
335
336        let alopex_sql::ExecutionResult::Query(before_commit) = database
337            .execute_sql("SELECT value FROM owned_sql WHERE id = 1")
338            .unwrap()
339        else {
340            panic!("database select must return a query")
341        };
342        assert!(before_commit.rows.is_empty());
343
344        transaction.commit().unwrap();
345        let alopex_sql::ExecutionResult::Query(after_commit) = database
346            .execute_sql("SELECT value FROM owned_sql WHERE id = 1")
347            .unwrap()
348        else {
349            panic!("database select must return a query")
350        };
351        assert_eq!(after_commit.rows.len(), 1);
352    }
353
354    #[test]
355    fn owned_embedded_transaction_preserves_vector_and_similarity_workflows() {
356        let database = Arc::new(Database::new());
357        let mut transaction = Arc::clone(&database)
358            .begin_owned_embedded_transaction(TxnMode::ReadWrite)
359            .unwrap();
360        transaction
361            .upsert_vector(b"first", b"one", &[1.0, 0.0], alopex_core::Metric::Cosine)
362            .unwrap();
363        transaction
364            .upsert_vector(b"second", b"two", &[0.0, 1.0], alopex_core::Metric::Cosine)
365            .unwrap();
366        let results = transaction
367            .search_similar(&[1.0, 0.0], alopex_core::Metric::Cosine, 1, None)
368            .unwrap();
369        assert_eq!(results.len(), 1);
370        assert_eq!(results[0].key, b"first".to_vec());
371        transaction.commit().unwrap();
372
373        let mut reader = Arc::clone(&database)
374            .begin_owned_embedded_transaction(TxnMode::ReadOnly)
375            .unwrap();
376        assert_eq!(
377            reader
378                .get_vector(b"second", alopex_core::Metric::Cosine)
379                .unwrap(),
380            Some(vec![0.0, 1.0])
381        );
382        reader.rollback().unwrap();
383    }
384}
alopex_embedded/owned_session.rs

alopex_embedded/
owned_session.rs
