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 IntegrityReport, MigrationPlan, MigrationRunOutcome, MissingRowPolicy, PersistedRow, Query,
20 QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
21 commit::EntityRuntimeHooks,
22 cursor::{decode_optional_cursor_token, decode_optional_grouped_cursor_token},
23 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
24 schema::{describe_entity_model, show_indexes_for_model},
25 },
26 error::InternalError,
27 metrics::sink::{MetricsSink, with_metrics_sink},
28 model::entity::EntityModel,
29 traits::{CanisterKind, EntityKind, EntityValue},
30 value::Value,
31};
32use std::thread::LocalKey;
33
34#[cfg(feature = "sql")]
35pub use sql::{SqlDispatchResult, SqlParsedStatement, SqlStatementRoute};
36
37fn decode_optional_cursor_bytes(cursor_token: Option<&str>) -> Result<Option<Vec<u8>>, QueryError> {
40 decode_optional_cursor_token(cursor_token).map_err(QueryError::from_cursor_plan_error)
41}
42
43fn decode_optional_grouped_cursor(
46 cursor_token: Option<&str>,
47) -> Result<Option<crate::db::cursor::GroupedContinuationToken>, QueryError> {
48 decode_optional_grouped_cursor_token(cursor_token).map_err(QueryError::from_cursor_plan_error)
49}
50
51pub struct DbSession<C: CanisterKind> {
58 db: Db<C>,
59 debug: bool,
60 metrics: Option<&'static dyn MetricsSink>,
61}
62
63impl<C: CanisterKind> DbSession<C> {
64 #[must_use]
66 pub(crate) const fn new(db: Db<C>) -> Self {
67 Self {
68 db,
69 debug: false,
70 metrics: None,
71 }
72 }
73
74 #[must_use]
76 pub const fn new_with_hooks(
77 store: &'static LocalKey<StoreRegistry>,
78 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
79 ) -> Self {
80 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
81 }
82
83 #[must_use]
85 pub const fn debug(mut self) -> Self {
86 self.debug = true;
87 self
88 }
89
90 #[must_use]
92 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
93 self.metrics = Some(sink);
94 self
95 }
96
97 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
98 if let Some(sink) = self.metrics {
99 with_metrics_sink(sink, f)
100 } else {
101 f()
102 }
103 }
104
105 fn execute_save_with<E, T, R>(
107 &self,
108 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
109 map: impl FnOnce(T) -> R,
110 ) -> Result<R, InternalError>
111 where
112 E: PersistedRow<Canister = C> + EntityValue,
113 {
114 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
115
116 Ok(map(value))
117 }
118
119 fn execute_save_entity<E>(
121 &self,
122 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
123 ) -> Result<E, InternalError>
124 where
125 E: PersistedRow<Canister = C> + EntityValue,
126 {
127 self.execute_save_with(op, std::convert::identity)
128 }
129
130 fn execute_save_batch<E>(
131 &self,
132 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
133 ) -> Result<WriteBatchResponse<E>, InternalError>
134 where
135 E: PersistedRow<Canister = C> + EntityValue,
136 {
137 self.execute_save_with(op, WriteBatchResponse::new)
138 }
139
140 #[must_use]
146 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
147 where
148 E: EntityKind<Canister = C>,
149 {
150 FluentLoadQuery::new(self, Query::new(MissingRowPolicy::Ignore))
151 }
152
153 #[must_use]
155 pub const fn load_with_consistency<E>(
156 &self,
157 consistency: MissingRowPolicy,
158 ) -> FluentLoadQuery<'_, E>
159 where
160 E: EntityKind<Canister = C>,
161 {
162 FluentLoadQuery::new(self, Query::new(consistency))
163 }
164
165 #[must_use]
167 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
168 where
169 E: PersistedRow<Canister = C>,
170 {
171 FluentDeleteQuery::new(self, Query::new(MissingRowPolicy::Ignore).delete())
172 }
173
174 #[must_use]
176 pub fn delete_with_consistency<E>(
177 &self,
178 consistency: MissingRowPolicy,
179 ) -> FluentDeleteQuery<'_, E>
180 where
181 E: PersistedRow<Canister = C>,
182 {
183 FluentDeleteQuery::new(self, Query::new(consistency).delete())
184 }
185
186 #[must_use]
190 pub const fn select_one(&self) -> Value {
191 Value::Int(1)
192 }
193
194 #[must_use]
201 pub fn show_indexes<E>(&self) -> Vec<String>
202 where
203 E: EntityKind<Canister = C>,
204 {
205 self.show_indexes_for_model(E::MODEL)
206 }
207
208 #[must_use]
210 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
211 show_indexes_for_model(model)
212 }
213
214 #[must_use]
216 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
217 where
218 E: EntityKind<Canister = C>,
219 {
220 self.show_columns_for_model(E::MODEL)
221 }
222
223 #[must_use]
225 pub fn show_columns_for_model(
226 &self,
227 model: &'static EntityModel,
228 ) -> Vec<EntityFieldDescription> {
229 describe_entity_model(model).fields().to_vec()
230 }
231
232 #[must_use]
234 pub fn show_entities(&self) -> Vec<String> {
235 self.db.runtime_entity_names()
236 }
237
238 #[must_use]
243 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
244 where
245 E: EntityKind<Canister = C>,
246 {
247 self.describe_entity_model(E::MODEL)
248 }
249
250 #[must_use]
252 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
253 describe_entity_model(model)
254 }
255
256 pub fn storage_report(
258 &self,
259 name_to_path: &[(&'static str, &'static str)],
260 ) -> Result<StorageReport, InternalError> {
261 self.db.storage_report(name_to_path)
262 }
263
264 pub fn storage_report_default(&self) -> Result<StorageReport, InternalError> {
266 self.db.storage_report_default()
267 }
268
269 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
271 self.db.integrity_report()
272 }
273
274 pub fn execute_migration_plan(
279 &self,
280 plan: &MigrationPlan,
281 max_steps: usize,
282 ) -> Result<MigrationRunOutcome, InternalError> {
283 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
284 }
285
286 #[must_use]
291 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
292 where
293 E: EntityKind<Canister = C> + EntityValue,
294 {
295 LoadExecutor::new(self.db, self.debug)
296 }
297
298 #[must_use]
299 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
300 where
301 E: PersistedRow<Canister = C> + EntityValue,
302 {
303 DeleteExecutor::new(self.db, self.debug)
304 }
305
306 #[must_use]
307 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
308 where
309 E: PersistedRow<Canister = C> + EntityValue,
310 {
311 SaveExecutor::new(self.db, self.debug)
312 }
313}