mod fetch;
mod statement_snapshot;
mod table_snapshot;
mod worker;
pub(crate) use statement_snapshot::StatementReadSnapshot;
use crate::{
DocumentStore, Engine, EpochCoordinator, InvertedIndex, PinnedPortalTransactionControl,
QueryRuntime, RelationIdentity, RuntimeExtensions, SQLError, SQLResult,
SessionPortalCatalogSnapshot, SessionPortalCommandDeclaration, SessionPortalData,
SessionPortalDeclaration, SessionPortalMaterialization, SessionPortalPosition,
SessionPortalRestart, SessionPortalSQLFunctionSnapshots, SessionPortalState,
SessionPortalTableSnapshots, SessionPortalViewSnapshots, StorageContext, TableState, Value,
};
use fetch::{
ensure_portal_rows_for_fetch, fetch_directional_query_portal, fetch_indices,
materialize_portal_to_end, select_portal_rows, uses_directional_query_execution,
};
use uqa_sql::ast::{CursorDirection, FetchCursorStmt};
use uqa_execution::query::document_changes::{DocumentChanges, DocumentSelection};
use uqa_sql::binding::portals::SessionPortalTableDependencies;
type SessionPortalTableSource = (
RelationIdentity,
std::sync::Arc<TableState>,
std::sync::Arc<TableState>,
);
impl Engine {
pub(crate) fn allocate_session_portal_name(&self) -> String {
self.session.portal_registry.allocate_name()
}
fn allocate_session_portal_transaction_origin(&self) -> u64 {
let mut next = self.session.next_portal_transaction_origin.lock();
let origin = *next;
*next = next.wrapping_add(1).max(1);
origin
}
pub(crate) fn open_pending_session_portal(
&self,
declaration: SessionPortalDeclaration,
) -> Result<(), SQLError> {
let SessionPortalDeclaration {
metadata,
mut query,
params,
columns,
column_types,
} = declaration;
let name = metadata.name.clone();
let scrollable = metadata.is_scrollable;
let holdable = metadata.is_holdable;
let table_dependencies =
uqa_sql::binding::portals::prepare_query(&self.portal_binding_context(), &mut query)?;
let snapshot_gate = self
.row_locks
.begin_change_snapshot(&self.runtime.cancellation)?;
let table_sources = {
let stack = self.session.transactions.lock();
let fixed_snapshot = stack
.first()
.and_then(|frame| frame.fixed_snapshot.as_ref());
self.capture_session_portal_table_sources(fixed_snapshot, &table_dependencies)
};
let transaction_overlay =
self.capture_session_portal_transaction_overlay(&table_sources)?;
snapshot_gate.baseline()?;
drop(snapshot_gate);
let table_snapshots =
self.detach_session_portal_table_snapshots(table_sources, transaction_overlay)?;
let mut catalog_snapshot = self.durable.snapshot();
catalog_snapshot.graphs = self.freeze_graph_read_handles(
table_dependencies.graphs.as_ref(),
table_dependencies.graph_catalog,
)?;
let catalog_snapshot = std::sync::Arc::new(catalog_snapshot);
let view_snapshots = std::sync::Arc::clone(&catalog_snapshot.views);
let sql_function_snapshots = std::sync::Arc::clone(&catalog_snapshot.sql_user_functions);
let restart = holdable.then(|| SessionPortalRestart {
query: query.clone(),
params: params.clone(),
table_snapshots: std::sync::Arc::clone(&table_snapshots),
view_snapshots: std::sync::Arc::clone(&view_snapshots),
sql_function_snapshots: std::sync::Arc::clone(&sql_function_snapshots),
catalog_snapshot: std::sync::Arc::clone(&catalog_snapshot),
});
let transaction_origin = self.allocate_session_portal_transaction_origin();
let registration = self.session.portal_registry.register(metadata)?;
let mut portals = self.session.portals.lock();
portals.insert(
name,
SessionPortalState {
data: SessionPortalData::Pending {
query,
params,
table_snapshots,
view_snapshots,
sql_function_snapshots,
catalog_snapshot,
restart,
},
columns,
column_types,
transaction_origin,
position: SessionPortalPosition::BeforeFirst,
scrollable,
holdable,
pinned_transaction_control: PinnedPortalTransactionControl::MakeHoldable,
pin_count: 0,
_registration: registration,
},
);
Ok(())
}
pub(crate) fn open_pending_command_session_portal(
&self,
declaration: SessionPortalCommandDeclaration,
) -> Result<(), SQLError> {
let SessionPortalCommandDeclaration {
metadata,
command,
params,
columns,
column_types,
null_returning_values,
} = declaration;
let name = metadata.name.clone();
let scrollable = metadata.is_scrollable;
let transaction_origin = self.allocate_session_portal_transaction_origin();
let registration = self.session.portal_registry.register(metadata)?;
let mut portals = self.session.portals.lock();
portals.insert(
name,
SessionPortalState {
data: SessionPortalData::PendingCommand {
command,
params,
null_returning_values,
},
columns,
column_types,
transaction_origin,
position: SessionPortalPosition::BeforeFirst,
scrollable,
holdable: false,
pinned_transaction_control: PinnedPortalTransactionControl::Reject,
pin_count: 0,
_registration: registration,
},
);
Ok(())
}
pub(crate) fn ensure_session_portal_available(&self, name: &str) -> Result<(), SQLError> {
self.session.portal_registry.ensure_available(name)
}
pub(crate) fn pin_session_portal(&self, name: &str) -> Result<(), SQLError> {
let mut portals = self.session.portals.lock();
let portal = portals
.get_mut(name)
.ok_or_else(|| cursor_error(name, "does not exist", "34000"))?;
portal.pin_count = portal
.pin_count
.checked_add(1)
.ok_or_else(|| SQLError::Internal(format!("cursor \"{name}\" pin count overflow")))?;
Ok(())
}
pub(crate) fn unpin_session_portal(&self, name: &str) -> Result<(), SQLError> {
let mut portals = self.session.portals.lock();
let portal = portals
.get_mut(name)
.ok_or_else(|| cursor_error(name, "does not exist", "34000"))?;
portal.pin_count = portal
.pin_count
.checked_sub(1)
.ok_or_else(|| SQLError::Internal(format!("cursor \"{name}\" is not pinned")))?;
Ok(())
}
pub(crate) fn fetch_session_portal(
&self,
fetch: &FetchCursorStmt,
) -> Result<SQLResult, SQLError> {
let mut state = self
.session
.portals
.lock()
.remove(&fetch.name)
.ok_or_else(|| cursor_error(&fetch.name, "does not exist", "34000"))?;
let result = (|| {
if uses_directional_query_execution(&state) {
return fetch_directional_query_portal(
self,
&mut state,
fetch.direction,
fetch.count,
fetch.move_only,
);
}
ensure_portal_rows_for_fetch(
self,
&mut state,
fetch.direction,
fetch.count,
fetch.move_only,
)?;
let indices = fetch_indices(&mut state, fetch.direction, fetch.count, fetch.move_only)?;
if fetch.move_only {
return Ok(SQLResult::from_affected(indices.len() as u64));
}
select_portal_rows(&mut state, &indices)
})();
self.session
.portals
.lock()
.insert(fetch.name.clone(), state);
result
}
pub(crate) fn materialize_holdable_session_portals(&self) -> Result<bool, SQLError> {
let names = self
.session
.portals
.lock()
.iter()
.filter_map(|(name, portal)| {
(portal.holdable
&& matches!(
&portal.data,
SessionPortalData::Pending { .. } | SessionPortalData::Streaming { .. }
))
.then_some(name.clone())
})
.collect::<Vec<_>>();
let materialized_any = !names.is_empty();
self.materialize_named_session_portals(&names)?;
Ok(materialized_any)
}
pub(crate) fn prepare_pinned_session_portals_for_transaction_control(
&self,
) -> Result<(), SQLError> {
let names = {
let mut portals = self.session.portals.lock();
if portals.values().any(|portal| {
portal.pin_count != 0
&& portal.pinned_transaction_control == PinnedPortalTransactionControl::Reject
}) {
return Err(SQLError::Routine {
sqlstate: "55000".into(),
message: "cannot perform transaction commands inside a cursor loop that is not read-only".into(),
});
}
let names = portals
.iter()
.filter_map(|(name, portal)| (portal.pin_count != 0).then_some(name.clone()))
.collect::<Vec<_>>();
for name in &names {
if let Some(portal) = portals.get_mut(name) {
portal.holdable = true;
}
}
names
};
self.materialize_named_session_portals(&names)?;
if !names.is_empty() {
let mut stack = self.session.transactions.lock();
let frame = stack.last_mut().ok_or_else(|| {
SQLError::Internal(
"pinned PL/pgSQL portal prepared without an active transaction".into(),
)
})?;
frame.session_snapshot.portal_names.extend(names);
}
Ok(())
}
fn materialize_named_session_portals(&self, names: &[String]) -> Result<(), SQLError> {
for name in names {
let mut state = self
.session
.portals
.lock()
.remove(name)
.ok_or_else(|| cursor_error(name, "does not exist", "34000"))?;
let _row_lock_statement = self.begin_row_lock_statement();
let result = materialize_portal_to_end(self, &mut state);
self.session.portals.lock().insert(name.clone(), state);
result?;
}
Ok(())
}
pub(crate) fn close_session_portal(&self, name: &str) -> Result<(), SQLError> {
let mut portals = self.session.portals.lock();
let Some(portal) = portals.get(name) else {
return Err(cursor_error(name, "does not exist", "34000"));
};
if portal.pin_count != 0 {
return Err(SQLError::Routine {
sqlstate: "24000".into(),
message: format!("cannot drop pinned portal \"{name}\""),
});
}
portals.remove(name);
Ok(())
}
pub(crate) fn close_all_session_portals(&self) {
self.session.portals.lock().clear();
}
fn capture_session_portal_table_sources(
&self,
fixed_snapshot: Option<&crate::FixedTransactionSnapshot>,
dependencies: &SessionPortalTableDependencies,
) -> Vec<SessionPortalTableSource> {
let live_tables = self
.storage
.tables
.read()
.iter()
.filter(|(relation, _)| dependencies.includes(relation))
.map(|(relation, metadata)| (relation.clone(), std::sync::Arc::clone(metadata)))
.collect::<Vec<_>>();
live_tables
.into_iter()
.map(|(relation, metadata)| {
let data = fixed_snapshot
.and_then(|snapshot| snapshot.table_for_live_relation(&relation, &metadata))
.unwrap_or_else(|| std::sync::Arc::clone(&metadata));
(relation, data, metadata)
})
.collect()
}
fn detach_session_portal_table_snapshots(
&self,
sources: Vec<SessionPortalTableSource>,
mut transaction_overlay: std::collections::BTreeMap<String, DocumentChanges>,
) -> Result<SessionPortalTableSnapshots, SQLError> {
let mut snapshots = std::collections::BTreeMap::new();
for (relation, data, metadata) in sources {
let canonical = relation.qualified_name();
let changes = if std::sync::Arc::ptr_eq(&data, &metadata) {
None
} else {
transaction_overlay.remove(&canonical)
};
snapshots.insert(
relation,
self.detach_query_table(&data, &metadata, changes)?,
);
}
Ok(std::sync::Arc::new(snapshots))
}
pub(crate) fn capture_detached_fixed_transaction_snapshot(
&self,
) -> Result<SessionPortalTableSnapshots, SQLError> {
let mut snapshots = std::collections::BTreeMap::new();
let live_tables = self
.storage
.tables
.read()
.iter()
.filter(|(_, table)| table.persistence != uqa_sql::ast::RelationPersistence::Temporary)
.map(|(relation, table)| (relation.clone(), std::sync::Arc::clone(table)))
.collect::<Vec<_>>();
for (relation, table) in live_tables {
snapshots.insert(relation, self.detach_query_table(&table, &table, None)?);
}
Ok(std::sync::Arc::new(snapshots))
}
pub(crate) fn detach_query_table(
&self,
data: &std::sync::Arc<TableState>,
metadata: &std::sync::Arc<TableState>,
changes: Option<DocumentChanges>,
) -> Result<std::sync::Arc<TableState>, SQLError> {
let control = self.query_retention_control()?;
if std::sync::Arc::ptr_eq(data, metadata)
&& changes
.as_ref()
.is_none_or(|changes| !changes.has_changes())
&& (self.storage.backend.is_none() || self.versioned_backend_transactions())
{
return Self::retain_query_table(data, &control);
}
let source_columns = data.columns.snapshot();
let source = data.document_store.read();
let storage = Self::with_query_snapshot_schema(metadata, |schema| {
if self.storage.backend.is_none() || self.versioned_backend_transactions() {
return uqa_execution::query::table_snapshot::retain(
source
.snapshot()
.map_err(|error| portal_snapshot_error("documents", &error))?,
&source_columns,
schema,
changes.unwrap_or_default(),
&control,
);
}
uqa_execution::query::table_snapshot::materialize(
source.as_ref(),
&source_columns,
schema,
changes.unwrap_or_default(),
&control,
)
})?;
Ok(Self::query_table_with_storage(
metadata,
storage.documents,
storage.text,
storage.vectors,
storage.document_count,
false,
))
}
pub(crate) fn detach_empty_query_table(
&self,
metadata: &std::sync::Arc<TableState>,
) -> Result<std::sync::Arc<TableState>, SQLError> {
let control = self.query_retention_control()?;
let storage = Self::with_query_snapshot_schema(metadata, |schema| {
uqa_execution::query::table_snapshot::empty(schema, &control)
})?;
Ok(Self::query_table_with_storage(
metadata,
storage.documents,
storage.text,
storage.vectors,
0,
false,
))
}
fn capture_session_portal_transaction_overlay(
&self,
sources: &[SessionPortalTableSource],
) -> Result<std::collections::BTreeMap<String, DocumentChanges>, SQLError> {
let relation_names = sources
.iter()
.filter(|(_, data, metadata)| !std::sync::Arc::ptr_eq(data, metadata))
.map(|(relation, _, metadata)| {
(metadata.storage_generation(), relation.qualified_name())
})
.collect::<std::collections::BTreeMap<_, _>>();
if relation_names.is_empty() {
return Ok(std::collections::BTreeMap::new());
}
let control = self.query_retention_control()?;
let desired = {
let stack = self.session.transactions.lock();
let mut desired = std::collections::BTreeMap::<String, DocumentSelection>::new();
for change in stack.iter().flat_map(|frame| frame.row_changes.iter()) {
if let Some(table) = relation_names.get(&change.source_generation) {
desired
.entry(table.clone())
.or_insert_with(|| DocumentSelection::new(&control))
.insert(
change.pending.key.doc_id,
!matches!(
change.pending.kind,
crate::row_locks::PendingRowChangeKind::Delete
| crate::row_locks::PendingRowChangeKind::Rewrite(_)
),
&control,
)
.map_err(|error| {
uqa_execution::storage_errors::storage_error(
"select portal rows",
&error,
)
})?;
}
if let crate::row_locks::PendingRowChangeKind::Rewrite(successor) =
change.pending.kind
{
if let Some(table) = change
.successor_generation
.and_then(|generation| relation_names.get(&generation))
{
desired
.entry(table.clone())
.or_insert_with(|| DocumentSelection::new(&control))
.insert(successor.doc_id, true, &control)
.map_err(|error| {
uqa_execution::storage_errors::storage_error(
"select portal rows",
&error,
)
})?;
}
}
}
desired
};
let mut overlay = std::collections::BTreeMap::new();
for (table_name, desired_documents) in desired {
let table = self.require_table(&table_name)?;
overlay.insert(
table_name,
self.capture_query_document_changes(&table, desired_documents)?,
);
}
Ok(overlay)
}
}
impl Engine {
pub(crate) fn fork_session_portal_worker_engine(&self) -> Result<Engine, SQLError> {
let table_snapshots = self.query_table_snapshots.clone().ok_or_else(|| {
SQLError::Internal("directional query branch has no table snapshot".into())
})?;
let view_snapshots = self.query_view_snapshots.clone().ok_or_else(|| {
SQLError::Internal("directional query branch has no view snapshot".into())
})?;
let sql_function_snapshots =
self.query_sql_function_snapshots.clone().ok_or_else(|| {
SQLError::Internal("directional query branch has no SQL function snapshot".into())
})?;
let catalog_snapshot = self.query_catalog_snapshot.clone().ok_or_else(|| {
SQLError::Internal("directional query branch has no catalog snapshot".into())
})?;
let transaction_origin = self.query_transaction_origin.ok_or_else(|| {
SQLError::Internal("directional query branch has no transaction origin".into())
})?;
Ok(self.session_portal_worker_engine(
table_snapshots,
view_snapshots,
sql_function_snapshots,
catalog_snapshot,
transaction_origin,
))
}
fn session_portal_worker_engine(
&self,
table_snapshots: SessionPortalTableSnapshots,
view_snapshots: SessionPortalViewSnapshots,
sql_function_snapshots: SessionPortalSQLFunctionSnapshots,
catalog_snapshot: SessionPortalCatalogSnapshot,
transaction_origin: u64,
) -> Engine {
let mut epochs = EpochCoordinator::new();
epochs.share_published_from(&self.epochs);
let mut runtime = QueryRuntime::new(self.sql_function_depth_limit());
runtime.statement_gate = std::sync::Arc::clone(&self.runtime.statement_gate);
runtime.cancellation = self.runtime.cancellation.clone();
runtime.notices = std::sync::Arc::clone(&self.runtime.notices);
runtime.notifications = std::sync::Arc::clone(&self.runtime.notifications);
Engine {
storage: StorageContext::shared_from(&self.storage),
durable: std::sync::Arc::clone(&self.durable),
session: std::sync::Arc::clone(&self.session),
extensions: RuntimeExtensions::shared_from(&self.extensions),
epochs,
runtime,
row_locks: std::sync::Arc::clone(&self.row_locks),
statistics: std::sync::Arc::clone(&self.statistics),
notification_hub: std::sync::Arc::clone(&self.notification_hub),
session_id: self.session_id,
owns_session_registration: false,
query_table_snapshots: Some(table_snapshots),
query_view_snapshots: Some(view_snapshots),
query_sql_function_snapshots: Some(sql_function_snapshots),
query_catalog_snapshot: Some(catalog_snapshot),
query_transaction_overlay: Some(std::sync::Arc::new(std::collections::BTreeMap::new())),
query_transaction_origin: Some(transaction_origin),
}
}
}
fn cursor_error(name: &str, message: &str, sqlstate: &str) -> SQLError {
SQLError::Routine {
sqlstate: sqlstate.into(),
message: format!("cursor \"{name}\" {message}"),
}
}
fn portal_snapshot_error(component: &str, error: &uqa_storage::StorageBackendError) -> SQLError {
uqa_execution::storage_errors::storage_error(
&format!("capture cursor {component} snapshot"),
error,
)
}