1mod accepted_schema;
7mod bounded_cache;
8mod catalog;
9mod query;
10mod response;
11#[cfg(feature = "sql")]
12mod sql;
13#[cfg(all(test, feature = "sql"))]
17mod tests;
18mod write;
19
20use crate::{
21 db::{
22 Db, EntityRuntimeHooks, FluentDeleteQuery, FluentLoadQuery, MissingRowPolicy, PersistedRow,
23 Query, StoreRegistry, WriteBatchResponse,
24 commit::CommitSchemaFingerprint,
25 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
26 schema::{AcceptedRowDecodeContract, CompiledAcceptedRowConstraints, SchemaInfo},
27 },
28 entity::{EntityKind, EntityValue},
29 error::InternalError,
30 metrics::sink::{ExecKind, MetricsSink, record_exec_error_for_path, with_metrics_sink},
31 traits::CanisterKind,
32 value::Value,
33};
34use std::thread::LocalKey;
35
36#[cfg(all(test, feature = "sql"))]
37pub(in crate::db) use accepted_schema::AcceptedCatalogRuntimeCounterSnapshot;
38pub(in crate::db) use accepted_schema::AcceptedSchemaCatalogContext;
39pub(in crate::db) use query::{
40 AcceptedExecutionOutput, AcceptedIdValuesOutput, AcceptedOptionalValueOutput,
41 AcceptedValuesOutput,
42};
43#[cfg(feature = "diagnostics")]
44pub use query::{
45 DirectDataRowAttribution, FluentTerminalExecutionAttribution, GroupedCountAttribution,
46 GroupedExecutionAttribution, KernelRowAttribution, QueryExecutionAttribution,
47 ScalarAggregateAttribution,
48};
49pub(in crate::db) use response::finalize_scalar_paged_execution;
50pub(in crate::db) use response::finalize_structural_grouped_projection_result;
51#[cfg(feature = "sql")]
52pub(in crate::db) use response::sql_grouped_cursor_from_bytes;
53#[cfg(all(feature = "sql", feature = "diagnostics"))]
54pub use sql::{
55 SqlCompileAttribution, SqlExecutionAttribution, SqlHybridCoveringAttribution,
56 SqlOutputBlobAttribution, SqlPureCoveringAttribution, SqlQueryCacheAttribution,
57 SqlQueryExecutionAttribution,
58};
59#[cfg(feature = "sql")]
60pub use sql::{
61 SqlConstraintValidationPage, SqlConstraintValidationRevisionStatus,
62 SqlConstraintValidationState, SqlDdlExecutionStatus, SqlDdlMutationKind,
63 SqlDdlPreparationReport, SqlStatementDispatch, SqlStatementResult, SqlStatementShellSurface,
64 SqlStatementSurface, TrustedResumableUpdateContinuation, TrustedResumableUpdatePhase,
65 TrustedResumableUpdateReceipt, TrustedResumableUpdateRestartReason, sql_statement_dispatch,
66 sql_statement_entity_name, sql_statement_shell_surface, sql_statement_surface,
67};
68#[cfg(all(feature = "sql", test))]
69pub(in crate::db) use sql::{
70 SqlDeleteExposurePolicy, SqlDeletePolicyContext, SqlPublicBoundedDeletePlan,
71 SqlPublicBoundedUpdatePlan, SqlPublicPrimaryKeyDeletePlan, SqlPublicPrimaryKeyUpdatePlan,
72 SqlUpdateExposurePolicy, SqlUpdatePolicyContext, SqlValidatedDeletePlan,
73 SqlValidatedUpdatePlan, classify_sql_delete_policy, classify_sql_update_policy,
74};
75#[cfg(all(feature = "sql", feature = "diagnostics"))]
76pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
77
78pub struct DbSession<C: CanisterKind> {
85 db: Db<C>,
86 debug: bool,
87 metrics: Option<&'static dyn MetricsSink>,
88}
89
90impl<C: CanisterKind> DbSession<C> {
91 #[must_use]
93 pub(crate) const fn new(db: Db<C>) -> Self {
94 Self {
95 db,
96 debug: false,
97 metrics: None,
98 }
99 }
100
101 #[must_use]
103 pub const fn new_with_hooks(
104 store: &'static LocalKey<StoreRegistry>,
105 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
106 ) -> Self {
107 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
108 }
109
110 #[must_use]
112 pub const fn debug(mut self) -> Self {
113 self.debug = true;
114 self
115 }
116
117 #[must_use]
119 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
120 self.metrics = Some(sink);
121 self
122 }
123
124 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
127 where
128 E: EntityKind<Canister = C>,
129 {
130 FluentLoadQuery::new(self, consistency)
131 }
132
133 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
137 where
138 E: PersistedRow<Canister = C>,
139 {
140 FluentDeleteQuery::new(self, Query::new(consistency).delete())
141 }
142
143 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
144 if let Some(sink) = self.metrics {
145 with_metrics_sink(sink, f)
146 } else {
147 f()
148 }
149 }
150
151 fn execute_save_with<E, T, R>(
153 &self,
154 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
155 map: impl FnOnce(T) -> R,
156 ) -> Result<R, InternalError>
157 where
158 E: PersistedRow<Canister = C>,
159 {
160 let (contract, schema_info, schema_fingerprint, accepted_row_constraints) = match self
161 .with_metrics(|| self.ensure_generated_compatible_accepted_save_schema::<E>())
162 {
163 Ok(authority) => authority,
164 Err(error) => {
165 self.with_metrics(|| record_exec_error_for_path(ExecKind::Save, E::PATH, &error));
166
167 return Err(error);
168 }
169 };
170 let value = self.with_metrics(|| {
171 op(self.save_executor::<E>(
172 contract,
173 schema_info,
174 schema_fingerprint,
175 accepted_row_constraints,
176 ))
177 })?;
178
179 Ok(map(value))
180 }
181
182 fn execute_save_with_checked_accepted_row_contract<E, T, R>(
187 &self,
188 accepted_row_decode_contract: AcceptedRowDecodeContract,
189 accepted_schema_info: SchemaInfo,
190 accepted_schema_fingerprint: CommitSchemaFingerprint,
191 accepted_row_constraints: CompiledAcceptedRowConstraints,
192 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
193 map: impl FnOnce(T) -> R,
194 ) -> Result<R, InternalError>
195 where
196 E: PersistedRow<Canister = C>,
197 {
198 let value = self.with_metrics(|| {
199 op(self.save_executor::<E>(
200 accepted_row_decode_contract,
201 accepted_schema_info,
202 accepted_schema_fingerprint,
203 accepted_row_constraints,
204 ))
205 })?;
206
207 Ok(map(value))
208 }
209
210 fn execute_save_entity<E>(
212 &self,
213 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
214 ) -> Result<E, InternalError>
215 where
216 E: PersistedRow<Canister = C>,
217 {
218 self.execute_save_with(op, std::convert::identity)
219 }
220
221 fn execute_save_batch<E>(
222 &self,
223 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
224 ) -> Result<WriteBatchResponse<E>, InternalError>
225 where
226 E: PersistedRow<Canister = C>,
227 {
228 self.execute_save_with(op, WriteBatchResponse::new)
229 }
230
231 #[must_use]
237 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
238 where
239 E: EntityKind<Canister = C>,
240 {
241 self.fluent_load_query(MissingRowPolicy::Ignore)
242 }
243
244 #[must_use]
246 pub const fn load_with_consistency<E>(
247 &self,
248 consistency: MissingRowPolicy,
249 ) -> FluentLoadQuery<'_, E>
250 where
251 E: EntityKind<Canister = C>,
252 {
253 self.fluent_load_query(consistency)
254 }
255
256 #[must_use]
258 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
259 where
260 E: PersistedRow<Canister = C>,
261 {
262 self.fluent_delete_query(MissingRowPolicy::Ignore)
263 }
264
265 #[must_use]
267 pub fn delete_with_consistency<E>(
268 &self,
269 consistency: MissingRowPolicy,
270 ) -> FluentDeleteQuery<'_, E>
271 where
272 E: PersistedRow<Canister = C>,
273 {
274 self.fluent_delete_query(consistency)
275 }
276
277 #[must_use]
281 pub const fn select_one(&self) -> Value {
282 Value::Int64(1)
283 }
284
285 #[must_use]
290 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
291 where
292 E: EntityKind<Canister = C> + EntityValue,
293 {
294 LoadExecutor::new(self.db, self.debug)
295 }
296
297 #[must_use]
298 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
299 where
300 E: PersistedRow<Canister = C>,
301 {
302 DeleteExecutor::new(self.db)
303 }
304
305 #[must_use]
306 pub(in crate::db) const fn save_executor<E>(
307 &self,
308 accepted_row_decode_contract: AcceptedRowDecodeContract,
309 accepted_schema_info: SchemaInfo,
310 accepted_schema_fingerprint: CommitSchemaFingerprint,
311 accepted_row_constraints: CompiledAcceptedRowConstraints,
312 ) -> SaveExecutor<E>
313 where
314 E: PersistedRow<Canister = C>,
315 {
316 SaveExecutor::new_with_accepted_contract(
317 self.db,
318 self.debug,
319 accepted_row_decode_contract,
320 accepted_schema_info,
321 accepted_schema_fingerprint,
322 accepted_row_constraints,
323 )
324 }
325}