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