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, MissingRowPolicy, PersistedRow, Query,
21 QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
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::{
39 DirectDataRowAttribution, GroupedCountAttribution, GroupedExecutionAttribution,
40 QueryExecutionAttribution,
41};
42pub(in crate::db) use response::finalize_structural_grouped_projection_result;
43pub(in crate::db) use response::{finalize_scalar_paged_execution, sql_grouped_cursor_from_bytes};
44#[cfg(feature = "sql")]
45pub use sql::SqlStatementResult;
46#[cfg(all(feature = "sql", feature = "diagnostics"))]
47pub use sql::{
48 SqlCompileAttribution, SqlExecutionAttribution, SqlPureCoveringAttribution,
49 SqlQueryCacheAttribution, SqlQueryExecutionAttribution, SqlScalarAggregateAttribution,
50};
51#[cfg(all(feature = "sql", feature = "diagnostics"))]
52pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
53
54pub struct DbSession<C: CanisterKind> {
61 db: Db<C>,
62 debug: bool,
63 metrics: Option<&'static dyn MetricsSink>,
64}
65
66impl<C: CanisterKind> DbSession<C> {
67 #[must_use]
69 pub(crate) const fn new(db: Db<C>) -> Self {
70 Self {
71 db,
72 debug: false,
73 metrics: None,
74 }
75 }
76
77 #[must_use]
79 pub const fn new_with_hooks(
80 store: &'static LocalKey<StoreRegistry>,
81 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
82 ) -> Self {
83 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
84 }
85
86 #[must_use]
88 pub const fn debug(mut self) -> Self {
89 self.debug = true;
90 self
91 }
92
93 #[must_use]
95 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
96 self.metrics = Some(sink);
97 self
98 }
99
100 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
103 where
104 E: EntityKind<Canister = C>,
105 {
106 FluentLoadQuery::new(self, Query::new(consistency))
107 }
108
109 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
113 where
114 E: PersistedRow<Canister = C>,
115 {
116 FluentDeleteQuery::new(self, Query::new(consistency).delete())
117 }
118
119 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
120 if let Some(sink) = self.metrics {
121 with_metrics_sink(sink, f)
122 } else {
123 f()
124 }
125 }
126
127 fn execute_save_with<E, T, R>(
129 &self,
130 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
131 map: impl FnOnce(T) -> R,
132 ) -> Result<R, InternalError>
133 where
134 E: PersistedRow<Canister = C> + EntityValue,
135 {
136 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
137
138 Ok(map(value))
139 }
140
141 fn execute_save_entity<E>(
143 &self,
144 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
145 ) -> Result<E, InternalError>
146 where
147 E: PersistedRow<Canister = C> + EntityValue,
148 {
149 self.execute_save_with(op, std::convert::identity)
150 }
151
152 fn execute_save_batch<E>(
153 &self,
154 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
155 ) -> Result<WriteBatchResponse<E>, InternalError>
156 where
157 E: PersistedRow<Canister = C> + EntityValue,
158 {
159 self.execute_save_with(op, WriteBatchResponse::new)
160 }
161
162 #[must_use]
168 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
169 where
170 E: EntityKind<Canister = C>,
171 {
172 self.fluent_load_query(MissingRowPolicy::Ignore)
173 }
174
175 #[must_use]
177 pub const fn load_with_consistency<E>(
178 &self,
179 consistency: MissingRowPolicy,
180 ) -> FluentLoadQuery<'_, E>
181 where
182 E: EntityKind<Canister = C>,
183 {
184 self.fluent_load_query(consistency)
185 }
186
187 #[must_use]
189 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
190 where
191 E: PersistedRow<Canister = C>,
192 {
193 self.fluent_delete_query(MissingRowPolicy::Ignore)
194 }
195
196 #[must_use]
198 pub fn delete_with_consistency<E>(
199 &self,
200 consistency: MissingRowPolicy,
201 ) -> FluentDeleteQuery<'_, E>
202 where
203 E: PersistedRow<Canister = C>,
204 {
205 self.fluent_delete_query(consistency)
206 }
207
208 #[must_use]
212 pub const fn select_one(&self) -> Value {
213 Value::Int(1)
214 }
215
216 #[must_use]
223 pub fn show_indexes<E>(&self) -> Vec<String>
224 where
225 E: EntityKind<Canister = C>,
226 {
227 self.show_indexes_for_store_model(E::Store::PATH, E::MODEL)
228 }
229
230 #[must_use]
236 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
237 show_indexes_for_model(model)
238 }
239
240 pub(in crate::db) fn show_indexes_for_store_model(
244 &self,
245 store_path: &str,
246 model: &'static EntityModel,
247 ) -> Vec<String> {
248 let runtime_state = self
249 .db
250 .with_store_registry(|registry| registry.try_get_store(store_path).ok())
251 .map(|store| store.index_state());
252
253 show_indexes_for_model_with_runtime_state(model, runtime_state)
254 }
255
256 #[must_use]
258 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
259 where
260 E: EntityKind<Canister = C>,
261 {
262 self.show_columns_for_model(E::MODEL)
263 }
264
265 #[must_use]
267 pub fn show_columns_for_model(
268 &self,
269 model: &'static EntityModel,
270 ) -> Vec<EntityFieldDescription> {
271 describe_entity_fields(model)
272 }
273
274 #[must_use]
276 pub fn show_entities(&self) -> Vec<String> {
277 self.db.runtime_entity_names()
278 }
279
280 #[must_use]
285 pub fn show_tables(&self) -> Vec<String> {
286 self.show_entities()
287 }
288
289 fn visible_indexes_for_store_model(
292 &self,
293 store_path: &str,
294 model: &'static EntityModel,
295 ) -> Result<VisibleIndexes<'static>, QueryError> {
296 let store = self
299 .db
300 .recovered_store(store_path)
301 .map_err(QueryError::execute)?;
302 let state = store.index_state();
303 if state != IndexState::Ready {
304 return Ok(VisibleIndexes::none());
305 }
306 debug_assert_eq!(state, IndexState::Ready);
307
308 Ok(VisibleIndexes::planner_visible(model.indexes()))
311 }
312
313 #[must_use]
318 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
319 where
320 E: EntityKind<Canister = C>,
321 {
322 self.describe_entity_model(E::MODEL)
323 }
324
325 #[must_use]
327 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
328 describe_entity_model(model)
329 }
330
331 pub fn storage_report(
333 &self,
334 name_to_path: &[(&'static str, &'static str)],
335 ) -> Result<StorageReport, InternalError> {
336 self.db.storage_report(name_to_path)
337 }
338
339 pub fn storage_report_default(&self) -> Result<StorageReport, InternalError> {
341 self.db.storage_report_default()
342 }
343
344 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
346 self.db.integrity_report()
347 }
348
349 #[must_use]
354 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
355 where
356 E: EntityKind<Canister = C> + EntityValue,
357 {
358 LoadExecutor::new(self.db, self.debug)
359 }
360
361 #[must_use]
362 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
363 where
364 E: PersistedRow<Canister = C> + EntityValue,
365 {
366 DeleteExecutor::new(self.db)
367 }
368
369 #[must_use]
370 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
371 where
372 E: PersistedRow<Canister = C> + EntityValue,
373 {
374 SaveExecutor::new(self.db, self.debug)
375 }
376}