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::finalize_structural_grouped_projection_result;
42pub(in crate::db) use response::{finalize_scalar_paged_execution, sql_grouped_cursor_from_bytes};
43#[cfg(all(feature = "sql", feature = "diagnostics"))]
44pub use sql::SqlQueryExecutionAttribution;
45#[cfg(feature = "sql")]
46pub use sql::SqlStatementResult;
47#[cfg(all(feature = "sql", feature = "diagnostics"))]
48pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
49
50pub struct DbSession<C: CanisterKind> {
57 db: Db<C>,
58 debug: bool,
59 metrics: Option<&'static dyn MetricsSink>,
60}
61
62impl<C: CanisterKind> DbSession<C> {
63 #[must_use]
65 pub(crate) const fn new(db: Db<C>) -> Self {
66 Self {
67 db,
68 debug: false,
69 metrics: None,
70 }
71 }
72
73 #[must_use]
75 pub const fn new_with_hooks(
76 store: &'static LocalKey<StoreRegistry>,
77 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
78 ) -> Self {
79 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
80 }
81
82 #[must_use]
84 pub const fn debug(mut self) -> Self {
85 self.debug = true;
86 self
87 }
88
89 #[must_use]
91 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
92 self.metrics = Some(sink);
93 self
94 }
95
96 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
99 where
100 E: EntityKind<Canister = C>,
101 {
102 FluentLoadQuery::new(self, Query::new(consistency))
103 }
104
105 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
109 where
110 E: PersistedRow<Canister = C>,
111 {
112 FluentDeleteQuery::new(self, Query::new(consistency).delete())
113 }
114
115 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
116 if let Some(sink) = self.metrics {
117 with_metrics_sink(sink, f)
118 } else {
119 f()
120 }
121 }
122
123 fn execute_save_with<E, T, R>(
125 &self,
126 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
127 map: impl FnOnce(T) -> R,
128 ) -> Result<R, InternalError>
129 where
130 E: PersistedRow<Canister = C> + EntityValue,
131 {
132 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
133
134 Ok(map(value))
135 }
136
137 fn execute_save_entity<E>(
139 &self,
140 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
141 ) -> Result<E, InternalError>
142 where
143 E: PersistedRow<Canister = C> + EntityValue,
144 {
145 self.execute_save_with(op, std::convert::identity)
146 }
147
148 fn execute_save_batch<E>(
149 &self,
150 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
151 ) -> Result<WriteBatchResponse<E>, InternalError>
152 where
153 E: PersistedRow<Canister = C> + EntityValue,
154 {
155 self.execute_save_with(op, WriteBatchResponse::new)
156 }
157
158 #[must_use]
164 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
165 where
166 E: EntityKind<Canister = C>,
167 {
168 self.fluent_load_query(MissingRowPolicy::Ignore)
169 }
170
171 #[must_use]
173 pub const fn load_with_consistency<E>(
174 &self,
175 consistency: MissingRowPolicy,
176 ) -> FluentLoadQuery<'_, E>
177 where
178 E: EntityKind<Canister = C>,
179 {
180 self.fluent_load_query(consistency)
181 }
182
183 #[must_use]
185 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
186 where
187 E: PersistedRow<Canister = C>,
188 {
189 self.fluent_delete_query(MissingRowPolicy::Ignore)
190 }
191
192 #[must_use]
194 pub fn delete_with_consistency<E>(
195 &self,
196 consistency: MissingRowPolicy,
197 ) -> FluentDeleteQuery<'_, E>
198 where
199 E: PersistedRow<Canister = C>,
200 {
201 self.fluent_delete_query(consistency)
202 }
203
204 #[must_use]
208 pub const fn select_one(&self) -> Value {
209 Value::Int(1)
210 }
211
212 #[must_use]
219 pub fn show_indexes<E>(&self) -> Vec<String>
220 where
221 E: EntityKind<Canister = C>,
222 {
223 self.show_indexes_for_store_model(E::Store::PATH, E::MODEL)
224 }
225
226 #[must_use]
232 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
233 show_indexes_for_model(model)
234 }
235
236 pub(in crate::db) fn show_indexes_for_store_model(
240 &self,
241 store_path: &str,
242 model: &'static EntityModel,
243 ) -> Vec<String> {
244 let runtime_state = self
245 .db
246 .with_store_registry(|registry| registry.try_get_store(store_path).ok())
247 .map(|store| store.index_state());
248
249 show_indexes_for_model_with_runtime_state(model, runtime_state)
250 }
251
252 #[must_use]
254 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
255 where
256 E: EntityKind<Canister = C>,
257 {
258 self.show_columns_for_model(E::MODEL)
259 }
260
261 #[must_use]
263 pub fn show_columns_for_model(
264 &self,
265 model: &'static EntityModel,
266 ) -> Vec<EntityFieldDescription> {
267 describe_entity_fields(model)
268 }
269
270 #[must_use]
272 pub fn show_entities(&self) -> Vec<String> {
273 self.db.runtime_entity_names()
274 }
275
276 #[must_use]
281 pub fn show_tables(&self) -> Vec<String> {
282 self.show_entities()
283 }
284
285 fn visible_indexes_for_store_model(
288 &self,
289 store_path: &str,
290 model: &'static EntityModel,
291 ) -> Result<VisibleIndexes<'static>, QueryError> {
292 let store = self
295 .db
296 .recovered_store(store_path)
297 .map_err(QueryError::execute)?;
298 let state = store.index_state();
299 if state != IndexState::Ready {
300 return Ok(VisibleIndexes::none());
301 }
302 debug_assert_eq!(state, IndexState::Ready);
303
304 Ok(VisibleIndexes::planner_visible(model.indexes()))
307 }
308
309 #[must_use]
314 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
315 where
316 E: EntityKind<Canister = C>,
317 {
318 self.describe_entity_model(E::MODEL)
319 }
320
321 #[must_use]
323 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
324 describe_entity_model(model)
325 }
326
327 pub fn storage_report(
329 &self,
330 name_to_path: &[(&'static str, &'static str)],
331 ) -> Result<StorageReport, InternalError> {
332 self.db.storage_report(name_to_path)
333 }
334
335 pub fn storage_report_default(&self) -> Result<StorageReport, InternalError> {
337 self.db.storage_report_default()
338 }
339
340 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
342 self.db.integrity_report()
343 }
344
345 pub fn execute_migration_plan(
350 &self,
351 plan: &MigrationPlan,
352 max_steps: usize,
353 ) -> Result<MigrationRunOutcome, InternalError> {
354 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
355 }
356
357 pub fn execute_schema_migration_descriptor(
363 &self,
364 registry: &mut MigrationRegistry,
365 planner: &SchemaMigrationPlanner,
366 descriptor: &SchemaMigrationDescriptor,
367 max_steps: usize,
368 ) -> Result<SchemaMigrationExecutionOutcome, InternalError> {
369 self.with_metrics(|| {
370 self.db
371 .execute_schema_migration_descriptor(registry, planner, descriptor, max_steps)
372 })
373 }
374
375 #[must_use]
380 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
381 where
382 E: EntityKind<Canister = C> + EntityValue,
383 {
384 LoadExecutor::new(self.db, self.debug)
385 }
386
387 #[must_use]
388 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
389 where
390 E: PersistedRow<Canister = C> + EntityValue,
391 {
392 DeleteExecutor::new(self.db)
393 }
394
395 #[must_use]
396 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
397 where
398 E: PersistedRow<Canister = C> + EntityValue,
399 {
400 SaveExecutor::new(self.db, self.debug)
401 }
402}