1mod query;
7mod sql;
8#[cfg(test)]
12mod tests;
13mod write;
14
15use crate::{
16 db::{
17 Db, EntityFieldDescription, EntitySchemaDescription, FluentDeleteQuery, FluentLoadQuery,
18 IntegrityReport, MigrationPlan, MigrationRunOutcome, MissingRowPolicy, PlanError, Query,
19 QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
20 commit::EntityRuntimeHooks,
21 cursor::decode_optional_cursor_token,
22 executor::{DeleteExecutor, ExecutorPlanError, LoadExecutor, SaveExecutor},
23 schema::{describe_entity_model, show_indexes_for_model},
24 },
25 error::InternalError,
26 metrics::sink::{MetricsSink, with_metrics_sink},
27 model::entity::EntityModel,
28 traits::{CanisterKind, EntityKind, EntityValue},
29 value::Value,
30};
31use std::thread::LocalKey;
32
33pub use sql::{SqlDispatchResult, SqlParsedStatement, SqlPreparedStatement, SqlStatementRoute};
34
35fn map_executor_plan_error(err: ExecutorPlanError) -> QueryError {
37 match err {
38 ExecutorPlanError::Cursor(err) => QueryError::from(PlanError::from(*err)),
39 }
40}
41
42fn decode_optional_cursor_bytes(cursor_token: Option<&str>) -> Result<Option<Vec<u8>>, QueryError> {
45 decode_optional_cursor_token(cursor_token).map_err(|err| QueryError::from(PlanError::from(err)))
46}
47
48pub struct DbSession<C: CanisterKind> {
55 db: Db<C>,
56 debug: bool,
57 metrics: Option<&'static dyn MetricsSink>,
58}
59
60impl<C: CanisterKind> DbSession<C> {
61 #[must_use]
63 pub(crate) const fn new(db: Db<C>) -> Self {
64 Self {
65 db,
66 debug: false,
67 metrics: None,
68 }
69 }
70
71 #[must_use]
73 pub const fn new_with_hooks(
74 store: &'static LocalKey<StoreRegistry>,
75 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
76 ) -> Self {
77 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
78 }
79
80 #[must_use]
82 pub const fn debug(mut self) -> Self {
83 self.debug = true;
84 self
85 }
86
87 #[must_use]
89 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
90 self.metrics = Some(sink);
91 self
92 }
93
94 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
95 if let Some(sink) = self.metrics {
96 with_metrics_sink(sink, f)
97 } else {
98 f()
99 }
100 }
101
102 fn execute_save_with<E, T, R>(
104 &self,
105 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
106 map: impl FnOnce(T) -> R,
107 ) -> Result<R, InternalError>
108 where
109 E: EntityKind<Canister = C> + EntityValue,
110 {
111 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
112
113 Ok(map(value))
114 }
115
116 fn execute_save_entity<E>(
118 &self,
119 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
120 ) -> Result<E, InternalError>
121 where
122 E: EntityKind<Canister = C> + EntityValue,
123 {
124 self.execute_save_with(op, std::convert::identity)
125 }
126
127 fn execute_save_batch<E>(
128 &self,
129 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
130 ) -> Result<WriteBatchResponse<E>, InternalError>
131 where
132 E: EntityKind<Canister = C> + EntityValue,
133 {
134 self.execute_save_with(op, WriteBatchResponse::new)
135 }
136
137 fn execute_save_view<E>(
138 &self,
139 op: impl FnOnce(SaveExecutor<E>) -> Result<E::ViewType, InternalError>,
140 ) -> Result<E::ViewType, InternalError>
141 where
142 E: EntityKind<Canister = C> + EntityValue,
143 {
144 self.execute_save_with(op, std::convert::identity)
145 }
146
147 #[must_use]
153 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
154 where
155 E: EntityKind<Canister = C>,
156 {
157 FluentLoadQuery::new(self, Query::new(MissingRowPolicy::Ignore))
158 }
159
160 #[must_use]
162 pub const fn load_with_consistency<E>(
163 &self,
164 consistency: MissingRowPolicy,
165 ) -> FluentLoadQuery<'_, E>
166 where
167 E: EntityKind<Canister = C>,
168 {
169 FluentLoadQuery::new(self, Query::new(consistency))
170 }
171
172 #[must_use]
174 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
175 where
176 E: EntityKind<Canister = C>,
177 {
178 FluentDeleteQuery::new(self, Query::new(MissingRowPolicy::Ignore).delete())
179 }
180
181 #[must_use]
183 pub fn delete_with_consistency<E>(
184 &self,
185 consistency: MissingRowPolicy,
186 ) -> FluentDeleteQuery<'_, E>
187 where
188 E: EntityKind<Canister = C>,
189 {
190 FluentDeleteQuery::new(self, Query::new(consistency).delete())
191 }
192
193 #[must_use]
197 pub const fn select_one(&self) -> Value {
198 Value::Int(1)
199 }
200
201 #[must_use]
208 pub fn show_indexes<E>(&self) -> Vec<String>
209 where
210 E: EntityKind<Canister = C>,
211 {
212 self.show_indexes_for_model(E::MODEL)
213 }
214
215 #[must_use]
217 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
218 show_indexes_for_model(model)
219 }
220
221 #[must_use]
223 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
224 where
225 E: EntityKind<Canister = C>,
226 {
227 self.show_columns_for_model(E::MODEL)
228 }
229
230 #[must_use]
232 pub fn show_columns_for_model(
233 &self,
234 model: &'static EntityModel,
235 ) -> Vec<EntityFieldDescription> {
236 describe_entity_model(model).fields().to_vec()
237 }
238
239 #[must_use]
241 pub fn show_entities(&self) -> Vec<String> {
242 self.db.runtime_entity_names()
243 }
244
245 #[must_use]
250 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
251 where
252 E: EntityKind<Canister = C>,
253 {
254 self.describe_entity_model(E::MODEL)
255 }
256
257 #[must_use]
259 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
260 describe_entity_model(model)
261 }
262
263 pub fn storage_report(
265 &self,
266 name_to_path: &[(&'static str, &'static str)],
267 ) -> Result<StorageReport, InternalError> {
268 self.db.storage_report(name_to_path)
269 }
270
271 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
273 self.db.integrity_report()
274 }
275
276 pub fn execute_migration_plan(
281 &self,
282 plan: &MigrationPlan,
283 max_steps: usize,
284 ) -> Result<MigrationRunOutcome, InternalError> {
285 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
286 }
287
288 #[must_use]
293 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
294 where
295 E: EntityKind<Canister = C> + EntityValue,
296 {
297 LoadExecutor::new(self.db, self.debug)
298 }
299
300 #[must_use]
301 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
302 where
303 E: EntityKind<Canister = C> + EntityValue,
304 {
305 DeleteExecutor::new(self.db, self.debug)
306 }
307
308 #[must_use]
309 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
310 where
311 E: EntityKind<Canister = C> + EntityValue,
312 {
313 SaveExecutor::new(self.db, self.debug)
314 }
315}