Skip to main content

powdb_query/executor/
statement.rs

1//! The outer query boundary owns transaction completion, including fast paths.
2//! Nested source queries (for example view materialization) share the boundary.
3
4use super::Engine;
5use crate::plan::PlanNode;
6use crate::result::{QueryError, QueryResult};
7use std::sync::atomic::Ordering;
8
9/// Test-only rendezvous points around an implicit statement commit.
10#[cfg(feature = "testing")]
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum StatementCommitPhase {
13    /// Row application finished, but no statement commit was written yet.
14    BeforeCommit,
15    /// The inline commit returned; a caller using deferred durability still
16    /// has to settle its ticket before acknowledging the result.
17    AfterCommit,
18}
19
20impl Engine {
21    /// Install a test-only hook for deterministic process-crash validation.
22    /// This API and its callback branches do not exist in shipping builds.
23    #[cfg(feature = "testing")]
24    pub fn set_statement_commit_hook_for_testing(
25        &mut self,
26        hook: Option<fn(StatementCommitPhase)>,
27    ) {
28        self.statement_commit_hook = hook;
29    }
30
31    pub(super) fn ensure_usable(&self) -> Result<(), QueryError> {
32        if self.poisoned || self.catalog.is_sync_poisoned() {
33            return Err(QueryError::EnginePoisoned);
34        }
35        Ok(())
36    }
37
38    pub(super) fn ensure_read_allowed(&self) -> Result<(), QueryError> {
39        self.ensure_usable()?;
40        if self.transaction_aborted.load(Ordering::Relaxed) {
41            return Err(QueryError::TransactionAborted);
42        }
43        Ok(())
44    }
45
46    pub(super) fn ensure_plan_allowed(&self, plan: &PlanNode) -> Result<(), QueryError> {
47        self.ensure_usable()?;
48        if self.transaction_aborted.load(Ordering::Relaxed) && !matches!(plan, PlanNode::Rollback) {
49            return Err(QueryError::TransactionAborted);
50        }
51        Ok(())
52    }
53
54    pub(super) fn begin_mutation_statement(&mut self) -> Result<(), QueryError> {
55        self.ensure_read_allowed()?;
56        if self.read_only {
57            return Err(QueryError::ReadonlyMode);
58        }
59        if !self.in_transaction {
60            if let Err(error) = self.catalog.begin_statement_transaction() {
61                if self.catalog.is_sync_poisoned() {
62                    self.poison_live_state();
63                }
64                return Err(QueryError::from_storage_io(error));
65            }
66            self.in_transaction = true;
67            self.implicit_transaction = true;
68        }
69        Ok(())
70    }
71
72    pub(super) fn begin_plan_mutation(&mut self, plan: &PlanNode) -> Result<(), QueryError> {
73        if let PlanNode::Insert { table, .. }
74        | PlanNode::Update { table, .. }
75        | PlanNode::Delete { table, .. }
76        | PlanNode::Upsert { table, .. } = plan
77        {
78            self.begin_mutation_statement()?;
79            // Persist a conservative dirty flag before any base row changes.
80            // A failed statement may leave the flag dirty, never falsely clean.
81            self.view_registry
82                .mark_dependents_dirty(table)
83                .map_err(QueryError::from_storage_io)?;
84        }
85        Ok(())
86    }
87
88    pub(super) fn poison_live_state(&mut self) {
89        self.poisoned = true;
90        self.catalog.abandon_untrusted_state();
91    }
92
93    /// Mark the current explicit transaction aborted after a rejected request.
94    /// Server callers must invoke this only for the connection owning the
95    /// transaction gate, never for another connection's admission/auth errors.
96    pub fn mark_transaction_aborted(&self) {
97        if self.in_transaction && !self.implicit_transaction {
98            self.transaction_aborted.store(true, Ordering::Relaxed);
99        }
100    }
101
102    pub(super) fn run_read_statement(
103        &self,
104        execute: impl FnOnce(&Self) -> Result<QueryResult, QueryError>,
105    ) -> Result<QueryResult, QueryError> {
106        self.ensure_read_allowed()?;
107        let result = execute(self);
108        if result.is_err() && !matches!(result, Err(QueryError::ReadonlyNeedsWrite)) {
109            self.mark_transaction_aborted();
110        }
111        result
112    }
113
114    pub(super) fn run_statement(
115        &mut self,
116        execute: impl FnOnce(&mut Self) -> Result<QueryResult, QueryError>,
117    ) -> Result<QueryResult, QueryError> {
118        self.ensure_usable()?;
119        if self.statement_depth != 0 {
120            return execute(self);
121        }
122        self.statement_depth = 1;
123        let result = execute(self);
124        self.statement_depth = 0;
125        if self.read_only {
126            return result;
127        }
128
129        if self.implicit_transaction {
130            self.implicit_transaction = false;
131            return match result {
132                Ok(result) => {
133                    #[cfg(feature = "testing")]
134                    if let Some(hook) = self.statement_commit_hook {
135                        hook(StatementCommitPhase::BeforeCommit);
136                    }
137                    if let Err(error) = self.catalog.commit_transaction() {
138                        tracing::error!(%error, "statement commit outcome is uncertain");
139                        self.poison_live_state();
140                        return Err(QueryError::CommitOutcomeUnknown);
141                    }
142                    self.in_transaction = false;
143                    if let Err(error) = self.commit_statement() {
144                        tracing::error!(%error, "statement durability completion failed");
145                        self.poison_live_state();
146                        return Err(QueryError::CommitOutcomeUnknown);
147                    }
148                    #[cfg(feature = "testing")]
149                    if let Some(hook) = self.statement_commit_hook {
150                        hook(StatementCommitPhase::AfterCommit);
151                    }
152                    Ok(result)
153                }
154                Err(error) => {
155                    if let Err(rollback_error) = self.rollback_transaction_preserving_wal_archive()
156                    {
157                        tracing::error!(%rollback_error, "failed statement could not be rolled back");
158                        self.poison_live_state();
159                        return Err(QueryError::EnginePoisoned);
160                    }
161                    Err(error)
162                }
163            };
164        }
165
166        match result {
167            Ok(result) => {
168                if let Err(error) = self.commit_statement() {
169                    tracing::error!(%error, "statement durability completion failed");
170                    self.poison_live_state();
171                    return Err(QueryError::CommitOutcomeUnknown);
172                }
173                Ok(result)
174            }
175            Err(error) => {
176                self.mark_transaction_aborted();
177                Err(error)
178            }
179        }
180    }
181
182    pub(super) fn commit_explicit_transaction(&mut self) -> Result<(), QueryError> {
183        self.ensure_read_allowed()?;
184        if let Err(error) = self.catalog.commit_transaction() {
185            tracing::error!(%error, "explicit transaction commit outcome is uncertain");
186            self.poison_live_state();
187            return Err(QueryError::CommitOutcomeUnknown);
188        }
189        self.in_transaction = false;
190        self.transaction_aborted.store(false, Ordering::Relaxed);
191        Ok(())
192    }
193}