1mod query;
7mod response;
8#[cfg(feature = "sql")]
9mod sql;
10#[cfg(all(test, feature = "sql"))]
14mod tests;
15mod write;
16
17use crate::{
18 db::{
19 Db, EntityFieldDescription, EntityRuntimeHooks, EntitySchemaDescription, FluentDeleteQuery,
20 FluentLoadQuery, IndexState, IntegrityReport, MigrationPlan, MigrationRegistry,
21 MigrationRunOutcome, MissingRowPolicy, PersistedRow, Query, QueryError,
22 SchemaMigrationDescriptor, SchemaMigrationExecutionOutcome, SchemaMigrationPlanner,
23 StorageReport, StoreRegistry, WriteBatchResponse,
24 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
25 query::plan::VisibleIndexes,
26 schema::{
27 describe_entity_fields, describe_entity_model, show_indexes_for_model,
28 show_indexes_for_model_with_runtime_state,
29 },
30 },
31 error::InternalError,
32 metrics::sink::{MetricsSink, with_metrics_sink},
33 model::entity::EntityModel,
34 traits::{CanisterKind, EntityKind, EntityValue, Path},
35 value::Value,
36};
37use std::thread::LocalKey;
38
39#[cfg(feature = "diagnostics")]
40pub use query::QueryExecutionAttribution;
41pub(in crate::db) use response::{
42 finalize_grouped_paged_execution, finalize_scalar_paged_execution,
43 sql_grouped_cursor_from_bytes,
44};
45#[cfg(all(feature = "sql", feature = "diagnostics"))]
46pub use sql::SqlQueryExecutionAttribution;
47#[cfg(feature = "sql")]
48pub use sql::SqlStatementResult;
49#[cfg(all(feature = "sql", feature = "diagnostics"))]
50pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
51
52pub struct DbSession<C: CanisterKind> {
59 db: Db<C>,
60 debug: bool,
61 metrics: Option<&'static dyn MetricsSink>,
62}
63
64impl<C: CanisterKind> DbSession<C> {
65 #[must_use]
67 pub(crate) const fn new(db: Db<C>) -> Self {
68 Self {
69 db,
70 debug: false,
71 metrics: None,
72 }
73 }
74
75 #[must_use]
77 pub const fn new_with_hooks(
78 store: &'static LocalKey<StoreRegistry>,
79 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
80 ) -> Self {
81 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
82 }
83
84 #[must_use]
86 pub const fn debug(mut self) -> Self {
87 self.debug = true;
88 self
89 }
90
91 #[must_use]
93 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
94 self.metrics = Some(sink);
95 self
96 }
97
98 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
101 where
102 E: EntityKind<Canister = C>,
103 {
104 FluentLoadQuery::new(self, Query::new(consistency))
105 }
106
107 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
111 where
112 E: PersistedRow<Canister = C>,
113 {
114 FluentDeleteQuery::new(self, Query::new(consistency).delete())
115 }
116
117 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
118 if let Some(sink) = self.metrics {
119 with_metrics_sink(sink, f)
120 } else {
121 f()
122 }
123 }
124
125 fn execute_save_with<E, T, R>(
127 &self,
128 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
129 map: impl FnOnce(T) -> R,
130 ) -> Result<R, InternalError>
131 where
132 E: PersistedRow<Canister = C> + EntityValue,
133 {
134 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
135
136 Ok(map(value))
137 }
138
139 fn execute_save_entity<E>(
141 &self,
142 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
143 ) -> Result<E, InternalError>
144 where
145 E: PersistedRow<Canister = C> + EntityValue,
146 {
147 self.execute_save_with(op, std::convert::identity)
148 }
149
150 fn execute_save_batch<E>(
151 &self,
152 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
153 ) -> Result<WriteBatchResponse<E>, InternalError>
154 where
155 E: PersistedRow<Canister = C> + EntityValue,
156 {
157 self.execute_save_with(op, WriteBatchResponse::new)
158 }
159
160 #[must_use]
166 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
167 where
168 E: EntityKind<Canister = C>,
169 {
170 self.fluent_load_query(MissingRowPolicy::Ignore)
171 }
172
173 #[must_use]
175 pub const fn load_with_consistency<E>(
176 &self,
177 consistency: MissingRowPolicy,
178 ) -> FluentLoadQuery<'_, E>
179 where
180 E: EntityKind<Canister = C>,
181 {
182 self.fluent_load_query(consistency)
183 }
184
185 #[must_use]
187 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
188 where
189 E: PersistedRow<Canister = C>,
190 {
191 self.fluent_delete_query(MissingRowPolicy::Ignore)
192 }
193
194 #[must_use]
196 pub fn delete_with_consistency<E>(
197 &self,
198 consistency: MissingRowPolicy,
199 ) -> FluentDeleteQuery<'_, E>
200 where
201 E: PersistedRow<Canister = C>,
202 {
203 self.fluent_delete_query(consistency)
204 }
205
206 #[must_use]
210 pub const fn select_one(&self) -> Value {
211 Value::Int(1)
212 }
213
214 #[must_use]
221 pub fn show_indexes<E>(&self) -> Vec<String>
222 where
223 E: EntityKind<Canister = C>,
224 {
225 self.show_indexes_for_store_model(E::Store::PATH, E::MODEL)
226 }
227
228 #[must_use]
234 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
235 show_indexes_for_model(model)
236 }
237
238 pub(in crate::db) fn show_indexes_for_store_model(
242 &self,
243 store_path: &str,
244 model: &'static EntityModel,
245 ) -> Vec<String> {
246 let runtime_state = self
247 .db
248 .with_store_registry(|registry| registry.try_get_store(store_path).ok())
249 .map(|store| store.index_state());
250
251 show_indexes_for_model_with_runtime_state(model, runtime_state)
252 }
253
254 #[must_use]
256 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
257 where
258 E: EntityKind<Canister = C>,
259 {
260 self.show_columns_for_model(E::MODEL)
261 }
262
263 #[must_use]
265 pub fn show_columns_for_model(
266 &self,
267 model: &'static EntityModel,
268 ) -> Vec<EntityFieldDescription> {
269 describe_entity_fields(model)
270 }
271
272 #[must_use]
274 pub fn show_entities(&self) -> Vec<String> {
275 self.db.runtime_entity_names()
276 }
277
278 #[must_use]
283 pub fn show_tables(&self) -> Vec<String> {
284 self.show_entities()
285 }
286
287 fn visible_indexes_for_store_model(
290 &self,
291 store_path: &str,
292 model: &'static EntityModel,
293 ) -> Result<VisibleIndexes<'static>, QueryError> {
294 let store = self
297 .db
298 .recovered_store(store_path)
299 .map_err(QueryError::execute)?;
300 let state = store.index_state();
301 if state != IndexState::Ready {
302 return Ok(VisibleIndexes::none());
303 }
304 debug_assert_eq!(state, IndexState::Ready);
305
306 Ok(VisibleIndexes::planner_visible(model.indexes()))
309 }
310
311 #[must_use]
316 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
317 where
318 E: EntityKind<Canister = C>,
319 {
320 self.describe_entity_model(E::MODEL)
321 }
322
323 #[must_use]
325 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
326 describe_entity_model(model)
327 }
328
329 pub fn storage_report(
331 &self,
332 name_to_path: &[(&'static str, &'static str)],
333 ) -> Result<StorageReport, InternalError> {
334 self.db.storage_report(name_to_path)
335 }
336
337 pub fn storage_report_default(&self) -> Result<StorageReport, InternalError> {
339 self.db.storage_report_default()
340 }
341
342 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
344 self.db.integrity_report()
345 }
346
347 pub fn execute_migration_plan(
352 &self,
353 plan: &MigrationPlan,
354 max_steps: usize,
355 ) -> Result<MigrationRunOutcome, InternalError> {
356 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
357 }
358
359 pub fn execute_schema_migration_descriptor(
365 &self,
366 registry: &mut MigrationRegistry,
367 planner: &SchemaMigrationPlanner,
368 descriptor: &SchemaMigrationDescriptor,
369 max_steps: usize,
370 ) -> Result<SchemaMigrationExecutionOutcome, InternalError> {
371 self.with_metrics(|| {
372 self.db
373 .execute_schema_migration_descriptor(registry, planner, descriptor, max_steps)
374 })
375 }
376
377 #[must_use]
382 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
383 where
384 E: EntityKind<Canister = C> + EntityValue,
385 {
386 LoadExecutor::new(self.db, self.debug)
387 }
388
389 #[must_use]
390 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
391 where
392 E: PersistedRow<Canister = C> + EntityValue,
393 {
394 DeleteExecutor::new(self.db)
395 }
396
397 #[must_use]
398 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
399 where
400 E: PersistedRow<Canister = C> + EntityValue,
401 {
402 SaveExecutor::new(self.db, self.debug)
403 }
404}