1mod query;
7#[cfg(feature = "sql")]
8mod sql;
9#[cfg(all(test, feature = "sql"))]
13mod tests;
14mod write;
15
16use crate::{
17 db::{
18 Db, EntityFieldDescription, EntitySchemaDescription, FluentDeleteQuery, FluentLoadQuery,
19 IndexState, IntegrityReport, MigrationPlan, MigrationRunOutcome, MissingRowPolicy,
20 PersistedRow, Query, QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
21 commit::EntityRuntimeHooks,
22 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
23 query::plan::VisibleIndexes,
24 schema::{
25 describe_entity_fields, describe_entity_model, show_indexes_for_model,
26 show_indexes_for_model_with_runtime_state,
27 },
28 },
29 error::InternalError,
30 metrics::sink::{MetricsSink, with_metrics_sink},
31 model::entity::EntityModel,
32 traits::{CanisterKind, EntityKind, EntityValue, Path},
33 value::Value,
34};
35use std::thread::LocalKey;
36
37#[cfg(feature = "diagnostics")]
38pub use query::QueryExecutionAttribution;
39#[cfg(all(feature = "sql", feature = "diagnostics"))]
40pub use sql::SqlQueryExecutionAttribution;
41#[cfg(feature = "sql")]
42pub use sql::SqlStatementResult;
43#[cfg(all(feature = "sql", feature = "diagnostics"))]
44pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
45
46pub struct DbSession<C: CanisterKind> {
53 db: Db<C>,
54 debug: bool,
55 metrics: Option<&'static dyn MetricsSink>,
56}
57
58impl<C: CanisterKind> DbSession<C> {
59 #[must_use]
61 pub(crate) const fn new(db: Db<C>) -> Self {
62 Self {
63 db,
64 debug: false,
65 metrics: None,
66 }
67 }
68
69 #[must_use]
71 pub const fn new_with_hooks(
72 store: &'static LocalKey<StoreRegistry>,
73 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
74 ) -> Self {
75 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
76 }
77
78 #[must_use]
80 pub const fn debug(mut self) -> Self {
81 self.debug = true;
82 self
83 }
84
85 #[must_use]
87 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
88 self.metrics = Some(sink);
89 self
90 }
91
92 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
95 where
96 E: EntityKind<Canister = C>,
97 {
98 FluentLoadQuery::new(self, Query::new(consistency))
99 }
100
101 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
105 where
106 E: PersistedRow<Canister = C>,
107 {
108 FluentDeleteQuery::new(self, Query::new(consistency).delete())
109 }
110
111 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
112 if let Some(sink) = self.metrics {
113 with_metrics_sink(sink, f)
114 } else {
115 f()
116 }
117 }
118
119 fn execute_save_with<E, T, R>(
121 &self,
122 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
123 map: impl FnOnce(T) -> R,
124 ) -> Result<R, InternalError>
125 where
126 E: PersistedRow<Canister = C> + EntityValue,
127 {
128 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
129
130 Ok(map(value))
131 }
132
133 fn execute_save_entity<E>(
135 &self,
136 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
137 ) -> Result<E, InternalError>
138 where
139 E: PersistedRow<Canister = C> + EntityValue,
140 {
141 self.execute_save_with(op, std::convert::identity)
142 }
143
144 fn execute_save_batch<E>(
145 &self,
146 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
147 ) -> Result<WriteBatchResponse<E>, InternalError>
148 where
149 E: PersistedRow<Canister = C> + EntityValue,
150 {
151 self.execute_save_with(op, WriteBatchResponse::new)
152 }
153
154 #[must_use]
160 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
161 where
162 E: EntityKind<Canister = C>,
163 {
164 self.fluent_load_query(MissingRowPolicy::Ignore)
165 }
166
167 #[must_use]
169 pub const fn load_with_consistency<E>(
170 &self,
171 consistency: MissingRowPolicy,
172 ) -> FluentLoadQuery<'_, E>
173 where
174 E: EntityKind<Canister = C>,
175 {
176 self.fluent_load_query(consistency)
177 }
178
179 #[must_use]
181 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
182 where
183 E: PersistedRow<Canister = C>,
184 {
185 self.fluent_delete_query(MissingRowPolicy::Ignore)
186 }
187
188 #[must_use]
190 pub fn delete_with_consistency<E>(
191 &self,
192 consistency: MissingRowPolicy,
193 ) -> FluentDeleteQuery<'_, E>
194 where
195 E: PersistedRow<Canister = C>,
196 {
197 self.fluent_delete_query(consistency)
198 }
199
200 #[must_use]
204 pub const fn select_one(&self) -> Value {
205 Value::Int(1)
206 }
207
208 #[must_use]
215 pub fn show_indexes<E>(&self) -> Vec<String>
216 where
217 E: EntityKind<Canister = C>,
218 {
219 self.show_indexes_for_store_model(E::Store::PATH, E::MODEL)
220 }
221
222 #[must_use]
228 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
229 show_indexes_for_model(model)
230 }
231
232 pub(in crate::db) fn show_indexes_for_store_model(
236 &self,
237 store_path: &str,
238 model: &'static EntityModel,
239 ) -> Vec<String> {
240 let runtime_state = self
241 .db
242 .with_store_registry(|registry| registry.try_get_store(store_path).ok())
243 .map(|store| store.index_state());
244
245 show_indexes_for_model_with_runtime_state(model, runtime_state)
246 }
247
248 #[must_use]
250 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
251 where
252 E: EntityKind<Canister = C>,
253 {
254 self.show_columns_for_model(E::MODEL)
255 }
256
257 #[must_use]
259 pub fn show_columns_for_model(
260 &self,
261 model: &'static EntityModel,
262 ) -> Vec<EntityFieldDescription> {
263 describe_entity_fields(model)
264 }
265
266 #[must_use]
268 pub fn show_entities(&self) -> Vec<String> {
269 self.db.runtime_entity_names()
270 }
271
272 #[must_use]
277 pub fn show_tables(&self) -> Vec<String> {
278 self.show_entities()
279 }
280
281 fn visible_indexes_for_store_model(
284 &self,
285 store_path: &str,
286 model: &'static EntityModel,
287 ) -> Result<VisibleIndexes<'static>, QueryError> {
288 let store = self
291 .db
292 .recovered_store(store_path)
293 .map_err(QueryError::execute)?;
294 let state = store.index_state();
295 if state != IndexState::Ready {
296 return Ok(VisibleIndexes::none());
297 }
298 debug_assert_eq!(state, IndexState::Ready);
299
300 Ok(VisibleIndexes::planner_visible(model.indexes()))
303 }
304
305 #[must_use]
310 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
311 where
312 E: EntityKind<Canister = C>,
313 {
314 self.describe_entity_model(E::MODEL)
315 }
316
317 #[must_use]
319 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
320 describe_entity_model(model)
321 }
322
323 pub fn storage_report(
325 &self,
326 name_to_path: &[(&'static str, &'static str)],
327 ) -> Result<StorageReport, InternalError> {
328 self.db.storage_report(name_to_path)
329 }
330
331 pub fn storage_report_default(&self) -> Result<StorageReport, InternalError> {
333 self.db.storage_report_default()
334 }
335
336 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
338 self.db.integrity_report()
339 }
340
341 pub fn execute_migration_plan(
346 &self,
347 plan: &MigrationPlan,
348 max_steps: usize,
349 ) -> Result<MigrationRunOutcome, InternalError> {
350 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
351 }
352
353 #[must_use]
358 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
359 where
360 E: EntityKind<Canister = C> + EntityValue,
361 {
362 LoadExecutor::new(self.db, self.debug)
363 }
364
365 #[must_use]
366 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
367 where
368 E: PersistedRow<Canister = C> + EntityValue,
369 {
370 DeleteExecutor::new(self.db)
371 }
372
373 #[must_use]
374 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
375 where
376 E: PersistedRow<Canister = C> + EntityValue,
377 {
378 SaveExecutor::new(self.db, self.debug)
379 }
380}