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, PlanError, Query,
20 QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
21 commit::EntityRuntimeHooks,
22 cursor::decode_optional_cursor_token,
23 executor::{DeleteExecutor, ExecutorPlanError, 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, SqlPreparedStatement, SqlStatementRoute};
36
37fn map_executor_plan_error(err: ExecutorPlanError) -> QueryError {
39 match err {
40 ExecutorPlanError::Cursor(err) => QueryError::from(PlanError::from(*err)),
41 }
42}
43
44fn decode_optional_cursor_bytes(cursor_token: Option<&str>) -> Result<Option<Vec<u8>>, QueryError> {
47 decode_optional_cursor_token(cursor_token).map_err(|err| QueryError::from(PlanError::from(err)))
48}
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 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
97 if let Some(sink) = self.metrics {
98 with_metrics_sink(sink, f)
99 } else {
100 f()
101 }
102 }
103
104 fn execute_save_with<E, T, R>(
106 &self,
107 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
108 map: impl FnOnce(T) -> R,
109 ) -> Result<R, InternalError>
110 where
111 E: EntityKind<Canister = C> + EntityValue,
112 {
113 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
114
115 Ok(map(value))
116 }
117
118 fn execute_save_entity<E>(
120 &self,
121 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
122 ) -> Result<E, InternalError>
123 where
124 E: EntityKind<Canister = C> + EntityValue,
125 {
126 self.execute_save_with(op, std::convert::identity)
127 }
128
129 fn execute_save_batch<E>(
130 &self,
131 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
132 ) -> Result<WriteBatchResponse<E>, InternalError>
133 where
134 E: EntityKind<Canister = C> + EntityValue,
135 {
136 self.execute_save_with(op, WriteBatchResponse::new)
137 }
138
139 fn execute_save_view<E>(
140 &self,
141 op: impl FnOnce(SaveExecutor<E>) -> Result<E::ViewType, InternalError>,
142 ) -> Result<E::ViewType, InternalError>
143 where
144 E: EntityKind<Canister = C> + EntityValue,
145 {
146 self.execute_save_with(op, std::convert::identity)
147 }
148
149 #[must_use]
155 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
156 where
157 E: EntityKind<Canister = C>,
158 {
159 FluentLoadQuery::new(self, Query::new(MissingRowPolicy::Ignore))
160 }
161
162 #[must_use]
164 pub const fn load_with_consistency<E>(
165 &self,
166 consistency: MissingRowPolicy,
167 ) -> FluentLoadQuery<'_, E>
168 where
169 E: EntityKind<Canister = C>,
170 {
171 FluentLoadQuery::new(self, Query::new(consistency))
172 }
173
174 #[must_use]
176 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
177 where
178 E: EntityKind<Canister = C>,
179 {
180 FluentDeleteQuery::new(self, Query::new(MissingRowPolicy::Ignore).delete())
181 }
182
183 #[must_use]
185 pub fn delete_with_consistency<E>(
186 &self,
187 consistency: MissingRowPolicy,
188 ) -> FluentDeleteQuery<'_, E>
189 where
190 E: EntityKind<Canister = C>,
191 {
192 FluentDeleteQuery::new(self, Query::new(consistency).delete())
193 }
194
195 #[must_use]
199 pub const fn select_one(&self) -> Value {
200 Value::Int(1)
201 }
202
203 #[must_use]
210 pub fn show_indexes<E>(&self) -> Vec<String>
211 where
212 E: EntityKind<Canister = C>,
213 {
214 self.show_indexes_for_model(E::MODEL)
215 }
216
217 #[must_use]
219 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
220 show_indexes_for_model(model)
221 }
222
223 #[must_use]
225 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
226 where
227 E: EntityKind<Canister = C>,
228 {
229 self.show_columns_for_model(E::MODEL)
230 }
231
232 #[must_use]
234 pub fn show_columns_for_model(
235 &self,
236 model: &'static EntityModel,
237 ) -> Vec<EntityFieldDescription> {
238 describe_entity_model(model).fields().to_vec()
239 }
240
241 #[must_use]
243 pub fn show_entities(&self) -> Vec<String> {
244 self.db.runtime_entity_names()
245 }
246
247 #[must_use]
252 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
253 where
254 E: EntityKind<Canister = C>,
255 {
256 self.describe_entity_model(E::MODEL)
257 }
258
259 #[must_use]
261 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
262 describe_entity_model(model)
263 }
264
265 pub fn storage_report(
267 &self,
268 name_to_path: &[(&'static str, &'static str)],
269 ) -> Result<StorageReport, InternalError> {
270 self.db.storage_report(name_to_path)
271 }
272
273 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
275 self.db.integrity_report()
276 }
277
278 pub fn execute_migration_plan(
283 &self,
284 plan: &MigrationPlan,
285 max_steps: usize,
286 ) -> Result<MigrationRunOutcome, InternalError> {
287 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
288 }
289
290 #[must_use]
295 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
296 where
297 E: EntityKind<Canister = C> + EntityValue,
298 {
299 LoadExecutor::new(self.db, self.debug)
300 }
301
302 #[must_use]
303 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
304 where
305 E: EntityKind<Canister = C> + EntityValue,
306 {
307 DeleteExecutor::new(self.db, self.debug)
308 }
309
310 #[must_use]
311 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
312 where
313 E: EntityKind<Canister = C> + EntityValue,
314 {
315 SaveExecutor::new(self.db, self.debug)
316 }
317}