powdb_query/executor/
statement.rs1use super::Engine;
5use crate::plan::PlanNode;
6use crate::result::{QueryError, QueryResult};
7use std::sync::atomic::Ordering;
8
9#[cfg(feature = "testing")]
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum StatementCommitPhase {
13 BeforeCommit,
15 AfterCommit,
18}
19
20impl Engine {
21 #[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 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 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}