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