use std::cell::Cell;
use std::rc::Rc;
use std::sync::Arc;
use parking_lot::RwLock;
use crate::io::{MemoryIO, OpenFlags};
use crate::result::LimboResult;
use crate::schema::{Schema, Table};
use crate::storage::database::{DatabaseStorage, FileMemoryStorage};
use crate::storage::pager::Pager;
use crate::util::normalize_ident;
use crate::{
maybe_init_database_file, Connection, Database, LimboError, PagerCacheflushStatus, Result,
TransactionState, IO,
};
pub(crate) const DB_MAIN: usize = 0;
pub(crate) const DB_TEMP: usize = 1;
pub(crate) const MAX_ATTACHED: usize = 10;
pub(crate) const MAIN_DB_NAME: &str = "main";
pub(crate) const TEMP_DB_NAME: &str = "temp";
pub(crate) const TEMP_DB_NAME_ALT: &str = "sqlite_temp";
const SCHEMA_TABLE_NAME: &str = "sqlite_schema";
const TEMP_SCHEMA_TABLE_NAME: &str = "sqlite_temp_schema";
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum AuxDbKind {
Temp,
Attached,
}
pub(crate) struct AuxDb {
pub(crate) name: String,
pub(crate) kind: AuxDbKind,
_db: Arc<Database>,
_conn: Arc<Connection>,
pub(crate) pager: Rc<Pager>,
pub(crate) schema: Arc<RwLock<Schema>>,
pub(crate) txn_state: Cell<TransactionState>,
}
pub(crate) enum AuxTxnEnd {
Done,
Io,
}
pub(crate) fn retag_schema_db_index(schema: &mut Schema, db: usize) {
#[allow(clippy::arc_with_non_send_sync)]
let retagged: Vec<(String, Arc<Table>)> = schema
.tables
.iter()
.filter_map(|(name, table)| match table.as_ref() {
Table::BTree(btree) if btree.db_index != db => {
let mut retagged_btree = (**btree).clone();
retagged_btree.db_index = db;
Some((
name.clone(),
Arc::new(Table::BTree(Rc::new(retagged_btree))),
))
}
_ => None,
})
.collect();
for (name, table) in retagged {
schema.tables.insert(name, table);
}
}
fn retagged_catalog(db: &Arc<Database>, index: usize) -> Result<Arc<RwLock<Schema>>> {
let mut tagged = db
.schema
.try_read()
.ok_or(LimboError::SchemaLocked)?
.clone();
retag_schema_db_index(&mut tagged, index);
Ok(Arc::new(RwLock::new(tagged)))
}
fn open_temp_database() -> Result<Arc<Database>> {
let io: Arc<dyn IO> = Arc::new(MemoryIO::new());
let file = io.open_file(":memory:", OpenFlags::default(), false)?;
maybe_init_database_file(&file, &io)?;
let db_file: Arc<dyn DatabaseStorage> = Arc::new(FileMemoryStorage::new(file));
Database::open(io, ":memory:", db_file, false)
}
fn open_attached_database(io: &Arc<dyn IO>, path: &str) -> Result<Arc<Database>> {
if path.is_empty() || path.eq_ignore_ascii_case(":memory:") {
return open_temp_database();
}
#[cfg(feature = "fs")]
{
Database::open_file(io.clone(), path, false)
}
#[cfg(not(feature = "fs"))]
{
let _ = io;
Err(LimboError::InvalidArgument(format!(
"cannot ATTACH '{path}': this build has no filesystem support (feature \"fs\" is off)"
)))
}
}
impl Connection {
pub(crate) fn has_aux_dbs(&self) -> bool {
self.aux_dbs.borrow().iter().any(|slot| slot.is_some())
}
pub(crate) fn pager_for_db(&self, db: usize) -> Result<Rc<Pager>> {
if db == DB_MAIN {
return Ok(self.pager.clone());
}
let slots = self.aux_dbs.borrow();
match slots.get(db).and_then(|slot| slot.as_ref()) {
Some(aux) => Ok(aux.pager.clone()),
None => Err(LimboError::InternalError(format!(
"no such database: index {db}"
))),
}
}
pub(crate) fn schema_for_db(&self, db: usize) -> Result<Arc<RwLock<Schema>>> {
if db == DB_MAIN {
return Ok(self.schema.clone());
}
let slots = self.aux_dbs.borrow();
match slots.get(db).and_then(|slot| slot.as_ref()) {
Some(aux) => Ok(aux.schema.clone()),
None => Err(LimboError::InternalError(format!(
"no such database: index {db}"
))),
}
}
pub(crate) fn db_index_by_name(&self, name: &str) -> Option<usize> {
let normalized = normalize_ident(name);
if normalized.eq_ignore_ascii_case(MAIN_DB_NAME) {
return Some(DB_MAIN);
}
if normalized.eq_ignore_ascii_case(TEMP_DB_NAME)
|| normalized.eq_ignore_ascii_case(TEMP_DB_NAME_ALT)
{
return Some(DB_TEMP);
}
self.aux_dbs
.borrow()
.iter()
.position(|slot| {
slot.as_ref().is_some_and(|aux| {
aux.kind == AuxDbKind::Attached && aux.name.eq_ignore_ascii_case(&normalized)
})
})
.filter(|idx| *idx >= 2)
}
pub(crate) fn db_name_for_index(&self, db: usize) -> Option<String> {
if db == DB_MAIN {
return Some(MAIN_DB_NAME.to_string());
}
if db == DB_TEMP {
return Some(TEMP_DB_NAME.to_string());
}
self.aux_dbs
.borrow()
.get(db)
.and_then(|slot| slot.as_ref())
.map(|aux| aux.name.clone())
}
pub(crate) fn ensure_temp_db(&self) -> Result<usize> {
{
let slots = self.aux_dbs.borrow();
if slots.get(DB_TEMP).and_then(|slot| slot.as_ref()).is_some() {
return Ok(DB_TEMP);
}
}
let db = open_temp_database()?;
let conn = db.connect()?;
let pager = conn.pager.clone();
let schema = retagged_catalog(&db, DB_TEMP)?;
let aux = AuxDb {
name: TEMP_DB_NAME.to_string(),
kind: AuxDbKind::Temp,
_db: db,
_conn: conn,
pager,
schema,
txn_state: Cell::new(TransactionState::None),
};
let mut slots = self.aux_dbs.borrow_mut();
while slots.len() <= DB_TEMP {
slots.push(None);
}
slots[DB_TEMP] = Some(aux);
Ok(DB_TEMP)
}
pub(crate) fn attach_db(&self, path: &str, alias: &str) -> Result<usize> {
let normalized_alias = normalize_ident(alias);
if normalized_alias.eq_ignore_ascii_case(MAIN_DB_NAME)
|| normalized_alias.eq_ignore_ascii_case(TEMP_DB_NAME)
{
return Err(LimboError::InvalidArgument(format!(
"database {normalized_alias} is already in use"
)));
}
if self.db_index_by_name(&normalized_alias).is_some() {
return Err(LimboError::InvalidArgument(format!(
"database {normalized_alias} is already in use"
)));
}
if !self.auto_commit.get() {
return Err(LimboError::TxError(
"cannot ATTACH database within transaction".to_string(),
));
}
let attached_count = self
.aux_dbs
.borrow()
.iter()
.skip(2)
.filter(|slot| slot.is_some())
.count();
if attached_count >= MAX_ATTACHED {
return Err(LimboError::InvalidArgument(format!(
"too many attached databases - max {MAX_ATTACHED}"
)));
}
let db = open_attached_database(&self._db.io, path)?;
let conn = db.connect()?;
let pager = conn.pager.clone();
let mut slots = self.aux_dbs.borrow_mut();
while slots.len() <= DB_TEMP {
slots.push(None);
}
let index = match slots.iter().skip(2).position(|slot| slot.is_none()) {
Some(free) => free + 2,
None => {
slots.push(None);
slots.len() - 1
}
};
let schema = retagged_catalog(&db, index)?;
slots[index] = Some(AuxDb {
name: normalized_alias,
kind: AuxDbKind::Attached,
_db: db,
_conn: conn,
pager,
schema,
txn_state: Cell::new(TransactionState::None),
});
Ok(index)
}
pub(crate) fn detach_db(&self, alias: &str) -> Result<()> {
let normalized_alias = normalize_ident(alias);
if normalized_alias.eq_ignore_ascii_case(MAIN_DB_NAME)
|| normalized_alias.eq_ignore_ascii_case(TEMP_DB_NAME)
{
return Err(LimboError::InvalidArgument(format!(
"cannot detach database {normalized_alias}"
)));
}
let Some(index) = self.db_index_by_name(&normalized_alias) else {
return Err(LimboError::InvalidArgument(format!(
"no such database: {normalized_alias}"
)));
};
if !self.auto_commit.get() {
return Err(LimboError::TxError(
"cannot DETACH database within transaction".to_string(),
));
}
let slot = self.aux_dbs.borrow_mut()[index].take();
if let Some(aux) = slot {
if !matches!(aux.txn_state.get(), TransactionState::None) {
return Err(LimboError::TxError(format!(
"database {normalized_alias} is locked"
)));
}
drop(aux);
}
Ok(())
}
pub(crate) fn begin_aux_txn(&self, db: usize, write: bool) -> Result<bool> {
if db == DB_MAIN {
return Ok(false);
}
let (pager, current) = {
let slots = self.aux_dbs.borrow();
let Some(aux) = slots.get(db).and_then(|slot| slot.as_ref()) else {
return Err(LimboError::InternalError(format!(
"no such database: index {db}"
)));
};
(aux.pager.clone(), aux.txn_state.get())
};
if matches!(current, TransactionState::None) {
if let LimboResult::Busy = pager.begin_read_tx()? {
return Ok(true);
}
self.set_aux_txn_state(db, TransactionState::Read)?;
}
if write && !matches!(self.aux_txn_state(db)?, TransactionState::Write) {
if let LimboResult::Busy = pager.begin_write_tx()? {
return Ok(true);
}
self.set_aux_txn_state(db, TransactionState::Write)?;
}
Ok(false)
}
fn aux_txn_state(&self, db: usize) -> Result<TransactionState> {
let slots = self.aux_dbs.borrow();
match slots.get(db).and_then(|slot| slot.as_ref()) {
Some(aux) => Ok(aux.txn_state.get()),
None => Err(LimboError::InternalError(format!(
"no such database: index {db}"
))),
}
}
fn set_aux_txn_state(&self, db: usize, state: TransactionState) -> Result<()> {
let slots = self.aux_dbs.borrow();
match slots.get(db).and_then(|slot| slot.as_ref()) {
Some(aux) => {
aux.txn_state.set(state);
Ok(())
}
None => Err(LimboError::InternalError(format!(
"no such database: index {db}"
))),
}
}
pub(crate) fn next_aux_db_with_txn(&self, from: usize) -> Option<usize> {
let slots = self.aux_dbs.borrow();
(from.max(1)..slots.len()).find(|idx| {
slots[*idx]
.as_ref()
.is_some_and(|aux| !matches!(aux.txn_state.get(), TransactionState::None))
})
}
pub(crate) fn step_end_aux_txn(&self, db: usize) -> Result<AuxTxnEnd> {
let (pager, state) = {
let slots = self.aux_dbs.borrow();
let Some(aux) = slots.get(db).and_then(|slot| slot.as_ref()) else {
return Ok(AuxTxnEnd::Done);
};
(aux.pager.clone(), aux.txn_state.get())
};
match state {
TransactionState::Write => match pager.end_tx()? {
PagerCacheflushStatus::Done(_) => {
self.set_aux_txn_state(db, TransactionState::None)?;
Ok(AuxTxnEnd::Done)
}
PagerCacheflushStatus::IO => Ok(AuxTxnEnd::Io),
},
TransactionState::Read => {
pager.end_read_tx()?;
self.set_aux_txn_state(db, TransactionState::None)?;
Ok(AuxTxnEnd::Done)
}
TransactionState::None => Ok(AuxTxnEnd::Done),
}
}
pub(crate) fn all_catalogs(&self) -> Vec<Arc<RwLock<Schema>>> {
let mut catalogs = vec![self.schema.clone()];
catalogs.extend(
self.aux_dbs
.borrow()
.iter()
.flatten()
.map(|aux| aux.schema.clone()),
);
catalogs
}
pub(crate) fn rollback_aux_txns(&self) {
let slots = self.aux_dbs.borrow();
for slot in slots.iter().flatten() {
if !matches!(slot.txn_state.get(), TransactionState::None) {
slot.pager.rollback();
slot.txn_state.set(TransactionState::None);
}
}
}
pub(crate) fn reparse_aux_schema(
&self,
db: usize,
mv_tx_id: Option<crate::mvcc::database::TxID>,
) -> Result<()> {
let (conn, schema) = {
let slots = self.aux_dbs.borrow();
let Some(aux) = slots.get(db).and_then(|slot| slot.as_ref()) else {
return Err(LimboError::InternalError(format!(
"no such database: index {db}"
)));
};
(aux._conn.clone(), aux.schema.clone())
};
let stmt = conn.prepare("SELECT * FROM sqlite_schema")?;
let mut parsed = Schema::new();
crate::util::parse_schema_rows(
Some(stmt),
&mut parsed,
conn.pager.io.clone(),
&conn.syms.borrow(),
mv_tx_id,
)?;
retag_schema_db_index(&mut parsed, db);
*schema.write() = parsed;
Ok(())
}
pub(crate) fn compile_schema(&self) -> Result<CompileSchema<'_>> {
match self.resolved_schema()? {
Some(merged) => Ok(CompileSchema::Merged(Box::new(merged))),
None => Ok(CompileSchema::Main(
self.schema.try_read().ok_or(LimboError::SchemaLocked)?,
)),
}
}
pub(crate) fn resolved_schema(&self) -> Result<Option<Schema>> {
if !self.has_aux_dbs() {
return Ok(None);
}
let main_schema = self
.schema
.try_read()
.ok_or(LimboError::SchemaLocked)?
.clone();
let mut merged = main_schema.clone();
overlay_schema(&mut merged, &main_schema, MAIN_DB_NAME, DB_MAIN, false);
let slots = self.aux_dbs.borrow();
for (index, slot) in slots.iter().enumerate().skip(2) {
let Some(slot) = slot.as_ref() else { continue };
let aux_schema = slot.schema.read();
overlay_schema(&mut merged, &aux_schema, &slot.name, index, false);
}
if let Some(temp) = slots.get(DB_TEMP).and_then(|slot| slot.as_ref()) {
let temp_schema = temp.schema.read();
overlay_schema(&mut merged, &temp_schema, TEMP_DB_NAME, DB_TEMP, true);
}
Ok(Some(merged))
}
}
pub(crate) enum CompileSchema<'conn> {
Main(parking_lot::RwLockReadGuard<'conn, Schema>),
Merged(Box<Schema>),
}
impl CompileSchema<'_> {
pub(crate) fn as_ref(&self) -> &Schema {
match self {
CompileSchema::Main(guard) => guard,
CompileSchema::Merged(schema) => schema,
}
}
}
fn overlay_schema(target: &mut Schema, source: &Schema, db_name: &str, db: usize, shadow: bool) {
let is_main = db == DB_MAIN;
for (name, table) in source.tables.iter() {
if name.contains('.') {
continue;
}
let qualified = format!("{db_name}.{name}");
target.tables.insert(qualified.clone(), table.clone());
target.object_db.insert(qualified, db);
let is_schema_table = name == SCHEMA_TABLE_NAME;
if is_schema_table && !is_main {
if db == DB_TEMP {
target
.tables
.insert(TEMP_SCHEMA_TABLE_NAME.to_string(), table.clone());
target
.object_db
.insert(TEMP_SCHEMA_TABLE_NAME.to_string(), db);
}
continue;
}
if shadow || !target.tables.contains_key(name) {
target.tables.insert(name.clone(), table.clone());
target.object_db.insert(name.clone(), db);
}
}
for (table_name, indexes) in source.indexes.iter() {
if table_name.contains('.') {
continue;
}
target
.indexes
.insert(format!("{db_name}.{table_name}"), indexes.clone());
if shadow || !target.indexes.contains_key(table_name) {
target.indexes.insert(table_name.clone(), indexes.clone());
}
for index in indexes {
let index_key = normalize_ident(&index.name);
target
.object_db
.insert(format!("{db_name}.{index_key}"), db);
if shadow || !target.object_db.contains_key(&index_key) {
target.object_db.insert(index_key, db);
}
}
}
#[cfg(not(feature = "index_experimental"))]
for table_name in source.has_indexes.iter() {
if table_name.contains('.') {
continue;
}
target.has_indexes.insert(format!("{db_name}.{table_name}"));
target.has_indexes.insert(table_name.clone());
}
for (trigger_name, trigger) in source.triggers.iter() {
target
.object_db
.insert(format!("{db_name}.{trigger_name}"), db);
if shadow || !target.triggers.contains_key(trigger_name) {
target
.triggers
.insert(trigger_name.clone(), trigger.clone());
target.object_db.insert(trigger_name.clone(), db);
}
}
}