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,
20 PlanError, Query, 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: PersistedRow<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: PersistedRow<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: PersistedRow<Canister = C> + EntityValue,
135 {
136 self.execute_save_with(op, WriteBatchResponse::new)
137 }
138
139 #[must_use]
145 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
146 where
147 E: EntityKind<Canister = C>,
148 {
149 FluentLoadQuery::new(self, Query::new(MissingRowPolicy::Ignore))
150 }
151
152 #[must_use]
154 pub const fn load_with_consistency<E>(
155 &self,
156 consistency: MissingRowPolicy,
157 ) -> FluentLoadQuery<'_, E>
158 where
159 E: EntityKind<Canister = C>,
160 {
161 FluentLoadQuery::new(self, Query::new(consistency))
162 }
163
164 #[must_use]
166 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
167 where
168 E: PersistedRow<Canister = C>,
169 {
170 FluentDeleteQuery::new(self, Query::new(MissingRowPolicy::Ignore).delete())
171 }
172
173 #[must_use]
175 pub fn delete_with_consistency<E>(
176 &self,
177 consistency: MissingRowPolicy,
178 ) -> FluentDeleteQuery<'_, E>
179 where
180 E: PersistedRow<Canister = C>,
181 {
182 FluentDeleteQuery::new(self, Query::new(consistency).delete())
183 }
184
185 #[must_use]
189 pub const fn select_one(&self) -> Value {
190 Value::Int(1)
191 }
192
193 #[must_use]
200 pub fn show_indexes<E>(&self) -> Vec<String>
201 where
202 E: EntityKind<Canister = C>,
203 {
204 self.show_indexes_for_model(E::MODEL)
205 }
206
207 #[must_use]
209 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
210 show_indexes_for_model(model)
211 }
212
213 #[must_use]
215 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
216 where
217 E: EntityKind<Canister = C>,
218 {
219 self.show_columns_for_model(E::MODEL)
220 }
221
222 #[must_use]
224 pub fn show_columns_for_model(
225 &self,
226 model: &'static EntityModel,
227 ) -> Vec<EntityFieldDescription> {
228 describe_entity_model(model).fields().to_vec()
229 }
230
231 #[must_use]
233 pub fn show_entities(&self) -> Vec<String> {
234 self.db.runtime_entity_names()
235 }
236
237 #[must_use]
242 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
243 where
244 E: EntityKind<Canister = C>,
245 {
246 self.describe_entity_model(E::MODEL)
247 }
248
249 #[must_use]
251 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
252 describe_entity_model(model)
253 }
254
255 pub fn storage_report(
257 &self,
258 name_to_path: &[(&'static str, &'static str)],
259 ) -> Result<StorageReport, InternalError> {
260 self.db.storage_report(name_to_path)
261 }
262
263 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
265 self.db.integrity_report()
266 }
267
268 pub fn execute_migration_plan(
273 &self,
274 plan: &MigrationPlan,
275 max_steps: usize,
276 ) -> Result<MigrationRunOutcome, InternalError> {
277 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
278 }
279
280 #[must_use]
285 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
286 where
287 E: EntityKind<Canister = C> + EntityValue,
288 {
289 LoadExecutor::new(self.db, self.debug)
290 }
291
292 #[must_use]
293 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
294 where
295 E: PersistedRow<Canister = C> + EntityValue,
296 {
297 DeleteExecutor::new(self.db, self.debug)
298 }
299
300 #[must_use]
301 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
302 where
303 E: PersistedRow<Canister = C> + EntityValue,
304 {
305 SaveExecutor::new(self.db, self.debug)
306 }
307}