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