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, SchemaInfo},
27 },
28 error::InternalError,
29 metrics::sink::{ExecKind, MetricsSink, record_exec_error_for_path, with_metrics_sink},
30 traits::{CanisterKind, EntityKind, EntityValue},
31 value::Value,
32};
33use std::thread::LocalKey;
34
35#[cfg(all(test, feature = "sql"))]
36pub(in crate::db) use accepted_schema::AcceptedCatalogRuntimeCounterSnapshot;
37pub(in crate::db) use accepted_schema::AcceptedSchemaCatalogContext;
38pub(in crate::db) use query::{
39 AcceptedExecutionOutput, AcceptedIdValuesOutput, AcceptedOptionalValueOutput,
40 AcceptedValuesOutput,
41};
42#[cfg(feature = "diagnostics")]
43pub use query::{
44 DirectDataRowAttribution, FluentTerminalExecutionAttribution, GroupedCountAttribution,
45 GroupedExecutionAttribution, KernelRowAttribution, QueryExecutionAttribution,
46 ScalarAggregateAttribution,
47};
48pub(in crate::db) use response::finalize_scalar_paged_execution;
49pub(in crate::db) use response::finalize_structural_grouped_projection_result;
50#[cfg(feature = "sql")]
51pub(in crate::db) use response::sql_grouped_cursor_from_bytes;
52#[cfg(feature = "sql")]
53pub use sql::{
54 SqlAdminBulkDeletePlan, SqlAdminBulkUpdatePlan, SqlDdlExecutionStatus, SqlDdlMutationKind,
55 SqlDdlPreparationReport, SqlDeleteExposurePolicy, SqlDeletePolicyContext,
56 SqlDeletePolicyRejection, SqlDeletePolicyReport, SqlDeleteStatementClassification,
57 SqlPublicBoundedDeletePlan, SqlPublicBoundedUpdatePlan, SqlPublicPrimaryKeyDeletePlan,
58 SqlPublicPrimaryKeyUpdatePlan, SqlSessionCurrentDeletePlan, SqlSessionCurrentUpdatePlan,
59 SqlStatementDispatch, SqlStatementResult, SqlStatementShellSurface, SqlStatementSurface,
60 SqlUpdateAssignmentPolicy, SqlUpdateExposurePolicy, SqlUpdatePolicyContext,
61 SqlUpdatePolicyRejection, SqlUpdatePolicyReport, SqlUpdateStatementClassification,
62 SqlValidatedDeletePlan, SqlValidatedUpdatePlan, SqlWriteExecutionBounds, SqlWriteOrderProof,
63 SqlWriteReturningBounds, SqlWriteReturningShape, SqlWriteStatementShape, SqlWriteWhereProof,
64 classify_sql_delete_policy, classify_sql_update_policy, sql_statement_dispatch,
65 sql_statement_entity_name, sql_statement_shell_surface, sql_statement_surface,
66};
67#[cfg(all(feature = "sql", feature = "diagnostics"))]
68pub use sql::{
69 SqlCompileAttribution, SqlExecutionAttribution, SqlHybridCoveringAttribution,
70 SqlOutputBlobAttribution, SqlPureCoveringAttribution, SqlQueryCacheAttribution,
71 SqlQueryExecutionAttribution, SqlScalarAggregateAttribution,
72};
73#[cfg(all(feature = "sql", feature = "diagnostics"))]
74pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
75
76pub struct DbSession<C: CanisterKind> {
83 db: Db<C>,
84 debug: bool,
85 metrics: Option<&'static dyn MetricsSink>,
86}
87
88impl<C: CanisterKind> DbSession<C> {
89 #[must_use]
91 pub(crate) const fn new(db: Db<C>) -> Self {
92 Self {
93 db,
94 debug: false,
95 metrics: None,
96 }
97 }
98
99 #[must_use]
101 pub const fn new_with_hooks(
102 store: &'static LocalKey<StoreRegistry>,
103 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
104 ) -> Self {
105 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
106 }
107
108 #[must_use]
110 pub const fn debug(mut self) -> Self {
111 self.debug = true;
112 self
113 }
114
115 #[must_use]
117 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
118 self.metrics = Some(sink);
119 self
120 }
121
122 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
125 where
126 E: EntityKind<Canister = C>,
127 {
128 FluentLoadQuery::new(self, Query::new(consistency))
129 }
130
131 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
135 where
136 E: PersistedRow<Canister = C>,
137 {
138 FluentDeleteQuery::new(self, Query::new(consistency).delete())
139 }
140
141 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
142 if let Some(sink) = self.metrics {
143 with_metrics_sink(sink, f)
144 } else {
145 f()
146 }
147 }
148
149 fn execute_save_with<E, T, R>(
151 &self,
152 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
153 map: impl FnOnce(T) -> R,
154 ) -> Result<R, InternalError>
155 where
156 E: PersistedRow<Canister = C> + EntityValue,
157 {
158 let (contract, schema_info, schema_fingerprint) = match self
159 .with_metrics(|| self.ensure_generated_compatible_accepted_save_schema::<E>())
160 {
161 Ok(authority) => authority,
162 Err(error) => {
163 self.with_metrics(|| record_exec_error_for_path(ExecKind::Save, E::PATH, &error));
164
165 return Err(error);
166 }
167 };
168 let value = self.with_metrics(|| {
169 op(self.save_executor::<E>(contract, schema_info, schema_fingerprint))
170 })?;
171
172 Ok(map(value))
173 }
174
175 fn execute_save_with_checked_accepted_row_contract<E, T, R>(
180 &self,
181 accepted_row_decode_contract: AcceptedRowDecodeContract,
182 accepted_schema_info: SchemaInfo,
183 accepted_schema_fingerprint: CommitSchemaFingerprint,
184 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
185 map: impl FnOnce(T) -> R,
186 ) -> Result<R, InternalError>
187 where
188 E: PersistedRow<Canister = C> + EntityValue,
189 {
190 let value = self.with_metrics(|| {
191 op(self.save_executor::<E>(
192 accepted_row_decode_contract,
193 accepted_schema_info,
194 accepted_schema_fingerprint,
195 ))
196 })?;
197
198 Ok(map(value))
199 }
200
201 fn execute_save_entity<E>(
203 &self,
204 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
205 ) -> Result<E, InternalError>
206 where
207 E: PersistedRow<Canister = C> + EntityValue,
208 {
209 self.execute_save_with(op, std::convert::identity)
210 }
211
212 fn execute_save_batch<E>(
213 &self,
214 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
215 ) -> Result<WriteBatchResponse<E>, InternalError>
216 where
217 E: PersistedRow<Canister = C> + EntityValue,
218 {
219 self.execute_save_with(op, WriteBatchResponse::new)
220 }
221
222 #[must_use]
228 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
229 where
230 E: EntityKind<Canister = C>,
231 {
232 self.fluent_load_query(MissingRowPolicy::Ignore)
233 }
234
235 #[must_use]
237 pub const fn load_with_consistency<E>(
238 &self,
239 consistency: MissingRowPolicy,
240 ) -> FluentLoadQuery<'_, E>
241 where
242 E: EntityKind<Canister = C>,
243 {
244 self.fluent_load_query(consistency)
245 }
246
247 #[must_use]
249 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
250 where
251 E: PersistedRow<Canister = C>,
252 {
253 self.fluent_delete_query(MissingRowPolicy::Ignore)
254 }
255
256 #[must_use]
258 pub fn delete_with_consistency<E>(
259 &self,
260 consistency: MissingRowPolicy,
261 ) -> FluentDeleteQuery<'_, E>
262 where
263 E: PersistedRow<Canister = C>,
264 {
265 self.fluent_delete_query(consistency)
266 }
267
268 #[must_use]
272 pub const fn select_one(&self) -> Value {
273 Value::Int64(1)
274 }
275
276 #[must_use]
281 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
282 where
283 E: EntityKind<Canister = C> + EntityValue,
284 {
285 LoadExecutor::new(self.db, self.debug)
286 }
287
288 #[must_use]
289 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
290 where
291 E: PersistedRow<Canister = C> + EntityValue,
292 {
293 DeleteExecutor::new(self.db)
294 }
295
296 #[must_use]
297 pub(in crate::db) const fn save_executor<E>(
298 &self,
299 accepted_row_decode_contract: AcceptedRowDecodeContract,
300 accepted_schema_info: SchemaInfo,
301 accepted_schema_fingerprint: CommitSchemaFingerprint,
302 ) -> SaveExecutor<E>
303 where
304 E: PersistedRow<Canister = C> + EntityValue,
305 {
306 SaveExecutor::new_with_accepted_contract(
307 self.db,
308 self.debug,
309 accepted_row_decode_contract,
310 accepted_schema_info,
311 accepted_schema_fingerprint,
312 )
313 }
314}